import gc from functools import wraps import fire import torch from diffusers import FlowMatchEulerDiscreteScheduler, SanaVideoPipeline from diffusers.pipelines.ltx2 import LTX2LatentUpsamplePipeline, LTX2Pipeline from diffusers.pipelines.ltx2.export_utils import encode_video from diffusers.pipelines.ltx2.latent_upsampler import LTX2LatentUpsamplerModel from diffusers.pipelines.ltx2.utils import STAGE_2_DISTILLED_SIGMA_VALUES from diffusers.utils import export_to_video def cli(func): wraps(func) def wrapper(*args, **kwargs): return fire.Fire(func) return wrapper @cli def sana_video_ltx2_refine( prompt: str = "A cat and a dog baking a cake together in a kitchen. The cat is carefully measuring flour, while the dog is stirring the batter with a wooden spoon. The kitchen is cozy, with sunlight streaming through the window.", negative_prompt: str = "A chaotic sequence with misshapen, deformed limbs in heavy motion blur, sudden disappearance, jump cuts, jerky movements, rapid shot changes, frames out of sync, inconsistent character shapes, temporal artifacts, jitter, and ghosting effects, creating a disorienting visual experience.", sana_model_id: str = "Efficient-Large-Model/SANA-Video_2B_720p_diffusers", ltx2_model_id: str = "Lightricks/LTX-2", sana_height: int = 704, sana_width: int = 1280, sana_frames: int = 81, motion_score: int = 30, sana_guidance_scale: float = 6.0, sana_num_steps: int = 50, frame_rate: float = 16.0, seed: int = 42, output_path: str = "sana_ltx2_refined.mp4", save_sana_output: str = "sana_original.mp4", skip_audio: bool = True, skip_upsampler: bool = True, ): """ Use SanaVideoPipeline to generate video latent, then use LTX2 Pipeline for Stage-2 refinement. Key technical points: 1. Manually pack video latent to skip diffusers' normalize step (consistent with original LTX-2 code) 2. Create audio latent such that it becomes zero after normalize (consistent with zero audio latent in original code) """ device = "cuda" dtype = torch.bfloat16 # Step 1: Generate latent using Sana Video Pipeline sana_pipe = SanaVideoPipeline.from_pretrained(sana_model_id, torch_dtype=dtype) sana_pipe.text_encoder.to(dtype) sana_pipe.enable_model_cpu_offload(gpu_id=0) full_prompt = prompt + f" motion score: {motion_score}." generator = torch.Generator(device=device).manual_seed(seed) sana_output = sana_pipe( prompt=full_prompt, negative_prompt=negative_prompt, height=sana_height, width=sana_width, frames=sana_frames, guidance_scale=sana_guidance_scale, num_inference_steps=sana_num_steps, generator=generator, output_type="latent", return_dict=True, ) # Extract video latent video_latent = None for key in ("latents", "video_latents", "latent", "dit_latents", "dit_latent", "frames", "videos"): if hasattr(sana_output, key): video_latent = getattr(sana_output, key) if video_latent is not None: break if video_latent is None: raise ValueError("Failed to extract latents from Sana output.") if isinstance(video_latent, (list, tuple)): video_latent = video_latent[0] if video_latent.dim() == 4: video_latent = video_latent.unsqueeze(0) del sana_pipe, sana_output gc.collect() torch.cuda.empty_cache() # Step 2: Load LTX2 Pipeline ltx2_pipe = LTX2Pipeline.from_pretrained(ltx2_model_id, torch_dtype=dtype) # Save Sana original output if save_sana_output: ltx2_pipe.vae.to(device) ltx2_pipe.vae.enable_tiling( tile_sample_min_height=512, tile_sample_min_width=512, tile_sample_stride_height=448, tile_sample_stride_width=448, ) sana_latent_for_decode = video_latent.to(device=device, dtype=ltx2_pipe.vae.dtype) latents_mean = ltx2_pipe.vae.latents_mean.view(1, -1, 1, 1, 1).to( sana_latent_for_decode.device, sana_latent_for_decode.dtype ) latents_std = ltx2_pipe.vae.latents_std.view(1, -1, 1, 1, 1).to( sana_latent_for_decode.device, sana_latent_for_decode.dtype ) sana_latent_denorm = sana_latent_for_decode * latents_std / ltx2_pipe.vae.config.scaling_factor + latents_mean timestep = ( None if not ltx2_pipe.vae.config.timestep_conditioning else torch.tensor([0.0], device=device, dtype=sana_latent_denorm.dtype) ) with torch.no_grad(): sana_decoded = ltx2_pipe.vae.decode(sana_latent_denorm, timestep, return_dict=False)[0] del sana_latent_for_decode, sana_latent_denorm sana_video = ltx2_pipe.video_processor.postprocess_video(sana_decoded, output_type="pil") export_to_video(sana_video[0], save_sana_output, fps=int(frame_rate)) del sana_decoded, sana_video ltx2_pipe.vae.to("cpu") gc.collect() torch.cuda.empty_cache() ltx2_pipe.enable_model_cpu_offload(gpu_id=0) # Latent Upsampler (optional) if skip_upsampler: print("Skipping latent upsampler") upscaled_video_latent = video_latent.to(device=device, dtype=dtype) else: latent_upsampler = LTX2LatentUpsamplerModel.from_pretrained( ltx2_model_id, subfolder="latent_upsampler", torch_dtype=dtype ) upsample_pipe = LTX2LatentUpsamplePipeline(vae=ltx2_pipe.vae, latent_upsampler=latent_upsampler) upsample_pipe.enable_model_cpu_offload(device=device) upscaled_video_latent = upsample_pipe( latents=video_latent.to(device=device, dtype=dtype), latents_normalized=True, height=sana_height, width=sana_width, num_frames=sana_frames, output_type="latent", return_dict=False, )[0] latents_mean = ltx2_pipe.vae.latents_mean.view(1, -1, 1, 1, 1).to( upscaled_video_latent.device, upscaled_video_latent.dtype ) latents_std = ltx2_pipe.vae.latents_std.view(1, -1, 1, 1, 1).to( upscaled_video_latent.device, upscaled_video_latent.dtype ) scaling_factor = ltx2_pipe.vae.config.scaling_factor upscaled_video_latent = (upscaled_video_latent - latents_mean) * scaling_factor / latents_std del latent_upsampler, upsample_pipe gc.collect() torch.cuda.empty_cache() # Step 3: Manually pack video latent (skip diffusers' normalize, consistent with original code) packed_video_latent = LTX2Pipeline._pack_latents( upscaled_video_latent, patch_size=ltx2_pipe.transformer_spatial_patch_size, patch_size_t=ltx2_pipe.transformer_temporal_patch_size, ) _, _, latent_num_frames, latent_height, latent_width = upscaled_video_latent.shape pixel_height = latent_height * ltx2_pipe.vae_spatial_compression_ratio pixel_width = latent_width * ltx2_pipe.vae_spatial_compression_ratio pixel_num_frames = (latent_num_frames - 1) * ltx2_pipe.vae_temporal_compression_ratio + 1 # Step 4: Create audio latent (becomes zero after normalize, consistent with original code) num_frames_pixel = (latent_num_frames - 1) * 8 + 1 duration_s = num_frames_pixel / frame_rate audio_num_frames = round(duration_s * 16000 / 160 / 4) num_channels_audio, latent_mel_bins = 8, 16 audio_latents_mean = ltx2_pipe.audio_vae.latents_mean packed_audio_latent = ( audio_latents_mean.unsqueeze(0) .unsqueeze(0) .expand(1, audio_num_frames, num_channels_audio * latent_mel_bins) .to(dtype=dtype, device=device) .contiguous() ) audio_latent = ( packed_audio_latent.unflatten(2, (num_channels_audio, latent_mel_bins)).permute(0, 2, 1, 3).contiguous() ) del video_latent, upscaled_video_latent gc.collect() torch.cuda.empty_cache() # Step 5: LTX2 Stage-2 Refinement ltx2_pipe.load_lora_weights( ltx2_model_id, adapter_name="stage_2_distilled", weight_name="ltx-2-19b-distilled-lora-384.safetensors" ) ltx2_pipe.set_adapters("stage_2_distilled", 1.0) ltx2_pipe.vae.enable_tiling() ltx2_pipe.scheduler = FlowMatchEulerDiscreteScheduler.from_config( ltx2_pipe.scheduler.config, use_dynamic_shifting=False, shift_terminal=None, ) generator = torch.Generator(device=device).manual_seed(seed) video, audio = ltx2_pipe( latents=packed_video_latent, audio_latents=audio_latent, prompt=prompt, negative_prompt=negative_prompt, height=pixel_height, width=pixel_width, num_frames=pixel_num_frames, num_inference_steps=3, noise_scale=STAGE_2_DISTILLED_SIGMA_VALUES[0], sigmas=STAGE_2_DISTILLED_SIGMA_VALUES, guidance_scale=1.0, frame_rate=frame_rate, generator=generator, output_type="np", return_dict=False, ) # Step 6: Save output video = (video * 255).round().astype("uint8") video = torch.from_numpy(video) encode_video( video[0], fps=frame_rate, audio=None if skip_audio else audio[0].float().cpu(), audio_sample_rate=None if skip_audio else ltx2_pipe.vocoder.config.output_sampling_rate, output_path=output_path, ) print(f"Done! Output saved to: {output_path}")