From 81872ee8866f29cf708ccc68ca196bd40c398f2f Mon Sep 17 00:00:00 2001 From: Will Lin Date: Sun, 25 Jan 2026 16:57:10 -0800 Subject: [PATCH] update --- docs/utilities/debugging.md | 116 ++++++ fastvideo/configs/sample/ltx2.py | 17 +- fastvideo/entrypoints/cli/main.py | 2 + fastvideo/entrypoints/cli/upsample.py | 134 +++++++ fastvideo/entrypoints/upsample.py | 202 +++++++++++ fastvideo/fastvideo_args.py | 152 ++++++++ fastvideo/layers/lora/linear.py | 37 +- fastvideo/models/dits/ltx2.py | 5 + fastvideo/models/loader/component_loader.py | 30 ++ fastvideo/models/registry.py | 5 + fastvideo/models/upsamplers/__init__.py | 23 ++ fastvideo/models/upsamplers/ltx2_upsampler.py | 319 +++++++++++++++++ fastvideo/models/vaes/ltx2vae.py | 3 + .../pipelines/basic/ltx2/ltx2_pipeline.py | 130 ++++++- fastvideo/pipelines/composed_pipeline_base.py | 96 +++++ fastvideo/pipelines/stages/__init__.py | 7 + fastvideo/pipelines/stages/base.py | 68 ++++ fastvideo/pipelines/stages/ltx2_denoising.py | 189 +++++++--- .../stages/ltx2_latent_preparation.py | 63 +++- fastvideo/pipelines/stages/ltx2_refine.py | 336 ++++++++++++++++++ tests/local_tests/transformers/test_ltx2.py | 34 +- .../transformers/test_ltx2_audio.py | 47 ++- .../vaes/test_ltx2_vae_official.py | 2 + 23 files changed, 1937 insertions(+), 80 deletions(-) create mode 100644 docs/utilities/debugging.md create mode 100644 fastvideo/entrypoints/cli/upsample.py create mode 100644 fastvideo/entrypoints/upsample.py create mode 100644 fastvideo/models/upsamplers/__init__.py create mode 100644 fastvideo/models/upsamplers/ltx2_upsampler.py create mode 100644 fastvideo/pipelines/stages/ltx2_refine.py diff --git a/docs/utilities/debugging.md b/docs/utilities/debugging.md new file mode 100644 index 000000000..7aee6a7ca --- /dev/null +++ b/docs/utilities/debugging.md @@ -0,0 +1,116 @@ +# Debugging Tools + +FastVideo provides opt-in debugging utilities designed to help track numerical parity +without impacting performance when disabled. These tools are available via CLI flags +and are no-ops unless explicitly enabled. + +## Stage-level sums + +Log sums of key tensors after each pipeline stage (latents, outputs, prompt embeds). + +Example: + +```bash +python -m fastvideo.entrypoints.cli.generate \ + --model-path converted/ltx2_diffusers \ + --prompt "A curious raccoon peers through a vibrant field of yellow sunflowers." \ + --height 64 --width 64 --num-frames 9 \ + --debug-stage-sums \ + --debug-stage-sums-path outputs/debug/ltx2_stage_sums.log +``` + +Log format: + +``` +fastvideo:stage=LTX2DenoisingStage latents=... prompt_embeds=... extra.ltx2_audio_latents=... +``` + +## Model-level sums (LTX-2) + +Enable LTX-2 DiT internal sum logging (blocks, modality prep, etc.). + +Example: + +```bash +python -m fastvideo.entrypoints.cli.generate \ + --model-path converted/ltx2_diffusers \ + --prompt "A curious raccoon peers through a vibrant field of yellow sunflowers." \ + --height 64 --width 64 --num-frames 9 \ + --debug-model-sums \ + --debug-model-sums-path outputs/debug/ltx2_model_sums.log +``` + +For even more detail (forward hooks inside the DiT blocks): + +```bash +python -m fastvideo.entrypoints.cli.generate \ + --model-path converted/ltx2_diffusers \ + --prompt "A curious raccoon peers through a vibrant field of yellow sunflowers." \ + --height 64 --width 64 --num-frames 9 \ + --debug-model-detail \ + --debug-model-detail-path outputs/debug/ltx2_model_detail.log +``` + +## Recursive module sums + +Attach forward hooks to all parameterized submodules (recursively) and log output +sums. This is useful for debugging where a mismatch begins inside encoders, VAEs, +or DiT blocks. + +Example: + +```bash +python -m fastvideo.entrypoints.cli.generate \ + --model-path converted/ltx2_diffusers \ + --prompt "A curious raccoon peers through a vibrant field of yellow sunflowers." \ + --height 64 --width 64 --num-frames 9 \ + --debug-module-sums \ + --debug-module-sums-path outputs/debug/ltx2_module_sums.log \ + --debug-module-sums-include transformer text_encoder vae +``` + +Use `--debug-module-sums-exclude` to skip noisy modules by name substring. + +## Deterministic noise replay (LTX-2) + +These flags let you save/load the exact noise used for reproducible runs and parity +testing. + +Stage 1: +- `--ltx2-initial-latent-path` (video) +- `--ltx2-audio-latent-path` (audio) + +Stage 2 (refinement): +- `--ltx2-refine-noise-path` (video noise) +- `--ltx2-refine-audio-noise-path` (audio noise) + +Example: + +```bash +python -m fastvideo.entrypoints.cli.generate \ + --model-path converted/ltx2_diffusers \ + --prompt "A curious raccoon peers through a vibrant field of yellow sunflowers." \ + --height 64 --width 64 --num-frames 9 \ + --ltx2-refine-enabled \ + --ltx2-refine-upsampler-path converted/ltx2_spatial_upscaler \ + --ltx2-initial-latent-path outputs/debug/ltx2_stage1_latent.pt \ + --ltx2-refine-noise-path outputs/debug/ltx2_stage2_noise.pt +``` + +If the file exists, it is loaded. If it does not exist, FastVideo will save the +generated tensor to that path for future runs. + +## Parity testing (tests) + +Some local tests allow debug logging via environment variables, for example: + +```bash +FASTVIDEO_DEBUG_STAGE_SUMS=1 \ +FASTVIDEO_DEBUG_STAGE_SUMS_PATH=outputs/debug/two_stage_stage_sums.log \ +FASTVIDEO_DEBUG_REF_SUMS=1 \ +FASTVIDEO_DEBUG_REF_SUMS_PATH=outputs/debug/two_stage_reference_sums.log \ +pytest tests/local_tests/pipelines/test_ltx2_two_stage_parity.py -q -rs +``` + +This writes both FastVideo and reference sums so you can compare where divergence +begins. diff --git a/fastvideo/configs/sample/ltx2.py b/fastvideo/configs/sample/ltx2.py index 4be8d40bd..9466a6ea4 100644 --- a/fastvideo/configs/sample/ltx2.py +++ b/fastvideo/configs/sample/ltx2.py @@ -16,5 +16,18 @@ class LTX2SamplingParam(SamplingParam): fps: int = 24 num_inference_steps: int = 8 guidance_scale: float = 1.0 - # No default negative_prompt for distilled models - negative_prompt: str = "" + + # Official LTX-2 negative prompt (used only when guidance_scale > 1) + negative_prompt: str = ( + "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." + ) diff --git a/fastvideo/entrypoints/cli/main.py b/fastvideo/entrypoints/cli/main.py index 17cba3cb7..0b4319180 100644 --- a/fastvideo/entrypoints/cli/main.py +++ b/fastvideo/entrypoints/cli/main.py @@ -3,6 +3,7 @@ from fastvideo.entrypoints.cli.cli_types import CLISubcommand from fastvideo.entrypoints.cli.generate import cmd_init as generate_cmd_init +from fastvideo.entrypoints.cli.upsample import cmd_init as upsample_cmd_init from fastvideo.utils import FlexibleArgumentParser @@ -10,6 +11,7 @@ def cmd_init() -> list[CLISubcommand]: """Initialize all commands from separate modules""" commands = [] commands.extend(generate_cmd_init()) + commands.extend(upsample_cmd_init()) return commands diff --git a/fastvideo/entrypoints/cli/upsample.py b/fastvideo/entrypoints/cli/upsample.py new file mode 100644 index 000000000..55a7b1dc4 --- /dev/null +++ b/fastvideo/entrypoints/cli/upsample.py @@ -0,0 +1,134 @@ +# SPDX-License-Identifier: Apache-2.0 + +import argparse +import os +from typing import cast + +from fastvideo.entrypoints.cli.cli_types import CLISubcommand +from fastvideo.entrypoints.upsample import upscale_video_file +from fastvideo.logger import init_logger +from fastvideo.utils import FlexibleArgumentParser, StoreBoolean + +logger = init_logger(__name__) + + +class UpsampleSubcommand(CLISubcommand): + """The `upsample` subcommand for the FastVideo CLI.""" + + def __init__(self) -> None: + self.name = "upsample" + super().__init__() + + def cmd(self, args: argparse.Namespace) -> None: + upscale_video_file( + input_video=args.input_video, + output_video=args.output_video, + vae_path=args.vae_path, + upsampler_path=args.upsampler_path, + precision=args.precision, + device=args.device, + max_frames=args.max_frames, + trim_frames=args.trim_frames, + pad_frames=args.pad_frames, + crop_multiple=args.crop_multiple, + output_fps=args.output_fps, + ) + + def validate(self, args: argparse.Namespace) -> None: + if not os.path.exists(args.input_video): + raise ValueError(f"Input video not found: {args.input_video}") + if args.crop_multiple is not None and args.crop_multiple < 0: + raise ValueError("crop_multiple must be >= 0") + if args.max_frames is not None and args.max_frames <= 0: + raise ValueError("max_frames must be positive") + if args.trim_frames and args.pad_frames: + raise ValueError( + "Only one of --trim-frames or --pad-frames can be enabled") + + def subparser_init( + self, + subparsers: argparse._SubParsersAction, + ) -> FlexibleArgumentParser: + parser = subparsers.add_parser( + "upsample", + help="Upscale an existing video using the LTX-2 spatial upsampler", + usage= + ("fastvideo upsample --input-video INPUT.mp4 --output-video OUTPUT.mp4 " + "[--vae-path PATH] [--upsampler-path PATH]"), + ) + + parser.add_argument( + "--input-video", + type=str, + required=True, + help="Path to the input video file", + ) + parser.add_argument( + "--output-video", + type=str, + required=True, + help="Path to save the upscaled video", + ) + parser.add_argument( + "--vae-path", + type=str, + default="converted/ltx2_diffusers/vae", + help="Path to LTX-2 VAE weights (diffusers-style)", + ) + parser.add_argument( + "--upsampler-path", + type=str, + default="converted/ltx2_spatial_upscaler", + help="Path to LTX-2 spatial upsampler weights", + ) + parser.add_argument( + "--precision", + type=str, + default="bf16", + choices=["fp32", "fp16", "bf16"], + help="Precision to use for VAE + upsampler", + ) + parser.add_argument( + "--device", + type=str, + default=None, + help="Torch device string (e.g. cuda, cuda:0, cpu)", + ) + parser.add_argument( + "--max-frames", + type=int, + default=None, + help="Maximum number of frames to read from the input video", + ) + parser.add_argument( + "--trim-frames", + action=StoreBoolean, + default=True, + help="Trim frames to satisfy the 1+8k requirement", + ) + parser.add_argument( + "--pad-frames", + action=StoreBoolean, + default=False, + help= + "Pad frames to satisfy the 1+8k requirement (repeats last frame)", + ) + parser.add_argument( + "--crop-multiple", + type=int, + default=32, + help= + "Center-crop to make H/W divisible by this value (0 to disable)", + ) + parser.add_argument( + "--output-fps", + type=float, + default=None, + help="Override output video FPS (defaults to input FPS)", + ) + + return cast(FlexibleArgumentParser, parser) + + +def cmd_init() -> list[CLISubcommand]: + return [UpsampleSubcommand()] diff --git a/fastvideo/entrypoints/upsample.py b/fastvideo/entrypoints/upsample.py new file mode 100644 index 000000000..a1e3448cf --- /dev/null +++ b/fastvideo/entrypoints/upsample.py @@ -0,0 +1,202 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Utilities for upscaling videos with LTX-2 spatial upsampler.""" + +from __future__ import annotations + +from pathlib import Path + +import av +import imageio +import numpy as np +import torch + +from fastvideo.configs.pipelines import PipelineConfig +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.logger import init_logger +from fastvideo.models.loader.component_loader import UpsamplerLoader, VAELoader +from fastvideo.models.upsamplers import upsample_video +from fastvideo.utils import PRECISION_TO_TYPE + +logger = init_logger(__name__) + + +def _read_video(path: str | Path, + max_frames: int | None = None) -> tuple[torch.Tensor, float]: + """Read video frames via PyAV. + + Returns a tensor of shape [F, C, H, W] in [0, 1] and the fps. + """ + path = Path(path) + if not path.exists(): + raise FileNotFoundError(f"Input video not found: {path}") + + frames: list[np.ndarray] = [] + with av.open(str(path)) as container: + video_stream = container.streams.video[0] + fps = float(video_stream.average_rate or video_stream.base_rate or 24) + for frame in container.decode(video=0): + if max_frames is not None and len(frames) >= max_frames: + break + frames.append(frame.to_ndarray(format="rgb24")) + + if not frames: + raise ValueError(f"No frames decoded from {path}") + + frames_np = np.stack(frames, axis=0) + video = torch.from_numpy(frames_np).float().div(255.0) + return video.permute(0, 3, 1, 2), fps + + +def _write_video(frames: torch.Tensor, output_path: str | Path, + fps: float) -> None: + """Write frames [F, C, H, W] in [0, 1] to a video file.""" + output_path = Path(output_path) + output_path.parent.mkdir(parents=True, exist_ok=True) + + frames = frames.clamp(0, 1) + frames = (frames * 255.0).to(torch.uint8) + frames_np = frames.permute(0, 2, 3, 1).cpu().numpy() + imageio.mimsave(str(output_path), list(frames_np), fps=fps, format="mp4") + + +def _prepare_video( + video: torch.Tensor, + *, + trim_frames: bool, + pad_frames: bool, + crop_multiple: int, +) -> torch.Tensor: + """Ensure frames count and resolution satisfy LTX-2 VAE constraints.""" + frames, _, height, width = video.shape + + if trim_frames and pad_frames: + raise ValueError( + "Only one of trim_frames or pad_frames can be enabled.") + + if trim_frames and ((frames - 1) % 8) != 0: + valid_frames = 1 + 8 * ((frames - 1) // 8) + if valid_frames < 1: + raise ValueError("Video must have at least 1 frame.") + if valid_frames != frames: + logger.warning( + "Trimming frames from %d to %d to satisfy 1+8k requirement.", + frames, + valid_frames, + ) + video = video[:valid_frames] + frames = valid_frames + elif pad_frames and ((frames - 1) % 8) != 0: + valid_frames = 1 + 8 * (((frames - 1) + 7) // 8) + pad_count = valid_frames - frames + if pad_count > 0: + logger.warning( + "Padding frames from %d to %d to satisfy 1+8k requirement.", + frames, + valid_frames, + ) + pad = video[-1:].repeat(pad_count, 1, 1, 1) + video = torch.cat([video, pad], dim=0) + frames = valid_frames + + if crop_multiple > 0: + new_height = height - (height % crop_multiple) + new_width = width - (width % crop_multiple) + if new_height != height or new_width != width: + top = max((height - new_height) // 2, 0) + left = max((width - new_width) // 2, 0) + logger.warning( + "Center-cropping from %dx%d to %dx%d to be divisible by %d.", + height, + width, + new_height, + new_width, + crop_multiple, + ) + video = video[:, :, top:top + new_height, left:left + new_width] + + return video + + +def upscale_video_file( + *, + input_video: str | Path, + output_video: str | Path, + vae_path: str | Path, + upsampler_path: str | Path, + precision: str = "bf16", + device: str | None = None, + max_frames: int | None = None, + trim_frames: bool = True, + pad_frames: bool = False, + crop_multiple: int = 32, + output_fps: float | None = None, +) -> None: + """Upscale an existing video using the LTX-2 spatial upsampler.""" + input_video = str(input_video) + output_video = str(output_video) + vae_path = str(vae_path) + upsampler_path = str(upsampler_path) + + video, fps = _read_video(input_video, max_frames=max_frames) + original_frames = video.shape[0] + video = _prepare_video( + video, + trim_frames=trim_frames, + pad_frames=pad_frames, + crop_multiple=crop_multiple, + ) + final_frames = video.shape[0] + + target_device = torch.device(device) if device else (torch.device( + "cuda") if torch.cuda.is_available() else torch.device("cpu")) + + precision = precision.lower() + dtype = PRECISION_TO_TYPE.get(precision, torch.bfloat16) + if target_device.type == "cpu" and dtype != torch.float32: + logger.warning("CPU device selected; overriding precision to fp32.") + dtype = torch.float32 + precision = "fp32" + + args = FastVideoArgs( + model_path=vae_path, + pipeline_config=PipelineConfig(vae_precision=precision), + vae_cpu_offload=False, + ) + + vae_loader = VAELoader() + upsampler_loader = UpsamplerLoader() + + vae = vae_loader.load(vae_path, args).to(device=target_device, dtype=dtype) + upsampler = upsampler_loader.load(upsampler_path, + args).to(device=target_device, + dtype=dtype) + + if hasattr(vae.decoder, "decode_noise_scale"): + vae.decoder.decode_noise_scale = 0.0 + + # [F, C, H, W] -> [B, C, F, H, W] + video = video.unsqueeze(0).permute(0, 2, 1, 3, 4).to(device=target_device, + dtype=dtype) + + with torch.no_grad(): + latents = vae.encoder(video) + up_latents = upsample_video(latents, vae.encoder, + getattr(upsampler, "model", upsampler)) + + timestep_value = getattr(vae.decoder, "decode_timestep", 0.05) + timestep = torch.full((video.shape[0], ), + float(timestep_value), + device=target_device, + dtype=dtype) + decoded = vae.decoder(up_latents, timestep=timestep) + + # [B, C, F, H, W] -> [F, C, H, W] + decoded = decoded[0].permute(1, 0, 2, 3).detach().cpu() + if pad_frames and final_frames != original_frames: + decoded = decoded[:original_frames] + final_fps = output_fps or fps + _write_video(decoded, output_video, final_fps) + logger.info("Upscaled video saved to %s", output_video) + + +__all__ = ["upscale_video_file"] diff --git a/fastvideo/fastvideo_args.py b/fastvideo/fastvideo_args.py index b17df7047..fcb709028 100644 --- a/fastvideo/fastvideo_args.py +++ b/fastvideo/fastvideo_args.py @@ -171,6 +171,28 @@ class FastVideoArgs: ltx2_vae_temporal_tile_size_in_frames: int | None = None ltx2_vae_temporal_tile_overlap_in_frames: int | None = None ltx2_initial_latent_path: str | None = None + ltx2_audio_latent_path: str | None = None + ltx2_refine_enabled: bool = False + ltx2_refine_upsampler_path: str | None = None + ltx2_refine_transformer_path: str | None = None + ltx2_refine_lora_path: str | None = None + ltx2_refine_num_inference_steps: int = 3 + ltx2_refine_guidance_scale: float = 1.0 + ltx2_refine_add_noise: bool = True + ltx2_refine_noise_path: str | None = None + ltx2_refine_audio_noise_path: str | None = None + + # Debugging (opt-in, minimal overhead when disabled) + debug_stage_sums: bool = False + debug_stage_sums_path: str | None = None + debug_model_sums: bool = False + debug_model_sums_path: str | None = None + debug_model_detail: bool = False + debug_model_detail_path: str | None = None + debug_module_sums: bool = False + debug_module_sums_path: str | None = None + debug_module_sums_include: list[str] | None = None + debug_module_sums_exclude: list[str] | None = None # model paths for correct deallocation model_paths: dict[str, str] = field(default_factory=dict) @@ -405,6 +427,136 @@ class FastVideoArgs: default=FastVideoArgs.ltx2_initial_latent_path, help="Path to load/save a precomputed LTX-2 initial latent.", ) + parser.add_argument( + "--ltx2-audio-latent-path", + type=str, + default=FastVideoArgs.ltx2_audio_latent_path, + help="Path to load/save a precomputed LTX-2 initial audio latent.", + ) + parser.add_argument( + "--ltx2-refine-enabled", + action=StoreBoolean, + default=FastVideoArgs.ltx2_refine_enabled, + help= + "Enable LTX-2 stage2 refinement (2x spatial upsample + distilled denoising).", + ) + parser.add_argument( + "--ltx2-refine-upsampler-path", + type=str, + default=FastVideoArgs.ltx2_refine_upsampler_path, + help= + "Path to the LTX-2 spatial upsampler weights (diffusers format).", + ) + parser.add_argument( + "--ltx2-refine-transformer-path", + type=str, + default=FastVideoArgs.ltx2_refine_transformer_path, + help= + "Optional path to a dedicated stage2 transformer (e.g., distilled LoRA weights).", + ) + parser.add_argument( + "--ltx2-refine-lora-path", + type=str, + default=FastVideoArgs.ltx2_refine_lora_path, + help= + "Optional LoRA path to apply only during LTX-2 refinement stage2.", + ) + parser.add_argument( + "--ltx2-refine-num-inference-steps", + type=int, + default=FastVideoArgs.ltx2_refine_num_inference_steps, + help="Number of refinement steps for stage2 denoising (default: 3).", + ) + parser.add_argument( + "--ltx2-refine-guidance-scale", + type=float, + default=FastVideoArgs.ltx2_refine_guidance_scale, + help="CFG guidance scale for refinement (1.0 disables CFG).", + ) + parser.add_argument( + "--ltx2-refine-add-noise", + action=StoreBoolean, + default=FastVideoArgs.ltx2_refine_add_noise, + help="Add noise at sigma0 before stage2 denoising.", + ) + parser.add_argument( + "--ltx2-refine-noise-path", + type=str, + default=FastVideoArgs.ltx2_refine_noise_path, + help="Path to load/save stage2 video noise before refinement.", + ) + parser.add_argument( + "--ltx2-refine-audio-noise-path", + type=str, + default=FastVideoArgs.ltx2_refine_audio_noise_path, + help="Path to load/save stage2 audio noise before refinement.", + ) + + # Debugging (opt-in) + parser.add_argument( + "--debug-stage-sums", + action=StoreBoolean, + default=FastVideoArgs.debug_stage_sums, + help="Log tensor sums after each pipeline stage (debug only).", + ) + parser.add_argument( + "--debug-stage-sums-path", + type=str, + default=FastVideoArgs.debug_stage_sums_path, + help="Path to write stage-level sum logs (appended).", + ) + parser.add_argument( + "--debug-model-sums", + action=StoreBoolean, + default=FastVideoArgs.debug_model_sums, + help="Enable model-level sum logging (e.g., LTX-2 transformer).", + ) + parser.add_argument( + "--debug-model-sums-path", + type=str, + default=FastVideoArgs.debug_model_sums_path, + help="Path to write model-level sum logs.", + ) + parser.add_argument( + "--debug-model-detail", + action=StoreBoolean, + default=FastVideoArgs.debug_model_detail, + help="Enable detailed model hooks for activation sums (debug only).", + ) + parser.add_argument( + "--debug-model-detail-path", + type=str, + default=FastVideoArgs.debug_model_detail_path, + help="Path to write detailed model hook logs.", + ) + parser.add_argument( + "--debug-module-sums", + action=StoreBoolean, + default=FastVideoArgs.debug_module_sums, + help="Enable recursive module-level output sum logging.", + ) + parser.add_argument( + "--debug-module-sums-path", + type=str, + default=FastVideoArgs.debug_module_sums_path, + help="Path to write recursive module-level sum logs.", + ) + parser.add_argument( + "--debug-module-sums-include", + nargs="+", + type=str, + default=FastVideoArgs.debug_module_sums_include, + help= + "Optional list of substrings; only module names containing these will be logged.", + ) + parser.add_argument( + "--debug-module-sums-exclude", + nargs="+", + type=str, + default=FastVideoArgs.debug_module_sums_exclude, + help= + "Optional list of substrings; module names containing these will be skipped.", + ) # LoRA parameters (inference-time adapter loading) parser.add_argument( diff --git a/fastvideo/layers/lora/linear.py b/fastvideo/layers/lora/linear.py index d1368242f..aa62cbb24 100644 --- a/fastvideo/layers/lora/linear.py +++ b/fastvideo/layers/lora/linear.py @@ -13,7 +13,7 @@ from fastvideo.distributed import (get_local_torch_device, get_tp_rank, split_tensor_along_last_dim, tensor_model_parallel_all_gather, tensor_model_parallel_all_reduce) -from fastvideo.layers.linear import (ColumnParallelLinear, LinearBase, +from fastvideo.layers.linear import (ColumnParallelLinear, MergedColumnParallelLinear, QKVParallelLinear, ReplicatedLinear, RowParallelLinear) @@ -82,10 +82,9 @@ class BaseLayerWithLoRA(nn.Module): lora_B_sliced = self.slice_lora_b_weights( lora_B.to(x, non_blocking=True)) delta = x @ lora_A_sliced.T @ lora_B_sliced.T - if self.lora_alpha != self.lora_rank: - delta = delta * ( - self.lora_alpha / self.lora_rank # type: ignore - ) # type: ignore + if (self.lora_alpha is not None and self.lora_rank is not None + and self.lora_alpha != self.lora_rank): + delta = delta * (self.lora_alpha / self.lora_rank) out, output_bias = self.base_layer(x) return out + delta, output_bias else: @@ -98,6 +97,31 @@ class BaseLayerWithLoRA(nn.Module): def slice_lora_b_weights(self, B: torch.Tensor) -> torch.Tensor: return B + +class TorchLinearWithLoRA(BaseLayerWithLoRA): + """LoRA wrapper for torch.nn.Linear modules.""" + + @torch.compile() + def forward(self, x: torch.Tensor) -> torch.Tensor: + lora_A = self.lora_A + lora_B = self.lora_B + if isinstance(self.lora_B, DTensor): + lora_B = self.lora_B.to_local() + lora_A = self.lora_A.to_local() + + if not self.merged and not self.disable_lora: + lora_A_sliced = self.slice_lora_a_weights( + lora_A.to(x, non_blocking=True)) + lora_B_sliced = self.slice_lora_b_weights( + lora_B.to(x, non_blocking=True)) + delta = x @ lora_A_sliced.T @ lora_B_sliced.T + if (self.lora_alpha is not None and self.lora_rank is not None + and self.lora_alpha != self.lora_rank): + delta = delta * (self.lora_alpha / self.lora_rank) + out = self.base_layer(x) + return out + delta + return self.base_layer(x) + def set_lora_weights(self, A: torch.Tensor, B: torch.Tensor, @@ -369,7 +393,7 @@ def get_lora_layer(layer: nn.Module, lora_rank: int | None = None, lora_alpha: int | None = None, training_mode: bool = False) -> BaseLayerWithLoRA | None: - supported_layer_types: dict[type[LinearBase], type[BaseLayerWithLoRA]] = { + supported_layer_types: dict[type[nn.Module], type[BaseLayerWithLoRA]] = { # the order matters # VocabParallelEmbedding: VocabParallelEmbeddingWithLoRA, QKVParallelLinear: QKVParallelLinearWithLoRA, @@ -377,6 +401,7 @@ def get_lora_layer(layer: nn.Module, ColumnParallelLinear: ColumnParallelLinearWithLoRA, RowParallelLinear: RowParallelLinearWithLoRA, ReplicatedLinear: BaseLayerWithLoRA, + nn.Linear: TorchLinearWithLoRA, } for src_layer_type, lora_layer_type in supported_layer_types.items(): if isinstance(layer, src_layer_type): # pylint: disable=unidiomatic-typecheck diff --git a/fastvideo/models/dits/ltx2.py b/fastvideo/models/dits/ltx2.py index df6426144..f1c63fd6c 100644 --- a/fastvideo/models/dits/ltx2.py +++ b/fastvideo/models/dits/ltx2.py @@ -2275,6 +2275,7 @@ class LTX2Transformer3DModel(CachableDiT): DEFAULT_LTX2_SCALE_FACTORS, fps=fps, causal_fix=True, +<<<<<<< HEAD ).to(hidden_states.dtype) # Pad positions to match padded sequence length for SP cross-attention @@ -2283,6 +2284,10 @@ class LTX2Transformer3DModel(CachableDiT): last_pos = positions[:, :, -1:, :].expand(-1, -1, padding_needed, -1) positions = torch.cat([positions, last_pos], dim=2) +======= + ).to(torch.float32) + latents = self.patchifier.patchify(hidden_states) +>>>>>>> a02c2daf (update) video_modality = Modality( enabled=True, latent=latents, diff --git a/fastvideo/models/loader/component_loader.py b/fastvideo/models/loader/component_loader.py index 8437e9821..c6d8f1872 100644 --- a/fastvideo/models/loader/component_loader.py +++ b/fastvideo/models/loader/component_loader.py @@ -84,6 +84,8 @@ class ComponentLoader(ABC): "audio_vae": (AudioDecoderLoader, "diffusers"), "audio_decoder": (AudioDecoderLoader, "diffusers"), "vocoder": (VocoderLoader, "diffusers"), + "upsampler": (UpsamplerLoader, "diffusers"), + "spatial_upsampler": (UpsamplerLoader, "diffusers"), "text_encoder": (TextEncoderLoader, "transformers"), "text_encoder_2": (TextEncoderLoader, "transformers"), "tokenizer": (TokenizerLoader, "transformers"), @@ -727,6 +729,34 @@ class VocoderLoader(ComponentLoader): return vocoder.eval() +class UpsamplerLoader(ComponentLoader): + """Loader for LTX-2 spatial/temporal upsampler.""" + + def load(self, model_path: str, fastvideo_args: FastVideoArgs): + config = get_diffusers_config(model=model_path) + class_name = config.pop("_class_name", None) or "LTX2LatentUpsampler" + + model_cls, _ = ModelRegistry.resolve_model_cls(class_name) + target_device = get_local_torch_device() + + precision = getattr( + fastvideo_args.pipeline_config, "vae_precision", "bf16" + ) + with set_default_torch_dtype(PRECISION_TO_TYPE[precision]): + upsampler = model_cls(config).to(target_device) + + safetensors_list = glob.glob( + os.path.join(str(model_path), "*.safetensors") + ) + loaded: dict[str, torch.Tensor] = {} + for sf_file in safetensors_list: + loaded.update(safetensors_load_file(sf_file)) + + target_module = getattr(upsampler, "model", upsampler) + target_module.load_state_dict(loaded, strict=False) + return upsampler.eval() + + class TransformerLoader(ComponentLoader): """Loader for transformer.""" diff --git a/fastvideo/models/registry.py b/fastvideo/models/registry.py index 350618805..c79740b6c 100644 --- a/fastvideo/models/registry.py +++ b/fastvideo/models/registry.py @@ -84,6 +84,10 @@ _AUDIO_MODELS = { "LTX2Vocoder": ("audio", "ltx2_audio_vae", "LTX2Vocoder"), } +_UPSAMPLER_MODELS = { + "LTX2LatentUpsampler": ("upsamplers", "ltx2_upsampler", "LTX2LatentUpsampler"), +} + _SCHEDULERS = { "FlowMatchEulerDiscreteScheduler": ("schedulers", "scheduling_flow_match_euler_discrete", @@ -111,6 +115,7 @@ _LEGACY_FAST_VIDEO_MODELS = { **_IMAGE_ENCODER_MODELS, **_VAE_MODELS, **_AUDIO_MODELS, + **_UPSAMPLER_MODELS, **_SCHEDULERS, **_UPSAMPLERS, } diff --git a/fastvideo/models/upsamplers/__init__.py b/fastvideo/models/upsamplers/__init__.py new file mode 100644 index 000000000..b5752ce52 --- /dev/null +++ b/fastvideo/models/upsamplers/__init__.py @@ -0,0 +1,23 @@ +# SPDX-License-Identifier: Apache-2.0 + +from fastvideo.models.upsamplers.ltx2_upsampler import ( + BlurDownsample, + LatentUpsampler, + LatentUpsamplerConfigurator, + LTX2LatentUpsampler, + PixelShuffleND, + ResBlock, + SpatialRationalResampler, + upsample_video, +) + +__all__ = [ + "BlurDownsample", + "LatentUpsampler", + "LatentUpsamplerConfigurator", + "LTX2LatentUpsampler", + "PixelShuffleND", + "ResBlock", + "SpatialRationalResampler", + "upsample_video", +] diff --git a/fastvideo/models/upsamplers/ltx2_upsampler.py b/fastvideo/models/upsamplers/ltx2_upsampler.py new file mode 100644 index 000000000..c6bef8aba --- /dev/null +++ b/fastvideo/models/upsamplers/ltx2_upsampler.py @@ -0,0 +1,319 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +LTX-2 latent upsampler (spatial/temporal) implementation. +""" + +from __future__ import annotations + +import math +from typing import Any, Optional, Tuple + +import torch +import torch.nn as nn +from einops import rearrange + + +class PixelShuffleND(nn.Module): + """N-dimensional pixel shuffle for upsampling.""" + + def __init__(self, dims: int, upscale_factors: Tuple[int, int, int] = (2, 2, 2)) -> None: + super().__init__() + if dims not in (1, 2, 3): + raise ValueError("dims must be 1, 2, or 3") + self.dims = dims + self.upscale_factors = upscale_factors + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if self.dims == 3: + return rearrange( + x, + "b (c p1 p2 p3) d h w -> b c (d p1) (h p2) (w p3)", + p1=self.upscale_factors[0], + p2=self.upscale_factors[1], + p3=self.upscale_factors[2], + ) + if self.dims == 2: + return rearrange( + x, + "b (c p1 p2) h w -> b c (h p1) (w p2)", + p1=self.upscale_factors[0], + p2=self.upscale_factors[1], + ) + if self.dims == 1: + return rearrange( + x, + "b (c p1) f h w -> b c (f p1) h w", + p1=self.upscale_factors[0], + ) + raise ValueError(f"Unsupported dims: {self.dims}") + + +class BlurDownsample(nn.Module): + """ + Anti-aliased spatial downsampling by integer stride using a fixed separable binomial kernel. + Applies only on H,W. Works for dims=2 or dims=3 (per-frame). + """ + + def __init__(self, dims: int, stride: int, kernel_size: int = 5) -> None: + super().__init__() + if dims not in (2, 3): + raise ValueError("dims must be 2 or 3") + if stride < 1: + raise ValueError("stride must be >= 1") + if kernel_size < 3 or kernel_size % 2 != 1: + raise ValueError("kernel_size must be an odd integer >= 3") + + self.dims = dims + self.stride = stride + self.kernel_size = kernel_size + + k = torch.tensor([math.comb(kernel_size - 1, idx) for idx in range(kernel_size)]) + k2d = k[:, None] @ k[None, :] + k2d = (k2d / k2d.sum()).float() + self.register_buffer("kernel", k2d[None, None, :, :]) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if self.stride == 1: + return x + if self.dims == 2: + return self._apply_2d(x) + b, _, f, _, _ = x.shape + x = rearrange(x, "b c f h w -> (b f) c h w") + x = self._apply_2d(x) + h2, w2 = x.shape[-2:] + return rearrange(x, "(b f) c h w -> b c f h w", b=b, f=f, h=h2, w=w2) + + def _apply_2d(self, x2d: torch.Tensor) -> torch.Tensor: + c = x2d.shape[1] + weight = self.kernel.expand(c, 1, self.kernel_size, self.kernel_size) + return nn.functional.conv2d( + x2d, + weight=weight, + bias=None, + stride=self.stride, + padding=self.kernel_size // 2, + groups=c, + ) + + +def _rational_for_scale(scale: float) -> Tuple[int, int]: + mapping = {0.75: (3, 4), 1.5: (3, 2), 2.0: (2, 1), 4.0: (4, 1)} + if float(scale) not in mapping: + raise ValueError(f"Unsupported scale {scale}. Choose from {list(mapping.keys())}") + return mapping[float(scale)] + + +class SpatialRationalResampler(nn.Module): + """ + Fully-learned rational spatial scaling: up by 'num' via PixelShuffle, then + anti-aliased downsample by 'den' using fixed blur + stride. Operates on H,W only. + For dims==3, work per-frame for spatial scaling (temporal axis untouched). + """ + + def __init__(self, mid_channels: int, scale: float) -> None: + super().__init__() + self.scale = float(scale) + self.num, self.den = _rational_for_scale(self.scale) + self.conv = nn.Conv2d(mid_channels, (self.num**2) * mid_channels, kernel_size=3, padding=1) + self.pixel_shuffle = PixelShuffleND(2, upscale_factors=(self.num, self.num)) + self.blur_down = BlurDownsample(dims=2, stride=self.den) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + b, _, f, _, _ = x.shape + x = rearrange(x, "b c f h w -> (b f) c h w") + x = self.conv(x) + x = self.pixel_shuffle(x) + x = self.blur_down(x) + return rearrange(x, "(b f) c h w -> b c f h w", b=b, f=f) + + +class ResBlock(nn.Module): + """Residual block with two convolutional layers, group norm, and SiLU.""" + + def __init__(self, channels: int, mid_channels: Optional[int] = None, dims: int = 3) -> None: + super().__init__() + if mid_channels is None: + mid_channels = channels + + conv = nn.Conv2d if dims == 2 else nn.Conv3d + + self.conv1 = conv(channels, mid_channels, kernel_size=3, padding=1) + self.norm1 = nn.GroupNorm(32, mid_channels) + self.conv2 = conv(mid_channels, channels, kernel_size=3, padding=1) + self.norm2 = nn.GroupNorm(32, channels) + self.activation = nn.SiLU() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + residual = x + x = self.conv1(x) + x = self.norm1(x) + x = self.activation(x) + x = self.conv2(x) + x = self.norm2(x) + x = self.activation(x + residual) + return x + + +class LatentUpsampler(nn.Module): + """ + Model to upsample VAE latents spatially and/or temporally. + """ + + def __init__( + self, + in_channels: int = 128, + mid_channels: int = 512, + num_blocks_per_stage: int = 4, + dims: int = 3, + spatial_upsample: bool = True, + temporal_upsample: bool = False, + spatial_scale: float = 2.0, + rational_resampler: bool = False, + ) -> None: + super().__init__() + + self.in_channels = in_channels + self.mid_channels = mid_channels + self.num_blocks_per_stage = num_blocks_per_stage + self.dims = dims + self.spatial_upsample = spatial_upsample + self.temporal_upsample = temporal_upsample + self.spatial_scale = float(spatial_scale) + self.rational_resampler = rational_resampler + + conv = nn.Conv2d if dims == 2 else nn.Conv3d + + self.initial_conv = conv(in_channels, mid_channels, kernel_size=3, padding=1) + self.initial_norm = nn.GroupNorm(32, mid_channels) + self.initial_activation = nn.SiLU() + + self.res_blocks = nn.ModuleList([ResBlock(mid_channels, dims=dims) for _ in range(num_blocks_per_stage)]) + + if spatial_upsample and temporal_upsample: + self.upsampler = nn.Sequential( + nn.Conv3d(mid_channels, 8 * mid_channels, kernel_size=3, padding=1), + PixelShuffleND(3), + ) + elif spatial_upsample: + if rational_resampler: + self.upsampler = SpatialRationalResampler(mid_channels=mid_channels, scale=self.spatial_scale) + else: + self.upsampler = nn.Sequential( + nn.Conv2d(mid_channels, 4 * mid_channels, kernel_size=3, padding=1), + PixelShuffleND(2), + ) + elif temporal_upsample: + self.upsampler = nn.Sequential( + nn.Conv3d(mid_channels, 2 * mid_channels, kernel_size=3, padding=1), + PixelShuffleND(1), + ) + else: + raise ValueError("Either spatial_upsample or temporal_upsample must be True") + + self.post_upsample_res_blocks = nn.ModuleList( + [ResBlock(mid_channels, dims=dims) for _ in range(num_blocks_per_stage)] + ) + + self.final_conv = conv(mid_channels, in_channels, kernel_size=3, padding=1) + + def forward(self, latent: torch.Tensor) -> torch.Tensor: + b, _, f, _, _ = latent.shape + + if self.dims == 2: + x = rearrange(latent, "b c f h w -> (b f) c h w") + x = self.initial_conv(x) + x = self.initial_norm(x) + x = self.initial_activation(x) + + for block in self.res_blocks: + x = block(x) + + x = self.upsampler(x) + + for block in self.post_upsample_res_blocks: + x = block(x) + + x = self.final_conv(x) + x = rearrange(x, "(b f) c h w -> b c f h w", b=b, f=f) + else: + x = self.initial_conv(latent) + x = self.initial_norm(x) + x = self.initial_activation(x) + + for block in self.res_blocks: + x = block(x) + + if self.temporal_upsample: + x = self.upsampler(x) + x = x[:, :, 1:, :, :] + elif isinstance(self.upsampler, SpatialRationalResampler): + x = self.upsampler(x) + else: + x = rearrange(x, "b c f h w -> (b f) c h w") + x = self.upsampler(x) + x = rearrange(x, "(b f) c h w -> b c f h w", b=b, f=f) + + for block in self.post_upsample_res_blocks: + x = block(x) + + x = self.final_conv(x) + + return x + + +class LatentUpsamplerConfigurator: + """Configurator for LatentUpsampler from a config dict.""" + + @classmethod + def from_config(cls, config: dict[str, Any]) -> LatentUpsampler: + cfg = dict(config) + cfg.pop("_class_name", None) + if "upsampler" in cfg and isinstance(cfg["upsampler"], dict): + cfg = cfg["upsampler"] + + return LatentUpsampler( + in_channels=cfg.get("in_channels", 128), + mid_channels=cfg.get("mid_channels", 512), + num_blocks_per_stage=cfg.get("num_blocks_per_stage", 4), + dims=cfg.get("dims", 3), + spatial_upsample=cfg.get("spatial_upsample", True), + temporal_upsample=cfg.get("temporal_upsample", False), + spatial_scale=cfg.get("spatial_scale", 2.0), + rational_resampler=cfg.get("rational_resampler", False), + ) + + +class LTX2LatentUpsampler(nn.Module): + """Public wrapper for the LTX-2 latent upsampler.""" + + def __init__(self, config: dict[str, Any]): + super().__init__() + self.model: LatentUpsampler = LatentUpsamplerConfigurator.from_config(config) + + def forward(self, latent: torch.Tensor) -> torch.Tensor: + return self.model(latent) + + +def upsample_video(latent: torch.Tensor, video_encoder: Any, upsampler: LatentUpsampler) -> torch.Tensor: + """ + Upsample a latent tensor with normalization based on the video encoder's per-channel statistics. + """ + if not hasattr(video_encoder, "per_channel_statistics"): + raise ValueError("video_encoder must expose per_channel_statistics for normalization") + stats = video_encoder.per_channel_statistics + latent = stats.un_normalize(latent) + latent = upsampler(latent) + latent = stats.normalize(latent) + return latent + + +__all__ = [ + "PixelShuffleND", + "BlurDownsample", + "SpatialRationalResampler", + "ResBlock", + "LatentUpsampler", + "LatentUpsamplerConfigurator", + "LTX2LatentUpsampler", + "upsample_video", +] diff --git a/fastvideo/models/vaes/ltx2vae.py b/fastvideo/models/vaes/ltx2vae.py index 5f921dccf..52866fdd0 100644 --- a/fastvideo/models/vaes/ltx2vae.py +++ b/fastvideo/models/vaes/ltx2vae.py @@ -17,7 +17,9 @@ import torch.nn.functional as F from einops import rearrange from fastvideo.models.vaes.common import DiagonalGaussianDistribution +from fastvideo.logger import init_logger +logger = init_logger(__name__) # ============================================================================= # Enums @@ -1285,6 +1287,7 @@ class VideoEncoder(nn.Module): def forward(self, sample: torch.Tensor) -> torch.Tensor: frames_count = sample.shape[2] + logger.info(f"Frames count: {frames_count}") if ((frames_count - 1) % 8) != 0: raise ValueError( "Invalid number of frames: Encode input must have 1 + 8 * x frames " diff --git a/fastvideo/pipelines/basic/ltx2/ltx2_pipeline.py b/fastvideo/pipelines/basic/ltx2/ltx2_pipeline.py index 4ba3ecd2e..0d5da0861 100644 --- a/fastvideo/pipelines/basic/ltx2/ltx2_pipeline.py +++ b/fastvideo/pipelines/basic/ltx2/ltx2_pipeline.py @@ -11,17 +11,17 @@ from transformers import AutoTokenizer from fastvideo.fastvideo_args import FastVideoArgs from fastvideo.logger import init_logger from fastvideo.models.loader.component_loader import PipelineComponentLoader -from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase -from fastvideo.pipelines.stages import (DecodingStage, InputValidationStage, - LTX2AudioDecodingStage, - LTX2DenoisingStage, - LTX2LatentPreparationStage, - LTX2TextEncodingStage) +from fastvideo.pipelines.lora_pipeline import LoRAPipeline +from fastvideo.pipelines.stages import ( + DecodingStage, InputValidationStage, LTX2AudioDecodingStage, + LTX2DenoisingStage, LTX2LatentPreparationStage, LTX2RefineInitStage, + LTX2RefineLoRAStage, LTX2UpsampleStage, STAGE_2_DISTILLED_SIGMA_VALUES, + LTX2TextEncodingStage) logger = init_logger(__name__) -class LTX2Pipeline(ComposedPipelineBase): +class LTX2Pipeline(LoRAPipeline): _required_config_modules = [ "text_encoder", @@ -33,6 +33,8 @@ class LTX2Pipeline(ComposedPipelineBase): ] def create_pipeline_stages(self, fastvideo_args: FastVideoArgs): + refine_enabled = fastvideo_args.ltx2_refine_enabled + self.add_stage( stage_name="input_validation_stage", stage=InputValidationStage(), @@ -46,6 +48,12 @@ class LTX2Pipeline(ComposedPipelineBase): ), ) + if refine_enabled: + self.add_stage( + stage_name="ltx2_refine_init_stage", + stage=LTX2RefineInitStage(), + ) + self.add_stage( stage_name="latent_preparation_stage", stage=LTX2LatentPreparationStage( @@ -58,6 +66,54 @@ class LTX2Pipeline(ComposedPipelineBase): transformer=self.get_module("transformer"), ), ) + if refine_enabled: + stage2_sigmas = STAGE_2_DISTILLED_SIGMA_VALUES + stage2_steps = fastvideo_args.ltx2_refine_num_inference_steps + expected_steps = len(stage2_sigmas) - 1 + if stage2_steps != expected_steps: + logger.warning( + "ltx2_refine_num_inference_steps=%s does not match distilled schedule; " + "using %s steps to align with stage2 sigmas.", + stage2_steps, + expected_steps, + ) + stage2_steps = expected_steps + + transformer_refine = self.get_module("transformer_refine", + self.get_module("transformer")) + + self.add_stage( + stage_name="ltx2_upsample_stage", + stage=LTX2UpsampleStage( + upsampler=self.get_module("spatial_upsampler"), + vae=self.get_module("vae"), + transformer=transformer_refine, + sigmas=stage2_sigmas, + add_noise=fastvideo_args.ltx2_refine_add_noise, + ), + ) + + if fastvideo_args.ltx2_refine_lora_path: + self.add_stage( + stage_name="ltx2_refine_lora_stage", + stage=LTX2RefineLoRAStage( + pipeline=self, + lora_path=fastvideo_args.ltx2_refine_lora_path, + ), + ) + + self.add_stage( + stage_name="ltx2_refine_denoising_stage", + stage=LTX2DenoisingStage( + transformer=transformer_refine, + sigmas_override=stage2_sigmas, + num_inference_steps_override=stage2_steps, + force_guidance_scale=fastvideo_args. + ltx2_refine_guidance_scale, + initial_audio_latents_key="ltx2_audio_latents", + ), + ) + self.add_stage( stage_name="audio_decoding_stage", stage=LTX2AudioDecodingStage( @@ -72,6 +128,31 @@ class LTX2Pipeline(ComposedPipelineBase): ) def initialize_pipeline(self, fastvideo_args: FastVideoArgs): + if fastvideo_args.debug_model_sums: + os.environ["LTX2_PIPELINE_DEBUG_LOG"] = "1" + if fastvideo_args.debug_model_sums_path: + os.environ[ + "LTX2_PIPELINE_DEBUG_PATH"] = fastvideo_args.debug_model_sums_path + else: + logger.warning( + "debug_model_sums is enabled but debug_model_sums_path is not set; no model sums will be logged." + ) + else: + os.environ.pop("LTX2_PIPELINE_DEBUG_LOG", None) + os.environ.pop("LTX2_PIPELINE_DEBUG_PATH", None) + if fastvideo_args.debug_model_detail: + os.environ["LTX2_DEBUG_DETAIL"] = "1" + if fastvideo_args.debug_model_detail_path: + os.environ[ + "LTX2_PIPELINE_DEBUG_DETAIL_PATH"] = fastvideo_args.debug_model_detail_path + else: + logger.warning( + "debug_model_detail is enabled but debug_model_detail_path is not set; no detailed hooks will be logged." + ) + else: + os.environ.pop("LTX2_DEBUG_DETAIL", None) + os.environ.pop("LTX2_PIPELINE_DEBUG_DETAIL_PATH", None) + tokenizer = self.get_module("tokenizer") if tokenizer is not None: tokenizer.padding_side = "left" @@ -144,6 +225,41 @@ class LTX2Pipeline(ComposedPipelineBase): raise ValueError( f"Required module {module_name} was not loaded properly") + if fastvideo_args.ltx2_refine_enabled: + upsampler_path = fastvideo_args.ltx2_refine_upsampler_path + if upsampler_path is None: + raise ValueError( + "ltx2_refine_enabled is True but ltx2_refine_upsampler_path was not provided." + ) + if loaded_modules is not None and "spatial_upsampler" in loaded_modules: + modules["spatial_upsampler"] = loaded_modules[ + "spatial_upsampler"] + else: + modules[ + "spatial_upsampler"] = PipelineComponentLoader.load_module( + module_name="spatial_upsampler", + component_model_path=upsampler_path, + transformers_or_diffusers="diffusers", + fastvideo_args=fastvideo_args, + ) + logger.info("Loaded module spatial_upsampler from %s", + upsampler_path) + + if loaded_modules is not None and "transformer_refine" in loaded_modules: + modules["transformer_refine"] = loaded_modules[ + "transformer_refine"] + elif fastvideo_args.ltx2_refine_transformer_path: + modules[ + "transformer_refine"] = PipelineComponentLoader.load_module( + module_name="transformer", + component_model_path=fastvideo_args. + ltx2_refine_transformer_path, + transformers_or_diffusers="diffusers", + fastvideo_args=fastvideo_args, + ) + logger.info("Loaded module transformer_refine from %s", + fastvideo_args.ltx2_refine_transformer_path) + return modules diff --git a/fastvideo/pipelines/composed_pipeline_base.py b/fastvideo/pipelines/composed_pipeline_base.py index d9a7f6419..c044aa859 100644 --- a/fastvideo/pipelines/composed_pipeline_base.py +++ b/fastvideo/pipelines/composed_pipeline_base.py @@ -155,6 +155,8 @@ class ComposedPipelineBase(ABC): transformer_module_2, **compile_kwargs) logger.info("Torch Compile enabled for DiT") + self._maybe_attach_module_sum_hooks() + if not self.fastvideo_args.training_mode: logger.info("Creating pipeline stages...") self.create_pipeline_stages(self.fastvideo_args) @@ -224,6 +226,100 @@ class ComposedPipelineBase(ABC): def add_module(self, module_name: str, module: Any): self.modules[module_name] = module + def _maybe_attach_module_sum_hooks(self) -> None: + args = self.fastvideo_args + if args is None or not getattr(args, "debug_module_sums", False): + return + if getattr(self, "_debug_module_sums_attached", False): + return + + log_path = getattr(args, "debug_module_sums_path", None) + if not log_path: + log_path = os.path.join("outputs", "debug", "module_sums.log") + logger.warning( + "debug_module_sums is enabled but debug_module_sums_path is not set; " + "defaulting to %s", + log_path, + ) + + include = getattr(args, "debug_module_sums_include", None) or [] + exclude = getattr(args, "debug_module_sums_exclude", None) or [] + + def _matches(name: str) -> bool: + include_ok = (not include) or any(key in name for key in include) + exclude_ok = (not exclude) or not any(key in name + for key in exclude) + return include_ok and exclude_ok + + def _sum_output(value: object) -> float | None: + if isinstance(value, torch.Tensor): + return float(value.detach().sum(dtype=torch.float32).item()) + if isinstance(value, dict): + total = 0.0 + found = False + for item in value.values(): + summed = _sum_output(item) + if summed is not None: + total += summed + found = True + return total if found else None + if isinstance(value, list | tuple): + total = 0.0 + found = False + for item in value: + summed = _sum_output(item) + if summed is not None: + total += summed + found = True + return total if found else None + return None + + def _write_line(line: str) -> None: + log_dir = os.path.dirname(log_path) + if log_dir: + os.makedirs(log_dir, exist_ok=True) + with open(log_path, "a", encoding="utf-8") as f: + f.write(line + "\n") + + handles: list[torch.utils.hooks.RemovableHandle] = [] + + for root_name, module in self.modules.items(): + if module is None or not isinstance(module, torch.nn.Module): + continue + for name, submodule in module.named_modules(): + full_name = root_name if not name else f"{root_name}.{name}" + if not _matches(full_name): + continue + if getattr(submodule, "_fastvideo_module_sum_hooked", False): + continue + # Only hook modules with direct parameters to avoid excessive noise. + is_root = name == "" + if not is_root and not any( + True for _ in submodule.parameters(recurse=False)): + continue + + def _hook_factory(module_name: str, module_type: str): + + def _hook(_module, _inputs, outputs): # noqa: ANN001 + summed = _sum_output(outputs) + if summed is None: + return + line = (f"fastvideo:module={module_name} " + f"class={module_type} out_sum={summed:.6f}") + _write_line(line) + + return _hook + + handle = submodule.register_forward_hook( + _hook_factory(full_name, submodule.__class__.__name__)) + handles.append(handle) + submodule._fastvideo_module_sum_hooked = True + + self._debug_module_sums_attached = True + self._debug_module_sum_handles = handles + logger.info("Attached %s module sum hooks for recursive logging", + len(handles)) + def _load_config(self, model_path: str) -> dict[str, Any]: model_path = maybe_download_model(self.model_path) self.model_path = model_path diff --git a/fastvideo/pipelines/stages/__init__.py b/fastvideo/pipelines/stages/__init__.py index f085d6749..6c2d9cc56 100644 --- a/fastvideo/pipelines/stages/__init__.py +++ b/fastvideo/pipelines/stages/__init__.py @@ -30,6 +30,9 @@ from fastvideo.pipelines.stages.ltx2_denoising import LTX2DenoisingStage from fastvideo.pipelines.stages.ltx2_latent_preparation import ( LTX2LatentPreparationStage) from fastvideo.pipelines.stages.ltx2_text_encoding import LTX2TextEncodingStage +from fastvideo.pipelines.stages.ltx2_refine import ( + LTX2RefineInitStage, LTX2RefineLoRAStage, LTX2UpsampleStage, + STAGE_2_DISTILLED_SIGMA_VALUES) from fastvideo.pipelines.stages.matrixgame_denoising import ( MatrixGameCausalDenoisingStage) from fastvideo.pipelines.stages.hyworld_denoising import HYWorldDenoisingStage @@ -58,6 +61,10 @@ __all__ = [ "Cosmos25AutoLatentPreparationStage", "LTX2LatentPreparationStage", "LTX2AudioDecodingStage", + "LTX2RefineInitStage", + "LTX2RefineLoRAStage", + "LTX2UpsampleStage", + "STAGE_2_DISTILLED_SIGMA_VALUES", "ConditioningStage", "DenoisingStage", "DmdDenoisingStage", diff --git a/fastvideo/pipelines/stages/base.py b/fastvideo/pipelines/stages/base.py index 97d818044..daba67f3c 100644 --- a/fastvideo/pipelines/stages/base.py +++ b/fastvideo/pipelines/stages/base.py @@ -6,6 +6,7 @@ This module defines the abstract base classes for pipeline stages that can be composed to create complete diffusion pipelines. """ +import os import time import traceback from abc import ABC, abstractmethod @@ -21,6 +22,28 @@ from fastvideo.pipelines.stages.validators import VerificationResult logger = init_logger(__name__) +def _sum_tensor(value: object) -> float | None: + if isinstance(value, torch.Tensor): + return float(value.detach().sum(dtype=torch.float32).item()) + if isinstance(value, list | tuple): + total = 0.0 + found = False + for item in value: + if isinstance(item, torch.Tensor): + total += item.detach().sum(dtype=torch.float32).item() + found = True + return total if found else None + return None + + +def _write_debug_line(path: str, line: str) -> None: + log_dir = os.path.dirname(path) + if log_dir: + os.makedirs(log_dir, exist_ok=True) + with open(path, "a", encoding="utf-8") as f: + f.write(line + "\n") + + class StageVerificationError(Exception): """Exception raised when stage verification fails.""" pass @@ -152,6 +175,7 @@ class PipelineStage(ABC): try: result = self.forward(batch, fastvideo_args) + self._maybe_log_stage_sums(stage_name, result, fastvideo_args) execution_time = time.perf_counter() - start_time logger.info("[%s] Execution completed in %s ms", stage_name, execution_time * 1000) @@ -167,6 +191,7 @@ class PipelineStage(ABC): else: # Direct execution (current behavior) result = self.forward(batch, fastvideo_args) + self._maybe_log_stage_sums(stage_name, result, fastvideo_args) if enable_verification: # Post-execution output verification @@ -180,6 +205,49 @@ class PipelineStage(ABC): return result + def _maybe_log_stage_sums( + self, + stage_name: str, + batch: ForwardBatch, + fastvideo_args: FastVideoArgs, + ) -> None: + if not getattr(fastvideo_args, "debug_stage_sums", False): + return + + entries: dict[str, float] = {} + + def _add(name: str, value: object) -> None: + summed = _sum_tensor(value) + if summed is not None: + entries[name] = summed + + _add("latents", batch.latents) + _add("image_latent", batch.image_latent) + _add("video_latent", batch.video_latent) + _add("noise_pred", batch.noise_pred) + _add("output", batch.output) + _add("prompt_embeds", batch.prompt_embeds) + _add("negative_prompt_embeds", batch.negative_prompt_embeds) + + if isinstance(batch.extra, dict): + for key, value in batch.extra.items(): + summed = _sum_tensor(value) + if summed is not None: + entries[f"extra.{key}"] = summed + + if not entries: + return + + batch.logging_info.add_stage_metric(stage_name, "sums", entries) + line = "fastvideo:stage={}".format(stage_name) + " " + " ".join( + f"{key}={value:.6f}" for key, value in entries.items()) + + path = getattr(fastvideo_args, "debug_stage_sums_path", None) + if path: + _write_debug_line(path, line) + else: + logger.info("%s", line) + @abstractmethod def forward( self, diff --git a/fastvideo/pipelines/stages/ltx2_denoising.py b/fastvideo/pipelines/stages/ltx2_denoising.py index 98d04f6e1..cee78bc9b 100644 --- a/fastvideo/pipelines/stages/ltx2_denoising.py +++ b/fastvideo/pipelines/stages/ltx2_denoising.py @@ -7,6 +7,7 @@ from __future__ import annotations import math import os +from pathlib import Path import torch from tqdm.auto import tqdm @@ -80,9 +81,21 @@ def _ltx2_sigmas( class LTX2DenoisingStage(PipelineStage): """Run the LTX-2 denoising loop over the sigma schedule.""" - def __init__(self, transformer) -> None: + def __init__( + self, + transformer, + *, + sigmas_override: list[float] | None = None, + num_inference_steps_override: int | None = None, + force_guidance_scale: float | None = None, + initial_audio_latents_key: str | None = None, + ) -> None: super().__init__() self.transformer = transformer + self.sigmas_override = sigmas_override + self.num_inference_steps_override = num_inference_steps_override + self.force_guidance_scale = force_guidance_scale + self.initial_audio_latents_key = initial_audio_latents_key def forward( self, @@ -98,8 +111,14 @@ class LTX2DenoisingStage(PipelineStage): neg_prompt_embeds = None neg_prompt_mask = None + guidance_scale = batch.guidance_scale + use_cfg = batch.do_classifier_free_guidance + if self.force_guidance_scale is not None: + guidance_scale = self.force_guidance_scale + use_cfg = guidance_scale > 1.0 + # Only load negative prompts if CFG is actually enabled - if batch.do_classifier_free_guidance: + if use_cfg: assert batch.negative_prompt_embeds is not None, ( "CFG is enabled but negative_prompt_embeds is None") neg_prompt_embeds = batch.negative_prompt_embeds[0] @@ -121,22 +140,31 @@ class LTX2DenoisingStage(PipelineStage): ) and not fastvideo_args.disable_autocast and ( not disable_autocast) - # Use official distilled sigma schedule for 8 steps (distilled models) - use_distilled_sigmas = os.getenv("LTX2_USE_DISTILLED_SIGMAS", - "1") == "1" - if use_distilled_sigmas and batch.num_inference_steps == 8: + num_steps = self.num_inference_steps_override or batch.num_inference_steps + if self.sigmas_override is not None: sigmas = torch.tensor( - DISTILLED_SIGMA_VALUES, + self.sigmas_override, device=latents.device, dtype=torch.float32, ) - logger.info("[LTX2] Using official distilled sigma schedule") + logger.info("[LTX2] Using overridden sigma schedule") else: - sigmas = _ltx2_sigmas( - steps=batch.num_inference_steps, - latent=None, - device=latents.device, - ) + # Use official distilled sigma schedule for 8 steps (distilled models) + use_distilled_sigmas = os.getenv("LTX2_USE_DISTILLED_SIGMAS", + "1") == "1" + if use_distilled_sigmas and num_steps == 8: + sigmas = torch.tensor( + DISTILLED_SIGMA_VALUES, + device=latents.device, + dtype=torch.float32, + ) + logger.info("[LTX2] Using official distilled sigma schedule") + else: + sigmas = _ltx2_sigmas( + steps=num_steps, + latent=None, + device=latents.device, + ) if hasattr(self.transformer, "patchifier"): video_shape = VideoLatentShape.from_torch_shape(latents.shape) token_count = self.transformer.patchifier.get_token_count( @@ -144,7 +172,7 @@ class LTX2DenoisingStage(PipelineStage): else: token_count = 1 timestep_template = torch.ones( - (latents.shape[0], token_count), + (latents.shape[0], token_count, 1), device=latents.device, dtype=torch.float32, ) @@ -154,7 +182,15 @@ class LTX2DenoisingStage(PipelineStage): audio_context_n = audio_neg_embeds[0] if audio_neg_embeds else None audio_latents = None audio_timestep_template = None - if audio_context_p is not None: + if self.initial_audio_latents_key is not None: + audio_latents = batch.extra.get(self.initial_audio_latents_key) + if audio_latents is not None: + audio_timestep_template = torch.ones( + (latents.shape[0], audio_latents.shape[2], 1), + device=latents.device, + dtype=torch.float32, + ) + if audio_context_p is not None and audio_latents is None: fps_value = batch.fps if isinstance(fps_value, list): fps_value = fps_value[0] if fps_value else None @@ -170,53 +206,69 @@ class LTX2DenoisingStage(PipelineStage): hop_length=DEFAULT_LTX2_AUDIO_HOP_LENGTH, audio_latent_downsample_factor=DEFAULT_LTX2_AUDIO_DOWNSAMPLE, ) - audio_generator = None - if fastvideo_args.ltx2_initial_latent_path and batch.seed is not None: - audio_generator = torch.Generator( - device=latents.device).manual_seed(batch.seed) - elif batch.generator is not None: - if isinstance(batch.generator, list): - audio_generator = batch.generator[0] - else: - audio_generator = batch.generator - if audio_generator is not None and audio_generator.device.type != latents.device.type: - if batch.seed is None: - audio_generator = torch.Generator(device=latents.device) - else: - audio_generator = torch.Generator( - device=latents.device).manual_seed(batch.seed) - audio_patch_shape = ( + expected_shape = ( audio_shape.batch, + audio_shape.channels, audio_shape.frames, - audio_shape.channels * audio_shape.mel_bins, + audio_shape.mel_bins, ) - audio_latents_patch = torch.randn( - audio_patch_shape, - generator=audio_generator, + audio_latent_path = fastvideo_args.ltx2_audio_latent_path + audio_latents = self._load_audio_latents( + audio_latent_path, device=latents.device, dtype=latents.dtype, - ) - if hasattr(self.transformer, "audio_patchifier"): - audio_latents = self.transformer.audio_patchifier.unpatchify( - audio_latents_patch, audio_shape) - else: - audio_latents = audio_latents_patch.view( + expected_shape=expected_shape, + ) if audio_latent_path else None + if audio_latents is None: + audio_generator = None + if fastvideo_args.ltx2_initial_latent_path and batch.seed is not None: + audio_generator = torch.Generator( + device=latents.device).manual_seed(batch.seed) + elif batch.generator is not None: + if isinstance(batch.generator, list): + audio_generator = batch.generator[0] + else: + audio_generator = batch.generator + if audio_generator is not None and audio_generator.device.type != latents.device.type: + if batch.seed is None: + audio_generator = torch.Generator(device=latents.device) + else: + audio_generator = torch.Generator( + device=latents.device).manual_seed(batch.seed) + audio_patch_shape = ( audio_shape.batch, audio_shape.frames, - audio_shape.channels, - audio_shape.mel_bins, - ).permute(0, 2, 1, 3).contiguous() + audio_shape.channels * audio_shape.mel_bins, + ) + audio_latents_patch = torch.randn( + audio_patch_shape, + generator=audio_generator, + device=latents.device, + dtype=latents.dtype, + ) + if hasattr(self.transformer, "audio_patchifier"): + audio_latents = self.transformer.audio_patchifier.unpatchify( + audio_latents_patch, audio_shape) + else: + audio_latents = audio_latents_patch.view( + audio_shape.batch, + audio_shape.frames, + audio_shape.channels, + audio_shape.mel_bins, + ).permute(0, 2, 1, 3).contiguous() + if audio_latent_path: + self._save_audio_latents(audio_latent_path, audio_latents) audio_timestep_template = torch.ones( - (latents.shape[0], audio_shape.frames), + (latents.shape[0], audio_shape.frames, 1), device=latents.device, dtype=torch.float32, ) logger.info( "[LTX2] Denoising start: steps=%d dtype=%s guidance=%s " "sigmas_shape=%s latents_shape=%s", - batch.num_inference_steps, + num_steps, target_dtype, - batch.guidance_scale, + guidance_scale, tuple(sigmas.shape), tuple(latents.shape), ) @@ -253,7 +305,7 @@ class LTX2DenoisingStage(PipelineStage): pos_audio = None # Only run negative pass if CFG is enabled - if batch.do_classifier_free_guidance: + if use_cfg: neg_outputs = self.transformer( hidden_states=latents.to(target_dtype), encoder_hidden_states=neg_prompt_embeds, @@ -268,10 +320,10 @@ class LTX2DenoisingStage(PipelineStage): else: neg_denoised = neg_outputs neg_audio = None - pos_denoised = pos_denoised + (batch.guidance_scale - 1) * ( + pos_denoised = pos_denoised + (guidance_scale - 1) * ( pos_denoised - neg_denoised) if pos_audio is not None and neg_audio is not None: - pos_audio = pos_audio + (batch.guidance_scale - + pos_audio = pos_audio + (guidance_scale - 1) * (pos_audio - neg_audio) sigma_value = sigma.to(torch.float32) if isinstance( @@ -297,6 +349,43 @@ class LTX2DenoisingStage(PipelineStage): logger.info("[LTX2] Denoising done.") return batch + def _load_audio_latents( + self, + latent_path: str | None, + *, + device: torch.device, + dtype: torch.dtype, + expected_shape: tuple[int, ...], + ) -> torch.Tensor | None: + if not latent_path: + return None + path = Path(latent_path) + if not path.exists(): + return None + payload = torch.load(path, map_location=device) + if isinstance(payload, dict): + latent = (payload.get("audio_latent") or payload.get("latent") + or payload.get("audio")) + else: + latent = payload + if not torch.is_tensor(latent): + raise TypeError(f"Expected tensor audio latent in {path}") + if tuple(latent.shape) != tuple(expected_shape): + raise ValueError( + f"Audio latent shape mismatch for {path}: expected {expected_shape}, got {tuple(latent.shape)}" + ) + logger.info("[LTX2] Loaded audio latent from %s", path) + return latent.to(device=device, dtype=dtype) + + def _save_audio_latents(self, latent_path: str, + latents: torch.Tensor) -> None: + path = Path(latent_path) + path.parent.mkdir(parents=True, exist_ok=True) + if path.exists(): + return + torch.save({"audio_latent": latents.detach().cpu()}, path) + logger.info("[LTX2] Saved audio latent to %s", path) + def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult: result = VerificationResult() diff --git a/fastvideo/pipelines/stages/ltx2_latent_preparation.py b/fastvideo/pipelines/stages/ltx2_latent_preparation.py index 3459f5f28..a9c50071d 100644 --- a/fastvideo/pipelines/stages/ltx2_latent_preparation.py +++ b/fastvideo/pipelines/stages/ltx2_latent_preparation.py @@ -3,6 +3,7 @@ Latent preparation stage for LTX-2 pipelines. """ +import math from pathlib import Path import torch @@ -11,6 +12,7 @@ from diffusers.utils.torch_utils import randn_tensor from fastvideo.distributed import get_local_torch_device from fastvideo.fastvideo_args import FastVideoArgs from fastvideo.logger import init_logger +from fastvideo.models.dits.ltx2 import VideoLatentShape from fastvideo.pipelines.pipeline_batch_info import ForwardBatch from fastvideo.pipelines.stages.base import PipelineStage from fastvideo.pipelines.stages.validators import StageValidators as V @@ -61,6 +63,26 @@ class LTX2LatentPreparationStage(PipelineStage): dtype = batch.prompt_embeds[0].dtype device = get_local_torch_device() generator = batch.generator + if generator is not None: + if isinstance(generator, list): + if generator and generator[0].device.type != device.type: + seeds = batch.seeds + if seeds is None and batch.seed is not None: + seeds = [batch.seed + i for i in range(len(generator))] + if seeds is not None: + generator = [ + torch.Generator(device=device).manual_seed(seed) + for seed in seeds + ] + batch.generator = generator + else: + if generator.device.type != device.type: + if batch.seed is not None: + generator = torch.Generator(device=device).manual_seed( + batch.seed) + else: + generator = torch.Generator(device=device) + batch.generator = generator latents = batch.latents num_frames = latent_num_frames if latent_num_frames is not None else batch.num_frames height = batch.height @@ -95,16 +117,16 @@ class LTX2LatentPreparationStage(PipelineStage): if loaded_latents is not None: latents = loaded_latents else: - latents = randn_tensor( - shape, + latents = self._generate_initial_latents( + shape=shape, generator=generator, device=device, dtype=dtype, ) self._save_initial_latent(latent_path, latents) else: - latents = randn_tensor( - shape, + latents = self._generate_initial_latents( + shape=shape, generator=generator, device=device, dtype=dtype, @@ -116,6 +138,39 @@ class LTX2LatentPreparationStage(PipelineStage): batch.raw_latent_shape = shape return batch + def _generate_initial_latents( + self, + shape: tuple[int, int, int, int, int], + *, + generator: torch.Generator | list[torch.Generator] | None, + device: torch.device, + dtype: torch.dtype, + ) -> torch.Tensor: + patchifier = getattr(self.transformer, "patchifier", None) + if patchifier is None: + return randn_tensor( + shape, + generator=generator, + device=device, + dtype=dtype, + ) + + video_shape = VideoLatentShape.from_torch_shape(shape) + patch_volume = math.prod(patchifier.patch_size) + token_count = patchifier.get_token_count(video_shape) + patch_shape = ( + shape[0], + token_count, + shape[1] * patch_volume, + ) + patch_noise = randn_tensor( + patch_shape, + generator=generator, + device=device, + dtype=dtype, + ) + return patchifier.unpatchify(patch_noise, video_shape) + def _adjust_video_length(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> int | None: if not fastvideo_args.pipeline_config.vae_config.use_temporal_scaling_frames: diff --git a/fastvideo/pipelines/stages/ltx2_refine.py b/fastvideo/pipelines/stages/ltx2_refine.py new file mode 100644 index 000000000..409da24e8 --- /dev/null +++ b/fastvideo/pipelines/stages/ltx2_refine.py @@ -0,0 +1,336 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +LTX-2 refinement stages for 2x spatial upscaling + distilled denoising. +""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any +import weakref +from diffusers.utils.torch_utils import randn_tensor +import torch + +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.logger import init_logger +from fastvideo.models.dits.ltx2 import AudioLatentShape, VideoLatentShape +from fastvideo.models.upsamplers import upsample_video +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.pipelines.stages.base import PipelineStage +from fastvideo.pipelines.stages.validators import StageValidators as V +from fastvideo.pipelines.stages.validators import VerificationResult + +logger = init_logger(__name__) + +# Reduced schedule for super-resolution stage 2 (subset of distilled values) +# From LTX-2/packages/ltx-pipelines/src/ltx_pipelines/utils/constants.py +STAGE_2_DISTILLED_SIGMA_VALUES = [0.909375, 0.725, 0.421875, 0.0] + + +class LTX2RefineInitStage(PipelineStage): + """Prepare stage-1 resolution for LTX-2 2x spatial refinement.""" + + def forward( + self, + batch: ForwardBatch, + fastvideo_args: FastVideoArgs, + ) -> ForwardBatch: + if not fastvideo_args.ltx2_refine_enabled: + return batch + + height = batch.height + width = batch.width + if height is None or width is None: + raise ValueError( + "Height and width must be provided for LTX-2 refinement.") + if isinstance(height, list) or isinstance(width, list): + raise ValueError("LTX-2 refinement expects scalar height/width.") + + if height % 2 != 0 or width % 2 != 0: + raise ValueError( + "LTX-2 refinement requires even height/width so stage1 can be half resolution." + ) + + spatial_ratio = fastvideo_args.pipeline_config.vae_config.arch_config.spatial_compression_ratio + stage1_height = height // 2 + stage1_width = width // 2 + if stage1_height % spatial_ratio != 0 or stage1_width % spatial_ratio != 0: + raise ValueError( + f"LTX-2 refinement requires height/width divisible by {2 * spatial_ratio} " + f"(got {height}x{width}).") + + batch.extra["ltx2_refine_target_height"] = height + batch.extra["ltx2_refine_target_width"] = width + batch.height = stage1_height + batch.width = stage1_width + + logger.info( + "[LTX2] Refinement enabled: stage1=%dx%d stage2=%dx%d", + stage1_width, + stage1_height, + width, + height, + ) + return batch + + def verify_output(self, batch: ForwardBatch, + fastvideo_args: FastVideoArgs) -> VerificationResult: + result = VerificationResult() + # if fastvideo_args.ltx2_refine_enabled: + # result.add_check( + # "height", batch.height, [V.is_int, V.positive_int]) + # result.add_check( + # "width", batch.width, [V.is_int, V.positive_int]) + return result + + +class LTX2UpsampleStage(PipelineStage): + """Upsample stage-1 latents to stage-2 resolution and add refinement noise.""" + + def __init__( + self, + *, + upsampler: Any, + vae: Any, + transformer: Any | None = None, + sigmas: list[float] | None = None, + add_noise: bool = True, + ) -> None: + super().__init__() + self.upsampler = upsampler + self.vae = vae + self.transformer = transformer + self.sigmas = sigmas or STAGE_2_DISTILLED_SIGMA_VALUES + self.add_noise = add_noise + + def forward( + self, + batch: ForwardBatch, + fastvideo_args: FastVideoArgs, + ) -> ForwardBatch: + if not fastvideo_args.ltx2_refine_enabled: + return batch + + if batch.latents is None: + raise ValueError( + "Latents must be available before LTX-2 upsample stage.") + + target_height = batch.extra.get("ltx2_refine_target_height") + target_width = batch.extra.get("ltx2_refine_target_width") + if target_height is None or target_width is None: + raise ValueError("Missing target resolution for LTX-2 refinement.") + + video_encoder = getattr(self.vae, "encoder", None) + if video_encoder is None: + raise ValueError( + "LTX-2 VAE encoder is required for latent upsampling.") + + upsampler_module = getattr(self.upsampler, "model", self.upsampler) + latents = upsample_video(batch.latents, video_encoder, upsampler_module) + + sigma0 = float(self.sigmas[0]) if self.sigmas else 1.0 + if self.add_noise: + patchifier = getattr(self.transformer, "patchifier", None) + if patchifier is not None: + video_shape = VideoLatentShape.from_torch_shape(latents.shape) + latents_patch = patchifier.patchify(latents) + noise_shape = latents_patch.shape + noise_path = fastvideo_args.ltx2_refine_noise_path + noise = self._load_noise( + noise_path, + device=latents.device, + dtype=latents.dtype, + expected_shape=noise_shape, + alternate_shape=latents.shape, + ) if noise_path else None + if noise is None: + noise = randn_tensor( + noise_shape, + generator=batch.generator, + device=latents.device, + dtype=latents.dtype, + ) + if noise_path: + self._save_noise(noise_path, noise) + elif noise.shape == latents.shape: + noise = patchifier.patchify(noise) + noised_patch = noise * sigma0 + latents_patch * (1.0 - sigma0) + latents = patchifier.unpatchify(noised_patch, video_shape) + else: + noise_path = fastvideo_args.ltx2_refine_noise_path + noise = self._load_noise( + noise_path, + device=latents.device, + dtype=latents.dtype, + expected_shape=latents.shape, + ) if noise_path else None + if noise is None: + noise = randn_tensor( + latents.shape, + generator=batch.generator, + device=latents.device, + dtype=latents.dtype, + ) + if noise_path: + self._save_noise(noise_path, noise) + # Match LTX-2 GaussianNoiser: noise * sigma + latent * (1 - sigma). + latents = noise * sigma0 + latents * (1.0 - sigma0) + + audio_latents = batch.extra.get("ltx2_audio_latents") + if audio_latents is not None: + audio_latents = audio_latents.to(device=latents.device) + if self.add_noise: + audio_patchifier = getattr(self.transformer, "audio_patchifier", + None) + if audio_patchifier is not None: + audio_shape = AudioLatentShape.from_torch_shape( + audio_latents.shape) + audio_patch = audio_patchifier.patchify(audio_latents) + audio_noise_shape = audio_patch.shape + audio_noise_path = fastvideo_args.ltx2_refine_audio_noise_path + audio_noise = self._load_noise( + audio_noise_path, + device=audio_latents.device, + dtype=audio_latents.dtype, + expected_shape=audio_noise_shape, + alternate_shape=audio_latents.shape, + ) if audio_noise_path else None + if audio_noise is None: + audio_noise = randn_tensor( + audio_noise_shape, + generator=batch.generator, + device=audio_latents.device, + dtype=audio_latents.dtype, + ) + if audio_noise_path: + self._save_noise(audio_noise_path, audio_noise) + elif audio_noise.shape == audio_latents.shape: + audio_noise = audio_patchifier.patchify(audio_noise) + audio_noised_patch = audio_noise * sigma0 + audio_patch * ( + 1.0 - sigma0) + audio_latents = audio_patchifier.unpatchify( + audio_noised_patch, audio_shape) + else: + audio_noise_path = fastvideo_args.ltx2_refine_audio_noise_path + audio_noise = self._load_noise( + audio_noise_path, + device=audio_latents.device, + dtype=audio_latents.dtype, + expected_shape=audio_latents.shape, + ) if audio_noise_path else None + if audio_noise is None: + audio_noise = randn_tensor( + audio_latents.shape, + generator=batch.generator, + device=audio_latents.device, + dtype=audio_latents.dtype, + ) + if audio_noise_path: + self._save_noise(audio_noise_path, audio_noise) + # Same noise mixing as video latents for distilled refinement. + audio_latents = audio_noise * sigma0 + audio_latents * ( + 1.0 - sigma0) + batch.extra["ltx2_audio_latents"] = audio_latents + + batch.latents = latents + batch.raw_latent_shape = latents.shape + batch.height = target_height + batch.width = target_width + return batch + + def verify_input(self, batch: ForwardBatch, + fastvideo_args: FastVideoArgs) -> VerificationResult: + result = VerificationResult() + if fastvideo_args.ltx2_refine_enabled: + result.add_check("latents", batch.latents, + [V.is_tensor, V.with_dims(5)]) + return result + + def _load_noise( + self, + noise_path: str | None, + *, + device: torch.device, + dtype: torch.dtype, + expected_shape: torch.Size | tuple[int, ...], + alternate_shape: torch.Size | tuple[int, ...] | None = None, + ) -> torch.Tensor | None: + if not noise_path: + return None + path = Path(noise_path) + if not path.exists(): + return None + payload = torch.load(path, map_location=device) + if isinstance(payload, dict): + noise = (payload.get("noise") or payload.get("latent_noise") + or payload.get("latent") or payload.get("video_noise")) + else: + noise = payload + if not torch.is_tensor(noise): + raise TypeError(f"Expected tensor noise in {path}") + noise_shape = tuple(noise.shape) + if noise_shape != tuple(expected_shape) and (alternate_shape is None + or noise_shape + != tuple(alternate_shape)): + raise ValueError( + f"Noise shape mismatch for {path}: expected {tuple(expected_shape)}, got {noise_shape}" + ) + logger.info("[LTX2] Loaded refine noise from %s", path) + return noise.to(device=device, dtype=dtype) + + def _save_noise(self, noise_path: str, noise: torch.Tensor) -> None: + path = Path(noise_path) + path.parent.mkdir(parents=True, exist_ok=True) + if path.exists(): + return + torch.save({"noise": noise.detach().cpu()}, path) + logger.info("[LTX2] Saved refine noise to %s", path) + + +class LTX2RefineLoRAStage(PipelineStage): + """Apply a refinement-specific LoRA before stage-2 denoising.""" + + def __init__( + self, + *, + pipeline: Any, + lora_path: str | None, + lora_nickname: str = "ltx2_refine", + ) -> None: + super().__init__() + self._pipeline_ref = weakref.ref( + pipeline) if pipeline is not None else None + self._lora_path = lora_path + self._lora_nickname = lora_nickname + self._applied = False + + def forward( + self, + batch: ForwardBatch, + fastvideo_args: FastVideoArgs, + ) -> ForwardBatch: + if not fastvideo_args.ltx2_refine_enabled: + return batch + lora_path = fastvideo_args.ltx2_refine_lora_path or self._lora_path + if not lora_path or self._applied: + return batch + + pipeline = self._pipeline_ref( + ) if self._pipeline_ref is not None else None + if pipeline is None or not hasattr(pipeline, "set_lora_adapter"): + raise ValueError( + "LTX2 refinement LoRA requested but pipeline does not support LoRA adapters." + ) + + pipeline.set_lora_adapter(self._lora_nickname, lora_path) + self._applied = True + logger.info("[LTX2] Applied refinement LoRA from %s", lora_path) + return batch + + +__all__ = [ + "STAGE_2_DISTILLED_SIGMA_VALUES", + "LTX2RefineInitStage", + "LTX2UpsampleStage", + "LTX2RefineLoRAStage", +] diff --git a/tests/local_tests/transformers/test_ltx2.py b/tests/local_tests/transformers/test_ltx2.py index 14a3d8366..32114f2a0 100644 --- a/tests/local_tests/transformers/test_ltx2.py +++ b/tests/local_tests/transformers/test_ltx2.py @@ -9,8 +9,11 @@ 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)) @@ -32,7 +35,7 @@ def _read_transformer_config(config: dict) -> dict: def _infer_patch_params(in_channels: int) -> tuple[int, int]: patch_size = 1 num_channels_latents = 128 - for candidate in (8, 16, 32, 64, 128): + for candidate in (128, 64, 32, 16, 8): if in_channels % candidate != 0: continue patch_volume = in_channels // candidate @@ -161,17 +164,27 @@ def test_ltx2_transformer_parity(): pytest.skip(f"FastVideo converted weights not found at {fastvideo_path}") try: - from ltx_core.components.patchifiers import VideoLatentPatchifier + from ltx_core.components.patchifiers import (VideoLatentPatchifier, + get_pixel_coords) from ltx_core.guidance.perturbations import BatchedPerturbationConfig from ltx_core.loader.sft_loader import SafetensorsModelStateDictLoader from ltx_core.loader.single_gpu_model_builder import SingleGPUModelBuilder from ltx_core.model.transformer import (LTXModelConfigurator, LTXV_MODEL_COMFY_RENAMING_MAP) from ltx_core.model.transformer.modality import Modality - from ltx_core.types import VideoLatentShape + from ltx_core.types import VideoLatentShape, VIDEO_SCALE_FACTORS + from ltx_core.utils import to_denoised except ImportError as exc: pytest.skip(f"LTX-2 import failed: {exc}") + # Force reference attention to use PyTorch SDPA for parity with FastVideo. + try: + from ltx_core.model.transformer import attention as ltx_attention + ltx_attention.memory_efficient_attention = None + ltx_attention.flash_attn_interface = None + except Exception: + pass + config_loader = SafetensorsModelStateDictLoader() metadata = config_loader.metadata(str(official_path)) transformer_config = _read_transformer_config(metadata) @@ -304,11 +317,21 @@ def test_ltx2_transformer_parity(): device=device, dtype=precision, ) - timestep = torch.tensor([500], device=device, dtype=precision) video_shape = VideoLatentShape.from_torch_shape(hidden_states.shape) - positions = patchifier.get_patch_grid_bounds(video_shape, device=hidden_states.device) + latent_coords = patchifier.get_patch_grid_bounds(video_shape, device=hidden_states.device) + positions = get_pixel_coords( + latent_coords=latent_coords, + scale_factors=VIDEO_SCALE_FACTORS, + causal_fix=True, + ) latents = patchifier.patchify(hidden_states) + timestep = torch.full( + (batch_size, latents.shape[1], 1), + 500, + device=device, + dtype=precision, + ) video = Modality( enabled=True, @@ -325,6 +348,7 @@ def test_ltx2_transformer_parity(): audio=None, perturbations=BatchedPerturbationConfig.empty(batch_size), ) + ref_out = to_denoised(latents, ref_out, timestep) ref_out = patchifier.unpatchify(ref_out, output_shape=video_shape) print(f"[LTX2 TEST] Reference model output shape: {ref_out.shape}") with set_forward_context( diff --git a/tests/local_tests/transformers/test_ltx2_audio.py b/tests/local_tests/transformers/test_ltx2_audio.py index c67c8cb74..cf839dd24 100644 --- a/tests/local_tests/transformers/test_ltx2_audio.py +++ b/tests/local_tests/transformers/test_ltx2_audio.py @@ -14,6 +14,8 @@ 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)) @@ -55,17 +57,33 @@ def test_ltx2_transformer_audio_parity(): pytest.skip(f"FastVideo converted weights not found at {fastvideo_path}") try: - from ltx_core.components.patchifiers import AudioPatchifier, VideoLatentPatchifier + from ltx_core.components.patchifiers import ( + AudioPatchifier, + VideoLatentPatchifier, + get_pixel_coords, + ) from ltx_core.guidance.perturbations import BatchedPerturbationConfig from ltx_core.loader.sft_loader import SafetensorsModelStateDictLoader from ltx_core.loader.single_gpu_model_builder import SingleGPUModelBuilder from ltx_core.model.transformer import (LTXModelConfigurator, LTXV_MODEL_COMFY_RENAMING_MAP) from ltx_core.model.transformer.modality import Modality - from ltx_core.types import AudioLatentShape, VideoLatentShape + from ltx_core.types import ( + AudioLatentShape, + VideoLatentShape, + VIDEO_SCALE_FACTORS, + ) except ImportError as exc: pytest.skip(f"LTX-2 import failed: {exc}") + # Force reference attention to use PyTorch SDPA for parity with FastVideo. + try: + from ltx_core.model.transformer import attention as ltx_attention + ltx_attention.memory_efficient_attention = None + ltx_attention.flash_attn_interface = None + except Exception: + pass + # Load config from metadata using same approach as test_ltx2.py config_loader = SafetensorsModelStateDictLoader() metadata = config_loader.metadata(str(official_path)) @@ -201,11 +219,22 @@ def test_ltx2_transformer_audio_parity(): device=device, dtype=precision, ) - timestep = torch.tensor([500], device=device, dtype=precision) + timestep = None video_shape = VideoLatentShape.from_torch_shape(hidden_states.shape) - positions = patchifier.get_patch_grid_bounds(video_shape, device=hidden_states.device) + latent_coords = patchifier.get_patch_grid_bounds(video_shape, device=hidden_states.device) + positions = get_pixel_coords( + latent_coords=latent_coords, + scale_factors=VIDEO_SCALE_FACTORS, + causal_fix=True, + ) latents = patchifier.patchify(hidden_states) + timestep = torch.full( + (batch_size, latents.shape[1], 1), + 500, + device=device, + dtype=precision, + ) audio_frames = 16 audio_channels = 8 @@ -221,6 +250,12 @@ def test_ltx2_transformer_audio_parity(): audio_shape = AudioLatentShape.from_torch_shape(audio_latents.shape) audio_positions = audio_patchifier.get_patch_grid_bounds(audio_shape, device=audio_latents.device) audio_tokens = audio_patchifier.patchify(audio_latents) + audio_timestep = torch.full( + (batch_size, audio_tokens.shape[1], 1), + 500, + device=device, + dtype=precision, + ) video = Modality( enabled=True, @@ -233,7 +268,7 @@ def test_ltx2_transformer_audio_parity(): audio = Modality( enabled=True, latent=audio_tokens, - timesteps=timestep, + timesteps=audio_timestep, positions=audio_positions, context=encoder_hidden_states, context_mask=None, @@ -250,7 +285,7 @@ def test_ltx2_transformer_audio_parity(): fastvideo_audio = FastVideoModality( enabled=True, latent=audio_tokens, - timesteps=timestep, + timesteps=audio_timestep, positions=audio_positions, context=encoder_hidden_states, context_mask=None, diff --git a/tests/local_tests/vaes/test_ltx2_vae_official.py b/tests/local_tests/vaes/test_ltx2_vae_official.py index c02da5875..a6b0b3387 100644 --- a/tests/local_tests/vaes/test_ltx2_vae_official.py +++ b/tests/local_tests/vaes/test_ltx2_vae_official.py @@ -8,6 +8,8 @@ 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))