From b076cd934e80058da324455def30df5b667647df Mon Sep 17 00:00:00 2001 From: Will Lin Date: Sun, 25 Jan 2026 16:57:40 -0800 Subject: [PATCH] test --- .../inference/basic/basic_ltx2_upscale.py | 53 +++ .../convert_ltx2_upsampler_weights.py | 112 ++++++ .../pipelines/test_ltx2_scheduler_parity.py | 36 ++ .../pipelines/test_ltx2_two_stage_parity.py | 379 ++++++++++++++++++ .../transformers/test_ltx2_lora_parity.py | 228 +++++++++++ tests/local_tests/upsamplers/__init__.py | 1 + .../upsamplers/test_ltx2_upsampler.py | 104 +++++ 7 files changed, 913 insertions(+) create mode 100644 examples/inference/basic/basic_ltx2_upscale.py create mode 100644 scripts/checkpoint_conversion/convert_ltx2_upsampler_weights.py create mode 100644 tests/local_tests/pipelines/test_ltx2_scheduler_parity.py create mode 100644 tests/local_tests/pipelines/test_ltx2_two_stage_parity.py create mode 100644 tests/local_tests/transformers/test_ltx2_lora_parity.py create mode 100644 tests/local_tests/upsamplers/__init__.py create mode 100644 tests/local_tests/upsamplers/test_ltx2_upsampler.py diff --git a/examples/inference/basic/basic_ltx2_upscale.py b/examples/inference/basic/basic_ltx2_upscale.py new file mode 100644 index 000000000..7c442e0f9 --- /dev/null +++ b/examples/inference/basic/basic_ltx2_upscale.py @@ -0,0 +1,53 @@ +# SPDX-License-Identifier: Apache-2.0 +from fastvideo import VideoGenerator + +PROMPT = ( + "A warm sunny backyard. The camera starts in a tight cinematic close-up " + "of a woman and a man in their 30s, facing each other with serious " + "expressions. The woman, emotional and dramatic, says softly, \"That's " + "it... Dad's lost it. And we've lost Dad.\" The man exhales, slightly " + "annoyed: \"Stop being so dramatic, Jess.\" A beat. He glances aside, " + "then mutters defensively, \"He's just having fun.\" The camera slowly " + "pans right, revealing the grandfather in the garden wearing enormous " + "butterfly wings, waving his arms in the air like he's trying to take " + "off. He shouts, \"Wheeeew!\" as he flaps his wings with full commitment. " + "The woman covers her face, on the verge of tears. The tone is deadpan, " + "absurd, and quietly tragic." +) + +# FastVideo-converted weights (align settings to the official distilled pipeline). +MODEL_PATH = "converted/ltx2_diffusers" +UPSAMPLER_PATH = "converted/ltx2_spatial_upscaler" + +OUTPUT_PATH = "outputs_video/ltx2_upscale/ltx2_official_two_stage2222.mp4" + + +def main() -> None: + generator = VideoGenerator.from_pretrained( + MODEL_PATH, + num_gpus=1, + ltx2_refine_enabled=True, + ltx2_refine_upsampler_path=UPSAMPLER_PATH, + ltx2_refine_num_inference_steps=3, + ltx2_refine_guidance_scale=1.0, + ltx2_refine_add_noise=True, + ) + + generator.generate_video( + prompt=PROMPT, + output_path=OUTPUT_PATH, + height=960, + width=1664, + num_frames=81, + fps=24, + seed=10, + # Distilled schedule in LTX-2 uses 8 steps; setting this keeps parity. + num_inference_steps=8, + guidance_scale=1.0, + save_video=True, + ) + generator.shutdown() + + +if __name__ == "__main__": + main() diff --git a/scripts/checkpoint_conversion/convert_ltx2_upsampler_weights.py b/scripts/checkpoint_conversion/convert_ltx2_upsampler_weights.py new file mode 100644 index 000000000..0090e6fe3 --- /dev/null +++ b/scripts/checkpoint_conversion/convert_ltx2_upsampler_weights.py @@ -0,0 +1,112 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +Convert LTX-2 latent upsampler weights to FastVideo format. +""" + +from __future__ import annotations + +import argparse +import glob +import json +from pathlib import Path + +from safetensors import safe_open +from safetensors.torch import load_file, save_file + + +def _find_shards(model_path: Path) -> list[Path]: + if model_path.is_file(): + return [model_path] + + index_files = list(model_path.glob("*.safetensors.index.json")) + if index_files: + with index_files[0].open("r", encoding="utf-8") as f: + index = json.load(f) + return sorted({model_path / shard for shard in index["weight_map"].values()}) + return sorted(Path(p) for p in glob.glob(str(model_path / "*.safetensors"))) + + +def _read_metadata_config(path: Path) -> dict: + with safe_open(str(path), framework="pt") as f: + metadata = f.metadata() + if not metadata or "config" not in metadata: + raise ValueError(f"Missing config metadata in {path}") + return json.loads(metadata["config"]) + + +def _load_weights(shards: list[Path]) -> dict: + weights: dict = {} + for shard in shards: + weights.update(load_file(str(shard))) + return weights + + +def _map_keys(weights: dict, add_model_prefix: bool) -> dict: + remapped = {} + for key, value in weights.items(): + new_key = key + if not add_model_prefix and new_key.startswith("model."): + new_key = new_key[len("model.") :] + if add_model_prefix and not new_key.startswith("model."): + new_key = f"model.{new_key}" + remapped[new_key] = value + return remapped + + +def main() -> None: + parser = argparse.ArgumentParser( + description="Convert LTX-2 spatial upsampler weights to FastVideo format" + ) + parser.add_argument( + "--source", + type=str, + required=True, + help="Path to the official LTX-2 upsampler safetensors file or directory", + ) + parser.add_argument( + "--output", + type=str, + required=True, + help="Output directory for converted weights", + ) + parser.add_argument( + "--class-name", + type=str, + default="LTX2LatentUpsampler", + help="_class_name to write into config.json", + ) + parser.add_argument( + "--add-model-prefix", + action="store_true", + help="Prefix all weights with 'model.' for loading into wrapper modules", + ) + + args = parser.parse_args() + + source_path = Path(args.source) + shards = _find_shards(source_path) + if not shards: + raise FileNotFoundError(f"No safetensors found in {source_path}") + + config = _read_metadata_config(shards[0]) + config["_class_name"] = args.class_name + + weights = _load_weights(shards) + weights = _map_keys(weights, add_model_prefix=args.add_model_prefix) + + output_dir = Path(args.output) + output_dir.mkdir(parents=True, exist_ok=True) + + output_weights = output_dir / "model.safetensors" + save_file(weights, str(output_weights)) + print(f"Saved weights to {output_weights}") + + config_path = output_dir / "config.json" + with config_path.open("w", encoding="utf-8") as f: + json.dump(config, f, indent=2) + f.write("\n") + print(f"Saved config to {config_path}") + + +if __name__ == "__main__": + main() diff --git a/tests/local_tests/pipelines/test_ltx2_scheduler_parity.py b/tests/local_tests/pipelines/test_ltx2_scheduler_parity.py new file mode 100644 index 000000000..5935a09c8 --- /dev/null +++ b/tests/local_tests/pipelines/test_ltx2_scheduler_parity.py @@ -0,0 +1,36 @@ +# SPDX-License-Identifier: Apache-2.0 +import sys +from pathlib import Path + +import pytest +import torch +from torch.testing import assert_close + +repo_root = Path(__file__).resolve().parents[3] +if str(repo_root) not in sys.path: + sys.path.insert(0, str(repo_root)) +ltx_core_path = repo_root / "LTX-2" / "packages" / "ltx-core" / "src" +if ltx_core_path.exists() and str(ltx_core_path) not in sys.path: + sys.path.insert(0, str(ltx_core_path)) + +from fastvideo.pipelines.stages.ltx2_denoising import _ltx2_sigmas + + +def test_ltx2_sigma_schedule_parity(): + try: + from ltx_core.components.schedulers import LTX2Scheduler + except ImportError as exc: + pytest.skip(f"LTX-2 import failed: {exc}") + + scheduler = LTX2Scheduler() + device = torch.device("cpu") + + for steps in (4, 8, 40): + ref = scheduler.execute(steps=steps, latent=None).to(torch.float32) + ours = _ltx2_sigmas(steps=steps, latent=None, device=device) + assert_close(ref, ours, atol=1e-6, rtol=1e-6) + + latent = torch.randn(1, 128, 4, 8, 8) + ref = scheduler.execute(steps=10, latent=latent).to(torch.float32) + ours = _ltx2_sigmas(steps=10, latent=latent, device=device) + assert_close(ref, ours, atol=1e-6, rtol=1e-6) diff --git a/tests/local_tests/pipelines/test_ltx2_two_stage_parity.py b/tests/local_tests/pipelines/test_ltx2_two_stage_parity.py new file mode 100644 index 000000000..b334a1878 --- /dev/null +++ b/tests/local_tests/pipelines/test_ltx2_two_stage_parity.py @@ -0,0 +1,379 @@ +# SPDX-License-Identifier: Apache-2.0 +import os +from pathlib import Path +import sys +import tempfile + +import pytest +import torch +from torch.testing import assert_close + +repo_root = Path(__file__).resolve().parents[3] +if str(repo_root) not in sys.path: + sys.path.insert(0, str(repo_root)) + +from fastvideo import VideoGenerator + + +def _log_tensor_stats(label: str, tensor: torch.Tensor) -> None: + tensor_f32 = tensor.float() + print( + f"[LTX2 TWO STAGE] {label}: shape={tuple(tensor.shape)} " + f"dtype={tensor.dtype} device={tensor.device} " + f"min={tensor_f32.min().item():.6f} max={tensor_f32.max().item():.6f} " + f"mean={tensor_f32.mean().item():.6f} sum={tensor_f32.sum().item():.6f}" + ) + + +@pytest.mark.skipif( + not torch.cuda.is_available(), + reason="LTX-2 two-stage parity test requires CUDA.", +) +def test_ltx2_two_stage_parity(): + repo_root = Path(__file__).resolve().parents[3] + if str(repo_root) not in sys.path: + sys.path.insert(0, str(repo_root)) + + ltx_core_path = repo_root / "LTX-2" / "packages" / "ltx-core" / "src" + if ltx_core_path.exists() and str(ltx_core_path) not in sys.path: + sys.path.insert(0, str(ltx_core_path)) + ltx_pipelines_path = repo_root / "LTX-2" / "packages" / "ltx-pipelines" / "src" + if ltx_pipelines_path.exists() and str(ltx_pipelines_path) not in sys.path: + sys.path.insert(0, str(ltx_pipelines_path)) + + os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "TORCH_SDPA") + os.environ.setdefault("LTX2_REFERENCE_ATTN", "pytorch") + torch.backends.cuda.enable_flash_sdp(False) + torch.backends.cuda.enable_mem_efficient_sdp(False) + torch.backends.cuda.enable_math_sdp(True) + + diffusers_path = os.getenv("LTX2_DIFFUSERS_PATH", "converted/ltx2_diffusers") + gemma_model_path = os.path.join(diffusers_path, "text_encoder", "gemma") + official_path = os.getenv( + "LTX2_OFFICIAL_PATH", + "official_ltx_weights/ltx-2-19b-distilled.safetensors", + ) + upsampler_official_path = os.getenv( + "LTX2_UPSAMPLER_OFFICIAL_PATH", + "official_ltx_weights/ltx-2-spatial-upscaler-x2-1.0.safetensors", + ) + upsampler_path = os.getenv( + "LTX2_UPSAMPLER_PATH", + "converted/ltx2_spatial_upscaler", + ) + lora_path = os.getenv( + "LTX2_REFINER_LORA_PATH", + "official_ltx_weights/ltx-2-19b-distilled-lora-384.safetensors", + ) + + if not os.path.isdir(diffusers_path): + pytest.skip(f"Missing LTX-2 diffusers repo at {diffusers_path}") + if not os.path.isfile(os.path.join(diffusers_path, "model_index.json")): + pytest.skip("Missing model_index.json in diffusers path") + if not gemma_model_path or not os.path.isdir(gemma_model_path): + pytest.skip("Gemma weights not found in text_encoder/gemma.") + if not os.path.isfile(official_path): + pytest.skip(f"Missing LTX-2 official weights at {official_path}") + if not os.path.isfile(upsampler_official_path): + pytest.skip(f"Missing official upsampler at {upsampler_official_path}") + if not os.path.isdir(upsampler_path): + pytest.skip(f"Missing FastVideo upsampler at {upsampler_path}") + if not os.path.isfile(lora_path): + pytest.skip(f"Missing distilled LoRA at {lora_path}") + + try: + from ltx_pipelines.ti2vid_two_stages import TI2VidTwoStagesPipeline + from ltx_core.loader import LoraPathStrengthAndSDOps + from ltx_core.loader.sd_ops import LTXV_LORA_COMFY_RENAMING_MAP + from ltx_core.model.transformer import attention as ltx_attention + from ltx_core.model.transformer.attention import Attention, AttentionFunction + except ImportError as exc: + pytest.skip(f"LTX-2 pipeline import failed: {exc}") + + # Force reference attention to use PyTorch SDPA + ltx_attention.memory_efficient_attention = None + ltx_attention.flash_attn_interface = None + + device = torch.device("cuda:0") + prompt = "A curious raccoon peers through a vibrant field of yellow sunflowers." + negative_prompt = "low quality, blurry, distorted, artifacts, jpeg compression" + seed = 42 + height = 64 + width = 64 + num_frames = 9 + fps = 12.0 + steps = 4 + guidance_scale = 4.0 + + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + + def _env_flag(name: str) -> bool: + return os.getenv(name, "0").lower() in ("1", "true", "yes") + + debug_kwargs = {} + if _env_flag("FASTVIDEO_DEBUG_STAGE_SUMS"): + debug_kwargs["debug_stage_sums"] = True + debug_kwargs["debug_stage_sums_path"] = os.getenv( + "FASTVIDEO_DEBUG_STAGE_SUMS_PATH") + if _env_flag("FASTVIDEO_DEBUG_MODULE_SUMS"): + debug_kwargs["debug_module_sums"] = True + debug_kwargs["debug_module_sums_path"] = os.getenv( + "FASTVIDEO_DEBUG_MODULE_SUMS_PATH") + include = os.getenv("FASTVIDEO_DEBUG_MODULE_SUMS_INCLUDE") + exclude = os.getenv("FASTVIDEO_DEBUG_MODULE_SUMS_EXCLUDE") + if include: + debug_kwargs["debug_module_sums_include"] = include.split(",") + if exclude: + debug_kwargs["debug_module_sums_exclude"] = exclude.split(",") + if _env_flag("FASTVIDEO_DEBUG_MODEL_SUMS"): + debug_kwargs["debug_model_sums"] = True + debug_kwargs["debug_model_sums_path"] = os.getenv( + "FASTVIDEO_DEBUG_MODEL_SUMS_PATH") + if _env_flag("FASTVIDEO_DEBUG_MODEL_DETAIL"): + debug_kwargs["debug_model_detail"] = True + debug_kwargs["debug_model_detail_path"] = os.getenv( + "FASTVIDEO_DEBUG_MODEL_DETAIL_PATH") + + ref_debug = _env_flag("FASTVIDEO_DEBUG_REF_SUMS") + ref_debug_path = os.getenv("FASTVIDEO_DEBUG_REF_SUMS_PATH") + + def _ref_log(line: str) -> None: + if ref_debug_path: + os.makedirs(os.path.dirname(ref_debug_path), exist_ok=True) + with open(ref_debug_path, "a", encoding="utf-8") as f: + f.write(line + "\n") + else: + print(line) + + with tempfile.TemporaryDirectory() as tmpdir: + # FastVideo two-stage + generator = VideoGenerator.from_pretrained( + diffusers_path, + num_gpus=1, + use_fsdp_inference=False, + dit_cpu_offload=False, + vae_cpu_offload=False, + text_encoder_cpu_offload=False, + pin_cpu_memory=False, + ltx2_vae_tiling=False, + ltx2_refine_enabled=True, + ltx2_refine_upsampler_path=upsampler_path, + ltx2_refine_lora_path=lora_path, + ltx2_refine_num_inference_steps=3, + ltx2_refine_guidance_scale=1.0, + ltx2_refine_add_noise=True, + **debug_kwargs, + ) + # Disable decoder noise for parity if possible. + if hasattr(generator, "executor") and hasattr(generator.executor, "pipeline"): + pipeline = generator.executor.pipeline + if hasattr(pipeline, "modules") and "vae" in pipeline.modules: + decoder = getattr(pipeline.modules["vae"], "decoder", None) + if decoder is not None and hasattr(decoder, "decode_noise_scale"): + decoder.decode_noise_scale = 0.0 + + result = generator.generate_video( + prompt=prompt, + negative_prompt=negative_prompt, + output_path=os.path.join(tmpdir, "fastvideo_two_stage"), + save_video=False, + height=height, + width=width, + num_frames=num_frames, + fps=fps, + num_inference_steps=steps, + guidance_scale=guidance_scale, + seed=seed, + ) + generator.shutdown() + + fastvideo_out = result["samples"].to(dtype=torch.float32).cpu() + _log_tensor_stats("fastvideo_video", fastvideo_out) + del result + torch.cuda.synchronize() + torch.cuda.empty_cache() + + # Reference two-stage + if ref_debug: + from ltx_pipelines.utils import helpers as ltx_helpers + import ltx_pipelines.ti2vid_two_stages as ltx_two_stage + from ltx_core.model import upsampler as ltx_upsampler + from ltx_core.tools import AudioLatentTools, VideoLatentTools + from ltx_core.text_encoders import gemma as ltx_gemma + + original_denoise = ltx_helpers.denoise_audio_video + original_upsample = ltx_upsampler.upsample_video + original_create_noised_state = ltx_helpers.create_noised_state + original_encode_text = ltx_gemma.encode_text + + def _ref_sum(tensor: torch.Tensor | None) -> float: + if tensor is None: + return 0.0 + return float(tensor.detach().sum(dtype=torch.float32).item()) + + def _wrapped_upsample(*args, **kwargs): + out = original_upsample(*args, **kwargs) + _ref_log( + f"reference:upsample_latent_sum={_ref_sum(out):.6f} " + f"shape={tuple(out.shape)}") + return out + + def _wrapped_denoise( + *, + output_shape, + conditionings, + noiser, + sigmas, + stepper, + denoising_loop_fn, + components, + dtype, + device, + noise_scale=1.0, + initial_video_latent=None, + initial_audio_latent=None, + skip_video_noiser=False, + skip_audio_noiser=False, + ): + stage = "stage2" if output_shape.width == width else "stage1" + _ref_log( + f"reference:{stage}:noise_scale={float(noise_scale):.6f} " + f"init_video_sum={_ref_sum(initial_video_latent):.6f} " + f"init_audio_sum={_ref_sum(initial_audio_latent):.6f}") + video_state, audio_state = original_denoise( + output_shape=output_shape, + conditionings=conditionings, + noiser=noiser, + sigmas=sigmas, + stepper=stepper, + denoising_loop_fn=denoising_loop_fn, + components=components, + dtype=dtype, + device=device, + noise_scale=noise_scale, + initial_video_latent=initial_video_latent, + initial_audio_latent=initial_audio_latent, + skip_video_noiser=skip_video_noiser, + skip_audio_noiser=skip_audio_noiser, + ) + _ref_log( + f"reference:{stage}:video_latent_sum={_ref_sum(video_state.latent):.6f} " + f"audio_latent_sum={_ref_sum(audio_state.latent):.6f}") + return video_state, audio_state + + def _wrapped_create_noised_state( + tools, + conditionings, + noiser, + dtype, + device, + noise_scale=1.0, + initial_latent=None, + skip_noiser=False, + ): + state = original_create_noised_state( + tools=tools, + conditionings=conditionings, + noiser=noiser, + dtype=dtype, + device=device, + noise_scale=noise_scale, + initial_latent=initial_latent, + skip_noiser=skip_noiser, + ) + stage = "stage2" if initial_latent is not None else "stage1" + kind = "video" if isinstance(tools, VideoLatentTools) else "audio" + _ref_log( + f"reference:{stage}:{kind}:noised_latent_sum={_ref_sum(state.latent):.6f} " + f"noise_scale={float(noise_scale):.6f}") + return state + + def _wrapped_encode_text(text_encoder, prompts): + context_p, context_n = original_encode_text(text_encoder, prompts=prompts) + v_context_p, a_context_p = context_p + v_context_n, a_context_n = context_n + _ref_log( + "reference:text:" + f"v_pos_sum={_ref_sum(v_context_p):.6f} " + f"a_pos_sum={_ref_sum(a_context_p):.6f} " + f"v_neg_sum={_ref_sum(v_context_n):.6f} " + f"a_neg_sum={_ref_sum(a_context_n):.6f}") + return context_p, context_n + + ltx_helpers.denoise_audio_video = _wrapped_denoise + ltx_upsampler.upsample_video = _wrapped_upsample + ltx_helpers.create_noised_state = _wrapped_create_noised_state + ltx_two_stage.denoise_audio_video = _wrapped_denoise + ltx_two_stage.upsample_video = _wrapped_upsample + ltx_gemma.encode_text = _wrapped_encode_text + ltx_two_stage.encode_text = _wrapped_encode_text + + ref_pipeline = TI2VidTwoStagesPipeline( + checkpoint_path=official_path, + distilled_lora=[ + LoraPathStrengthAndSDOps( + lora_path, + 1.0, + LTXV_LORA_COMFY_RENAMING_MAP, + ) + ], + spatial_upsampler_path=upsampler_official_path, + gemma_root=gemma_model_path, + loras=[], + device=device, + fp8transformer=False, + ) + + original_text_encoder = ref_pipeline.stage_1_model_ledger.text_encoder + original_transformer = ref_pipeline.stage_1_model_ledger.transformer + original_video_decoder = ref_pipeline.stage_2_model_ledger.video_decoder + + def _patched_text_encoder(): + encoder = original_text_encoder() + if hasattr(encoder, "model") and hasattr(encoder.model, "config"): + if hasattr(encoder.model.config, "attn_implementation"): + encoder.model.config.attn_implementation = "sdpa" + if hasattr(encoder.model.config, "_attn_implementation"): + encoder.model.config._attn_implementation = "sdpa" + for module in encoder.modules(): + if isinstance(module, Attention): + module.attention_function = AttentionFunction.PYTORCH + return encoder + + def _patched_video_decoder(): + decoder = original_video_decoder() + if hasattr(decoder, "decode_noise_scale"): + decoder.decode_noise_scale = 0.0 + return decoder + + ref_pipeline.stage_1_model_ledger.text_encoder = _patched_text_encoder + ref_pipeline.stage_2_model_ledger.video_decoder = _patched_video_decoder + ref_pipeline.stage_1_model_ledger.transformer = original_transformer + + with torch.no_grad(): + ref_video_iter, _ = ref_pipeline( + prompt=prompt, + negative_prompt=negative_prompt, + seed=seed, + height=height, + width=width, + num_frames=num_frames, + frame_rate=fps, + num_inference_steps=steps, + cfg_guidance_scale=guidance_scale, + images=[], + enhance_prompt=False, + ) + ref_chunks = list(ref_video_iter) + ref_video = torch.cat( + [chunk if torch.is_tensor(chunk) else torch.from_numpy(chunk) for chunk in ref_chunks], + dim=0, + ) + ref_video = ref_video.to(torch.float32) / 255.0 + ref_video = ref_video.permute(3, 0, 1, 2).unsqueeze(0) + ref_video = ref_video.cpu() + _log_tensor_stats("reference_video", ref_video) + + assert ref_video.shape == fastvideo_out.shape + assert_close(ref_video, fastvideo_out, atol=2 / 255, rtol=1e-3) diff --git a/tests/local_tests/transformers/test_ltx2_lora_parity.py b/tests/local_tests/transformers/test_ltx2_lora_parity.py new file mode 100644 index 000000000..feeb7694e --- /dev/null +++ b/tests/local_tests/transformers/test_ltx2_lora_parity.py @@ -0,0 +1,228 @@ +# SPDX-License-Identifier: Apache-2.0 +import os +from pathlib import Path +import sys + +import pytest +import torch +from torch.testing import assert_close + +os.environ.setdefault("MASTER_ADDR", "localhost") +os.environ.setdefault("MASTER_PORT", "29513") +os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "TORCH_SDPA") + +repo_root = Path(__file__).resolve().parents[3] +if str(repo_root) not in sys.path: + sys.path.insert(0, str(repo_root)) +ltx_core_path = repo_root / "LTX-2" / "packages" / "ltx-core" / "src" +if ltx_core_path.exists() and str(ltx_core_path) not in sys.path: + sys.path.insert(0, str(ltx_core_path)) + +from fastvideo.configs.models.dits import LTX2VideoConfig +from fastvideo.configs.pipelines import LTX2T2VConfig +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.pipelines.basic.ltx2.ltx2_pipeline import LTX2Pipeline +from fastvideo.pipelines.lora_pipeline import LoRAPipeline +from .test_ltx2 import _infer_patch_params, _read_transformer_config + + +def test_ltx2_transformer_lora_parity(): + torch.manual_seed(42) + diffusers_root = Path( + os.getenv("LTX2_DIFFUSERS_PATH", "converted/ltx2_diffusers") + ) + official_path = Path( + os.getenv( + "LTX2_OFFICIAL_PATH", + "official_ltx_weights/ltx-2-19b-distilled.safetensors", + ) + ) + lora_path = Path( + os.getenv( + "LTX2_LORA_PATH", + "official_ltx_weights/ltx-2-19b-distilled-lora-384.safetensors", + ) + ) + if not official_path.exists(): + pytest.skip(f"LTX-2 official weights not found at {official_path}") + if not lora_path.exists(): + pytest.skip(f"LTX-2 distilled LoRA not found at {lora_path}") + if not diffusers_root.exists(): + pytest.skip(f"LTX-2 diffusers weights not found at {diffusers_root}") + + try: + from ltx_core.loader.sft_loader import SafetensorsModelStateDictLoader + except ImportError as exc: + pytest.skip(f"LTX-2 import failed: {exc}") + + config_loader = SafetensorsModelStateDictLoader() + metadata = config_loader.metadata(str(official_path)) + transformer_config = _read_transformer_config(metadata) + + config = LTX2VideoConfig() + cfg = config.arch_config + cfg.num_attention_heads = transformer_config.get("num_attention_heads", + cfg.num_attention_heads) + cfg.attention_head_dim = transformer_config.get("attention_head_dim", + cfg.attention_head_dim) + cfg.num_layers = transformer_config.get("num_layers", cfg.num_layers) + cfg.cross_attention_dim = transformer_config.get( + "cross_attention_dim", cfg.cross_attention_dim) + cfg.caption_channels = transformer_config.get("caption_channels", + cfg.caption_channels) + cfg.norm_eps = transformer_config.get("norm_eps", cfg.norm_eps) + cfg.attention_type = transformer_config.get("attention_type", + cfg.attention_type) + cfg.positional_embedding_theta = transformer_config.get( + "positional_embedding_theta", cfg.positional_embedding_theta) + cfg.positional_embedding_max_pos = transformer_config.get( + "positional_embedding_max_pos", cfg.positional_embedding_max_pos) + cfg.timestep_scale_multiplier = transformer_config.get( + "timestep_scale_multiplier", cfg.timestep_scale_multiplier) + cfg.use_middle_indices_grid = transformer_config.get( + "use_middle_indices_grid", cfg.use_middle_indices_grid) + cfg.rope_type = transformer_config.get("rope_type", cfg.rope_type) + cfg.double_precision_rope = transformer_config.get( + "double_precision_rope", + transformer_config.get("frequencies_precision", "") + == "float64", + ) + cfg.audio_num_attention_heads = transformer_config.get( + "audio_num_attention_heads", cfg.audio_num_attention_heads) + cfg.audio_attention_head_dim = transformer_config.get( + "audio_attention_head_dim", cfg.audio_attention_head_dim) + cfg.audio_in_channels = transformer_config.get("audio_in_channels", + cfg.audio_in_channels) + cfg.audio_out_channels = transformer_config.get("audio_out_channels", + cfg.audio_out_channels) + cfg.audio_cross_attention_dim = transformer_config.get( + "audio_cross_attention_dim", cfg.audio_cross_attention_dim) + cfg.audio_positional_embedding_max_pos = transformer_config.get( + "audio_positional_embedding_max_pos", + cfg.audio_positional_embedding_max_pos, + ) + cfg.av_ca_timestep_scale_multiplier = transformer_config.get( + "av_ca_timestep_scale_multiplier", cfg.av_ca_timestep_scale_multiplier) + cfg.in_channels = transformer_config.get("in_channels", cfg.in_channels) + cfg.out_channels = transformer_config.get("out_channels", cfg.out_channels) + + patch_size, num_channels_latents = _infer_patch_params(cfg.in_channels) + cfg.patch_size = (1, patch_size, patch_size) + cfg.num_channels_latents = num_channels_latents + + if not torch.cuda.is_available(): + pytest.skip("LTX-2 LoRA parity test requires CUDA for attention backends.") + + device = torch.device("cuda:0") + precision = torch.bfloat16 + precision_str = "bf16" + + pipeline_config = LTX2T2VConfig() + pipeline_config.dit_config = config + pipeline_config.dit_precision = precision_str + + args = FastVideoArgs( + model_path=str(diffusers_root), + num_gpus=1, + use_fsdp_inference=False, + dit_cpu_offload=False, + dit_layerwise_offload=False, + text_encoder_cpu_offload=True, + vae_cpu_offload=True, + pin_cpu_memory=False, + lora_path=str(lora_path), + lora_nickname="refine", + pipeline_config=pipeline_config, + ) + args.device = device + + fastvideo_pipeline = LTX2Pipeline( + model_path=str(diffusers_root), + fastvideo_args=args, + required_config_modules=["transformer"], + ) + fastvideo_model = fastvideo_pipeline.modules["transformer"].to( + device=device, dtype=precision) + fastvideo_model.eval() + + lora_layers = fastvideo_pipeline.lora_layers.get("transformer") + assert lora_layers is not None, "LoRA layers were not created." + applied_layers = [ + layer for _, layer in lora_layers.all_lora_layers() + if not layer.disable_lora + ] + assert applied_layers, "LoRA adapter did not match any layers." + + fastvideo_state = fastvideo_model.state_dict() + + def _canonical(name: str) -> str: + while name.startswith("model."): + name = name[len("model."):] + return name + + fastvideo_weights: dict[str, torch.Tensor] = {} + for key, value in fastvideo_state.items(): + if key.endswith(".base_layer.weight"): + trimmed = key[:-len(".base_layer.weight")] + fastvideo_weights[_canonical(trimmed)] = value + continue + if key.endswith(".weight") and ".base_layer." not in key: + fastvideo_weights[_canonical(key[:-len(".weight")])] = value + + def _iter_lora_layers(path: Path) -> list[str]: + from safetensors import safe_open + + with safe_open(str(path), framework="pt") as f: + layers = set() + for key in f.keys(): + if key.endswith(".lora_A.weight"): + layer = key.replace("diffusion_model.", "").replace( + ".lora_A.weight", "") + layers.add(layer) + return sorted(layers) + + def _get_tensor(path: Path, key: str) -> torch.Tensor | None: + from safetensors import safe_open + + with safe_open(str(path), framework="pt") as f: + if key not in f.keys(): + return None + return f.get_tensor(key) + + lora_layers = _iter_lora_layers(lora_path) + assert lora_layers, "No LoRA layers found in the adapter." + + checked = 0 + for layer in lora_layers: + base_key = f"model.diffusion_model.{layer}.weight" + lora_a_key = f"diffusion_model.{layer}.lora_A.weight" + lora_b_key = f"diffusion_model.{layer}.lora_B.weight" + fast_weight = fastvideo_weights.get(layer) + if fast_weight is None: + continue + base_weight = _get_tensor(official_path, base_key) + lora_a = _get_tensor(lora_path, lora_a_key) + lora_b = _get_tensor(lora_path, lora_b_key) + if base_weight is None or lora_a is None or lora_b is None: + continue + expected = (base_weight.to(torch.bfloat16) + + torch.matmul(lora_b.to(torch.bfloat16), + lora_a.to(torch.bfloat16))) + actual = fast_weight.detach().cpu().to(torch.bfloat16) + assert_close(actual.float(), expected.float(), atol=5e-4, rtol=5e-4) + checked += 1 + if checked >= 10: + break + + assert checked > 0, "No matching LoRA layers found for parity checks." + + del fastvideo_model + del fastvideo_pipeline + LoRAPipeline.lora_layers.clear() + LoRAPipeline.lora_adapters.clear() + LoRAPipeline.cur_adapter_name = "" + LoRAPipeline.cur_adapter_path = "" + LoRAPipeline.lora_initialized = False + import gc + gc.collect() + torch.cuda.empty_cache() diff --git a/tests/local_tests/upsamplers/__init__.py b/tests/local_tests/upsamplers/__init__.py new file mode 100644 index 000000000..988131360 --- /dev/null +++ b/tests/local_tests/upsamplers/__init__.py @@ -0,0 +1 @@ +# SPDX-License-Identifier: Apache-2.0 diff --git a/tests/local_tests/upsamplers/test_ltx2_upsampler.py b/tests/local_tests/upsamplers/test_ltx2_upsampler.py new file mode 100644 index 000000000..b47db8475 --- /dev/null +++ b/tests/local_tests/upsamplers/test_ltx2_upsampler.py @@ -0,0 +1,104 @@ +# SPDX-License-Identifier: Apache-2.0 +import json +import os +from pathlib import Path +import sys + +import pytest +import torch +from safetensors import safe_open +from safetensors.torch import load_file +from torch.testing import assert_close + +repo_root = Path(__file__).resolve().parents[3] +if str(repo_root) not in sys.path: + sys.path.insert(0, str(repo_root)) +ltx_core_path = repo_root / "LTX-2" / "packages" / "ltx-core" / "src" +if ltx_core_path.exists() and str(ltx_core_path) not in sys.path: + sys.path.insert(0, str(ltx_core_path)) + +from fastvideo.configs.pipelines import PipelineConfig +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.models.loader.component_loader import UpsamplerLoader + + +def _load_metadata(path: Path) -> dict: + with safe_open(str(path), framework="pt") as f: + meta = f.metadata() + if not meta or "config" not in meta: + raise KeyError("Missing config metadata in safetensors file.") + return json.loads(meta["config"]) + + +def test_ltx2_upsampler_parity(): + official_path = Path( + os.getenv( + "LTX2_UPSAMPLER_OFFICIAL_PATH", + "official_ltx_weights/ltx-2-spatial-upscaler-x2-1.0.safetensors", + ) + ) + fastvideo_path = Path( + os.getenv( + "LTX2_UPSAMPLER_PATH", + "converted/ltx2_spatial_upscaler", + ) + ) + + if not official_path.exists(): + pytest.skip(f"LTX-2 upsampler weights not found at {official_path}") + if not fastvideo_path.exists(): + pytest.skip(f"FastVideo upsampler weights not found at {fastvideo_path}") + + try: + from ltx_core.model.upsampler import LatentUpsamplerConfigurator + except ImportError as exc: + pytest.skip(f"LTX-2 import failed: {exc}") + + config = _load_metadata(official_path) + + device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") + precision = torch.bfloat16 if torch.cuda.is_available() else torch.float32 + precision_str = "bf16" if torch.cuda.is_available() else "fp32" + + ref_model = LatentUpsamplerConfigurator.from_config(config).to( + device=device, dtype=precision + ) + ref_weights = load_file(str(official_path)) + ref_model.load_state_dict(ref_weights, strict=False) + + args = FastVideoArgs( + model_path=str(fastvideo_path), + pipeline_config=PipelineConfig(vae_precision=precision_str), + ) + loader = UpsamplerLoader() + fastvideo_model = loader.load(str(fastvideo_path), args).to( + device=device, dtype=precision + ) + + ref_model.eval() + fastvideo_model.eval() + + batch = 1 + channels = config.get("in_channels", 128) + frames = 3 + height = 8 + width = 8 + latent = torch.randn( + batch, + channels, + frames, + height, + width, + device=device, + dtype=precision, + ) + + with torch.no_grad(): + ref_out = ref_model(latent) + fast_out = fastvideo_model(latent) + + assert ref_out.shape == fast_out.shape + assert ref_out.dtype == fast_out.dtype + assert torch.isfinite(ref_out).all(), "Reference upsampler produced non-finite output." + assert torch.isfinite(fast_out).all(), "FastVideo upsampler produced non-finite output." + assert_close(ref_out, fast_out, atol=1e-4, rtol=1e-4)