Add LTX 2.0 Video Pipelines

scripts/ltx2_test_full_pipeline.py

@@ -0,0 +1,108 @@
+import argparse
+import os
+
+import torch
+
+from diffusers import LTX2Pipeline
+from diffusers.pipelines.ltx2.export_utils import encode_video
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument("--model_id", type=str, default="diffusers-internal-dev/new-ltx-model")
+    parser.add_argument("--revision", type=str, default="main")
+
+    parser.add_argument(
+        "--prompt",
+        type=str,
+        default="A video of a dog dancing to energetic electronic dance music",
+    )
+    parser.add_argument(
+        "--negative_prompt",
+        type=str,
+        default=(
+            "blurry, out of focus, overexposed, underexposed, low contrast, washed out colors, excessive noise, "
+            "grainy texture, poor lighting, flickering, motion blur, distorted proportions, unnatural skin tones, "
+            "deformed facial features, asymmetrical face, missing facial features, extra limbs, disfigured hands, "
+            "wrong hand count, artifacts around text, inconsistent perspective, camera shake, incorrect depth of "
+            "field, background too sharp, background clutter, distracting reflections, harsh shadows, inconsistent "
+            "lighting direction, color banding, cartoonish rendering, 3D CGI look, unrealistic materials, uncanny "
+            "valley effect, incorrect ethnicity, wrong gender, exaggerated expressions, wrong gaze direction, "
+            "mismatched lip sync, silent or muted audio, distorted voice, robotic voice, echo, background noise, "
+            "off-sync audio,incorrect dialogue, added dialogue, repetitive speech, jittery movement, awkward "
+            "pauses, incorrect timing, unnatural transitions, inconsistent framing, tilted camera, flat lighting, "
+            "inconsistent tone, cinematic oversaturation, stylized filters, or AI artifacts."
+        ),
+    )
+
+    parser.add_argument("--num_inference_steps", type=int, default=40)
+    parser.add_argument("--height", type=int, default=512)
+    parser.add_argument("--width", type=int, default=768)
+    parser.add_argument("--num_frames", type=int, default=121)
+    parser.add_argument("--frame_rate", type=float, default=25.0)
+    parser.add_argument("--guidance_scale", type=float, default=3.0)
+    parser.add_argument("--seed", type=int, default=42)
+
+    parser.add_argument("--device", type=str, default="cuda:0")
+    parser.add_argument("--dtype", type=str, default="bf16")
+    parser.add_argument("--cpu_offload", action="store_true")
+
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        default="/home/daniel_gu/samples",
+        help="Output directory for generated video",
+    )
+    parser.add_argument(
+        "--output_filename",
+        type=str,
+        default="ltx2_sample_video.mp4",
+        help="Filename of the exported generated video",
+    )
+
+    args = parser.parse_args()
+    args.dtype = torch.bfloat16 if args.dtype == "bf16" else torch.float32
+    return args
+
+
+def main(args):
+    pipeline = LTX2Pipeline.from_pretrained(
+        args.model_id,
+        revision=args.revision,
+        torch_dtype=args.dtype,
+    )
+    pipeline.to(device=args.device)
+    if args.cpu_offload:
+        pipeline.enable_model_cpu_offload()
+
+    video, audio = pipeline(
+        prompt=args.prompt,
+        negative_prompt=args.negative_prompt,
+        height=args.height,
+        width=args.width,
+        num_frames=args.num_frames,
+        frame_rate=args.frame_rate,
+        num_inference_steps=args.num_inference_steps,
+        guidance_scale=args.guidance_scale,
+        generator=torch.Generator(device=args.device).manual_seed(args.seed),
+        output_type="np",
+        return_dict=False,
+    )
+
+    # Convert video to uint8 (but keep as NumPy array)
+    video = (video * 255).round().astype("uint8")
+    video = torch.from_numpy(video)
+
+    encode_video(
+        video[0],
+        fps=args.frame_rate,
+        audio=audio[0].float().cpu(),
+        audio_sample_rate=pipeline.vocoder.config.output_sampling_rate,  # should be 24000
+        output_path=os.path.join(args.output_dir, args.output_filename),
+    )
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    main(args)

scripts/ltx2_test_full_pipeline_i2v.py

@@ -0,0 +1,102 @@
+import argparse
+import os
+
+import torch
+
+from diffusers.pipelines.ltx2 import LTX2ImageToVideoPipeline
+from diffusers.pipelines.ltx2.export_utils import encode_video
+from diffusers.utils import load_image
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument("--model_id", type=str, default="diffusers-internal-dev/new-ltx-model")
+    parser.add_argument("--revision", type=str, default="main")
+
+    parser.add_argument("--image_path", required=True, type=str)
+    parser.add_argument(
+        "--prompt",
+        type=str,
+        default="An astronaut hatches from a fragile egg on the surface of the Moon, the shell cracking and peeling apart in gentle low-gravity motion. Fine lunar dust lifts and drifts outward with each movement, floating in slow arcs before settling back onto the ground. The astronaut pushes free in a deliberate, weightless motion, small fragments of the egg tumbling and spinning through the air. In the background, the deep darkness of space subtly shifts as stars glide with the camera's movement, emphasizing vast depth and scale. The camera performs a smooth, cinematic slow push-in, with natural parallax between the foreground dust, the astronaut, and the distant starfield. Ultra-realistic detail, physically accurate low-gravity motion, cinematic lighting, and a breath-taking, movie-like shot.",
+    )
+    parser.add_argument(
+        "--negative_prompt",
+        type=str,
+        default="shaky, glitchy, low quality, worst quality, deformed, distorted, disfigured, motion smear, motion artifacts, fused fingers, bad anatomy, weird hand, ugly, transition, static.",
+    )
+
+    parser.add_argument("--num_inference_steps", type=int, default=40)
+    parser.add_argument("--height", type=int, default=512)
+    parser.add_argument("--width", type=int, default=768)
+    parser.add_argument("--num_frames", type=int, default=121)
+    parser.add_argument("--frame_rate", type=float, default=25.0)
+    parser.add_argument("--guidance_scale", type=float, default=3.0)
+    parser.add_argument("--seed", type=int, default=42)
+
+    parser.add_argument("--device", type=str, default="cuda:0")
+    parser.add_argument("--dtype", type=str, default="bf16")
+    parser.add_argument("--cpu_offload", action="store_true")
+
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        default="samples",
+        help="Output directory for generated video",
+    )
+    parser.add_argument(
+        "--output_filename",
+        type=str,
+        default="ltx2_sample_video.mp4",
+        help="Filename of the exported generated video",
+    )
+
+    args = parser.parse_args()
+    args.dtype = torch.bfloat16 if args.dtype == "bf16" else torch.float32
+    return args
+
+
+def main(args):
+    pipeline = LTX2ImageToVideoPipeline.from_pretrained(
+        args.model_id,
+        revision=args.revision,
+        torch_dtype=args.dtype,
+    )
+    if args.cpu_offload:
+        pipeline.enable_model_cpu_offload()
+    else:
+        pipeline.to(device=args.device)
+
+    image = load_image(args.image_path)
+
+    video, audio = pipeline(
+        image=image,
+        prompt=args.prompt,
+        negative_prompt=args.negative_prompt,
+        height=args.height,
+        width=args.width,
+        num_frames=args.num_frames,
+        frame_rate=args.frame_rate,
+        num_inference_steps=args.num_inference_steps,
+        guidance_scale=args.guidance_scale,
+        generator=torch.Generator(device=args.device).manual_seed(args.seed),
+        output_type="np",
+        return_dict=False,
+    )
+
+    # Convert video to uint8 (but keep as NumPy array)
+    video = (video * 255).round().astype("uint8")
+    video = torch.from_numpy(video)
+
+    encode_video(
+        video[0],
+        fps=args.frame_rate,
+        audio=audio[0].float().cpu(),
+        audio_sample_rate=pipeline.vocoder.config.output_sampling_rate,  # should be 24000
+        output_path=os.path.join(args.output_dir, args.output_filename),
+    )
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    main(args)

scripts/test_ltx2_audio_conversion.py

@@ -0,0 +1,119 @@
+import argparse
+from pathlib import Path
+
+import torch
+from huggingface_hub import hf_hub_download
+
+
+def download_checkpoint(
+    repo_id="diffusers-internal-dev/new-ltx-model",
+    filename="ltx-av-step-1932500-interleaved-new-vae.safetensors",
+):
+    ckpt_path = hf_hub_download(repo_id=repo_id, filename=filename)
+    return ckpt_path
+
+
+def convert_state_dict(state_dict: dict) -> dict:
+    converted = {}
+    for key, value in state_dict.items():
+        if not isinstance(value, torch.Tensor):
+            continue
+        new_key = key
+        if new_key.startswith("decoder."):
+            new_key = new_key[len("decoder.") :]
+        converted[f"decoder.{new_key}"] = value
+
+    converted["latents_mean"] = converted.pop("decoder.per_channel_statistics.mean-of-means")
+    converted["latents_std"] = converted.pop("decoder.per_channel_statistics.std-of-means")
+    return converted
+
+
+def load_original_decoder(device: torch.device, dtype: torch.dtype):
+    from ltx_core.loader.single_gpu_model_builder import SingleGPUModelBuilder as Builder
+    from ltx_core.model.audio_vae.model_configurator import AUDIO_VAE_DECODER_COMFY_KEYS_FILTER
+    from ltx_core.model.audio_vae.model_configurator import VAEDecoderConfigurator as AudioDecoderConfigurator
+
+    checkpoint_path = download_checkpoint()
+
+    # The code below comes from `ltx-pipelines/src/ltx_pipelines/txt2vid.py`
+    decoder = Builder(
+        model_path=checkpoint_path,
+        model_class_configurator=AudioDecoderConfigurator,
+        model_sd_key_ops=AUDIO_VAE_DECODER_COMFY_KEYS_FILTER,
+    ).build(device=device)
+
+    decoder.eval()
+    return decoder
+
+
+def build_diffusers_decoder():
+    from diffusers.models.autoencoders import AutoencoderKLLTX2Audio
+
+    with torch.device("meta"):
+        model = AutoencoderKLLTX2Audio()
+
+    model.eval()
+    return model
+
+
+@torch.no_grad()
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Validate LTX2 audio decoder conversion.")
+    parser.add_argument("--device", type=str, default="cpu")
+    parser.add_argument("--dtype", type=str, default="bfloat16", choices=["float32", "bfloat16", "float16"])
+    parser.add_argument("--batch", type=int, default=2)
+    parser.add_argument("--output-path", type=Path, required=True)
+    args = parser.parse_args()
+
+    device = torch.device(args.device)
+    dtype_map = {"float32": torch.float32, "bfloat16": torch.bfloat16, "float16": torch.float16}
+    dtype = dtype_map[args.dtype]
+
+    original_decoder = load_original_decoder(device, dtype)
+    diffusers_model = build_diffusers_decoder()
+
+    converted_state_dict = convert_state_dict(original_decoder.state_dict())
+    diffusers_model.load_state_dict(converted_state_dict, assign=True, strict=False)
+
+    per_channel_len = original_decoder.per_channel_statistics.get_buffer("std-of-means").numel()
+    latent_channels = diffusers_model.decoder.latent_channels
+    mel_bins_for_match = per_channel_len // latent_channels if per_channel_len % latent_channels == 0 else None
+
+    levels = len(diffusers_model.decoder.channel_multipliers)
+    latent_height = diffusers_model.decoder.resolution // (2 ** (levels - 1))
+    latent_width = mel_bins_for_match or latent_height
+
+    dummy = torch.randn(
+        args.batch,
+        diffusers_model.decoder.latent_channels,
+        latent_height,
+        latent_width,
+        device=device,
+        dtype=dtype,
+        generator=torch.Generator(device).manual_seed(42),
+    )
+
+    original_out = original_decoder(dummy)
+
+    from diffusers.pipelines.ltx2.pipeline_ltx2 import LTX2Pipeline
+
+    _, a_channels, a_time, a_freq = dummy.shape
+    dummy = dummy.permute(0, 2, 1, 3).reshape(-1, a_time, a_channels * a_freq)
+    dummy = LTX2Pipeline._denormalize_audio_latents(
+        dummy,
+        diffusers_model.latents_mean,
+        diffusers_model.latents_std,
+    )
+    dummy = dummy.view(-1, a_time, a_channels, a_freq).permute(0, 2, 1, 3)
+    diffusers_out = diffusers_model.decode(dummy).sample
+
+    torch.testing.assert_close(diffusers_out, original_out, rtol=1e-4, atol=1e-4)
+    max_diff = (diffusers_out - original_out).abs().max().item()
+    print(f"Conversion successful. Max diff: {max_diff:.6f}")
+
+    diffusers_model.to(dtype).save_pretrained(args.output_path)
+    print(f"Serialized model to {args.output_path}")
+
+
+if __name__ == "__main__":
+    main()

src/diffusers/__init__.py

@@ -193,6 +193,8 @@
             "AutoencoderKLHunyuanImageRefiner",
             "AutoencoderKLHunyuanVideo",
             "AutoencoderKLHunyuanVideo15",
+            "AutoencoderKLLTX2Audio",
+            "AutoencoderKLLTX2Video",
             "AutoencoderKLLTXVideo",
             "AutoencoderKLMagvit",
             "AutoencoderKLMochi",
@@ -236,6 +238,7 @@
             "Kandinsky5Transformer3DModel",
             "LatteTransformer3DModel",
             "LongCatImageTransformer2DModel",
+            "LTX2VideoTransformer3DModel",
             "LTXVideoTransformer3DModel",
             "Lumina2Transformer2DModel",
             "LuminaNextDiT2DModel",
@@ -537,6 +540,8 @@
             "LEditsPPPipelineStableDiffusionXL",
             "LongCatImageEditPipeline",
             "LongCatImagePipeline",
+            "LTX2ImageToVideoPipeline",
+            "LTX2Pipeline",
             "LTXConditionPipeline",
             "LTXImageToVideoPipeline",
             "LTXLatentUpsamplePipeline",
@@ -937,6 +942,8 @@
             AutoencoderKLHunyuanImageRefiner,
             AutoencoderKLHunyuanVideo,
             AutoencoderKLHunyuanVideo15,
+            AutoencoderKLLTX2Audio,
+            AutoencoderKLLTX2Video,
             AutoencoderKLLTXVideo,
             AutoencoderKLMagvit,
             AutoencoderKLMochi,
@@ -980,6 +987,7 @@
             Kandinsky5Transformer3DModel,
             LatteTransformer3DModel,
             LongCatImageTransformer2DModel,
+            LTX2VideoTransformer3DModel,
             LTXVideoTransformer3DModel,
             Lumina2Transformer2DModel,
             LuminaNextDiT2DModel,
@@ -1251,6 +1259,8 @@
             LEditsPPPipelineStableDiffusionXL,
             LongCatImageEditPipeline,
             LongCatImagePipeline,
+            LTX2ImageToVideoPipeline,
+            LTX2Pipeline,
             LTXConditionPipeline,
             LTXImageToVideoPipeline,
             LTXLatentUpsamplePipeline,