update

Updated the number of GPUs used for training from 1 to 4.

.github/workflows/fastvideo-kernel-publish.yml

@@ -156,23 +156,8 @@ jobs:
           
           # Fix the wheel to be manylinux compliant
           pip install auditwheel
-          # Point auditwheel at torch libs, but do not vendor them into the wheel.
-          TORCH_LIB_DIR=$(python - <<'PY'
-          import os
-          import torch
-
-          print(os.path.join(os.path.dirname(torch.__file__), "lib"))
-          PY
-          )
-          export LD_LIBRARY_PATH="${TORCH_LIB_DIR}:${LD_LIBRARY_PATH}"
           # Target manylinux_2_35 (Ubuntu 22.04 native)
-          auditwheel repair dist/*.whl --plat manylinux_2_35_x86_64 -w fixed_dist \
-            --exclude libtorch_cuda.so \
-            --exclude libtorch_cpu.so \
-            --exclude libtorch.so \
-            --exclude libc10.so \
-            --exclude libc10_cuda.so \
-            --exclude libtorch_python.so
+          auditwheel repair dist/*.whl --plat manylinux_2_35_x86_64 -w fixed_dist
           # Move fixed wheels back to dist for upload consistency
           rm dist/*.whl
           mv fixed_dist/*.whl dist/

examples/inference/basic/basic_ltx2.py

@@ -1,34 +0,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."
-)
-
-
-def main() -> None:
-    generator = VideoGenerator.from_pretrained(
-        "FastVideo/LTX2-Distilled-Diffusers",
-        num_gpus=1,
-    )
-
-    output_path = "outputs_video/ltx2_basic/output_ltx2_distilled_t2v.mp4"
-    generator.generate_video(
-        prompt=PROMPT,
-        output_path=output_path,
-        save_video=True,
-    )
-    generator.shutdown()
-
-
-if __name__ == "__main__":
-    main()

examples/training/finetune/hunyuan15_t2v/download_dataset.sh

@@ -0,0 +1,3 @@
+#!/bin/bash
+
+python scripts/huggingface/download_hf.py --repo_id "wlsaidhi/crush-smol-merged" --local_dir "data/crush-smol" --repo_type "dataset"

examples/training/finetune/hunyuan15_t2v/preprocess_hunyuan_data_t2v.sh

@@ -0,0 +1,25 @@
+#!/bin/bash
+
+GPU_NUM=1 # 2,4,8
+MODEL_PATH="hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v"
+DATASET_PATH="data/crush-smol"
+OUTPUT_DIR="data/crush-smol_processed_t2v_hunyuan15/"
+
+torchrun --nproc_per_node=$GPU_NUM --master_port=29513 \
+    -m fastvideo.pipelines.preprocess.v1_preprocessing_new \
+    --model_path $MODEL_PATH \
+    --mode preprocess \
+    --workload_type t2v \
+    --preprocess.video_loader_type torchvision \
+    --preprocess.dataset_type merged \
+    --preprocess.dataset_path $DATASET_PATH \
+    --preprocess.dataset_output_dir $OUTPUT_DIR \
+    --preprocess.preprocess_video_batch_size 2 \
+    --preprocess.dataloader_num_workers 0 \
+    --preprocess.max_height 480 \
+    --preprocess.max_width 832 \
+    --preprocess.num_frames 61 \
+    --preprocess.train_fps 15 \
+    --preprocess.samples_per_file 8 \
+    --preprocess.flush_frequency 8 \
+    --preprocess.video_length_tolerance_range 5

examples/training/finetune/hunyuan_t2v/download_dataset.sh

@@ -0,0 +1,3 @@
+#!/bin/bash
+
+python scripts/huggingface/download_hf.py --repo_id "wlsaidhi/crush-smol-merged" --local_dir "data/crush-smol" --repo_type "dataset"

examples/training/finetune/hunyuan_t2v/preprocess_hunyuan_data_t2v.sh

@@ -0,0 +1,25 @@
+#!/bin/bash
+
+GPU_NUM=1 # 2,4,8
+MODEL_PATH="hunyuanvideo-community/HunyuanVideo"
+DATASET_PATH="data/crush-smol"
+OUTPUT_DIR="data/crush-smol_processed_t2v_hunyuan/"
+
+torchrun --nproc_per_node=$GPU_NUM \
+    -m fastvideo.pipelines.preprocess.v1_preprocessing_new \
+    --model_path $MODEL_PATH \
+    --mode preprocess \
+    --workload_type t2v \
+    --preprocess.video_loader_type torchvision \
+    --preprocess.dataset_type merged \
+    --preprocess.dataset_path $DATASET_PATH \
+    --preprocess.dataset_output_dir $OUTPUT_DIR \
+    --preprocess.preprocess_video_batch_size 2 \
+    --preprocess.dataloader_num_workers 0 \
+    --preprocess.max_height 480 \
+    --preprocess.max_width 832 \
+    --preprocess.num_frames 77 \
+    --preprocess.train_fps 16 \
+    --preprocess.samples_per_file 8 \
+    --preprocess.flush_frequency 8 \
+    --preprocess.video_length_tolerance_range 5

fastvideo-kernel/csrc/turbodiffusion/norm/layernorm.cu

@@ -4,7 +4,6 @@
 #include <torch/all.h>
 #include <torch/python.h>
 #include <cutlass/cutlass.h>
-#include <cutlass/numeric_types.h>
 #include "common/common.hpp"
 #include "norm/layernorm.hpp"
 
@@ -15,6 +14,10 @@ auto layer_norm(
   std::optional<at::Tensor const> const B,
   std::optional<at::Tensor> Output
 ) {
+  using ElementIn = float;
+  using ElementOut = float;
+  using ElementWeight = float;
+
   int64_t const m = Input.size(0);
   int64_t const n = Input.size(1);
   torch::Device const input_device = Input.device();
@@ -23,70 +26,31 @@ auto layer_norm(
     Output.emplace(
       torch::empty(
         {m, n},
-        torch::TensorOptions().device(input_device).dtype(Input.scalar_type())
+        torch::TensorOptions().device(input_device).dtype(torch::kFloat32)
       )
     );
   }
 
-  TORCH_CHECK(Output.value().scalar_type() == Input.scalar_type(),
-              "Output dtype must match Input dtype. Got Output=",
-              Output.value().scalar_type(), ", Input=", Input.scalar_type());
-  if (W.has_value()) {
-    TORCH_CHECK(W.value().scalar_type() == Input.scalar_type(),
-                "W dtype must match Input dtype. Got W=",
-                W.value().scalar_type(), ", Input=", Input.scalar_type());
-  }
-  if (B.has_value()) {
-    TORCH_CHECK(B.value().scalar_type() == Input.scalar_type(),
-                "B dtype must match Input dtype. Got B=",
-                B.value().scalar_type(), ", Input=", Input.scalar_type());
-  }
 
   void *Iptr = Input.data_ptr();
   void *Wptr = W.has_value() ? W.value().data_ptr() : nullptr;
   void *Bptr = B.has_value() ? B.value().data_ptr() : nullptr;
   void *Optr = Output.value().data_ptr();
 
-  auto stream = at::cuda::getCurrentCUDAStream().stream();
-  if (Input.scalar_type() == at::kHalf) {
-    using ElementIn = cutlass::half_t;
-    using ElementOut = cutlass::half_t;
-    using ElementWeight = cutlass::half_t;
-    BOOL_SWITCH(B.has_value(), BIAS, [&]{
-      BOOL_SWITCH(W.has_value(), AFFINE, [&]{
-        CONFIG_SWITCH(n, [&]{
-          layernorm<ElementIn, ElementOut, ElementWeight, AFFINE, BIAS, MAX_HIDDEN_SIZE, NUM_THR_PER_CTA>(
-              Iptr, Wptr, Bptr, Optr, eps, m, n, stream);
-        });
+  BOOL_SWITCH(B.has_value(), BIAS, [&]{
+    BOOL_SWITCH(W.has_value(), AFFINE, [&]{
+    CONFIG_SWITCH(n, [&]{
+    layernorm<
+    ElementIn, ElementOut, ElementWeight,
+    AFFINE, BIAS,
+    MAX_HIDDEN_SIZE, NUM_THR_PER_CTA> (
+      Iptr, Wptr, Bptr,
+      Optr, eps, m, n,
+      at::cuda::getCurrentCUDAStream().stream()
+        );
       });
     });
-  } else if (Input.scalar_type() == at::kBFloat16) {
-    using ElementIn = cutlass::bfloat16_t;
-    using ElementOut = cutlass::bfloat16_t;
-    using ElementWeight = cutlass::bfloat16_t;
-    BOOL_SWITCH(B.has_value(), BIAS, [&]{
-      BOOL_SWITCH(W.has_value(), AFFINE, [&]{
-        CONFIG_SWITCH(n, [&]{
-          layernorm<ElementIn, ElementOut, ElementWeight, AFFINE, BIAS, MAX_HIDDEN_SIZE, NUM_THR_PER_CTA>(
-              Iptr, Wptr, Bptr, Optr, eps, m, n, stream);
-        });
-      });
-    });
-  } else if (Input.scalar_type() == at::kFloat) {
-    using ElementIn = float;
-    using ElementOut = float;
-    using ElementWeight = float;
-    BOOL_SWITCH(B.has_value(), BIAS, [&]{
-      BOOL_SWITCH(W.has_value(), AFFINE, [&]{
-        CONFIG_SWITCH(n, [&]{
-          layernorm<ElementIn, ElementOut, ElementWeight, AFFINE, BIAS, MAX_HIDDEN_SIZE, NUM_THR_PER_CTA>(
-              Iptr, Wptr, Bptr, Optr, eps, m, n, stream);
-        });
-      });
-    });
-  } else {
-    TORCH_CHECK(false, "Unsupported dtype for layer_norm_cuda: ", Input.scalar_type());
-  }
+  });
   
     
 

fastvideo-kernel/csrc/turbodiffusion/norm/layernorm.hpp

@@ -68,22 +68,9 @@ public:
     // mean reduction
     float u = _reduce_sum(x, shared_data) / params.n;
 
-    // IMPORTANT:
-    // Loader pads out-of-range lanes with 0. That is OK for the sum, but after
-    // subtracting mean, those padded lanes become -u and would incorrectly
-    // contribute to the variance. Mask them back to 0 before variance reduction.
-    // We launch exactly NumThrPerCta threads for a 1xMaxHiddenSize tile,
-    // so each thread is responsible for a contiguous chunk in N.
-    int thr_n_offset = tidx * NumElementPerThread;
     CUTLASS_PRAGMA_UNROLL
-    for (int i = 0; i < NumElementPerThread; ++i) {
-      int idx = thr_n_offset + i;
-      if (idx < params.n) {
-        x[i] -= u;
-      } else {
-        x[i] = 0.f;
-      }
-    }
+    for (int i = 0; i < NumElementPerThread; ++i)
+      x[i] -= u;
 
     __syncthreads();
     // var reduction

fastvideo-kernel/csrc/turbodiffusion/norm/rmsnorm.cu

@@ -4,7 +4,6 @@
 #include <torch/all.h>
 #include <torch/python.h>
 #include <cutlass/cutlass.h>
-#include <cutlass/numeric_types.h>
 #include <pybind11/pybind11.h>
 
 #include "common/common.hpp"
@@ -17,6 +16,10 @@ auto rms_norm(
   std::optional<at::Tensor>& Output
 ) {
   
+  using ElementIn = float;
+  using ElementOut = float;
+  using ElementWeight = float;
+  
   int64_t const m = Input.size(0);
   int64_t const n = Input.size(1);
   torch::Device const input_device = Input.device();
@@ -25,51 +28,27 @@ auto rms_norm(
     Output.emplace(
       torch::empty(
         {m, n},
-        torch::TensorOptions().device(input_device).dtype(Input.scalar_type())
+        torch::TensorOptions().device(input_device).dtype(torch::kFloat32)
       )
     );
   }
 
-  TORCH_CHECK(Output.value().scalar_type() == Input.scalar_type(),
-              "Output dtype must match Input dtype. Got Output=",
-              Output.value().scalar_type(), ", Input=", Input.scalar_type());
-  if (Weight.has_value()) {
-    TORCH_CHECK(Weight.value().scalar_type() == Input.scalar_type(),
-                "Weight dtype must match Input dtype. Got Weight=",
-                Weight.value().scalar_type(), ", Input=", Input.scalar_type());
-  }
-
   void *Iptr = Input.data_ptr();
   void *Wptr = Weight.has_value() ? Weight.value().data_ptr() : nullptr;
   void *Optr = Output.value().data_ptr();
 
-  if (Input.scalar_type() == at::kHalf) {
-    using ElementIn = cutlass::half_t;
-    using ElementOut = cutlass::half_t;
-    using ElementWeight = cutlass::half_t;
-    CONFIG_SWITCH(n, [&]{
-      rmsnorm<ElementIn, ElementOut, ElementWeight, MAX_HIDDEN_SIZE, NUM_THR_PER_CTA>(
-          Iptr, Wptr, Optr, eps, m, n, at::cuda::getCurrentCUDAStream().stream());
-    });
-  } else if (Input.scalar_type() == at::kBFloat16) {
-    using ElementIn = cutlass::bfloat16_t;
-    using ElementOut = cutlass::bfloat16_t;
-    using ElementWeight = cutlass::bfloat16_t;
-    CONFIG_SWITCH(n, [&]{
-      rmsnorm<ElementIn, ElementOut, ElementWeight, MAX_HIDDEN_SIZE, NUM_THR_PER_CTA>(
-          Iptr, Wptr, Optr, eps, m, n, at::cuda::getCurrentCUDAStream().stream());
-    });
-  } else if (Input.scalar_type() == at::kFloat) {
-    using ElementIn = float;
-    using ElementOut = float;
-    using ElementWeight = float;
-    CONFIG_SWITCH(n, [&]{
-      rmsnorm<ElementIn, ElementOut, ElementWeight, MAX_HIDDEN_SIZE, NUM_THR_PER_CTA>(
-          Iptr, Wptr, Optr, eps, m, n, at::cuda::getCurrentCUDAStream().stream());
-    });
-  } else {
-    TORCH_CHECK(false, "Unsupported dtype for rms_norm_cuda: ", Input.scalar_type());
-  }
+
+  CONFIG_SWITCH(n, [&]{
+    rmsnorm<
+    ElementIn, ElementOut, ElementWeight,
+    MAX_HIDDEN_SIZE, NUM_THR_PER_CTA
+    > (
+      Iptr, Wptr,
+      Optr, 
+      eps, m, n, 
+      at::cuda::getCurrentCUDAStream().stream()
+    );
+});
   
 
   return Output;

fastvideo-kernel/pyproject.toml

@@ -9,7 +9,7 @@ build-backend = "scikit_build_core.build"
 
 [project]
 name = "fastvideo-kernel"
-version = "0.2.4"
+version = "0.2.3"
 description = "Unified CUDA kernels for FastVideo"
 readme = "README.md"
 requires-python = ">=3.10"

fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import os
 from typing import Tuple
 
 import torch
@@ -24,12 +23,6 @@ def _is_sm90() -> bool:
     return major == 9 and minor == 0
 
 
-def _force_triton() -> bool:
-    # Force Triton even on SM90 and even if the compiled extension is available.
-    # Useful for CI / debugging / parity testing.
-    return os.environ.get("FASTVIDEO_KERNEL_VSA_FORCE_TRITON", "0") == "1"
-
-
 def _map_to_index_torch(block_map: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
     """
     Pure-torch (no triton) conversion:
@@ -288,7 +281,7 @@ def block_sparse_attn(
     - Otherwise: uses Triton implementation (requires q/k/v to have same padded length today).
     """
     block_sparse_fwd, block_sparse_bwd = _get_sm90_ops()
-    if (not _force_triton()) and _is_sm90() and (block_sparse_fwd is not None) and (block_sparse_bwd is not None):
+    if _is_sm90() and (block_sparse_fwd is not None) and (block_sparse_bwd is not None):
         return block_sparse_attn_sm90(q, k, v, block_map, variable_block_sizes)
     # Triton path: generally assumes q/k/v share the same padded length
     if q.shape[2] != k.shape[2] or q.shape[2] != v.shape[2]:

fastvideo-kernel/python/fastvideo_kernel/ops.py

@@ -46,22 +46,14 @@ def sliding_tile_attention(
     flag = shape_map[seq_shape]
 
     for head_idx, (t, h, w) in enumerate(window_size):
-        # Per-head slices are not contiguous in the batch dimension when batch>1
-        # (they keep the original head-stride). The TK kernel assumes contiguous
-        # [B, H, S, D] layout, so we materialize a contiguous [B,1,S,D] view.
-        q_h = q[:, head_idx:head_idx + 1].contiguous()
-        k_h = k[:, head_idx:head_idx + 1].contiguous()
-        v_h = v[:, head_idx:head_idx + 1].contiguous()
-        o_h = torch.empty_like(q_h)
         sta_fwd(
-            q_h, k_h,
-            v_h, o_h,
+            q[:, head_idx:head_idx + 1], k[:, head_idx:head_idx + 1],
+            v[:, head_idx:head_idx + 1], output[:, head_idx:head_idx + 1],
             t, h, w, text_length, False, has_text, flag
         )
-        output[:, head_idx:head_idx + 1] = o_h
 
     if has_text:
-        sta_fwd(q.contiguous(), k.contiguous(), v.contiguous(), output, 3, 3, 3, text_length, True, True, flag)
+        sta_fwd(q, k, v, output, 3, 3, 3, text_length, True, True, flag)
 
     return output[:, :, :seq_length]
 

fastvideo-kernel/python/fastvideo_kernel/version.py

@@ -1 +1 @@
-__version__ = "0.2.4"
+__version__ = "0.2.3"

fastvideo/configs/models/__init__.py

@@ -2,16 +2,5 @@ from fastvideo.configs.models.base import ModelConfig
 from fastvideo.configs.models.dits.base import DiTConfig
 from fastvideo.configs.models.encoders.base import EncoderConfig
 from fastvideo.configs.models.vaes.base import VAEConfig
-from fastvideo.configs.models.audio import (LTX2AudioDecoderConfig,
-                                            LTX2AudioEncoderConfig,
-                                            LTX2VocoderConfig)
 
-__all__ = [
-    "ModelConfig",
-    "VAEConfig",
-    "DiTConfig",
-    "EncoderConfig",
-    "LTX2AudioEncoderConfig",
-    "LTX2AudioDecoderConfig",
-    "LTX2VocoderConfig",
-]
+__all__ = ["ModelConfig", "VAEConfig", "DiTConfig", "EncoderConfig"]

fastvideo/configs/models/audio/__init__.py

@@ -1,13 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-
-from fastvideo.configs.models.audio.ltx2_audio_vae import (
-    LTX2AudioDecoderConfig,
-    LTX2AudioEncoderConfig,
-    LTX2VocoderConfig,
-)
-
-__all__ = [
-    "LTX2AudioEncoderConfig",
-    "LTX2AudioDecoderConfig",
-    "LTX2VocoderConfig",
-]

fastvideo/configs/models/audio/ltx2_audio_vae.py

@@ -1,31 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""
-LTX-2 audio VAE and vocoder configuration.
-"""
-
-from dataclasses import dataclass, field
-
-from fastvideo.configs.models.base import ArchConfig, ModelConfig
-
-
-@dataclass
-class LTX2AudioArchConfig(ArchConfig):
-    architectures: list[str] = field(default_factory=list)
-
-
-@dataclass
-class LTX2AudioEncoderConfig(ModelConfig):
-    arch_config: ArchConfig = field(default_factory=lambda: LTX2AudioArchConfig(
-        architectures=["LTX2AudioEncoder"]))
-
-
-@dataclass
-class LTX2AudioDecoderConfig(ModelConfig):
-    arch_config: ArchConfig = field(default_factory=lambda: LTX2AudioArchConfig(
-        architectures=["LTX2AudioDecoder"]))
-
-
-@dataclass
-class LTX2VocoderConfig(ModelConfig):
-    arch_config: ArchConfig = field(default_factory=lambda: LTX2AudioArchConfig(
-        architectures=["LTX2Vocoder"]))

fastvideo/configs/models/dits/__init__.py

@@ -3,12 +3,11 @@ from fastvideo.configs.models.dits.cosmos2_5 import Cosmos25VideoConfig
 from fastvideo.configs.models.dits.hunyuanvideo import HunyuanVideoConfig
 from fastvideo.configs.models.dits.hunyuanvideo15 import HunyuanVideo15Config
 from fastvideo.configs.models.dits.longcat import LongCatVideoConfig
-from fastvideo.configs.models.dits.ltx2 import LTX2VideoConfig
 from fastvideo.configs.models.dits.stepvideo import StepVideoConfig
 from fastvideo.configs.models.dits.wanvideo import WanVideoConfig
 
 __all__ = [
     "HunyuanVideoConfig", "HunyuanVideo15Config", "WanVideoConfig",
     "StepVideoConfig", "CosmosVideoConfig", "Cosmos25VideoConfig",
-    "LongCatVideoConfig", "LTX2VideoConfig"
+    "LongCatVideoConfig"
 ]

fastvideo/configs/models/dits/ltx2.py

@@ -1,84 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""
-LTX-2 Transformer configuration for native FastVideo integration.
-"""
-
-from dataclasses import dataclass, field
-
-from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig
-
-
-def is_ltx2_blocks(name: str, _module) -> bool:
-    """FSDP shard condition for LTX-2 transformer blocks."""
-    return "transformer_blocks" in name
-
-
-@dataclass
-class LTX2VideoArchConfig(DiTArchConfig):
-    """Architecture configuration for LTX-2 video transformer."""
-
-    _fsdp_shard_conditions: list = field(
-        default_factory=lambda: [is_ltx2_blocks])
-    _compile_conditions: list = field(default_factory=lambda: [is_ltx2_blocks])
-
-    # Parameter name mapping for weight conversion (hf/comfy -> FastVideo)
-    param_names_mapping: dict = field(
-        default_factory=lambda: {
-            r"^model\.diffusion_model\.(.*)$": r"model.\1",
-            r"^diffusion_model\.(.*)$": r"model.\1",
-            r"^model\.(.*)$": r"model.\1",
-            r"^(.*)$": r"model.\1",
-        })
-
-    reverse_param_names_mapping: dict = field(default_factory=lambda: {})
-    lora_param_names_mapping: dict = field(default_factory=lambda: {})
-
-    # Core transformer settings (defaults from LTX-2 metadata)
-    num_attention_heads: int = 32
-    attention_head_dim: int = 128
-    num_layers: int = 48
-    cross_attention_dim: int = 4096
-    caption_channels: int = 3840
-    norm_eps: float = 1e-6
-    attention_type: str = "default"
-    rope_type: str = "split"
-    double_precision_rope: bool = True
-
-    positional_embedding_theta: float = 10000.0
-    positional_embedding_max_pos: list[int] = field(
-        default_factory=lambda: [20, 2048, 2048])
-    timestep_scale_multiplier: int = 1000
-    use_middle_indices_grid: bool = True
-
-    # Patchification (video-only path)
-    patch_size: tuple[int, int, int] = (1, 1, 1)
-    num_channels_latents: int = 128
-    in_channels: int | None = None
-    out_channels: int | None = None
-
-    # Audio defaults (reserved for joint AV ports)
-    audio_num_attention_heads: int = 32
-    audio_attention_head_dim: int = 64
-    audio_in_channels: int = 128
-    audio_out_channels: int = 128
-    audio_cross_attention_dim: int = 2048
-    audio_positional_embedding_max_pos: list[int] = field(
-        default_factory=lambda: [20])
-    av_ca_timestep_scale_multiplier: int = 1
-
-    def __post_init__(self):
-        super().__post_init__()
-        patch_volume = self.patch_size[0] * self.patch_size[
-            1] * self.patch_size[2]
-        if self.in_channels is None:
-            self.in_channels = self.num_channels_latents * patch_volume
-        if self.out_channels is None:
-            self.out_channels = self.in_channels
-
-
-@dataclass
-class LTX2VideoConfig(DiTConfig):
-    """Main configuration for LTX-2 transformer."""
-
-    arch_config: DiTArchConfig = field(default_factory=LTX2VideoArchConfig)
-    prefix: str = "ltx2"

fastvideo/configs/models/encoders/__init__.py

@@ -8,11 +8,10 @@ from fastvideo.configs.models.encoders.llama import LlamaConfig
 from fastvideo.configs.models.encoders.t5 import T5Config, T5LargeConfig
 from fastvideo.configs.models.encoders.qwen2_5 import Qwen2_5_VLConfig
 from fastvideo.configs.models.encoders.reason1 import Reason1ArchConfig, Reason1Config
-from fastvideo.configs.models.encoders.gemma import LTX2GemmaConfig
 
 __all__ = [
     "EncoderConfig", "TextEncoderConfig", "ImageEncoderConfig",
     "BaseEncoderOutput", "CLIPTextConfig", "CLIPVisionConfig",
     "WAN2_1ControlCLIPVisionConfig", "LlamaConfig", "T5Config", "T5LargeConfig",
-    "Qwen2_5_VLConfig", "Reason1ArchConfig", "Reason1Config", "LTX2GemmaConfig"
+    "Qwen2_5_VLConfig", "Reason1ArchConfig", "Reason1Config"
 ]

fastvideo/configs/models/encoders/clip.py

@@ -44,6 +44,15 @@ class CLIPTextArchConfig(TextEncoderArchConfig):
     _fsdp_shard_conditions: list = field(
         default_factory=lambda: [_is_transformer_layer, _is_embeddings])
 
+    tokenizer_kwargs: dict = field(
+        default_factory=lambda: {
+            "padding": "max_length",
+            "truncation": True,
+            "max_length": 77,
+            "return_tensors": "pt",
+        }
+    )
+
 
 @dataclass
 class CLIPVisionArchConfig(ImageEncoderArchConfig):

fastvideo/configs/models/encoders/gemma.py

@@ -1,48 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-from dataclasses import dataclass, field
-
-from fastvideo.configs.models.encoders.base import (
-    TextEncoderArchConfig,
-    TextEncoderConfig,
-)
-
-
-@dataclass
-class LTX2GemmaArchConfig(TextEncoderArchConfig):
-    architectures: list[str] = field(
-        default_factory=lambda: ["LTX2GemmaTextEncoderModel"])
-    hidden_size: int = 3840
-    num_hidden_layers: int = 48
-    num_attention_heads: int = 30
-    text_len: int = 1024
-    pad_token_id: int = 0
-    eos_token_id: int = 2
-
-    gemma_model_path: str = ""
-    gemma_dtype: str = "bfloat16"
-    padding_side: str = "left"
-
-    feature_extractor_in_features: int = 3840 * 49
-    feature_extractor_out_features: int = 3840
-
-    connector_num_attention_heads: int = 30
-    connector_attention_head_dim: int = 128
-    connector_num_layers: int = 2
-    connector_positional_embedding_theta: float = 10000.0
-    connector_positional_embedding_max_pos: list[int] = field(
-        default_factory=lambda: [4096])
-    connector_rope_type: str = "split"
-    connector_double_precision_rope: bool = False
-    connector_num_learnable_registers: int | None = 128
-
-    def __post_init__(self) -> None:
-        super().__post_init__()
-        self.tokenizer_kwargs["padding"] = "max_length"
-
-
-@dataclass
-class LTX2GemmaConfig(TextEncoderConfig):
-    arch_config: TextEncoderArchConfig = field(
-        default_factory=LTX2GemmaArchConfig)
-
-    prefix: str = "ltx2_gemma"

fastvideo/configs/models/encoders/llama.py

@@ -56,6 +56,15 @@ class LlamaArchConfig(TextEncoderArchConfig):
         default_factory=lambda:
         [_is_transformer_layer, _is_embeddings, _is_final_norm])
 
+    tokenizer_kwargs: dict = field(
+        default_factory=lambda: {
+            "padding": "max_length",
+            "truncation": True,
+            "max_length": 256,
+            "return_tensors": "pt",
+        }
+    )
+
 
 @dataclass
 class LlamaConfig(TextEncoderConfig):

fastvideo/configs/models/vaes/__init__.py

@@ -2,7 +2,6 @@ from fastvideo.configs.models.vaes.cosmosvae import CosmosVAEConfig
 from fastvideo.configs.models.vaes.cosmos2_5vae import Cosmos25VAEConfig
 from fastvideo.configs.models.vaes.hunyuanvae import HunyuanVAEConfig
 from fastvideo.configs.models.vaes.hunyuan15vae import Hunyuan15VAEConfig
-from fastvideo.configs.models.vaes.ltx2vae import LTX2VAEConfig
 from fastvideo.configs.models.vaes.stepvideovae import StepVideoVAEConfig
 from fastvideo.configs.models.vaes.wanvae import WanVAEConfig
 
@@ -13,5 +12,4 @@ __all__ = [
     "CosmosVAEConfig",
     "Cosmos25VAEConfig",
     "Hunyuan15VAEConfig",
-    "LTX2VAEConfig",
 ]

fastvideo/configs/models/vaes/ltx2vae.py

@@ -1,45 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""
-LTX-2 VAE configuration.
-"""
-
-from dataclasses import dataclass, field
-
-from fastvideo.configs.models.vaes.base import VAEArchConfig, VAEConfig
-
-
-@dataclass
-class LTX2VAEArchConfig(VAEArchConfig):
-    # Mirrors LTX-2 safetensors metadata config under "vae"
-    _class_name: str = "CausalVideoAutoencoder"
-    dims: int = 3
-    in_channels: int = 3
-    out_channels: int = 3
-    latent_channels: int = 128
-    encoder_blocks: list = field(default_factory=list)
-    decoder_blocks: list = field(default_factory=list)
-    patch_size: int = 4
-    norm_layer: str = "pixel_norm"
-    latent_log_var: str = "uniform"
-    encoder_spatial_padding_mode: str = "zeros"
-    decoder_spatial_padding_mode: str = "reflect"
-    causal_decoder: bool = False
-    timestep_conditioning: bool = True
-    use_quant_conv: bool = False
-    scaling_factor: float = 1.0
-    normalize_latent_channels: bool = False
-
-    # Match FastVideo naming for compression ratios (LTX-2 default)
-    temporal_compression_ratio: int = 8
-    spatial_compression_ratio: int = 32
-
-
-@dataclass
-class LTX2VAEConfig(VAEConfig):
-    arch_config: VAEArchConfig = field(default_factory=LTX2VAEArchConfig)
-
-    # LTX-2 tiling defaults (match ltx_core.video_vae.TilingConfig.default()).
-    ltx2_spatial_tile_size_in_pixels: int = 512
-    ltx2_spatial_tile_overlap_in_pixels: int = 64
-    ltx2_temporal_tile_size_in_frames: int = 64
-    ltx2_temporal_tile_overlap_in_frames: int = 24

fastvideo/configs/pipelines/__init__.py

@@ -4,7 +4,6 @@ from fastvideo.configs.pipelines.cosmos import CosmosConfig
 from fastvideo.configs.pipelines.cosmos2_5 import Cosmos25Config
 from fastvideo.configs.pipelines.hunyuan import FastHunyuanConfig, HunyuanConfig
 from fastvideo.configs.pipelines.hunyuan15 import Hunyuan15T2V480PConfig, Hunyuan15T2V720PConfig
-from fastvideo.configs.pipelines.ltx2 import LTX2T2VConfig
 from fastvideo.configs.pipelines.registry import (
     get_pipeline_config_cls_from_name)
 from fastvideo.configs.pipelines.stepvideo import StepVideoT2VConfig
@@ -17,6 +16,5 @@ __all__ = [
     "Hunyuan15T2V480PConfig", "Hunyuan15T2V720PConfig", "SlidingTileAttnConfig",
     "WanT2V480PConfig", "WanI2V480PConfig", "WanT2V720PConfig",
     "WanI2V720PConfig", "StepVideoT2VConfig", "SelfForcingWanT2V480PConfig",
-    "CosmosConfig", "Cosmos25Config", "LTX2T2VConfig",
-    "get_pipeline_config_cls_from_name"
+    "CosmosConfig", "Cosmos25Config", "get_pipeline_config_cls_from_name"
 ]

fastvideo/configs/pipelines/ltx2.py

@@ -1,50 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-from collections.abc import Callable
-from dataclasses import dataclass, field
-
-import torch
-
-from fastvideo.configs.models import (DiTConfig, EncoderConfig, ModelConfig,
-                                      LTX2AudioDecoderConfig, LTX2VocoderConfig,
-                                      VAEConfig)
-from fastvideo.configs.models.dits import LTX2VideoConfig
-from fastvideo.configs.models.encoders import BaseEncoderOutput, LTX2GemmaConfig
-from fastvideo.configs.models.vaes import LTX2VAEConfig
-from fastvideo.configs.pipelines.base import PipelineConfig, preprocess_text
-
-
-def ltx2_postprocess_text(outputs: BaseEncoderOutput) -> torch.Tensor:
-    return outputs.last_hidden_state
-
-
-@dataclass
-class LTX2T2VConfig(PipelineConfig):
-    """Configuration for LTX-2 T2V pipeline."""
-
-    dit_config: DiTConfig = field(default_factory=LTX2VideoConfig)
-    vae_config: VAEConfig = field(default_factory=LTX2VAEConfig)
-    vae_tiling: bool = True
-    vae_sp: bool = False
-
-    text_encoder_configs: tuple[EncoderConfig, ...] = field(
-        default_factory=lambda: (LTX2GemmaConfig(), ))
-    preprocess_text_funcs: tuple[Callable[[str], str], ...] = field(
-        default_factory=lambda: (preprocess_text, ))
-    postprocess_text_funcs: tuple[Callable[[BaseEncoderOutput], torch.Tensor],
-                                  ...] = field(default_factory=lambda:
-                                               (ltx2_postprocess_text, ))
-
-    dit_precision: str = "bf16"
-    vae_precision: str = "bf16"
-    text_encoder_precisions: tuple[str, ...] = field(
-        default_factory=lambda: ("bf16", ))
-
-    audio_decoder_config: ModelConfig = field(
-        default_factory=LTX2AudioDecoderConfig)
-    vocoder_config: ModelConfig = field(default_factory=LTX2VocoderConfig)
-    audio_decoder_precision: str = "bf16"
-    vocoder_precision: str = "bf16"
-
-    def __post_init__(self) -> None:
-        self.vae_config.load_encoder = False
-        self.vae_config.load_decoder = True

fastvideo/configs/pipelines/registry.py

@@ -9,7 +9,6 @@ from fastvideo.configs.pipelines.cosmos import CosmosConfig
 from fastvideo.configs.pipelines.cosmos2_5 import Cosmos25Config
 from fastvideo.configs.pipelines.hunyuan import FastHunyuanConfig, HunyuanConfig
 from fastvideo.configs.pipelines.hunyuan15 import Hunyuan15T2V480PConfig, Hunyuan15T2V720PConfig
-from fastvideo.configs.pipelines.ltx2 import LTX2T2VConfig
 from fastvideo.configs.pipelines.stepvideo import StepVideoT2VConfig
 from fastvideo.configs.pipelines.longcat import LongCatT2V480PConfig
 from fastvideo.configs.pipelines.turbodiffusion import (
@@ -65,9 +64,6 @@ PIPE_NAME_TO_CONFIG: dict[str, type[PipelineConfig]] = {
     "FastVideo/LongCat-Video-T2V-Diffusers": LongCatT2V480PConfig,
     "FastVideo/LongCat-Video-I2V-Diffusers": LongCatT2V480PConfig,
     "FastVideo/LongCat-Video-VC-Diffusers": LongCatT2V480PConfig,
-    # LTX-2 models
-    "Lightricks/LTX-2": LTX2T2VConfig,
-    "converted/ltx2_diffusers": LTX2T2VConfig,
     # TurboDiffusion models
     "loayrashid/TurboWan2.1-T2V-1.3B-Diffusers": TurboDiffusionT2V_1_3B_Config,
     "loayrashid/TurboWan2.1-T2V-14B-Diffusers": TurboDiffusionT2V_14B_Config,
@@ -106,8 +102,6 @@ PIPELINE_DETECTOR: dict[str, Callable[[str], bool]] = {
     lambda id: "cosmos25" in id.lower(),
     "turbodiffusion":
     lambda id: "turbodiffusion" in id.lower() or "turbowan" in id.lower(),
-    "ltx2":
-    lambda id: "ltx2" in id.lower() or "ltx-2" in id.lower(),
     # Add other pipeline architecture detectors
 }
 
@@ -129,7 +123,6 @@ PIPELINE_FALLBACK_CONFIG: dict[str, type[PipelineConfig]] = {
     "wancausaldmdpipeline": SelfForcingWanT2V480PConfig,
     "stepvideo": StepVideoT2VConfig,
     "turbodiffusion": TurboDiffusionT2V_1_3B_Config,
-    "ltx2": LTX2T2VConfig,
     # Other fallbacks by architecture
 }
 

fastvideo/configs/sample/ltx2.py

@@ -1,20 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-from dataclasses import dataclass
-
-from fastvideo.configs.sample.base import SamplingParam
-
-
-@dataclass
-class LTX2SamplingParam(SamplingParam):
-    """Default sampling parameters for LTX-2 distilled T2V.
-    """
-
-    seed: int = 10
-    num_frames: int = 121
-    height: int = 1024
-    width: int = 1536
-    fps: int = 24
-    num_inference_steps: int = 8
-    guidance_scale: float = 1.0
-    # No default negative_prompt for distilled models
-    negative_prompt: str = ""

fastvideo/configs/sample/registry.py

@@ -10,7 +10,6 @@ from fastvideo.configs.sample.stepvideo import StepVideoT2VSamplingParam
 
 from fastvideo.configs.sample.cosmos import Cosmos_Predict2_2B_Video2World_SamplingParam
 from fastvideo.configs.sample.cosmos2_5 import Cosmos_Predict2_5_2B_Diffusers_SamplingParam
-from fastvideo.configs.sample.ltx2 import LTX2SamplingParam
 
 # isort: off
 from fastvideo.configs.sample.wan import (
@@ -41,36 +40,48 @@ from fastvideo.utils import (maybe_download_model_index,
 logger = init_logger(__name__)
 # Registry maps specific model weights to their config classes
 SAMPLING_PARAM_REGISTRY: dict[str, Any] = {
-    "FastVideo/FastHunyuan-diffusers": FastHunyuanSamplingParam,
-    "hunyuanvideo-community/HunyuanVideo": HunyuanSamplingParam,
+    "FastVideo/FastHunyuan-diffusers":
+    FastHunyuanSamplingParam,
+    "hunyuanvideo-community/HunyuanVideo":
+    HunyuanSamplingParam,
     "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v":
     Hunyuan15_480P_SamplingParam,
     "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v":
     Hunyuan15_720P_SamplingParam,
-    "FastVideo/stepvideo-t2v-diffusers": StepVideoT2VSamplingParam,
+    "FastVideo/stepvideo-t2v-diffusers":
+    StepVideoT2VSamplingParam,
 
     # Wan2.1
-    "Wan-AI/Wan2.1-T2V-1.3B-Diffusers": WanT2V_1_3B_SamplingParam,
-    "Wan-AI/Wan2.1-T2V-14B-Diffusers": WanT2V_14B_SamplingParam,
-    "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers": WanI2V_14B_480P_SamplingParam,
-    "Wan-AI/Wan2.1-I2V-14B-720P-Diffusers": WanI2V_14B_720P_SamplingParam,
+    "Wan-AI/Wan2.1-T2V-1.3B-Diffusers":
+    WanT2V_1_3B_SamplingParam,
+    "Wan-AI/Wan2.1-T2V-14B-Diffusers":
+    WanT2V_14B_SamplingParam,
+    "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers":
+    WanI2V_14B_480P_SamplingParam,
+    "Wan-AI/Wan2.1-I2V-14B-720P-Diffusers":
+    WanI2V_14B_720P_SamplingParam,
     "weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers":
     Wan2_1_Fun_1_3B_InP_SamplingParam,
     "IRMChen/Wan2.1-Fun-1.3B-Control-Diffusers":
     Wan2_1_Fun_1_3B_Control_SamplingParam,
 
     # Wan2.2
-    "Wan-AI/Wan2.2-TI2V-5B-Diffusers": Wan2_2_TI2V_5B_SamplingParam,
+    "Wan-AI/Wan2.2-TI2V-5B-Diffusers":
+    Wan2_2_TI2V_5B_SamplingParam,
     "FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers":
     Wan2_2_TI2V_5B_SamplingParam,
-    "Wan-AI/Wan2.2-T2V-A14B-Diffusers": Wan2_2_T2V_A14B_SamplingParam,
-    "Wan-AI/Wan2.2-I2V-A14B-Diffusers": Wan2_2_I2V_A14B_SamplingParam,
+    "Wan-AI/Wan2.2-T2V-A14B-Diffusers":
+    Wan2_2_T2V_A14B_SamplingParam,
+    "Wan-AI/Wan2.2-I2V-A14B-Diffusers":
+    Wan2_2_I2V_A14B_SamplingParam,
 
     # FastWan2.1
-    "FastVideo/FastWan2.1-T2V-1.3B-Diffusers": FastWanT2V480P_SamplingParam,
+    "FastVideo/FastWan2.1-T2V-1.3B-Diffusers":
+    FastWanT2V480P_SamplingParam,
 
     # FastWan2.2
-    "FastVideo/FastWan2.2-TI2V-5B-Diffusers": Wan2_2_TI2V_5B_SamplingParam,
+    "FastVideo/FastWan2.2-TI2V-5B-Diffusers":
+    Wan2_2_TI2V_5B_SamplingParam,
 
     # Causal Self-Forcing Wan2.1
     "wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers":
@@ -91,9 +102,12 @@ SAMPLING_PARAM_REGISTRY: dict[str, Any] = {
     Cosmos_Predict2_5_2B_Diffusers_SamplingParam,
 
     # MatrixGame2.0 models
-    "FastVideo/Matrix-Game-2.0-Base-Diffusers": MatrixGame2_SamplingParam,
-    "FastVideo/Matrix-Game-2.0-GTA-Diffusers": MatrixGame2_SamplingParam,
-    "FastVideo/Matrix-Game-2.0-TempleRun-Diffusers": MatrixGame2_SamplingParam,
+    "FastVideo/Matrix-Game-2.0-Base-Diffusers":
+    MatrixGame2_SamplingParam,
+    "FastVideo/Matrix-Game-2.0-GTA-Diffusers":
+    MatrixGame2_SamplingParam,
+    "FastVideo/Matrix-Game-2.0-TempleRun-Diffusers":
+    MatrixGame2_SamplingParam,
 
     # TurboDiffusion models
     "loayrashid/TurboWan2.1-T2V-1.3B-Diffusers":
@@ -103,10 +117,6 @@ SAMPLING_PARAM_REGISTRY: dict[str, Any] = {
     "loayrashid/TurboWan2.2-I2V-A14B-Diffusers":
     TurboDiffusionI2V_A14B_SamplingParam,
 
-    # LTX-2 models
-    "Lightricks/LTX-2": LTX2SamplingParam,
-    "FastVideo/LTX2-Distilled-Diffusers": LTX2SamplingParam,
-
     # Add other specific weight variants
 }
 
@@ -134,8 +144,6 @@ SAMPLING_PARAM_DETECTOR: dict[str, Callable[[str], bool]] = {
     lambda id: "cosmos2_5" in id.lower(),
     "cosmos":
     lambda id: "cosmos" in id.lower() and "2_5" not in id.lower(),
-    "ltx2":
-    lambda id: "ltx2" in id.lower() or "ltx-2" in id.lower(),
     # Add other pipeline architecture detectors
 }
 
@@ -156,7 +164,6 @@ SAMPLING_FALLBACK_PARAM: dict[str, Any] = {
     TurboDiffusionT2V_1_3B_SamplingParam,  # Default to T2V for fallback
     "cosmos25": Cosmos_Predict2_5_2B_Diffusers_SamplingParam,
     "cosmos": Cosmos_Predict2_2B_Video2World_SamplingParam,
-    "ltx2": LTX2SamplingParam,
     # Other fallbacks by architecture
 }
 

fastvideo/dataset/dataloader/record_schema.py

@@ -72,9 +72,16 @@ def i2v_record_creator(batch: PreprocessBatch) -> list[dict[str, Any]]:
     """Create a record for the Parquet dataset with CLIP features."""
     records = basic_t2v_record_creator(batch)
 
-    assert len(
-        batch.image_embeds) == 1, "image embedding should be a single tensor"
-    image_embeds = batch.image_embeds[0]
+    # Adapt for model doesn't have image encoder, e.g., Hunyuan
+    if len(batch.image_embeds) == 0:
+        image_embeds = None
+    elif len(batch.image_embeds) == 1:
+        image_embeds = batch.image_embeds[0]
+    else:
+        raise ValueError(
+            "Unexpected number of image_embeds in batch: expected 0 or 1, got {}".format(
+                len(batch.image_embeds)))
+        
     image_latent = batch.image_latent
     pil_image = batch.pil_image
 

fastvideo/dataset/utils.py

@@ -30,12 +30,13 @@ def get_torch_tensors_from_row_dict(row_dict,
     """
     return_dict = {}
     for key in keys:
-        shape, bytes = None, None
+        shape, bytes, dtype_str = None, None, None
         if isinstance(key, tuple):
             for k in key:
                 try:
                     shape = row_dict[f"{k}_shape"]
                     bytes = row_dict[f"{k}_bytes"]
+                    dtype_str = row_dict.get(f"{k}_dtype", 'float32')
                 except KeyError:
                     continue
             key = key[0]
@@ -44,13 +45,25 @@ def get_torch_tensors_from_row_dict(row_dict,
         else:
             shape = row_dict[f"{key}_shape"]
             bytes = row_dict[f"{key}_bytes"]
+            dtype_str = row_dict.get(f"{key}_dtype", 'float32')
+
+        # Convert dtype string to numpy dtype
+        if 'float16' in dtype_str:
+            np_dtype = np.float16
+        elif 'float32' in dtype_str:
+            np_dtype = np.float32
+        elif 'float64' in dtype_str:
+            np_dtype = np.float64
+        else:
+            # Default to float32 if dtype is unrecognized
+            np_dtype = np.float32
 
         # TODO (peiyuan): read precision
         if key == 'text_embedding' and (rng.random()
                                         if rng else random.random()) < cfg_rate:
             data = np.zeros((512, 4096), dtype=np.float32)
         else:
-            data = np.frombuffer(bytes, dtype=np.float32).reshape(shape).copy()
+            data = np.frombuffer(bytes, dtype=np_dtype).reshape(shape).copy()
         data = torch.from_numpy(data)
         if len(data.shape) == 3:
             B, L, D = data.shape
@@ -141,21 +154,35 @@ def collate_rows_from_parquet_schema(rows,
             # Get tensor data from row using the existing helper function pattern
             shape_key = f"{tensor_name}_shape"
             bytes_key = f"{tensor_name}_bytes"
+            dtype_key = f"{tensor_name}_dtype"
 
             if shape_key in row and bytes_key in row:
                 shape = row[shape_key]
                 bytes_data = row[bytes_key]
+                
+                # Get dtype from row, default to float32 for backward compatibility
+                dtype_str = row.get(dtype_key, 'float32')
+                # Convert dtype string to numpy dtype
+                if 'float16' in dtype_str:
+                    np_dtype = np.float16
+                elif 'float32' in dtype_str:
+                    np_dtype = np.float32
+                elif 'float64' in dtype_str:
+                    np_dtype = np.float64
+                else:
+                    # Default to float32 if dtype is unrecognized
+                    np_dtype = np.float32
 
                 if len(bytes_data) == 0:
                     tensor = torch.zeros(0, dtype=torch.bfloat16)
                 else:
-                    # Convert bytes to tensor using float32 as default
+                    # Convert bytes to tensor using the correct dtype
                     if tensor_name == 'text_embedding' and (rng.random(
                     ) if rng else random.random()) < cfg_rate:
                         data = np.zeros((512, 4096), dtype=np.float32)
                     else:
                         data = np.frombuffer(
-                            bytes_data, dtype=np.float32).reshape(shape).copy()
+                            bytes_data, dtype=np_dtype).reshape(shape).copy()
                     tensor = torch.from_numpy(data)
                     # if len(data.shape) == 3:
                     #     B, L, D = tensor.shape

fastvideo/distributed/communication_op.py

@@ -53,6 +53,9 @@ def sequence_model_parallel_all_gather_with_unpad(
         Tensor: Gathered and unpadded tensor
     """
 
+    # NCCL all_gather expects contiguous inputs.
+    if not input_.is_contiguous():
+        input_ = input_.contiguous()
     # First gather across all ranks
     gathered = get_sp_group().all_gather(input_, dim)
 

fastvideo/entrypoints/video_generator.py

@@ -18,8 +18,6 @@ import numpy as np
 import torch
 import torchvision
 from einops import rearrange
-import shutil
-import tempfile
 
 from fastvideo.configs.sample import SamplingParam
 from fastvideo.fastvideo_args import FastVideoArgs
@@ -391,11 +389,6 @@ class VideoGenerator:
         if batch.save_video:
             imageio.mimsave(output_path, frames, fps=batch.fps, format="mp4")
             logger.info("Saved video to %s", output_path)
-            audio = output_batch.extra.get("audio")
-            audio_sample_rate = output_batch.extra.get("audio_sample_rate")
-            if (audio is not None and audio_sample_rate is not None and
-                    not self._mux_audio(output_path, audio, audio_sample_rate)):
-                logger.warning("Audio mux failed; saved video without audio.")
 
         if batch.return_frames:
             return frames
@@ -403,7 +396,6 @@ class VideoGenerator:
             return {
                 "samples": samples,
                 "frames": frames,
-                "audio": output_batch.extra.get("audio"),
                 "prompts": prompt,
                 "size": (target_height, target_width, batch.num_frames),
                 "generation_time": gen_time,
@@ -413,98 +405,6 @@ class VideoGenerator:
                 "trajectory_decoded": output_batch.trajectory_decoded,
             }
 
-    @staticmethod
-    def _mux_audio(
-        video_path: str,
-        audio: torch.Tensor | np.ndarray,
-        sample_rate: int,
-    ) -> bool:
-        """Mux audio into video using PyAV."""
-        try:
-            import av
-        except ImportError:
-            logger.warning("PyAV not installed; cannot mux audio. "
-                           "Install with: pip install av")
-            return False
-
-        if torch.is_tensor(audio):
-            audio_np = audio.detach().cpu().float().numpy()
-        else:
-            audio_np = np.asarray(audio, dtype=np.float32)
-
-        if audio_np.ndim == 1:
-            audio_np = audio_np[:, None]
-        elif audio_np.ndim == 2:
-            if audio_np.shape[0] <= 8 and audio_np.shape[1] > audio_np.shape[0]:
-                audio_np = audio_np.T
-        else:
-            logger.warning("Unexpected audio shape %s; skipping mux.",
-                           audio_np.shape)
-            return False
-
-        audio_np = np.clip(audio_np, -1.0, 1.0)
-        audio_int16 = (audio_np * 32767.0).astype(np.int16)
-        num_channels = audio_int16.shape[1]
-        layout = "stereo" if num_channels == 2 else "mono"
-
-        try:
-            import wave
-            with tempfile.TemporaryDirectory() as tmpdir:
-                out_path = os.path.join(tmpdir, "muxed.mp4")
-                wav_path = os.path.join(tmpdir, "audio.wav")
-
-                # Write audio to WAV file
-                with wave.open(wav_path, "wb") as wav_file:
-                    wav_file.setnchannels(num_channels)
-                    wav_file.setsampwidth(2)
-                    wav_file.setframerate(sample_rate)
-                    wav_file.writeframes(audio_int16.tobytes())
-
-                # Open input video and audio
-                input_video = av.open(video_path)
-                input_audio = av.open(wav_path)
-
-                # Create output with both streams
-                output = av.open(out_path, mode="w")
-
-                # Add video stream (copy codec from input)
-                in_video_stream = input_video.streams.video[0]
-                out_video_stream = output.add_stream(
-                    codec_name=in_video_stream.codec_context.name,
-                    rate=in_video_stream.average_rate,
-                )
-                out_video_stream.width = in_video_stream.width
-                out_video_stream.height = in_video_stream.height
-                out_video_stream.pix_fmt = in_video_stream.pix_fmt
-
-                # Add audio stream (AAC)
-                out_audio_stream = output.add_stream("aac", rate=sample_rate)
-                out_audio_stream.layout = layout
-
-                # Remux video (decode and re-encode to be safe)
-                for frame in input_video.decode(video=0):
-                    for packet in out_video_stream.encode(frame):
-                        output.mux(packet)
-                for packet in out_video_stream.encode():
-                    output.mux(packet)
-
-                # Encode audio
-                for frame in input_audio.decode(audio=0):
-                    frame.pts = None  # Let encoder assign PTS
-                    for packet in out_audio_stream.encode(frame):
-                        output.mux(packet)
-                for packet in out_audio_stream.encode():
-                    output.mux(packet)
-
-                input_video.close()
-                input_audio.close()
-                output.close()
-                shutil.move(out_path, video_path)
-            return True
-        except Exception as e:
-            logger.warning("Audio mux failed: %s", e)
-            return False
-
     def set_lora_adapter(self,
                          lora_nickname: str,
                          lora_path: str | None = None) -> None:

fastvideo/fastvideo_args.py

@@ -133,7 +133,7 @@ class FastVideoArgs:
     # CPU offload parameters
     dit_cpu_offload: bool = True
     use_fsdp_inference: bool = False
-    dit_layerwise_offload: bool = True
+    dit_layerwise_offload: bool = False
     text_encoder_cpu_offload: bool = True
     image_encoder_cpu_offload: bool = True
     vae_cpu_offload: bool = True
@@ -166,14 +166,6 @@ class FastVideoArgs:
     # Prompt text file for batch processing
     prompt_txt: str | None = None
 
-    # LTX-2 VAE tiling overrides
-    ltx2_vae_tiling: bool | None = None
-    ltx2_vae_spatial_tile_size_in_pixels: int | None = None
-    ltx2_vae_spatial_tile_overlap_in_pixels: int | None = None
-    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
-
     # model paths for correct deallocation
     model_paths: dict[str, str] = field(default_factory=dict)
     model_loaded: dict[str, bool] = field(default_factory=lambda: {
@@ -211,44 +203,8 @@ class FastVideoArgs:
                 logger.error("Failed to load V-MoBA config from %s: %s",
                              self.moba_config_path, e)
                 raise
-        self._apply_ltx2_vae_overrides()
         self.check_fastvideo_args()
 
-    def _apply_ltx2_vae_overrides(self) -> None:
-        if self.pipeline_config is None:
-            return
-        vae_config = self.pipeline_config.vae_config
-        has_any = any(value is not None for value in (
-            self.ltx2_vae_spatial_tile_size_in_pixels,
-            self.ltx2_vae_spatial_tile_overlap_in_pixels,
-            self.ltx2_vae_temporal_tile_size_in_frames,
-            self.ltx2_vae_temporal_tile_overlap_in_frames,
-        ))
-        if self.ltx2_vae_tiling is not None and hasattr(self.pipeline_config,
-                                                        "vae_tiling"):
-            self.pipeline_config.vae_tiling = self.ltx2_vae_tiling
-        elif has_any and hasattr(self.pipeline_config, "vae_tiling"):
-            self.pipeline_config.vae_tiling = True
-
-        if hasattr(vae_config, "ltx2_spatial_tile_size_in_pixels"
-                   ) and self.ltx2_vae_spatial_tile_size_in_pixels is not None:
-            vae_config.ltx2_spatial_tile_size_in_pixels = (
-                self.ltx2_vae_spatial_tile_size_in_pixels)
-        if hasattr(
-                vae_config, "ltx2_spatial_tile_overlap_in_pixels"
-        ) and self.ltx2_vae_spatial_tile_overlap_in_pixels is not None:
-            vae_config.ltx2_spatial_tile_overlap_in_pixels = (
-                self.ltx2_vae_spatial_tile_overlap_in_pixels)
-        if hasattr(vae_config, "ltx2_temporal_tile_size_in_frames"
-                   ) and self.ltx2_vae_temporal_tile_size_in_frames is not None:
-            vae_config.ltx2_temporal_tile_size_in_frames = (
-                self.ltx2_vae_temporal_tile_size_in_frames)
-        if hasattr(
-                vae_config, "ltx2_temporal_tile_overlap_in_frames"
-        ) and self.ltx2_vae_temporal_tile_overlap_in_frames is not None:
-            vae_config.ltx2_temporal_tile_overlap_in_frames = (
-                self.ltx2_vae_temporal_tile_overlap_in_frames)
-
     @staticmethod
     def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser:
         # Model and path configuration
@@ -369,44 +325,6 @@ class FastVideoArgs:
             "Path to a text file containing prompts (one per line) for batch processing",
         )
 
-        # LTX-2 VAE tiling overrides
-        parser.add_argument(
-            "--ltx2-vae-tiling",
-            action=StoreBoolean,
-            default=FastVideoArgs.ltx2_vae_tiling,
-            help="Enable LTX-2 VAE tiling overrides.",
-        )
-        parser.add_argument(
-            "--ltx2-vae-spatial-tile-size-in-pixels",
-            type=int,
-            default=FastVideoArgs.ltx2_vae_spatial_tile_size_in_pixels,
-            help="LTX-2 VAE spatial tile size in pixels.",
-        )
-        parser.add_argument(
-            "--ltx2-vae-spatial-tile-overlap-in-pixels",
-            type=int,
-            default=FastVideoArgs.ltx2_vae_spatial_tile_overlap_in_pixels,
-            help="LTX-2 VAE spatial tile overlap in pixels.",
-        )
-        parser.add_argument(
-            "--ltx2-vae-temporal-tile-size-in-frames",
-            type=int,
-            default=FastVideoArgs.ltx2_vae_temporal_tile_size_in_frames,
-            help="LTX-2 VAE temporal tile size in frames.",
-        )
-        parser.add_argument(
-            "--ltx2-vae-temporal-tile-overlap-in-frames",
-            type=int,
-            default=FastVideoArgs.ltx2_vae_temporal_tile_overlap_in_frames,
-            help="LTX-2 VAE temporal tile overlap in frames.",
-        )
-        parser.add_argument(
-            "--ltx2-initial-latent-path",
-            type=str,
-            default=FastVideoArgs.ltx2_initial_latent_path,
-            help="Path to load/save a precomputed LTX-2 initial latent.",
-        )
-
         # LoRA parameters (inference-time adapter loading)
         parser.add_argument(
             "--lora-path",

fastvideo/hooks/hooks.py

@@ -1,122 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-
-import functools
-from typing import Any
-from torch import nn
-
-
-class ForwardHook:
-    """
-    Base class for forward hooks.
-    Hooks are used in the way:
-        modified_args, modified_kwargs = hook.pre_forward(module, *args, **kwargs)
-        output = module.forward(*modified_args, **modified_kwargs)
-        modified_output = hook.post_forward(module, output)
-    """
-
-    @classmethod
-    def name(cls) -> str:
-        raise NotImplementedError
-
-    def on_attach(self, module: nn.Module):  # noqa: B027
-        """Called once when the hook is attached to the module."""
-        pass
-
-    def on_detach(self, module: nn.Module):  # noqa: B027
-        """
-        Called once when the hook is detached from the module.
-        Note: this function is not guaranteed to be called if the module is
-              deleted before the hook is detached.
-        """
-        pass
-
-    def pre_forward(self, module: nn.Module, *args,
-                    **kwargs) -> tuple[tuple[Any, ...], dict[str, Any]]:
-        """Called before the module's forward method is executed."""
-        return args, kwargs
-
-    def post_forward(self, module: nn.Module, output: Any) -> Any:
-        """Called after the module's forward method is executed."""
-        return output
-
-
-class ModuleHookManager:
-    module_hook_attribute = "_hook_manager"
-
-    def __init__(self, module: nn.Module):
-        self.module = module
-        self.forward_hooks: dict[str, ForwardHook] = {}
-        self.original_forward = module.forward
-
-    @classmethod
-    def get_from(cls, module: nn.Module) -> "ModuleHookManager | None":
-        if hasattr(module, cls.module_hook_attribute):
-            return getattr(module, cls.module_hook_attribute)
-        return None
-
-    @classmethod
-    def get_from_or_default(cls, module: nn.Module) -> "ModuleHookManager":
-        if not hasattr(module, cls.module_hook_attribute):
-            setattr(module, cls.module_hook_attribute, cls(module))
-
-            def forward_hook_wrapper(mod: nn.Module, *args, **kwargs):
-                manager: ModuleHookManager = getattr(mod,
-                                                     cls.module_hook_attribute)
-                for hook in manager.forward_hooks.values():
-                    args, kwargs = hook.pre_forward(mod, *args, **kwargs)
-                output = manager.original_forward(*args, **kwargs)
-                for hook in reversed(manager.forward_hooks.values()):
-                    output = hook.post_forward(mod, output)
-                return output
-
-            module.forward = functools.partial(forward_hook_wrapper, module)
-
-        return getattr(module, cls.module_hook_attribute)
-
-    @staticmethod
-    def remove_from_manager(module: nn.Module) -> None:
-        if hasattr(module, ModuleHookManager.module_hook_attribute):
-            manager: ModuleHookManager = getattr(
-                module, ModuleHookManager.module_hook_attribute)
-            module.forward = manager.original_forward
-            delattr(module, ModuleHookManager.module_hook_attribute)
-
-    def _check_manager_attached(self) -> None:
-        if not hasattr(self.module, self.module_hook_attribute):
-            raise ValueError("ModuleHookManager is not attached to the module.")
-        if getattr(self.module, self.module_hook_attribute) is not self:
-            raise ValueError(
-                "ModuleHookManager attached to the module is different.")
-
-    def append_forward_hook(self, hook: ForwardHook):
-        self._check_manager_attached()
-        if hook.name() in self.forward_hooks:
-            raise ValueError(
-                f"Hook with name {hook.name()} is already registered.")
-        # after python 3.7, dicts maintain insertion order
-        self.forward_hooks[hook.name()] = hook
-        hook.on_attach(self.module)
-
-    def replace_forward_hook(self,
-                             hook_name: str,
-                             new_hook: ForwardHook,
-                             run_on_attach: bool = True):
-        self._check_manager_attached()
-        if hook_name not in self.forward_hooks:
-            raise ValueError(f"No hook with name {hook_name} found.")
-        old_hook = self.forward_hooks[hook_name]
-        if run_on_attach:
-            old_hook.on_detach(self.module)
-        self.forward_hooks[hook_name] = new_hook
-        new_hook.on_attach(self.module)
-
-    def remove_forward_hook(self, hook_name: str, run_detach: bool = True):
-        self._check_manager_attached()
-        if hook_name not in self.forward_hooks:
-            raise ValueError(f"No hook with name {hook_name} found.")
-        if run_detach:
-            self.forward_hooks[hook_name].on_detach(self.module)
-        del self.forward_hooks[hook_name]
-
-    def get_forward_hook(self, hook_name: str) -> ForwardHook | None:
-        return self.forward_hooks.get(hook_name, None)

fastvideo/hooks/layerwise_offload.py

@@ -1,164 +0,0 @@
-from contextlib import contextmanager
-from typing import Any
-import torch
-from torch import nn
-from fastvideo.hooks.hooks import ForwardHook, ModuleHookManager
-from fastvideo.logger import init_logger
-
-logger = init_logger(__name__)
-
-
-def _tensor_placeholder(tensor: torch.Tensor,
-                        device: torch.device) -> torch.Tensor:
-    """Create a rank-preserving empty placeholder on the specified device."""
-    shape = (0, ) if tensor.ndim <= 0 else (0, ) * tensor.ndim
-    return torch.empty(shape, device=device, dtype=tensor.dtype)
-
-
-class LayerwiseOffloadState:
-
-    def __init__(
-        self,
-        async_copy_stream: torch.cuda.Stream,
-        device: torch.device,
-        next_state: "LayerwiseOffloadState | None" = None,
-    ) -> None:
-        self.async_copy_stream = async_copy_stream
-        self.next_state = next_state
-        self.gpu_named_parameters: dict[str, torch.Tensor] = {}
-        self.cpu_named_parameters: dict[str, torch.Tensor] = {}
-        self.module_ref: nn.Module = None  # type: ignore
-        self.device: torch.device = device
-
-    def _will_offload(self, name: str) -> bool:
-        return True
-
-    @torch.compiler.disable
-    def on_init(self, module: nn.Module):
-        self.module_ref = module
-        for name, param in self.module_ref.named_parameters():
-            if self._will_offload(name):
-                self.cpu_named_parameters[name] = (
-                    param.data.detach().to("cpu").pin_memory())
-                param.data = _tensor_placeholder(param.data, self.device)
-
-    @torch.compiler.disable
-    def wait_and_replace_params(self):
-        torch.cuda.current_stream().wait_stream(self.async_copy_stream)
-        # now gpu_named_parameters are ready
-        for name, param in self.module_ref.named_parameters():
-            if not self._will_offload(name):
-                continue
-            if name not in self.gpu_named_parameters:
-                # first load with blocking load
-                self.gpu_named_parameters[name] = self.cpu_named_parameters[
-                    name].to(self.device)
-            param.data = self.gpu_named_parameters[name]
-
-    @torch.compiler.disable
-    def prefetch_params(self):
-        compute_stream = torch.cuda.current_stream()
-        with torch.cuda.stream(self.async_copy_stream):
-            for name, param in self.module_ref.named_parameters():
-                if not self._will_offload(name):
-                    continue
-                assert name not in self.gpu_named_parameters
-                gpu_param = self.cpu_named_parameters[name].to(
-                    self.device, non_blocking=True)
-                gpu_param.record_stream(
-                    compute_stream
-                )  # ensure tensor will not be freed until forward is completed
-                self.gpu_named_parameters[name] = gpu_param
-
-    @torch.compiler.disable
-    def release_gpu_params(self):
-        for name, param in self.module_ref.named_parameters():
-            if self._will_offload(name):
-                param.data = _tensor_placeholder(param.data, self.device)
-                del self.gpu_named_parameters[name]
-        assert len(self.gpu_named_parameters) == 0
-
-
-class LayerwiseOffloadHook(ForwardHook):
-    """A hook that enables layerwise CPU offloading during forward pass."""
-
-    def __init__(self, state: LayerwiseOffloadState) -> None:
-        self.state = state
-
-    def on_attach(self, module: nn.Module):
-        self.state.on_init(module)  # pyright: ignore
-
-    def on_detach(self, module: nn.Module):
-        named_parameters = dict(module.named_parameters())
-        for name, cpu_tensor in self.state.cpu_named_parameters.items():
-            if name not in self.state.gpu_named_parameters:
-                if name in named_parameters:
-                    named_parameters[name].data = cpu_tensor.to(
-                        device=self.state.device)
-                else:
-                    logger.warning(
-                        "Parameter {} not found in module during detachment.",
-                        name,
-                    )
-
-    @classmethod
-    def name(cls) -> str:
-        return "LayerwiseOffloadHook"
-
-    def pre_forward(self, module: nn.Module, *args, **kwargs):
-        self.state.wait_and_replace_params()  # pyright: ignore
-        if self.state.next_state is not None:
-            self.state.next_state.prefetch_params()  # pyright: ignore
-        return args, kwargs
-
-    def post_forward(self, module: torch.nn.Module, output: Any):
-        self.state.release_gpu_params()  # pyright: ignore
-        return output
-
-    @contextmanager
-    def mutate_params_scope(self):
-        try:
-            # load params to GPU and keep them there
-            self.state.wait_and_replace_params()  # pyright: ignore
-            yield
-        finally:
-            # instead of releasing, we should overwrite the original params since they have been modified
-            self.state.cpu_named_parameters.clear()
-            self.state.gpu_named_parameters.clear()
-            self.state.on_init(self.state.module_ref)  # pyright: ignore
-
-
-def enable_layerwise_offload(model: nn.Module, is_replace: bool = False):
-    if torch.cuda.is_available():
-        device = torch.device("cuda", torch.cuda.current_device())
-    else:
-        logger.warning(
-            "CUDA is not available. Layerwise offloading is disabled.")
-        return
-    state_list = []
-    async_stream = torch.cuda.Stream()
-    for name, submodule in model.named_children():
-        if isinstance(submodule, nn.ModuleList):
-            for idx, module_entry in enumerate(submodule):
-                state = LayerwiseOffloadState(async_copy_stream=async_stream,
-                                              device=device)
-                state_list.append(state)
-                hook_mgr = ModuleHookManager.get_from_or_default(module_entry)
-                hook = LayerwiseOffloadHook(state)
-                if is_replace:
-                    existing_hook = hook_mgr.forward_hooks.get(hook.name())
-                    if existing_hook is not None:
-                        hook_mgr.replace_forward_hook(hook.name(), hook)
-                    else:
-                        raise AssertionError(
-                            f"Expect hook exists in {name} for replacement.")
-                else:
-                    hook_mgr.append_forward_hook(hook)
-            break
-    if len(state_list) == 0:
-        raise ValueError(
-            "No nn.ModuleList found in the model for layerwise offloading.")
-
-    # circular linking of states
-    for i in range(len(state_list)):
-        state_list[i].next_state = state_list[(i + 1) % len(state_list)]

fastvideo/layers/rotary_embedding.py

@@ -194,14 +194,14 @@ class RotaryEmbedding(CustomOp):
         cos, sin = cos_sin.chunk(2, dim=-1)
 
         query_shape = query.shape
-        query = query.view(num_tokens, -1, self.head_size)
+        query = query.reshape(num_tokens, -1, self.head_size)
         query_rot = query[..., :self.rotary_dim]
         query_pass = query[..., self.rotary_dim:]
         query_rot = _apply_rotary_emb(query_rot, cos, sin, self.is_neox_style)
         query = torch.cat((query_rot, query_pass), dim=-1).reshape(query_shape)
 
         key_shape = key.shape
-        key = key.view(num_tokens, -1, self.head_size)
+        key = key.reshape(num_tokens, -1, self.head_size)
         key_rot = key[..., :self.rotary_dim]
         key_pass = key[..., self.rotary_dim:]
         key_rot = _apply_rotary_emb(key_rot, cos, sin, self.is_neox_style)

fastvideo/models/audio/__init__.py

@@ -1,9 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-
-from fastvideo.models.audio.ltx2_audio_vae import (
-    LTX2AudioDecoder,
-    LTX2AudioEncoder,
-    LTX2Vocoder,
-)
-
-__all__ = ["LTX2AudioEncoder", "LTX2AudioDecoder", "LTX2Vocoder"]

fastvideo/models/dits/hunyuanvideo15.py

@@ -690,12 +690,15 @@ class SingleTokenRefiner(nn.Module):
         timestep_aware_representations = self.t_embedder(t)
 
         # Get context-aware representations
-        original_dtype = x.dtype
+        target_dtype = self.c_embedder.fc_in.weight.dtype
+        if x.dtype != target_dtype:
+            x = x.to(dtype=target_dtype)
         if mask is None:
             context_aware_representations = x.mean(dim=1)
         else:
-            mask_float = mask.float().unsqueeze(-1)  # [B, L, 1]
-            context_aware_representations = (x * mask_float).sum(dim=1) / mask_float.sum(dim=1)
+            mask_float = mask.to(dtype=target_dtype).unsqueeze(-1)  # [B, L, 1]
+            context_aware_representations = (
+                x * mask_float).sum(dim=1) / mask_float.sum(dim=1)
 
         context_aware_representations = self.c_embedder(
             context_aware_representations)
@@ -850,4 +853,4 @@ class FinalLayer(nn.Module):
         scale, shift = self.adaLN_modulation(c).chunk(2, dim=-1)
         x = self.norm_final(x) * (1.0 + scale.unsqueeze(1)) + shift.unsqueeze(1)
         x, _ = self.linear(x)
-        return x
+        return x

fastvideo/models/dits/wanvideo.py

@@ -734,8 +734,18 @@ class WanTransformer3DModel(CachableDiT):
                         block, hidden_states, encoder_hidden_states,
                         timestep_proj, freqs_cis, attention_mask)
             else:
-                for block in self.blocks:
-                    hidden_states = block(hidden_states, encoder_hidden_states,
+                offload_mgr = getattr(self, "_layerwise_offload_manager", None)
+                use_offload = offload_mgr is not None and getattr(offload_mgr, "enabled", False)
+
+                for i, block in enumerate(self.blocks):
+                    scope = offload_mgr.layer_scope(
+                        prefetch_layer_idx=i + 1 if i + 1 < len(self.blocks) else None,
+                        release_layer_idx=i,
+                        non_blocking=True,
+                    ) if use_offload else nullcontext()
+                    
+                    with scope:
+                        hidden_states = block(hidden_states, encoder_hidden_states,
                                               timestep_proj, freqs_cis, attention_mask)
             # if teacache is enabled, we need to cache the original hidden states
 

fastvideo/models/encoders/gemma.py

@@ -1,563 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-from __future__ import annotations
-
-from dataclasses import dataclass
-import os
-from typing import Iterable
-
-import torch
-from torch import nn
-from transformers import Gemma3ForConditionalGeneration
-
-from fastvideo.configs.models.encoders import BaseEncoderOutput, TextEncoderConfig
-from fastvideo.models.encoders.base import TextEncoder
-from fastvideo.models.dits.ltx2 import (
-    FeedForward,
-    LTXRopeType,
-    apply_ltx_rotary_emb,
-    generate_ltx_freq_grid_np,
-    generate_ltx_freq_grid_pytorch,
-    precompute_ltx_freqs_cis,
-)
-from fastvideo.models.loader.weight_utils import default_weight_loader
-from fastvideo.platforms import AttentionBackendEnum
-
-
-def _debug_log_line(message: str) -> None:
-    if os.getenv("LTX2_PIPELINE_DEBUG_LOG", "0") != "1":
-        return
-    log_path = os.getenv("LTX2_PIPELINE_DEBUG_PATH", "")
-    if not log_path:
-        return
-    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(message + "\n")
-
-
-def _debug_gemma_log_line(message: str) -> None:
-    log_path = os.getenv("LTX2_FASTVIDEO_GEMMA_LOG", "")
-    if not log_path:
-        return
-    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(message + "\n")
-
-
-@dataclass(frozen=True)
-class GemmaConnectorConfig:
-    num_attention_heads: int
-    attention_head_dim: int
-    num_layers: int
-    positional_embedding_theta: float
-    positional_embedding_max_pos: list[int]
-    rope_type: LTXRopeType
-    double_precision_rope: bool
-    num_learnable_registers: int | None
-
-
-class GemmaFeaturesExtractorProjLinear(nn.Module):
-    """Linear projection that aggregates stacked Gemma hidden states."""
-
-    def __init__(self, in_features: int, out_features: int) -> None:
-        super().__init__()
-        self.aggregate_embed = nn.Linear(in_features, out_features, bias=False)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        return self.aggregate_embed(x)
-
-
-class _BasicTransformerBlock1D(nn.Module):
-    """1D transformer block for connector processing."""
-
-    def __init__(
-        self,
-        dim: int,
-        heads: int,
-        dim_head: int,
-        rope_type: LTXRopeType,
-        norm_eps: float = 1e-6,
-    ) -> None:
-        super().__init__()
-        self.attn1 = _GemmaAttention(
-            query_dim=dim,
-            context_dim=None,
-            heads=heads,
-            dim_head=dim_head,
-            norm_eps=norm_eps,
-            rope_type=rope_type,
-        )
-        self.ff = FeedForward(dim, dim_out=dim)
-        self.norm_eps = norm_eps
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        attention_mask: torch.Tensor | None = None,
-        pe: tuple[torch.Tensor, torch.Tensor] | None = None,
-    ) -> torch.Tensor:
-        norm_hidden_states = torch.nn.functional.rms_norm(
-            hidden_states, (hidden_states.shape[-1],), eps=self.norm_eps
-        )
-        if norm_hidden_states.ndim == 4:
-            norm_hidden_states = norm_hidden_states.squeeze(1)
-
-        attn_output = self.attn1(
-            norm_hidden_states,
-            mask=attention_mask,
-            pe=pe,
-        )
-        hidden_states = attn_output + hidden_states
-        if hidden_states.ndim == 4:
-            hidden_states = hidden_states.squeeze(1)
-
-        norm_hidden_states = torch.nn.functional.rms_norm(
-            hidden_states, (hidden_states.shape[-1],), eps=self.norm_eps
-        )
-        ff_output = self.ff(norm_hidden_states)
-        hidden_states = ff_output + hidden_states
-        if hidden_states.ndim == 4:
-            hidden_states = hidden_states.squeeze(1)
-        return hidden_states
-
-
-class _GemmaAttention(nn.Module):
-    """Attention implementation aligned with LTX-2 text encoder."""
-
-    def __init__(
-        self,
-        query_dim: int,
-        context_dim: int | None,
-        heads: int,
-        dim_head: int,
-        norm_eps: float,
-        rope_type: LTXRopeType,
-    ) -> None:
-        super().__init__()
-        inner_dim = dim_head * heads
-        context_dim = query_dim if context_dim is None else context_dim
-
-        self.heads = heads
-        self.dim_head = dim_head
-        self.rope_type = rope_type
-
-        self.q_norm = torch.nn.RMSNorm(inner_dim, eps=norm_eps)
-        self.k_norm = torch.nn.RMSNorm(inner_dim, eps=norm_eps)
-        self.to_q = nn.Linear(query_dim, inner_dim, bias=True)
-        self.to_k = nn.Linear(context_dim, inner_dim, bias=True)
-        self.to_v = nn.Linear(context_dim, inner_dim, bias=True)
-        self.to_out = nn.Sequential(nn.Linear(inner_dim, query_dim, bias=True), nn.Identity())
-
-    def forward(
-        self,
-        x: torch.Tensor,
-        context: torch.Tensor | None = None,
-        mask: torch.Tensor | None = None,
-        pe: tuple[torch.Tensor, torch.Tensor] | None = None,
-        k_pe: tuple[torch.Tensor, torch.Tensor] | None = None,
-    ) -> torch.Tensor:
-        q = self.to_q(x)
-        context = x if context is None else context
-        k = self.to_k(context)
-        v = self.to_v(context)
-
-        q = self.q_norm(q)
-        k = self.k_norm(k)
-
-        if pe is not None:
-            q = apply_ltx_rotary_emb(q, pe, self.rope_type)
-            k = apply_ltx_rotary_emb(k, pe if k_pe is None else k_pe, self.rope_type)
-
-        b, q_len, _ = q.shape
-        k_len = k.shape[1]
-        q = q.view(b, q_len, self.heads, self.dim_head).transpose(1, 2)
-        k = k.view(b, k_len, self.heads, self.dim_head).transpose(1, 2)
-        v = v.view(b, k_len, self.heads, self.dim_head).transpose(1, 2)
-
-        if mask is not None:
-            if mask.ndim == 2:
-                mask = mask.unsqueeze(0)
-            if mask.ndim == 3:
-                mask = mask.unsqueeze(1)
-
-        out = torch.nn.functional.scaled_dot_product_attention(
-            q,
-            k,
-            v,
-            attn_mask=mask,
-            dropout_p=0.0,
-            is_causal=False,
-        )
-        out = out.transpose(1, 2).reshape(b, q_len, -1)
-        return self.to_out(out)
-
-
-class Embeddings1DConnector(nn.Module):
-    """Transformer connector that refines Gemma embeddings for LTX-2."""
-
-    _supports_gradient_checkpointing = True
-
-    def __init__(self, config: GemmaConnectorConfig) -> None:
-        super().__init__()
-        self.num_attention_heads = config.num_attention_heads
-        self.inner_dim = config.num_attention_heads * config.attention_head_dim
-        self.positional_embedding_theta = config.positional_embedding_theta
-        self.positional_embedding_max_pos = config.positional_embedding_max_pos
-        self.rope_type = config.rope_type
-        self.double_precision_rope = config.double_precision_rope
-        self.transformer_1d_blocks = nn.ModuleList(
-            [
-                _BasicTransformerBlock1D(
-                    dim=self.inner_dim,
-                    heads=config.num_attention_heads,
-                    dim_head=config.attention_head_dim,
-                    rope_type=config.rope_type,
-                )
-                for _ in range(config.num_layers)
-            ]
-        )
-        self.num_learnable_registers = config.num_learnable_registers
-        if self.num_learnable_registers:
-            self.learnable_registers = nn.Parameter(
-                torch.rand(
-                    self.num_learnable_registers,
-                    self.inner_dim,
-                    dtype=torch.bfloat16,
-                )
-                * 2.0
-                - 1.0
-            )
-
-    def _replace_padded_with_learnable_registers(
-        self,
-        hidden_states: torch.Tensor,
-        attention_mask: torch.Tensor,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        assert hidden_states.shape[1] % self.num_learnable_registers == 0, (
-            f"Hidden states sequence length {hidden_states.shape[1]} must be divisible by "
-            f"num_learnable_registers {self.num_learnable_registers}."
-        )
-
-        num_registers_duplications = (
-            hidden_states.shape[1] // self.num_learnable_registers
-        )
-        learnable_registers = torch.tile(
-            self.learnable_registers, (num_registers_duplications, 1)
-        )
-        attention_mask_binary = (
-            attention_mask.squeeze(1).squeeze(1).unsqueeze(-1) >= -9000.0
-        ).int()
-
-        non_zero_hidden_states = hidden_states[
-            :, attention_mask_binary.squeeze().bool(), :
-        ]
-        non_zero_nums = non_zero_hidden_states.shape[1]
-        pad_length = hidden_states.shape[1] - non_zero_nums
-        adjusted_hidden_states = torch.nn.functional.pad(
-            non_zero_hidden_states, pad=(0, 0, 0, pad_length), value=0
-        )
-        flipped_mask = torch.flip(attention_mask_binary, dims=[1])
-        hidden_states = flipped_mask * adjusted_hidden_states + (
-            1 - flipped_mask
-        ) * learnable_registers
-
-        attention_mask = torch.full_like(
-            attention_mask,
-            0.0,
-            dtype=attention_mask.dtype,
-            device=attention_mask.device,
-        )
-
-        return hidden_states, attention_mask
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        attention_mask: torch.Tensor | None = None,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        if self.num_learnable_registers:
-            hidden_states, attention_mask = (
-                self._replace_padded_with_learnable_registers(
-                    hidden_states, attention_mask
-                )
-            )
-
-        indices_grid = torch.arange(
-            hidden_states.shape[1],
-            dtype=torch.float32,
-            device=hidden_states.device,
-        )
-        indices_grid = indices_grid[None, None, :]
-        freq_grid_generator = (
-            generate_ltx_freq_grid_np
-            if self.double_precision_rope
-            else generate_ltx_freq_grid_pytorch
-        )
-        freqs_cis = precompute_ltx_freqs_cis(
-            indices_grid=indices_grid,
-            dim=self.inner_dim,
-            out_dtype=hidden_states.dtype,
-            theta=self.positional_embedding_theta,
-            max_pos=self.positional_embedding_max_pos,
-            num_attention_heads=self.num_attention_heads,
-            rope_type=self.rope_type,
-            freq_grid_generator=freq_grid_generator,
-        )
-
-        for block in self.transformer_1d_blocks:
-            hidden_states = block(
-                hidden_states, attention_mask=attention_mask, pe=freqs_cis
-            )
-
-        hidden_states = torch.nn.functional.rms_norm(
-            hidden_states, (hidden_states.shape[-1],), eps=1e-6
-        )
-
-        return hidden_states, attention_mask
-
-
-class LTX2GemmaTextEncoderModel(TextEncoder):
-    _supported_attention_backends = (
-        AttentionBackendEnum.FLASH_ATTN,
-        AttentionBackendEnum.TORCH_SDPA,
-    )
-
-    def __init__(self, config: TextEncoderConfig) -> None:
-        super().__init__(config)
-        arch = config.arch_config
-
-        self.feature_extractor_linear = GemmaFeaturesExtractorProjLinear(
-            in_features=arch.feature_extractor_in_features,
-            out_features=arch.feature_extractor_out_features,
-        )
-
-        connector_config = GemmaConnectorConfig(
-            num_attention_heads=arch.connector_num_attention_heads,
-            attention_head_dim=arch.connector_attention_head_dim,
-            num_layers=arch.connector_num_layers,
-            positional_embedding_theta=arch.connector_positional_embedding_theta,
-            positional_embedding_max_pos=arch.connector_positional_embedding_max_pos,
-            rope_type=LTXRopeType(arch.connector_rope_type),
-            double_precision_rope=arch.connector_double_precision_rope,
-            num_learnable_registers=arch.connector_num_learnable_registers,
-        )
-        self.embeddings_connector = Embeddings1DConnector(connector_config)
-        self.audio_embeddings_connector = Embeddings1DConnector(connector_config)
-
-        self.gemma_model_path = arch.gemma_model_path
-        self.gemma_dtype = arch.gemma_dtype
-        self.padding_side = arch.padding_side
-        self._gemma_model: Gemma3ForConditionalGeneration | None = None
-
-    def named_parameters(self, prefix: str = "", recurse: bool = True):
-        for name, param in super().named_parameters(
-            prefix=prefix, recurse=recurse
-        ):
-            if name.startswith("gemma_model."):
-                continue
-            yield name, param
-
-    @property
-    def gemma_model(self) -> Gemma3ForConditionalGeneration:
-        if self._gemma_model is None:
-            gemma_path = self.gemma_model_path
-            if not gemma_path:
-                raise ValueError(
-                    "gemma_model_path must be set (expected text_encoder/gemma)."
-                )
-            dtype = getattr(torch, self.gemma_dtype, torch.bfloat16)
-            self._gemma_model = Gemma3ForConditionalGeneration.from_pretrained(
-                gemma_path,
-                local_files_only=True,
-                torch_dtype=dtype,
-            )
-            # Configure model-level attention implementation when using TORCH_SDPA.
-            # Note: torch.backends.cuda.enable_*_sdp() settings should be configured
-            # at application/pipeline initialization level, not here, to avoid
-            # unexpected side effects across the application.
-            if os.getenv("FASTVIDEO_ATTENTION_BACKEND") == "TORCH_SDPA":
-                if hasattr(self._gemma_model.config, "attn_implementation"):
-                    self._gemma_model.config.attn_implementation = "sdpa"
-                if hasattr(self._gemma_model.config, "_attn_implementation"):
-                    self._gemma_model.config._attn_implementation = "sdpa"
-            device = next(self.feature_extractor_linear.parameters()).device
-            self._gemma_model.to(device=device)
-            self._gemma_model.eval()
-        return self._gemma_model
-
-    def _run_feature_extractor(
-        self,
-        hidden_states: tuple[torch.Tensor, ...],
-        attention_mask: torch.Tensor,
-        padding_side: str,
-    ) -> torch.Tensor:
-        encoded_text_features = torch.stack(hidden_states, dim=-1)
-        if os.getenv("LTX2_FASTVIDEO_GEMMA_LOG", ""):
-            for idx, layer in enumerate(hidden_states):
-                _debug_gemma_log_line(
-                    f"fastvideo:gemma_hidden_state_{idx}"
-                    f":sum={layer.float().sum().item():.6f}"
-                )
-            _debug_gemma_log_line(
-                "fastvideo:gemma_hidden_states_stack"
-                f":sum={encoded_text_features.float().sum().item():.6f}"
-            )
-        encoded_text_features_dtype = encoded_text_features.dtype
-        sequence_lengths = attention_mask.sum(dim=-1)
-        normed_text_features = _norm_and_concat_padded_batch(
-            encoded_text_features, sequence_lengths, padding_side=padding_side
-        )
-        return self.feature_extractor_linear(
-            normed_text_features.to(encoded_text_features_dtype)
-        )
-
-    def _convert_to_additive_mask(
-        self, attention_mask: torch.Tensor, dtype: torch.dtype
-    ) -> torch.Tensor:
-        return (attention_mask - 1).to(dtype).reshape(
-            (attention_mask.shape[0], 1, -1, attention_mask.shape[-1])
-        ) * torch.finfo(dtype).max
-
-    def _run_connectors(
-        self,
-        encoded_input: torch.Tensor,
-        attention_mask: torch.Tensor,
-    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-        connector_attention_mask = self._convert_to_additive_mask(
-            attention_mask, encoded_input.dtype
-        )
-        encoded, encoded_connector_attention_mask = self.embeddings_connector(
-            encoded_input, connector_attention_mask
-        )
-
-        attention_mask = (encoded_connector_attention_mask < 0.000001).to(
-            torch.int64
-        )
-        attention_mask = attention_mask.reshape(
-            [encoded.shape[0], encoded.shape[1], 1]
-        )
-        encoded = encoded * attention_mask
-
-        encoded_for_audio, _ = self.audio_embeddings_connector(
-            encoded_input, connector_attention_mask
-        )
-
-        return encoded, encoded_for_audio, attention_mask.squeeze(-1)
-
-    def forward(
-        self,
-        input_ids: torch.Tensor | None,
-        position_ids: torch.Tensor | None = None,
-        attention_mask: torch.Tensor | None = None,
-        inputs_embeds: torch.Tensor | None = None,
-        output_hidden_states: bool | None = None,
-        **kwargs,
-    ) -> BaseEncoderOutput:
-        if input_ids is None:
-            raise ValueError("input_ids is required for Gemma text encoding.")
-        if attention_mask is None:
-            attention_mask = torch.ones_like(input_ids)
-
-        model = self.gemma_model
-        input_ids = input_ids.to(device=model.device)
-        attention_mask = attention_mask.to(device=model.device)
-        outputs = model(
-            input_ids=input_ids,
-            attention_mask=attention_mask,
-            output_hidden_states=True,
-            return_dict=True,
-        )
-
-        encoded_inputs = self._run_feature_extractor(
-            outputs.hidden_states,
-            attention_mask,
-            padding_side=self.padding_side,
-        )
-        if os.getenv("LTX2_PIPELINE_DEBUG_LOG", "0") == "1":
-            _debug_log_line(
-                "fastvideo:gemma_feature"
-                f":sum={encoded_inputs.float().sum().item():.6f} "
-                f"shape={tuple(encoded_inputs.shape)}"
-            )
-        video_encoding, audio_encoding, attention_mask = self._run_connectors(
-            encoded_inputs, attention_mask
-        )
-        if os.getenv("LTX2_PIPELINE_DEBUG_LOG", "0") == "1":
-            _debug_log_line(
-                "fastvideo:gemma_video_encoding"
-                f":sum={video_encoding.float().sum().item():.6f} "
-                f"shape={tuple(video_encoding.shape)}"
-            )
-            _debug_log_line(
-                "fastvideo:gemma_audio_encoding"
-                f":sum={audio_encoding.float().sum().item():.6f} "
-                f"shape={tuple(audio_encoding.shape)}"
-            )
-
-        hidden_states = (audio_encoding, ) if output_hidden_states else None
-        return BaseEncoderOutput(
-            last_hidden_state=video_encoding,
-            hidden_states=hidden_states,
-            attention_mask=attention_mask,
-        )
-
-    def load_weights(
-        self, weights: Iterable[tuple[str, torch.Tensor]]
-    ) -> set[str]:
-        params_dict = dict(self.named_parameters())
-        loaded_params: set[str] = set()
-        for name, loaded_weight in weights:
-            if name == "aggregate_embed.weight":
-                name = "feature_extractor_linear.aggregate_embed.weight"
-            if name not in params_dict:
-                continue
-            param = params_dict[name]
-            weight_loader = getattr(param, "weight_loader", default_weight_loader)
-            weight_loader(param, loaded_weight)
-            loaded_params.add(name)
-        return loaded_params
-
-
-def _norm_and_concat_padded_batch(
-    encoded_text: torch.Tensor,
-    sequence_lengths: torch.Tensor,
-    padding_side: str = "right",
-) -> torch.Tensor:
-    b, t, d, l = encoded_text.shape
-    device = encoded_text.device
-
-    token_indices = torch.arange(t, device=device)[None, :]
-    if padding_side == "right":
-        mask = token_indices < sequence_lengths[:, None]
-    elif padding_side == "left":
-        start_indices = t - sequence_lengths[:, None]
-        mask = token_indices >= start_indices
-    else:
-        raise ValueError(
-            f"padding_side must be 'left' or 'right', got {padding_side}"
-        )
-
-    mask = mask.reshape(b, t, 1, 1)
-    eps = 1e-6
-
-    masked = encoded_text.masked_fill(~mask, 0.0)
-    denom = (sequence_lengths * d).view(b, 1, 1, 1)
-    mean = masked.sum(dim=(1, 2), keepdim=True) / (denom + eps)
-
-    x_min = encoded_text.masked_fill(~mask, float("inf")).amin(
-        dim=(1, 2), keepdim=True
-    )
-    x_max = encoded_text.masked_fill(~mask, float("-inf")).amax(
-        dim=(1, 2), keepdim=True
-    )
-    range_ = x_max - x_min
-
-    normed = 8 * (encoded_text - mean) / (range_ + eps)
-    normed = normed.reshape(b, t, -1)
-
-    mask_flattened = mask.reshape(b, t, 1).expand(-1, -1, d * l)
-    normed = normed.masked_fill(~mask_flattened, 0.0)
-    return normed

fastvideo/models/loader/component_loader.py

@@ -35,7 +35,7 @@ from fastvideo.models.loader.weight_utils import (
 )
 from fastvideo.models.registry import ModelRegistry
 from fastvideo.utils import PRECISION_TO_TYPE, is_pin_memory_available
-from fastvideo.hooks.layerwise_offload import enable_layerwise_offload
+from fastvideo.models.layerwise_offload import LayerwiseOffloadManager
 
 logger = init_logger(__name__)
 
@@ -80,9 +80,6 @@ class ComponentLoader(ABC):
             "transformer": (TransformerLoader, "diffusers"),
             "transformer_2": (TransformerLoader, "diffusers"),
             "vae": (VAELoader, "diffusers"),
-            "audio_vae": (AudioDecoderLoader, "diffusers"),
-            "audio_decoder": (AudioDecoderLoader, "diffusers"),
-            "vocoder": (VocoderLoader, "diffusers"),
             "text_encoder": (TextEncoderLoader, "transformers"),
             "text_encoder_2": (TextEncoderLoader, "transformers"),
             "tokenizer": (TokenizerLoader, "transformers"),
@@ -245,47 +242,6 @@ class TextEncoderLoader(ComponentLoader):
         model_config.pop("model_type", None)
         model_config.pop("tokenizer_class", None)
         model_config.pop("torch_dtype", None)
-        repo_root = os.path.dirname(model_path)
-        index_path = os.path.join(repo_root, "model_index.json")
-        gemma_path = ""
-        gemma_path_from_candidate = False
-        if os.path.isfile(index_path):
-            try:
-                with open(index_path, encoding="utf-8") as f:
-                    model_index = json.load(f)
-                gemma_path = model_index.get("gemma_model_path", "")
-            except json.JSONDecodeError:
-                gemma_path = ""
-        if not gemma_path:
-            candidate = os.path.normpath(os.path.join(model_path, "gemma"))
-            if os.path.isdir(candidate):
-                gemma_path = candidate
-                gemma_path_from_candidate = True
-                model_config["gemma_model_path"] = gemma_path
-        if gemma_path and not gemma_path_from_candidate:
-            if not os.path.isabs(gemma_path):
-                model_config["gemma_model_path"] = os.path.normpath(
-                    os.path.join(repo_root, gemma_path)
-                )
-        transformer_config_path = os.path.join(
-            repo_root, "transformer", "config.json"
-        )
-        if os.path.isfile(transformer_config_path):
-            try:
-                with open(transformer_config_path, encoding="utf-8") as f:
-                    transformer_config = json.load(f)
-                if (
-                    "connector_double_precision_rope" not in model_config
-                    or not model_config["connector_double_precision_rope"]
-                ):
-                    if transformer_config.get("double_precision_rope") is True:
-                        model_config["connector_double_precision_rope"] = True
-                if "connector_rope_type" not in model_config:
-                    rope_type = transformer_config.get("rope_type")
-                    if rope_type is not None:
-                        model_config["connector_rope_type"] = rope_type
-            except json.JSONDecodeError:
-                pass
         logger.info("HF Model config: %s", model_config)
 
         # @TODO(Wei): Better way to handle this?
@@ -324,7 +280,7 @@ class TextEncoderLoader(ComponentLoader):
         target_device: torch.device,
         fastvideo_args: FastVideoArgs,
         dtype: str = "fp16",
-        use_text_encoder_override: bool = False,  # prevent subclasses from misusing
+        use_text_encoder_override: bool = False, # prevent subclasses from misusing
     ):
         use_cpu_offload = (
             fastvideo_args.text_encoder_cpu_offload
@@ -341,10 +297,7 @@ class TextEncoderLoader(ComponentLoader):
             )
 
         # Set quantization config if specified
-        if (
-            use_text_encoder_override
-            and fastvideo_args.override_text_encoder_quant is not None
-        ):
+        if use_text_encoder_override and fastvideo_args.override_text_encoder_quant is not None:
             if fastvideo_args.override_text_encoder_safetensors is None:
                 raise ValueError(
                     "override_text_encoder_quant is set but override_text_encoder_safetensors is None"
@@ -361,10 +314,7 @@ class TextEncoderLoader(ComponentLoader):
                 model: TextEncoder = model_cls(model_config)  # type: ignore
 
             weights_to_load = {name for name, _ in model.named_parameters()}
-            if (
-                use_text_encoder_override
-                and fastvideo_args.override_text_encoder_safetensors is not None
-            ):
+            if use_text_encoder_override and fastvideo_args.override_text_encoder_safetensors is not None:
                 loaded_weights: set[str] = model.load_weights(
                     safetensors_weights_iterator(
                         [fastvideo_args.override_text_encoder_safetensors],
@@ -392,6 +342,7 @@ class TextEncoderLoader(ComponentLoader):
 
             if use_cpu_offload:
                 pin_cpu_memory = fastvideo_args.pin_cpu_memory and is_pin_memory_available()
+                pin_cpu_memory = False
                 # Disable FSDP for MPS as it's not compatible
                 if current_platform.is_mps():
                     logger.info(
@@ -533,20 +484,8 @@ class TokenizerLoader(ComponentLoader):
             # in v0, this was same string as encoder_name "ClipTextModel"
             # TODO(will): pass these tokenizer kwargs from inference args? Maybe
             # other method of config?
+            padding_size="right",
         )
-        padding_side = None
-        if hasattr(fastvideo_args.pipeline_config, "text_encoder_configs"):
-            try:
-                arch_config = fastvideo_args.pipeline_config.text_encoder_configs[
-                    0
-                ].arch_config
-                padding_side = getattr(arch_config, "padding_side", None)
-            except Exception:
-                padding_side = None
-        if padding_side:
-            tokenizer.padding_side = padding_side
-        if tokenizer.pad_token is None and tokenizer.eos_token is not None:
-            tokenizer.pad_token = tokenizer.eos_token
         logger.info("Loaded tokenizer: %s", tokenizer.__class__.__name__)
         return tokenizer
 
@@ -557,12 +496,15 @@ class VAELoader(ComponentLoader):
     def load(self, model_path: str, fastvideo_args: FastVideoArgs):
         """Load the VAE based on the model path, and inference args."""
         config = get_diffusers_config(model=model_path)
-        class_name = config.get("_class_name")
+        class_name = config.pop("_class_name")
         assert class_name is not None, (
             "Model config does not contain a _class_name attribute. Only diffusers format is supported."
         )
         fastvideo_args.model_paths["vae"] = model_path
 
+        vae_config = fastvideo_args.pipeline_config.vae_config
+        vae_config.update_model_arch(config)
+
         from fastvideo.platforms import current_platform
 
         if fastvideo_args.vae_cpu_offload:
@@ -596,29 +538,8 @@ class VAELoader(ComponentLoader):
                 vae.load_state_dict(sd, strict=False)
                 return vae.eval()
 
-            # LTX-2 uses CausalVideoAutoencoder with nested "vae" config
-            if class_name == "CausalVideoAutoencoder" and "vae" in config:
-                vae_cls, _ = ModelRegistry.resolve_model_cls(class_name)
-                vae = vae_cls(config).to(target_device)
-                if hasattr(vae, "set_tiling_config"):
-                    vae_config = fastvideo_args.pipeline_config.vae_config
-                    vae.set_tiling_config(
-                        spatial_tile_size_in_pixels=getattr(
-                            vae_config, "ltx2_spatial_tile_size_in_pixels", 512),
-                        spatial_tile_overlap_in_pixels=getattr(
-                            vae_config, "ltx2_spatial_tile_overlap_in_pixels", 64),
-                        temporal_tile_size_in_frames=getattr(
-                            vae_config, "ltx2_temporal_tile_size_in_frames", 64),
-                        temporal_tile_overlap_in_frames=getattr(
-                            vae_config,
-                            "ltx2_temporal_tile_overlap_in_frames", 24),
-                    )
-            else:
-                config.pop("_class_name", None)
-                vae_config = fastvideo_args.pipeline_config.vae_config
-                vae_config.update_model_arch(config)
-                vae_cls, _ = ModelRegistry.resolve_model_cls(class_name)
-                vae = vae_cls(vae_config).to(target_device)
+            vae_cls, _ = ModelRegistry.resolve_model_cls(class_name)
+            vae = vae_cls(vae_config).to(target_device)
 
         # Find all safetensors files
         safetensors_list = glob.glob(
@@ -627,101 +548,17 @@ class VAELoader(ComponentLoader):
             raise ValueError(f"No safetensors files found in {model_path}")
         # Common case: a single `.safetensors` checkpoint file.
         # Some models may be sharded into multiple files; in that case we merge.
-        loaded = {}
-        for sf_file in safetensors_list:
-            loaded.update(safetensors_load_file(sf_file))
-
-        # LTX-2 CausalVideoAutoencoder needs per_channel_statistics remapping
-        if class_name == "CausalVideoAutoencoder" and "vae" in config:
-            per_channel_prefixes = (
-                "per_channel_statistics.",
-                "vae.per_channel_statistics.",
-            )
-            remapped = {}
-            for key, tensor in loaded.items():
-                remapped[key] = tensor
-                for prefix in per_channel_prefixes:
-                    if key.startswith(prefix):
-                        suffix = key[len(prefix):]
-                        remapped.setdefault(
-                            f"encoder.per_channel_statistics.{suffix}",
-                            tensor,
-                        )
-                        remapped.setdefault(
-                            f"decoder.per_channel_statistics.{suffix}",
-                            tensor,
-                        )
-                        break
-            loaded = remapped
-
+        if len(safetensors_list) == 1:
+            loaded = safetensors_load_file(safetensors_list[0])
+        else:
+            loaded = {}
+            for sf_file in safetensors_list:
+                loaded.update(safetensors_load_file(sf_file))
         vae.load_state_dict(loaded, strict=False)
 
         return vae.eval()
 
 
-class AudioDecoderLoader(ComponentLoader):
-    """Loader for LTX-2 audio decoder (audio_vae component)."""
-
-    def load(self, model_path: str, fastvideo_args: FastVideoArgs):
-        config = get_diffusers_config(model=model_path)
-        class_name = config.pop("_class_name", None) or "LTX2AudioDecoder"
-
-        model_cls, _ = ModelRegistry.resolve_model_cls(class_name)
-        target_device = get_local_torch_device()
-
-        precision = getattr(
-            fastvideo_args.pipeline_config, "audio_decoder_precision", "bf16"
-        )
-        with set_default_torch_dtype(PRECISION_TO_TYPE[precision]):
-            audio_decoder = 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))
-
-        decoder_state = {}
-        for name, tensor in loaded.items():
-            if name.startswith("decoder."):
-                decoder_state[name.replace("decoder.", "")] = tensor
-            elif name.startswith("per_channel_statistics."):
-                decoder_state[name] = tensor
-
-        target_module = getattr(audio_decoder, "model", audio_decoder)
-        target_module.load_state_dict(decoder_state, strict=False)
-        return audio_decoder.eval()
-
-
-class VocoderLoader(ComponentLoader):
-    """Loader for LTX-2 vocoder."""
-
-    def load(self, model_path: str, fastvideo_args: FastVideoArgs):
-        config = get_diffusers_config(model=model_path)
-        class_name = config.pop("_class_name", None) or "LTX2Vocoder"
-
-        model_cls, _ = ModelRegistry.resolve_model_cls(class_name)
-        target_device = get_local_torch_device()
-
-        precision = getattr(
-            fastvideo_args.pipeline_config, "vocoder_precision", "bf16"
-        )
-        with set_default_torch_dtype(PRECISION_TO_TYPE[precision]):
-            vocoder = 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(vocoder, "model", vocoder)
-        target_module.load_state_dict(loaded, strict=False)
-        return vocoder.eval()
-
-
 class TransformerLoader(ComponentLoader):
     """Loader for transformer."""
 
@@ -836,19 +673,35 @@ class TransformerLoader(ComponentLoader):
 
         model = model.eval()
 
-        if fastvideo_args.inference_mode and fastvideo_args.dit_layerwise_offload:
-            # Check if model has nn.ModuleList for layerwise offload compatibility
-            has_module_list = any(
-                isinstance(m, nn.ModuleList) for m in model.children()
-            )
-            if has_module_list:
-                enable_layerwise_offload(model)
-            else:
+        if fastvideo_args.dit_layerwise_offload and hasattr(model, "blocks"):
+            # Check if this is a Wan model (only Wan models support layerwise offload)
+            is_wan_model = "Wan" in cls_name
+            if not is_wan_model:
                 logger.warning(
-                    "Layerwise offload requested but model %s does not have "
-                    "nn.ModuleList structure. Skipping layerwise offload.",
+                    "Layerwise offload is currently only supported for Wan models. "
+                    "Model class '%s' does not support layerwise offload. "
+                    "Disabling layerwise offload for this model.",
                     cls_name
                 )
+            else:
+                try:
+                    num_layers = len(getattr(model, "blocks"))
+                except TypeError:
+                    num_layers = None
+                if isinstance(num_layers, int) and num_layers > 0:
+                    # Ensure model is on the correct device (CUDA) before initializing manager
+                    # This ensures non-managed parameters (embeddings, final norms) are on GPU
+                    model = model.to(get_local_torch_device())
+                    mgr = LayerwiseOffloadManager(
+                        model,
+                        module_list_attr="blocks",
+                        num_layers=num_layers,
+                        enabled=True,
+                        pin_cpu_memory=fastvideo_args.pin_cpu_memory,
+                        auto_initialize=True,
+                    )
+                    setattr(model, "_layerwise_offload_manager", mgr)
+
         return model
 
 

fastvideo/models/loader/fsdp_load.py

@@ -108,6 +108,7 @@ def maybe_load_fsdp_model(
 
     if use_fsdp:
         pin_cpu_memory = pin_cpu_memory and is_pin_memory_available()
+        pin_cpu_memory = False
         world_size = hsdp_replicate_dim * hsdp_shard_dim
         if not training_mode and not fsdp_inference:
             hsdp_replicate_dim = world_size

fastvideo/models/registry.py

@@ -33,7 +33,6 @@ _TEXT_TO_VIDEO_DIT_MODELS = {
     "Cosmos25Transformer3DModel": ("dits", "cosmos2_5", "Cosmos25Transformer3DModel"),
     "LongCatVideoTransformer3DModel": ("dits", "longcat_video_dit", "LongCatVideoTransformer3DModel"),  # Wrapper (Phase 1)
     "LongCatTransformer3DModel": ("dits", "longcat", "LongCatTransformer3DModel"),  # Native (Phase 2)
-    "LTX2Transformer3DModel": ("dits", "ltx2", "LTX2Transformer3DModel"),
 }
 
 _IMAGE_TO_VIDEO_DIT_MODELS = {
@@ -55,7 +54,6 @@ _TEXT_ENCODER_MODELS = {
     "Reason1TextEncoder": ("encoders", "reason1", "Reason1TextEncoder"),
     "Qwen2_5_VLForConditionalGeneration":
     ("encoders", "reason1", "Reason1TextEncoder"),
-    "LTX2GemmaTextEncoderModel": ("encoders", "gemma", "LTX2GemmaTextEncoderModel"),
 }
 
 _IMAGE_ENCODER_MODELS: dict[str, tuple] = {
@@ -69,14 +67,7 @@ _VAE_MODELS = {
     ("vaes", "hunyuanvae", "AutoencoderKLHunyuanVideo"),
     "AutoencoderKLHunyuanVideo15": ("vaes", "hunyuan15vae", "AutoencoderKLHunyuanVideo15"),
     "AutoencoderKLWan": ("vaes", "wanvae", "AutoencoderKLWan"),
-    "AutoencoderKLStepvideo": ("vaes", "stepvideovae", "AutoencoderKLStepvideo"),
-    "CausalVideoAutoencoder": ("vaes", "ltx2vae", "LTX2CausalVideoAutoencoder"),
-}
-
-_AUDIO_MODELS = {
-    "LTX2AudioEncoder": ("audio", "ltx2_audio_vae", "LTX2AudioEncoder"),
-    "LTX2AudioDecoder": ("audio", "ltx2_audio_vae", "LTX2AudioDecoder"),
-    "LTX2Vocoder": ("audio", "ltx2_audio_vae", "LTX2Vocoder"),
+    "AutoencoderKLStepvideo": ("vaes", "stepvideovae", "AutoencoderKLStepvideo")
 }
 
 _SCHEDULERS = {
@@ -100,7 +91,6 @@ _FAST_VIDEO_MODELS = {
     **_TEXT_ENCODER_MODELS,
     **_IMAGE_ENCODER_MODELS,
     **_VAE_MODELS,
-    **_AUDIO_MODELS,
     **_SCHEDULERS,
 }
 

fastvideo/models/vaes/hunyuan15vae.py

@@ -114,11 +114,14 @@ class HunyuanVideo15AttnBlock(nn.Module):
         """
         seq_len = n_frame * n_hw
         mask = torch.full((seq_len, seq_len), float("-inf"), dtype=dtype, device=device)
+        
         for i in range(seq_len):
             i_frame = i // n_hw
             mask[i, : (i_frame + 1) * n_hw] = 0
         if batch_size is not None:
-            mask = mask.unsqueeze(0).expand(batch_size, -1, -1)
+            # mask = mask.unsqueeze(0).expand(batch_size, -1, -1)
+            mask = mask.unsqueeze(0).unsqueeze(1).expand(batch_size, 1, -1, -1)
+
         return mask
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
@@ -131,6 +134,7 @@ class HunyuanVideo15AttnBlock(nn.Module):
         value = self.to_v(x)
 
         batch_size, channels, frames, height, width = query.shape
+        
 
         query = query.reshape(batch_size, channels, frames * height * width).permute(0, 2, 1).unsqueeze(1).contiguous()
         key = key.reshape(batch_size, channels, frames * height * width).permute(0, 2, 1).unsqueeze(1).contiguous()

fastvideo/models/vaes/hunyuanvae.py

@@ -92,6 +92,11 @@ class HunyuanVAEAttention(nn.Module):
         value = value.view(batch_size, -1, self.heads, head_dim).transpose(1, 2)
 
         # Perform scaled dot-product attention
+        # If attention_mask is provided, it needs to be expanded for multi-head attention
+        if attention_mask is not None:
+            # Expand mask from [batch_size, seq_len, seq_len] to [batch_size, num_heads, seq_len, seq_len]
+            attention_mask = attention_mask.unsqueeze(1).expand(-1, self.heads, -1, -1)
+        
         hidden_states = F.scaled_dot_product_attention(query,
                                                        key,
                                                        value,

fastvideo/pipelines/basic/ltx2/ltx2_pipeline.py

@@ -1,150 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""
-LTX-2 text-to-video pipeline.
-"""
-
-import os
-from typing import Any
-
-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,
-                                        TextEncodingStage)
-
-logger = init_logger(__name__)
-
-
-class LTX2Pipeline(ComposedPipelineBase):
-
-    _required_config_modules = [
-        "text_encoder",
-        "tokenizer",
-        "transformer",
-        "vae",
-        "audio_vae",
-        "vocoder",
-    ]
-
-    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs):
-        self.add_stage(
-            stage_name="input_validation_stage",
-            stage=InputValidationStage(),
-        )
-
-        self.add_stage(
-            stage_name="prompt_encoding_stage",
-            stage=TextEncodingStage(
-                text_encoders=[self.get_module("text_encoder")],
-                tokenizers=[self.get_module("tokenizer")],
-            ),
-        )
-
-        self.add_stage(
-            stage_name="latent_preparation_stage",
-            stage=LTX2LatentPreparationStage(
-                transformer=self.get_module("transformer"), ),
-        )
-
-        self.add_stage(
-            stage_name="denoising_stage",
-            stage=LTX2DenoisingStage(
-                transformer=self.get_module("transformer"), ),
-        )
-
-        self.add_stage(
-            stage_name="audio_decoding_stage",
-            stage=LTX2AudioDecodingStage(
-                audio_decoder=self.get_module("audio_vae"),
-                vocoder=self.get_module("vocoder"),
-            ),
-        )
-
-        self.add_stage(
-            stage_name="decoding_stage",
-            stage=DecodingStage(vae=self.get_module("vae")),
-        )
-
-    def initialize_pipeline(self, fastvideo_args: FastVideoArgs):
-        tokenizer = self.get_module("tokenizer")
-        if tokenizer is not None:
-            tokenizer.padding_side = "left"
-            if tokenizer.pad_token is None:
-                tokenizer.pad_token = tokenizer.eos_token
-
-    def load_modules(
-        self,
-        fastvideo_args: FastVideoArgs,
-        loaded_modules: dict[str, Any] | None = None,
-    ) -> dict[str, Any]:
-        model_index = self._load_config(self.model_path)
-        logger.info("Loading pipeline modules from config: %s", model_index)
-
-        model_index.pop("_class_name")
-        model_index.pop("_diffusers_version")
-        model_index.pop("workload_type", None)
-
-        if len(model_index) <= 1:
-            raise ValueError(
-                "model_index.json must contain at least one pipeline module")
-
-        required_modules = self.required_config_modules
-        modules: dict[str, Any] = {}
-
-        for module_name, module_spec in model_index.items():
-            if not isinstance(module_spec, list) or len(module_spec) < 1:
-                continue
-            transformers_or_diffusers = module_spec[0]
-            if transformers_or_diffusers is None:
-                if module_name in self.required_config_modules:
-                    self.required_config_modules.remove(module_name)
-                continue
-            if module_name not in required_modules:
-                continue
-            if loaded_modules is not None and module_name in loaded_modules:
-                modules[module_name] = loaded_modules[module_name]
-                continue
-
-            component_model_path = os.path.join(self.model_path, module_name)
-            if module_name == "tokenizer" and not os.path.isdir(
-                    component_model_path):
-                gemma_path = os.path.join(self.model_path, "text_encoder",
-                                          "gemma")
-                if os.path.isdir(gemma_path):
-                    component_model_path = gemma_path
-                else:
-                    raise ValueError(
-                        "Tokenizer directory missing and Gemma weights were not found."
-                    )
-
-            module = PipelineComponentLoader.load_module(
-                module_name=module_name,
-                component_model_path=component_model_path,
-                transformers_or_diffusers=transformers_or_diffusers,
-                fastvideo_args=fastvideo_args,
-            )
-            logger.info("Loaded module %s from %s", module_name,
-                        component_model_path)
-            modules[module_name] = module
-
-        if "tokenizer" in required_modules and "tokenizer" not in modules:
-            gemma_path = os.path.join(self.model_path, "text_encoder", "gemma")
-            if os.path.isdir(gemma_path):
-                modules["tokenizer"] = AutoTokenizer.from_pretrained(
-                    gemma_path, local_files_only=True)
-
-        for module_name in required_modules:
-            if module_name not in modules or modules[module_name] is None:
-                raise ValueError(
-                    f"Required module {module_name} was not loaded properly")
-
-        return modules
-
-
-EntryClass = LTX2Pipeline

fastvideo/pipelines/lora_pipeline.py

@@ -1,9 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 from collections import defaultdict
 from collections.abc import Hashable
-from contextlib import nullcontext
 from typing import Any
-from collections.abc import Generator
 
 import torch
 import torch.distributed as dist
@@ -14,13 +12,8 @@ from torch.distributed.tensor import DTensor
 
 from fastvideo.distributed import get_local_torch_device
 from fastvideo.fastvideo_args import FastVideoArgs, TrainingArgs
-from fastvideo.hooks.hooks import ModuleHookManager
-from fastvideo.hooks.layerwise_offload import LayerwiseOffloadHook
-from fastvideo.layers.lora.linear import (
-    BaseLayerWithLoRA,
-    get_lora_layer,
-    replace_submodule,
-)
+from fastvideo.layers.lora.linear import (BaseLayerWithLoRA, get_lora_layer,
+                                          replace_submodule)
 from fastvideo.logger import init_logger
 from fastvideo.models.loader.utils import get_param_names_mapping
 from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase
@@ -29,89 +22,17 @@ from fastvideo.utils import maybe_download_lora
 logger = init_logger(__name__)
 
 
-def _get_hook_ctx(module: nn.Module | None):
-    if module is None:
-        return nullcontext()
-    hook_mgr = ModuleHookManager.get_from(module)
-    if hook_mgr is not None:
-        offload_hook = hook_mgr.forward_hooks.get(LayerwiseOffloadHook.name())
-        if offload_hook is not None:
-            return offload_hook.mutate_params_scope()  # type: ignore
-    return nullcontext()
-
-
-def _named_module_by_prefix(
-    module: nn.Module, prefixes: list[str]
-) -> list[tuple[str | None, list[tuple[str, nn.Module]]]]:
-    none_list: list[tuple[str, nn.Module]] = []
-    prefix_list: list[tuple[str, list[tuple[str, nn.Module]]]] = [
-        (prefix, []) for prefix in prefixes
-    ]
-    for name, submodule in module.named_modules():
-        for cur_prefix, cur_list in prefix_list:
-            # we should exclude e.g. block.1 and block.12.attn
-            if name.startswith(cur_prefix + "."):
-                cur_list.append((name, submodule))
-                break
-        else:
-            none_list.append((name, submodule))
-    return prefix_list + [(None, none_list)]  # type: ignore
-
-
-class LoRAModelLayers:
-
-    def __init__(self, block_list: list[tuple[str, nn.Module]]) -> None:
-        # block_name -> {layer_name -> layer}
-        self.block_to_lora_layers: dict[str, dict[str, BaseLayerWithLoRA]] = {}
-        # layer_name -> block_name
-        self.lora_layers_to_block: dict[str, str | None] = {}
-        self.other_lora_layers: dict[str, BaseLayerWithLoRA] = {}
-        self.block_mapping = dict(block_list)
-
-    def add_lora_layer(self, block_name: str | None, layer_name: str,
-                       layer: BaseLayerWithLoRA):
-        if block_name is None:
-            self.other_lora_layers[layer_name] = layer
-            self.lora_layers_to_block[layer_name] = None
-        else:
-            if block_name not in self.block_to_lora_layers:
-                self.block_to_lora_layers[block_name] = {}
-            self.block_to_lora_layers[block_name][layer_name] = layer
-            self.lora_layers_to_block[layer_name] = block_name
-
-    def all_lora_layers(
-        self, ) -> Generator[tuple[str, BaseLayerWithLoRA], Any, None]:
-        for block_layers in self.block_to_lora_layers.values():
-            for name, layer in block_layers.items():
-                yield name, layer
-        for name, layer in self.other_lora_layers.items():
-            yield name, layer
-
-    def lora_layers_by_block(
-        self,
-    ) -> Generator[
-            tuple[nn.Module | None, dict[str, BaseLayerWithLoRA]],
-            Any,
-            None,
-    ]:
-        for block_name, layers in self.block_to_lora_layers.items():
-            yield self.block_mapping[block_name], layers
-        yield None, self.other_lora_layers
-
-
 class LoRAPipeline(ComposedPipelineBase):
     """
     Pipeline that supports injecting LoRA adapters into the diffusion transformer.
     TODO: support training.
     """
-
     lora_adapters: dict[str, dict[str, torch.Tensor]] = defaultdict(
         dict
     )  # state dicts of loaded lora adapters (includes lora_A, lora_B, and lora_alpha)
     cur_adapter_name: str = ""
     cur_adapter_path: str = ""
-    # model_name -> layers
-    lora_layers: dict[str, LoRAModelLayers] = {}
+    lora_layers: dict[str, dict[str, BaseLayerWithLoRA]] = {}
     fastvideo_args: FastVideoArgs | TrainingArgs
     exclude_lora_layers: dict[str, list[str]] = {}
     device: torch.device = get_local_torch_device()
@@ -127,10 +48,10 @@ class LoRAPipeline(ComposedPipelineBase):
         self.device = get_local_torch_device()
         # build list of trainable transformers
         for transformer_name in self.trainable_transformer_names:
-            if (transformer_name in self.modules
-                    and self.modules[transformer_name] is not None):
-                self.trainable_transformer_modules[transformer_name] = (
-                    self.modules[transformer_name])
+            if transformer_name in self.modules and self.modules[
+                    transformer_name] is not None:
+                self.trainable_transformer_modules[
+                    transformer_name] = self.modules[transformer_name]
             # check for transformer_2 in case of Wan2.2 MoE or fake_score_transformer_2
             if transformer_name.endswith("_2"):
                 raise ValueError(
@@ -138,23 +59,19 @@ class LoRAPipeline(ComposedPipelineBase):
                 )
 
             secondary_transformer_name = transformer_name + "_2"
-            if (secondary_transformer_name in self.modules
-                    and self.modules[secondary_transformer_name] is not None):
+            if secondary_transformer_name in self.modules and self.modules[
+                    secondary_transformer_name] is not None:
                 self.trainable_transformer_modules[
                     secondary_transformer_name] = self.modules[
                         secondary_transformer_name]
 
-        logger.info(
-            "trainable_transformer_modules: %s",
-            self.trainable_transformer_modules.keys(),
-        )
+        logger.info("trainable_transformer_modules: %s",
+                    self.trainable_transformer_modules.keys())
 
-        for (
-                transformer_name,
-                transformer_module,
-        ) in self.trainable_transformer_modules.items():
-            self.exclude_lora_layers[transformer_name] = (
-                transformer_module.config.arch_config.exclude_lora_layers)
+        for transformer_name, transformer_module in self.trainable_transformer_modules.items(
+        ):
+            self.exclude_lora_layers[
+                transformer_name] = transformer_module.config.arch_config.exclude_lora_layers
         self.lora_target_modules = self.fastvideo_args.lora_target_modules
         self.lora_path = self.fastvideo_args.lora_path
         self.lora_nickname = self.fastvideo_args.lora_nickname
@@ -166,33 +83,20 @@ class LoRAPipeline(ComposedPipelineBase):
                 self.fastvideo_args.lora_alpha = self.fastvideo_args.lora_rank
             self.lora_rank = self.fastvideo_args.lora_rank  # type: ignore
             self.lora_alpha = self.fastvideo_args.lora_alpha  # type: ignore
-            logger.info(
-                "Using LoRA training with rank %d and alpha %d",
-                self.lora_rank,
-                self.lora_alpha,
-            )
+            logger.info("Using LoRA training with rank %d and alpha %d",
+                        self.lora_rank, self.lora_alpha)
             if self.lora_target_modules is None:
                 self.lora_target_modules = [
-                    "q_proj",
-                    "k_proj",
-                    "v_proj",
-                    "o_proj",
-                    "to_q",
-                    "to_k",
-                    "to_v",
-                    "to_out",
-                    "to_qkv",
-                    "to_gate_compress",
+                    "q_proj", "k_proj", "v_proj", "o_proj", "to_q", "to_k",
+                    "to_v", "to_out", "to_qkv", "to_gate_compress"
                 ]
                 logger.info(
                     "Using default lora_target_modules for all transformers: %s",
-                    self.lora_target_modules,
-                )
+                    self.lora_target_modules)
             else:
                 logger.warning(
                     "Using custom lora_target_modules for all transformers, which may not be intended: %s",
-                    self.lora_target_modules,
-                )
+                    self.lora_target_modules)
 
             self.convert_to_lora_layers()
         # Inference
@@ -200,8 +104,7 @@ class LoRAPipeline(ComposedPipelineBase):
             self.convert_to_lora_layers()
             self.set_lora_adapter(
                 self.lora_nickname,  # type: ignore
-                self.lora_path,
-            )  # type: ignore
+                self.lora_path)  # type: ignore
 
     def is_target_layer(self, module_name: str) -> bool:
         if self.lora_target_modules is None:
@@ -211,9 +114,9 @@ class LoRAPipeline(ComposedPipelineBase):
 
     def set_trainable(self) -> None:
 
-        def set_lora_grads(lora_layers: LoRAModelLayers,
+        def set_lora_grads(lora_layers: dict[str, BaseLayerWithLoRA],
                            device_mesh: DeviceMesh):
-            for name, layer in lora_layers.all_lora_layers():
+            for name, layer in lora_layers.items():
                 layer.lora_A.requires_grad_(True)
                 layer.lora_B.requires_grad_(True)
                 layer.base_layer.requires_grad_(False)
@@ -228,15 +131,10 @@ class LoRAPipeline(ComposedPipelineBase):
             super().set_trainable()
             return
 
-        device_mesh = init_device_mesh(
-            "cuda",
-            (dist.get_world_size(), 1),
-            mesh_dim_names=["fake", "replicate"],
-        )
-        for (
-                transformer_name,
-                transformer_module,
-        ) in self.trainable_transformer_modules.items():
+        device_mesh = init_device_mesh("cuda", (dist.get_world_size(), 1),
+                                       mesh_dim_names=["fake", "replicate"])
+        for transformer_name, transformer_module in self.trainable_transformer_modules.items(
+        ):
             transformer_module.train()
             transformer_module.requires_grad_(False)
             if transformer_name in self.lora_layers:
@@ -253,71 +151,32 @@ class LoRAPipeline(ComposedPipelineBase):
         if self.lora_initialized:
             return
         self.lora_initialized = True
-        for (
-                transformer_name,
-                transformer_module,
-        ) in self.trainable_transformer_modules.items():
+        for transformer_name, transformer_module in self.trainable_transformer_modules.items(
+        ):
             converted_count = 0
-            # init bookkeeping structures
             if transformer_name not in self.lora_layers:
-                # get block list
-                block_list = []
-                for name, submodule in transformer_module.named_children():
-                    if isinstance(submodule, nn.ModuleList):
-                        block_list = [(f"{name}.{i}", m)
-                                      for i, m in enumerate(submodule)]
-                        break
-                self.lora_layers[transformer_name] = LoRAModelLayers(block_list)
+                self.lora_layers[transformer_name] = {}
             logger.info("Converting %s to LoRA Transformer", transformer_name)
-            # scan every module and convert to LoRA layer if applicable
+            for name, layer in transformer_module.named_modules():
+                if not self.is_target_layer(name):
+                    continue
 
-            for block_name, block_modules in _named_module_by_prefix(
-                    transformer_module,
-                    list(self.lora_layers[transformer_name].block_mapping),
-            ):
-                if block_name is not None and (
-                        not self.fastvideo_args.training_mode
-                        and self.fastvideo_args.dit_layerwise_offload):
-                    scope_ctx = _get_hook_ctx(
-                        self.lora_layers[transformer_name].
-                        block_mapping[block_name])
-                else:
-                    scope_ctx = nullcontext()
-                with scope_ctx:
-                    for name, layer in block_modules:
-                        if not self.is_target_layer(name):
-                            continue
+                excluded = False
+                for exclude_layer in self.exclude_lora_layers[transformer_name]:
+                    if exclude_layer in name:
+                        excluded = True
+                        break
+                if excluded:
+                    continue
 
-                        excluded = False
-                        for exclude_layer in self.exclude_lora_layers[
-                                transformer_name]:
-                            if exclude_layer in name:
-                                excluded = True
-                                break
-                        if excluded:
-                            continue
-
-                        layer = get_lora_layer(
-                            layer,
-                            lora_rank=self.lora_rank,
-                            lora_alpha=self.lora_alpha,
-                            training_mode=self.training_mode,
-                        )
-                        if layer is not None:
-                            block_name_split = name.split(".", 2)
-                            if len(block_name_split) > 2:
-                                block_name = (block_name_split[0] + "." +
-                                              block_name_split[1])
-                            else:
-                                block_name = None
-                            if (block_name
-                                    not in self.lora_layers[transformer_name].
-                                    block_mapping):
-                                block_name = None
-                            self.lora_layers[transformer_name].add_lora_layer(
-                                block_name, name, layer)
-                            replace_submodule(transformer_module, name, layer)
-                            converted_count += 1
+                layer = get_lora_layer(layer,
+                                       lora_rank=self.lora_rank,
+                                       lora_alpha=self.lora_alpha,
+                                       training_mode=self.training_mode)
+                if layer is not None:
+                    self.lora_layers[transformer_name][name] = layer
+                    replace_submodule(transformer_module, name, layer)
+                    converted_count += 1
             logger.info("Converted %d layers to LoRA layers", converted_count)
 
     def set_lora_adapter(self,
@@ -350,7 +209,7 @@ class LoRAPipeline(ComposedPipelineBase):
 
             # Extract alpha values and weights in a single pass
             to_merge_params: defaultdict[Hashable,
-                                         dict[Any, Any]] = (defaultdict(dict))
+                                         dict[Any, Any]] = defaultdict(dict)
             for name, weight in lora_state_dict.items():
                 # Extract weights (lora_A, lora_B, and lora_alpha)
                 name = name.replace("diffusion_model.", "")
@@ -364,14 +223,13 @@ class LoRAPipeline(ComposedPipelineBase):
                     target_name, _, _ = param_names_mapping_fn(layer_name)
                     # Store alpha alongside weights with same target_name base
                     alpha_key = target_name + ".lora_alpha"
-                    self.lora_adapters[lora_nickname][alpha_key] = (
-                        weight.item()
-                        if weight.numel() == 1 else float(weight.mean()))
+                    self.lora_adapters[lora_nickname][alpha_key] = weight.item(
+                    ) if weight.numel() == 1 else float(weight.mean())
                     continue
 
                 name, _, _ = lora_param_names_mapping_fn(name)
-                target_name, merge_index, num_params_to_merge = (
-                    param_names_mapping_fn(name))
+                target_name, merge_index, num_params_to_merge = param_names_mapping_fn(
+                    name)
                 # for (in_dim, r) @ (r, out_dim), we only merge (r, out_dim * n) where n is the number of linear layers to fuse
                 # see param mapping in HunyuanVideoArchConfig
                 if merge_index is not None and "lora_B" in name:
@@ -403,84 +261,45 @@ class LoRAPipeline(ComposedPipelineBase):
 
         # Merge the new adapter
         adapted_count = 0
-        for (
-                transformer_name,
-                transformer_lora_layers,
-        ) in self.lora_layers.items():
-            for (
-                    module,
-                    layers,
-            ) in transformer_lora_layers.lora_layers_by_block():
-                with _get_hook_ctx(module):
-                    for name, layer in layers.items():
-                        lora_A_name = name + ".lora_A"
-                        lora_B_name = name + ".lora_B"
-                        lora_alpha_name = name + ".lora_alpha"
-                        if (lora_A_name in self.lora_adapters[lora_nickname]
-                                and lora_B_name
-                                in self.lora_adapters[lora_nickname]):
-                            # Get alpha value for this layer (defaults to None if not present)
-                            lora_A = self.lora_adapters[lora_nickname][
-                                lora_A_name]
-                            lora_B = self.lora_adapters[lora_nickname][
-                                lora_B_name]
-                            # Simple lookup - alpha stored with same naming scheme as lora_A/lora_B
-                            alpha = (self.lora_adapters[lora_nickname].get(
-                                lora_alpha_name) if adapter_updated else None)
-                            try:
-                                layer.set_lora_weights(
-                                    lora_A,
-                                    lora_B,
-                                    lora_alpha=alpha,
-                                    training_mode=self.fastvideo_args.
-                                    training_mode,
-                                    lora_path=lora_path,
-                                )
-                            except Exception as e:
-                                logger.error(
-                                    "Error setting LoRA weights for layer %s: %s",
-                                    name,
-                                    str(e),
-                                )
-                                raise e
-                            adapted_count += 1
-                        else:
-                            if rank == 0:
-                                logger.warning(
-                                    "LoRA adapter %s does not contain the weights for layer %s. LoRA will not be applied to it.",
-                                    lora_path,
-                                    name,
-                                )
-                            layer.disable_lora = True
-        logger.info(
-            "Rank %d: LoRA adapter %s applied to %d layers",
-            rank,
-            lora_path,
-            adapted_count,
-        )
+        for transformer_name, transformer_lora_layers in self.lora_layers.items(
+        ):
+            for name, layer in transformer_lora_layers.items():
+                lora_A_name = name + ".lora_A"
+                lora_B_name = name + ".lora_B"
+                lora_alpha_name = name + ".lora_alpha"
+                if lora_A_name in self.lora_adapters[lora_nickname]\
+                    and lora_B_name in self.lora_adapters[lora_nickname]:
+                    # Get alpha value for this layer (defaults to None if not present)
+                    lora_A = self.lora_adapters[lora_nickname][lora_A_name]
+                    lora_B = self.lora_adapters[lora_nickname][lora_B_name]
+                    # Simple lookup - alpha stored with same naming scheme as lora_A/lora_B
+                    alpha = self.lora_adapters[lora_nickname].get(
+                        lora_alpha_name) if adapter_updated else None
+
+                    layer.set_lora_weights(
+                        lora_A,
+                        lora_B,
+                        lora_alpha=alpha,
+                        training_mode=self.fastvideo_args.training_mode,
+                        lora_path=lora_path)
+                    adapted_count += 1
+                else:
+                    if rank == 0:
+                        logger.warning(
+                            "LoRA adapter %s does not contain the weights for layer %s. LoRA will not be applied to it.",
+                            lora_path, name)
+                    layer.disable_lora = True
+        logger.info("Rank %d: LoRA adapter %s applied to %d layers", rank,
+                    lora_path, adapted_count)
 
     def merge_lora_weights(self) -> None:
-        for (
-                transformer_name,
-                transformer_lora_layers,
-        ) in self.lora_layers.items():
-            for (
-                    module,
-                    layers,
-            ) in transformer_lora_layers.lora_layers_by_block():
-                with _get_hook_ctx(module):
-                    for name, layer in layers.items():
-                        layer.merge_lora_weights()
+        for transformer_name, transformer_lora_layers in self.lora_layers.items(
+        ):
+            for name, layer in transformer_lora_layers.items():
+                layer.merge_lora_weights()
 
     def unmerge_lora_weights(self) -> None:
-        for (
-                transformer_name,
-                transformer_lora_layers,
-        ) in self.lora_layers.items():
-            for (
-                    module,
-                    layers,
-            ) in transformer_lora_layers.lora_layers_by_block():
-                with _get_hook_ctx(module):
-                    for name, layer in layers.items():
-                        layer.unmerge_lora_weights()
+        for transformer_name, transformer_lora_layers in self.lora_layers.items(
+        ):
+            for name, layer in transformer_lora_layers.items():
+                layer.unmerge_lora_weights()

fastvideo/pipelines/pipeline_registry.py

@@ -35,7 +35,6 @@ _PIPELINE_NAME_TO_ARCHITECTURE_NAME: dict[str, str] = {
     "LongCatPipeline": "longcat",
     "LongCatImageToVideoPipeline": "longcat",
     "LongCatVideoContinuationPipeline": "longcat",
-    "LTX2Pipeline": "ltx2",
 }
 
 _PREPROCESS_WORKLOAD_TYPE_TO_PIPELINE_NAME: dict[WorkloadType, str] = {

fastvideo/pipelines/preprocess/hunyuan/hunyuan_preprocess_pipelins.py

@@ -0,0 +1,59 @@
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase
+from fastvideo.pipelines.preprocess.preprocess_stages import (
+    TextTransformStage, VideoTransformStage)
+from fastvideo.pipelines.stages import (EncodingStage, ImageEncodingStage,
+                                        TextEncodingStage)
+from fastvideo.pipelines.stages.image_encoding import ImageVAEEncodingStage
+        
+
+class PreprocessPipelineT2V(ComposedPipelineBase):
+    _required_config_modules = [
+        "text_encoder", "tokenizer", "text_encoder_2", "tokenizer_2", "vae"
+    ]
+
+    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs):
+        assert fastvideo_args.preprocess_config is not None
+        
+        self.add_stage(
+            stage_name="text_transform_stage",
+            stage=TextTransformStage(
+                cfg_uncondition_drop_rate=fastvideo_args.preprocess_config.training_cfg_rate,
+                seed=fastvideo_args.preprocess_config.seed,
+            )
+        )
+        
+        text_encoders = [
+            self.get_module("text_encoder"),
+            self.get_module("text_encoder_2")
+        ]
+        tokenizers = [
+            self.get_module("tokenizer"),
+            self.get_module("tokenizer_2")
+        ]
+        
+        self.add_stage(
+            stage_name="prompt_encoding_stage",
+            stage=TextEncodingStage(
+                text_encoders=text_encoders,
+                tokenizers=tokenizers,
+            )
+        )
+        
+        self.add_stage(
+            stage_name="video_transform_stage",
+            stage=VideoTransformStage(
+                train_fps=fastvideo_args.preprocess_config.train_fps,
+                num_frames=fastvideo_args.preprocess_config.num_frames,
+                max_height=fastvideo_args.preprocess_config.max_height,
+                max_width=fastvideo_args.preprocess_config.max_width,
+                do_temporal_sample=fastvideo_args.preprocess_config.do_temporal_sample,
+            )
+        )
+        
+        self.add_stage(
+            stage_name="video_encoding_stage",
+            stage=EncodingStage(vae=self.get_module("vae"))
+        )
+
+EntryClass = [PreprocessPipelineT2V]

fastvideo/pipelines/preprocess/hunyuan15/hunyuan_preprocess_pipelins.py

@@ -0,0 +1,59 @@
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase
+from fastvideo.pipelines.preprocess.preprocess_stages import (
+    TextTransformStage, VideoTransformStage)
+from fastvideo.pipelines.stages import (EncodingStage, ImageEncodingStage,
+                                        TextEncodingStage)
+from fastvideo.pipelines.stages.image_encoding import ImageVAEEncodingStage
+        
+
+class PreprocessPipelineT2V(ComposedPipelineBase):
+    _required_config_modules = [
+        "text_encoder", "tokenizer", "text_encoder_2", "tokenizer_2", "vae"
+    ]
+
+    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs):
+        assert fastvideo_args.preprocess_config is not None
+        
+        self.add_stage(
+            stage_name="text_transform_stage",
+            stage=TextTransformStage(
+                cfg_uncondition_drop_rate=fastvideo_args.preprocess_config.training_cfg_rate,
+                seed=fastvideo_args.preprocess_config.seed,
+            )
+        )
+        
+        text_encoders = [
+            self.get_module("text_encoder"),
+            self.get_module("text_encoder_2")
+        ]
+        tokenizers = [
+            self.get_module("tokenizer"),
+            self.get_module("tokenizer_2")
+        ]
+        
+        self.add_stage(
+            stage_name="prompt_encoding_stage",
+            stage=TextEncodingStage(
+                text_encoders=text_encoders,
+                tokenizers=tokenizers,
+            )
+        )
+        
+        self.add_stage(
+            stage_name="video_transform_stage",
+            stage=VideoTransformStage(
+                train_fps=fastvideo_args.preprocess_config.train_fps,
+                num_frames=fastvideo_args.preprocess_config.num_frames,
+                max_height=fastvideo_args.preprocess_config.max_height,
+                max_width=fastvideo_args.preprocess_config.max_width,
+                do_temporal_sample=fastvideo_args.preprocess_config.do_temporal_sample,
+            )
+        )
+        
+        self.add_stage(
+            stage_name="video_encoding_stage",
+            stage=EncodingStage(vae=self.get_module("vae"))
+        )
+
+EntryClass = [PreprocessPipelineT2V]

fastvideo/pipelines/preprocess/v1_preprocessing_new.py

@@ -5,10 +5,6 @@ from fastvideo.logger import init_logger
 from fastvideo.utils import FlexibleArgumentParser
 from fastvideo.workflow.workflow_base import WorkflowBase
 
-import os
-
-os.environ["MASTER_PORT"] = "29513"
-
 logger = init_logger(__name__)
 
 

fastvideo/pipelines/stages/__init__.py

@@ -22,10 +22,6 @@ from fastvideo.pipelines.stages.input_validation import InputValidationStage
 from fastvideo.pipelines.stages.latent_preparation import (
     Cosmos25LatentPreparationStage, CosmosLatentPreparationStage,
     LatentPreparationStage)
-from fastvideo.pipelines.stages.ltx2_audio_decoding import LTX2AudioDecodingStage
-from fastvideo.pipelines.stages.ltx2_denoising import LTX2DenoisingStage
-from fastvideo.pipelines.stages.ltx2_latent_preparation import (
-    LTX2LatentPreparationStage)
 from fastvideo.pipelines.stages.matrixgame_denoising import (
     MatrixGameCausalDenoisingStage)
 from fastvideo.pipelines.stages.stepvideo_encoding import (
@@ -48,8 +44,6 @@ __all__ = [
     "LatentPreparationStage",
     "CosmosLatentPreparationStage",
     "Cosmos25LatentPreparationStage",
-    "LTX2LatentPreparationStage",
-    "LTX2AudioDecodingStage",
     "ConditioningStage",
     "DenoisingStage",
     "DmdDenoisingStage",
@@ -57,7 +51,6 @@ __all__ = [
     "MatrixGameCausalDenoisingStage",
     "CosmosDenoisingStage",
     "Cosmos25DenoisingStage",
-    "LTX2DenoisingStage",
     "EncodingStage",
     "DecodingStage",
     "ImageEncodingStage",

fastvideo/pipelines/stages/ltx2_audio_decoding.py

@@ -1,67 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""
-Audio decoding stage for LTX-2 pipelines.
-"""
-
-from __future__ import annotations
-
-import os
-
-import torch
-
-from fastvideo.models.dits.ltx2 import DEFAULT_LTX2_VOCODER_OUTPUT_SAMPLE_RATE
-
-from fastvideo.distributed import get_local_torch_device
-from fastvideo.fastvideo_args import FastVideoArgs
-from fastvideo.logger import init_logger
-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__)
-
-
-class LTX2AudioDecodingStage(PipelineStage):
-    """Decode LTX-2 audio latents into a waveform."""
-
-    def __init__(self, audio_decoder, vocoder) -> None:
-        super().__init__()
-        self.audio_decoder = audio_decoder
-        self.vocoder = vocoder
-
-    def forward(
-        self,
-        batch: ForwardBatch,
-        fastvideo_args: FastVideoArgs,
-    ) -> ForwardBatch:
-        audio_latents = batch.extra.get("ltx2_audio_latents")
-        if audio_latents is None:
-            return batch
-
-        device = get_local_torch_device()
-        self.audio_decoder = self.audio_decoder.to(device)
-        self.vocoder = self.vocoder.to(device)
-        audio_latents = audio_latents.to(device)
-
-        disable_autocast = os.getenv("LTX2_DISABLE_AUDIO_AUTOCAST", "1") == "1"
-        with torch.no_grad(), torch.autocast(
-                device_type="cuda",
-                dtype=audio_latents.dtype,
-                enabled=not disable_autocast,
-        ):
-            decoded_spec = self.audio_decoder(audio_latents)
-            audio_wave = self.vocoder(decoded_spec).squeeze(0).float()
-
-        # Move to CPU for pickling across process boundary
-        batch.extra["audio"] = audio_wave.cpu()
-        batch.extra[
-            "audio_sample_rate"] = DEFAULT_LTX2_VOCODER_OUTPUT_SAMPLE_RATE
-        return batch
-
-    def verify_input(self, batch: ForwardBatch,
-                     fastvideo_args: FastVideoArgs) -> VerificationResult:
-        result = VerificationResult()
-        result.add_check("audio_latents", batch.extra.get("ltx2_audio_latents"),
-                         V.none_or_tensor)
-        return result

fastvideo/pipelines/stages/ltx2_denoising.py

@@ -1,308 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""
-LTX-2 denoising stage using the native sigma schedule.
-"""
-
-from __future__ import annotations
-
-import math
-import os
-
-import torch
-from tqdm.auto import tqdm
-
-from fastvideo.fastvideo_args import FastVideoArgs
-from fastvideo.forward_context import set_forward_context
-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
-from fastvideo.logger import init_logger
-from fastvideo.models.dits.ltx2 import (
-    AudioLatentShape, DEFAULT_LTX2_AUDIO_CHANNELS,
-    DEFAULT_LTX2_AUDIO_DOWNSAMPLE, DEFAULT_LTX2_AUDIO_HOP_LENGTH,
-    DEFAULT_LTX2_AUDIO_MEL_BINS, DEFAULT_LTX2_AUDIO_SAMPLE_RATE,
-    VideoLatentShape)
-from fastvideo.utils import PRECISION_TO_TYPE
-
-BASE_SHIFT_ANCHOR = 1024
-MAX_SHIFT_ANCHOR = 4096
-
-# Official distilled sigma schedule (8 denoising steps)
-# From LTX-2/packages/ltx-pipelines/src/ltx_pipelines/utils/constants.py
-DISTILLED_SIGMA_VALUES = [
-    1.0, 0.99375, 0.9875, 0.98125, 0.975, 0.909375, 0.725, 0.421875, 0.0
-]
-
-logger = init_logger(__name__)
-
-
-def _ltx2_sigmas(
-    steps: int,
-    latent: torch.Tensor | None,
-    device: torch.device,
-    max_shift: float = 2.05,
-    base_shift: float = 0.95,
-    stretch: bool = True,
-    terminal: float = 0.1,
-) -> torch.Tensor:
-    tokens = math.prod(
-        latent.shape[2:]) if latent is not None else MAX_SHIFT_ANCHOR
-    sigmas = torch.linspace(1.0,
-                            0.0,
-                            steps + 1,
-                            device=device,
-                            dtype=torch.float32)
-
-    mm = (max_shift - base_shift) / (MAX_SHIFT_ANCHOR - BASE_SHIFT_ANCHOR)
-    b = base_shift - mm * BASE_SHIFT_ANCHOR
-    sigma_shift = tokens * mm + b
-
-    numerator = math.exp(sigma_shift)
-    sigmas = torch.where(
-        sigmas != 0,
-        numerator / (numerator + (1 / sigmas - 1)),
-        torch.zeros_like(sigmas),
-    )
-
-    if stretch:
-        non_zero_mask = sigmas != 0
-        non_zero_sigmas = sigmas[non_zero_mask]
-        one_minus_z = 1.0 - non_zero_sigmas
-        scale_factor = one_minus_z[-1] / (1.0 - terminal)
-        stretched = 1.0 - (one_minus_z / scale_factor)
-        sigmas = sigmas.clone()
-        sigmas[non_zero_mask] = stretched
-
-    return sigmas
-
-
-class LTX2DenoisingStage(PipelineStage):
-    """Run the LTX-2 denoising loop over the sigma schedule."""
-
-    def __init__(self, transformer) -> None:
-        super().__init__()
-        self.transformer = transformer
-
-    def forward(
-        self,
-        batch: ForwardBatch,
-        fastvideo_args: FastVideoArgs,
-    ) -> ForwardBatch:
-        if batch.latents is None:
-            raise ValueError("Latents must be provided before denoising.")
-
-        latents = batch.latents
-        prompt_embeds = batch.prompt_embeds[0]
-        prompt_mask = None
-
-        neg_prompt_embeds = None
-        neg_prompt_mask = None
-        # Only load negative prompts if CFG is actually enabled
-        if batch.do_classifier_free_guidance:
-            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]
-
-        # Ensure text conditioning is on the same device as latents.
-        if prompt_embeds.device != latents.device:
-            prompt_embeds = prompt_embeds.to(latents.device)
-        if prompt_mask is not None and prompt_mask.device != latents.device:
-            prompt_mask = prompt_mask.to(latents.device)
-        if neg_prompt_embeds is not None and neg_prompt_embeds.device != latents.device:
-            neg_prompt_embeds = neg_prompt_embeds.to(latents.device)
-        if neg_prompt_mask is not None and neg_prompt_mask.device != latents.device:
-            neg_prompt_mask = neg_prompt_mask.to(latents.device)
-
-        target_dtype = PRECISION_TO_TYPE[
-            fastvideo_args.pipeline_config.dit_precision]
-        disable_autocast = os.getenv("LTX2_DISABLE_AUTOCAST", "1") == "1"
-        autocast_enabled = (target_dtype != torch.float32
-                            ) 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:
-            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=batch.num_inference_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(
-                video_shape)
-        else:
-            token_count = 1
-        timestep_template = torch.ones(
-            (latents.shape[0], token_count),
-            device=latents.device,
-            dtype=torch.float32,
-        )
-        audio_prompt_embeds = batch.extra.get("ltx2_audio_prompt_embeds")
-        audio_neg_embeds = batch.extra.get("ltx2_audio_negative_embeds")
-        audio_context_p = audio_prompt_embeds[0] if audio_prompt_embeds else None
-        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:
-            fps_value = batch.fps
-            if isinstance(fps_value, list):
-                fps_value = fps_value[0] if fps_value else None
-            if fps_value is None:
-                fps_value = 1.0
-            duration = float(batch.num_frames) / float(fps_value)
-            audio_shape = AudioLatentShape.from_duration(
-                batch=latents.shape[0],
-                duration=duration,
-                channels=DEFAULT_LTX2_AUDIO_CHANNELS,
-                mel_bins=DEFAULT_LTX2_AUDIO_MEL_BINS,
-                sample_rate=DEFAULT_LTX2_AUDIO_SAMPLE_RATE,
-                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 = (
-                audio_shape.batch,
-                audio_shape.frames,
-                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()
-            audio_timestep_template = torch.ones(
-                (latents.shape[0], audio_shape.frames),
-                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,
-            target_dtype,
-            batch.guidance_scale,
-            tuple(sigmas.shape),
-            tuple(latents.shape),
-        )
-
-        for step_index in tqdm(range(len(sigmas) - 1)):
-            sigma = sigmas[step_index]
-            sigma_next = sigmas[step_index + 1]
-            timestep = timestep_template * sigma
-            audio_timestep = (audio_timestep_template * sigma
-                              if audio_timestep_template is not None else None)
-
-            with torch.autocast(
-                    device_type="cuda",
-                    dtype=target_dtype,
-                    enabled=autocast_enabled,
-            ), set_forward_context(
-                    current_timestep=sigma.item(),
-                    attn_metadata=None,
-                    forward_batch=batch,
-            ):
-                pos_outputs = self.transformer(
-                    hidden_states=latents.to(target_dtype),
-                    encoder_hidden_states=prompt_embeds,
-                    encoder_attention_mask=prompt_mask,
-                    timestep=timestep,
-                    audio_hidden_states=audio_latents,
-                    audio_encoder_hidden_states=audio_context_p,
-                    audio_timestep=audio_timestep,
-                )
-                if isinstance(pos_outputs, tuple):
-                    pos_denoised, pos_audio = pos_outputs
-                else:
-                    pos_denoised = pos_outputs
-                    pos_audio = None
-
-                # Only run negative pass if CFG is enabled
-                if batch.do_classifier_free_guidance:
-                    neg_outputs = self.transformer(
-                        hidden_states=latents.to(target_dtype),
-                        encoder_hidden_states=neg_prompt_embeds,
-                        encoder_attention_mask=neg_prompt_mask,
-                        timestep=timestep,
-                        audio_hidden_states=audio_latents,
-                        audio_encoder_hidden_states=audio_context_n,
-                        audio_timestep=audio_timestep,
-                    )
-                    if isinstance(neg_outputs, tuple):
-                        neg_denoised, neg_audio = neg_outputs
-                    else:
-                        neg_denoised = neg_outputs
-                        neg_audio = None
-                    pos_denoised = pos_denoised + (batch.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 -
-                                                 1) * (pos_audio - neg_audio)
-
-            sigma_value = sigma.to(torch.float32) if isinstance(
-                sigma, torch.Tensor) else torch.tensor(
-                    float(sigma),
-                    device=latents.device,
-                    dtype=torch.float32,
-                )
-            dt = sigma_next - sigma
-            velocity = ((latents.float() - pos_denoised.float()) /
-                        sigma_value).to(latents.dtype)
-            latents = (latents.float() + velocity.float() * dt).to(
-                latents.dtype)
-            if pos_audio is not None and audio_latents is not None:
-                audio_velocity = ((audio_latents.float() - pos_audio.float()) /
-                                  sigma_value).to(audio_latents.dtype)
-                audio_latents = (audio_latents.float() +
-                                 audio_velocity.float() * dt).to(
-                                     audio_latents.dtype)
-
-        batch.latents = latents
-        batch.extra["ltx2_audio_latents"] = audio_latents
-        logger.info("[LTX2] Denoising done.")
-        return batch
-
-    def verify_input(self, batch: ForwardBatch,
-                     fastvideo_args: FastVideoArgs) -> VerificationResult:
-        result = VerificationResult()
-        result.add_check("latents", batch.latents,
-                         [V.is_tensor, V.with_dims(5)])
-        result.add_check("prompt_embeds", batch.prompt_embeds, V.list_not_empty)
-        result.add_check("num_inference_steps", batch.num_inference_steps,
-                         V.positive_int)
-        return result

fastvideo/pipelines/stages/ltx2_latent_preparation.py

@@ -1,189 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""
-Latent preparation stage for LTX-2 pipelines.
-"""
-
-from pathlib import Path
-
-import torch
-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.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__)
-
-
-class LTX2LatentPreparationStage(PipelineStage):
-    """Prepare initial LTX-2 latents without relying on a diffusers scheduler."""
-
-    def __init__(self, transformer) -> None:
-        super().__init__()
-        self.transformer = transformer
-
-    def forward(
-        self,
-        batch: ForwardBatch,
-        fastvideo_args: FastVideoArgs,
-    ) -> ForwardBatch:
-        latent_num_frames = self._adjust_video_length(batch, fastvideo_args)
-
-        if not batch.prompt_embeds:
-            batch_size = 1
-        elif isinstance(batch.prompt, list):
-            batch_size = len(batch.prompt)
-        elif batch.prompt is not None:
-            batch_size = 1
-        else:
-            batch_size = batch.prompt_embeds[0].shape[0]
-
-        batch_size *= batch.num_videos_per_prompt
-
-        if not batch.prompt_embeds:
-            transformer_dtype = next(self.transformer.parameters()).dtype
-            device = get_local_torch_device()
-            dummy_prompt = torch.zeros(
-                batch_size,
-                0,
-                self.transformer.hidden_size,
-                device=device,
-                dtype=transformer_dtype,
-            )
-            batch.prompt_embeds = [dummy_prompt]
-            batch.negative_prompt_embeds = []
-            batch.do_classifier_free_guidance = False
-
-        dtype = batch.prompt_embeds[0].dtype
-        device = get_local_torch_device()
-        generator = batch.generator
-        latents = batch.latents
-        num_frames = latent_num_frames if latent_num_frames is not None else batch.num_frames
-        height = batch.height
-        width = batch.width
-        latent_path = fastvideo_args.ltx2_initial_latent_path
-
-        if height is None or width is None:
-            raise ValueError("Height and width must be provided")
-
-        spatial_ratio = fastvideo_args.pipeline_config.vae_config.arch_config.spatial_compression_ratio
-        if height % spatial_ratio != 0 or width % spatial_ratio != 0:
-            raise ValueError(
-                f"Height and width must be divisible by {spatial_ratio} "
-                f"but are {height} and {width}.")
-        shape = (
-            batch_size,
-            self.transformer.num_channels_latents,
-            num_frames,
-            height // spatial_ratio,
-            width // spatial_ratio,
-        )
-
-        if isinstance(generator, list) and len(generator) != batch_size:
-            raise ValueError(
-                f"You have passed a list of generators of length {len(generator)}, "
-                f"but requested an effective batch size of {batch_size}.")
-
-        if latents is None:
-            if latent_path:
-                loaded_latents = self._load_initial_latent(
-                    latent_path, device, dtype)
-                if loaded_latents is not None:
-                    latents = loaded_latents
-                else:
-                    latents = randn_tensor(
-                        shape,
-                        generator=generator,
-                        device=device,
-                        dtype=dtype,
-                    )
-                    self._save_initial_latent(latent_path, latents)
-            else:
-                latents = randn_tensor(
-                    shape,
-                    generator=generator,
-                    device=device,
-                    dtype=dtype,
-                )
-        else:
-            latents = latents.to(device)
-
-        batch.latents = latents
-        batch.raw_latent_shape = shape
-        return batch
-
-    def _adjust_video_length(self, batch: ForwardBatch,
-                             fastvideo_args: FastVideoArgs) -> int | None:
-        if not fastvideo_args.pipeline_config.vae_config.use_temporal_scaling_frames:
-            return None
-        temporal_scale_factor = (fastvideo_args.pipeline_config.vae_config.
-                                 arch_config.temporal_compression_ratio)
-        video_length = batch.num_frames
-        return int((video_length - 1) // temporal_scale_factor + 1)
-
-    def _load_initial_latent(
-        self,
-        latent_path: str,
-        device: torch.device,
-        dtype: torch.dtype,
-    ) -> torch.Tensor | None:
-        path = Path(latent_path)
-        if not path.exists():
-            return None
-        payload = torch.load(path, map_location=device)
-        if isinstance(payload, dict):
-            if "video_latent" in payload:
-                latent = payload["video_latent"]
-            elif "latent" in payload:
-                latent = payload["latent"]
-            else:
-                latent = None
-        else:
-            latent = payload
-        if not torch.is_tensor(latent):
-            raise TypeError(f"Expected tensor for initial latent in {path}")
-        logger.info("[LTX2] Loaded initial latent from %s", path)
-        return latent.to(device=device, dtype=dtype)
-
-    def _save_initial_latent(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({"video_latent": latents.detach().cpu()}, path)
-        logger.info("[LTX2] Saved initial latent to %s", path)
-
-    def verify_input(self, batch: ForwardBatch,
-                     fastvideo_args: FastVideoArgs) -> VerificationResult:
-        result = VerificationResult()
-        result.add_check(
-            "prompt_or_embeds",
-            None,
-            lambda _: V.string_or_list_strings(batch.prompt) or not batch.
-            prompt_embeds or V.list_not_empty(batch.prompt_embeds),
-        )
-        if batch.prompt_embeds:
-            result.add_check("prompt_embeds", batch.prompt_embeds,
-                             V.list_of_tensors)
-        result.add_check("num_videos_per_prompt", batch.num_videos_per_prompt,
-                         V.positive_int)
-        result.add_check("generator", batch.generator,
-                         V.generator_or_list_generators)
-        result.add_check("num_frames", batch.num_frames, V.positive_int)
-        result.add_check("height", batch.height, V.positive_int)
-        result.add_check("width", batch.width, V.positive_int)
-        result.add_check("latents", batch.latents, V.none_or_tensor)
-        return result
-
-    def verify_output(self, batch: ForwardBatch,
-                      fastvideo_args: FastVideoArgs) -> VerificationResult:
-        result = VerificationResult()
-        result.add_check("latents", batch.latents,
-                         [V.is_tensor, V.with_dims(5)])
-        result.add_check("raw_latent_shape", batch.raw_latent_shape, V.is_tuple)
-        return result

fastvideo/pipelines/stages/text_encoding.py

@@ -11,11 +11,14 @@ from typing import Any
 from fastvideo.distributed import get_local_torch_device
 from fastvideo.fastvideo_args import FastVideoArgs
 from fastvideo.forward_context import set_forward_context
+from fastvideo.logger import init_logger
 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__)
+
 
 class TextEncodingStage(PipelineStage):
     """
@@ -36,7 +39,6 @@ class TextEncodingStage(PipelineStage):
         super().__init__()
         self.tokenizers = tokenizers
         self.text_encoders = text_encoders
-        self._last_audio_embeds: list[torch.Tensor] | None = None
 
     @torch.no_grad()
     def forward(
@@ -68,8 +70,6 @@ class TextEncodingStage(PipelineStage):
             encoder_index=all_indices,
             return_attention_mask=True,
         )
-        if self._last_audio_embeds is not None:
-            batch.extra["ltx2_audio_prompt_embeds"] = self._last_audio_embeds
 
         for pe in prompt_embeds_list:
             batch.prompt_embeds.append(pe)
@@ -86,9 +86,6 @@ class TextEncodingStage(PipelineStage):
                 encoder_index=all_indices,
                 return_attention_mask=True,
             )
-            if self._last_audio_embeds is not None:
-                batch.extra[
-                    "ltx2_audio_negative_embeds"] = self._last_audio_embeds
 
             assert batch.negative_prompt_embeds is not None
             for ne in neg_embeds_list:
@@ -187,13 +184,10 @@ class TextEncodingStage(PipelineStage):
 
         embeds_list: list[torch.Tensor] = []
         attn_masks_list: list[torch.Tensor] = []
-        audio_embeds_list: list[torch.Tensor] = []
 
         preprocess_funcs = fastvideo_args.pipeline_config.preprocess_text_funcs
         postprocess_funcs = fastvideo_args.pipeline_config.postprocess_text_funcs
         encoder_cfgs = fastvideo_args.pipeline_config.text_encoder_configs
-        is_ltx2 = getattr(fastvideo_args.pipeline_config.dit_config, "prefix",
-                          "") == "ltx2"
 
         if return_type not in ("list", "dict", "stack"):
             raise ValueError(
@@ -265,11 +259,6 @@ class TextEncodingStage(PipelineStage):
             except Exception:
                 prompt_embeds, attention_mask = postprocess_func(
                     outputs, attention_mask)
-            if is_ltx2 and getattr(outputs, "hidden_states", None):
-                audio_embed = outputs.hidden_states[0]
-                if dtype is not None:
-                    audio_embed = audio_embed.to(dtype=dtype)
-                audio_embeds_list.append(audio_embed)
 
             if dtype is not None:
                 prompt_embeds = prompt_embeds.to(dtype=dtype)
@@ -277,7 +266,6 @@ class TextEncodingStage(PipelineStage):
             if return_attention_mask:
                 attn_masks_list.append(attention_mask)
 
-        self._last_audio_embeds = audio_embeds_list if is_ltx2 else None
         return self.return_embeds(embeds_list, attn_masks_list, return_type,
                                   return_attention_mask, indices)
 

fastvideo/tests/hooks/test_hooks_api.py

@@ -1,60 +0,0 @@
-from fastvideo.hooks.hooks import ForwardHook, ModuleHookManager
-from torch import nn
-from typing import Any
-import torch
-
-
-class EventHook(ForwardHook):
-    def __init__(self, content: str, event_list: list[str]):
-        self.content = content
-        self.event_list = event_list
-
-    def name(self) -> str:
-        return f"EventHook_{self.content}"
-
-    def pre_forward(self, module: nn.Module, *args, **kwargs):
-        print(
-            f"[{self.content}] Pre-forward called with args[0].shape: {args[0].shape}"
-        )
-        self.event_list.append(f"[pre]{self.content}")
-        return args, kwargs
-
-    def post_forward(self, module: nn.Module, output: Any):
-        print(
-            f"[{self.content}] Post-forward called with outputs.shape: {output.shape}"
-        )
-        self.event_list.append(f"[post]{self.content}")
-        return output
-
-
-def test_hook_execution_order():
-    """Test that hooks are executed in the correct order: LIFO for pre-hooks, FIFO for post-hooks."""
-    # Create a simple model
-    model = nn.Linear(10, 20)
-
-    # Create event list to track hook execution order
-    events = []
-
-    # Create and push hooks in order: A then B
-
-    manager = ModuleHookManager.get_from_or_default(model)
-
-    hook_a = EventHook("A", events)
-    hook_b = EventHook("B", events)
-
-    manager.append_forward_hook(hook_a)
-    manager.append_forward_hook(hook_b)
-
-    # Perform a forward pass
-    input_tensor = torch.randn(2, 10)
-    model(input_tensor)
-
-    # Verify the execution order is [pre_a, pre_b, post_b, post_a]
-    # Pre-hooks should be FILO (First In Last Out): A then B
-    # Post-hooks should be LIFO (Last In First Out): B then A
-    expected_events = ["[pre]A", "[pre]B", "[post]B", "[post]A"]
-
-    assert events == expected_events, (
-        f"Expected {expected_events}, but got {events}"
-    )
-    print(f"✓ Hook execution order test passed: {events}")

fastvideo/tests/hooks/test_layerwise_offload.py

@@ -1,395 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-
-import pytest
-import torch
-import torch.nn as nn
-
-from fastvideo.hooks.layerwise_offload import (
-    LayerwiseOffloadHook,
-    enable_layerwise_offload,
-)
-from fastvideo.hooks.hooks import ModuleHookManager
-
-
-class SimpleBlock(nn.Module):
-    """A simple block with linear layers for testing."""
-
-    def __init__(self, hidden_size: int, dtype: torch.dtype = torch.float32):
-        super().__init__()
-        self.linear1 = nn.Linear(hidden_size, hidden_size, dtype=dtype)
-        self.linear2 = nn.Linear(hidden_size, hidden_size, dtype=dtype)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        x = self.linear1(x)
-        x = torch.relu(x)
-        x = self.linear2(x)
-        return x
-
-
-class SimpleModelWithModuleList(nn.Module):
-    """A simple model with ModuleList for testing layerwise offloading."""
-
-    def __init__(
-        self,
-        num_blocks: int = 4,
-        hidden_size: int = 128,
-        dtype: torch.dtype = torch.float32,
-    ):
-        super().__init__()
-        self.blocks = nn.ModuleList(
-            [SimpleBlock(hidden_size, dtype=dtype) for _ in range(num_blocks)]
-        )
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        for block in self.blocks:
-            x = block(x)
-        return x
-
-
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
-def test_layerwise_offload_basic():
-    """Test basic functionality of layerwise offloading."""
-    device = torch.device("cuda")
-    hidden_size = 128
-    batch_size = 2
-    seq_len = 16
-    num_blocks = 4
-
-    # Create model
-    model = SimpleModelWithModuleList(
-        num_blocks=num_blocks, hidden_size=hidden_size, dtype=torch.float32
-    ).to(device)
-
-    # Get reference output without offloading
-    input_tensor = torch.randn(batch_size, seq_len, hidden_size, device=device)
-    with torch.no_grad():
-        reference_output = model(input_tensor.clone())
-
-    # Enable layerwise offloading
-    enable_layerwise_offload(model)
-
-    # Verify parameters are offloaded to CPU
-    for block in model.blocks:
-        for param in block.parameters():
-            # Parameters should be placeholder tensors (empty)
-            assert param.numel() == 0, (
-                "Parameters should be offloaded (empty tensors)"
-            )
-
-    # Run forward pass with offloading
-    with torch.no_grad():
-        offloaded_output = model(input_tensor.clone())
-
-    # Check output correctness
-    assert torch.allclose(
-        reference_output, offloaded_output, rtol=1e-4, atol=1e-5
-    ), "Output with offloading should match reference output"
-
-    print(
-        f" Layerwise offload basic test passed: max diff = {(reference_output - offloaded_output).abs().max().item()}"
-    )
-
-
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
-def test_layerwise_offload_bf16():
-    """Test layerwise offloading with bfloat16 precision."""
-    device = torch.device("cuda")
-    hidden_size = 256
-    batch_size = 1
-    seq_len = 32
-    num_blocks = 3
-
-    # Create model
-    model = SimpleModelWithModuleList(
-        num_blocks=num_blocks, hidden_size=hidden_size, dtype=torch.bfloat16
-    ).to(device)
-
-    # Get reference output without offloading
-    input_tensor = torch.randn(
-        batch_size, seq_len, hidden_size, device=device, dtype=torch.bfloat16
-    )
-    with torch.no_grad():
-        reference_output = model(input_tensor.clone())
-
-    # Enable layerwise offloading
-    enable_layerwise_offload(model)
-
-    # Run forward pass with offloading
-    with torch.no_grad():
-        offloaded_output = model(input_tensor.clone())
-
-    # Check output correctness (looser tolerance for bf16)
-    assert torch.allclose(
-        reference_output, offloaded_output, rtol=1e-2, atol=1e-3
-    ), "Output with offloading should match reference output for bf16"
-
-    print(
-        f" Layerwise offload bf16 test passed: max diff = {(reference_output - offloaded_output).abs().max().item()}"
-    )
-
-
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
-def test_layerwise_offload_multiple_forward_passes():
-    """Test that layerwise offloading works correctly across multiple forward passes."""
-    device = torch.device("cuda")
-    hidden_size = 64
-    batch_size = 2
-    seq_len = 8
-    num_blocks = 3
-    num_iterations = 5
-
-    # Create model
-    model = SimpleModelWithModuleList(
-        num_blocks=num_blocks, hidden_size=hidden_size, dtype=torch.float32
-    ).to(device)
-
-    # Get reference outputs without offloading
-    torch.manual_seed(42)
-    reference_outputs = []
-    for i in range(num_iterations):
-        input_tensor = torch.randn(
-            batch_size, seq_len, hidden_size, device=device
-        )
-        with torch.no_grad():
-            reference_outputs.append(model(input_tensor.clone()))
-
-    # Enable layerwise offloading
-    enable_layerwise_offload(model)
-
-    # Run multiple forward passes with offloading
-    torch.manual_seed(42)
-    for i in range(num_iterations):
-        input_tensor = torch.randn(
-            batch_size, seq_len, hidden_size, device=device
-        )
-        with torch.no_grad():
-            offloaded_output = model(input_tensor.clone())
-
-        # Check output correctness for each iteration
-        assert torch.allclose(
-            reference_outputs[i], offloaded_output, rtol=1e-4, atol=1e-5
-        ), f"Output mismatch in iteration {i}"
-
-    print(
-        f" Multiple forward passes test passed for {num_iterations} iterations"
-    )
-
-
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
-def test_layerwise_offload_parameter_integrity():
-    """Test that parameters are correctly restored during forward pass."""
-    device = torch.device("cuda")
-    hidden_size = 128
-    num_blocks = 2
-
-    # Create model
-    model = SimpleModelWithModuleList(
-        num_blocks=num_blocks, hidden_size=hidden_size, dtype=torch.float32
-    ).to(device)
-
-    # Store original parameter values per block
-    original_params_per_block = []
-    for block in model.blocks:
-        block_params = {}
-        for name, param in block.named_parameters():
-            block_params[name] = param.data.clone()
-        original_params_per_block.append(block_params)
-
-    # Enable layerwise offloading
-    enable_layerwise_offload(model)
-
-    # Create hook managers and verify they exist
-    for block_idx, block in enumerate(model.blocks):
-        manager = ModuleHookManager.get_from(block)
-        assert manager is not None, "Hook manager should be attached to blocks"
-        hook: LayerwiseOffloadHook | None = manager.get_forward_hook(
-            "LayerwiseOffloadHook"
-        )
-        assert hook is not None, "LayerwiseOffloadHook should be registered"
-
-        # Verify parameters are stored in CPU
-        state = hook.state
-        assert len(state.cpu_named_parameters) > 0, (
-            "CPU parameters should be stored"
-        )
-
-        # Verify CPU parameters match original values for this block
-        original_params = original_params_per_block[block_idx]
-        for name, cpu_param in state.cpu_named_parameters.items():
-            assert name in original_params, (
-                f"CPU parameter {name} not found in original params for block {block_idx}"
-            )
-            assert torch.allclose(
-                cpu_param.cpu(), original_params[name].cpu(), rtol=1e-5, atol=1e-6
-            ), f"CPU parameter {name} should match original for block {block_idx}"
-
-    print(" Parameter integrity test passed")
-
-
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
-def test_layerwise_offload_no_modulelist_error():
-    """Test that enabling offload on a model without ModuleList raises an error."""
-
-    class ModelWithoutModuleList(nn.Module):
-        def __init__(self):
-            super().__init__()
-            self.linear = nn.Linear(128, 128)
-
-        def forward(self, x):
-            return self.linear(x)
-
-    model = ModelWithoutModuleList().to("cuda")
-
-    with pytest.raises(
-        ValueError,
-        match="No nn.ModuleList found in the model for layerwise offloading",
-    ):
-        enable_layerwise_offload(model)
-
-    print(" No ModuleList error test passed")
-
-
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
-def test_layerwise_offload_memory_reduction():
-    """Test that layerwise offloading reduces GPU memory usage."""
-    device = torch.device("cuda")
-    hidden_size = 512
-    num_blocks = 8
-    batch_size = 1
-    seq_len = 64
-
-    # Create model
-    model = SimpleModelWithModuleList(
-        num_blocks=num_blocks, hidden_size=hidden_size, dtype=torch.float32
-    ).to(device)
-
-    # Measure initial GPU memory
-    torch.cuda.empty_cache()
-    torch.cuda.reset_peak_memory_stats()
-
-    input_tensor = torch.randn(batch_size, seq_len, hidden_size, device=device)
-    with torch.no_grad():
-        _ = model(input_tensor)
-
-    memory_without_offload = torch.cuda.max_memory_allocated()
-
-    # Reset model
-    model = SimpleModelWithModuleList(
-        num_blocks=num_blocks, hidden_size=hidden_size, dtype=torch.float32
-    ).to(device)
-
-    # Enable offloading
-    enable_layerwise_offload(model)
-
-    # Measure GPU memory with offloading
-    torch.cuda.empty_cache()
-    torch.cuda.reset_peak_memory_stats()
-
-    input_tensor = torch.randn(batch_size, seq_len, hidden_size, device=device)
-    with torch.no_grad():
-        _ = model(input_tensor)
-
-    memory_with_offload = torch.cuda.max_memory_allocated()
-
-    # Memory with offload should be less (parameters are offloaded)
-    # Note: This is a weak check as memory usage depends on many factors
-    print(
-        f"Memory without offload: {memory_without_offload / 1024**2:.2f} MB, "
-        f"with offload: {memory_with_offload / 1024**2:.2f} MB"
-    )
-
-    # We expect some reduction but this test is more informational
-    assert memory_with_offload < memory_without_offload * 1.5, (
-        "Memory usage should not increase significantly with offloading"
-    )
-
-    print(" Memory reduction test passed")
-
-
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
-def test_layerwise_offload_gradient_disabled():
-    """Test that layerwise offloading works correctly with gradients disabled."""
-    device = torch.device("cuda")
-    hidden_size = 128
-    batch_size = 2
-    seq_len = 16
-    num_blocks = 3
-
-    # Create model
-    model = SimpleModelWithModuleList(
-        num_blocks=num_blocks, hidden_size=hidden_size, dtype=torch.float32
-    ).to(device)
-    model.eval()
-
-    # Enable layerwise offloading
-    enable_layerwise_offload(model)
-
-    input_tensor = torch.randn(batch_size, seq_len, hidden_size, device=device)
-
-    # Forward pass should work without gradients
-    with torch.no_grad():
-        output = model(input_tensor)
-
-    assert output.requires_grad is False, "Output should not require gradients"
-    assert output.shape == (
-        batch_size,
-        seq_len,
-        hidden_size,
-    ), "Output shape should match input"
-
-    print(" Gradient disabled test passed")
-
-
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
-def test_layerwise_offload_different_batch_sizes():
-    """Test layerwise offloading with different batch sizes."""
-    device = torch.device("cuda")
-    hidden_size = 128
-    seq_len = 16
-    num_blocks = 3
-
-    model = SimpleModelWithModuleList(
-        num_blocks=num_blocks, hidden_size=hidden_size, dtype=torch.float32
-    ).to(device)
-
-    # Get reference model without offloading
-    reference_model = SimpleModelWithModuleList(
-        num_blocks=num_blocks, hidden_size=hidden_size, dtype=torch.float32
-    ).to(device)
-    reference_model.load_state_dict(model.state_dict())
-
-    # Enable layerwise offloading
-    enable_layerwise_offload(model)
-
-    # Test different batch sizes
-    for batch_size in [1, 2, 4, 8]:
-        input_tensor = torch.randn(
-            batch_size, seq_len, hidden_size, device=device
-        )
-
-        with torch.no_grad():
-            reference_output = reference_model(input_tensor.clone())
-            offloaded_output = model(input_tensor.clone())
-
-        assert torch.allclose(
-            reference_output, offloaded_output, rtol=1e-4, atol=1e-5
-        ), f"Output mismatch for batch_size={batch_size}"
-
-    print(" Different batch sizes test passed")
-
-
-if __name__ == "__main__":
-    # Run tests manually for debugging
-    if torch.cuda.is_available():
-        print("Running layerwise offloading tests...")
-        test_layerwise_offload_basic()
-        test_layerwise_offload_bf16()
-        test_layerwise_offload_multiple_forward_passes()
-        test_layerwise_offload_parameter_integrity()
-        test_layerwise_offload_no_modulelist_error()
-        test_layerwise_offload_memory_reduction()
-        test_layerwise_offload_gradient_disabled()
-        test_layerwise_offload_different_batch_sizes()
-        print("\n All tests passed!")
-    else:
-        print("CUDA not available, skipping tests")

fastvideo/tests/inference/lora/test_lora_inference_similarity.py

@@ -93,11 +93,8 @@ def test_merge_lora_weights(model_id):
 
     lora_nickname = lora_config["lora_nickname"]
     lora_path = lora_config["lora_path"]
-    # When layerwise offload is enabled, placeholder tensors cannot be compared directly.
     args = FastVideoArgs.from_kwargs(
         model_path=model_id,
-        dit_layerwise_offload=False,
-        use_fsdp_inference=True,
         dit_cpu_offload=True,
         dit_precision="bf16",
     )

fastvideo/tests/modal/pr_test.py

@@ -67,23 +67,17 @@ def run_test(pytest_command: str):
     
     sys.exit(result.returncode)
 
-@app.function(gpu="H100:1", 
-              image=image, 
-              timeout=1200, 
-              secrets=[modal.Secret.from_dict({"HF_API_KEY": os.environ.get("HF_API_KEY", "")})],
-              volumes={"/root/data": model_vol})
+@app.function(gpu="H100:1", image=image, timeout=1200, secrets=[modal.Secret.from_dict({"HF_API_KEY": os.environ.get("HF_API_KEY", "")})])
 def run_encoder_tests():
-    run_test("export HF_HOME='/root/data/.cache' && hf auth login --token $HF_API_KEY && pytest ./fastvideo/tests/encoders -vs")
+    run_test("hf auth login --token $HF_API_KEY && pytest ./fastvideo/tests/encoders -vs")
 
-@app.function(gpu="L40S:1", image=image, timeout=1200, secrets=[modal.Secret.from_dict({"HF_API_KEY": os.environ.get("HF_API_KEY", "")})],
-              volumes={"/root/data": model_vol})
+@app.function(gpu="L40S:1", image=image, timeout=1200, secrets=[modal.Secret.from_dict({"HF_API_KEY": os.environ.get("HF_API_KEY", "")})])
 def run_vae_tests():
-    run_test("export HF_HOME='/root/data/.cache' && hf auth login --token $HF_API_KEY && pytest ./fastvideo/tests/vaes -vs")
+    run_test("hf auth login --token $HF_API_KEY && pytest ./fastvideo/tests/vaes -vs")
 
-@app.function(gpu="L40S:1", image=image, timeout=900, secrets=[modal.Secret.from_dict({"HF_API_KEY": os.environ.get("HF_API_KEY", "")})],
-              volumes={"/root/data": model_vol})
+@app.function(gpu="L40S:1", image=image, timeout=900, secrets=[modal.Secret.from_dict({"HF_API_KEY": os.environ.get("HF_API_KEY", "")})])
 def run_transformer_tests():
-    run_test("export HF_HOME='/root/data/.cache' && hf auth login --token $HF_API_KEY && pytest ./fastvideo/tests/transformers -vs")
+    run_test("hf auth login --token $HF_API_KEY && pytest ./fastvideo/tests/transformers -vs")
 
 @app.function(
     gpu="L40S:4", 
@@ -95,21 +89,13 @@ def run_transformer_tests():
 def run_ssim_tests():
     run_test("export HF_HOME='/root/data/.cache' && export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True && hf auth login --token $HF_API_KEY && pytest ./fastvideo/tests/ssim -vs")
 
-@app.function(gpu="L40S:4", 
-    image=image, 
-    timeout=900, 
-    secrets=[modal.Secret.from_dict({"WANDB_API_KEY": os.environ.get("WANDB_API_KEY", "")})],
-    volumes={"/root/data": model_vol})
+@app.function(gpu="L40S:4", image=image, timeout=900, secrets=[modal.Secret.from_dict({"WANDB_API_KEY": os.environ.get("WANDB_API_KEY", "")})])
 def run_training_tests():
-    run_test("export HF_HOME='/root/data/.cache' && wandb login $WANDB_API_KEY && pytest ./fastvideo/tests/training/Vanilla -srP")
+    run_test("wandb login $WANDB_API_KEY && pytest ./fastvideo/tests/training/Vanilla -srP")
 
-@app.function(gpu="L40S:2", 
-    image=image, 
-    timeout=900, 
-    secrets=[modal.Secret.from_dict({"WANDB_API_KEY": os.environ.get("WANDB_API_KEY", "")})],
-    volumes={"/root/data": model_vol})
+@app.function(gpu="L40S:2", image=image, timeout=900, secrets=[modal.Secret.from_dict({"WANDB_API_KEY": os.environ.get("WANDB_API_KEY", "")})])
 def run_training_lora_tests():
-    run_test("export HF_HOME='/root/data/.cache' && wandb login $WANDB_API_KEY && pytest ./fastvideo/tests/training/lora/test_lora_training.py -srP")
+    run_test("wandb login $WANDB_API_KEY && pytest ./fastvideo/tests/training/lora/test_lora_training.py -srP")
 
 @app.function(gpu="H100:2", image=image, timeout=900, secrets=[modal.Secret.from_dict({"WANDB_API_KEY": os.environ.get("WANDB_API_KEY", "")})])
 def run_training_tests_VSA():

fastvideo/tests/nightly/test_e2e_hunyuan15_overfit_single_sample.py

@@ -0,0 +1,186 @@
+import os
+from pathlib import Path
+import shutil
+import subprocess
+import sys
+
+from huggingface_hub import snapshot_download
+from fastvideo.tests.ssim.test_inference_similarity import compute_video_ssim_torchvision
+
+sys.path.append(str(Path(__file__).parent.parent.parent.parent.parent))
+
+NUM_NODES = "1"
+MODEL_PATH = "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v"
+
+DATA_DIR = "data"
+LOCAL_RAW_DATA_DIR = Path(DATA_DIR) / "cats15"
+LOCAL_PREPROCESSED_DATA_DIR = Path(DATA_DIR) / "cats_processed_t2v_hunyuan15"
+LOCAL_OUTPUT_DIR = Path(DATA_DIR) / "outputs_hunyuan15"
+
+NUM_GPUS_PER_NODE_PREPROCESSING = "1"
+NUM_GPUS_PER_NODE_TRAINING = "8"
+
+# entrypoints (adjust to what hunyuan15 scripts use)
+PREPROCESS_ENTRY = ["-m", "fastvideo.pipelines.preprocess.v1_preprocessing_new"]
+TRAIN_ENTRY_FILE_PATH = "fastvideo/training/hunyuan15_training_pipeline.py"  # change if hunyuan15 uses a different file
+
+LOCAL_TRAINING_DATA_DIR = os.path.join(LOCAL_PREPROCESSED_DATA_DIR, "training_dataset", "worker_0", "worker_0")
+LOCAL_VALIDATION_DATASET_FILE = os.path.join(LOCAL_RAW_DATA_DIR, "validation_prompt.json")
+
+
+def download_data():
+    os.makedirs(DATA_DIR, exist_ok=True)
+    snapshot_download(
+        repo_id="wlsaidhi/cats-overfit-merged",
+        local_dir=str(LOCAL_RAW_DATA_DIR),
+        repo_type="dataset",
+        resume_download=True,
+        token=os.environ.get("HF_TOKEN"),
+    )
+
+    # normalize dataset layout like your hunyuan test
+    video_dir = LOCAL_RAW_DATA_DIR / "video"
+    videos_dir = LOCAL_RAW_DATA_DIR / "videos"
+    if video_dir.exists() and not videos_dir.exists():
+        video_dir.rename(videos_dir)
+
+    src_val = LOCAL_RAW_DATA_DIR / "validation_prompt_1_sample.json"
+    shutil.copy2(src_val, LOCAL_VALIDATION_DATASET_FILE)
+
+    src_v2c = LOCAL_RAW_DATA_DIR / "videos2caption_1_sample.json"
+    shutil.copy2(src_v2c, LOCAL_RAW_DATA_DIR / "videos2caption.json")
+
+
+def run_preprocessing():
+    if LOCAL_PREPROCESSED_DATA_DIR.exists():
+        shutil.rmtree(LOCAL_PREPROCESSED_DATA_DIR)
+
+    env = os.environ.copy()
+    env["PYTHONPATH"] = os.getcwd()
+
+    cmd = [
+        "torchrun",
+        "--nnodes", NUM_NODES,
+        "--nproc_per_node", NUM_GPUS_PER_NODE_PREPROCESSING,
+        *PREPROCESS_ENTRY,
+        "--model-path", MODEL_PATH,
+        "--mode", "preprocess",
+        "--workload-type", "t2v",
+        "--preprocess.video_loader_type", "torchvision",
+        "--preprocess.dataset_type", "merged",
+        "--preprocess.dataset_path", str(LOCAL_RAW_DATA_DIR),
+        "--preprocess.dataset_output_dir", str(LOCAL_PREPROCESSED_DATA_DIR),
+        "--preprocess.preprocess_video_batch_size", "1",
+        "--preprocess.dataloader_num_workers", "0",
+        "--preprocess.max_height", "480",
+        "--preprocess.max_width", "832",
+        "--preprocess.num_frames", "61",
+        "--preprocess.train_fps", "24",
+        "--preprocess.samples_per_file", "1",
+        "--preprocess.flush_frequency", "1",
+        "--preprocess.video_length_tolerance_range", "5",
+
+        # IMPORTANT: add a dtype/precision flag here if available in your codebase,
+        # so saved arrays are fp16/fp32 instead of bf16.
+        # Example (ONLY if supported by args):
+        # "--preprocess.save_dtype", "fp16",
+        # "--preprocess.embedding_dtype", "fp16",
+    ]
+
+    subprocess.run(cmd, check=True, env=env)
+
+
+def run_training():
+    env = os.environ.copy()
+    env["PYTHONPATH"] = os.getcwd()
+
+    cmd = [
+        "torchrun",
+        "--nnodes", NUM_NODES,
+        "--nproc_per_node", NUM_GPUS_PER_NODE_TRAINING,
+        TRAIN_ENTRY_FILE_PATH,
+        "--model_path", MODEL_PATH,
+        "--inference_mode", "False",  # Required: must be False for training mode
+        "--pretrained_model_name_or_path", MODEL_PATH,
+        "--data_path", LOCAL_TRAINING_DATA_DIR,
+        "--validation_dataset_file", LOCAL_VALIDATION_DATASET_FILE,
+
+        "--train_batch_size", "1",
+        "--num_latent_t", "16",  # (61-1)//4 + 1 = 16 for 61 frames with temporal_compression=4
+        "--max_train_steps", "901",
+        "--learning_rate", "1e-5",
+        "--gradient_accumulation_steps", "1",  # Required: default 0 causes empty training loop
+
+        # If your pipeline produces bf16 artifacts that later go to numpy/parquet,
+        # switch mixed_precision to fp16 to avoid bf16 -> numpy issues.
+        "--mixed_precision", "fp16",
+
+        # Distributed training configuration - must match nproc_per_node
+        "--num_gpus", NUM_GPUS_PER_NODE_TRAINING,
+        "--sp_size", NUM_GPUS_PER_NODE_TRAINING,
+        "--tp_size", "1",
+        "--hsdp_replicate_dim", "1",
+        "--hsdp_shard_dim", NUM_GPUS_PER_NODE_TRAINING,
+        "--train_sp_batch_size", "1",
+
+        # Dataloader
+        "--dataloader_num_workers", "4",
+
+        # Validation & logging
+        "--log_validation",
+        "--validation_steps", "100",
+        "--validation_sampling_steps", "50",
+        "--validation_guidance_scale", "6.0",  # Match Hunyuan15_480P_SamplingParam default
+
+        # Checkpointing
+        "--weight_only_checkpointing_steps", "6000",
+        "--training_state_checkpointing_steps", "6000",
+        "--checkpoints_total_limit", "3",
+
+        # Optimizer settings
+        "--weight_decay", "0.01",
+        "--max_grad_norm", "1.0",
+
+        # Training config
+        "--ema_start_step", "0",
+        "--training_cfg_rate", "0.0",
+        "--dit_precision", "fp32",
+        "--enable_gradient_checkpointing_type", "full",
+
+        # Output
+        "--output_dir", str(LOCAL_OUTPUT_DIR),
+        "--tracker_project_name", "hunyuan15_finetune_overfit_ci",
+        "--num_height", "480",
+        "--num_width", "832",
+        "--num_frames", "61",
+        "--embedded_cfg_scale", "6.0",
+        "--num_euler_timesteps", "50",
+    ]
+
+    subprocess.run(cmd, check=True, env=env)
+
+
+def test_e2e_hunyuan15_overfit_single_sample():
+    os.environ["WANDB_MODE"] = "online"
+    os.environ["WANDB_API_KEY"] = "8d9f4b39abd68eb4e29f6fc010b7ee71a2207cde"
+
+    # download_data()
+    # run_preprocessing()
+    run_training()
+
+    reference_video_file = os.path.join(os.path.dirname(__file__), "reference_video_1_sample_hy15.mp4")
+    final_validation_video_file = os.path.join(LOCAL_OUTPUT_DIR, "validation_step_900_inference_steps_50_video_0.mp4")
+
+    assert os.path.exists(reference_video_file)
+    assert os.path.exists(final_validation_video_file)
+
+    mean_ssim, min_ssim, max_ssim = compute_video_ssim_torchvision(
+        reference_video_file,
+        final_validation_video_file,
+        use_ms_ssim=True,
+    )
+    assert max_ssim > 0.5, f"Max SSIM is below threshold: {max_ssim}"
+
+
+if __name__ == "__main__":
+    test_e2e_hunyuan15_overfit_single_sample()

fastvideo/tests/nightly/test_e2e_hunyuan_overfit_single_sample.py

@@ -0,0 +1,218 @@
+import os
+from pathlib import Path
+from huggingface_hub import snapshot_download
+import shutil
+import subprocess
+import sys
+from fastvideo.tests.ssim.test_inference_similarity import compute_video_ssim_torchvision
+
+# Example of using Hunyuan preprocessing and training pipeline
+
+# Import the training pipeline
+sys.path.append(str(Path(__file__).parent.parent.parent.parent.parent))
+
+NUM_NODES = "1"
+MODEL_PATH = "hunyuanvideo-community/HunyuanVideo"
+
+# preprocessing
+DATA_DIR = "data"
+LOCAL_RAW_DATA_DIR = Path(os.path.join(DATA_DIR, "cats"))
+NUM_GPUS_PER_NODE_PREPROCESSING = "1"
+PREPROCESSING_ENTRY_FILE_PATH = "fastvideo/pipelines/preprocess/v1_preprocessing_new.py"
+
+LOCAL_PREPROCESSED_DATA_DIR = Path(os.path.join(DATA_DIR, "cats_processed_t2v_hunyuan"))
+
+
+# training
+NUM_GPUS_PER_NODE_TRAINING = "4"
+TRAINING_ENTRY_FILE_PATH = "fastvideo/training/hunyuan_training_pipeline.py"
+# New preprocessing pipeline creates files in training_dataset/worker_0/worker_0/
+LOCAL_TRAINING_DATA_DIR = os.path.join(LOCAL_PREPROCESSED_DATA_DIR, "training_dataset", "worker_0", "worker_0")
+LOCAL_VALIDATION_DATASET_FILE = os.path.join(LOCAL_RAW_DATA_DIR, "validation_prompt.json")
+LOCAL_OUTPUT_DIR = Path(os.path.join(DATA_DIR, "outputs_hunyuan"))
+
+def download_data():
+    # create the data dir if it doesn't exist
+    data_dir = Path(DATA_DIR)
+
+    print(f"Creating data directory at {data_dir}")
+    os.makedirs(data_dir, exist_ok=True)
+
+    print(f"Downloading raw dataset to {LOCAL_RAW_DATA_DIR}...")
+    try:
+        result = snapshot_download(
+            repo_id="wlsaidhi/cats-overfit-merged",
+            local_dir=str(LOCAL_RAW_DATA_DIR),
+            repo_type="dataset",
+            resume_download=True,
+            token=os.environ.get("HF_TOKEN"),  # In case authentication is needed
+        )
+        print(f"Download completed successfully. Files downloaded to: {result}")
+        
+        # Verify the download
+        if not LOCAL_RAW_DATA_DIR.exists():
+            raise RuntimeError(f"Download appeared to succeed but {LOCAL_RAW_DATA_DIR} does not exist")
+            
+        # List downloaded files
+        print("Downloaded files:")
+        for file in LOCAL_RAW_DATA_DIR.rglob("*"):
+            if file.is_file():
+                print(f"  - {file.relative_to(LOCAL_RAW_DATA_DIR)}")
+        
+        # Rename video directory if needed (dataset has 'video' but preprocessing expects 'videos')
+        video_dir = os.path.join(LOCAL_RAW_DATA_DIR, "video")
+        videos_dir = os.path.join(LOCAL_RAW_DATA_DIR, "videos")
+        if os.path.exists(video_dir) and not os.path.exists(videos_dir):
+            print(f"Renaming video directory to videos...")
+            os.rename(video_dir, videos_dir)
+            
+        # Copy validation file to expected name
+        source_validation_file = os.path.join(LOCAL_RAW_DATA_DIR, "validation_prompt_1_sample.json")
+        target_validation_file = LOCAL_VALIDATION_DATASET_FILE
+        
+        if os.path.exists(source_validation_file):
+            print(f"Copying {source_validation_file} to {target_validation_file}...")
+            shutil.copy2(source_validation_file, target_validation_file)
+        else:
+            raise FileNotFoundError(f"Source validation file not found: {source_validation_file}")
+        
+        # Override videos2caption.json with the 1-sample version for preprocessing
+        # The new preprocessing pipeline automatically reads videos2caption.json from the dataset path
+        source_videos2caption = os.path.join(LOCAL_RAW_DATA_DIR, "videos2caption_1_sample.json")
+        target_videos2caption = os.path.join(LOCAL_RAW_DATA_DIR, "videos2caption.json")
+        
+        if os.path.exists(source_videos2caption):
+            print(f"Overriding videos2caption.json with 1-sample version...")
+            shutil.copy2(source_videos2caption, target_videos2caption)
+        else:
+            raise FileNotFoundError(f"Source videos2caption file not found: {source_videos2caption}")
+        
+    except Exception as e:
+        print(f"Error during download: {str(e)}")
+        raise
+
+
+def run_preprocessing():
+    # remove the local_preprocessed_data_dir if it exists
+    if LOCAL_PREPROCESSED_DATA_DIR.exists():
+        print(f"Removing local_preprocessed_data_dir: {LOCAL_PREPROCESSED_DATA_DIR}")
+        shutil.rmtree(LOCAL_PREPROCESSED_DATA_DIR)
+    env = os.environ.copy()
+    env['PYTHONPATH'] = os.getcwd()
+    # Run torchrun command using the new preprocessing pipeline
+    cmd = [
+        "torchrun",
+        "--nnodes", NUM_NODES,
+        "--nproc_per_node", NUM_GPUS_PER_NODE_PREPROCESSING,
+        "-m", "fastvideo.pipelines.preprocess.v1_preprocessing_new",
+        "--model-path", MODEL_PATH,
+        "--mode", "preprocess",
+        "--workload-type", "t2v",
+        "--preprocess.video_loader_type", "torchvision",
+        "--preprocess.dataset_type", "merged",
+        "--preprocess.dataset_path", str(LOCAL_RAW_DATA_DIR),
+        "--preprocess.dataset_output_dir", str(LOCAL_PREPROCESSED_DATA_DIR),
+        "--preprocess.preprocess_video_batch_size", "1",
+        "--preprocess.dataloader_num_workers", "0",
+        "--preprocess.max_height", "480",
+        "--preprocess.max_width", "832",
+        "--preprocess.num_frames", "77",
+        "--preprocess.train_fps", "16",
+        "--preprocess.samples_per_file", "1",
+        "--preprocess.flush_frequency", "1",
+        "--preprocess.video_length_tolerance_range", "5",
+    ]
+
+    process = subprocess.run(cmd, check=True, env=env)
+
+
+def run_training():
+    env = os.environ.copy()
+    env['PYTHONPATH'] = os.getcwd()
+
+    cmd = [
+        "torchrun",
+        "--nnodes", NUM_NODES,
+        "--nproc_per_node", NUM_GPUS_PER_NODE_TRAINING,
+        TRAINING_ENTRY_FILE_PATH,
+        "--model_path", MODEL_PATH,
+        "--inference_mode", "False",
+        "--pretrained_model_name_or_path", MODEL_PATH,
+        "--data_path", LOCAL_TRAINING_DATA_DIR,
+        "--validation_dataset_file", LOCAL_VALIDATION_DATASET_FILE,
+        "--train_batch_size", "1",
+        "--num_latent_t", "8",
+        "--num_gpus", NUM_GPUS_PER_NODE_TRAINING,
+        "--sp_size", NUM_GPUS_PER_NODE_TRAINING,
+        "--tp_size", "1",
+        "--hsdp_replicate_dim", "1",
+        "--hsdp_shard_dim", NUM_GPUS_PER_NODE_TRAINING,
+        "--train_sp_batch_size", "1",
+        "--dataloader_num_workers", "10",
+        "--gradient_accumulation_steps", "1",
+        "--max_train_steps", "901",
+        "--learning_rate", "1e-5",
+        "--mixed_precision", "bf16",
+        "--weight_only_checkpointing_steps", "6000",
+        "--training_state_checkpointing_steps", "6000",
+        "--validation_steps", "100",
+        "--validation_sampling_steps", "50",
+        "--log_validation",
+        "--checkpoints_total_limit", "3",
+        "--ema_start_step", "0",
+        "--training_cfg_rate", "0.0",
+        "--output_dir", str(LOCAL_OUTPUT_DIR),
+        "--tracker_project_name", "hunyuan_finetune_overfit_ci",
+        "--num_height", "480",
+        "--num_width", "832",
+        "--num_frames", "81",
+        "--validation_guidance_scale", "1.0",
+        "--embedded_cfg_scale", "6.0",
+        "--num_euler_timesteps", "50",
+        "--multi_phased_distill_schedule", "4000-1",
+        "--weight_decay", "0.01",
+        "--not_apply_cfg_solver",
+        "--dit_precision", "fp32",
+        "--max_grad_norm", "1.0",
+        "--flow_shift", "7",
+    ]
+
+    print(f"Running training with command: {cmd}")
+    process = subprocess.run(cmd, check=True, env=env)
+
+
+def test_e2e_hunyuan_overfit_single_sample():
+    os.environ["WANDB_MODE"] = "online"
+
+    #download_data()
+    run_preprocessing()
+    run_training()
+
+    reference_video_file = os.path.join(os.path.dirname(__file__), "reference_video_1_sample_v0.mp4")
+    print(f"reference_video_file: {reference_video_file}")
+    final_validation_video_file = os.path.join(LOCAL_OUTPUT_DIR, "validation_step_900_inference_steps_50_video_0.mp4")
+    print(f"final_validation_video_file: {final_validation_video_file}")
+
+
+    # Ensure both files exist
+    assert os.path.exists(reference_video_file), f"Reference video not found at {reference_video_file}"
+    assert os.path.exists(final_validation_video_file), f"Validation video not found at {final_validation_video_file}"
+
+    #Compute SSIM
+    mean_ssim, min_ssim, max_ssim = compute_video_ssim_torchvision(
+        reference_video_file,
+        final_validation_video_file,
+        use_ms_ssim=True  # Using MS-SSIM for better quality assessment
+    )
+
+    print("\n===== SSIM Results for Step 900 Validation (Hunyuan) =====")
+    print(f"Mean MS-SSIM: {mean_ssim:.4f}")
+    print(f"Min MS-SSIM: {min_ssim:.4f}")
+    print(f"Max MS-SSIM: {max_ssim:.4f}")
+
+    assert max_ssim > 0.5, f"Max SSIM is below 0.5: {max_ssim}"
+
+
+
+if __name__ == "__main__":
+    test_e2e_hunyuan_overfit_single_sample()

fastvideo/tests/ssim/test_inference_similarity.py

@@ -19,12 +19,16 @@ if "A40" in device_name:
     device_reference_folder = "A40" + device_reference_folder_suffix
 elif "L40S" in device_name:
     device_reference_folder = "L40S" + device_reference_folder_suffix
+elif "H100" in device_name or "NVIDIA H100" in device_name:
+    device_reference_folder = "H100" + device_reference_folder_suffix
+
 else:
     # device_reference_folder = "L40S" + device_reference_folder_suffix
     logger.warning(f"Unsupported device for ssim tests: {device_name}")
     # raise ValueError(f"Unsupported device for ssim tests: {device_name}")
 
 # Base parameters from the shell script
+
 HUNYUAN_PARAMS = {
     "num_gpus": 4,
     "model_path": "FastVideo/FastHunyuan-diffusers",
@@ -80,48 +84,9 @@ WAN_I2V_PARAMS = {
     "text-encoder-precision": ("fp32",)
 }
 
-# LTX-2 distilled one-stage params (no refine/upscale)
-# Official defaults: height=512, width=768, num_frames=121, fps=24, seed=10
-# Using num_frames=41 for faster CI (still valid: 41 = 8×5 + 1)
-LTX2_T2V_PARAMS = {
-    "num_gpus": 2,
-    "model_path": "FastVideo/LTX2-Distilled-Diffusers",
-    "height": 512,
-    "width": 768,
-    "num_frames": 41,  # Shorter for CI; official default is 121
-    "num_inference_steps": 8,  # Distilled uses 8 steps
-    "guidance_scale": 1.0,  # No CFG for distilled
-    "embedded_cfg_scale": 6,
-    "seed": 1024,
-    "sp_size": 2,
-    "tp_size": 1,
-    "fps": 24,
-    "neg_prompt": (
-        "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."
-    ),
-    "ltx2_vae_tiling": True,
-    "ltx2_vae_spatial_tile_size_in_pixels": 512,
-    "ltx2_vae_spatial_tile_overlap_in_pixels": 64,
-    "ltx2_vae_temporal_tile_size_in_frames": 64,
-    "ltx2_vae_temporal_tile_overlap_in_frames": 24,
-}
-
 MODEL_TO_PARAMS = {
     "FastHunyuan-diffusers": HUNYUAN_PARAMS,
     "Wan2.1-T2V-1.3B-Diffusers": WAN_T2V_PARAMS,
-    # "ltx2_diffusers": LTX2_T2V_PARAMS,
 }
 
 I2V_MODEL_TO_PARAMS = {
@@ -268,26 +233,17 @@ def test_inference_similarity(prompt, ATTENTION_BACKEND, model_id):
 
     init_kwargs = {
         "num_gpus": BASE_PARAMS["num_gpus"],
+        "flow_shift": BASE_PARAMS["flow_shift"],
         "sp_size": BASE_PARAMS["sp_size"],
         "tp_size": BASE_PARAMS["tp_size"],
         "use_fsdp_inference": True,
         "dit_cpu_offload": False,
-        "dit_layerwise_offload": False,
     }
-    if "flow_shift" in BASE_PARAMS:
-        init_kwargs["flow_shift"] = BASE_PARAMS["flow_shift"]
     if BASE_PARAMS.get("vae_sp"):
         init_kwargs["vae_sp"] = True
         init_kwargs["vae_tiling"] = True
     if "text-encoder-precision" in BASE_PARAMS:
         init_kwargs["text_encoder_precisions"] = BASE_PARAMS["text-encoder-precision"]
-    # LTX2-specific VAE tiling parameters
-    if BASE_PARAMS.get("ltx2_vae_tiling"):
-        init_kwargs["ltx2_vae_tiling"] = True
-        init_kwargs["ltx2_vae_spatial_tile_size_in_pixels"] = BASE_PARAMS.get("ltx2_vae_spatial_tile_size_in_pixels", 512)
-        init_kwargs["ltx2_vae_spatial_tile_overlap_in_pixels"] = BASE_PARAMS.get("ltx2_vae_spatial_tile_overlap_in_pixels", 64)
-        init_kwargs["ltx2_vae_temporal_tile_size_in_frames"] = BASE_PARAMS.get("ltx2_vae_temporal_tile_size_in_frames", 64)
-        init_kwargs["ltx2_vae_temporal_tile_overlap_in_frames"] = BASE_PARAMS.get("ltx2_vae_temporal_tile_overlap_in_frames", 24)
 
     generation_kwargs = {
         "num_inference_steps": num_inference_steps,

fastvideo/tests/ssim/test_matrixgame_similarity.py

@@ -88,7 +88,6 @@ def test_matrixgame_similarity(prompt, ATTENTION_BACKEND, model_id):
     init_kwargs = {
         "num_gpus": BASE_PARAMS["num_gpus"],
         "use_fsdp_inference": True,
-        "dit_layerwise_offload": False,
         "dit_cpu_offload": False,
         "vae_cpu_offload": False,
         "text_encoder_cpu_offload": True,

fastvideo/tests/ssim/test_turbodiffusion_similarity.py

@@ -218,7 +218,6 @@ def test_turbodiffusion_i2v_inference_similarity(prompt, model_id):
         # Keep both transformers in VRAM - avoids CPU RAM bottleneck
         "dit_cpu_offload": False,
         "use_fsdp_inference": True,
-        "dit_layerwise_offload": False,
     }
 
     generation_kwargs = {

fastvideo/tests/training/VSA/test_training_loss_VSA.py

@@ -6,11 +6,6 @@ import json
 from huggingface_hub import snapshot_download
 import torch
 
-# Ensure backend selection happens during import-time initialization
-os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "VIDEO_SPARSE_ATTN"
-# Force VSA to use Triton implementation even on H100 / when CUDA extension is available
-# os.environ["FASTVIDEO_KERNEL_VSA_FORCE_TRITON"] = "1"
-
 # Import the training pipeline
 sys.path.append(str(Path(__file__).parent.parent.parent.parent.parent))
 from fastvideo.training.wan_training_pipeline import main
@@ -24,6 +19,8 @@ h200_reference_wandb_summary_file = "fastvideo/tests/training/VSA/h200_reference
 NUM_NODES = "1"
 NUM_GPUS_PER_NODE = "2"
 
+os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "VIDEO_SPARSE_ATTN"
+
 def run_worker():
     """Worker function that will be run on each GPU"""
     # Create and populate args
@@ -121,10 +118,10 @@ def test_distributed_training():
     wandb_summary = json.load(open(summary_file))
 
     fields_and_thresholds = {
-        'avg_step_time': 3,
-        'grad_norm': 0.2,
-        'step_time': 2.5,
-        'train_loss': 0.04
+        'avg_step_time': 1.0,
+        'grad_norm': 0.1,
+        'step_time': 1.0,
+        'train_loss': 0.02
     }
 
     failures = []

fastvideo/training/hunyuan15_training_pipeline.py

@@ -0,0 +1,81 @@
+# SPDX-License-Identifier: Apache-2.0
+import sys
+from copy import deepcopy
+
+from fastvideo.fastvideo_args import FastVideoArgs, TrainingArgs
+from fastvideo.logger import init_logger
+from fastvideo.training.training_pipeline import TrainingPipeline
+
+from fastvideo.pipelines.basic.hunyuan15.hunyuan15_pipeline import HunyuanVideo15Pipeline
+
+logger = init_logger(__name__)
+
+
+class Hunyuan15TrainingPipeline(TrainingPipeline):
+    """
+    A training pipeline for HunyuanVideo-1.5.
+    """
+    _required_config_modules = ["scheduler", "transformer", "vae"]
+
+    def initialize_training_pipeline(self, training_args: TrainingArgs):
+        if training_args.enable_gradient_checkpointing_type is None:
+            training_args.enable_gradient_checkpointing_type = "full"
+        super().initialize_training_pipeline(training_args)
+
+    def initialize_pipeline(self, fastvideo_args: FastVideoArgs):
+        pass
+        # self.modules["scheduler"] = FlowMatchEulerDiscreteScheduler(
+        #     shift=fastvideo_args.pipeline_config.flow_shift)
+
+    def create_training_stages(self, training_args: TrainingArgs):
+        # reserved for future refactors
+        pass
+
+    def initialize_validation_pipeline(self, training_args: TrainingArgs):
+        logger.info("Initializing validation pipeline (Hunyuan15)...")
+        args_copy = deepcopy(training_args)
+        args_copy.inference_mode = True
+
+        validation_pipeline = HunyuanVideo15Pipeline.from_pretrained(
+            training_args.model_path,
+            args=args_copy,  # type: ignore
+            inference_mode=True,
+            loaded_modules={
+                # reuse the training transformer weights for validation sampling
+                "transformer": self.get_module("transformer"),
+            },
+            tp_size=training_args.tp_size,
+            sp_size=training_args.sp_size,
+            num_gpus=training_args.num_gpus,
+            pin_cpu_memory=training_args.pin_cpu_memory,
+            dit_cpu_offload=True,
+            # dit_layerwise_offload=True,
+            use_fsdp_inference=True,
+        )
+
+        self.validation_pipeline = validation_pipeline
+
+
+def main(args) -> None:
+    logger.info("Starting Hunyuan15 training pipeline...")
+
+    pipeline = Hunyuan15TrainingPipeline.from_pretrained(
+        args.pretrained_model_name_or_path,
+        args=args,
+    )
+    args = pipeline.training_args
+    pipeline.train()
+    logger.info("Training pipeline done")
+
+
+if __name__ == "__main__":
+    argv = sys.argv
+    from fastvideo.utils import FlexibleArgumentParser
+
+    parser = FlexibleArgumentParser()
+    parser = TrainingArgs.add_cli_args(parser)
+    parser = FastVideoArgs.add_cli_args(parser)
+    args = parser.parse_args()
+
+    args.dit_cpu_offload = False
+    main(args)

fastvideo/training/hunyuan_training_pipeline.py

@@ -0,0 +1,74 @@
+# SPDX-License-Identifier: Apache-2.0
+import sys
+from copy import deepcopy
+
+from fastvideo.fastvideo_args import FastVideoArgs, TrainingArgs
+from fastvideo.logger import init_logger
+from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import (
+    FlowMatchEulerDiscreteScheduler)
+from fastvideo.pipelines.basic.hunyuan.hunyuan_pipeline import HunyuanVideoPipeline
+from fastvideo.training.training_pipeline import TrainingPipeline
+from fastvideo.utils import is_vsa_available
+
+vsa_available = is_vsa_available()
+
+logger = init_logger(__name__)
+
+
+class HunyuanTrainingPipeline(TrainingPipeline):
+    """
+    A training pipeline for Hunyuan.
+    """
+    _required_config_modules = ["scheduler", "transformer", "vae"]
+
+    def initialize_pipeline(self, fastvideo_args: FastVideoArgs):
+        self.modules["scheduler"] = FlowMatchEulerDiscreteScheduler(
+            shift=fastvideo_args.pipeline_config.flow_shift)
+
+    def create_training_stages(self, training_args: TrainingArgs):
+        """
+        May be used in future refactors.
+        """
+        pass
+
+    def initialize_validation_pipeline(self, training_args: TrainingArgs):
+        logger.info("Initializing validation pipeline...")
+        args_copy = deepcopy(training_args)
+
+        args_copy.inference_mode = True
+        validation_pipeline = HunyuanVideoPipeline.from_pretrained(
+            training_args.model_path,
+            args=args_copy,  # type: ignore
+            inference_mode=True,
+            loaded_modules={
+                "transformer": self.get_module("transformer"),
+            },
+            tp_size=training_args.tp_size,
+            sp_size=training_args.sp_size,
+            num_gpus=training_args.num_gpus,
+            pin_cpu_memory=training_args.pin_cpu_memory,
+            dit_cpu_offload=True)
+
+        self.validation_pipeline = validation_pipeline
+
+
+def main(args) -> None:
+    logger.info("Starting training pipeline...")
+
+    pipeline = HunyuanTrainingPipeline.from_pretrained(
+        args.pretrained_model_name_or_path, args=args)
+    args = pipeline.training_args
+    pipeline.train()
+    logger.info("Training pipeline done")
+
+
+if __name__ == "__main__":
+    argv = sys.argv
+    from fastvideo.fastvideo_args import TrainingArgs
+    from fastvideo.utils import FlexibleArgumentParser
+    parser = FlexibleArgumentParser()
+    parser = TrainingArgs.add_cli_args(parser)
+    parser = FastVideoArgs.add_cli_args(parser)
+    args = parser.parse_args()
+    args.dit_cpu_offload = False
+    main(args)

fastvideo/training/training_pipeline.py

@@ -1,5 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 from dataclasses import asdict
+import inspect
 import math
 import os
 import time
@@ -174,6 +175,16 @@ class TrainingPipeline(LoRAPipeline, ABC):
                 last_epoch=self.init_steps - 1,
             )
 
+        text_padding_length = training_args.pipeline_config.text_encoder_configs[
+            0].arch_config.text_len  # type: ignore[attr-defined]
+        if not text_padding_length:
+            text_max_lengths = getattr(training_args.pipeline_config,
+                                       "text_encoder_max_lengths", None)
+            if text_max_lengths:
+                text_padding_length = text_max_lengths[0]
+            else:
+                text_padding_length = 512
+
         self.train_dataset, self.train_dataloader = build_parquet_map_style_dataloader(
             training_args.data_path,
             training_args.train_batch_size,
@@ -181,9 +192,7 @@ class TrainingPipeline(LoRAPipeline, ABC):
             num_data_workers=training_args.dataloader_num_workers,
             cfg_rate=training_args.training_cfg_rate,
             drop_last=True,
-            text_padding_length=training_args.pipeline_config.
-            text_encoder_configs[0].arch_config.
-            text_len,  # type: ignore[attr-defined]
+            text_padding_length=text_padding_length,
             seed=self.seed)
 
         self.noise_scheduler = noise_scheduler
@@ -272,11 +281,27 @@ class TrainingPipeline(LoRAPipeline, ABC):
     def _normalize_dit_input(self,
                              training_batch: TrainingBatch) -> TrainingBatch:
         # TODO(will): support other models
+        # Automatically detect model type based on VAE class name
+        vae = self.get_module("vae")
+        vae_class_name = vae.__class__.__name__
+
+        # Map VAE class names to model types
+        if 'Hunyuan' in vae_class_name:
+            model_type = 'hunyuan'
+        elif 'Wan' in vae_class_name or 'AutoencoderKLCausal3D' in vae_class_name:
+            model_type = 'wan'
+        else:
+            # Default to checking for latents_mean attribute
+            if hasattr(vae, 'latents_mean'):
+                model_type = 'wan'
+            else:
+                model_type = 'hunyuan'
+
         with self.tracker.timed("timing/normalize_input"):
             training_batch.latents = normalize_dit_input(
-                'wan',
+                model_type,
                 training_batch.latents,
-                self.get_module("vae"),
+                vae,
             )
         return training_batch
 
@@ -428,6 +453,87 @@ class TrainingPipeline(LoRAPipeline, ABC):
         #         device=training_batch.noisy_model_input.device,
         #         dtype=torch.bfloat16)
         current_model = self.transformer_2 if self.train_transformer_2 else self.transformer
+        if "return_dict" in input_kwargs:
+            try:
+                sig = inspect.signature(current_model.forward)
+            except (TypeError, ValueError):
+                sig = None
+            if sig is not None:
+                params = sig.parameters.values()
+                accepts_kwargs = any(param.kind == inspect.Parameter.VAR_KEYWORD
+                                     for param in params)
+                if ("return_dict" not in sig.parameters and not accepts_kwargs):
+                    input_kwargs = dict(input_kwargs)
+                    input_kwargs.pop("return_dict", None)
+        base_model = current_model
+        for attr in ("module", "_orig_mod"):
+            if hasattr(base_model, attr):
+                base_model = getattr(base_model, attr)
+        if hasattr(base_model, "_fsdp_wrapped_module"):
+            base_model = base_model._fsdp_wrapped_module
+            if hasattr(base_model, "module"):
+                base_model = base_model.module
+        model_config = getattr(base_model, "config", None)
+        if model_config is None:
+            model_config = getattr(current_model, "config", None)
+        if model_config is not None and hasattr(model_config,
+                                                "text_embed_2_dim"):
+            input_kwargs = dict(input_kwargs)
+            text_states = training_batch.encoder_hidden_states
+            text_mask = training_batch.encoder_attention_mask
+            if text_states is None or text_mask is None:
+                raise ValueError(
+                    "HunyuanVideo15 training requires text embeddings and masks"
+                )
+            if not isinstance(input_kwargs.get("encoder_hidden_states"),
+                              (list, tuple)):
+                batch_size = text_states.shape[0]
+                text_2_len = 256
+                if hasattr(self.training_args.pipeline_config,
+                           "text_encoder_max_lengths"):
+                    text_2_len = self.training_args.pipeline_config.text_encoder_max_lengths[
+                        1]
+                text_2_dim = model_config.text_embed_2_dim
+                text_2 = torch.zeros((batch_size, text_2_len, text_2_dim),
+                                     device=text_states.device,
+                                     dtype=text_states.dtype)
+                input_kwargs["encoder_hidden_states"] = [text_states, text_2]
+            if not isinstance(input_kwargs.get("encoder_attention_mask"),
+                              (list, tuple)):
+                batch_size = text_mask.shape[0]
+                text_2_len = 256
+                if hasattr(self.training_args.pipeline_config,
+                           "text_encoder_max_lengths"):
+                    text_2_len = self.training_args.pipeline_config.text_encoder_max_lengths[
+                        1]
+                text_mask_2 = torch.zeros((batch_size, text_2_len),
+                                          device=text_mask.device,
+                                          dtype=text_mask.dtype)
+                input_kwargs["encoder_attention_mask"] = [
+                    text_mask, text_mask_2
+                ]
+            if "encoder_hidden_states_image" not in input_kwargs:
+                image_embeds = training_batch.image_embeds
+                if image_embeds is None or image_embeds.numel() == 0:
+                    batch_size = text_states.shape[0]
+                    image_len = 729
+                    image_dim = model_config.image_embed_dim
+                    image_embeds = torch.zeros(
+                        (batch_size, image_len, image_dim),
+                        device=text_states.device,
+                        dtype=text_states.dtype,
+                    )
+                input_kwargs["encoder_hidden_states_image"] = [image_embeds]
+            hidden_states = input_kwargs.get("hidden_states")
+            if (isinstance(hidden_states, torch.Tensor)
+                    and hidden_states.shape[1] != model_config.in_channels):
+                batch_size, _, t, h, w = hidden_states.shape
+                video_latent = torch.zeros((batch_size, 1, t, h, w),
+                                           device=hidden_states.device,
+                                           dtype=hidden_states.dtype)
+                zeros_latent = torch.zeros_like(hidden_states)
+                input_kwargs["hidden_states"] = torch.cat(
+                    [hidden_states, video_latent, zeros_latent], dim=1)
 
         with self.tracker.timed("timing/forward_backward"), set_forward_context(
                 current_timestep=training_batch.current_timestep,
@@ -650,7 +756,13 @@ class TrainingPipeline(LoRAPipeline, ABC):
                 }
                 metrics["batch_size"] = int(training_batch.raw_latent_shape[0])
 
-                patch_t, patch_h, patch_w = self.training_args.pipeline_config.dit_config.patch_size
+                patch_size = self.training_args.pipeline_config.dit_config.patch_size
+                if isinstance(patch_size, tuple):
+                    patch_t, patch_h, patch_w = patch_size
+                else:
+                    patch_t = patch_size
+                    patch_h = patch_size
+                    patch_w = patch_size
                 seq_len = (training_batch.raw_latent_shape[2] // patch_t) * (
                     training_batch.raw_latent_shape[3] //
                     patch_h) * (training_batch.raw_latent_shape[4] // patch_w)
@@ -663,7 +775,13 @@ class TrainingPipeline(LoRAPipeline, ABC):
 
                 metrics["hidden_dim"] = arch_config.hidden_size
                 metrics["num_layers"] = arch_config.num_layers
-                metrics["ffn_dim"] = arch_config.ffn_dim
+                ffn_dim = getattr(arch_config, "ffn_dim", None)
+                if ffn_dim is None:
+                    mlp_ratio = getattr(arch_config, "mlp_ratio", None)
+                    if mlp_ratio and arch_config.hidden_size:
+                        ffn_dim = int(arch_config.hidden_size * mlp_ratio)
+                if ffn_dim is not None:
+                    metrics["ffn_dim"] = ffn_dim
 
                 self.tracker.log(metrics, step)
             if step % self.training_args.training_state_checkpointing_steps == 0:
@@ -838,8 +956,12 @@ class TrainingPipeline(LoRAPipeline, ABC):
                 step_captions.append(batch.prompt)
 
                 # Run validation inference
+                # Use the validation pipeline args to match its loaded modules.
+                validation_args = self.validation_pipeline.fastvideo_args
+                if validation_args is None:
+                    validation_args = training_args
                 output_batch = self.validation_pipeline.forward(
-                    batch, training_args)
+                    batch, validation_args)
                 samples = output_batch.output
 
                 if self.rank_in_sp_group != 0:

fastvideo/training/training_utils.py

@@ -849,7 +849,14 @@ def load_distillation_checkpoint(
 
 def normalize_dit_input(model_type, latents, vae) -> torch.Tensor:
     if model_type == "hunyuan_hf" or model_type == "hunyuan":
-        return latents * 0.476986
+        scaling_factor = vae.scaling_factor
+        if scaling_factor is None:
+            raise AttributeError(
+                "Hunyuan VAE must define scaling_factor for training.")
+        if isinstance(scaling_factor, torch.Tensor):
+            scaling_factor = scaling_factor.to(device=latents.device,
+                                               dtype=latents.dtype)
+        return latents * scaling_factor
     elif model_type == "wan":
         latents_mean = torch.tensor(vae.latents_mean)
         latents_std = 1.0 / torch.tensor(vae.latents_std)

fastvideo/worker/multiproc_executor.py

@@ -132,13 +132,9 @@ class MultiprocExecutor(Executor):
         else:
             logging_info = None
 
-        # Get extra dict (contains audio, etc.)
-        extra = responses[0].get("extra", {})
-
         result_batch = ForwardBatch(data_type=forward_batch.data_type,
                                     output=output,
-                                    logging_info=logging_info,
-                                    extra=extra)
+                                    logging_info=logging_info)
 
         return result_batch
 
@@ -652,8 +648,7 @@ class WorkerMultiprocProc:
                             logging_info = output_batch.logging_info
                         self.pipe.send({
                             "output_batch": output_batch.output.cpu(),
-                            "logging_info": logging_info,
-                            "extra": output_batch.extra,
+                            "logging_info": logging_info
                         })
                     else:
                         result = self.worker.execute_method(

fastvideo/workflow/preprocess/components.py

@@ -200,8 +200,14 @@ class ParquetDatasetSaver:
             value = getattr(batch, key.name)
             if isinstance(value, list):
                 for idx in range(len(value)):
-                    if isinstance(value[idx], torch.Tensor):
-                        value[idx] = value[idx].cpu().numpy()
+                    # if isinstance(value[idx], torch.Tensor):
+                    #     value[idx] = value[idx].cpu().numpy()
+                    t = value[idx]
+                    if isinstance(t, torch.Tensor):
+                        if t.dtype == torch.bfloat16:
+                            t = t.float()   # or t.to(torch.float16) if you want smaller
+                        value[idx] = t.cpu().numpy()
+
             elif isinstance(value, torch.Tensor):
                 value = value.cpu().numpy()
                 setattr(batch, key.name, value)
@@ -288,10 +294,17 @@ def build_dataset(preprocess_config: PreprocessConfig, split: str,
                                split=split)
         column_names = dataset.column_names
         # rename columns to match the schema
-        if "cap" in column_names:
+        if "cap" in column_names and "caption" not in column_names:
             dataset = dataset.rename_column("cap", "caption")
         if "path" in column_names:
             dataset = dataset.rename_column("path", "name")
+        
+        # Add num_frames column if it doesn't exist but fps and duration do
+        if 'num_frames' not in column_names and 'fps' in column_names and 'duration' in column_names:
+            def add_num_frames(item: dict[str, Any]) -> dict[str, Any]:
+                item['num_frames'] = int(item['fps'] * item['duration'])
+                return item
+            dataset = dataset.map(add_num_frames)
 
         dataset = dataset.filter(validator)
         dataset = dataset.shard(num_shards=get_world_size(),

kill.sh

@@ -0,0 +1,14 @@
+#!/bin/bash
+# Kill all GPU processes shown by nvidia-smi
+
+pids=$(nvidia-smi --query-compute-apps=pid --format=csv,noheader 2>/dev/null | tr -d ' ')
+
+if [ -z "$pids" ]; then
+    echo "No GPU processes found."
+    exit 0
+fi
+
+echo "Killing GPU processes: $pids"
+for pid in $pids; do
+    kill -9 "$pid" 2>/dev/null && echo "Killed PID $pid" || echo "Failed to kill PID $pid"
+done

pyproject.toml

@@ -29,7 +29,6 @@ dependencies = [
     "diffusers>=0.33.1",
     "torch>=2.9.1",
     "torchvision",
-    "torchaudio",
 
     # Acceleration & Optimization
     "accelerate==1.0.1",
@@ -64,7 +63,7 @@ dependencies = [
     "remote-pdb",
 
     # Kernel & Packaging
-    "fastvideo-kernel==0.2.4",
+    "fastvideo-kernel==0.2.3",
     "wheel",
 
     # Training Dependencies

pyproject_other.toml

@@ -63,7 +63,7 @@ dependencies = [
     "remote-pdb",
 
     # Kernel & Packaging
-    "fastvideo-kernel==0.2.4",
+    "fastvideo-kernel==0.2.3",
     "wheel",
 
     # Training Dependencies

scripts/checkpoint_conversion/convert_ltx2_weights.py

@@ -1,444 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""
-Convert LTX-2 weights to FastVideo naming conventions and split by component.
-"""
-
-from __future__ import annotations
-
-import argparse
-import glob
-import json
-import os
-import re
-import shutil
-from collections import OrderedDict
-from pathlib import Path
-
-import torch
-from safetensors import safe_open
-from safetensors.torch import load_file, save_file
-
-try:
-    from huggingface_hub import snapshot_download
-except ImportError:  # pragma: no cover - optional dependency
-    snapshot_download = None
-
-
-PARAM_NAME_MAP: dict[str, str] = {
-    r"^model\.diffusion_model\.(.*)$": r"\1",
-}
-
-COMPONENT_PREFIXES: dict[str, tuple[str, ...]] = {
-    "transformer": ("model.diffusion_model.",),
-    "vae": ("vae.",),
-    "audio_vae": ("audio_vae.",),
-    "vocoder": ("vocoder.",),
-    "text_embedding_projection": ("text_embedding_projection.", "model.text_embedding_projection."),
-}
-
-
-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 _apply_mapping(key: str) -> str:
-    for pattern, replacement in PARAM_NAME_MAP.items():
-        if re.match(pattern, key):
-            return re.sub(pattern, replacement, key)
-    return key
-
-
-def _load_weights(shards: list[Path]) -> dict[str, torch.Tensor]:
-    weights: dict[str, torch.Tensor] = {}
-    for shard in shards:
-        weights.update(load_file(str(shard)))
-    return weights
-
-
-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:
-        return {}
-    return json.loads(metadata["config"])
-
-
-def _filter_transformer_config(config: dict) -> dict:
-    transformer = config.get("transformer", {})
-    allowed = {
-        "num_attention_heads",
-        "attention_head_dim",
-        "num_layers",
-        "cross_attention_dim",
-        "caption_channels",
-        "norm_eps",
-        "attention_type",
-        "positional_embedding_theta",
-        "positional_embedding_max_pos",
-        "timestep_scale_multiplier",
-        "use_middle_indices_grid",
-        "rope_type",
-        "frequencies_precision",
-        "in_channels",
-        "out_channels",
-        "audio_num_attention_heads",
-        "audio_attention_head_dim",
-        "audio_in_channels",
-        "audio_out_channels",
-        "audio_cross_attention_dim",
-        "audio_positional_embedding_max_pos",
-        "av_ca_timestep_scale_multiplier",
-    }
-    filtered = {k: v for k, v in transformer.items() if k in allowed}
-    if "frequencies_precision" in filtered:
-        filtered["double_precision_rope"] = filtered["frequencies_precision"] == "float64"
-        del filtered["frequencies_precision"]
-    return filtered
-
-
-def _build_text_embedding_projection_config(
-    gemma_model_path: str = "",
-) -> dict:
-    return {
-        "architectures": ["LTX2GemmaTextEncoderModel"],
-        "hidden_size": 3840,
-        "num_hidden_layers": 48,
-        "num_attention_heads": 30,
-        "text_len": 1024,
-        "pad_token_id": 0,
-        "eos_token_id": 2,
-        "gemma_model_path": gemma_model_path,
-        "gemma_dtype": "bfloat16",
-        "padding_side": "left",
-        "feature_extractor_in_features": 3840 * 49,
-        "feature_extractor_out_features": 3840,
-        "connector_num_attention_heads": 30,
-        "connector_attention_head_dim": 128,
-        "connector_num_layers": 2,
-        "connector_positional_embedding_theta": 10000.0,
-        "connector_positional_embedding_max_pos": [4096],
-        "connector_rope_type": "split",
-        "connector_double_precision_rope": True,
-        "connector_num_learnable_registers": 128,
-    }
-
-
-def _wrap_component_config(
-    component_name: str,
-    component_config: dict | None,
-    class_name: str | None = None,
-) -> dict | None:
-    if component_config is None:
-        return None
-    wrapped = {component_name: component_config}
-    if class_name is not None:
-        wrapped["_class_name"] = class_name
-    return wrapped
-
-
-def _split_component_weights(weights: dict[str, torch.Tensor]) -> dict[str, OrderedDict]:
-    components: dict[str, OrderedDict] = {name: OrderedDict() for name in COMPONENT_PREFIXES}
-    for key, value in weights.items():
-        if key.startswith("model.diffusion_model.audio_embeddings_connector."):
-            new_key = key.replace("model.diffusion_model.audio_embeddings_connector.", "audio_embeddings_connector.")
-            components["text_embedding_projection"][new_key] = value
-            continue
-        if key.startswith("model.diffusion_model.video_embeddings_connector."):
-            new_key = key.replace("model.diffusion_model.video_embeddings_connector.", "embeddings_connector.")
-            components["text_embedding_projection"][new_key] = value
-            continue
-
-        matched = False
-        for component, prefixes in COMPONENT_PREFIXES.items():
-            for prefix in prefixes:
-                if key.startswith(prefix):
-                    new_key = key[len(prefix):]
-                    components[component][new_key] = value
-                    matched = True
-                    break
-            if matched:
-                break
-    return {name: weights for name, weights in components.items() if weights}
-
-
-def _write_component(
-    output_dir: Path,
-    name: str,
-    weights: OrderedDict,
-    config: dict | None,
-    dir_name: str | None = None,
-) -> None:
-    component_dir = output_dir / (dir_name or name)
-    component_dir.mkdir(parents=True, exist_ok=True)
-    output_file = component_dir / "model.safetensors"
-    save_file(weights, str(output_file))
-    print(f"Saved {name} weights to {output_file}")
-
-    if config is not None:
-        config_path = component_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 {name} config to {config_path}")
-
-
-def _build_model_index(
-    transformer_class_name: str,
-    vae_class_name: str,
-    pipeline_class_name: str,
-    diffusers_version: str,
-) -> dict:
-    return {
-        "_class_name": pipeline_class_name,
-        "_diffusers_version": diffusers_version,
-        "transformer": ["diffusers", transformer_class_name],
-        "vae": ["diffusers", vae_class_name],
-        "text_encoder": ["transformers", "LTX2GemmaTextEncoderModel"],
-        "tokenizer": ["transformers", "AutoTokenizer"],
-        "audio_vae": ["diffusers", "LTX2AudioDecoder"],
-        "vocoder": ["diffusers", "LTX2Vocoder"],
-    }
-
-
-def _write_model_index(output_dir: Path, model_index: dict) -> None:
-    output_dir.mkdir(parents=True, exist_ok=True)
-    model_index_path = output_dir / "model_index.json"
-    with model_index_path.open("w", encoding="utf-8") as f:
-        json.dump(model_index, f, indent=2)
-        f.write("\n")
-    print(f"Saved model_index.json to {model_index_path}")
-
-
-def convert_components(
-    source_path: Path,
-    output_dir: Path,
-    metadata_config: dict,
-    transformer_class_name: str,
-    components_to_write: set[str] | None = None,
-    emit_diffusers_repo: bool = True,
-    pipeline_class_name: str = "LTX2Pipeline",
-    diffusers_version: str = "0.33.0.dev0",
-    gemma_model_path: str = "",
-) -> None:
-    shards = _find_shards(source_path)
-    if not shards:
-        raise FileNotFoundError(f"No safetensors found in {source_path}")
-
-    weights = _load_weights(shards)
-    split_weights = _split_component_weights(weights)
-    if components_to_write is not None:
-        split_weights = {name: weights for name, weights in split_weights.items() if name in components_to_write}
-
-    transformer_weights = split_weights.get("transformer", OrderedDict())
-    converted_transformer = OrderedDict()
-    for key, value in transformer_weights.items():
-        new_key = _apply_mapping(f"model.diffusion_model.{key}")
-        converted_transformer[new_key] = value
-    split_weights["transformer"] = converted_transformer
-
-    transformer_config = _filter_transformer_config(metadata_config)
-    if transformer_config:
-        transformer_config["_class_name"] = transformer_class_name
-
-    component_configs: dict[str, dict | None] = {
-        "transformer": transformer_config or None,
-        "vae": _wrap_component_config(
-            "vae",
-            metadata_config.get("vae"),
-            class_name="CausalVideoAutoencoder",
-        ),
-        "audio_vae": _wrap_component_config(
-            "audio_vae",
-            metadata_config.get("audio_vae"),
-            class_name="LTX2AudioDecoder",
-        ),
-        "vocoder": _wrap_component_config(
-            "vocoder",
-            metadata_config.get("vocoder"),
-            class_name="LTX2Vocoder",
-        ),
-        "text_embedding_projection": _build_text_embedding_projection_config(
-            gemma_model_path=gemma_model_path
-        ),
-    }
-
-    output_dir.mkdir(parents=True, exist_ok=True)
-    for name, component_weights in split_weights.items():
-        _write_component(output_dir, name, component_weights, component_configs.get(name))
-        if emit_diffusers_repo and name == "text_embedding_projection":
-            _write_component(
-                output_dir,
-                name,
-                component_weights,
-                component_configs.get(name),
-                dir_name="text_encoder",
-            )
-    if emit_diffusers_repo:
-        required_for_index = {
-            "transformer",
-            "vae",
-            "audio_vae",
-            "vocoder",
-            "text_embedding_projection",
-        }
-        if components_to_write is not None and not required_for_index.issubset(components_to_write):
-            print("Skipping model_index.json; not all diffusers components were written.")
-            return
-        if not required_for_index.issubset(split_weights.keys()):
-            print("Skipping model_index.json; missing diffusers components in weights.")
-            return
-        vae_class_name = (component_configs.get("vae") or {}).get(
-            "_class_name", "CausalVideoAutoencoder"
-        )
-        model_index = _build_model_index(
-            transformer_class_name=transformer_class_name,
-            vae_class_name=vae_class_name,
-            pipeline_class_name=pipeline_class_name,
-            diffusers_version=diffusers_version,
-        )
-        _write_model_index(output_dir, model_index)
-
-
-def update_transformer_config(config_path: Path, class_name: str) -> None:
-    if not config_path.exists():
-        print(f"Config file not found: {config_path}")
-        return
-
-    with config_path.open("r", encoding="utf-8") as f:
-        config = json.load(f)
-
-    config["_class_name"] = class_name
-    with config_path.open("w", encoding="utf-8") as f:
-        json.dump(config, f, indent=2)
-        f.write("\n")
-    print(f"Updated _class_name in {config_path} -> {class_name}")
-
-
-def maybe_download(repo_id: str, target_dir: Path, token: str | None, allow_patterns: str | None) -> Path:
-    if snapshot_download is None:
-        raise RuntimeError("huggingface_hub is required for --download")
-    target_dir.mkdir(parents=True, exist_ok=True)
-    snapshot_download(
-        repo_id=repo_id,
-        local_dir=str(target_dir),
-        local_dir_use_symlinks=False,
-        token=token,
-        allow_patterns=allow_patterns,
-    )
-    return target_dir
-
-
-def main() -> None:
-    parser = argparse.ArgumentParser(description="Convert LTX-2 weights to FastVideo format")
-    parser.add_argument("--source", type=str, help="Path to transformer weights directory")
-    parser.add_argument("--output", type=str, required=True, help="Output directory for converted weights")
-    parser.add_argument("--download", type=str, help="HF repo id to download before conversion")
-    parser.add_argument("--allow-patterns", type=str, help="Limit HF download to matching files")
-    parser.add_argument("--token", type=str, default=os.getenv("HF_TOKEN"), help="HF token (or set HF_TOKEN)")
-    parser.add_argument("--update-config", action="store_true", help="Update source config.json _class_name")
-    parser.add_argument("--class-name", type=str, default="LTX2Transformer3DModel")
-    parser.add_argument(
-        "--diffusers-repo",
-        action=argparse.BooleanOptionalAction,
-        default=True,
-        help="Emit a diffusers-style repo layout with model_index.json.",
-    )
-    parser.add_argument(
-        "--pipeline-class-name",
-        type=str,
-        default="LTX2Pipeline",
-        help="Pipeline class name for model_index.json.",
-    )
-    parser.add_argument(
-        "--diffusers-version",
-        type=str,
-        default="0.33.0.dev0",
-        help="Diffusers version for model_index.json.",
-    )
-    parser.add_argument(
-        "--transformer-only",
-        action="store_true",
-        help="Only convert transformer weights (no component split).",
-    )
-    parser.add_argument(
-        "--components",
-        type=str,
-        default="",
-        help=(
-            "Comma-separated component list to write "
-            "(transformer,vae,audio_vae,vocoder,text_embedding_projection)."
-        ),
-    )
-    parser.add_argument(
-        "--gemma-path",
-        type=str,
-        default="",
-        help="Optional local Gemma model path to copy into the output repo.",
-    )
-
-    args = parser.parse_args()
-
-    if args.download:
-        if args.source:
-            raise ValueError("Use either --download or --source, not both.")
-        source_dir = maybe_download(args.download, Path(args.output) / "download", args.token, args.allow_patterns)
-    else:
-        if not args.source:
-            raise ValueError("--source is required when not using --download")
-        source_dir = Path(args.source)
-
-    output_dir = Path(args.output)
-    shards = _find_shards(source_dir)
-    if not shards:
-        raise FileNotFoundError(f"No safetensors found in {source_dir}")
-    metadata_path = shards[0]
-    metadata_config = _read_metadata_config(metadata_path)
-    components_to_write: set[str] | None = None
-    if args.transformer_only:
-        components_to_write = {"transformer"}
-    elif args.components:
-        components_to_write = {
-            component.strip()
-            for component in args.components.split(",")
-            if component.strip()
-        }
-
-    gemma_model_path = ""
-    if args.gemma_path:
-        gemma_src = Path(args.gemma_path)
-        if not gemma_src.is_dir():
-            raise ValueError(f"--gemma-path must be a directory: {gemma_src}")
-        gemma_dest = output_dir / "text_encoder" / "gemma"
-        if gemma_dest.exists():
-            shutil.rmtree(gemma_dest)
-        gemma_dest.parent.mkdir(parents=True, exist_ok=True)
-        shutil.copytree(gemma_src, gemma_dest)
-        gemma_model_path = "gemma"
-
-    convert_components(
-        source_dir,
-        output_dir,
-        metadata_config,
-        args.class_name,
-        components_to_write=components_to_write,
-        emit_diffusers_repo=args.diffusers_repo,
-        pipeline_class_name=args.pipeline_class_name,
-        diffusers_version=args.diffusers_version,
-        gemma_model_path=gemma_model_path,
-    )
-
-    if args.update_config:
-        if source_dir.is_dir():
-            update_transformer_config(source_dir / "config.json", args.class_name)
-
-
-if __name__ == "__main__":
-    main()

tests/README.md

@@ -1,7 +0,0 @@
-# Tests
-
-- `tests/local_tests/` are local-only tests that require a checked-out `LTX-2/`
-  directory under the repo root (`FastVideo/LTX-2`). Without that repo, they
-  will skip or fail.
-- The CI-backed test suite still lives in `fastvideo/tests/`.
-- Eventually, all tests will move under `tests/`.

tests/local_tests/README.md

@@ -1,7 +0,0 @@
-# Local LTX-2 Tests
-
-These tests depend on a checked-out `LTX-2/` directory under the repo root
-(`FastVideo/LTX-2`). Without that local repo, the tests will skip or fail.
-
-For the CI-backed test suite, see `fastvideo/tests/`. Those are the tests
-currently exercised in CI. (Eventually all tests will move to `tests/`.)

tests/local_tests/encoders/test_ltx2_gemma_encoder.py

@@ -1,126 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-import os
-import sys
-from pathlib import Path
-
-import pytest
-import torch
-from safetensors.torch import load_file
-from torch.testing import assert_close
-
-from fastvideo.configs.models.encoders import LTX2GemmaConfig
-from fastvideo.configs.pipelines import PipelineConfig
-from fastvideo.fastvideo_args import FastVideoArgs
-from fastvideo.forward_context import set_forward_context
-from fastvideo.models.loader.component_loader import TextEncoderLoader
-
-repo_root = Path(__file__).resolve().parents[3]
-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))
-
-
-def _load_connector_weights(path: str) -> dict[str, torch.Tensor]:
-    weights = load_file(path)
-    mapped: dict[str, torch.Tensor] = {}
-    for name, tensor in weights.items():
-        if name == "aggregate_embed.weight":
-            mapped["feature_extractor_linear.aggregate_embed.weight"] = tensor
-        elif name.startswith("embeddings_connector."):
-            mapped[name] = tensor
-        elif name.startswith("audio_embeddings_connector."):
-            mapped[name] = tensor
-    return mapped
-
-
-@pytest.mark.skipif(
-    not torch.cuda.is_available(),
-    reason="LTX-2 Gemma encoder parity test requires CUDA.",
-)
-def test_ltx2_gemma_text_encoder_parity():
-    diffusers_root = Path(
-        os.getenv("LTX2_DIFFUSERS_PATH", "converted/ltx2_diffusers")
-    )
-    text_encoder_path = os.getenv(
-        "LTX2_TEXT_ENCODER_PATH",
-        str(diffusers_root / "text_encoder"),
-    )
-    gemma_model_path = str(Path(text_encoder_path) / "gemma")
-    if not os.path.isdir(text_encoder_path):
-        pytest.skip(f"LTX-2 text encoder weights not found at {text_encoder_path}")
-    if not gemma_model_path or not os.path.isdir(gemma_model_path):
-        pytest.skip("Gemma weights not found in text_encoder/gemma.")
-
-    try:
-        from ltx_core.text_encoders.gemma.embeddings_connector import (
-            Embeddings1DConnector,
-        )
-        from ltx_core.text_encoders.gemma.encoders.av_encoder import (
-            AVGemmaTextEncoderModel,
-        )
-        from ltx_core.text_encoders.gemma.feature_extractor import (
-            GemmaFeaturesExtractorProjLinear,
-        )
-        from ltx_core.text_encoders.gemma.tokenizer import LTXVGemmaTokenizer
-        from transformers import Gemma3ForConditionalGeneration
-    except Exception as exc:
-        pytest.skip(f"LTX-2 Gemma import failed: {exc}")
-
-    device = torch.device("cuda:0")
-    precision = torch.bfloat16
-
-    tokenizer = LTXVGemmaTokenizer(gemma_model_path, max_length=1024)
-    gemma_model = Gemma3ForConditionalGeneration.from_pretrained(
-        gemma_model_path,
-        local_files_only=True,
-        torch_dtype=precision,
-    ).to(device)
-    gemma_model.eval()
-
-    ref_model = AVGemmaTextEncoderModel(
-        feature_extractor_linear=GemmaFeaturesExtractorProjLinear(),
-        embeddings_connector=Embeddings1DConnector(),
-        audio_embeddings_connector=Embeddings1DConnector(),
-        tokenizer=tokenizer,
-        model=gemma_model,
-        dtype=precision,
-    ).to(device)
-    ref_model.eval()
-
-    connector_weights = _load_connector_weights(
-        os.path.join(text_encoder_path, "model.safetensors")
-    )
-    ref_model.load_state_dict(connector_weights, strict=False)
-
-    prompt = "A fast moving train in a snowy landscape."
-    token_pairs = tokenizer.tokenize_with_weights(prompt)["gemma"]
-    input_ids = torch.tensor(
-        [[t[0] for t in token_pairs]], device=device, dtype=torch.long
-    )
-    attention_mask = torch.tensor(
-        [[t[1] for t in token_pairs]], device=device, dtype=torch.long
-    )
-
-    args = FastVideoArgs(
-        model_path=text_encoder_path,
-        pipeline_config=PipelineConfig(
-            text_encoder_configs=(LTX2GemmaConfig(),),
-            text_encoder_precisions=("bf16",),
-        ),
-    )
-    loader = TextEncoderLoader()
-    fastvideo_model = loader.load(text_encoder_path, args).to(device)
-    fastvideo_model.eval()
-
-    with torch.no_grad():
-        ref_video, ref_audio, ref_mask = ref_model(prompt, padding_side="left")
-        with set_forward_context(current_timestep=0, attn_metadata=None):
-            fastvideo_out = fastvideo_model(
-                input_ids=input_ids,
-                attention_mask=attention_mask,
-                output_hidden_states=True,
-            )
-
-    assert_close(ref_video, fastvideo_out.last_hidden_state, atol=1e-2, rtol=1e-2)
-    assert_close(ref_audio, fastvideo_out.hidden_states[0], atol=1e-2, rtol=1e-2)
-    assert torch.equal(ref_mask, fastvideo_out.attention_mask)

tests/local_tests/encoders/test_ltx2_gemma_parity.py

@@ -1,378 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-import os
-from pathlib import Path
-import sys
-
-import pytest
-import torch
-from torch.testing import assert_close
-
-from safetensors.torch import load_file
-
-from fastvideo.configs.models.encoders import LTX2GemmaConfig
-from fastvideo.models.encoders.gemma import LTX2GemmaTextEncoderModel
-from fastvideo.models.loader.component_loader import get_diffusers_config
-
-
-repo_root = Path(__file__).resolve().parents[3]
-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))
-
-
-def _init_log_paths() -> tuple[Path, Path]:
-    base_dir = Path(os.getenv("LTX2_DEBUG_DIR", "ltx2_debug"))
-    fastvideo_log = Path(os.getenv("LTX2_FASTVIDEO_GEMMA_LOG", base_dir / "fastvideo_gemma.log"))
-    reference_log = Path(os.getenv("LTX2_REFERENCE_GEMMA_LOG", base_dir / "reference_gemma.log"))
-    fastvideo_log.parent.mkdir(parents=True, exist_ok=True)
-    reference_log.parent.mkdir(parents=True, exist_ok=True)
-    fastvideo_log.write_text("")
-    reference_log.write_text("")
-    return fastvideo_log, reference_log
-
-
-def _log_line(path: Path, message: str) -> None:
-    with path.open("a", encoding="utf-8") as f:
-        f.write(message + "\n")
-
-
-def _attach_encoder_logging(
-    encoder: torch.nn.Module,
-    log_path: Path,
-    label: str,
-) -> None:
-    def _format_sum(tensor: torch.Tensor | None) -> str:
-        if tensor is None:
-            return "None"
-        return f"{tensor.float().sum().item():.6f}"
-
-    def _hook_factory(name: str):
-        def _hook(_module, _inputs, outputs):  # noqa: ANN001
-            out = outputs[0] if isinstance(outputs, tuple) else outputs
-            out_sum = _format_sum(out if torch.is_tensor(out) else None)
-            _log_line(log_path, f"{label}:{name}:sum={out_sum}")
-        return _hook
-
-    def _pre_hook_factory(name: str):
-        def _hook(_module, inputs):  # noqa: ANN001
-            tensor = inputs[0] if inputs else None
-            in_sum = _format_sum(tensor if torch.is_tensor(tensor) else None)
-            _log_line(log_path, f"{label}:{name}:in_sum={in_sum}")
-        return _hook
-
-    def _attach_block_detail(block: torch.nn.Module, prefix: str) -> None:
-        block.register_forward_pre_hook(_pre_hook_factory(f"{prefix}:input"))
-        block.register_forward_hook(_hook_factory(f"{prefix}:output"))
-        if hasattr(block, "attn1"):
-            block.attn1.register_forward_hook(
-                _hook_factory(f"{prefix}:attn1"))
-        if hasattr(block, "ff"):
-            block.ff.register_forward_hook(_hook_factory(f"{prefix}:ff"))
-
-    if hasattr(encoder, "feature_extractor_linear"):
-        encoder.feature_extractor_linear.register_forward_pre_hook(
-            _pre_hook_factory("feature_extractor_linear"))
-        encoder.feature_extractor_linear.register_forward_hook(_hook_factory("feature_extractor_linear"))
-    if hasattr(encoder, "embeddings_connector"):
-        encoder.embeddings_connector.register_forward_hook(_hook_factory("embeddings_connector"))
-        for idx, block in enumerate(encoder.embeddings_connector.transformer_1d_blocks):
-            _attach_block_detail(block, f"embeddings_block_{idx}")
-    if hasattr(encoder, "audio_embeddings_connector"):
-        encoder.audio_embeddings_connector.register_forward_hook(_hook_factory("audio_embeddings_connector"))
-        for idx, block in enumerate(encoder.audio_embeddings_connector.transformer_1d_blocks):
-            _attach_block_detail(block, f"audio_embeddings_block_{idx}")
-
-
-def _log_register_sums(encoder: torch.nn.Module, log_path: Path, label: str) -> None:
-    def _sum_param(module: torch.nn.Module, name: str) -> float | None:
-        if not hasattr(module, "learnable_registers"):
-            return None
-        param = getattr(module, "learnable_registers")
-        if not torch.is_tensor(param):
-            return None
-        return param.float().sum().item()
-
-    if hasattr(encoder, "embeddings_connector"):
-        reg_sum = _sum_param(encoder.embeddings_connector, "learnable_registers")
-        if reg_sum is not None:
-            _log_line(log_path, f"{label}:embeddings_registers:sum={reg_sum:.6f}")
-    if hasattr(encoder, "audio_embeddings_connector"):
-        reg_sum = _sum_param(encoder.audio_embeddings_connector, "learnable_registers")
-        if reg_sum is not None:
-            _log_line(log_path, f"{label}:audio_registers:sum={reg_sum:.6f}")
-
-
-def _log_param_sums(encoder: torch.nn.Module, log_path: Path, label: str) -> None:
-    def _log_param(name: str, tensor: torch.Tensor | None) -> None:
-        if tensor is None:
-            _log_line(log_path, f"{label}:param:{name}:sum=None")
-            return
-        _log_line(log_path, f"{label}:param:{name}:sum={tensor.float().sum().item():.6f}")
-
-    if hasattr(encoder, "feature_extractor_linear"):
-        _log_param(
-            "feature_extractor_linear.aggregate_embed.weight",
-            encoder.feature_extractor_linear.aggregate_embed.weight,
-        )
-    if hasattr(encoder, "embeddings_connector"):
-        block0 = encoder.embeddings_connector.transformer_1d_blocks[0]
-        _log_param("embeddings_block0.attn1.to_q.weight", block0.attn1.to_q.weight)
-        _log_param("embeddings_block0.attn1.to_k.weight", block0.attn1.to_k.weight)
-        _log_param("embeddings_block0.attn1.to_v.weight", block0.attn1.to_v.weight)
-        _log_param("embeddings_block0.ff.net.0.proj.weight", block0.ff.net[0].proj.weight)
-    if hasattr(encoder, "audio_embeddings_connector"):
-        block0 = encoder.audio_embeddings_connector.transformer_1d_blocks[0]
-        _log_param("audio_block0.attn1.to_q.weight", block0.attn1.to_q.weight)
-        _log_param("audio_block0.attn1.to_k.weight", block0.attn1.to_k.weight)
-        _log_param("audio_block0.attn1.to_v.weight", block0.attn1.to_v.weight)
-        _log_param("audio_block0.ff.net.0.proj.weight", block0.ff.net[0].proj.weight)
-
-
-def _log_gemma_param_sums(
-    gemma_model: torch.nn.Module | None,
-    log_path: Path,
-    label: str,
-) -> None:
-    if gemma_model is None:
-        _log_line(log_path, f"{label}:gemma_param:embed_tokens.weight:sum=None")
-        return
-    tokens = None
-    if hasattr(gemma_model, "get_input_embeddings"):
-        try:
-            tokens = gemma_model.get_input_embeddings()
-        except Exception:
-            tokens = None
-    if tokens is None:
-        embed = getattr(gemma_model, "model", None)
-        if embed is not None:
-            tokens = getattr(embed, "embed_tokens", None)
-    if tokens is None or not hasattr(tokens, "weight"):
-        _log_line(log_path, f"{label}:gemma_param:embed_tokens.weight:sum=None")
-        return
-    _log_line(
-        log_path,
-        f"{label}:gemma_param:embed_tokens.weight:sum={tokens.weight.float().sum().item():.6f}",
-    )
-
-
-@pytest.mark.skipif(
-    not torch.cuda.is_available(),
-    reason="LTX-2 Gemma parity test requires CUDA.",
-)
-def test_ltx2_gemma_parity():
-    os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "TORCH_SDPA"
-    torch.backends.cuda.enable_flash_sdp(False)
-    torch.backends.cuda.enable_mem_efficient_sdp(False)
-    torch.backends.cuda.enable_math_sdp(True)
-    fastvideo_log, reference_log = _init_log_paths()
-    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",
-        )
-    )
-    text_encoder_path = diffusers_root / "text_encoder"
-    gemma_path = text_encoder_path / "gemma"
-
-    if not official_path.exists():
-        pytest.skip(f"LTX-2 weights not found at {official_path}")
-    if not text_encoder_path.exists():
-        pytest.skip(f"LTX-2 text encoder not found at {text_encoder_path}")
-    if not gemma_path.exists():
-        pytest.skip(f"LTX-2 Gemma weights not found at {gemma_path}")
-
-    try:
-        from ltx_core.loader.single_gpu_model_builder import SingleGPUModelBuilder
-        from ltx_core.model.transformer.attention import Attention, AttentionFunction
-        from ltx_core.text_encoders.gemma import (
-            AV_GEMMA_TEXT_ENCODER_KEY_OPS,
-            AVGemmaTextEncoderModelConfigurator,
-            module_ops_from_gemma_root,
-        )
-    except ImportError as exc:
-        pytest.skip(f"LTX-2 import failed: {exc}")
-
-    device = torch.device("cuda:0")
-    precision = torch.bfloat16
-
-    reference_builder = SingleGPUModelBuilder(
-        model_path=str(official_path),
-        model_class_configurator=AVGemmaTextEncoderModelConfigurator,
-        model_sd_ops=AV_GEMMA_TEXT_ENCODER_KEY_OPS,
-        module_ops=module_ops_from_gemma_root(str(gemma_path)),
-    )
-    reference_encoder = reference_builder.build(
-        device=device, dtype=precision
-    ).to(device=device, dtype=precision)
-    if hasattr(reference_encoder.model, "config"):
-        if hasattr(reference_encoder.model.config, "attn_implementation"):
-            reference_encoder.model.config.attn_implementation = "sdpa"
-        if hasattr(reference_encoder.model.config, "_attn_implementation"):
-            reference_encoder.model.config._attn_implementation = "sdpa"
-    for module in reference_encoder.modules():
-        if isinstance(module, Attention):
-            module.attention_function = AttentionFunction.PYTORCH
-    reference_encoder.eval()
-    _attach_encoder_logging(reference_encoder, reference_log, "reference")
-    _log_register_sums(reference_encoder, reference_log, "reference")
-    _log_param_sums(reference_encoder, reference_log, "reference")
-    _log_gemma_param_sums(reference_encoder.model, reference_log, "reference")
-    _log_line(
-        reference_log,
-        "reference:gemma_config:attn_impl="
-        f"{getattr(reference_encoder.model.config, 'attn_implementation', None)} "
-        f"dtype={reference_encoder.model.dtype}",
-    )
-
-    diffusers_config = get_diffusers_config(model=str(text_encoder_path))
-    encoder_config = LTX2GemmaConfig()
-    encoder_config.update_model_arch(diffusers_config)
-    encoder_config.arch_config.gemma_model_path = str(gemma_path)
-    fastvideo_encoder = LTX2GemmaTextEncoderModel(encoder_config).to(
-        device=device, dtype=precision
-    )
-    if hasattr(fastvideo_encoder.gemma_model, "config"):
-        if hasattr(fastvideo_encoder.gemma_model.config, "attn_implementation"):
-            fastvideo_encoder.gemma_model.config.attn_implementation = "sdpa"
-        if hasattr(fastvideo_encoder.gemma_model.config, "_attn_implementation"):
-            fastvideo_encoder.gemma_model.config._attn_implementation = "sdpa"
-
-    official_weights = load_file(str(official_path))
-    fastvideo_weights: dict[str, torch.Tensor] = {}
-    for name, tensor in official_weights.items():
-        mapped_name = AV_GEMMA_TEXT_ENCODER_KEY_OPS.apply_to_key(name)
-        if mapped_name is None:
-            continue
-        fastvideo_weights[mapped_name] = tensor
-    fastvideo_encoder.load_weights(fastvideo_weights.items())
-    fastvideo_encoder.eval()
-    _attach_encoder_logging(fastvideo_encoder, fastvideo_log, "fastvideo")
-    _log_register_sums(fastvideo_encoder, fastvideo_log, "fastvideo")
-    _log_param_sums(fastvideo_encoder, fastvideo_log, "fastvideo")
-    _log_gemma_param_sums(fastvideo_encoder.gemma_model, fastvideo_log, "fastvideo")
-    _log_line(
-        fastvideo_log,
-        "fastvideo:gemma_config:attn_impl="
-        f"{getattr(fastvideo_encoder.gemma_model.config, 'attn_implementation', None)} "
-        f"dtype={fastvideo_encoder.gemma_model.dtype}",
-    )
-
-    prompt = "A curious raccoon peers through a vibrant field of yellow sunflowers."
-    ref_tokenizer = reference_encoder.tokenizer
-    if ref_tokenizer is None:
-        pytest.skip("Reference tokenizer is not initialized.")
-    token_pairs = ref_tokenizer.tokenize_with_weights(prompt)["gemma"]
-    input_ids = torch.tensor([[t[0] for t in token_pairs]], device=device)
-    attention_mask = torch.tensor([[w[1] for w in token_pairs]], device=device)
-
-    with torch.no_grad(), torch.backends.cuda.sdp_kernel(
-        enable_flash=False,
-        enable_mem_efficient=False,
-        enable_math=True,
-    ):
-        ref_outputs = reference_encoder.model(
-            input_ids=input_ids,
-            attention_mask=attention_mask,
-            output_hidden_states=True,
-            return_dict=True,
-            use_cache=False,
-        )
-        _log_line(
-            reference_log,
-            "reference:gemma_hidden_last:sum="
-            f"{ref_outputs.hidden_states[-1].float().sum().item():.6f}",
-        )
-        ref_projected = reference_encoder._run_feature_extractor(
-            ref_outputs.hidden_states,
-            attention_mask=attention_mask,
-            padding_side="left",
-        )
-        # Compare rotary embeddings between implementations.
-        from ltx_core.model.transformer.rope import precompute_freqs_cis
-        from fastvideo.models.dits.ltx2 import precompute_ltx_freqs_cis
-
-        seq_len = ref_projected.shape[1]
-        indices_grid = torch.arange(seq_len, device=device, dtype=torch.float32)[None, None, :]
-        ref_cos, ref_sin = precompute_freqs_cis(
-            indices_grid=indices_grid,
-            dim=reference_encoder.embeddings_connector.inner_dim,
-            out_dtype=ref_projected.dtype,
-            theta=reference_encoder.embeddings_connector.positional_embedding_theta,
-            max_pos=reference_encoder.embeddings_connector.positional_embedding_max_pos,
-            num_attention_heads=reference_encoder.embeddings_connector.num_attention_heads,
-            rope_type=reference_encoder.embeddings_connector.rope_type,
-        )
-        fast_cos, fast_sin = precompute_ltx_freqs_cis(
-            indices_grid=indices_grid,
-            dim=fastvideo_encoder.embeddings_connector.inner_dim,
-            out_dtype=ref_projected.dtype,
-            theta=fastvideo_encoder.embeddings_connector.positional_embedding_theta,
-            max_pos=fastvideo_encoder.embeddings_connector.positional_embedding_max_pos,
-            num_attention_heads=fastvideo_encoder.embeddings_connector.num_attention_heads,
-            rope_type=fastvideo_encoder.embeddings_connector.rope_type,
-        )
-        _log_line(
-            reference_log,
-            "reference:rope:cos_sum="
-            f"{ref_cos.float().sum().item():.6f} sin_sum={ref_sin.float().sum().item():.6f} "
-            f"theta={reference_encoder.embeddings_connector.positional_embedding_theta} "
-            f"max_pos={reference_encoder.embeddings_connector.positional_embedding_max_pos} "
-            f"rope_type={reference_encoder.embeddings_connector.rope_type}"
-        )
-        _log_line(
-            fastvideo_log,
-            "fastvideo:rope:cos_sum="
-            f"{fast_cos.float().sum().item():.6f} sin_sum={fast_sin.float().sum().item():.6f} "
-            f"theta={fastvideo_encoder.embeddings_connector.positional_embedding_theta} "
-            f"max_pos={fastvideo_encoder.embeddings_connector.positional_embedding_max_pos} "
-            f"rope_type={fastvideo_encoder.embeddings_connector.rope_type}"
-        )
-        ref_video, ref_audio, _ = reference_encoder._run_connectors(
-            ref_projected, attention_mask
-        )
-        fast_video_from_ref, fast_audio_from_ref, _ = fastvideo_encoder._run_connectors(
-            ref_projected, attention_mask
-        )
-        _log_line(
-            fastvideo_log,
-            "fastvideo:connector_on_ref:video_sum="
-            f"{fast_video_from_ref.float().sum().item():.6f}",
-        )
-        _log_line(
-            fastvideo_log,
-            "fastvideo:connector_on_ref:audio_sum="
-            f"{fast_audio_from_ref.float().sum().item():.6f}",
-        )
-
-        fast_outputs = fastvideo_encoder.gemma_model(
-            input_ids=input_ids,
-            attention_mask=attention_mask,
-            output_hidden_states=True,
-            return_dict=True,
-            use_cache=False,
-        )
-        _log_line(
-            fastvideo_log,
-            "fastvideo:gemma_hidden_last:sum="
-            f"{fast_outputs.hidden_states[-1].float().sum().item():.6f}",
-        )
-        fast_projected = fastvideo_encoder._run_feature_extractor(
-            fast_outputs.hidden_states,
-            attention_mask=attention_mask,
-            padding_side=fastvideo_encoder.padding_side,
-        )
-        fast_video, fast_audio, _ = fastvideo_encoder._run_connectors(
-            fast_projected, attention_mask
-        )
-
-    assert ref_video.shape == fast_video.shape
-    assert ref_audio.shape == fast_audio.shape
-    assert torch.isfinite(ref_video).all(), "Reference Gemma produced non-finite video embeddings."
-    assert torch.isfinite(ref_audio).all(), "Reference Gemma produced non-finite audio embeddings."
-    assert torch.isfinite(fast_video).all(), "FastVideo Gemma produced non-finite video embeddings."
-    assert torch.isfinite(fast_audio).all(), "FastVideo Gemma produced non-finite audio embeddings."
-    assert_close(ref_video, fast_video, atol=3e-1, rtol=5e-2)
-    assert_close(ref_audio, fast_audio, atol=3e-1, rtol=5e-2)

tests/local_tests/pipelines/test_ltx2_pipeline_smoke.py

@@ -1,296 +0,0 @@
-# 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
-
-from fastvideo import VideoGenerator
-from fastvideo.models.dits.ltx2 import (
-    AudioLatentShape,
-    DEFAULT_LTX2_AUDIO_CHANNELS,
-    DEFAULT_LTX2_AUDIO_DOWNSAMPLE,
-    DEFAULT_LTX2_AUDIO_HOP_LENGTH,
-    DEFAULT_LTX2_AUDIO_MEL_BINS,
-    DEFAULT_LTX2_AUDIO_SAMPLE_RATE,
-)
-from fastvideo.models.loader.component_loader import PipelineComponentLoader
-
-
-def _log_tensor_stats(label: str, tensor: torch.Tensor) -> None:
-    tensor_f32 = tensor.float()
-    print(
-        f"[LTX2 SMOKE] {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}"
-    )
-
-def _truncate_debug_logs() -> None:
-    for env_var in (
-        "LTX2_PIPELINE_DEBUG_PATH",
-        "LTX2_REFERENCE_DEBUG_PATH",
-        "LTX2_PIPELINE_DEBUG_DETAIL_PATH",
-        "LTX2_REFERENCE_DEBUG_DETAIL_PATH",
-    ):
-        log_path = os.getenv(env_var, "")
-        if not log_path:
-            continue
-        log_dir = os.path.dirname(log_path)
-        if log_dir:
-            os.makedirs(log_dir, exist_ok=True)
-        with open(log_path, "w", encoding="utf-8") as f:
-            f.write("")
-
-
-def _run_audio_decode_smoke(
-    diffusers_path: str,
-    fastvideo_args,
-    device: torch.device,
-    num_frames: int,
-    fps: float,
-) -> None:
-    audio_decoder_path = os.path.join(diffusers_path, "audio_vae")
-    vocoder_path = os.path.join(diffusers_path, "vocoder")
-    if not os.path.isdir(audio_decoder_path):
-        pytest.skip(f"Missing LTX-2 audio decoder at {audio_decoder_path}")
-    if not os.path.isdir(vocoder_path):
-        pytest.skip(f"Missing LTX-2 vocoder at {vocoder_path}")
-
-    audio_decoder = PipelineComponentLoader.load_module(
-        module_name="audio_decoder",
-        component_model_path=audio_decoder_path,
-        transformers_or_diffusers="diffusers",
-        fastvideo_args=fastvideo_args,
-    )
-    vocoder = PipelineComponentLoader.load_module(
-        module_name="vocoder",
-        component_model_path=vocoder_path,
-        transformers_or_diffusers="diffusers",
-        fastvideo_args=fastvideo_args,
-    )
-
-    duration = float(num_frames) / float(fps)
-    audio_shape = AudioLatentShape.from_duration(
-        batch=1,
-        duration=duration,
-        channels=DEFAULT_LTX2_AUDIO_CHANNELS,
-        mel_bins=DEFAULT_LTX2_AUDIO_MEL_BINS,
-        sample_rate=DEFAULT_LTX2_AUDIO_SAMPLE_RATE,
-        hop_length=DEFAULT_LTX2_AUDIO_HOP_LENGTH,
-        audio_latent_downsample_factor=DEFAULT_LTX2_AUDIO_DOWNSAMPLE,
-    )
-    audio_dtype = next(audio_decoder.parameters()).dtype
-    audio_latents = torch.randn(
-        audio_shape.to_torch_shape(),
-        device=device,
-        dtype=audio_dtype,
-    )
-
-    with torch.no_grad():
-        decoded_spec = audio_decoder(audio_latents)
-        audio_wave = vocoder(decoded_spec)
-
-    assert audio_wave.ndim == 3
-
-
-@pytest.mark.skipif(
-    not torch.cuda.is_available(),
-    reason="LTX-2 pipeline smoke test requires CUDA.",
-)
-def test_ltx2_pipeline_smoke():
-    repo_root = Path(__file__).resolve().parents[3]
-    debug_dir = repo_root / "ltx2_debug"
-    os.environ.setdefault(
-        "LTX2_PIPELINE_DEBUG_PATH",
-        str(debug_dir / "fastvideo_pipeline.log"),
-    )
-    os.environ.setdefault(
-        "LTX2_REFERENCE_DEBUG_PATH",
-        str(debug_dir / "reference_pipeline.log"),
-    )
-    os.environ.setdefault(
-        "LTX2_PIPELINE_DEBUG_DETAIL_PATH",
-        str(debug_dir / "fastvideo_pipeline_detail.log"),
-    )
-    os.environ.setdefault(
-        "LTX2_REFERENCE_DEBUG_DETAIL_PATH",
-        str(debug_dir / "reference_pipeline_detail.log"),
-    )
-    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)
-    _truncate_debug_logs()
-
-    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["PYTHONPATH"] = str(repo_root)
-
-    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",
-    )
-
-    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(f"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}")
-
-    try:
-        from ltx_pipelines.ti2vid_one_stage import TI2VidOneStagePipeline
-        from ltx_core.model.transformer import attention as ltx_attention
-    except ImportError as exc:
-        pytest.skip(f"LTX-2 pipeline import failed: {exc}")
-    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 = 96
-    num_frames = 9
-    fps = 12.0
-    steps = 4
-    guidance_scale = 4.0
-
-    with tempfile.TemporaryDirectory() as tmpdir:
-        latent_path = str(Path(tmpdir) / "ltx2_initial_latent.pt")
-
-        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_initial_latent_path=latent_path,
-        )
-        result = generator.generate_video(
-            prompt=prompt,
-            negative_prompt=negative_prompt,
-            output_path="outputs_video/ltx2_smoke",
-            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()
-        _run_audio_decode_smoke(
-            diffusers_path=diffusers_path,
-            fastvideo_args=generator.fastvideo_args,
-            device=device,
-            num_frames=num_frames,
-            fps=fps,
-        )
-
-        fastvideo_out = result["samples"]
-        fastvideo_out = fastvideo_out.to(device=device, dtype=torch.float32)
-        _log_tensor_stats("fastvideo_video", fastvideo_out)
-
-        ref_pipeline = TI2VidOneStagePipeline(
-            checkpoint_path=official_path,
-            gemma_root=gemma_model_path,
-            loras=[],
-            device=device,
-            fp8transformer=False,
-        )
-        original_text_encoder = ref_pipeline.model_ledger.text_encoder
-        original_transformer = ref_pipeline.model_ledger.transformer
-        original_video_decoder = ref_pipeline.model_ledger.video_decoder
-
-        def _patched_text_encoder():
-            encoder = original_text_encoder()
-            try:
-                from ltx_core.model.transformer.attention import (  # type: ignore
-                    Attention,
-                    AttentionFunction,
-                )
-            except ImportError:
-                return 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
-
-        ref_pipeline.model_ledger.text_encoder = _patched_text_encoder
-        if os.getenv("LTX2_DISABLE_VAE_NOISE", "1") == "1":
-            def _patched_video_decoder():
-                decoder = original_video_decoder()
-                if hasattr(decoder, "decode_noise_scale"):
-                    decoder.decode_noise_scale = 0.0
-                return decoder
-
-            ref_pipeline.model_ledger.video_decoder = _patched_video_decoder
-        if os.getenv("LTX2_DEBUG_DETAIL", "0") == "1":
-            from ..transformers.test_ltx2 import (
-                _attach_block_detail_logging,
-            )
-
-            def _patched_transformer():
-                model = original_transformer()
-                core = getattr(model, "velocity_model", model)
-                _attach_block_detail_logging(
-                    core,
-                    Path(os.environ["LTX2_REFERENCE_DEBUG_DETAIL_PATH"]),
-                    "reference",
-                    True,
-                )
-                return model
-
-            ref_pipeline.model_ledger.transformer = _patched_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,
-                initial_video_latent_path=latent_path,
-            )
-            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)
-        _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)

tests/local_tests/transformers/test_ltx2.py

@@ -1,343 +0,0 @@
-# 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")
-
-repo_root = Path(__file__).resolve().parents[3]
-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 PipelineConfig
-from fastvideo.fastvideo_args import FastVideoArgs
-from fastvideo.forward_context import set_forward_context
-from fastvideo.models.loader.component_loader import TransformerLoader
-
-
-def _read_transformer_config(config: dict) -> dict:
-    transformer_config = config.get("transformer", {})
-    if not transformer_config:
-        raise ValueError("Missing transformer config in LTX-2 metadata.")
-    return transformer_config
-
-
-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):
-        if in_channels % candidate != 0:
-            continue
-        patch_volume = in_channels // candidate
-        root = int(round(patch_volume**0.5))
-        if root * root == patch_volume:
-            patch_size = root
-            num_channels_latents = candidate
-            break
-    print(
-        f"[LTX2 TEST] Inferred patch_size={patch_size}, "
-        f"num_channels_latents={num_channels_latents}"
-    )
-    return patch_size, num_channels_latents
-
-
-def _attach_block_sum_logging(
-    model: torch.nn.Module,
-    log_path: Path,
-    label: str,
-    enabled: bool,
-) -> None:
-    if not enabled:
-        return
-
-    log_path.parent.mkdir(parents=True, exist_ok=True)
-    if log_path.exists():
-        log_path.unlink()
-
-    def _format_sum(tensor: torch.Tensor | None) -> str:
-        if tensor is None:
-            return "None"
-        return f"{tensor.float().sum().item():.6f}"
-
-    def _hook(module, inputs, outputs):  # noqa: ANN001
-        if isinstance(outputs, tuple):
-            video_args, audio_args = outputs
-            video_sum = _format_sum(video_args.x if video_args is not None else None)
-            audio_sum = _format_sum(audio_args.x if audio_args is not None else None)
-        else:
-            video_sum = _format_sum(outputs)
-            audio_sum = "None"
-        with log_path.open("a", encoding="utf-8") as f:
-            f.write(f"{label}:{module.idx}:video_sum={video_sum},audio_sum={audio_sum}\n")
-
-    for block in model.transformer_blocks:
-        block.register_forward_hook(_hook)
-
-
-def _attach_block_detail_logging(
-    model: torch.nn.Module,
-    log_path: Path,
-    label: str,
-    enabled: bool,
-) -> None:
-    if not enabled:
-        return
-
-    log_path.parent.mkdir(parents=True, exist_ok=True)
-    if log_path.exists():
-        log_path.unlink()
-
-    def _format_sum(tensor: torch.Tensor | None) -> str:
-        if tensor is None:
-            return "None"
-        return f"{tensor.float().sum().item():.6f}"
-
-    def _hook_factory(block_idx: int, name: str):
-        def _hook(_module, _inputs, outputs):  # noqa: ANN001
-            if isinstance(outputs, tuple):
-                out = outputs[0]
-            else:
-                out = outputs
-            out_sum = _format_sum(out if torch.is_tensor(out) else None)
-            with log_path.open("a", encoding="utf-8") as f:
-                f.write(f"{label}:{block_idx}:{name}:out_sum={out_sum}\n")
-        return _hook
-
-    for block in model.transformer_blocks:
-        idx = block.idx
-        for name in (
-            "attn1",
-            "attn2",
-            "ff",
-            "audio_attn1",
-            "audio_attn2",
-            "audio_ff",
-            "audio_to_video_attn",
-            "video_to_audio_attn",
-        ):
-            if hasattr(block, name):
-                getattr(block, name).register_forward_hook(_hook_factory(idx, name))
-
-    def _output_hook(name: str):
-        def _hook(_module, _inputs, outputs):  # noqa: ANN001
-            out = outputs[0] if isinstance(outputs, tuple) else outputs
-            out_sum = _format_sum(out if torch.is_tensor(out) else None)
-            with log_path.open("a", encoding="utf-8") as f:
-                f.write(f"{label}:output:{name}:out_sum={out_sum}\n")
-        return _hook
-
-    for name in ("proj_out", "audio_proj_out"):
-        if hasattr(model, name):
-            getattr(model, name).register_forward_hook(_output_hook(name))
-
-
-def test_ltx2_transformer_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",
-        )
-    )
-    fastvideo_path = Path(
-        os.getenv(
-            "LTX2_FASTVIDEO_PATH",
-            str(diffusers_root / "transformer"),
-        )
-    )
-    if not official_path.exists():
-        pytest.skip(f"LTX-2 official weights not found at {official_path}")
-    if not fastvideo_path.exists():
-        pytest.skip(f"FastVideo converted weights not found at {fastvideo_path}")
-
-    try:
-        from ltx_core.components.patchifiers import VideoLatentPatchifier
-        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
-    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 transformer parity test requires CUDA for attention backends.")
-
-    device = torch.device("cuda:0")
-    precision = torch.bfloat16
-    precision_str = "bf16"
-
-    args = FastVideoArgs(
-        model_path=str(fastvideo_path),
-        dit_cpu_offload=True,
-        use_fsdp_inference=False,
-        pipeline_config=PipelineConfig(dit_config=config, dit_precision=precision_str),
-    )
-    args.device = device
-
-    loader = TransformerLoader()
-    fastvideo_model = loader.load(str(fastvideo_path), args).to(device=device, dtype=precision)
-
-    reference_builder = SingleGPUModelBuilder(
-        model_class_configurator=LTXModelConfigurator,
-        model_path=str(official_path),
-        model_sd_ops=LTXV_MODEL_COMFY_RENAMING_MAP,
-    )
-    reference_model = reference_builder.build(
-        device=device, dtype=precision).to(device=device, dtype=precision)
-    reference_model.set_gradient_checkpointing(False)
-
-    fastvideo_model.eval()
-    reference_model.eval()
-
-    debug_logs = os.getenv("LTX2_DEBUG_LOGS", "0") == "1"
-    _attach_block_sum_logging(
-        fastvideo_model.model,
-        repo_root / "ltx2_debug" / "fastvideo.log",
-        "fastvideo",
-        debug_logs,
-    )
-    _attach_block_sum_logging(
-        reference_model,
-        repo_root / "ltx2_debug" / "reference.log",
-        "reference",
-        debug_logs,
-    )
-    _attach_block_detail_logging(
-        fastvideo_model.model,
-        repo_root / "ltx2_debug" / "fastvideo_detail.log",
-        "fastvideo",
-        os.getenv("LTX2_DEBUG_DETAIL", "0") == "1",
-    )
-    _attach_block_detail_logging(
-        reference_model,
-        repo_root / "ltx2_debug" / "reference_detail.log",
-        "reference",
-        os.getenv("LTX2_DEBUG_DETAIL", "0") == "1",
-    )
-
-    patchifier = VideoLatentPatchifier(patch_size=cfg.patch_size[1])
-    batch_size = 1
-    frames = 4
-    height = cfg.patch_size[1] * 4
-    width = cfg.patch_size[2] * 4
-    hidden_states = torch.randn(
-        batch_size,
-        cfg.num_channels_latents,
-        frames,
-        height,
-        width,
-        device=device,
-        dtype=precision,
-    )
-    encoder_hidden_states = torch.randn(
-        batch_size,
-        16,
-        cfg.caption_channels,
-        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)
-    latents = patchifier.patchify(hidden_states)
-
-    video = Modality(
-        enabled=True,
-        latent=latents,
-        timesteps=timestep,
-        positions=positions,
-        context=encoder_hidden_states,
-        context_mask=None,
-    )
-
-    with torch.no_grad():
-        ref_out, _ = reference_model(
-            video=video,
-            audio=None,
-            perturbations=BatchedPerturbationConfig.empty(batch_size),
-        )
-        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(
-            current_timestep=0,
-            attn_metadata=None,
-            forward_batch=None,
-        ):
-            fastvideo_out = fastvideo_model(
-                hidden_states=hidden_states,
-                encoder_hidden_states=encoder_hidden_states,
-                timestep=timestep,
-            )
-            print(f"[LTX2 TEST] FastVideo model output shape: {fastvideo_out.shape}")
-    assert ref_out.shape == fastvideo_out.shape
-    assert ref_out.dtype == fastvideo_out.dtype
-    assert_close(ref_out, fastvideo_out, atol=1e-4, rtol=1e-4)

tests/local_tests/transformers/test_ltx2_audio.py

@@ -1,280 +0,0 @@
-# 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")
-# Force TORCH_SDPA backend for parity testing - both FastVideo and LTX-2 reference
-# will use PyTorch's scaled_dot_product_attention for consistent results
-os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "TORCH_SDPA")
-
-repo_root = Path(__file__).resolve().parents[3]
-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 PipelineConfig
-from fastvideo.fastvideo_args import FastVideoArgs
-from fastvideo.forward_context import set_forward_context
-from fastvideo.models.dits.ltx2 import Modality as FastVideoModality
-from fastvideo.models.loader.component_loader import TransformerLoader
-from .test_ltx2 import (
-    _attach_block_detail_logging,
-    _attach_block_sum_logging,
-    _infer_patch_params,
-    _read_transformer_config,
-)
-
-
-def test_ltx2_transformer_audio_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",
-        )
-    )
-    fastvideo_path = Path(
-        os.getenv(
-            "LTX2_FASTVIDEO_PATH",
-            str(diffusers_root / "transformer"),
-        )
-    )
-    if not official_path.exists():
-        pytest.skip(f"LTX-2 official weights not found at {official_path}")
-    if not fastvideo_path.exists():
-        pytest.skip(f"FastVideo converted weights not found at {fastvideo_path}")
-
-    try:
-        from ltx_core.components.patchifiers import AudioPatchifier, VideoLatentPatchifier
-        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
-    except ImportError as exc:
-        pytest.skip(f"LTX-2 import failed: {exc}")
-
-    # Load config from metadata using same approach as test_ltx2.py
-    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 transformer parity test requires CUDA for attention backends.")
-
-    device = torch.device("cuda:0")
-    precision = torch.bfloat16
-    precision_str = "bf16"
-
-    args = FastVideoArgs(
-        model_path=str(fastvideo_path),
-        dit_cpu_offload=True,
-        use_fsdp_inference=False,
-        pipeline_config=PipelineConfig(dit_config=config, dit_precision=precision_str),
-    )
-    args.device = device
-
-    loader = TransformerLoader()
-    fastvideo_model = loader.load(str(fastvideo_path), args).to(device=device, dtype=precision)
-
-    # Use SingleGPUModelBuilder to load the reference model (same as test_ltx2.py)
-    reference_builder = SingleGPUModelBuilder(
-        model_class_configurator=LTXModelConfigurator,
-        model_path=str(official_path),
-        model_sd_ops=LTXV_MODEL_COMFY_RENAMING_MAP,
-    )
-    reference_model = reference_builder.build(
-        device=device, dtype=precision).to(device=device, dtype=precision)
-    reference_model.set_gradient_checkpointing(False)
-
-    fastvideo_model.eval()
-    reference_model.eval()
-
-    debug_logs = os.getenv("LTX2_DEBUG_LOGS", "0") == "1"
-    _attach_block_sum_logging(
-        fastvideo_model.model,
-        repo_root / "ltx2_debug" / "fastvideo_audio.log",
-        "fastvideo",
-        debug_logs,
-    )
-    _attach_block_sum_logging(
-        reference_model,
-        repo_root / "ltx2_debug" / "reference_audio.log",
-        "reference",
-        debug_logs,
-    )
-    _attach_block_detail_logging(
-        fastvideo_model.model,
-        repo_root / "ltx2_debug" / "fastvideo_audio_detail.log",
-        "fastvideo",
-        os.getenv("LTX2_DEBUG_DETAIL", "0") == "1",
-    )
-    _attach_block_detail_logging(
-        reference_model,
-        repo_root / "ltx2_debug" / "reference_audio_detail.log",
-        "reference",
-        os.getenv("LTX2_DEBUG_DETAIL", "0") == "1",
-    )
-
-    patchifier = VideoLatentPatchifier(patch_size=cfg.patch_size[1])
-    audio_patchifier = AudioPatchifier(patch_size=1)
-    batch_size = 1
-    frames = 4
-    height = cfg.patch_size[1] * 4
-    width = cfg.patch_size[2] * 4
-    hidden_states = torch.randn(
-        batch_size,
-        cfg.num_channels_latents,
-        frames,
-        height,
-        width,
-        device=device,
-        dtype=precision,
-    )
-    encoder_hidden_states = torch.randn(
-        batch_size,
-        16,
-        cfg.caption_channels,
-        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)
-    latents = patchifier.patchify(hidden_states)
-
-    audio_frames = 16
-    audio_channels = 8
-    audio_mel_bins = 16
-    audio_latents = torch.randn(
-        batch_size,
-        audio_channels,
-        audio_frames,
-        audio_mel_bins,
-        device=device,
-        dtype=precision,
-    )
-    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)
-
-    video = Modality(
-        enabled=True,
-        latent=latents,
-        timesteps=timestep,
-        positions=positions,
-        context=encoder_hidden_states,
-        context_mask=None,
-    )
-    audio = Modality(
-        enabled=True,
-        latent=audio_tokens,
-        timesteps=timestep,
-        positions=audio_positions,
-        context=encoder_hidden_states,
-        context_mask=None,
-    )
-
-    fastvideo_video = FastVideoModality(
-        enabled=True,
-        latent=latents,
-        timesteps=timestep,
-        positions=positions,
-        context=encoder_hidden_states,
-        context_mask=None,
-    )
-    fastvideo_audio = FastVideoModality(
-        enabled=True,
-        latent=audio_tokens,
-        timesteps=timestep,
-        positions=audio_positions,
-        context=encoder_hidden_states,
-        context_mask=None,
-    )
-
-    with torch.no_grad():
-        _, ref_audio_out = reference_model(
-            video=video,
-            audio=audio,
-            perturbations=BatchedPerturbationConfig.empty(batch_size),
-        )
-        ref_audio_out = audio_patchifier.unpatchify(ref_audio_out, output_shape=audio_shape)
-        with set_forward_context(
-            current_timestep=0,
-            attn_metadata=None,
-            forward_batch=None,
-        ):
-            _, fastvideo_audio_out = fastvideo_model.model(
-                video=fastvideo_video,
-                audio=fastvideo_audio,
-            )
-            fastvideo_audio_out = audio_patchifier.unpatchify(fastvideo_audio_out, output_shape=audio_shape)
-
-    assert ref_audio_out.shape == fastvideo_audio_out.shape
-    assert ref_audio_out.dtype == fastvideo_audio_out.dtype
-    # With TORCH_SDPA backend for both, use same tolerance as video parity test
-    assert_close(ref_audio_out, fastvideo_audio_out, atol=1e-4, rtol=1e-4)

tests/local_tests/vaes/test_ltx2_audio_vae.py

@@ -1,272 +0,0 @@
-# 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]
-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.models.audio.ltx2_audio_vae import (
-    LTX2AudioDecoder,
-    LTX2AudioEncoder,
-    LTX2Vocoder,
-)
-
-
-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 _load_weights(path: Path) -> dict[str, torch.Tensor]:
-    print(f"[LTX2 AUDIO VAE TEST] Loading weights from {path}")
-    return load_file(str(path))
-
-
-def _select_audio_vae_weights(
-    weights: dict[str, torch.Tensor], prefix: str
-) -> dict[str, torch.Tensor]:
-    filtered: dict[str, torch.Tensor] = {}
-    alt_prefix = prefix.replace("audio_vae.", "")
-    for name, tensor in weights.items():
-        if name.startswith(prefix):
-            filtered[name.replace(prefix, "")] = tensor
-        elif alt_prefix and name.startswith(alt_prefix):
-            filtered[name.replace(alt_prefix, "")] = tensor
-        elif name.startswith("audio_vae.per_channel_statistics."):
-            filtered[name.replace("audio_vae.", "")] = tensor
-        elif name.startswith("per_channel_statistics."):
-            filtered[name] = tensor
-    print(f"[LTX2 AUDIO VAE TEST] Selected {len(filtered)} tensors for {prefix}")
-    return filtered
-
-
-def _select_vocoder_weights(
-    weights: dict[str, torch.Tensor]
-) -> dict[str, torch.Tensor]:
-    if any(name.startswith("vocoder.") for name in weights):
-        filtered = {
-            name.replace("vocoder.", ""): tensor
-            for name, tensor in weights.items()
-            if name.startswith("vocoder.")
-        }
-    else:
-        filtered = dict(weights)
-    print(f"[LTX2 AUDIO VAE TEST] Selected {len(filtered)} tensors for vocoder.")
-    return filtered
-
-
-def _load_into_model(
-    model: torch.nn.Module, weights: dict[str, torch.Tensor]
-) -> tuple[int, list[str]]:
-    model_state = model.state_dict()
-    filtered = {
-        k: v
-        for k, v in weights.items()
-        if k in model_state and model_state[k].shape == v.shape
-    }
-    missing = [k for k in model_state.keys() if k not in filtered]
-    print(
-        f"[LTX2 AUDIO VAE TEST] Loading {len(filtered)} / {len(model_state)} tensors "
-        f"from {len(weights)} available"
-    )
-    if not filtered:
-        return 0, missing
-    model.load_state_dict(filtered, strict=False)
-    return len(filtered), missing
-
-
-def test_ltx2_audio_vae_vocoder_parity():
-    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",
-        )
-    )
-    audio_vae_path = Path(
-        os.getenv("LTX2_AUDIO_VAE_PATH", str(diffusers_root / "audio_vae"))
-    )
-    vocoder_path = Path(
-        os.getenv("LTX2_VOCODER_PATH", str(diffusers_root / "vocoder"))
-    )
-    if not official_path.exists():
-        pytest.skip(f"LTX-2 weights not found at {official_path}")
-    if not audio_vae_path.exists():
-        pytest.skip(f"LTX-2 audio VAE weights not found at {audio_vae_path}")
-    if not vocoder_path.exists():
-        pytest.skip(f"LTX-2 vocoder weights not found at {vocoder_path}")
-
-    config = _load_metadata(official_path)
-    if "audio_vae" not in config or "vocoder" not in config:
-        pytest.skip("Audio VAE or vocoder config not found in safetensors metadata.")
-
-    try:
-        from ltx_core.model.audio_vae import (
-            AudioDecoderConfigurator,
-            AudioEncoderConfigurator,
-            VocoderConfigurator,
-        )
-    except ImportError as exc:
-        pytest.skip(f"LTX-2 import failed: {exc}")
-
-    ref_weights = _load_weights(official_path)
-    encoder_weights = _select_audio_vae_weights(ref_weights, "audio_vae.encoder.")
-    decoder_weights = _select_audio_vae_weights(ref_weights, "audio_vae.decoder.")
-    vocoder_weights = _select_vocoder_weights(ref_weights)
-    if not encoder_weights or not decoder_weights or not vocoder_weights:
-        pytest.skip("Audio VAE or vocoder weights not found in safetensors file.")
-
-    fastvideo_audio_weights_path = audio_vae_path / "model.safetensors"
-    fastvideo_vocoder_weights_path = vocoder_path / "model.safetensors"
-    if not fastvideo_audio_weights_path.exists():
-        pytest.skip(
-            f"FastVideo audio VAE weights not found at {fastvideo_audio_weights_path}"
-        )
-    if not fastvideo_vocoder_weights_path.exists():
-        pytest.skip(
-            f"FastVideo vocoder weights not found at {fastvideo_vocoder_weights_path}"
-        )
-    fastvideo_audio_weights = _load_weights(fastvideo_audio_weights_path)
-    fastvideo_encoder_weights = _select_audio_vae_weights(
-        fastvideo_audio_weights, "encoder."
-    )
-    fastvideo_decoder_weights = _select_audio_vae_weights(
-        fastvideo_audio_weights, "decoder."
-    )
-    fastvideo_vocoder_weights = _select_vocoder_weights(
-        _load_weights(fastvideo_vocoder_weights_path)
-    )
-    if (not fastvideo_encoder_weights or not fastvideo_decoder_weights
-            or not fastvideo_vocoder_weights):
-        pytest.skip("FastVideo audio VAE/vocoder weights not found in diffusers files.")
-
-    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-    precision = torch.bfloat16 if torch.cuda.is_available() else torch.float32
-
-    fastvideo_encoder = LTX2AudioEncoder(config).to(device=device, dtype=precision)
-    fastvideo_decoder = LTX2AudioDecoder(config).to(device=device, dtype=precision)
-    fastvideo_vocoder = LTX2Vocoder(config).to(device=device, dtype=precision)
-
-    ref_encoder = AudioEncoderConfigurator.from_config(config).to(
-        device=device, dtype=precision
-    )
-    ref_decoder = AudioDecoderConfigurator.from_config(config).to(
-        device=device, dtype=precision
-    )
-    ref_vocoder = VocoderConfigurator.from_config(config).to(
-        device=device, dtype=precision
-    )
-
-    loaded_fastvideo_encoder, missing_fastvideo_encoder = _load_into_model(
-        fastvideo_encoder.model, fastvideo_encoder_weights
-    )
-    loaded_ref_encoder, missing_ref_encoder = _load_into_model(
-        ref_encoder, encoder_weights
-    )
-    loaded_fastvideo_decoder, missing_fastvideo_decoder = _load_into_model(
-        fastvideo_decoder.model, fastvideo_decoder_weights
-    )
-    loaded_ref_decoder, missing_ref_decoder = _load_into_model(
-        ref_decoder, decoder_weights
-    )
-    loaded_fastvideo_vocoder, missing_fastvideo_vocoder = _load_into_model(
-        fastvideo_vocoder.model, fastvideo_vocoder_weights
-    )
-    loaded_ref_vocoder, missing_ref_vocoder = _load_into_model(
-        ref_vocoder, vocoder_weights
-    )
-
-    if min(
-        loaded_fastvideo_encoder,
-        loaded_ref_encoder,
-        loaded_fastvideo_decoder,
-        loaded_ref_decoder,
-        loaded_fastvideo_vocoder,
-        loaded_ref_vocoder,
-    ) == 0:
-        pytest.skip("Failed to load audio VAE or vocoder weights.")
-    if (
-        missing_fastvideo_encoder
-        or missing_ref_encoder
-        or missing_fastvideo_decoder
-        or missing_ref_decoder
-        or missing_fastvideo_vocoder
-        or missing_ref_vocoder
-    ):
-        print(
-            f"[LTX2 AUDIO VAE TEST] Missing encoder keys: {len(missing_fastvideo_encoder)}"
-        )
-        print(
-            f"[LTX2 AUDIO VAE TEST] Missing decoder keys: {len(missing_fastvideo_decoder)}"
-        )
-        print(
-            f"[LTX2 AUDIO VAE TEST] Missing vocoder keys: {len(missing_fastvideo_vocoder)}"
-        )
-        pytest.skip("Missing audio VAE/vocoder keys; cannot ensure parity.")
-
-    fastvideo_encoder.model.eval()
-    fastvideo_decoder.model.eval()
-    fastvideo_vocoder.model.eval()
-    ref_encoder.eval()
-    ref_decoder.eval()
-    ref_vocoder.eval()
-
-    ddconfig = config["audio_vae"]["model"]["params"]["ddconfig"]
-    in_channels = ddconfig.get("in_channels", 2)
-    resolution = ddconfig.get("resolution", 256)
-    mel_bins = ddconfig.get("mel_bins", 64)
-    batch_size = 1
-    spectrogram = torch.randn(
-        batch_size,
-        in_channels,
-        resolution,
-        mel_bins,
-        device=device,
-        dtype=precision,
-    )
-
-    with torch.no_grad():
-        ref_latents = ref_encoder(spectrogram)
-        fast_latents = fastvideo_encoder(spectrogram)
-
-    assert ref_latents.shape == fast_latents.shape
-    assert ref_latents.dtype == fast_latents.dtype
-    assert torch.isfinite(ref_latents).all(), "Reference encoder produced non-finite latents."
-    assert torch.isfinite(fast_latents).all(), "FastVideo encoder produced non-finite latents."
-    assert_close(ref_latents, fast_latents, atol=1e-2, rtol=1e-2)
-
-    with torch.no_grad():
-        ref_decoded = ref_decoder(ref_latents)
-        fast_decoded = fastvideo_decoder(ref_latents)
-
-    assert ref_decoded.shape == fast_decoded.shape
-    assert ref_decoded.dtype == fast_decoded.dtype
-    assert torch.isfinite(ref_decoded).all(), "Reference decoder produced non-finite output."
-    assert torch.isfinite(fast_decoded).all(), "FastVideo decoder produced non-finite output."
-    assert_close(ref_decoded, fast_decoded, atol=1e-2, rtol=1e-2)
-
-    with torch.no_grad():
-        ref_audio = ref_vocoder(ref_decoded)
-        fast_audio = fastvideo_vocoder(ref_decoded)
-
-    assert ref_audio.shape == fast_audio.shape
-    assert ref_audio.dtype == fast_audio.dtype
-    assert torch.isfinite(ref_audio).all(), "Reference vocoder produced non-finite audio."
-    assert torch.isfinite(fast_audio).all(), "FastVideo vocoder produced non-finite audio."
-    assert_close(ref_audio, fast_audio, atol=1e-2, rtol=1e-2)

tests/local_tests/vaes/test_ltx2_vae.py

@@ -1,190 +0,0 @@
-# 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]
-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.models.vaes.ltx2vae import LTX2VideoDecoder, LTX2VideoEncoder
-
-
-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 _load_weights(path: Path) -> dict[str, torch.Tensor]:
-    print(f"[LTX2 VAE TEST] Loading weights from {path}")
-    return load_file(str(path))
-
-
-def _select_vae_weights(weights: dict[str, torch.Tensor], prefix: str) -> dict[str, torch.Tensor]:
-    filtered: dict[str, torch.Tensor] = {}
-    alt_prefix = prefix.replace("vae.", "")
-    for name, tensor in weights.items():
-        if name.startswith(prefix):
-            filtered[name.replace(prefix, "")] = tensor
-        elif alt_prefix and name.startswith(alt_prefix):
-            filtered[name.replace(alt_prefix, "")] = tensor
-        elif name.startswith("vae.per_channel_statistics."):
-            filtered[name.replace("vae.", "")] = tensor
-        elif name.startswith("per_channel_statistics."):
-            filtered[name] = tensor
-    print(f"[LTX2 VAE TEST] Selected {len(filtered)} tensors for {prefix}")
-    return filtered
-
-
-def _load_into_model(model: torch.nn.Module, weights: dict[str, torch.Tensor]) -> tuple[int, list[str]]:
-    model_state = model.state_dict()
-    filtered = {
-        k: v
-        for k, v in weights.items()
-        if k in model_state and model_state[k].shape == v.shape
-    }
-    missing = [k for k in model_state.keys() if k not in filtered]
-    print(
-        f"[LTX2 VAE TEST] Loading {len(filtered)} / {len(model_state)} tensors "
-        f"from {len(weights)} available"
-    )
-    if not filtered:
-        return 0, missing
-    model.load_state_dict(filtered, strict=False)
-    return len(filtered), missing
-
-
-def test_ltx2_vae_parity():
-    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",
-        )
-    )
-    fastvideo_path = Path(
-        os.getenv("LTX2_VAE_PATH", str(diffusers_root / "vae"))
-    )
-    if not official_path.exists():
-        pytest.skip(f"LTX-2 weights not found at {official_path}")
-    if not fastvideo_path.exists():
-        pytest.skip(f"LTX-2 diffusers VAE not found at {fastvideo_path}")
-
-    config = _load_metadata(official_path)
-    if "vae" not in config:
-        pytest.skip("VAE config not found in safetensors metadata.")
-
-    try:
-        from ltx_core.model.video_vae import VideoDecoderConfigurator, VideoEncoderConfigurator
-    except ImportError as exc:
-        pytest.skip(f"LTX-2 import failed: {exc}")
-
-    ref_weights = _load_weights(official_path)
-    encoder_weights = _select_vae_weights(ref_weights, "vae.encoder.")
-    decoder_weights = _select_vae_weights(ref_weights, "vae.decoder.")
-    if not encoder_weights or not decoder_weights:
-        pytest.skip("VAE weights not found in safetensors file.")
-
-    fastvideo_weights_path = fastvideo_path / "model.safetensors"
-    if not fastvideo_weights_path.exists():
-        pytest.skip(f"FastVideo VAE weights not found at {fastvideo_weights_path}")
-    fastvideo_weights = _load_weights(fastvideo_weights_path)
-    fastvideo_encoder_weights = _select_vae_weights(
-        fastvideo_weights, "encoder."
-    )
-    fastvideo_decoder_weights = _select_vae_weights(
-        fastvideo_weights, "decoder."
-    )
-    if not fastvideo_encoder_weights or not fastvideo_decoder_weights:
-        pytest.skip("FastVideo VAE weights not found in diffusers file.")
-
-    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-    precision = torch.bfloat16 if torch.cuda.is_available() else torch.float32
-
-    fastvideo_encoder = LTX2VideoEncoder(config).to(device=device, dtype=precision)
-    fastvideo_decoder = LTX2VideoDecoder(config).to(device=device, dtype=precision)
-    ref_encoder = VideoEncoderConfigurator.from_config(config).to(device=device, dtype=precision)
-    ref_decoder = VideoDecoderConfigurator.from_config(config).to(device=device, dtype=precision)
-
-    loaded_fastvideo_encoder, missing_fastvideo_encoder = _load_into_model(
-        fastvideo_encoder.model, fastvideo_encoder_weights
-    )
-    loaded_ref_encoder, missing_ref_encoder = _load_into_model(ref_encoder, encoder_weights)
-    loaded_fastvideo_decoder, missing_fastvideo_decoder = _load_into_model(
-        fastvideo_decoder.model, fastvideo_decoder_weights
-    )
-    loaded_ref_decoder, missing_ref_decoder = _load_into_model(ref_decoder, decoder_weights)
-
-    if min(
-        loaded_fastvideo_encoder,
-        loaded_ref_encoder,
-        loaded_fastvideo_decoder,
-        loaded_ref_decoder,
-    ) == 0:
-        pytest.skip("Failed to load VAE weights into one or more models.")
-    if (
-        missing_fastvideo_encoder
-        or missing_ref_encoder
-        or missing_fastvideo_decoder
-        or missing_ref_decoder
-    ):
-        print(f"[LTX2 VAE TEST] Missing encoder keys: {len(missing_fastvideo_encoder)}")
-        print(f"[LTX2 VAE TEST] Missing decoder keys: {len(missing_fastvideo_decoder)}")
-        pytest.skip("Missing VAE keys; cannot ensure parity.")
-
-    fastvideo_encoder.model.eval()
-    fastvideo_decoder.model.eval()
-    ref_encoder.eval()
-    ref_decoder.eval()
-
-    fastvideo_decoder.model.decode_noise_scale = 0.0
-    ref_decoder.decode_noise_scale = 0.0
-
-    batch_size = 1
-    frames = 9
-    height = 64
-    width = 64
-    video = torch.randn(
-        batch_size,
-        3,
-        frames,
-        height,
-        width,
-        device=device,
-        dtype=precision,
-    )
-
-    with torch.no_grad():
-        ref_latents = ref_encoder(video)
-        fast_latents = fastvideo_encoder(video)
-
-    assert ref_latents.shape == fast_latents.shape
-    assert ref_latents.dtype == fast_latents.dtype
-    assert torch.isfinite(ref_latents).all(), "Reference encoder produced non-finite latents."
-    assert torch.isfinite(fast_latents).all(), "FastVideo encoder produced non-finite latents."
-    assert_close(ref_latents, fast_latents, atol=1e-2, rtol=1e-2)
-
-    timestep = torch.tensor([0.05], device=device, dtype=precision)
-    with torch.no_grad():
-        ref_decoded = ref_decoder(ref_latents, timestep=timestep)
-        fast_decoded = fastvideo_decoder(fast_latents, timestep=timestep)
-
-    assert ref_decoded.shape == fast_decoded.shape
-    assert ref_decoded.dtype == fast_decoded.dtype
-    assert torch.isfinite(ref_decoded).all(), "Reference decoder produced non-finite output."
-    assert torch.isfinite(fast_decoded).all(), "FastVideo decoder produced non-finite output."
-    assert_close(ref_decoded, fast_decoded, atol=1e-2, rtol=1e-2)

tests/local_tests/vaes/test_ltx2_vae_official.py

@@ -1,132 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-import os
-from pathlib import Path
-import sys
-
-import pytest
-import torch
-from torch.testing import assert_close
-
-repo_root = Path(__file__).resolve().parents[3]
-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.vaes import LTX2VAEConfig
-from fastvideo.configs.pipelines import PipelineConfig
-from fastvideo.fastvideo_args import FastVideoArgs
-from fastvideo.models.loader.component_loader import VAELoader
-
-
-@pytest.mark.skipif(
-    not torch.cuda.is_available(),
-    reason="LTX-2 VAE parity test requires CUDA.",
-)
-def test_ltx2_vae_parity_official():
-    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",
-        )
-    )
-    fastvideo_path = Path(
-        os.getenv("LTX2_VAE_PATH", str(diffusers_root / "vae"))
-    )
-    if not official_path.exists():
-        pytest.skip(f"LTX-2 weights not found at {official_path}")
-    if not fastvideo_path.exists():
-        pytest.skip(f"LTX-2 diffusers VAE not found at {fastvideo_path}")
-
-    try:
-        from ltx_core.loader.single_gpu_model_builder import SingleGPUModelBuilder
-        from ltx_core.model.video_vae import (
-            VAE_DECODER_COMFY_KEYS_FILTER,
-            VAE_ENCODER_COMFY_KEYS_FILTER,
-            VideoDecoderConfigurator,
-            VideoEncoderConfigurator,
-        )
-    except ImportError as exc:
-        pytest.skip(f"LTX-2 import failed: {exc}")
-
-    device = torch.device("cuda:0")
-    precision = torch.bfloat16
-    precision_str = "bf16"
-
-    args = FastVideoArgs(
-        model_path=str(fastvideo_path),
-        vae_cpu_offload=False,
-        pipeline_config=PipelineConfig(
-            vae_config=LTX2VAEConfig(),
-            vae_precision=precision_str,
-        ),
-    )
-
-    loader = VAELoader()
-    fastvideo_vae = loader.load(str(fastvideo_path), args).to(
-        device=device, dtype=precision
-    )
-
-    encoder_builder = SingleGPUModelBuilder(
-        model_class_configurator=VideoEncoderConfigurator,
-        model_path=str(official_path),
-        model_sd_ops=VAE_ENCODER_COMFY_KEYS_FILTER,
-    )
-    decoder_builder = SingleGPUModelBuilder(
-        model_class_configurator=VideoDecoderConfigurator,
-        model_path=str(official_path),
-        model_sd_ops=VAE_DECODER_COMFY_KEYS_FILTER,
-    )
-    ref_encoder = encoder_builder.build(
-        device=device, dtype=precision
-    ).to(device=device, dtype=precision)
-    ref_decoder = decoder_builder.build(
-        device=device, dtype=precision
-    ).to(device=device, dtype=precision)
-
-    fastvideo_vae.encoder.eval()
-    fastvideo_vae.decoder.eval()
-    ref_encoder.eval()
-    ref_decoder.eval()
-
-    if hasattr(fastvideo_vae.decoder, "decode_noise_scale"):
-        fastvideo_vae.decoder.decode_noise_scale = 0.0
-    if hasattr(ref_decoder, "decode_noise_scale"):
-        ref_decoder.decode_noise_scale = 0.0
-
-    batch_size = 1
-    frames = 9
-    height = 64
-    width = 64
-    video = torch.randn(
-        batch_size,
-        3,
-        frames,
-        height,
-        width,
-        device=device,
-        dtype=precision,
-    )
-
-    with torch.no_grad():
-        ref_latents = ref_encoder(video)
-        fast_latents = fastvideo_vae.encoder(video)
-
-    assert ref_latents.shape == fast_latents.shape
-    assert ref_latents.dtype == fast_latents.dtype
-    assert torch.isfinite(ref_latents).all(), "Reference encoder produced non-finite latents."
-    assert torch.isfinite(fast_latents).all(), "FastVideo encoder produced non-finite latents."
-    assert_close(ref_latents, fast_latents, atol=1e-2, rtol=1e-2)
-
-    timestep = torch.tensor([0.05], device=device, dtype=precision)
-    with torch.no_grad():
-        ref_decoded = ref_decoder(ref_latents, timestep=timestep)
-        fast_decoded = fastvideo_vae.decoder(fast_latents, timestep=timestep)
-
-    assert ref_decoded.shape == fast_decoded.shape
-    assert ref_decoded.dtype == fast_decoded.dtype
-    assert torch.isfinite(ref_decoded).all(), "Reference decoder produced non-finite output."
-    assert torch.isfinite(fast_decoded).all(), "FastVideo decoder produced non-finite output."
-    assert_close(ref_decoded, fast_decoded, atol=1e-2, rtol=1e-2)