1import torch
2from diffusers import FlowMatchEulerDiscreteScheduler
3from diffusers.pipelines.ltx2 import LTX2Pipeline, LTX2LatentUpsamplePipeline
4from diffusers.pipelines.ltx2.latent_upsampler import LTX2LatentUpsamplerModel
5from diffusers.pipelines.ltx2.utils import STAGE_2_DISTILLED_SIGMA_VALUES
6from diffusers.pipelines.ltx2.export_utils import encode_video
7
8device = "cuda:0"
9width = 768
10height = 512
11
12pipe = LTX2Pipeline.from_pretrained(
13 "Lightricks/LTX-2", torch_dtype=torch.bfloat16
14)
15pipe.enable_sequential_cpu_offload(device=device)
16
17prompt = "A beautiful sunset over the ocean"
18negative_prompt = "shaky, glitchy, low quality, worst quality, deformed, distorted, disfigured, motion smear, motion artifacts, fused fingers, bad anatomy, weird hand, ugly, transition, static."
19
20# Stage 1 default (non-distilled) inference
21frame_rate = 24.0
22video_latent, audio_latent = pipe(
23 prompt=prompt,
24 negative_prompt=negative_prompt,
25 width=width,
26 height=height,
27 num_frames=121,
28 frame_rate=frame_rate,
29 num_inference_steps=40,
30 sigmas=None,
31 guidance_scale=4.0,
32 output_type="latent",
33 return_dict=False,
34)
35
36latent_upsampler = LTX2LatentUpsamplerModel.from_pretrained(
37 "Lightricks/LTX-2",
38 subfolder="latent_upsampler",
39 torch_dtype=torch.bfloat16,
40)
41upsample_pipe = LTX2LatentUpsamplePipeline(vae=pipe.vae, latent_upsampler=latent_upsampler)
42upsample_pipe.enable_model_cpu_offload(device=device)
43upscaled_video_latent = upsample_pipe(
44 latents=video_latent,
45 output_type="latent",
46 return_dict=False,
47)[0]
48
49# Load Stage 2 distilled LoRA
50pipe.load_lora_weights(
51 "Lightricks/LTX-2", adapter_name="stage_2_distilled", weight_name="ltx-2-19b-distilled-lora-384.safetensors"
52)
53pipe.set_adapters("stage_2_distilled", 1.0)
54# VAE tiling is usually necessary to avoid OOM error when VAE decoding
55pipe.vae.enable_tiling()
56# Change scheduler to use Stage 2 distilled sigmas as is
57new_scheduler = FlowMatchEulerDiscreteScheduler.from_config(
58 pipe.scheduler.config, use_dynamic_shifting=False, shift_terminal=None
59)
60pipe.scheduler = new_scheduler
61# Stage 2 inference with distilled LoRA and sigmas
62video, audio = pipe(
63 latents=upscaled_video_latent,
64 audio_latents=audio_latent,
65 prompt=prompt,
66 negative_prompt=negative_prompt,
67 num_inference_steps=3,
68 noise_scale=STAGE_2_DISTILLED_SIGMA_VALUES[0], # renoise with first sigma value https://github.com/Lightricks/LTX-2/blob/main/packages/ltx-pipelines/src/ltx_pipelines/ti2vid_two_stages.py#L218
69 sigmas=STAGE_2_DISTILLED_SIGMA_VALUES,
70 guidance_scale=1.0,
71 output_type="np",
72 return_dict=False,
73)
74
75encode_video(
76 video[0],
77 fps=frame_rate,
78 audio=audio[0].float().cpu(),
79 audio_sample_rate=pipe.vocoder.config.output_sampling_rate,
80 output_path="ltx2_lora_distilled_sample.mp4",
81)