update

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

README.md

@@ -3,7 +3,7 @@
 </div>
 
 <p align="center">
-     | <a href="https://hao-ai-lab.github.io/FastVideo"><b>Documentation</b></a> | <a href="https://hao-ai-lab.github.io/FastVideo/inference/inference_quick_start/"><b> Quick Start</b></a> | <a href="https://github.com/hao-ai-lab/FastVideo/discussions/982"  target="_blank"><b>Weekly Dev Meeting</b></a>  | 🟣💬 <a href="https://join.slack.com/t/fastvideo/shared_invite/zt-3f4lao1uq-u~Ipx6Lt4J27AlD2y~IdLQ" target="_blank"> <b>Slack</b> </a> |  🟣💬 <a href="https://ibb.co/c7g1qdD" target="_blank"> <b> WeChat </b> </a> |
+     | <a href="https://hao-ai-lab.github.io/FastVideo"><b>Documentation</b></a> | <a href="https://hao-ai-lab.github.io/FastVideo/inference/inference_quick_start/"><b> Quick Start</b></a> | <a href="https://github.com/hao-ai-lab/FastVideo/discussions/982"  target="_blank"><b>Weekly Dev Meeting</b></a>  | 🟣💬 <a href="https://join.slack.com/t/fastvideo/shared_invite/zt-3f4lao1uq-u~Ipx6Lt4J27AlD2y~IdLQ" target="_blank"> <b>Slack</b> </a> |  🟣💬 <a href="https://ibb.co/sv3MMKyv" target="_blank"> <b> WeChat </b> </a> |
 </p>
 
 **FastVideo is a unified post-training and inference framework for accelerated video generation.**

docs/attention/index.md

@@ -41,7 +41,7 @@ Clone the repository and build the kernel:
 
 ```bash
 # Clone recursively to get ThunderKittens submodule
-git clone --recursive https://github.com/hao-ai-lab/FastVideo.git
+git clone https://github.com/hao-ai-lab/FastVideo.git
 cd FastVideo/fastvideo-kernel
 
 # Build and install

docs/getting_started/installation.md

@@ -11,15 +11,13 @@ FastVideo supports the following hardware platforms:
 ### Using pip
 
 ```bash
+# Create and activate a new conda environment
+conda create -n fastvideo python=3.12
+conda activate fastvideo
+
 pip install fastvideo
 ```
 
-### Using conda
-
-```bash
-conda install -c conda-forge fastvideo
-```
-
 ### From source
 
 ```bash
@@ -28,6 +26,12 @@ cd FastVideo
 pip install -e .
 ```
 
+Also optionally install flash-attn:
+
+```bash
+pip install flash-attn --no-build-isolation
+```
+
 ## Hardware Requirements
 
 - **NVIDIA GPUs**: CUDA 11.8+ with compute capability 7.0+

docs/getting_started/quick_start.md

@@ -15,6 +15,12 @@ conda activate fastvideo
 pip install fastvideo
 ```
 
+Also optionally install flash-attn:
+
+```bash
+pip install flash-attn --no-build-isolation
+```
+
 ## Basic Usage
 
 ### Text-to-Video Generation

docs/inference/add_pipeline.md

@@ -168,6 +168,10 @@ Pipelines are composed of stages, each handling a specific part of the diffusion
 - **DenoisingStage**: Performs denoising diffusion
 - **DecodingStage**: Converts latents to pixels
 
+Note: `DenoisingStage` uses the unified denoising engine under the hood. You
+can inject a custom strategy via `strategy_cls` if a pipeline needs specialized
+denoising behavior.
+
 ### Creating Your Pipeline
 
 ```python
@@ -230,7 +234,7 @@ class MyCustomPipeline(ComposedPipelineBase):
             stage_name="denoising_stage",
             stage=DenoisingStage(
                 transformer=self.get_module("transformer"),
-                scheduler=self.get_module("scheduler")
+                scheduler=self.get_module("scheduler"),
             )
         )
         
@@ -245,6 +249,24 @@ class MyCustomPipeline(ComposedPipelineBase):
 EntryClass = MyCustomPipeline
 ```
 
+### Customizing Denoising Strategies
+
+If your model requires a custom denoising loop, pass a strategy class:
+
+```python
+from fastvideo.pipelines.stages.denoising_cosmos_strategy import (
+    CosmosStrategy)
+
+        self.add_stage(
+            stage_name="denoising_stage",
+            stage=DenoisingStage(
+                transformer=self.get_module("transformer"),
+                scheduler=self.get_module("scheduler"),
+                strategy_cls=CosmosStrategy,
+            )
+        )
+```
+
 ### Creating Custom Stages (Optional)
 
 If existing stages don't meet your needs, create custom ones:

docs/inference/inference_quick_start.md

@@ -106,6 +106,8 @@ If you encounter CUDA out of memory errors:
 - Enable memory optimization with `enable_model_cpu_offload`
 - Try a smaller model or use distilled versions
 - Use `num_gpus` > 1 if multiple GPUs are available
+- Try enabling FSDP inference with `use_fsdp_inference=True` (may slow down generation)
+- Try enabling DiT layerwise offload with `dit_layerwise_offload=True` (now only a few models support this, but may introduce less overhead than FSDP)
 
 ### Slow Generation
 

docs/utilities/lora.md

@@ -0,0 +1,71 @@
+# LoRA Extraction and Merging
+
+Tools for extracting and merging LoRA adapters for FastVideo models.
+
+## Extract LoRA Adapter
+
+```bash
+python scripts/lora_extraction/extract_lora.py \
+  --base Wan-AI/Wan2.2-TI2V-5B-Diffusers \
+  --ft FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers \
+  --out adapter_r32.safetensors \
+  --rank 32
+```
+
+**Options:**
+
+- `--base`: Base model (HuggingFace ID or local path)
+- `--ft`: Fine-tuned model (HuggingFace ID or local path)
+- `--out`: Output adapter file
+- `--rank`: LoRA rank (16, 32, 64, 128)
+- `--full-rank`: Extract full-rank adapter (optional)
+
+## Merge Adapter
+
+```bash
+python scripts/lora_extraction/merge_lora.py \
+  --base Wan-AI/Wan2.2-TI2V-5B-Diffusers \
+  --adapter adapter_r32.safetensors \
+  --ft FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers \
+  --output merged_model
+```
+
+**Options:**
+
+- `--base`: Base model (HuggingFace ID or local path)
+- `--adapter`: LoRA adapter file (.safetensors)
+- `--ft`: Fine-tuned model (for configuration)
+- `--output`: Output directory
+
+## Validate Quality (Optional)
+
+```bash
+python scripts/lora_extraction/lora_inference_comparison.py \
+  --base merged_model \
+  --ft FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers \
+  --adapter NONE \
+  --output-dir results \
+  --prompt "A cat sitting on a windowsill" \
+  --seed 42 \
+  --height 480 \
+  --width 480 \
+  --num-frames 49 \
+  --num-inference-steps 32 \
+  --compute-ssim \
+  --compute-lpips
+```
+
+**Options:**
+
+- `--base`: Merged model or base model path
+- `--ft`: Fine-tuned model (reference)
+- `--adapter`: Path to adapter or NONE
+- `--output-dir`: Output directory
+- `--prompt`: Text prompt (default: "A cat sitting on a windowsill")
+- `--seed`: Random seed (default: 42)
+- `--height`: Video height (default: 480)
+- `--width`: Video width (default: 832)
+- `--num-frames`: Number of frames (default: 49)
+- `--num-inference-steps`: Inference steps (default: 32)
+- `--compute-ssim`: Compute SSIM metric
+- `--compute-lpips`: Compute LPIPS metric

examples/inference/basic/basic.py

@@ -12,7 +12,7 @@ def main():
         "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
         # FastVideo will automatically handle distributed setup
         num_gpus=1,
-        use_fsdp_inference=True,
+        use_fsdp_inference=False, # set to True if GPU is out of memory
         dit_cpu_offload=False,
         vae_cpu_offload=False,
         text_encoder_cpu_offload=True,

examples/inference/basic/basic_dmd.py

@@ -14,7 +14,7 @@ def main():
         model_name,
         # FastVideo will automatically handle distributed setup
         num_gpus=1,
-        use_fsdp_inference=True,
+        use_fsdp_inference=False, # set to True if GPU is out of memory
         # Adjust these offload parameters if you have < 32GB of VRAM
         text_encoder_cpu_offload=True,
         pin_cpu_memory=True, # set to false if low CPU RAM or hit obscure "CUDA error: Invalid argument"

examples/inference/basic/basic_hy15.py

@@ -12,7 +12,7 @@ def main():
         "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v",
         # FastVideo will automatically handle distributed setup
         num_gpus=1,
-        use_fsdp_inference=True,
+        use_fsdp_inference=False, # set to True if GPU is out of memory
         dit_cpu_offload=True,
         vae_cpu_offload=True,
         text_encoder_cpu_offload=True,

examples/inference/basic/basic_longcat_i2v.py

@@ -0,0 +1,210 @@
+"""
+LongCat Image-to-Video (I2V) Example Script
+
+This script demonstrates LongCat I2V inference using the FastVideo Python API.
+LongCat I2V takes an input image and generates a video from it.
+
+It runs both basic generation (50 steps) and distill+refine generation
+(16 steps distill + 50 steps refinement to 720p with BSA).
+
+Usage:
+    python examples/inference/basic/basic_longcat_i2v.py
+
+Note:
+    Refinement uses 768x768 dimensions where latent (48x48) is divisible by 8,
+    compatible with BSA chunks [4, 4, 8].
+"""
+
+import glob
+import os
+
+from fastvideo import VideoGenerator
+
+# Common prompts and settings matching the shell script examples
+PROMPT = (
+    "A woman sits at a wooden table by the window in a cozy café. She reaches out "
+    "with her right hand, picks up the white coffee cup from the saucer, and gently "
+    "brings it to her lips to take a sip. After drinking, she places the cup back on "
+    "the table and looks out the window, enjoying the peaceful atmosphere."
+)
+
+NEGATIVE_PROMPT = (
+    "Bright tones, overexposed, static, blurred details, subtitles, style, works, "
+    "paintings, images, static, overall gray, worst quality, low quality, JPEG compression "
+    "residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, "
+    "deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, "
+    "three legs, many people in the background, walking backwards"
+)
+
+# Input image path
+IMAGE_PATH = "assets/girl.png"
+
+SEED = 42
+
+
+def basic_generation():
+    """
+    Run basic LongCat I2V generation (50 steps at 480p).
+    
+    This uses the full 50-step denoising process for highest quality.
+    """
+    print("=" * 60)
+    print("LongCat I2V: Basic Generation (50 steps, 480p)")
+    print("=" * 60)
+    
+    generator = VideoGenerator.from_pretrained(
+        "FastVideo/LongCat-Video-I2V-Diffusers",
+        num_gpus=1,
+        use_fsdp_inference=False, # set to True if GPU is out of memory
+        dit_cpu_offload=False,
+        vae_cpu_offload=True,
+        text_encoder_cpu_offload=True,
+        pin_cpu_memory=False,
+        enable_bsa=False,
+    )
+    
+    output_path = "outputs_video/longcat_i2v_basic"
+    
+    generator.generate_video(
+        prompt=PROMPT,
+        negative_prompt=NEGATIVE_PROMPT,
+        image_path=IMAGE_PATH,
+        output_path=output_path,
+        save_video=True,
+        height=480,
+        width=480,  # Square
+        num_frames=93,
+        num_inference_steps=50,
+        fps=15,
+        guidance_scale=4.0,
+        seed=SEED,
+    )
+    
+    print(f"\nBasic generation complete! Video saved to: {output_path}")
+    generator.shutdown()
+
+
+def distill_refine_generation():
+    """
+    Run LongCat I2V with distill+refine pipeline (16 steps + refinement to 768p).
+    
+    This uses the distilled LoRA for fast 480p generation (16 steps),
+    then refines to 768p using the refinement LoRA with BSA enabled.
+    """
+    print("\n" + "=" * 60)
+    print("LongCat I2V: Distill + Refine Pipeline")
+    print("=" * 60)
+    
+    # Stage 1: Distilled generation (16 steps at 480p)
+    print("\n[Stage 1] Distilled generation (16 steps, 480p)")
+    print("-" * 40)
+    
+    generator = VideoGenerator.from_pretrained(
+        "FastVideo/LongCat-Video-I2V-Diffusers",
+        num_gpus=1,
+        use_fsdp_inference=True,
+        dit_cpu_offload=False,
+        vae_cpu_offload=True,
+        text_encoder_cpu_offload=True,
+        pin_cpu_memory=False,
+        enable_bsa=False,
+        lora_path="FastVideo/LongCat-Video-T2V-Distilled-LoRA",
+        lora_nickname="distilled",
+    )
+    
+    distill_output_path = "outputs_video/longcat_i2v_distill"
+    
+    generator.generate_video(
+        prompt=PROMPT,
+        negative_prompt=NEGATIVE_PROMPT,
+        image_path=IMAGE_PATH,
+        output_path=distill_output_path,
+        save_video=True,
+        height=480,
+        width=480,  # Square
+        num_frames=93,
+        num_inference_steps=16,
+        fps=15,
+        guidance_scale=1.0,
+        seed=SEED,
+    )
+    
+    print(f"Distilled generation complete! Video saved to: {distill_output_path}")
+    generator.shutdown()
+    
+    # Stage 2: Refinement (480p -> 768p)
+    print("\n[Stage 2] Refinement (480p -> 768p with BSA)")
+    print("-" * 40)
+    
+    # Find the actual saved video file from stage 1
+    video_files = glob.glob(os.path.join(distill_output_path, "*.mp4"))
+    if not video_files:
+        raise FileNotFoundError(f"No video file found in {distill_output_path}")
+    # Use the most recently created video file
+    distill_video_path = max(video_files, key=os.path.getmtime)
+    print(f"Using stage 1 video: {distill_video_path}")
+    
+    # Create a new generator with refinement LoRA and BSA enabled
+    # Note: Refinement uses the T2V model (not I2V) since it's upscaling the generated video
+    # For BSA [4, 4, 8]: latent must be divisible by 8
+    # 768x768: latent 48x48, 48%8=0 ✓
+    refine_generator = VideoGenerator.from_pretrained(
+        "FastVideo/LongCat-Video-T2V-Diffusers",
+        num_gpus=1,
+        use_fsdp_inference=True,
+        dit_cpu_offload=True,
+        vae_cpu_offload=True,
+        text_encoder_cpu_offload=True,
+        pin_cpu_memory=False,
+        enable_bsa=True,
+        bsa_sparsity=0.875,
+        bsa_chunk_q=[4, 4, 4],
+        bsa_chunk_k=[4, 4, 4],
+        lora_path="FastVideo/LongCat-Video-T2V-Refinement-LoRA",
+        lora_nickname="refinement",
+    )
+    
+    refine_output_path = "outputs_video/longcat_i2v_refine_720p"
+    
+    refine_generator.generate_video(
+        prompt=PROMPT,
+        negative_prompt=NEGATIVE_PROMPT,
+        output_path=refine_output_path,
+        save_video=True,
+        refine_from=distill_video_path,
+        t_thresh=0.5,
+        spatial_refine_only=False,
+        num_cond_frames=0,
+        height=720,
+        width=720,
+        num_inference_steps=50,
+        fps=30,
+        guidance_scale=1.0,
+        seed=SEED,
+    )
+    
+    print(f"Refinement complete! Video saved to: {refine_output_path}")
+    refine_generator.shutdown()
+
+
+def main():
+    """Run both basic and distill+refine generation pipelines."""
+    print("\n" + "=" * 60)
+    print("LongCat Image-to-Video Example")
+    print("=" * 60 + "\n")
+    
+    # Run basic generation
+    basic_generation()
+    
+    # Run distill+refine pipeline
+    distill_refine_generation()
+    
+    print("\n" + "=" * 60)
+    print("All generations complete!")
+    print("=" * 60)
+
+
+if __name__ == "__main__":
+    main()
+
+

examples/inference/basic/basic_longcat_t2v.py

@@ -0,0 +1,198 @@
+"""
+LongCat Text-to-Video (T2V) Example Script
+
+This script demonstrates LongCat T2V inference using the FastVideo Python API.
+It runs both basic generation (50 steps) and distill+refine generation
+(16 steps distill + 50 steps refinement to 720p).
+
+Usage:
+    python examples/inference/basic/basic_longcat_t2v.py
+"""
+
+import glob
+import os
+
+from fastvideo import VideoGenerator
+
+# Common prompts and settings matching the shell script examples
+PROMPT = (
+    "In a realistic photography style, a white boy around seven or eight years old "
+    "sits on a park bench, wearing a light blue T-shirt, denim shorts, and white sneakers. "
+    "He holds an ice cream cone with vanilla and chocolate flavors, and beside him is a "
+    "medium-sized golden Labrador. Smiling, the boy offers the ice cream to the dog, "
+    "who eagerly licks it with its tongue. The sun is shining brightly, and the background "
+    "features a green lawn and several tall trees, creating a warm and loving scene."
+)
+
+NEGATIVE_PROMPT = (
+    "Bright tones, overexposed, static, blurred details, subtitles, style, works, "
+    "paintings, images, static, overall gray, worst quality, low quality, JPEG compression "
+    "residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, "
+    "deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, "
+    "three legs, many people in the background, walking backwards"
+)
+
+SEED = 42
+
+
+def basic_generation():
+    """
+    Run basic LongCat T2V generation (50 steps at 480p).
+    
+    This uses the full 50-step denoising process for highest quality.
+    """
+    print("=" * 60)
+    print("LongCat T2V: Basic Generation (50 steps, 480p)")
+    print("=" * 60)
+    
+    generator = VideoGenerator.from_pretrained(
+        "FastVideo/LongCat-Video-T2V-Diffusers",
+        num_gpus=1,
+        use_fsdp_inference=False, # set to True if GPU is out of memory
+        dit_cpu_offload=False,
+        vae_cpu_offload=True,
+        text_encoder_cpu_offload=True,
+        pin_cpu_memory=False,
+        enable_bsa=False,
+    )
+    
+    output_path = "outputs_video/longcat_t2v_basic"
+    
+    generator.generate_video(
+        prompt=PROMPT,
+        negative_prompt=NEGATIVE_PROMPT,
+        output_path=output_path,
+        save_video=True,
+        height=480,
+        width=832,
+        num_frames=93,
+        num_inference_steps=50,
+        fps=15,
+        guidance_scale=4.0,
+        seed=SEED,
+    )
+    
+    print(f"\nBasic generation complete! Video saved to: {output_path}")
+    generator.shutdown()
+
+
+def distill_refine_generation():
+    """
+    Run LongCat T2V with distill+refine pipeline (16 steps + refinement to 720p).
+    
+    This uses the distilled LoRA for fast 480p generation (16 steps),
+    then refines to 720p using the refinement LoRA with BSA enabled.
+    """
+    print("\n" + "=" * 60)
+    print("LongCat T2V: Distill + Refine Pipeline")
+    print("=" * 60)
+    
+    # Stage 1: Distilled generation (16 steps at 480p)
+    print("\n[Stage 1] Distilled generation (16 steps, 480p)")
+    print("-" * 40)
+    
+    generator = VideoGenerator.from_pretrained(
+        "FastVideo/LongCat-Video-T2V-Diffusers",
+        num_gpus=1,
+        use_fsdp_inference=True,
+        dit_cpu_offload=False,
+        vae_cpu_offload=True,
+        text_encoder_cpu_offload=True,
+        pin_cpu_memory=False,
+        enable_bsa=False,
+        lora_path="FastVideo/LongCat-Video-T2V-Distilled-LoRA",
+        lora_nickname="distilled",
+    )
+    
+    distill_output_path = "outputs_video/longcat_t2v_distill"
+    
+    generator.generate_video(
+        prompt=PROMPT,
+        negative_prompt=NEGATIVE_PROMPT,
+        output_path=distill_output_path,
+        save_video=True,
+        height=480,
+        width=832,
+        num_frames=93,
+        num_inference_steps=16,
+        fps=15,
+        guidance_scale=1.0,
+        seed=SEED,
+    )
+    
+    print(f"Distilled generation complete! Video saved to: {distill_output_path}")
+    generator.shutdown()
+    
+    # Stage 2: Refinement (480p -> 720p)
+    print("\n[Stage 2] Refinement (480p -> 720p with BSA)")
+    print("-" * 40)
+    
+    # Find the actual saved video file from stage 1
+    video_files = glob.glob(os.path.join(distill_output_path, "*.mp4"))
+    if not video_files:
+        raise FileNotFoundError(f"No video file found in {distill_output_path}")
+    # Use the most recently created video file
+    distill_video_path = max(video_files, key=os.path.getmtime)
+    print(f"Using stage 1 video: {distill_video_path}")
+    
+    # Create a new generator with refinement LoRA and BSA enabled
+    refine_generator = VideoGenerator.from_pretrained(
+        "FastVideo/LongCat-Video-T2V-Diffusers",
+        num_gpus=1,
+        use_fsdp_inference=True,
+        dit_cpu_offload=True,
+        vae_cpu_offload=True,
+        text_encoder_cpu_offload=True,
+        pin_cpu_memory=False,
+        enable_bsa=True,
+        bsa_sparsity=0.875,
+        bsa_chunk_q=[4, 4, 8],
+        bsa_chunk_k=[4, 4, 8],
+        lora_path="FastVideo/LongCat-Video-T2V-Refinement-LoRA",
+        lora_nickname="refinement",
+    )
+    
+    refine_output_path = "outputs_video/longcat_t2v_refine_720p"
+    
+    refine_generator.generate_video(
+        prompt=PROMPT,
+        negative_prompt=NEGATIVE_PROMPT,
+        output_path=refine_output_path,
+        save_video=True,
+        refine_from=distill_video_path,
+        t_thresh=0.5,
+        spatial_refine_only=False,
+        num_cond_frames=0,
+        height=720,
+        width=1280,
+        num_inference_steps=50,
+        fps=30,
+        guidance_scale=1.0,
+        seed=SEED,
+    )
+    
+    print(f"Refinement complete! Video saved to: {refine_output_path}")
+    refine_generator.shutdown()
+
+
+def main():
+    """Run both basic and distill+refine generation pipelines."""
+    print("\n" + "=" * 60)
+    print("LongCat Text-to-Video Example")
+    print("=" * 60 + "\n")
+    
+    # Run basic generation
+    basic_generation()
+    
+    # Run distill+refine pipeline
+    distill_refine_generation()
+    
+    print("\n" + "=" * 60)
+    print("All generations complete!")
+    print("=" * 60)
+
+
+if __name__ == "__main__":
+    main()
+
+

examples/inference/basic/basic_longcat_vc.py

@@ -0,0 +1,228 @@
+"""
+LongCat Video Continuation (VC) Example Script
+
+This script demonstrates LongCat VC inference using the FastVideo Python API.
+LongCat VC takes an input video and generates a continuation of it.
+
+It runs both basic generation (50 steps) and distill+refine generation
+(16 steps distill + 50 steps refinement to 720p).
+
+Usage:
+    python examples/inference/basic/basic_longcat_vc.py
+
+Prerequisites:
+    - Ensure the input video exists at assets/motorcycle.mp4
+      (or provide your own video)
+"""
+
+import glob
+import os
+
+from fastvideo import VideoGenerator
+
+# Common prompts and settings matching the shell script examples
+PROMPT = (
+    "A person rides a motorcycle along a long, straight road that stretches between "
+    "a body of water and a forested hillside. The rider steadily accelerates, keeping "
+    "the motorcycle centered between the guardrails, while the scenery passes by on "
+    "both sides. The video captures the journey from the rider's perspective, emphasizing "
+    "the sense of motion and adventure."
+)
+
+NEGATIVE_PROMPT = (
+    "Bright tones, overexposed, static, blurred details, subtitles, style, works, "
+    "paintings, images, static, overall gray, worst quality, low quality, JPEG compression "
+    "residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, "
+    "deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, "
+    "three legs, many people in the background, walking backwards"
+)
+
+# Input video path
+VIDEO_PATH = "assets/motorcycle.mp4"
+
+# Number of conditioning frames from the input video
+NUM_COND_FRAMES = 13
+
+SEED = 42
+
+
+def basic_generation():
+    """
+    Run basic LongCat VC generation (50 steps at 480p).
+    
+    This uses the full 50-step denoising process for highest quality.
+    """
+    print("=" * 60)
+    print("LongCat VC: Basic Generation (50 steps, 480p)")
+    print("=" * 60)
+    
+    # Check if video exists
+    if not os.path.exists(VIDEO_PATH):
+        raise FileNotFoundError(
+            f"Video not found at {VIDEO_PATH}. "
+            "Please provide a valid video path."
+        )
+    
+    generator = VideoGenerator.from_pretrained(
+        "FastVideo/LongCat-Video-VC-Diffusers",
+        num_gpus=1,
+        use_fsdp_inference=False, # set to True if GPU is out of memory
+        dit_cpu_offload=False,
+        vae_cpu_offload=True,
+        text_encoder_cpu_offload=True,
+        pin_cpu_memory=False,
+        enable_bsa=False,
+    )
+    
+    output_path = "outputs_video/longcat_vc_basic"
+    
+    generator.generate_video(
+        prompt=PROMPT,
+        negative_prompt=NEGATIVE_PROMPT,
+        video_path=VIDEO_PATH,
+        num_cond_frames=NUM_COND_FRAMES,
+        output_path=output_path,
+        save_video=True,
+        height=480,
+        width=832,
+        num_frames=93,
+        num_inference_steps=50,
+        fps=15,
+        guidance_scale=4.0,
+        seed=SEED,
+    )
+    
+    print(f"\nBasic generation complete! Video saved to: {output_path}")
+    generator.shutdown()
+
+
+def distill_refine_generation():
+    """
+    Run LongCat VC with distill+refine pipeline (16 steps + refinement to 720p).
+    
+    This uses the distilled LoRA for fast 480p generation (16 steps),
+    then refines to 720p using the refinement LoRA with BSA enabled.
+    """
+    print("\n" + "=" * 60)
+    print("LongCat VC: Distill + Refine Pipeline")
+    print("=" * 60)
+    
+    # Check if video exists
+    if not os.path.exists(VIDEO_PATH):
+        raise FileNotFoundError(
+            f"Video not found at {VIDEO_PATH}. "
+            "Please provide a valid video path."
+        )
+    
+    # Stage 1: Distilled generation (16 steps at 480p)
+    print("\n[Stage 1] Distilled generation (16 steps, 480p)")
+    print("-" * 40)
+    
+    generator = VideoGenerator.from_pretrained(
+        "FastVideo/LongCat-Video-VC-Diffusers",
+        num_gpus=1,
+        use_fsdp_inference=True,
+        dit_cpu_offload=False,
+        vae_cpu_offload=True,
+        text_encoder_cpu_offload=True,
+        pin_cpu_memory=False,
+        enable_bsa=False,
+        lora_path="FastVideo/LongCat-Video-T2V-Distilled-LoRA",
+        lora_nickname="distilled",
+    )
+    
+    distill_output_path = "outputs_video/longcat_vc_distill"
+    
+    generator.generate_video(
+        prompt=PROMPT,
+        negative_prompt=NEGATIVE_PROMPT,
+        video_path=VIDEO_PATH,
+        num_cond_frames=NUM_COND_FRAMES,
+        output_path=distill_output_path,
+        save_video=True,
+        height=480,
+        width=832,
+        num_frames=93,
+        num_inference_steps=16,
+        fps=15,
+        guidance_scale=1.0,
+        seed=SEED,
+    )
+    
+    print(f"Distilled generation complete! Video saved to: {distill_output_path}")
+    generator.shutdown()
+    
+    # Stage 2: Refinement (480p -> 720p)
+    print("\n[Stage 2] Refinement (480p -> 720p with BSA)")
+    print("-" * 40)
+    
+    # Find the actual saved video file from stage 1
+    video_files = glob.glob(os.path.join(distill_output_path, "*.mp4"))
+    if not video_files:
+        raise FileNotFoundError(f"No video file found in {distill_output_path}")
+    # Use the most recently created video file
+    distill_video_path = max(video_files, key=os.path.getmtime)
+    print(f"Using stage 1 video: {distill_video_path}")
+    
+    # Create a new generator with refinement LoRA and BSA enabled
+    # Note: Refinement uses the T2V model (not VC) since it's upscaling the generated video
+    refine_generator = VideoGenerator.from_pretrained(
+        "FastVideo/LongCat-Video-T2V-Diffusers",
+        num_gpus=1,
+        use_fsdp_inference=True,
+        dit_cpu_offload=True,
+        vae_cpu_offload=True,
+        text_encoder_cpu_offload=True,
+        pin_cpu_memory=False,
+        enable_bsa=True,
+        bsa_sparsity=0.875,
+        bsa_chunk_q=[4, 4, 8],
+        bsa_chunk_k=[4, 4, 8],
+        lora_path="FastVideo/LongCat-Video-T2V-Refinement-LoRA",
+        lora_nickname="refinement",
+    )
+    
+    refine_output_path = "outputs_video/longcat_vc_refine_720p"
+    
+    refine_generator.generate_video(
+        prompt=PROMPT,
+        negative_prompt=NEGATIVE_PROMPT,
+        output_path=refine_output_path,
+        save_video=True,
+        refine_from=distill_video_path,
+        t_thresh=0.5,
+        spatial_refine_only=False,
+        num_cond_frames=0,  # For refinement, no conditioning frames
+        height=720,
+        width=1280,
+        num_inference_steps=50,
+        fps=30,
+        guidance_scale=1.0,
+        seed=SEED,
+    )
+    
+    print(f"Refinement complete! Video saved to: {refine_output_path}")
+    refine_generator.shutdown()
+
+
+def main():
+    """Run both basic and distill+refine generation pipelines."""
+    print("\n" + "=" * 60)
+    print("LongCat Video Continuation Example")
+    print("=" * 60 + "\n")
+    
+    # Run basic generation
+    basic_generation()
+    
+    # Run distill+refine pipeline
+    distill_refine_generation()
+    
+    print("\n" + "=" * 60)
+    print("All generations complete!")
+    print("=" * 60)
+
+
+if __name__ == "__main__":
+    main()
+
+

examples/inference/basic/basic_matrixgame.py

@@ -43,7 +43,7 @@ def main():
         config["model_path"],
         # FastVideo will automatically handle distributed setup
         num_gpus=1,
-        use_fsdp_inference=True,
+        use_fsdp_inference=False, # set to True if GPU is out of memory
         dit_cpu_offload=True, # DiT need to be offloaded for MoE
         vae_cpu_offload=False,
         text_encoder_cpu_offload=True,

examples/inference/basic/basic_matrixgame_streaming.py

@@ -46,7 +46,7 @@ async def main():
         config["model_path"],
         # FastVideo will automatically handle distributed setup
         num_gpus=1,
-        use_fsdp_inference=True,
+        use_fsdp_inference=False, # set to True if GPU is out of memory
         dit_cpu_offload=True, # DiT need to be offloaded for MoE
         vae_cpu_offload=False,
         text_encoder_cpu_offload=True,

examples/inference/basic/basic_self_forcing_causal.py

@@ -13,7 +13,7 @@ def main():
         model_name,
         # FastVideo will automatically handle distributed setup
         num_gpus=1,
-        use_fsdp_inference=True,
+        use_fsdp_inference=False, # set to True if GPU is out of memory
         text_encoder_cpu_offload=False,
         dit_cpu_offload=False,
     )

examples/inference/basic/basic_self_forcing_causal_wan2_2_i2v.py

@@ -14,7 +14,7 @@ def main():
         "FastVideo/SFWan2.2-I2V-A14B-Preview-Diffusers",
         # FastVideo will automatically handle distributed setup
         num_gpus=1,
-        use_fsdp_inference=True,
+        use_fsdp_inference=False, # set to True if GPU is out of memory
         dit_cpu_offload=True, # DiT need to be offloaded for MoE
         dit_precision="fp32",
         vae_cpu_offload=False,

examples/inference/basic/basic_self_forcing_causal_wan2_2_t2v.py

@@ -14,7 +14,7 @@ def main():
         "rand0nmr/SFWan2.2-T2V-A14B-Diffusers",
         # FastVideo will automatically handle distributed setup
         num_gpus=1,
-        use_fsdp_inference=True,
+        use_fsdp_inference=False, # set to True if GPU is out of memory
         dit_cpu_offload=True, # DiT need to be offloaded for MoE
         vae_cpu_offload=False,
         text_encoder_cpu_offload=True,

examples/inference/basic/basic_turbodiffusion.py

@@ -15,12 +15,14 @@ def main() -> None:
         "loayrashid/TurboWan2.1-T2V-1.3B-Diffusers",
         # FastVideo will automatically handle distributed setup
         num_gpus=1,
-        # TurboDiffusion uses a custom pipeline with RCM scheduler
-        override_pipeline_cls_name="TurboDiffusionPipeline",
+        use_fsdp_inference=False, # set to True if GPU is out of memory
+
+        # set to false if using RTX 4090 
+        # pin_cpu_memory=False,
     )
 
     # Generate videos with the same simple API, regardless of GPU count
-    # TurboDiffusion uses guidance_scale=1.0 (no CFG) and only 4 steps
+    # TurboDiffusion defaults: guidance_scale=1.0 and num_inference_steps=4 (from config)
     prompt = (
         "A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes "
         "wide with interest. The playful yet serene atmosphere is complemented by soft "
@@ -30,9 +32,7 @@ def main() -> None:
         prompt,
         output_path=OUTPUT_PATH,
         save_video=True,
-        num_inference_steps=4,
         seed=42,
-        guidance_scale=1.0,
     )
 
     # Generate another video with a different prompt, without reloading the model!
@@ -47,9 +47,7 @@ def main() -> None:
         prompt2,
         output_path=OUTPUT_PATH,
         save_video=True,
-        num_inference_steps=4,
         seed=42,
-        guidance_scale=1.0,
     )
 
 

examples/inference/basic/basic_turbodiffusion_14b.py

@@ -15,8 +15,6 @@ def main() -> None:
         "loayrashid/TurboWan2.1-T2V-14B-Diffusers",
         # 14B model needs more GPUs
         num_gpus=2,
-        # TurboDiffusion uses a custom pipeline with RCM scheduler
-        override_pipeline_cls_name="TurboDiffusionPipeline",
     )
 
     prompt = (
@@ -28,9 +26,7 @@ def main() -> None:
         prompt,
         output_path=OUTPUT_PATH,
         save_video=True,
-        num_inference_steps=4,
         seed=42,
-        guidance_scale=1.0,
     )
 
     # Generate another video with a different prompt, without reloading the model!
@@ -45,9 +41,7 @@ def main() -> None:
         prompt2,
         output_path=OUTPUT_PATH,
         save_video=True,
-        num_inference_steps=4,
         seed=42,
-        guidance_scale=1.0,
     )
 
 

examples/inference/basic/basic_turbodiffusion_i2v.py

@@ -0,0 +1,37 @@
+import os
+
+# Set SLA attention backend BEFORE fastvideo imports
+os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "SLA_ATTN"
+
+from fastvideo import VideoGenerator
+
+# Use local model path
+MODEL_PATH = "loayrashid/TurboWan2.2-I2V-A14B-Diffusers"
+OUTPUT_PATH = "video_samples_turbodiffusion_i2v"
+
+
+def main() -> None:
+    # TurboDiffusion I2V: 1-4 step image-to-video generation
+    generator = VideoGenerator.from_pretrained(
+        MODEL_PATH,
+        num_gpus=2,
+    )
+
+    # Example prompt and image for I2V
+    prompt = ("Summer beach vacation style, a white cat wearing sunglasses sits on a surfboard. The fluffy-furred feline gazes directly at the camera with a relaxed expression. Blurred beach scenery forms the background featuring crystal-clear waters, distant green hills, and a blue sky dotted with white clouds. The cat assumes a naturally relaxed posture, as if savoring the sea breeze and warm sunlight. A close-up shot highlights the feline's intricate details and the refreshing atmosphere of the seaside.")
+
+    # Use an example image path
+    image_path = "https://huggingface.co/datasets/YiYiXu/testing-images/resolve/main/wan_i2v_input.JPG"
+
+
+    video = generator.generate_video(
+        prompt,
+        image_path=image_path,
+        output_path=OUTPUT_PATH,
+        save_video=True,
+        seed=42,
+    )
+
+
+if __name__ == "__main__":
+    main()

examples/inference/basic/basic_wan2_2.py

@@ -12,7 +12,7 @@ def main():
         "Wan-AI/Wan2.2-T2V-A14B-Diffusers",
         # FastVideo will automatically handle distributed setup
         num_gpus=1,
-        use_fsdp_inference=True,
+        use_fsdp_inference=False, # set to True if GPU is out of memory 
         dit_cpu_offload=True, # DiT need to be offloaded for MoE
         vae_cpu_offload=False,
         text_encoder_cpu_offload=True,

examples/inference/basic/basic_wan2_2_Fun.py

@@ -14,7 +14,7 @@ def main():
         # "alibaba-pai/Wan2.2-Fun-A14B-Control",
         # FastVideo will automatically handle distributed setup
         num_gpus=1,
-        use_fsdp_inference=True,
+        use_fsdp_inference=False, # set to True if GPU is out of memory
         dit_cpu_offload=True, # DiT need to be offloaded for MoE
         vae_cpu_offload=False,
         text_encoder_cpu_offload=True,

examples/inference/basic/basic_wan2_2_i2v.py

@@ -12,7 +12,7 @@ def main():
         "Wan-AI/Wan2.2-I2V-A14B-Diffusers",
         # FastVideo will automatically handle distributed setup
         num_gpus=1,
-        use_fsdp_inference=True,
+        use_fsdp_inference=False, # set to True if GPU is out of memory
         dit_cpu_offload=True, # DiT need to be offloaded for MoE
         vae_cpu_offload=False,
         text_encoder_cpu_offload=True,

examples/inference/basic/basic_wan2_2_ti2v.py

@@ -11,7 +11,7 @@ def main():
         model_name,
         # FastVideo will automatically handle distributed setup
         num_gpus=1,
-        use_fsdp_inference=True,
+        use_fsdp_inference=False, # set to True if GPU is out of memory
         dit_cpu_offload=True,
         vae_cpu_offload=False,
         text_encoder_cpu_offload=True,

fastvideo-kernel/pyproject.toml

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

fastvideo-kernel/python/fastvideo_kernel/version.py

@@ -1 +1 @@
-__version__ = "0.2.1"
+__version__ = "0.2.2"

fastvideo/configs/models/dits/longcat.py

@@ -26,7 +26,6 @@ class LongCatVideoArchConfig(DiTArchConfig):
         default_factory=lambda: [is_longcat_blocks])
 
     # Parameter name mapping for weight conversion
-    # Maps original LongCat third_party names -> native FastVideo names
     param_names_mapping: dict = field(
         default_factory=lambda: {
             # Embedders

fastvideo/configs/pipelines/longcat.py

@@ -17,11 +17,7 @@ from fastvideo.configs.pipelines.base import PipelineConfig
 
 @dataclass
 class LongCatDiTArchConfig(DiTArchConfig):
-    """Extended DiTArchConfig with LongCat-specific fields.
-    
-    NOTE: This is for Phase 1 wrapper compatibility. For native model (Phase 2),
-    use LongCatVideoConfig from fastvideo.configs.models.dits.longcat instead.
-    """
+    """Extended DiTArchConfig with LongCat-specific fields."""
     # LongCat-specific architecture parameters
     adaln_tembed_dim: int = 512
     caption_channels: int = 4096
@@ -88,20 +84,16 @@ def umt5_postprocess_text(outputs: BaseEncoderOutput) -> torch.Tensor:
 
 @dataclass
 class LongCatT2V480PConfig(PipelineConfig):
-    """Configuration for LongCat pipeline (480p) aligned to LongCat-Video modules.
+    """Configuration for LongCat pipeline (480p).
 
     Components expected by loaders:
       - tokenizer: AutoTokenizer
       - text_encoder: UMT5EncoderModel
-      - transformer: LongCatVideoTransformer3DModel (Phase 1 wrapper)
-                  OR LongCatTransformer3DModel (Phase 2 native)
+      - transformer: LongCatTransformer3DModel
       - vae: AutoencoderKLWan (Wan VAE, 4x8 compression)
       - scheduler: FlowMatchEulerDiscreteScheduler
     """
 
-    # DiT config with LongCat-specific arch_config
-    # NOTE: For Phase 1 wrapper, uses LongCatDiTArchConfig
-    # For Phase 2 native model, can use LongCatVideoConfig directly
     dit_config: DiTConfig = field(
         default_factory=lambda: DiTConfig(arch_config=LongCatDiTArchConfig()))
 

fastvideo/configs/pipelines/registry.py

@@ -10,6 +10,9 @@ from fastvideo.configs.pipelines.hunyuan import FastHunyuanConfig, HunyuanConfig
 from fastvideo.configs.pipelines.hunyuan15 import Hunyuan15T2V480PConfig, Hunyuan15T2V720PConfig
 from fastvideo.configs.pipelines.stepvideo import StepVideoT2VConfig
 from fastvideo.configs.pipelines.longcat import LongCatT2V480PConfig
+from fastvideo.configs.pipelines.turbodiffusion import (
+    TurboDiffusionT2V_1_3B_Config, TurboDiffusionT2V_14B_Config,
+    TurboDiffusionI2V_A14B_Config)
 
 # isort: off
 from fastvideo.configs.pipelines.wan import (
@@ -55,11 +58,25 @@ PIPE_NAME_TO_CONFIG: dict[str, type[PipelineConfig]] = {
     "FastVideo/Matrix-Game-2.0-Base-Diffusers": MatrixGameI2V480PConfig,
     "FastVideo/Matrix-Game-2.0-GTA-Diffusers": MatrixGameI2V480PConfig,
     "FastVideo/Matrix-Game-2.0-TempleRun-Diffusers": MatrixGameI2V480PConfig,
+    # LongCat Video models
+    "FastVideo/LongCat-Video-T2V-Diffusers": LongCatT2V480PConfig,
+    "FastVideo/LongCat-Video-I2V-Diffusers": LongCatT2V480PConfig,
+    "FastVideo/LongCat-Video-VC-Diffusers": LongCatT2V480PConfig,
+    # TurboDiffusion models
+    "loayrashid/TurboWan2.1-T2V-1.3B-Diffusers": TurboDiffusionT2V_1_3B_Config,
+    "loayrashid/TurboWan2.1-T2V-14B-Diffusers": TurboDiffusionT2V_14B_Config,
+    "loayrashid/TurboWan2.2-I2V-A14B-Diffusers": TurboDiffusionI2V_A14B_Config,
     # Add other specific weight variants
 }
 
 # For determining pipeline type from model ID
 PIPELINE_DETECTOR: dict[str, Callable[[str], bool]] = {
+    "longcatimagetovideo":
+    lambda id: "longcatimagetovideo" in id.lower(),
+    "longcatvideocontinuation":
+    lambda id: "longcatvideocontinuation" in id.lower(),
+    "longcat":
+    lambda id: "longcat" in id.lower(),
     "hunyuan":
     lambda id: "hunyuan" in id.lower(),
     "hunyuan15":
@@ -78,13 +95,15 @@ PIPELINE_DETECTOR: dict[str, Callable[[str], bool]] = {
     lambda id: "stepvideo" in id.lower(),
     "cosmos":
     lambda id: "cosmos" in id.lower(),
-    "longcat":
-    lambda id: "longcat" in id.lower(),
+    "turbodiffusion":
+    lambda id: "turbodiffusion" in id.lower() or "turbowan" in id.lower(),
     # Add other pipeline architecture detectors
 }
 
 # Fallback configs when exact match isn't found but architecture is detected
 PIPELINE_FALLBACK_CONFIG: dict[str, type[PipelineConfig]] = {
+    "longcatimagetovideo": LongCatT2V480PConfig,
+    "longcatvideocontinuation": LongCatT2V480PConfig,
     "longcat": LongCatT2V480PConfig,
     "hunyuan":
     HunyuanConfig,  # Base Hunyuan config as fallback for any Hunyuan variant
@@ -96,7 +115,8 @@ PIPELINE_FALLBACK_CONFIG: dict[str, type[PipelineConfig]] = {
     "wanimagetovideo": WanI2V480PConfig,
     "wandmdpipeline": FastWan2_1_T2V_480P_Config,
     "wancausaldmdpipeline": SelfForcingWanT2V480PConfig,
-    "stepvideo": StepVideoT2VConfig
+    "stepvideo": StepVideoT2VConfig,
+    "turbodiffusion": TurboDiffusionT2V_1_3B_Config,
     # Other fallbacks by architecture
 }
 

fastvideo/configs/pipelines/turbodiffusion.py

@@ -0,0 +1,128 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+TurboDiffusion pipeline configurations.
+
+TurboDiffusion uses RCM (recurrent Consistency Model) scheduler with
+SLA (Sparse-Linear Attention) for fast 1-4 step video generation.
+"""
+from dataclasses import dataclass, field
+
+from fastvideo.configs.models import DiTConfig, EncoderConfig, VAEConfig
+from fastvideo.configs.models.dits import WanVideoConfig
+from fastvideo.configs.models.encoders import CLIPVisionConfig
+from fastvideo.configs.models.vaes import WanVAEConfig
+from fastvideo.configs.pipelines.base import PipelineConfig
+from fastvideo.configs.pipelines.wan import t5_postprocess_text, T5Config, BaseEncoderOutput
+
+import torch
+from collections.abc import Callable
+
+
+@dataclass
+class TurboDiffusionT2VConfig(PipelineConfig):
+    """Base configuration for TurboDiffusion T2V pipeline.
+    
+    Uses RCM scheduler with sigma_max=80 for 1-4 step generation.
+    No boundary_ratio (single model, no switching).
+    """
+    # DiT
+    dit_config: DiTConfig = field(default_factory=WanVideoConfig)
+    # VAE
+    vae_config: VAEConfig = field(default_factory=WanVAEConfig)
+    vae_tiling: bool = False
+    vae_sp: bool = False
+
+    # Denoising stage
+    flow_shift: float | None = 3.0
+
+    # No boundary_ratio for T2V (single model)
+    boundary_ratio: float | None = None
+
+    # Text encoding stage
+    text_encoder_configs: tuple[EncoderConfig, ...] = field(
+        default_factory=lambda: (T5Config(), ))
+    postprocess_text_funcs: tuple[Callable[[BaseEncoderOutput], torch.Tensor],
+                                  ...] = field(default_factory=lambda:
+                                               (t5_postprocess_text, ))
+
+    # Precision for each component
+    precision: str = "bf16"
+    vae_precision: str = "fp32"
+    text_encoder_precisions: tuple[str, ...] = field(
+        default_factory=lambda: ("fp32", ))
+
+    # self-forcing params
+    warp_denoising_step: bool = True
+
+    def __post_init__(self):
+        self.vae_config.load_encoder = False
+        self.vae_config.load_decoder = True
+        # Ensure no boundary_ratio is set in dit_config
+        self.dit_config.boundary_ratio = None
+
+
+@dataclass
+class TurboDiffusionT2V_1_3B_Config(TurboDiffusionT2VConfig):
+    """Configuration for TurboDiffusion T2V 1.3B model."""
+    pass
+
+
+@dataclass
+class TurboDiffusionT2V_14B_Config(TurboDiffusionT2VConfig):
+    """Configuration for TurboDiffusion T2V 14B model.
+    
+    Uses same config as 1.3B but with higher flow_shift for 14B model.
+    """
+    flow_shift: float | None = 5.0
+
+
+@dataclass
+class TurboDiffusionI2VConfig(PipelineConfig):
+    """Base configuration for TurboDiffusion I2V pipeline.
+    
+    Uses RCM scheduler with sigma_max=200 for 1-4 step generation.
+    Uses boundary_ratio=0.9 for high-noise to low-noise model switching.
+    """
+    # DiT
+    dit_config: DiTConfig = field(default_factory=WanVideoConfig)
+    # VAE
+    vae_config: VAEConfig = field(default_factory=WanVAEConfig)
+    vae_tiling: bool = False
+    vae_sp: bool = False
+
+    # Denoising stage
+    flow_shift: float | None = 5.0
+
+    boundary_ratio: float | None = 0.9
+
+    # Text encoding stage
+    text_encoder_configs: tuple[EncoderConfig, ...] = field(
+        default_factory=lambda: (T5Config(), ))
+    postprocess_text_funcs: tuple[Callable[[BaseEncoderOutput], torch.Tensor],
+                                  ...] = field(default_factory=lambda:
+                                               (t5_postprocess_text, ))
+
+    # Image encoder for I2V
+    image_encoder_config: EncoderConfig = field(
+        default_factory=CLIPVisionConfig)
+    image_encoder_precision: str = "fp32"
+
+    # Precision for each component
+    precision: str = "bf16"
+    vae_precision: str = "fp32"
+    text_encoder_precisions: tuple[str, ...] = field(
+        default_factory=lambda: ("fp32", ))
+
+    # self-forcing params
+    warp_denoising_step: bool = True
+
+    def __post_init__(self):
+        self.vae_config.load_encoder = True
+        self.vae_config.load_decoder = True
+        self.dit_config.boundary_ratio = self.boundary_ratio
+
+
+@dataclass
+class TurboDiffusionI2V_A14B_Config(TurboDiffusionI2VConfig):
+    """Configuration for TurboDiffusion I2V A14B model."""
+    pass

fastvideo/configs/sample/base.py

@@ -223,7 +223,7 @@ class SamplingParam:
             help="Path to input image for image-to-video generation",
         )
         parser.add_argument(
-            "--video_path",
+            "--video-path",
             type=str,
             default=SamplingParam.video_path,
             help="Path to input video for video-to-video generation",

fastvideo/configs/sample/registry.py

@@ -26,6 +26,11 @@ from fastvideo.configs.sample.wan import (
     SelfForcingWan2_2_T2V_A14B_480P_SamplingParam,
     MatrixGame2_SamplingParam,
 )
+from fastvideo.configs.sample.turbodiffusion import (
+    TurboDiffusionT2V_1_3B_SamplingParam,
+    TurboDiffusionT2V_14B_SamplingParam,
+    TurboDiffusionI2V_A14B_SamplingParam,
+)
 # isort: on
 from fastvideo.logger import init_logger
 from fastvideo.utils import (maybe_download_model_index,
@@ -34,36 +39,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":
@@ -80,9 +97,20 @@ SAMPLING_PARAM_REGISTRY: dict[str, Any] = {
     Cosmos_Predict2_2B_Video2World_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":
+    TurboDiffusionT2V_1_3B_SamplingParam,
+    "loayrashid/TurboWan2.1-T2V-14B-Diffusers":
+    TurboDiffusionT2V_14B_SamplingParam,
+    "loayrashid/TurboWan2.2-I2V-A14B-Diffusers":
+    TurboDiffusionI2V_A14B_SamplingParam,
 
     # Add other specific weight variants
 }
@@ -105,6 +133,8 @@ SAMPLING_PARAM_DETECTOR: dict[str, Callable[[str], bool]] = {
     lambda id: "wancausaldmdpipeline" in id.lower(),
     "matrixgame":
     lambda id: "matrixgame" in id.lower() or "matrix-game" in id.lower(),
+    "turbodiffusion":
+    lambda id: "turbodiffusion" in id.lower() or "turbowan" in id.lower(),
     # Add other pipeline architecture detectors
 }
 
@@ -121,6 +151,8 @@ SAMPLING_FALLBACK_PARAM: dict[str, Any] = {
     "wancausaldmdpipeline": SelfForcingWan2_1_T2V_1_3B_480P_SamplingParam,
     "stepvideo": StepVideoT2VSamplingParam,
     "matrixgame": MatrixGame2_SamplingParam,
+    "turbodiffusion":
+    TurboDiffusionT2V_1_3B_SamplingParam,  # Default to T2V for fallback
     # Other fallbacks by architecture
 }
 

fastvideo/configs/sample/turbodiffusion.py

@@ -0,0 +1,73 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+TurboDiffusion sampling parameters.
+
+TurboDiffusion uses RCM (recurrent Consistency Model) scheduler for
+1-4 step video generation with no classifier-free guidance.
+"""
+from dataclasses import dataclass
+
+from fastvideo.configs.sample.base import SamplingParam
+
+
+@dataclass
+class TurboDiffusionT2V_1_3B_SamplingParam(SamplingParam):
+    """Sampling parameters for TurboDiffusion T2V 1.3B model.
+    
+    Uses 4-step RCM sampling with guidance_scale=1.0 (no CFG).
+    """
+    # Video parameters
+    height: int = 480
+    width: int = 832
+    num_frames: int = 81
+    fps: int = 16
+
+    # Denoising stage - TurboDiffusion uses 1-4 steps with no CFG
+    guidance_scale: float = 1.0
+    num_inference_steps: int = 4
+
+    # No negative prompt needed for TurboDiffusion (no CFG)
+    negative_prompt: str | None = None
+
+
+@dataclass
+class TurboDiffusionT2V_14B_SamplingParam(SamplingParam):
+    """Sampling parameters for TurboDiffusion T2V 14B model.
+    
+    Uses 4-step RCM sampling with guidance_scale=1.0 (no CFG).
+    """
+    # Video parameters (720p for 14B)
+    height: int = 720
+    width: int = 1280
+    num_frames: int = 81
+    fps: int = 16
+
+    # Denoising stage - TurboDiffusion uses 1-4 steps with no CFG
+    guidance_scale: float = 1.0
+    num_inference_steps: int = 4
+
+    # No negative prompt needed for TurboDiffusion (no CFG)
+    negative_prompt: str | None = None
+
+
+@dataclass
+class TurboDiffusionI2V_A14B_SamplingParam(SamplingParam):
+    """Sampling parameters for TurboDiffusion I2V A14B model.
+    
+    Uses 4-step RCM sampling with dual-model switching (high/low noise).
+    """
+    # Video parameters (720p for A14B I2V)
+    height: int = 720
+    width: int = 1280
+    num_frames: int = 81
+    fps: int = 16
+
+    # Denoising stage - TurboDiffusion uses 1-4 steps with no CFG
+    guidance_scale: float = 1.0
+    num_inference_steps: int = 4
+
+    # Note: boundary_ratio is set in the pipeline config (TurboDiffusionI2VConfig),
+    # not here. This keeps sampling params and pipeline config separate.
+
+    # No negative prompt needed for TurboDiffusion (no CFG)
+    negative_prompt: str | None = None

fastvideo/envs.py

@@ -38,8 +38,10 @@ if TYPE_CHECKING:
     FASTVIDEO_TORCH_PROFILE_REGIONS: str = ""
     FASTVIDEO_SERVER_DEV_MODE: bool = False
     FASTVIDEO_STAGE_LOGGING: bool = False
+    FASTVIDEO_DENOISING_PERF_LOGGING: bool = False
     FASTVIDEO_HOST_IP: str = ""
     FASTVIDEO_LOOPBACK_IP: str = ""
+    FASTVIDEO_DISABLE_PIN_MEMORY: str | None = None
 
 
 def get_default_cache_root() -> str:
@@ -136,6 +138,10 @@ environment_variables: dict[str, Callable[[], Any]] = {
     "FASTVIDEO_LOOPBACK_IP":
     lambda: os.getenv("FASTVIDEO_LOOPBACK_IP", ""),
 
+    # Disable pinned memory (e.g., on platforms that do not support it)
+    "FASTVIDEO_DISABLE_PIN_MEMORY":
+    lambda: os.getenv("FASTVIDEO_DISABLE_PIN_MEMORY", None),
+
     # Number of GPUs per worker in Ray, if it is set to be a fraction,
     # it allows ray to schedule multiple actors on a single GPU,
     # so that users can colocate other actors on the same GPUs as FastVideo.
@@ -275,6 +281,10 @@ environment_variables: dict[str, Callable[[], Any]] = {
     # taken for each stage
     "FASTVIDEO_STAGE_LOGGING":
     lambda: bool(int(os.getenv("FASTVIDEO_STAGE_LOGGING", "0"))),
+
+    # Enable per-step denoising perf logging hooks
+    "FASTVIDEO_DENOISING_PERF_LOGGING":
+    lambda: bool(int(os.getenv("FASTVIDEO_DENOISING_PERF_LOGGING", "0"))),
 }
 
 # end-env-vars-definition

fastvideo/fastvideo_args.py

@@ -132,7 +132,7 @@ class FastVideoArgs:
 
     # CPU offload parameters
     dit_cpu_offload: bool = True
-    use_fsdp_inference: bool = True
+    use_fsdp_inference: bool = False
     dit_layerwise_offload: bool = False
     text_encoder_cpu_offload: bool = True
     image_encoder_cpu_offload: bool = True
@@ -431,7 +431,9 @@ class FastVideoArgs:
             "--use-fsdp-inference",
             action=StoreBoolean,
             help=
-            "Use FSDP for inference by sharding the model weights. Latency is very low due to prefetch--enable if run out of memory.",
+            "Use FSDP for inference by sharding the model weights. FSDP helps reduce GPU memory usage but may introduce"
+            +
+            " weight transfer overhead depending on the specific setup. Enable if run out of memory.",
         )
         parser.add_argument(
             "--text-encoder-cpu-offload",
@@ -606,6 +608,7 @@ class FastVideoArgs:
     def check_fastvideo_args(self) -> None:
         """Validate inference arguments for consistency"""
         from fastvideo.platforms import current_platform
+        from fastvideo.pin_memory import is_pin_memory_available
 
         if current_platform.is_mps():
             self.use_fsdp_inference = False
@@ -686,6 +689,11 @@ class FastVideoArgs:
                 self.pipeline_config.vae_config.load_encoder = True
             self.preprocess_config.check_preprocess_config()
 
+        if self.pin_cpu_memory and not is_pin_memory_available():
+            logger.warning("Pinned memory is unavailable on this system; "
+                           "disabling pin_cpu_memory.")
+            self.pin_cpu_memory = False
+
 
 _current_fastvideo_args = None
 

fastvideo/layers/lora/linear.py

@@ -18,6 +18,7 @@ from fastvideo.layers.linear import (ColumnParallelLinear, LinearBase,
                                      QKVParallelLinear, ReplicatedLinear,
                                      RowParallelLinear)
 from fastvideo.layers.vocab_parallel_embedding import VocabParallelEmbedding
+from fastvideo.pin_memory import is_pin_memory_available
 from fastvideo.utils import get_mixed_precision_state
 
 torch._dynamo.config.recompile_limit = 16
@@ -155,8 +156,11 @@ class BaseLayerWithLoRA(nn.Module):
                         get_local_torch_device(),
                         non_blocking=True).full_tensor().to(current_device))
 
-            offload_policy = CPUOffloadPolicy() if "cpu" in str(
-                current_device) else OffloadPolicy()
+            if "cpu" in str(current_device):
+                offload_policy = CPUOffloadPolicy(
+                    pin_memory=is_pin_memory_available())
+            else:
+                offload_policy = OffloadPolicy()
             mp_policy = get_mixed_precision_state().mp_policy
 
             self.base_layer = fully_shard(unsharded_base_layer,

fastvideo/models/dits/longcat.py

@@ -1,9 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 """
 Native LongCat Video DiT implementation using FastVideo conventions.
-
-This is a Phase 2 reimplementation that replaces the third_party wrapper
-with native FastVideo layers for better performance and integration.
 """
 
 from typing import Any
@@ -129,7 +126,7 @@ class TimestepEmbedder(nn.Module):
         # Sinusoidal embedding in FP32
         t_freq = self.timestep_embedding(t.flatten(), self.frequency_embedding_size)
         
-        # Cast to model dtype before MLP
+        # Cast to model dtype before MLP (matching original LongCat)
         # Handle LoRA wrapper if present
         linear_layer = self.linear_1.base_layer if hasattr(self.linear_1, 'base_layer') else self.linear_1
         target_dtype = linear_layer.weight.dtype
@@ -166,13 +163,14 @@ class CaptionEmbedder(nn.Module):
         self.text_tokens_zero_pad = text_tokens_zero_pad
 
         # Two-layer MLP using ReplicatedLinear
+        # CRITICAL: Original LongCat uses GELU(approximate="tanh"), NOT SiLU!
         self.linear_1 = ReplicatedLinear(
             caption_channels,
             hidden_size,
             bias=True,
             params_dtype=dtype,
         )
-        self.act = nn.SiLU()
+        self.act = nn.GELU(approximate="tanh")  # Match original LongCat
         self.linear_2 = ReplicatedLinear(
             hidden_size,
             hidden_size,
@@ -268,10 +266,19 @@ class LongCatSelfAttention(nn.Module):
         self,
         x: torch.Tensor,              # [B, N, C]
         latent_shape: tuple,          # (T, H, W)
+        num_cond_latents: int = 0,    # Number of conditioning latent frames (for I2V)
+        return_kv: bool = False,      # Return K/V for caching
         **kwargs
-    ) -> torch.Tensor:
+    ) -> torch.Tensor | tuple:
         """
         Forward pass with 3D RoPE and optional BSA.
+        
+        For I2V mode (num_cond_latents > 0):
+        - Conditioned tokens only attend to themselves
+        - Noise tokens attend to ALL tokens (cond + noise)
+        
+        Args:
+            return_kv: If True, return (output, (k_cache, v_cache)) for KV caching
         """
         B, N, C = x.shape
         T, H, W = latent_shape
@@ -290,6 +297,12 @@ class LongCatSelfAttention(nn.Module):
         q = self.q_norm(q)
         k = self.k_norm(k)
 
+        # Save pre-RoPE K/V for cache if requested (before RoPE is applied)
+        if return_kv:
+            # [B, N, num_heads, head_dim] -> [B, num_heads, N, head_dim]
+            k_cache = k.transpose(1, 2).clone()
+            v_cache = v.transpose(1, 2).clone()
+
         # For RoPE: need [B, num_heads, N, head_dim]
         q_rope = q.transpose(1, 2)
         k_rope = k.transpose(1, 2)
@@ -301,6 +314,50 @@ class LongCatSelfAttention(nn.Module):
         q = q_rope.transpose(1, 2)
         k = k_rope.transpose(1, 2)
 
+        # === I2V Split Attention ===
+        # For I2V, conditioned tokens and noise tokens are processed separately
+        if num_cond_latents > 0:
+            # Calculate number of conditioned tokens (cond_latents * spatial_tokens_per_frame)
+            num_cond_tokens = num_cond_latents * (N // T)
+            
+            # Conditioned tokens: only attend to themselves (same seq length, use self.attn)
+            q_cond = q[:, :num_cond_tokens].contiguous()
+            k_cond = k[:, :num_cond_tokens].contiguous()
+            v_cond = v[:, :num_cond_tokens].contiguous()
+            out_cond, _ = self.attn(q_cond, k_cond, v_cond)
+            
+            # Noise tokens: attend to ALL tokens (different seq lengths!)
+            # Need to use flash attention directly since q has different length than k/v
+            q_noise = q[:, num_cond_tokens:].contiguous()  # [B, N_noise, num_heads, head_dim]
+            # k, v are full: [B, N, num_heads, head_dim]
+            
+            # Transpose for flash attention: [B, num_heads, seq, head_dim]
+            q_noise_t = q_noise.transpose(1, 2)
+            k_t = k.transpose(1, 2)
+            v_t = v.transpose(1, 2)
+            
+            # Use scaled dot product attention (handles different q/kv lengths)
+            out_noise_t = torch.nn.functional.scaled_dot_product_attention(
+                q_noise_t, k_t, v_t, 
+                attn_mask=None,
+                dropout_p=0.0,
+                is_causal=False
+            )  # [B, num_heads, N_noise, head_dim]
+            
+            # Transpose back: [B, N_noise, num_heads, head_dim]
+            out_noise = out_noise_t.transpose(1, 2)
+            
+            # Merge conditioned and noise outputs
+            out = torch.cat([out_cond, out_noise], dim=1)
+            
+            # Reshape and project out
+            out = out.reshape(B, N, C)
+            out, _ = self.to_out(out)
+            
+            if return_kv:
+                return out, (k_cache, v_cache)
+            return out
+
         # === Attention: BSA or standard ===
         if self.enable_bsa and T > 1:  # Only use BSA for multi-frame videos
             # BSA expects [B, H, S, D] format
@@ -348,6 +405,96 @@ class LongCatSelfAttention(nn.Module):
         out = out.reshape(B, N, C)
         out, _ = self.to_out(out)
 
+        if return_kv:
+            return out, (k_cache, v_cache)
+        return out
+
+    def forward_with_kv_cache(
+        self,
+        x: torch.Tensor,              # [B, N_noise, C] - only noise tokens
+        latent_shape: tuple,          # (T_noise, H, W) - shape for noise only
+        num_cond_latents: int,        # Number of conditioning latent frames
+        kv_cache: tuple,              # (k_cond, v_cond) - [B, heads, N_cond, head_dim]
+    ) -> torch.Tensor:
+        """
+        Forward using cached K/V from conditioning frames.
+        
+        x contains only NOISE tokens.
+        kv_cache contains pre-computed K/V for CONDITIONING tokens.
+        
+        CRITICAL: RoPE positions for noise tokens must start AFTER conditioning.
+        We achieve this by padding Q with dummy tokens for conditioning positions,
+        applying RoPE to the full sequence, then extracting only noise token Q.
+        """
+        B, N, C = x.shape
+        T, H, W = latent_shape
+
+        k_cache, v_cache = kv_cache
+
+        # Handle batch size mismatch (cache might be smaller for CFG)
+        # When using CFG, latent_model_input is doubled [neg, pos], but cache is for original batch
+        if k_cache.shape[0] != B:
+            # Expand cache to match input batch size
+            # For CFG: repeat the cache for both negative and positive branches
+            repeat_factor = B // k_cache.shape[0]
+            k_cache = k_cache.repeat(repeat_factor, 1, 1, 1)
+            v_cache = v_cache.repeat(repeat_factor, 1, 1, 1)
+
+        # Project to Q/K/V for noise tokens
+        q, _ = self.to_q(x)
+        k, _ = self.to_k(x)
+        v, _ = self.to_v(x)
+
+        # Reshape to heads: [B, N, num_heads, head_dim]
+        q = q.view(B, N, self.num_heads, self.head_dim)
+        k = k.view(B, N, self.num_heads, self.head_dim)
+        v = v.view(B, N, self.num_heads, self.head_dim)
+
+        # Per-head RMS normalization
+        q = self.q_norm(q)
+        k = self.k_norm(k)
+
+        # Transpose for RoPE: [B, heads, N, head_dim]
+        q_rope = q.transpose(1, 2)
+        k_rope = k.transpose(1, 2)
+        v = v.transpose(1, 2)
+
+        # CRITICAL: Apply RoPE with correct positional offset
+        # Noise frame queries need positions starting from num_cond_latents
+        # Following the original LongCat approach:
+        # 1. Pad Q with dummy tokens matching k_cache shape
+        # 2. Apply RoPE to full sequence (T_cond + T_noise)
+        # 3. Extract only the noise portion of Q
+        
+        # Create dummy Q padding to fill conditioning positions
+        # k_cache shape: [B, heads, N_cond, head_dim]
+        q_padding = torch.cat([torch.empty_like(k_cache), q_rope], dim=2).contiguous()
+        
+        # Concatenate cached K with noise K for RoPE
+        k_full = torch.cat([k_cache, k_rope], dim=2)
+        v_full = torch.cat([v_cache, v], dim=2)
+        
+        # Apply RoPE to full sequence (includes both cond and noise positions)
+        # Grid size: (T_cond + T_noise, H, W)
+        full_T = num_cond_latents + T
+        q_padding, k_full = self.rope_3d(q_padding, k_full, grid_size=(full_T, H, W))
+        
+        # Extract only the noise portion of Q (last N tokens)
+        q_rope = q_padding[:, :, -N:].contiguous()
+
+        # Run attention: Q_noise attends to full K/V (cond + noise)
+        out = torch.nn.functional.scaled_dot_product_attention(
+            q_rope, k_full, v_full,
+            attn_mask=None,
+            dropout_p=0.0,
+            is_causal=False
+        )  # [B, heads, N_noise, head_dim]
+
+        # Transpose back: [B, N_noise, heads, head_dim]
+        out = out.transpose(1, 2)
+        out = out.reshape(B, N, C)
+        out, _ = self.to_out(out)
+
         return out
 
 
@@ -394,6 +541,8 @@ class LongCatCrossAttention(nn.Module):
         self,
         x: torch.Tensor,              # [B, N_img, C]
         context: torch.Tensor,        # [B, N_text, C]
+        latent_shape: tuple = None,   # (T, H, W) - needed for I2V
+        num_cond_latents: int = 0,    # Number of conditioning latent frames (for I2V)
         **kwargs
     ) -> torch.Tensor:
         """
@@ -402,9 +551,57 @@ class LongCatCrossAttention(nn.Module):
         Args:
             x: Image tokens [B, N_img, C]
             context: Text tokens [B, N_text, C] (standard padded format)
+            latent_shape: (T, H, W) - needed for calculating num_cond_tokens
+            num_cond_latents: Number of conditioning latent frames (for I2V)
+            
+        For I2V mode (num_cond_latents > 0):
+        - Conditioned tokens get ZERO cross-attention output
+        - Only noise tokens get cross-attention with text
         """
         B, N_img, C = x.shape
 
+        # === I2V: Only noise tokens get cross-attention ===
+        if num_cond_latents > 0 and latent_shape is not None:
+            T, H, W = latent_shape
+            num_cond_tokens = num_cond_latents * (N_img // T)
+            
+            # Only process noise tokens
+            x_noise = x[:, num_cond_tokens:]  # [B, N_noise, C]
+            
+            # Project Q, K, V for noise tokens only
+            q, _ = self.to_q(x_noise)
+            k, _ = self.to_k(context)
+            v, _ = self.to_v(context)
+            
+            N_text = context.shape[1]
+            N_noise = x_noise.shape[1]
+
+            # Reshape to heads
+            q = q.view(B, N_noise, self.num_heads, self.head_dim)
+            k = k.view(B, N_text, self.num_heads, self.head_dim)
+            v = v.view(B, N_text, self.num_heads, self.head_dim)
+
+            # Per-head RMS normalization
+            q = self.q_norm(q)
+            k = self.k_norm(k)
+
+            # Run cross-attention
+            out_noise = self.attn(q, k, v)  # [B, N_noise, num_heads, head_dim]
+            out_noise = out_noise.reshape(B, N_noise, C)
+            out_noise, _ = self.to_out(out_noise)
+            
+            # Conditioned tokens get zero output
+            out_cond = torch.zeros(
+                (B, num_cond_tokens, C), 
+                dtype=out_noise.dtype, 
+                device=out_noise.device
+            )
+            
+            # Merge
+            out = torch.cat([out_cond, out_noise], dim=1)
+            return out
+
+        # === Standard cross-attention ===
         # Project Q, K, V (standard cross-attention like WanVideo/StepVideo/Cosmos)
         q, _ = self.to_q(x)
         k, _ = self.to_k(context)
@@ -475,19 +672,19 @@ class LongCatSwiGLUFFN(nn.Module):
 
 def modulate_fp32(norm: nn.Module, x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
     """
-    Apply modulation in FP32 for numerical stability (matching original LongCat).
+    Apply modulation in FP32 for numerical stability.
     
-    shift and scale should already be FP32 from torch.amp.autocast context.
+    Converts inputs to FP32 for the modulation operation, then casts back.
     """
-    # Ensure modulation params are FP32 (should be from autocast)
-    assert shift.dtype == torch.float32 and scale.dtype == torch.float32, \
-        f"shift and scale must be FP32, got {shift.dtype} and {scale.dtype}"
-    
     orig_dtype = x.dtype
     
+    # Convert to FP32 for numerical stability
+    shift_fp32 = shift.float()
+    scale_fp32 = scale.float()
+    
     # Normalize and modulate in FP32
     x_norm = norm(x.to(torch.float32))
-    x_mod = x_norm * (scale + 1) + shift
+    x_mod = x_norm * (scale_fp32 + 1) + shift_fp32
     
     return x_mod.to(orig_dtype)
 
@@ -568,10 +765,21 @@ class LongCatTransformerBlock(nn.Module):
         context: torch.Tensor,        # [B, N_text, C]
         t: torch.Tensor,              # [B, T, C_t]
         latent_shape: tuple,          # (T, H, W)
+        num_cond_latents: int = 0,    # Number of conditioning latent frames (for I2V)
+        return_kv: bool = False,      # Return K/V for caching
+        kv_cache: tuple | None = None,  # Pre-computed K/V cache
+        skip_crs_attn: bool = False,  # Skip cross-attention (for cache init)
         **kwargs
-    ) -> torch.Tensor:
+    ) -> torch.Tensor | tuple:
         """
         Forward pass with AdaLN modulation.
+        
+        Args:
+            num_cond_latents: For I2V, number of conditioning latent frames.
+                              These frames use split attention behavior.
+            return_kv: If True, return (x, (k_cache, v_cache))
+            kv_cache: Pre-computed K/V from conditioning frames
+            skip_crs_attn: If True, skip cross-attention (used during cache init)
         """
         B, N, C = x.shape
         T, H, W = latent_shape
@@ -592,17 +800,47 @@ class LongCatTransformerBlock(nn.Module):
         x_norm = modulate_fp32(self.norm_attn, x.view(B, T, -1, C), shift_msa, scale_msa)
         x_norm = x_norm.view(B, N, C)
 
-        attn_out = self.self_attn(x_norm, latent_shape=latent_shape)
+        # Handle KV cache
+        if kv_cache is not None:
+            # Move cache to device if offloaded
+            kv_cache = (kv_cache[0].to(x.device), kv_cache[1].to(x.device))
+            attn_out = self.self_attn.forward_with_kv_cache(
+                x_norm,
+                latent_shape=latent_shape,
+                num_cond_latents=num_cond_latents,
+                kv_cache=kv_cache,
+            )
+            kv_cache_new = None  # Don't return cache when using cache
+        else:
+            attn_result = self.self_attn(
+                x_norm, 
+                latent_shape=latent_shape,
+                num_cond_latents=num_cond_latents,
+                return_kv=return_kv,
+            )
+            if return_kv:
+                attn_out, kv_cache_new = attn_result
+            else:
+                attn_out = attn_result
+                kv_cache_new = None
 
         # Residual with gating (CRITICAL: FP32 like original, then cast back)
         with torch.amp.autocast(device_type='cuda', dtype=torch.float32):
             x = x + (gate_msa * attn_out.view(B, T, -1, C)).view(B, N, C)
         x = x.to(x_orig_dtype)
 
-        # === Cross-Attention ===
-        x_norm_cross = self.norm_cross(x)
-        cross_out = self.cross_attn(x_norm_cross, context)
-        x = x + cross_out
+        # === Cross-Attention (skip if requested) ===
+        if not skip_crs_attn:
+            x_norm_cross = self.norm_cross(x)
+            # When using KV cache, no need for num_cond_latents in cross-attn
+            cross_num_cond = 0 if kv_cache is not None else num_cond_latents
+            cross_out = self.cross_attn(
+                x_norm_cross, 
+                context,
+                latent_shape=latent_shape,
+                num_cond_latents=cross_num_cond
+            )
+            x = x + cross_out
 
         # === FFN ===
         x_norm_ffn = modulate_fp32(self.norm_ffn, x.view(B, T, -1, C), shift_mlp, scale_mlp)
@@ -615,6 +853,8 @@ class LongCatTransformerBlock(nn.Module):
             x = x + (gate_mlp * ffn_out.view(B, T, -1, C)).view(B, N, C)
         x = x.to(x_orig_dtype)
 
+        if return_kv:
+            return x, kv_cache_new
         return x
 
 
@@ -670,16 +910,12 @@ class FinalLayer(nn.Module):
         B, N, C = x.shape
         T, _, _ = latent_shape
 
-        # AdaLN modulation (FP32 for stability like original)
-        with torch.amp.autocast(device_type='cuda', dtype=torch.float32):
-            t_mod = self.adaln_act(t)
-            mod_params, _ = self.adaln_linear(t_mod)
-            # Ensure FP32 output (needed when LoRA is applied)
-            if mod_params.dtype != torch.float32:
-                mod_params = mod_params.float()
-            shift, scale = mod_params.unsqueeze(2).chunk(2, dim=-1)
+        # AdaLN modulation
+        t_mod = self.adaln_act(t)
+        mod_params, _ = self.adaln_linear(t_mod)
+        shift, scale = mod_params.unsqueeze(2).chunk(2, dim=-1)
 
-        # Modulate
+        # Modulate (converts to FP32 internally for stability)
         x = modulate_fp32(self.norm, x.view(B, T, -1, C), shift, scale)
         x = x.reshape(B, N, C)
 
@@ -696,8 +932,6 @@ class FinalLayer(nn.Module):
 class LongCatTransformer3DModel(CachableDiT):
     """
     Native LongCat Video Transformer using FastVideo layers.
-    
-    This is a Phase 2 implementation that replaces third_party dependencies.
     """
 
     # FSDP sharding: shard at each transformer block
@@ -789,13 +1023,28 @@ class LongCatTransformer3DModel(CachableDiT):
         encoder_attention_mask: torch.Tensor | None = None,  # [B, N_text]
         encoder_hidden_states_image: torch.Tensor | list[torch.Tensor] | None = None,
         guidance: float | None = None,         # Unused, for API compatibility
+        num_cond_latents: int = 0,             # For I2V: number of conditioning latent frames
+        # === KV Cache Parameters ===
+        return_kv: bool = False,               # If True, return (output, kv_cache_dict)
+        kv_cache_dict: dict | None = None,     # Pre-computed {block_idx: (k, v)}
+        skip_crs_attn: bool = False,           # Skip cross-attention (for cache init)
+        offload_kv_cache: bool = False,        # Move cache to CPU after compute
         **kwargs
-    ) -> torch.Tensor:
+    ) -> torch.Tensor | tuple[torch.Tensor, dict]:
         """
         Forward pass with FastVideo parameter ordering.
         
         NOTE: This follows FastVideo convention:
               (hidden_states, encoder_hidden_states, timestep)
+              
+        Args:
+            num_cond_latents: For I2V, number of conditioning latent frames.
+                              These frames are treated as "clean" (timestep=0)
+                              and use split attention behavior.
+            return_kv: If True, return (output, kv_cache_dict)
+            kv_cache_dict: Pre-computed K/V cache {block_idx: (k, v)}
+            skip_crs_attn: If True, skip cross-attention (for cache init)
+            offload_kv_cache: If True, move cache to CPU after compute
         """
         B, _, T, H, W = hidden_states.shape
 
@@ -825,12 +1074,31 @@ class LongCatTransformer3DModel(CachableDiT):
             encoder_attention_mask=encoder_attention_mask
         )  # [B, N_text, C]
 
-        # 4. Transformer blocks
+        # 4. Transformer blocks with optional KV cache
+        kv_cache_dict_ret = {} if return_kv else None
+        
         for i, block in enumerate(self.blocks):
-            x = block(
+            # Get cache for this block if available
+            block_kv_cache = kv_cache_dict.get(i, None) if kv_cache_dict else None
+            
+            block_out = block(
                 x, context, t,
-                latent_shape=(N_t, N_h, N_w)
+                latent_shape=(N_t, N_h, N_w),
+                num_cond_latents=num_cond_latents,
+                return_kv=return_kv,
+                kv_cache=block_kv_cache,
+                skip_crs_attn=skip_crs_attn,
             )
+            
+            if return_kv:
+                x, kv_cache = block_out
+                # Store cache
+                if offload_kv_cache:
+                    kv_cache_dict_ret[i] = (kv_cache[0].cpu(), kv_cache[1].cpu())
+                else:
+                    kv_cache_dict_ret[i] = (kv_cache[0].contiguous(), kv_cache[1].contiguous())
+            else:
+                x = block_out
 
         # 5. Output projection
         output = self.final_layer(x, t, latent_shape=(N_t, N_h, N_w))
@@ -841,6 +1109,8 @@ class LongCatTransformer3DModel(CachableDiT):
         # Cast to float32 for better accuracy (as per original)
         output = output.to(torch.float32)
 
+        if return_kv:
+            return output, kv_cache_dict_ret
         return output
 
     def unpatchify(self, x: torch.Tensor, N_t: int, N_h: int, N_w: int) -> torch.Tensor:

fastvideo/models/layerwise_offload.py

@@ -11,6 +11,8 @@ from typing import Dict, Set, Optional, Tuple
 
 import torch
 
+from fastvideo.pin_memory import is_pin_memory_available
+
 
 class LayerwiseOffloadManager:
     """A lightweight layerwise CPU offload manager.
@@ -33,6 +35,8 @@ class LayerwiseOffloadManager:
         self.module_list_attr = module_list_attr
         self.num_layers = int(num_layers)
         self.pin_cpu_memory = bool(pin_cpu_memory)
+        if self.pin_cpu_memory and not is_pin_memory_available():
+            self.pin_cpu_memory = False
 
         self.enabled = bool(enabled and torch.cuda.is_available())
         self.device = (

fastvideo/models/loader/fsdp_load.py

@@ -20,6 +20,7 @@ from torch.distributed.fsdp import (CPUOffloadPolicy, FSDPModule,
 from torch.nn.modules.module import _IncompatibleKeys
 
 from fastvideo.logger import init_logger
+from fastvideo.pin_memory import is_pin_memory_available
 from fastvideo.models.loader.utils import (get_param_names_mapping,
                                            hf_to_custom_state_dict)
 from fastvideo.models.loader.weight_utils import safetensors_weights_iterator
@@ -212,6 +213,11 @@ def shard_model(
         "mp_policy": mp_policy,
     }
     if cpu_offload:
+        if pin_cpu_memory and not is_pin_memory_available():
+            logger.warning(
+                "Pinned memory is unavailable; disabling pin_cpu_memory for "
+                "FSDP offload.")
+            pin_cpu_memory = False
         fsdp_kwargs["offload_policy"] = CPUOffloadPolicy(
             pin_memory=pin_cpu_memory)
 

fastvideo/models/registry.py

@@ -30,8 +30,8 @@ _TEXT_TO_VIDEO_DIT_MODELS = {
     "CausalWanTransformer3DModel": ("dits", "causal_wanvideo", "CausalWanTransformer3DModel"),
     "StepVideoModel": ("dits", "stepvideo", "StepVideoModel"),
     "CosmosTransformer3DModel": ("dits", "cosmos", "CosmosTransformer3DModel"),
-    "LongCatVideoTransformer3DModel": ("dits", "longcat_video_dit", "LongCatVideoTransformer3DModel"),  # Wrapper (Phase 1)
-    "LongCatTransformer3DModel": ("dits", "longcat", "LongCatTransformer3DModel"),  # Native (Phase 2)
+    "LongCatVideoTransformer3DModel": ("dits", "longcat_video_dit", "LongCatVideoTransformer3DModel"),
+    "LongCatTransformer3DModel": ("dits", "longcat", "LongCatTransformer3DModel"),
 }
 
 _IMAGE_TO_VIDEO_DIT_MODELS = {

fastvideo/models/schedulers/adapter.py

@@ -0,0 +1,49 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+Scheduler adapter interfaces for unified denoising.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Protocol
+
+import torch
+
+
+class SchedulerAdapter(Protocol):
+    def scale_model_input(self, latents: torch.Tensor,
+                          t: torch.Tensor) -> torch.Tensor:
+        ...
+
+    def step(self, noise_pred: torch.Tensor, t: torch.Tensor,
+             latents: torch.Tensor, **kwargs: Any) -> Any:
+        ...
+
+    def add_noise(self, latents: torch.Tensor, noise: torch.Tensor,
+                  t: torch.Tensor) -> torch.Tensor:
+        ...
+
+    def set_timesteps(self, num_steps: int, device: torch.device | None = None,
+                      **kwargs: Any) -> Any:
+        ...
+
+
+class DefaultSchedulerAdapter:
+    def __init__(self, scheduler: Any) -> None:
+        self.scheduler = scheduler
+
+    def scale_model_input(self, latents: torch.Tensor,
+                          t: torch.Tensor) -> torch.Tensor:
+        return self.scheduler.scale_model_input(latents, t)
+
+    def step(self, noise_pred: torch.Tensor, t: torch.Tensor,
+             latents: torch.Tensor, **kwargs: Any) -> Any:
+        return self.scheduler.step(noise_pred, t, latents, **kwargs)
+
+    def add_noise(self, latents: torch.Tensor, noise: torch.Tensor,
+                  t: torch.Tensor) -> torch.Tensor:
+        return self.scheduler.add_noise(latents, noise, t)
+
+    def set_timesteps(self, num_steps: int, device: torch.device | None = None,
+                      **kwargs: Any) -> Any:
+        return self.scheduler.set_timesteps(num_steps, device=device, **kwargs)

fastvideo/pin_memory.py

@@ -0,0 +1,42 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import os
+
+from functools import cache
+
+import torch
+
+from fastvideo.logger import init_logger
+from fastvideo.platforms import current_platform
+
+logger = init_logger(__name__)
+
+
+def _probe_pin_memory() -> bool:
+    if os.getenv("FASTVIDEO_DISABLE_PIN_MEMORY",
+                 "") not in ("", "0", "false", "False"):
+        return False
+    if current_platform.is_cpu() or current_platform.is_mps():
+        return False
+
+    try:
+        if torch.cuda.is_available():
+            torch.cuda.current_device()
+        _ = torch.empty(1024, device="cpu").pin_memory()
+        _ = torch.empty(1024, device="cpu", pin_memory=True)
+    except Exception as exc:
+        logger.warning("Pinned memory is unavailable: %s", exc)
+        return False
+
+    return True
+
+
+@cache
+def _cached_pin_memory_available(pid: int) -> bool:
+    return _probe_pin_memory()
+
+
+def is_pin_memory_available() -> bool:
+    return _cached_pin_memory_available(os.getpid())

fastvideo/pipelines/basic/cosmos/cosmos_pipeline.py

@@ -11,11 +11,12 @@ from fastvideo.logger import init_logger
 from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import (
     FlowMatchEulerDiscreteScheduler)
 from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase
-from fastvideo.pipelines.stages import (ConditioningStage, CosmosDenoisingStage,
+from fastvideo.pipelines.stages import (ConditioningStage, DenoisingStage,
                                         CosmosLatentPreparationStage,
                                         DecodingStage, InputValidationStage,
                                         TextEncodingStage,
                                         TimestepPreparationStage)
+from fastvideo.pipelines.stages.denoising_cosmos_strategy import CosmosStrategy
 
 logger = init_logger(__name__)
 
@@ -73,9 +74,10 @@ class Cosmos2VideoToWorldPipeline(ComposedPipelineBase):
                            vae=self.get_module("vae")))
 
         self.add_stage(stage_name="denoising_stage",
-                       stage=CosmosDenoisingStage(
+                       stage=DenoisingStage(
                            transformer=self.get_module("transformer"),
-                           scheduler=self.get_module("scheduler")))
+                           scheduler=self.get_module("scheduler"),
+                           strategy_cls=CosmosStrategy))
 
         self.add_stage(stage_name="decoding_stage",
                        stage=DecodingStage(vae=self.get_module("vae")))

fastvideo/pipelines/basic/longcat/__init__.py

@@ -2,5 +2,10 @@
 """LongCat pipeline module."""
 
 from fastvideo.pipelines.basic.longcat.longcat_pipeline import LongCatPipeline
+from fastvideo.pipelines.basic.longcat.longcat_i2v_pipeline import LongCatImageToVideoPipeline
+from fastvideo.pipelines.basic.longcat.longcat_vc_pipeline import LongCatVideoContinuationPipeline
 
-__all__ = ["LongCatPipeline"]
+__all__ = [
+    "LongCatPipeline", "LongCatImageToVideoPipeline",
+    "LongCatVideoContinuationPipeline"
+]

fastvideo/pipelines/basic/longcat/longcat_i2v_pipeline.py

@@ -0,0 +1,151 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+LongCat Image-to-Video pipeline implementation.
+
+This module implements I2V (Image-to-Video) generation for LongCat using Tier 3
+conditioning with timestep masking, num_cond_latents support, and RoPE skipping.
+
+Supports:
+- Basic I2V (50 steps, guidance_scale=4.0)
+- Distilled I2V with LoRA (16 steps, guidance_scale=1.0)
+- Refinement I2V for 720p upscaling (with refinement LoRA + BSA)
+"""
+
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.logger import init_logger
+from fastvideo.pipelines import ComposedPipelineBase, LoRAPipeline
+from fastvideo.pipelines.stages import (
+    DecodingStage,
+    DenoisingStage,
+    InputValidationStage,
+    TextEncodingStage,
+    TimestepPreparationStage,
+)
+from fastvideo.pipelines.stages.denoising_longcat_strategy import (
+    LongCatI2VStrategy)
+from fastvideo.pipelines.stages.longcat_image_vae_encoding import LongCatImageVAEEncodingStage
+from fastvideo.pipelines.stages.longcat_i2v_latent_preparation import LongCatI2VLatentPreparationStage
+from fastvideo.pipelines.stages.longcat_refine_init import LongCatRefineInitStage
+from fastvideo.pipelines.stages.longcat_refine_timestep import LongCatRefineTimestepStage
+
+logger = init_logger(__name__)
+
+
+class LongCatImageToVideoPipeline(LoRAPipeline, ComposedPipelineBase):
+    """
+    LongCat Image-to-Video pipeline.
+    
+    Generates video from a single input image using Tier 3 I2V conditioning:
+    - Per-frame timestep masking (timestep[:, 0] = 0)
+    - num_cond_latents parameter to transformer
+    - RoPE skipping for conditioning frames
+    - Selective denoising (skip first frame in scheduler)
+    """
+
+    _required_config_modules = [
+        "text_encoder", "tokenizer", "vae", "transformer", "scheduler"
+    ]
+
+    def initialize_pipeline(self, fastvideo_args: FastVideoArgs):
+        """Initialize LongCat-specific components."""
+        # Same BSA initialization as base LongCat pipeline
+        pipeline_config = fastvideo_args.pipeline_config
+        transformer = self.get_module("transformer", None)
+        if transformer is None:
+            return
+
+        # Enable BSA if configured
+        if pipeline_config.enable_bsa:
+            bsa_params_cfg = getattr(pipeline_config, 'bsa_params', None) or {}
+            sparsity = getattr(pipeline_config, 'bsa_sparsity', None)
+            cdf_threshold = getattr(pipeline_config, 'bsa_cdf_threshold', None)
+            chunk_q = getattr(pipeline_config, 'bsa_chunk_q', None)
+            chunk_k = getattr(pipeline_config, 'bsa_chunk_k', None)
+
+            effective_bsa_params = dict(bsa_params_cfg) if isinstance(
+                bsa_params_cfg, dict) else {}
+            if sparsity is not None:
+                effective_bsa_params['sparsity'] = sparsity
+            if cdf_threshold is not None:
+                effective_bsa_params['cdf_threshold'] = cdf_threshold
+            if chunk_q is not None:
+                effective_bsa_params['chunk_3d_shape_q'] = chunk_q
+            if chunk_k is not None:
+                effective_bsa_params['chunk_3d_shape_k'] = chunk_k
+
+            # Provide defaults
+            effective_bsa_params.setdefault('sparsity', 0.9375)
+            effective_bsa_params.setdefault('chunk_3d_shape_q', [4, 4, 4])
+            effective_bsa_params.setdefault('chunk_3d_shape_k', [4, 4, 4])
+
+            if hasattr(transformer, 'enable_bsa'):
+                logger.info("Enabling BSA for LongCat I2V transformer")
+                transformer.enable_bsa()
+                if hasattr(transformer, 'blocks'):
+                    try:
+                        for blk in transformer.blocks:
+                            if hasattr(blk, 'self_attn'):
+                                blk.self_attn.bsa_params = effective_bsa_params
+                    except Exception as e:
+                        logger.warning("Failed to set BSA params: %s", e)
+                logger.info("BSA parameters: %s", effective_bsa_params)
+        else:
+            if hasattr(transformer, 'disable_bsa'):
+                transformer.disable_bsa()
+
+    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs):
+        """Set up I2V-specific pipeline stages."""
+
+        # 1. Input validation
+        self.add_stage(stage_name="input_validation_stage",
+                       stage=InputValidationStage())
+
+        # 2. Text encoding (same as T2V)
+        self.add_stage(stage_name="prompt_encoding_stage",
+                       stage=TextEncodingStage(
+                           text_encoders=[self.get_module("text_encoder")],
+                           tokenizers=[self.get_module("tokenizer")],
+                       ))
+
+        # 3. Image VAE encoding (for I2V - skipped in refinement mode)
+        self.add_stage(
+            stage_name="image_vae_encoding_stage",
+            stage=LongCatImageVAEEncodingStage(vae=self.get_module("vae")))
+
+        # 4. Refinement initialization (skipped if not refining)
+        self.add_stage(stage_name="longcat_refine_init_stage",
+                       stage=LongCatRefineInitStage(vae=self.get_module("vae")))
+
+        # 5. Timestep preparation (generic)
+        self.add_stage(stage_name="timestep_preparation_stage",
+                       stage=TimestepPreparationStage(
+                           scheduler=self.get_module("scheduler")))
+
+        # 6. Refinement timestep override (skipped if not refining)
+        self.add_stage(stage_name="longcat_refine_timestep_stage",
+                       stage=LongCatRefineTimestepStage(
+                           scheduler=self.get_module("scheduler")))
+
+        # 7. Latent preparation with I2V conditioning
+        self.add_stage(stage_name="latent_preparation_stage",
+                       stage=LongCatI2VLatentPreparationStage(
+                           scheduler=self.get_module("scheduler"),
+                           transformer=self.get_module("transformer")))
+
+        # 8. Denoising with I2V support
+        self.add_stage(stage_name="denoising_stage",
+                       stage=DenoisingStage(
+                           transformer=self.get_module("transformer"),
+                           transformer_2=self.get_module("transformer_2", None),
+                           scheduler=self.get_module("scheduler"),
+                           vae=self.get_module("vae"),
+                           pipeline=self,
+                           strategy_cls=LongCatI2VStrategy))
+
+        # 9. Decoding
+        self.add_stage(stage_name="decoding_stage",
+                       stage=DecodingStage(vae=self.get_module("vae"),
+                                           pipeline=self))
+
+
+EntryClass = LongCatImageToVideoPipeline

fastvideo/pipelines/basic/longcat/longcat_pipeline.py

@@ -1,9 +1,9 @@
 # SPDX-License-Identifier: Apache-2.0
 """
-LongCat video diffusion pipeline implementation (Phase 1: Wrapper).
+LongCat video diffusion pipeline implementation.
 
-This module contains a wrapper implementation of the LongCat video diffusion pipeline
-using FastVideo's modular pipeline architecture with the original LongCat modules.
+This module implements the LongCat video diffusion pipeline using FastVideo's
+modular pipeline architecture.
 """
 
 from fastvideo.fastvideo_args import FastVideoArgs
@@ -11,12 +11,14 @@ from fastvideo.logger import init_logger
 from fastvideo.pipelines import ComposedPipelineBase, LoRAPipeline
 from fastvideo.pipelines.stages import (
     DecodingStage,
+    DenoisingStage,
     InputValidationStage,
     LatentPreparationStage,
     TextEncodingStage,
     TimestepPreparationStage,
 )
-from fastvideo.pipelines.stages.longcat_denoising import LongCatDenoisingStage
+from fastvideo.pipelines.stages.denoising_longcat_strategy import (
+    LongCatStrategy)
 from fastvideo.pipelines.stages.longcat_refine_init import LongCatRefineInitStage
 from fastvideo.pipelines.stages.longcat_refine_timestep import LongCatRefineTimestepStage
 
@@ -26,9 +28,6 @@ logger = init_logger(__name__)
 class LongCatPipeline(LoRAPipeline, ComposedPipelineBase):
     """
     LongCat video diffusion pipeline with LoRA support.
-    
-    Phase 1 implementation using wrapper modules from third_party/longcat_video.
-    This validates the pipeline infrastructure before full FastVideo integration.
     """
 
     _required_config_modules = [
@@ -130,12 +129,13 @@ class LongCatPipeline(LoRAPipeline, ComposedPipelineBase):
                            transformer=self.get_module("transformer", None)))
 
         self.add_stage(stage_name="denoising_stage",
-                       stage=LongCatDenoisingStage(
+                       stage=DenoisingStage(
                            transformer=self.get_module("transformer"),
                            transformer_2=self.get_module("transformer_2", None),
                            scheduler=self.get_module("scheduler"),
                            vae=self.get_module("vae"),
-                           pipeline=self))
+                           pipeline=self,
+                           strategy_cls=LongCatStrategy))
 
         self.add_stage(stage_name="decoding_stage",
                        stage=DecodingStage(vae=self.get_module("vae"),

fastvideo/pipelines/basic/longcat/longcat_vc_pipeline.py

@@ -0,0 +1,171 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+LongCat Video Continuation (VC) pipeline implementation.
+
+This module implements VC (Video Continuation) generation for LongCat with
+KV cache optimization for 2-3x speedup.
+
+Supports:
+- Basic VC (50 steps, guidance_scale=4.0)
+- Distilled VC with LoRA (16 steps, guidance_scale=1.0)
+- KV cache for conditioning frames
+"""
+
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.logger import init_logger
+from fastvideo.pipelines import ComposedPipelineBase, LoRAPipeline
+from fastvideo.pipelines.stages import (
+    DecodingStage,
+    DenoisingStage,
+    InputValidationStage,
+    TextEncodingStage,
+    TimestepPreparationStage,
+)
+from fastvideo.pipelines.stages.denoising_longcat_strategy import (
+    LongCatVCStrategy)
+from fastvideo.pipelines.stages.longcat_video_vae_encoding import LongCatVideoVAEEncodingStage
+from fastvideo.pipelines.stages.longcat_i2v_latent_preparation import LongCatI2VLatentPreparationStage
+from fastvideo.pipelines.stages.longcat_kv_cache_init import LongCatKVCacheInitStage
+
+logger = init_logger(__name__)
+
+
+class LongCatVideoContinuationPipeline(LoRAPipeline, ComposedPipelineBase):
+    """
+    LongCat Video Continuation pipeline.
+    
+    Generates video continuation from multiple conditioning frames using
+    optional KV cache for 2-3x speedup.
+    
+    Key features:
+    - Takes video input (13+ frames typically)
+    - Encodes conditioning frames via VAE
+    - Optionally pre-computes KV cache for conditioning
+    - Uses cached K/V during denoising for speedup
+    - Concatenates conditioning back after denoising
+    """
+
+    _required_config_modules = [
+        "text_encoder", "tokenizer", "vae", "transformer", "scheduler"
+    ]
+
+    def initialize_pipeline(self, fastvideo_args: FastVideoArgs):
+        """Initialize LongCat-specific components."""
+        pipeline_config = fastvideo_args.pipeline_config
+        transformer = self.get_module("transformer", None)
+        if transformer is None:
+            return
+
+        # Enable BSA if configured (for VC, BSA may not be needed)
+        if getattr(pipeline_config, 'enable_bsa', False):
+            bsa_params_cfg = getattr(pipeline_config, 'bsa_params', None) or {}
+            sparsity = getattr(pipeline_config, 'bsa_sparsity', None)
+            cdf_threshold = getattr(pipeline_config, 'bsa_cdf_threshold', None)
+            chunk_q = getattr(pipeline_config, 'bsa_chunk_q', None)
+            chunk_k = getattr(pipeline_config, 'bsa_chunk_k', None)
+
+            effective_bsa_params = dict(bsa_params_cfg) if isinstance(
+                bsa_params_cfg, dict) else {}
+            if sparsity is not None:
+                effective_bsa_params['sparsity'] = sparsity
+            if cdf_threshold is not None:
+                effective_bsa_params['cdf_threshold'] = cdf_threshold
+            if chunk_q is not None:
+                effective_bsa_params['chunk_3d_shape_q'] = chunk_q
+            if chunk_k is not None:
+                effective_bsa_params['chunk_3d_shape_k'] = chunk_k
+
+            # Provide defaults
+            effective_bsa_params.setdefault('sparsity', 0.9375)
+            effective_bsa_params.setdefault('chunk_3d_shape_q', [4, 4, 4])
+            effective_bsa_params.setdefault('chunk_3d_shape_k', [4, 4, 4])
+
+            if hasattr(transformer, 'enable_bsa'):
+                logger.info("Enabling BSA for LongCat VC transformer")
+                transformer.enable_bsa()
+                if hasattr(transformer, 'blocks'):
+                    try:
+                        for blk in transformer.blocks:
+                            if hasattr(blk, 'self_attn'):
+                                blk.self_attn.bsa_params = effective_bsa_params
+                    except Exception as e:
+                        logger.warning("Failed to set BSA params: %s", e)
+                logger.info("BSA parameters: %s", effective_bsa_params)
+        else:
+            if hasattr(transformer, 'disable_bsa'):
+                transformer.disable_bsa()
+
+    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs):
+        """Set up VC-specific pipeline stages."""
+
+        # 1. Input validation
+        self.add_stage(stage_name="input_validation_stage",
+                       stage=InputValidationStage())
+
+        # 2. Text encoding
+        self.add_stage(stage_name="prompt_encoding_stage",
+                       stage=TextEncodingStage(
+                           text_encoders=[self.get_module("text_encoder")],
+                           tokenizers=[self.get_module("tokenizer")],
+                       ))
+
+        # 3. Video VAE encoding (encodes conditioning frames)
+        self.add_stage(
+            stage_name="video_vae_encoding_stage",
+            stage=LongCatVideoVAEEncodingStage(vae=self.get_module("vae")))
+
+        # 4. Timestep preparation
+        self.add_stage(stage_name="timestep_preparation_stage",
+                       stage=TimestepPreparationStage(
+                           scheduler=self.get_module("scheduler")))
+
+        # 5. Latent preparation (reuse I2V stage - it handles video_latent too)
+        self.add_stage(stage_name="latent_preparation_stage",
+                       stage=LongCatVCLatentPreparationStage(
+                           scheduler=self.get_module("scheduler"),
+                           transformer=self.get_module("transformer")))
+
+        # 6. KV cache initialization (optional, based on config)
+        # This is always added but will skip if use_kv_cache=False
+        self.add_stage(stage_name="kv_cache_init_stage",
+                       stage=LongCatKVCacheInitStage(
+                           transformer=self.get_module("transformer")))
+
+        # 7. Denoising with VC and KV cache support
+        self.add_stage(stage_name="denoising_stage",
+                       stage=DenoisingStage(
+                           transformer=self.get_module("transformer"),
+                           transformer_2=self.get_module("transformer_2", None),
+                           scheduler=self.get_module("scheduler"),
+                           vae=self.get_module("vae"),
+                           pipeline=self,
+                           strategy_cls=LongCatVCStrategy))
+
+        # 8. Decoding
+        self.add_stage(stage_name="decoding_stage",
+                       stage=DecodingStage(vae=self.get_module("vae"),
+                                           pipeline=self))
+
+
+class LongCatVCLatentPreparationStage(LongCatI2VLatentPreparationStage):
+    """
+    Prepare latents with video conditioning for first N frames.
+    
+    Extends I2V latent preparation to handle video_latent (multiple frames)
+    instead of image_latent (single frame).
+    """
+
+    def forward(self, batch, fastvideo_args):
+        """Prepare latents with VC conditioning."""
+
+        # Check if we have video_latent (from VC encoding stage)
+        video_latent = getattr(batch, 'video_latent', None)
+        if video_latent is not None:
+            # Set image_latent to video_latent for parent class compatibility
+            batch.image_latent = video_latent
+
+        # Call parent class forward
+        return super().forward(batch, fastvideo_args)
+
+
+EntryClass = LongCatVideoContinuationPipeline

fastvideo/pipelines/basic/turbodiffusion/__init__.py

@@ -0,0 +1,6 @@
+from fastvideo.pipelines.basic.turbodiffusion.turbodiffusion_pipeline import (
+    TurboDiffusionPipeline, )
+from fastvideo.pipelines.basic.turbodiffusion.turbodiffusion_i2v_pipeline import (
+    TurboDiffusionI2VPipeline, )
+
+__all__ = ["TurboDiffusionPipeline", "TurboDiffusionI2VPipeline"]

fastvideo/pipelines/basic/turbodiffusion/turbodiffusion_i2v_pipeline.py

@@ -0,0 +1,89 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+TurboDiffusion I2V (Image-to-Video) Pipeline Implementation.
+
+This module contains an implementation of the TurboDiffusion I2V pipeline
+for 1-4 step image-to-video generation using rCM (recurrent Consistency Model)
+sampling with SLA (Sparse-Linear Attention).
+
+Key differences from T2V:
+- Uses dual models (high/low noise) with boundary switching
+- sigma_max=200 (vs 80 for T2V)
+- Mask conditioning with encoded first frame
+"""
+
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.logger import init_logger
+from fastvideo.models.schedulers.scheduling_rcm import RCMScheduler
+from fastvideo.pipelines import ComposedPipelineBase, LoRAPipeline
+from fastvideo.pipelines.stages import (ConditioningStage, DecodingStage,
+                                        DenoisingStage, ImageVAEEncodingStage,
+                                        InputValidationStage,
+                                        LatentPreparationStage,
+                                        TextEncodingStage,
+                                        TimestepPreparationStage)
+
+logger = init_logger(__name__)
+
+
+class TurboDiffusionI2VPipeline(LoRAPipeline, ComposedPipelineBase):
+    """
+    TurboDiffusion I2V pipeline for 1-4 step image-to-video generation.
+
+    Uses RCM scheduler, SLA attention, and dual model switching for
+    high-quality I2V generation.
+    """
+
+    _required_config_modules = [
+        "text_encoder", "tokenizer", "vae", "transformer", "transformer_2",
+        "scheduler"
+    ]
+
+    def initialize_pipeline(self, fastvideo_args: FastVideoArgs):
+        # Use RCM scheduler with higher sigma_max for I2V
+        logger.info(
+            "Initializing RCM scheduler for TurboDiffusion I2V (sigma_max=200)")
+        self.modules["scheduler"] = RCMScheduler(sigma_max=200.0)
+
+    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None:
+        """Set up pipeline stages with proper dependency injection."""
+
+        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="conditioning_stage",
+                       stage=ConditioningStage())
+
+        self.add_stage(stage_name="timestep_preparation_stage",
+                       stage=TimestepPreparationStage(
+                           scheduler=self.get_module("scheduler")))
+
+        self.add_stage(stage_name="latent_preparation_stage",
+                       stage=LatentPreparationStage(
+                           scheduler=self.get_module("scheduler"),
+                           transformer=self.get_module("transformer", None)))
+
+        # I2V: Encode initial image to latent space
+        self.add_stage(stage_name="image_latent_preparation_stage",
+                       stage=ImageVAEEncodingStage(vae=self.get_module("vae")))
+
+        self.add_stage(stage_name="denoising_stage",
+                       stage=DenoisingStage(
+                           transformer=self.get_module("transformer"),
+                           transformer_2=self.get_module("transformer_2", None),
+                           scheduler=self.get_module("scheduler"),
+                           vae=self.get_module("vae"),
+                           pipeline=self))
+
+        self.add_stage(stage_name="decoding_stage",
+                       stage=DecodingStage(vae=self.get_module("vae"),
+                                           pipeline=self))
+
+
+EntryClass = TurboDiffusionI2VPipeline

fastvideo/pipelines/basic/turbodiffusion/turbodiffusion_pipeline.py

@@ -36,11 +36,6 @@ class TurboDiffusionPipeline(LoRAPipeline, ComposedPipelineBase):
         logger.info("Initializing RCM scheduler for TurboDiffusion")
         self.modules["scheduler"] = RCMScheduler(sigma_max=80.0)
 
-        # Store checkpoint path for later loading
-        self._turbodiffusion_checkpoint = getattr(fastvideo_args,
-                                                  'turbodiffusion_checkpoint',
-                                                  None)
-
     def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None:
         """Set up pipeline stages with proper dependency injection."""
 

fastvideo/pipelines/basic/wan/wan_causal_dmd_pipeline.py

@@ -11,10 +11,11 @@ from fastvideo.pipelines import ComposedPipelineBase, LoRAPipeline
 
 # isort: off
 from fastvideo.pipelines.stages import (ConditioningStage, DecodingStage,
-                                        CausalDMDDenosingStage,
-                                        InputValidationStage,
+                                        DenoisingStage, InputValidationStage,
                                         LatentPreparationStage,
                                         TextEncodingStage)
+from fastvideo.pipelines.stages.denoising_causal_strategy import (
+    CausalBlockStrategy)
 # isort: on
 
 logger = init_logger(__name__)
@@ -47,11 +48,12 @@ class WanCausalDMDPipeline(LoRAPipeline, ComposedPipelineBase):
                            transformer=self.get_module("transformer", None)))
 
         self.add_stage(stage_name="denoising_stage",
-                       stage=CausalDMDDenosingStage(
+                       stage=DenoisingStage(
                            transformer=self.get_module("transformer"),
                            transformer_2=self.get_module("transformer_2", None),
                            scheduler=self.get_module("scheduler"),
-                           vae=self.get_module("vae")))
+                           vae=self.get_module("vae"),
+                           strategy_cls=CausalBlockStrategy))
 
         self.add_stage(stage_name="decoding_stage",
                        stage=DecodingStage(vae=self.get_module("vae")))

fastvideo/pipelines/basic/wan/wan_dmd_pipeline.py

@@ -14,10 +14,11 @@ from fastvideo.pipelines import ComposedPipelineBase, LoRAPipeline
 
 # isort: off
 from fastvideo.pipelines.stages import (ConditioningStage, DecodingStage,
-                                        DmdDenoisingStage, InputValidationStage,
+                                        DenoisingStage, InputValidationStage,
                                         LatentPreparationStage,
                                         TextEncodingStage,
                                         TimestepPreparationStage)
+from fastvideo.pipelines.stages.denoising_dmd_strategy import DmdStrategy
 # isort: on
 
 logger = init_logger(__name__)
@@ -63,9 +64,10 @@ class WanDMDPipeline(LoRAPipeline, ComposedPipelineBase):
                            use_btchw_layout=True))
 
         self.add_stage(stage_name="denoising_stage",
-                       stage=DmdDenoisingStage(
+                       stage=DenoisingStage(
                            transformer=self.get_module("transformer"),
-                           scheduler=self.get_module("scheduler")))
+                           scheduler=FlowMatchEulerDiscreteScheduler(shift=8.0),
+                           strategy_cls=DmdStrategy))
 
         self.add_stage(stage_name="decoding_stage",
                        stage=DecodingStage(vae=self.get_module("vae")))

fastvideo/pipelines/basic/wan/wan_i2v_dmd_pipeline.py

@@ -13,12 +13,13 @@ from fastvideo.pipelines.lora_pipeline import LoRAPipeline
 
 # isort: off
 from fastvideo.pipelines.stages import (
-    ImageEncodingStage, ConditioningStage, DecodingStage, DmdDenoisingStage,
+    ImageEncodingStage, ConditioningStage, DecodingStage, DenoisingStage,
     ImageVAEEncodingStage, InputValidationStage, LatentPreparationStage,
     TextEncodingStage, TimestepPreparationStage)
 # isort: on
 from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import (
     FlowMatchEulerDiscreteScheduler)
+from fastvideo.pipelines.stages.denoising_dmd_strategy import DmdStrategy
 
 logger = init_logger(__name__)
 
@@ -69,9 +70,10 @@ class WanImageToVideoDmdPipeline(LoRAPipeline, ComposedPipelineBase):
                        stage=ImageVAEEncodingStage(vae=self.get_module("vae")))
 
         self.add_stage(stage_name="denoising_stage",
-                       stage=DmdDenoisingStage(
+                       stage=DenoisingStage(
                            transformer=self.get_module("transformer"),
-                           scheduler=self.get_module("scheduler")))
+                           scheduler=FlowMatchEulerDiscreteScheduler(shift=8.0),
+                           strategy_cls=DmdStrategy))
 
         self.add_stage(stage_name="decoding_stage",
                        stage=DecodingStage(vae=self.get_module("vae")))

fastvideo/pipelines/pipeline_registry.py

@@ -24,6 +24,7 @@ _PIPELINE_NAME_TO_ARCHITECTURE_NAME: dict[str, str] = {
     "WanVideoToVideoPipeline": "wan",
     "WanCausalDMDPipeline": "wan",
     "TurboDiffusionPipeline": "turbodiffusion",
+    "TurboDiffusionI2VPipeline": "turbodiffusion",
     "StepVideoPipeline": "stepvideo",
     "HunyuanVideoPipeline": "hunyuan",
     "HunyuanVideo15Pipeline": "hunyuan15",
@@ -31,6 +32,8 @@ _PIPELINE_NAME_TO_ARCHITECTURE_NAME: dict[str, str] = {
     "MatrixGamePipeline": "matrixgame",
     "MatrixGameCausalDMDPipeline": "matrixgame",
     "LongCatPipeline": "longcat",
+    "LongCatImageToVideoPipeline": "longcat",
+    "LongCatVideoContinuationPipeline": "longcat",
 }
 
 _PREPROCESS_WORKLOAD_TYPE_TO_PIPELINE_NAME: dict[WorkloadType, str] = {

fastvideo/pipelines/stages/__init__.py

@@ -7,12 +7,9 @@ complete diffusion pipelines.
 """
 
 from fastvideo.pipelines.stages.base import PipelineStage
-from fastvideo.pipelines.stages.causal_denoising import CausalDMDDenosingStage
 from fastvideo.pipelines.stages.conditioning import ConditioningStage
 from fastvideo.pipelines.stages.decoding import DecodingStage
-from fastvideo.pipelines.stages.denoising import (CosmosDenoisingStage,
-                                                  DenoisingStage,
-                                                  DmdDenoisingStage)
+from fastvideo.pipelines.stages.denoising import DenoisingStage
 from fastvideo.pipelines.stages.encoding import EncodingStage
 from fastvideo.pipelines.stages.image_encoding import (
     ImageEncodingStage, MatrixGameImageEncodingStage, RefImageEncodingStage,
@@ -28,6 +25,10 @@ from fastvideo.pipelines.stages.text_encoding import TextEncodingStage
 from fastvideo.pipelines.stages.timestep_preparation import (
     TimestepPreparationStage)
 
+# LongCat stages
+from fastvideo.pipelines.stages.longcat_video_vae_encoding import LongCatVideoVAEEncodingStage
+from fastvideo.pipelines.stages.longcat_kv_cache_init import LongCatKVCacheInitStage
+
 __all__ = [
     "PipelineStage",
     "InputValidationStage",
@@ -36,10 +37,7 @@ __all__ = [
     "CosmosLatentPreparationStage",
     "ConditioningStage",
     "DenoisingStage",
-    "DmdDenoisingStage",
-    "CausalDMDDenosingStage",
     "MatrixGameCausalDenoisingStage",
-    "CosmosDenoisingStage",
     "EncodingStage",
     "DecodingStage",
     "ImageEncodingStage",
@@ -50,4 +48,7 @@ __all__ = [
     "VideoVAEEncodingStage",
     "TextEncodingStage",
     "StepvideoPromptEncodingStage",
+    # LongCat stages
+    "LongCatVideoVAEEncodingStage",
+    "LongCatKVCacheInitStage",
 ]

fastvideo/pipelines/stages/causal_denoising.py

@@ -1,497 +0,0 @@
-import torch  # type: ignore
-
-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.models.utils import pred_noise_to_pred_video, pred_noise_to_x_bound
-from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
-from fastvideo.pipelines.stages.denoising import DenoisingStage
-from fastvideo.pipelines.stages.validators import StageValidators as V
-from fastvideo.pipelines.stages.validators import VerificationResult
-
-try:
-    from fastvideo.attention.backends.sliding_tile_attn import (
-        SlidingTileAttentionBackend)
-    st_attn_available = True
-except ImportError:
-    st_attn_available = False
-    SlidingTileAttentionBackend = None  # type: ignore
-
-try:
-    from fastvideo.attention.backends.video_sparse_attn import (
-        VideoSparseAttentionBackend)
-    vsa_available = True
-except ImportError:
-    vsa_available = False
-    VideoSparseAttentionBackend = None  # type: ignore
-
-logger = init_logger(__name__)
-
-
-class CausalDMDDenosingStage(DenoisingStage):
-    """
-    Denoising stage for causal diffusion.
-    """
-
-    def __init__(self,
-                 transformer,
-                 scheduler,
-                 transformer_2=None,
-                 vae=None) -> None:
-        super().__init__(transformer, scheduler, transformer_2)
-        # KV and cross-attention cache state (initialized on first forward)
-        self.transformer = transformer
-        self.transformer_2 = transformer_2
-        self.vae = vae
-        # Model-dependent constants (aligned with causal_inference.py assumptions)
-        self.num_transformer_blocks = len(self.transformer.blocks)
-        self.num_frames_per_block = self.transformer.config.arch_config.num_frames_per_block
-        self.sliding_window_num_frames = self.transformer.config.arch_config.sliding_window_num_frames
-
-        try:
-            self.local_attn_size = getattr(self.transformer.model,
-                                           "local_attn_size",
-                                           -1)  # type: ignore
-        except Exception:
-            self.local_attn_size = -1
-
-    def forward(
-        self,
-        batch: ForwardBatch,
-        fastvideo_args: FastVideoArgs,
-    ) -> ForwardBatch:
-        target_dtype = torch.bfloat16
-        autocast_enabled = (target_dtype != torch.float32
-                            ) and not fastvideo_args.disable_autocast
-
-        latent_seq_length = batch.latents.shape[-1] * batch.latents.shape[-2]
-        patch_ratio = self.transformer.config.arch_config.patch_size[
-            -1] * self.transformer.config.arch_config.patch_size[-2]
-        self.frame_seq_length = latent_seq_length // patch_ratio
-        # TODO(will): make this a parameter once we add i2v support
-        independent_first_frame = self.transformer.independent_first_frame if hasattr(
-            self.transformer, 'independent_first_frame') else False
-        # Timesteps for DMD
-        timesteps = torch.tensor(
-            fastvideo_args.pipeline_config.dmd_denoising_steps,
-            dtype=torch.long).cpu()
-        if fastvideo_args.pipeline_config.warp_denoising_step:
-            scheduler_timesteps = torch.cat((self.scheduler.timesteps.cpu(),
-                                             torch.tensor([0],
-                                                          dtype=torch.float32)))
-            timesteps = scheduler_timesteps[1000 - timesteps]
-        timesteps = timesteps.to(get_local_torch_device())
-
-        if fastvideo_args.pipeline_config.dit_config.boundary_ratio is not None:
-            boundary_timestep = fastvideo_args.pipeline_config.dit_config.boundary_ratio * self.scheduler.num_train_timesteps
-            high_noise_timesteps = timesteps[timesteps >= boundary_timestep]
-        else:
-            boundary_timestep = None
-            high_noise_timesteps = None
-
-        # Image kwargs (kept empty unless caller provides compatible args)
-        image_kwargs: dict = {}
-
-        pos_cond_kwargs = self.prepare_extra_func_kwargs(
-            self.transformer.forward,
-            {
-                # "encoder_hidden_states_2": batch.clip_embedding_pos,
-                "encoder_attention_mask": batch.prompt_attention_mask,
-            },
-        )
-
-        # STA
-        if st_attn_available and self.attn_backend == SlidingTileAttentionBackend:
-            self.prepare_sta_param(batch, fastvideo_args)
-
-        # Latents and prompts
-        assert batch.latents is not None, "latents must be provided"
-        latents = batch.latents  # [B, C, T, H, W]
-        b, c, t, h, w = latents.shape
-        prompt_embeds = batch.prompt_embeds
-        assert torch.isnan(prompt_embeds[0]).sum() == 0
-
-        # Initialize or reset caches
-        kv_cache1 = self._initialize_kv_cache(batch_size=latents.shape[0],
-                                              dtype=target_dtype,
-                                              device=latents.device)
-        kv_cache2 = None
-        if boundary_timestep is not None:
-            # Initialize the low noise kv cache
-            kv_cache2 = self._initialize_kv_cache(batch_size=latents.shape[0],
-                                                  dtype=target_dtype,
-                                                  device=latents.device)
-
-        def _get_kv_cache(timestep: float) -> list[dict]:
-            if boundary_timestep is not None:
-                if timestep >= boundary_timestep:
-                    return kv_cache1
-                else:
-                    assert kv_cache2 is not None, "kv_cache2 is not initialized"
-                    return kv_cache2
-            return kv_cache1
-
-        crossattn_cache = self._initialize_crossattn_cache(
-            batch_size=latents.shape[0],
-            max_text_len=fastvideo_args.pipeline_config.text_encoder_configs[0].
-            arch_config.text_len,
-            dtype=target_dtype,
-            device=latents.device)
-
-        pos_start_base = 0
-
-        # Determine block sizes
-        if t % self.num_frames_per_block != 0:
-            raise ValueError(
-                "num_frames must be divisible by num_frames_per_block for causal DMD denoising"
-            )
-        num_blocks = t // self.num_frames_per_block
-        block_sizes = [self.num_frames_per_block] * num_blocks
-        start_index = 0
-
-        # For now hardcode the first block to be 1 frame assuming the model is Wan2.2-MoE
-        if boundary_timestep is not None:
-            block_sizes[0] = 1
-
-        first_frame_latent = None
-        if batch.pil_image is not None:
-            # Causal video gen directly replaces the first frame of the latent with
-            # the image latent instead of appending along the channel dim
-            assert self.vae is not None, "VAE is not provided for causal video gen task"
-            self.vae = self.vae.to(get_local_torch_device())
-            first_frame_latent = self.vae.encode(batch.pil_image).mean.float()
-            if (hasattr(self.vae, "shift_factor")
-                    and self.vae.shift_factor is not None):
-                if isinstance(self.vae.shift_factor, torch.Tensor):
-                    first_frame_latent -= self.vae.shift_factor.to(
-                        first_frame_latent.device, first_frame_latent.dtype)
-                else:
-                    first_frame_latent -= self.vae.shift_factor
-
-            if isinstance(self.vae.scaling_factor, torch.Tensor):
-                first_frame_latent = first_frame_latent * self.vae.scaling_factor.to(
-                    first_frame_latent.device, first_frame_latent.dtype)
-            else:
-                first_frame_latent = first_frame_latent * self.vae.scaling_factor
-
-            if fastvideo_args.vae_cpu_offload:
-                self.vae = self.vae.to("cpu")
-
-            # Fill the low noise and high noise kv cache with first_frame_latent and timestep 0
-            t_zero = torch.zeros([latents.shape[0], 1],
-                                 device=latents.device,
-                                 dtype=torch.long)
-            with torch.autocast(device_type="cuda",
-                                dtype=target_dtype,
-                                enabled=autocast_enabled), \
-                set_forward_context(current_timestep=0,
-                                    attn_metadata=None,
-                                    forward_batch=batch):
-                self.transformer(
-                    first_frame_latent.to(target_dtype),
-                    prompt_embeds,
-                    t_zero,
-                    kv_cache=kv_cache1,
-                    crossattn_cache=crossattn_cache,
-                    current_start=(pos_start_base + start_index) *
-                    self.frame_seq_length,
-                    start_frame=start_index,
-                    **image_kwargs,
-                    **pos_cond_kwargs,
-                )
-                if boundary_timestep is not None:
-                    self.transformer_2(
-                        first_frame_latent.to(target_dtype),
-                        prompt_embeds,
-                        t_zero,
-                        kv_cache=kv_cache2,
-                        crossattn_cache=crossattn_cache,
-                        current_start=(pos_start_base + start_index) *
-                        self.frame_seq_length,
-                        start_frame=start_index,
-                        **image_kwargs,
-                        **pos_cond_kwargs,
-                    )
-
-            start_index += 1
-            block_sizes.pop(0)
-            latents[:, :, :1, :, :] = first_frame_latent
-
-        # DMD loop in causal blocks
-        with self.progress_bar(total=len(block_sizes) *
-                               len(timesteps)) as progress_bar:
-            for current_num_frames in block_sizes:
-                current_latents = latents[:, :, start_index:start_index +
-                                          current_num_frames, :, :]
-                # use BTCHW for DMD conversion routines
-                noise_latents_btchw = current_latents.permute(0, 2, 1, 3, 4)
-                video_raw_latent_shape = noise_latents_btchw.shape
-
-                for i, t_cur in enumerate(timesteps):
-                    if boundary_timestep is not None and t_cur < boundary_timestep:
-                        current_model = self.transformer_2
-                    else:
-                        current_model = self.transformer
-                    # Copy for pred conversion
-                    noise_latents = noise_latents_btchw.clone()
-                    latent_model_input = current_latents.to(target_dtype)
-
-                    if batch.image_latent is not None and independent_first_frame and start_index == 0:
-                        latent_model_input = torch.cat([
-                            latent_model_input,
-                            batch.image_latent.to(target_dtype)
-                        ],
-                                                       dim=2)
-
-                    # Prepare inputs
-                    t_expand = t_cur.repeat(latent_model_input.shape[0])
-
-                    # Attention metadata if needed
-                    if (vsa_available and self.attn_backend
-                            == VideoSparseAttentionBackend):
-                        self.attn_metadata_builder_cls = self.attn_backend.get_builder_cls(
-                        )
-                        if self.attn_metadata_builder_cls is not None:
-                            self.attn_metadata_builder = self.attn_metadata_builder_cls(
-                            )
-                            attn_metadata = self.attn_metadata_builder.build(  # type: ignore
-                                current_timestep=i,  # type: ignore
-                                raw_latent_shape=(current_num_frames, h,
-                                                  w),  # type: ignore
-                                patch_size=fastvideo_args.pipeline_config.
-                                dit_config.patch_size,  # type: ignore
-                                STA_param=batch.STA_param,  # type: ignore
-                                VSA_sparsity=fastvideo_args.
-                                VSA_sparsity,  # type: ignore
-                                device=get_local_torch_device(),  # type: ignore
-                            )  # type: ignore
-                            assert attn_metadata is not None, "attn_metadata cannot be None"
-                        else:
-                            attn_metadata = None
-                    else:
-                        attn_metadata = None
-
-                    with torch.autocast(device_type="cuda",
-                                        dtype=target_dtype,
-                                        enabled=autocast_enabled), \
-                        set_forward_context(current_timestep=i,
-                                            attn_metadata=attn_metadata,
-                                            forward_batch=batch):
-                        # Run transformer; follow DMD stage pattern
-                        t_expanded_noise = t_cur * torch.ones(
-                            (latent_model_input.shape[0], 1),
-                            device=latent_model_input.device,
-                            dtype=torch.long)
-                        pred_noise_btchw = current_model(
-                            latent_model_input,
-                            prompt_embeds,
-                            t_expanded_noise,
-                            kv_cache=_get_kv_cache(t_cur),
-                            crossattn_cache=crossattn_cache,
-                            current_start=(pos_start_base + start_index) *
-                            self.frame_seq_length,
-                            start_frame=start_index,
-                            **image_kwargs,
-                            **pos_cond_kwargs,
-                        ).permute(0, 2, 1, 3, 4)
-
-                    # Convert pred noise to pred video with FM Euler scheduler utilities
-                    if boundary_timestep is not None and t_cur >= boundary_timestep:
-                        pred_video_btchw = pred_noise_to_x_bound(
-                            pred_noise=pred_noise_btchw.flatten(0, 1),
-                            noise_input_latent=noise_latents.flatten(0, 1),
-                            timestep=t_expand,
-                            boundary_timestep=torch.ones_like(t_expand) *
-                            boundary_timestep,
-                            scheduler=self.scheduler).unflatten(
-                                0, pred_noise_btchw.shape[:2])
-                    else:
-                        pred_video_btchw = pred_noise_to_pred_video(
-                            pred_noise=pred_noise_btchw.flatten(0, 1),
-                            noise_input_latent=noise_latents.flatten(0, 1),
-                            timestep=t_expand,
-                            scheduler=self.scheduler).unflatten(
-                                0, pred_noise_btchw.shape[:2])
-
-                    if i < len(timesteps) - 1:
-                        next_timestep = timesteps[i + 1] * torch.ones(
-                            [1],
-                            dtype=torch.long,
-                            device=pred_video_btchw.device)
-                        noise = torch.randn(
-                            video_raw_latent_shape,
-                            dtype=pred_video_btchw.dtype,
-                            generator=(batch.generator[0] if isinstance(
-                                batch.generator, list) else
-                                       batch.generator)).to(self.device)
-                        noise_btchw = noise
-                        if boundary_timestep is not None and i < len(
-                                high_noise_timesteps) - 1:
-                            noise_latents_btchw = self.scheduler.add_noise_high(
-                                pred_video_btchw.flatten(0, 1),
-                                noise_btchw.flatten(0, 1), next_timestep,
-                                torch.ones_like(next_timestep) *
-                                boundary_timestep).unflatten(
-                                    0, pred_video_btchw.shape[:2])
-                        elif boundary_timestep is not None and i == len(
-                                high_noise_timesteps) - 1:
-                            noise_latents_btchw = pred_video_btchw
-                        else:
-                            noise_latents_btchw = self.scheduler.add_noise(
-                                pred_video_btchw.flatten(0, 1),
-                                noise_btchw.flatten(0, 1),
-                                next_timestep).unflatten(
-                                    0, pred_video_btchw.shape[:2])
-                        current_latents = noise_latents_btchw.permute(
-                            0, 2, 1, 3, 4)
-                    else:
-                        current_latents = pred_video_btchw.permute(
-                            0, 2, 1, 3, 4)
-
-                    if progress_bar is not None:
-                        progress_bar.update()
-
-                # Write back and advance
-                latents[:, :, start_index:start_index +
-                        current_num_frames, :, :] = current_latents
-
-                # Re-run with context timestep to update KV cache using clean context
-                context_noise = getattr(fastvideo_args.pipeline_config,
-                                        "context_noise", 0)
-                t_context = torch.ones([latents.shape[0]],
-                                       device=latents.device,
-                                       dtype=torch.long) * int(context_noise)
-                context_bcthw = current_latents.to(target_dtype)
-                with torch.autocast(device_type="cuda",
-                                    dtype=target_dtype,
-                                    enabled=autocast_enabled), \
-                    set_forward_context(current_timestep=0,
-                                        attn_metadata=attn_metadata,
-                                        forward_batch=batch):
-                    t_expanded_context = t_context.unsqueeze(1)
-
-                    if boundary_timestep is not None:
-                        self.transformer_2(
-                            context_bcthw,
-                            prompt_embeds,
-                            t_expanded_context,
-                            kv_cache=kv_cache2,
-                            crossattn_cache=crossattn_cache,
-                            current_start=(pos_start_base + start_index) *
-                            self.frame_seq_length,
-                            start_frame=start_index,
-                            **image_kwargs,
-                            **pos_cond_kwargs,
-                        )
-
-                    self.transformer(
-                        context_bcthw,
-                        prompt_embeds,
-                        t_expanded_context,
-                        kv_cache=kv_cache1,
-                        crossattn_cache=crossattn_cache,
-                        current_start=(pos_start_base + start_index) *
-                        self.frame_seq_length,
-                        start_frame=start_index,
-                        **image_kwargs,
-                        **pos_cond_kwargs,
-                    )
-
-                start_index += current_num_frames
-
-        if boundary_timestep is not None:
-            num_frames_to_remove = self.num_frames_per_block - 1
-            latents = latents[:, :, :-num_frames_to_remove, :, :]
-
-        batch.latents = latents
-        return batch
-
-    def _initialize_kv_cache(self, batch_size, dtype, device) -> list[dict]:
-        """
-        Initialize a Per-GPU KV cache aligned with the Wan model assumptions.
-        """
-        kv_cache1 = []
-        num_attention_heads = self.transformer.num_attention_heads
-        attention_head_dim = self.transformer.attention_head_dim
-        if self.local_attn_size != -1:
-            kv_cache_size = self.local_attn_size * self.frame_seq_length
-        else:
-            kv_cache_size = self.frame_seq_length * self.sliding_window_num_frames
-
-        for _ in range(self.num_transformer_blocks):
-            kv_cache1.append({
-                "k":
-                torch.zeros([
-                    batch_size, kv_cache_size, num_attention_heads,
-                    attention_head_dim
-                ],
-                            dtype=dtype,
-                            device=device),
-                "v":
-                torch.zeros([
-                    batch_size, kv_cache_size, num_attention_heads,
-                    attention_head_dim
-                ],
-                            dtype=dtype,
-                            device=device),
-                "global_end_index":
-                torch.tensor([0], dtype=torch.long, device=device),
-                "local_end_index":
-                torch.tensor([0], dtype=torch.long, device=device),
-            })
-
-        return kv_cache1
-
-    def _initialize_crossattn_cache(self, batch_size, max_text_len, dtype,
-                                    device) -> list[dict]:
-        """
-        Initialize a Per-GPU cross-attention cache aligned with the Wan model assumptions.
-        """
-        crossattn_cache = []
-        num_attention_heads = self.transformer.num_attention_heads
-        attention_head_dim = self.transformer.attention_head_dim
-        for _ in range(self.num_transformer_blocks):
-            crossattn_cache.append({
-                "k":
-                torch.zeros([
-                    batch_size, max_text_len, num_attention_heads,
-                    attention_head_dim
-                ],
-                            dtype=dtype,
-                            device=device),
-                "v":
-                torch.zeros([
-                    batch_size, max_text_len, num_attention_heads,
-                    attention_head_dim
-                ],
-                            dtype=dtype,
-                            device=device),
-                "is_init":
-                False,
-            })
-        return crossattn_cache
-
-    def verify_input(self, batch: ForwardBatch,
-                     fastvideo_args: FastVideoArgs) -> VerificationResult:
-        """Verify denoising stage inputs."""
-        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("image_embeds", batch.image_embeds, V.is_list)
-        result.add_check("image_latent", batch.image_latent,
-                         V.none_or_tensor_with_dims(5))
-        result.add_check("num_inference_steps", batch.num_inference_steps,
-                         V.positive_int)
-        result.add_check("guidance_scale", batch.guidance_scale,
-                         V.positive_float)
-        result.add_check("eta", batch.eta, V.non_negative_float)
-        result.add_check("generator", batch.generator,
-                         V.generator_or_list_generators)
-        result.add_check("do_classifier_free_guidance",
-                         batch.do_classifier_free_guidance, V.bool_value)
-        result.add_check(
-            "negative_prompt_embeds", batch.negative_prompt_embeds, lambda x:
-            not batch.do_classifier_free_guidance or V.list_not_empty(x))
-        return result

fastvideo/pipelines/stages/denoising.py

@@ -13,41 +13,14 @@ from tqdm.auto import tqdm
 
 from fastvideo.attention import get_attn_backend
 from fastvideo.configs.pipelines.base import STA_Mode
-from fastvideo.distributed import (get_local_torch_device, get_world_group)
+from fastvideo.distributed import get_world_group
 from fastvideo.fastvideo_args import FastVideoArgs
-from fastvideo.forward_context import set_forward_context
 from fastvideo.logger import init_logger
-from fastvideo.models.loader.component_loader import TransformerLoader
-from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import (
-    FlowMatchEulerDiscreteScheduler)
-from fastvideo.models.utils import pred_noise_to_pred_video
 from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
 from fastvideo.pipelines.stages.base import PipelineStage
 from fastvideo.pipelines.stages.validators import StageValidators as V
 from fastvideo.pipelines.stages.validators import VerificationResult
 from fastvideo.platforms import AttentionBackendEnum
-from fastvideo.utils import dict_to_3d_list, masks_like
-
-try:
-    from fastvideo.attention.backends.sliding_tile_attn import (
-        SlidingTileAttentionBackend)
-    st_attn_available = True
-except ImportError:
-    st_attn_available = False
-
-try:
-    from fastvideo.attention.backends.vmoba import VMOBAAttentionBackend
-    from fastvideo.utils import is_vmoba_available
-    vmoba_attn_available = is_vmoba_available()
-except ImportError:
-    vmoba_attn_available = False
-
-try:
-    from fastvideo.attention.backends.video_sparse_attn import (
-        VideoSparseAttentionBackend)
-    vsa_available = True
-except ImportError:
-    vsa_available = False
 
 logger = init_logger(__name__)
 
@@ -65,13 +38,15 @@ class DenoisingStage(PipelineStage):
                  scheduler,
                  pipeline=None,
                  transformer_2=None,
-                 vae=None) -> None:
+                 vae=None,
+                 strategy_cls=None) -> None:
         super().__init__()
         self.transformer = transformer
         self.transformer_2 = transformer_2
         self.scheduler = scheduler
         self.vae = vae
         self.pipeline = weakref.ref(pipeline) if pipeline else None
+        self.strategy_cls = strategy_cls
         attn_head_size = self.transformer.hidden_size // self.transformer.num_attention_heads
         self.attn_backend = get_attn_backend(
             head_size=attn_head_size,
@@ -92,416 +67,23 @@ class DenoisingStage(PipelineStage):
     ) -> ForwardBatch:
         """
         Run the denoising loop.
-        
-        Args:
-            batch: The current batch information.
-            fastvideo_args: The inference arguments.
-            
-        Returns:
-            The batch with denoised latents.
         """
-        pipeline = self.pipeline() if self.pipeline else None
-        if not fastvideo_args.model_loaded["transformer"]:
-            loader = TransformerLoader()
-            self.transformer = loader.load(
-                fastvideo_args.model_paths["transformer"], fastvideo_args)
-            if pipeline:
-                pipeline.add_module("transformer", self.transformer)
-            fastvideo_args.model_loaded["transformer"] = True
+        from fastvideo.pipelines.stages.denoising_engine import DenoisingEngine
+        from fastvideo.pipelines.stages.denoising_standard_strategy import (
+            StandardStrategy)
 
-        # Prepare extra step kwargs for scheduler
-        extra_step_kwargs = self.prepare_extra_func_kwargs(
-            self.scheduler.step,
-            {
-                "generator": batch.generator,
-                "eta": batch.eta
-            },
-        )
+        strategy_cls = self.strategy_cls or StandardStrategy
+        engine = DenoisingEngine(strategy_cls(self),
+                                 hooks=self._build_engine_hooks())
+        return engine.run(batch, fastvideo_args)
 
-        # Setup precision and autocast settings
-        # TODO(will): make the precision configurable for inference
-        # target_dtype = PRECISION_TO_TYPE[fastvideo_args.precision]
-        target_dtype = torch.bfloat16
-        autocast_enabled = (target_dtype != torch.float32
-                            ) and not fastvideo_args.disable_autocast
-
-        # Get timesteps and calculate warmup steps
-        timesteps = batch.timesteps
-        # TODO(will): remove this once we add input/output validation for stages
-        if timesteps is None:
-            raise ValueError("Timesteps must be provided")
-        num_inference_steps = batch.num_inference_steps
-        num_warmup_steps = len(
-            timesteps) - num_inference_steps * self.scheduler.order
-
-        # Prepare image latents and embeddings for I2V generation
-        image_embeds = batch.image_embeds
-        if len(image_embeds) > 0:
-            assert not torch.isnan(
-                image_embeds[0]).any(), "image_embeds contains nan"
-            image_embeds = [
-                image_embed.to(target_dtype) for image_embed in image_embeds
-            ]
-
-        image_kwargs = self.prepare_extra_func_kwargs(
-            self.transformer.forward,
-            {
-                "encoder_hidden_states_image": image_embeds,
-                "mask_strategy": dict_to_3d_list(
-                    None, t_max=50, l_max=60, h_max=24)
-            },
-        )
-
-        pos_cond_kwargs = self.prepare_extra_func_kwargs(
-            self.transformer.forward,
-            {
-                "encoder_hidden_states_2": batch.clip_embedding_pos,
-                "encoder_attention_mask": batch.prompt_attention_mask,
-            },
-        )
-
-        neg_cond_kwargs = self.prepare_extra_func_kwargs(
-            self.transformer.forward,
-            {
-                "encoder_hidden_states_2": batch.clip_embedding_neg,
-                "encoder_attention_mask": batch.negative_attention_mask,
-            },
-        )
-
-        action_kwargs = self.prepare_extra_func_kwargs(
-            self.transformer.forward,
-            {
-                "mouse_cond": batch.mouse_cond,
-                "keyboard_cond": batch.keyboard_cond,
-            },
-        )
-
-        # Prepare STA parameters
-        if st_attn_available and self.attn_backend == SlidingTileAttentionBackend:
-            self.prepare_sta_param(batch, fastvideo_args)
-
-        # Get latents and embeddings
-        latents = batch.latents
-        prompt_embeds = batch.prompt_embeds
-        assert not torch.isnan(
-            prompt_embeds[0]).any(), "prompt_embeds contains nan"
-        if batch.do_classifier_free_guidance:
-            neg_prompt_embeds = batch.negative_prompt_embeds
-            assert neg_prompt_embeds is not None
-            assert not torch.isnan(
-                neg_prompt_embeds[0]).any(), "neg_prompt_embeds contains nan"
-
-        # (Wan2.2) Calculate timestep to switch from high noise expert to low noise expert
-        boundary_ratio = fastvideo_args.pipeline_config.dit_config.boundary_ratio
-        if batch.boundary_ratio is not None:
-            logger.info("Overriding boundary ratio from %s to %s",
-                        boundary_ratio, batch.boundary_ratio)
-            boundary_ratio = batch.boundary_ratio
-
-        if boundary_ratio is not None:
-            boundary_timestep = boundary_ratio * self.scheduler.num_train_timesteps
-        else:
-            boundary_timestep = None
-        latent_model_input = latents.to(target_dtype)
-        assert latent_model_input.shape[0] == 1, "only support batch size 1"
-
-        if fastvideo_args.pipeline_config.ti2v_task and batch.pil_image is not None:
-            # TI2V directly replaces the first frame of the latent with
-            # the image latent instead of appending along the channel dim
-            assert batch.image_latent is None, "TI2V task should not have image latents"
-            assert self.vae is not None, "VAE is not provided for TI2V task"
-            z = self.vae.encode(batch.pil_image).mean.float()
-            if (hasattr(self.vae, "shift_factor")
-                    and self.vae.shift_factor is not None):
-                if isinstance(self.vae.shift_factor, torch.Tensor):
-                    z -= self.vae.shift_factor.to(z.device, z.dtype)
-                else:
-                    z -= self.vae.shift_factor
-
-            if isinstance(self.vae.scaling_factor, torch.Tensor):
-                z = z * self.vae.scaling_factor.to(z.device, z.dtype)
-            else:
-                z = z * self.vae.scaling_factor
-
-            latent_model_input = latent_model_input.squeeze(0)
-            _, mask2 = masks_like([latent_model_input], zero=True)
-
-            latent_model_input = (1. -
-                                  mask2[0]) * z + mask2[0] * latent_model_input
-            # latent_model_input = latent_model_input.unsqueeze(0)
-            latent_model_input = latent_model_input.to(get_local_torch_device())
-            latents = latent_model_input
-            F = batch.num_frames
-            temporal_scale = fastvideo_args.pipeline_config.vae_config.arch_config.scale_factor_temporal
-            spatial_scale = fastvideo_args.pipeline_config.vae_config.arch_config.scale_factor_spatial
-            patch_size = fastvideo_args.pipeline_config.dit_config.arch_config.patch_size
-            seq_len = ((F - 1) // temporal_scale +
-                       1) * (batch.height // spatial_scale) * (
-                           batch.width // spatial_scale) // (patch_size[1] *
-                                                             patch_size[2])
-
-        # Initialize lists for ODE trajectory
-        trajectory_timesteps: list[torch.Tensor] = []
-        trajectory_latents: list[torch.Tensor] = []
-
-        # Run denoising loop
-        with self.progress_bar(total=num_inference_steps) as progress_bar:
-            for i, t in enumerate(timesteps):
-                # Skip if interrupted
-                if hasattr(self, 'interrupt') and self.interrupt:
-                    continue
-
-                if boundary_timestep is None or t >= boundary_timestep:
-                    if (fastvideo_args.dit_cpu_offload
-                            and not fastvideo_args.dit_layerwise_offload
-                            and self.transformer_2 is not None and next(
-                                self.transformer_2.parameters()).device.type
-                            == 'cuda'):
-                        self.transformer_2.to('cpu')
-                    current_model = self.transformer
-                    if (fastvideo_args.dit_cpu_offload
-                            and not fastvideo_args.dit_layerwise_offload
-                            and not fastvideo_args.use_fsdp_inference
-                            and current_model is not None):
-                        transformer_device = next(
-                            current_model.parameters()).device.type
-                        if transformer_device == 'cpu':
-                            current_model.to(get_local_torch_device())
-                    current_guidance_scale = batch.guidance_scale
-                else:
-                    # low-noise stage in wan2.2
-                    if (fastvideo_args.dit_cpu_offload
-                            and not fastvideo_args.dit_layerwise_offload
-                            and next(self.transformer.parameters()).device.type
-                            == 'cuda'):
-                        self.transformer.to('cpu')
-                    current_model = self.transformer_2
-                    if (fastvideo_args.dit_cpu_offload
-                            and not fastvideo_args.dit_layerwise_offload
-                            and not fastvideo_args.use_fsdp_inference
-                            and current_model is not None):
-                        transformer_2_device = next(
-                            current_model.parameters()).device.type
-                        if transformer_2_device == 'cpu':
-                            current_model.to(get_local_torch_device())
-                    current_guidance_scale = batch.guidance_scale_2
-                assert current_model is not None, "current_model is None"
-
-                # Expand latents for V2V/I2V
-                latent_model_input = latents.to(target_dtype)
-                if batch.video_latent is not None:
-                    latent_model_input = torch.cat([
-                        latent_model_input, batch.video_latent,
-                        torch.zeros_like(latents)
-                    ],
-                                                   dim=1).to(target_dtype)
-                elif batch.image_latent is not None:
-                    assert not fastvideo_args.pipeline_config.ti2v_task, "image latents should not be provided for TI2V task"
-                    latent_model_input = torch.cat(
-                        [latent_model_input, batch.image_latent],
-                        dim=1).to(target_dtype)
-
-                assert not torch.isnan(
-                    latent_model_input).any(), "latent_model_input contains nan"
-                if fastvideo_args.pipeline_config.ti2v_task and batch.pil_image is not None:
-                    timestep = torch.stack([t]).to(get_local_torch_device())
-                    temp_ts = (mask2[0][0][:, ::2, ::2] * timestep).flatten()
-                    temp_ts = torch.cat([
-                        temp_ts,
-                        temp_ts.new_ones(seq_len - temp_ts.size(0)) * timestep
-                    ])
-                    timestep = temp_ts.unsqueeze(0)
-                    t_expand = timestep.repeat(latent_model_input.shape[0], 1)
-                else:
-                    t_expand = t.repeat(latent_model_input.shape[0])
-
-                latent_model_input = self.scheduler.scale_model_input(
-                    latent_model_input, t)
-
-                # Prepare inputs for transformer
-                guidance_expand = (
-                    torch.tensor(
-                        [fastvideo_args.pipeline_config.embedded_cfg_scale] *
-                        latent_model_input.shape[0],
-                        dtype=torch.float32,
-                        device=get_local_torch_device(),
-                    ).to(target_dtype) *
-                    1000.0 if fastvideo_args.pipeline_config.embedded_cfg_scale
-                    is not None else None)
-
-                # Predict noise residual
-                with torch.autocast(device_type="cuda",
-                                    dtype=target_dtype,
-                                    enabled=autocast_enabled):
-                    if (st_attn_available
-                            and self.attn_backend == SlidingTileAttentionBackend
-                        ) or (vsa_available and self.attn_backend
-                              == VideoSparseAttentionBackend):
-                        self.attn_metadata_builder_cls = self.attn_backend.get_builder_cls(
-                        )
-
-                        if self.attn_metadata_builder_cls is not None:
-                            self.attn_metadata_builder = self.attn_metadata_builder_cls(
-                            )
-                            # TODO(will): clean this up
-                            attn_metadata = self.attn_metadata_builder.build(  # type: ignore
-                                current_timestep=i,  # type: ignore
-                                raw_latent_shape=batch.
-                                raw_latent_shape[2:5],  # type: ignore
-                                patch_size=fastvideo_args.
-                                pipeline_config.  # type: ignore
-                                dit_config.patch_size,  # type: ignore
-                                STA_param=batch.STA_param,  # type: ignore
-                                VSA_sparsity=fastvideo_args.
-                                VSA_sparsity,  # type: ignore
-                                device=get_local_torch_device(),
-                            )
-                            assert attn_metadata is not None, "attn_metadata cannot be None"
-                        else:
-                            attn_metadata = None
-                    elif (vmoba_attn_available
-                          and self.attn_backend == VMOBAAttentionBackend):
-                        self.attn_metadata_builder_cls = self.attn_backend.get_builder_cls(
-                        )
-                        if self.attn_metadata_builder_cls is not None:
-                            self.attn_metadata_builder = self.attn_metadata_builder_cls(
-                            )
-                            # Prepare V-MoBA parameters from config
-                            moba_params = fastvideo_args.moba_config.copy()
-                            moba_params.update({
-                                "current_timestep":
-                                i,
-                                "raw_latent_shape":
-                                batch.raw_latent_shape[2:5],
-                                "patch_size":
-                                fastvideo_args.pipeline_config.dit_config.
-                                patch_size,
-                                "device":
-                                get_local_torch_device(),
-                            })
-                            attn_metadata = self.attn_metadata_builder.build(
-                                **moba_params)
-                            assert attn_metadata is not None, "attn_metadata cannot be None"
-                        else:
-                            attn_metadata = None
-                    else:
-                        attn_metadata = None
-                    # TODO(will): finalize the interface. vLLM uses this to
-                    # support torch dynamo compilation. They pass in
-                    # attn_metadata, vllm_config, and num_tokens. We can pass in
-                    # fastvideo_args or training_args, and attn_metadata.
-                    batch.is_cfg_negative = False
-                    with set_forward_context(
-                            current_timestep=i,
-                            attn_metadata=attn_metadata,
-                            forward_batch=batch,
-                            # fastvideo_args=fastvideo_args
-                    ):
-                        # Run transformer
-                        noise_pred = current_model(
-                            latent_model_input,
-                            prompt_embeds,
-                            t_expand,
-                            guidance=guidance_expand,
-                            **image_kwargs,
-                            **pos_cond_kwargs,
-                            **action_kwargs,
-                        )
-
-                    if batch.do_classifier_free_guidance:
-                        batch.is_cfg_negative = True
-                        with set_forward_context(
-                                current_timestep=i,
-                                attn_metadata=attn_metadata,
-                                forward_batch=batch,
-                        ):
-                            noise_pred_uncond = current_model(
-                                latent_model_input,
-                                neg_prompt_embeds,
-                                t_expand,
-                                guidance=guidance_expand,
-                                **image_kwargs,
-                                **neg_cond_kwargs,
-                                **action_kwargs,
-                            )
-
-                        noise_pred_text = noise_pred
-                        noise_pred = noise_pred_uncond + current_guidance_scale * (
-                            noise_pred_text - noise_pred_uncond)
-
-                        # Apply guidance rescale if needed
-                        if batch.guidance_rescale > 0.0:
-                            # Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf
-                            noise_pred = self.rescale_noise_cfg(
-                                noise_pred,
-                                noise_pred_text,
-                                guidance_rescale=batch.guidance_rescale,
-                            )
-                    # Compute the previous noisy sample
-                    latents = self.scheduler.step(noise_pred,
-                                                  t,
-                                                  latents,
-                                                  **extra_step_kwargs,
-                                                  return_dict=False)[0]
-                    if fastvideo_args.pipeline_config.ti2v_task and batch.pil_image is not None:
-                        latents = latents.squeeze(0)
-                        latents = (1. - mask2[0]) * z + mask2[0] * latents
-                        # latents = latents.unsqueeze(0)
-
-                # save trajectory latents if needed
-                if batch.return_trajectory_latents:
-                    trajectory_timesteps.append(t)
-                    trajectory_latents.append(latents)
-
-                # Update progress bar
-                if i == len(timesteps) - 1 or (
-                    (i + 1) > num_warmup_steps and
-                    (i + 1) % self.scheduler.order == 0
-                        and progress_bar is not None):
-                    progress_bar.update()
-
-        trajectory_tensor: torch.Tensor | None = None
-        if trajectory_latents:
-            trajectory_tensor = torch.stack(trajectory_latents, dim=1)
-            trajectory_timesteps_tensor = torch.stack(trajectory_timesteps,
-                                                      dim=0)
-        else:
-            trajectory_tensor = None
-            trajectory_timesteps_tensor = None
-
-        if trajectory_tensor is not None and trajectory_timesteps_tensor is not None:
-            batch.trajectory_timesteps = trajectory_timesteps_tensor.cpu()
-            batch.trajectory_latents = trajectory_tensor.cpu()
-
-        # Update batch with final latents
-        batch.latents = latents
-
-        if fastvideo_args.dit_layerwise_offload:
-            mgr = getattr(self.transformer, "_layerwise_offload_manager", None)
-            if mgr is not None and getattr(mgr, "enabled", False):
-                mgr.release_all()
-            if self.transformer_2 is not None:
-                mgr2 = getattr(self.transformer_2, "_layerwise_offload_manager",
-                               None)
-                if mgr2 is not None and getattr(mgr2, "enabled", False):
-                    mgr2.release_all()
-
-        # Save STA mask search results if needed
-        if st_attn_available and self.attn_backend == SlidingTileAttentionBackend and fastvideo_args.STA_mode == STA_Mode.STA_SEARCHING:
-            self.save_sta_search_results(batch)
-
-        # deallocate transformer if on mps
-        if torch.backends.mps.is_available():
-            logger.info("Memory before deallocating transformer: %s",
-                        torch.mps.current_allocated_memory())
-            del self.transformer
-            if pipeline is not None and "transformer" in pipeline.modules:
-                del pipeline.modules["transformer"]
-            fastvideo_args.model_loaded["transformer"] = False
-            logger.info("Memory after deallocating transformer: %s",
-                        torch.mps.current_allocated_memory())
-
-        return batch
+    def _build_engine_hooks(self):
+        import fastvideo.envs as envs
+        if not envs.FASTVIDEO_DENOISING_PERF_LOGGING:
+            return []
+        from fastvideo.pipelines.stages.denoising_engine_hooks import (
+            PerfLoggingHook)
+        return [PerfLoggingHook()]
 
     def prepare_extra_func_kwargs(self, func, kwargs) -> dict[str, Any]:
         """
@@ -712,8 +294,9 @@ class DenoisingStage(PipelineStage):
                      fastvideo_args: FastVideoArgs) -> VerificationResult:
         """Verify denoising stage inputs."""
         result = VerificationResult()
-        result.add_check("timesteps", batch.timesteps,
-                         [V.is_tensor, V.min_dims(1)])
+        if batch.timesteps is not None:
+            result.add_check("timesteps", batch.timesteps,
+                             [V.is_tensor, V.min_dims(1)])
         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)
@@ -741,475 +324,3 @@ class DenoisingStage(PipelineStage):
         result.add_check("latents", batch.latents,
                          [V.is_tensor, V.with_dims(5)])
         return result
-
-
-class CosmosDenoisingStage(DenoisingStage):
-    """
-    Denoising stage for Cosmos models using FlowMatchEulerDiscreteScheduler.
-    """
-
-    def __init__(self, transformer, scheduler, pipeline=None) -> None:
-        super().__init__(transformer, scheduler, pipeline)
-
-    def forward(
-        self,
-        batch: ForwardBatch,
-        fastvideo_args: FastVideoArgs,
-    ) -> ForwardBatch:
-        pipeline = self.pipeline() if self.pipeline else None
-        if not fastvideo_args.model_loaded["transformer"]:
-            loader = TransformerLoader()
-            self.transformer = loader.load(
-                fastvideo_args.model_paths["transformer"], fastvideo_args)
-            if pipeline:
-                pipeline.add_module("transformer", self.transformer)
-            fastvideo_args.model_loaded["transformer"] = True
-
-        extra_step_kwargs = self.prepare_extra_func_kwargs(
-            self.scheduler.step,
-            {
-                "generator": batch.generator,
-                "eta": batch.eta
-            },
-        )
-
-        if hasattr(self.transformer, 'module'):
-            transformer_dtype = next(self.transformer.module.parameters()).dtype
-        else:
-            transformer_dtype = next(self.transformer.parameters()).dtype
-        target_dtype = transformer_dtype
-        autocast_enabled = (target_dtype != torch.float32
-                            ) and not fastvideo_args.disable_autocast
-
-        latents = batch.latents
-        num_inference_steps = batch.num_inference_steps
-        guidance_scale = batch.guidance_scale
-
-        sigma_max = 80.0
-        sigma_min = 0.002
-        sigma_data = 1.0
-        final_sigmas_type = "sigma_min"
-
-        if self.scheduler is not None:
-            self.scheduler.register_to_config(
-                sigma_max=sigma_max,
-                sigma_min=sigma_min,
-                sigma_data=sigma_data,
-                final_sigmas_type=final_sigmas_type,
-            )
-
-        self.scheduler.set_timesteps(num_inference_steps, device=latents.device)
-        timesteps = self.scheduler.timesteps
-
-        if (hasattr(self.scheduler.config, 'final_sigmas_type')
-                and self.scheduler.config.final_sigmas_type == "sigma_min"
-                and len(self.scheduler.sigmas) > 1):
-            self.scheduler.sigmas[-1] = self.scheduler.sigmas[-2]
-
-        conditioning_latents = getattr(batch, 'conditioning_latents', None)
-        unconditioning_latents = conditioning_latents
-
-        with self.progress_bar(total=num_inference_steps) as progress_bar:
-            for i, t in enumerate(timesteps):
-                if hasattr(self, 'interrupt') and self.interrupt:
-                    continue
-
-                current_sigma = self.scheduler.sigmas[i]
-                current_t = current_sigma / (current_sigma + 1)
-                c_in = 1 - current_t
-                c_skip = 1 - current_t
-                c_out = -current_t
-
-                timestep = current_t.view(1, 1, 1, 1,
-                                          1).expand(latents.size(0), -1,
-                                                    latents.size(2), -1,
-                                                    -1)  # [B, 1, T, 1, 1]
-
-                with torch.autocast(device_type="cuda",
-                                    dtype=target_dtype,
-                                    enabled=autocast_enabled):
-
-                    cond_latent = latents * c_in
-
-                    if hasattr(
-                            batch, 'cond_indicator'
-                    ) and batch.cond_indicator is not None and conditioning_latents is not None:
-                        cond_latent = batch.cond_indicator * conditioning_latents + (
-                            1 - batch.cond_indicator) * cond_latent
-                    else:
-                        logger.warning(
-                            "Step %s: Missing conditioning data - cond_indicator: %s, conditioning_latents: %s",
-                            i, hasattr(batch, 'cond_indicator'),
-                            conditioning_latents is not None)
-
-                    cond_latent = cond_latent.to(target_dtype)
-
-                    cond_timestep = timestep
-                    if hasattr(batch, 'cond_indicator'
-                               ) and batch.cond_indicator is not None:
-                        sigma_conditioning = 0.0001
-                        t_conditioning = sigma_conditioning / (
-                            sigma_conditioning + 1)
-                        cond_timestep = batch.cond_indicator * t_conditioning + (
-                            1 - batch.cond_indicator) * timestep
-                        cond_timestep = cond_timestep.to(target_dtype)
-
-                    with set_forward_context(
-                            current_timestep=i,
-                            attn_metadata=None,
-                            forward_batch=batch,
-                    ):
-                        # Use conditioning masks from CosmosLatentPreparationStage
-                        condition_mask = batch.cond_mask.to(
-                            target_dtype) if hasattr(batch,
-                                                     'cond_mask') else None
-                        padding_mask = torch.zeros(1,
-                                                   1,
-                                                   batch.height,
-                                                   batch.width,
-                                                   device=cond_latent.device,
-                                                   dtype=target_dtype)
-
-                        # Fallback if masks not available
-                        if condition_mask is None:
-                            batch_size, num_channels, num_frames, height, width = cond_latent.shape
-                            condition_mask = torch.zeros(
-                                batch_size,
-                                1,
-                                num_frames,
-                                height,
-                                width,
-                                device=cond_latent.device,
-                                dtype=target_dtype)
-
-                        noise_pred = self.transformer(
-                            hidden_states=cond_latent,
-                            timestep=cond_timestep.to(target_dtype),
-                            encoder_hidden_states=batch.prompt_embeds[0].to(
-                                target_dtype),
-                            fps=24,  # TODO: get fps from batch or config
-                            condition_mask=condition_mask,
-                            padding_mask=padding_mask,
-                            return_dict=False,
-                        )[0]
-
-                    cond_pred = (c_skip * latents +
-                                 c_out * noise_pred.float()).to(target_dtype)
-
-                    if hasattr(
-                            batch, 'cond_indicator'
-                    ) and batch.cond_indicator is not None and conditioning_latents is not None:
-                        cond_pred = batch.cond_indicator * conditioning_latents + (
-                            1 - batch.cond_indicator) * cond_pred
-
-                    if batch.do_classifier_free_guidance and batch.negative_prompt_embeds is not None:
-                        uncond_latent = latents * c_in
-
-                        if hasattr(
-                                batch, 'uncond_indicator'
-                        ) and batch.uncond_indicator is not None and unconditioning_latents is not None:
-                            uncond_latent = batch.uncond_indicator * unconditioning_latents + (
-                                1 - batch.uncond_indicator) * uncond_latent
-
-                        with set_forward_context(
-                                current_timestep=i,
-                                attn_metadata=None,
-                                forward_batch=batch,
-                        ):
-                            uncond_condition_mask = batch.uncond_mask.to(
-                                target_dtype
-                            ) if hasattr(
-                                batch, 'uncond_mask'
-                            ) and batch.uncond_mask is not None else condition_mask
-
-                            uncond_timestep = timestep
-                            if hasattr(batch, 'uncond_indicator'
-                                       ) and batch.uncond_indicator is not None:
-                                sigma_conditioning = 0.0001
-                                t_conditioning = sigma_conditioning / (
-                                    sigma_conditioning + 1)
-                                uncond_timestep = batch.uncond_indicator * t_conditioning + (
-                                    1 - batch.uncond_indicator) * timestep
-                                uncond_timestep = uncond_timestep.to(
-                                    target_dtype)
-
-                            noise_pred_uncond = self.transformer(
-                                hidden_states=uncond_latent.to(target_dtype),
-                                timestep=uncond_timestep.to(target_dtype),
-                                encoder_hidden_states=batch.
-                                negative_prompt_embeds[0].to(target_dtype),
-                                fps=24,  # TODO: get fps from batch or config
-                                condition_mask=uncond_condition_mask,
-                                padding_mask=padding_mask,
-                                return_dict=False,
-                            )[0]
-
-                        uncond_pred = (
-                            c_skip * latents +
-                            c_out * noise_pred_uncond.float()).to(target_dtype)
-
-                        if hasattr(
-                                batch, 'uncond_indicator'
-                        ) and batch.uncond_indicator is not None and unconditioning_latents is not None:
-                            uncond_pred = batch.uncond_indicator * unconditioning_latents + (
-                                1 - batch.uncond_indicator) * uncond_pred
-
-                        guidance_diff = cond_pred - uncond_pred
-                        final_pred = cond_pred + guidance_scale * guidance_diff
-                    else:
-                        final_pred = cond_pred
-
-                # Convert to noise for scheduler step
-                if current_sigma > 1e-8:
-                    noise_for_scheduler = (latents - final_pred) / current_sigma
-                else:
-                    logger.warning(
-                        "Step %s: current_sigma too small (%s), using final_pred directly",
-                        i, current_sigma)
-                    noise_for_scheduler = final_pred
-
-                if torch.isnan(noise_for_scheduler).sum() > 0:
-                    logger.error(
-                        "Step %s: NaN detected in noise_for_scheduler, sum: %s",
-                        i,
-                        noise_for_scheduler.float().sum().item())
-                    logger.error(
-                        "Step %s: latents sum: %s, final_pred sum: %s, current_sigma: %s",
-                        i,
-                        latents.float().sum().item(),
-                        final_pred.float().sum().item(), current_sigma)
-
-                latents = self.scheduler.step(noise_for_scheduler,
-                                              t,
-                                              latents,
-                                              **extra_step_kwargs,
-                                              return_dict=False)[0]
-
-                progress_bar.update()
-
-        batch.latents = latents
-
-        return batch
-
-    def verify_input(self, batch: ForwardBatch,
-                     fastvideo_args: FastVideoArgs) -> VerificationResult:
-        """Verify Cosmos denoising stage inputs."""
-        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)
-        result.add_check("guidance_scale", batch.guidance_scale,
-                         V.positive_float)
-        result.add_check("do_classifier_free_guidance",
-                         batch.do_classifier_free_guidance, V.bool_value)
-        result.add_check(
-            "negative_prompt_embeds", batch.negative_prompt_embeds, lambda x:
-            not batch.do_classifier_free_guidance or V.list_not_empty(x))
-        return result
-
-    def verify_output(self, batch: ForwardBatch,
-                      fastvideo_args: FastVideoArgs) -> VerificationResult:
-        """Verify Cosmos denoising stage outputs."""
-        result = VerificationResult()
-        result.add_check("latents", batch.latents,
-                         [V.is_tensor, V.with_dims(5)])
-        return result
-
-
-class DmdDenoisingStage(DenoisingStage):
-    """
-    Denoising stage for DMD.
-    """
-
-    def __init__(self, transformer, scheduler) -> None:
-        super().__init__(transformer, scheduler)
-        self.scheduler = FlowMatchEulerDiscreteScheduler(shift=8.0)
-
-    def forward(
-        self,
-        batch: ForwardBatch,
-        fastvideo_args: FastVideoArgs,
-    ) -> ForwardBatch:
-        """
-        Run the denoising loop.
-        
-        Args:
-            batch: The current batch information.
-            fastvideo_args: The inference arguments.
-            
-        Returns:
-            The batch with denoised latents.
-        """
-        # Setup precision and autocast settings
-        # TODO(will): make the precision configurable for inference
-        # target_dtype = PRECISION_TO_TYPE[fastvideo_args.precision]
-        target_dtype = torch.bfloat16
-        autocast_enabled = (target_dtype != torch.float32
-                            ) and not fastvideo_args.disable_autocast
-
-        # Get timesteps and calculate warmup steps
-        timesteps = batch.timesteps
-
-        # TODO(will): remove this once we add input/output validation for stages
-        if timesteps is None:
-            raise ValueError("Timesteps must be provided")
-        num_inference_steps = batch.num_inference_steps
-        num_warmup_steps = len(
-            timesteps) - num_inference_steps * self.scheduler.order
-
-        # Prepare image latents and embeddings for I2V generation
-        image_embeds = batch.image_embeds
-        if len(image_embeds) > 0:
-            assert torch.isnan(image_embeds[0]).sum() == 0
-            image_embeds = [
-                image_embed.to(target_dtype) for image_embed in image_embeds
-            ]
-
-        image_kwargs = self.prepare_extra_func_kwargs(
-            self.transformer.forward,
-            {
-                "encoder_hidden_states_image": image_embeds,
-                "mask_strategy": dict_to_3d_list(
-                    None, t_max=50, l_max=60, h_max=24)
-            },
-        )
-
-        pos_cond_kwargs = self.prepare_extra_func_kwargs(
-            self.transformer.forward,
-            {
-                "encoder_hidden_states_2": batch.clip_embedding_pos,
-                "encoder_attention_mask": batch.prompt_attention_mask,
-            },
-        )
-
-        # Prepare STA parameters
-        if st_attn_available and self.attn_backend == SlidingTileAttentionBackend:
-            self.prepare_sta_param(batch, fastvideo_args)
-
-        # Get latents and embeddings
-        assert batch.latents is not None, "latents must be provided"
-        latents = batch.latents
-
-        video_raw_latent_shape = latents.shape
-        prompt_embeds = batch.prompt_embeds
-        assert not torch.isnan(
-            prompt_embeds[0]).any(), "prompt_embeds contains nan"
-        timesteps = torch.tensor(
-            fastvideo_args.pipeline_config.dmd_denoising_steps,
-            dtype=torch.long,
-            device=get_local_torch_device())
-
-        # Run denoising loop
-        with self.progress_bar(total=len(timesteps)) as progress_bar:
-            for i, t in enumerate(timesteps):
-                # Skip if interrupted
-                if hasattr(self, 'interrupt') and self.interrupt:
-                    continue
-                # Expand latents for I2V
-                noise_latents = latents.clone()
-                latent_model_input = latents.to(target_dtype)
-
-                if batch.image_latent is not None:
-                    latent_model_input = torch.cat([
-                        latent_model_input,
-                        batch.image_latent.permute(0, 2, 1, 3, 4)
-                    ],
-                                                   dim=2).to(target_dtype)
-                assert not torch.isnan(
-                    latent_model_input).any(), "latent_model_input contains nan"
-
-                # Prepare inputs for transformer
-                t_expand = t.repeat(latent_model_input.shape[0])
-                guidance_expand = (
-                    torch.tensor(
-                        [fastvideo_args.pipeline_config.embedded_cfg_scale] *
-                        latent_model_input.shape[0],
-                        dtype=torch.float32,
-                        device=get_local_torch_device(),
-                    ).to(target_dtype) *
-                    1000.0 if fastvideo_args.pipeline_config.embedded_cfg_scale
-                    is not None else None)
-
-                # Predict noise residual
-                with torch.autocast(device_type="cuda",
-                                    dtype=target_dtype,
-                                    enabled=autocast_enabled):
-                    if (vsa_available and self.attn_backend
-                            == VideoSparseAttentionBackend):
-                        self.attn_metadata_builder_cls = self.attn_backend.get_builder_cls(
-                        )
-
-                        if self.attn_metadata_builder_cls is not None:
-                            self.attn_metadata_builder = self.attn_metadata_builder_cls(
-                            )
-                            # TODO(will): clean this up
-                            attn_metadata = self.attn_metadata_builder.build(  # type: ignore
-                                current_timestep=i,  # type: ignore
-                                raw_latent_shape=batch.
-                                raw_latent_shape[2:5],  # type: ignore
-                                patch_size=fastvideo_args.
-                                pipeline_config.  # type: ignore
-                                dit_config.patch_size,  # type: ignore
-                                STA_param=batch.STA_param,  # type: ignore
-                                VSA_sparsity=fastvideo_args.
-                                VSA_sparsity,  # type: ignore
-                                device=get_local_torch_device(),  # type: ignore
-                            )  # type: ignore
-                            assert attn_metadata is not None, "attn_metadata cannot be None"
-                        else:
-                            attn_metadata = None
-                    else:
-                        attn_metadata = None
-
-                    batch.is_cfg_negative = False
-                    with set_forward_context(
-                            current_timestep=i,
-                            attn_metadata=attn_metadata,
-                            forward_batch=batch,
-                            # fastvideo_args=fastvideo_args
-                    ):
-                        # Run transformer
-                        pred_noise = self.transformer(
-                            latent_model_input.permute(0, 2, 1, 3, 4),
-                            prompt_embeds,
-                            t_expand,
-                            guidance=guidance_expand,
-                            **image_kwargs,
-                            **pos_cond_kwargs,
-                        ).permute(0, 2, 1, 3, 4)
-
-                    pred_video = pred_noise_to_pred_video(
-                        pred_noise=pred_noise.flatten(0, 1),
-                        noise_input_latent=noise_latents.flatten(0, 1),
-                        timestep=t_expand,
-                        scheduler=self.scheduler).unflatten(
-                            0, pred_noise.shape[:2])
-
-                    if i < len(timesteps) - 1:
-                        next_timestep = timesteps[i + 1] * torch.ones(
-                            [1], dtype=torch.long, device=pred_video.device)
-                        noise = torch.randn(video_raw_latent_shape,
-                                            dtype=pred_video.dtype,
-                                            generator=batch.generator[0]).to(
-                                                self.device)
-                        latents = self.scheduler.add_noise(
-                            pred_video.flatten(0, 1), noise.flatten(0, 1),
-                            next_timestep).unflatten(0, pred_video.shape[:2])
-                    else:
-                        latents = pred_video
-
-                    # Update progress bar
-                    if i == len(timesteps) - 1 or (
-                        (i + 1) > num_warmup_steps and
-                        (i + 1) % self.scheduler.order == 0
-                            and progress_bar is not None):
-                        progress_bar.update()
-
-        # Gather results if using sequence parallelism
-        latents = latents.permute(0, 2, 1, 3, 4)
-        # Update batch with final latents
-        batch.latents = latents
-
-        return batch

fastvideo/pipelines/stages/denoising_causal_strategy.py

@@ -0,0 +1,623 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+Causal block denoising strategy (DMD).
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+
+from fastvideo.distributed import get_local_torch_device
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.forward_context import set_forward_context
+from fastvideo.models.utils import pred_noise_to_pred_video, pred_noise_to_x_bound
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+from fastvideo.pipelines.stages.denoising_strategies import (
+    BlockContext,
+    BlockDenoisingStrategy,
+    BlockPlan,
+    BlockPlanItem,
+    ModelInputs,
+    StrategyState,
+)
+
+try:
+    from fastvideo.attention.backends.sliding_tile_attn import (
+        SlidingTileAttentionBackend)
+    st_attn_available = True
+except ImportError:
+    st_attn_available = False
+    SlidingTileAttentionBackend = None  # type: ignore
+
+try:
+    from fastvideo.attention.backends.video_sparse_attn import (
+        VideoSparseAttentionBackend)
+    vsa_available = True
+except ImportError:
+    vsa_available = False
+    VideoSparseAttentionBackend = None  # type: ignore
+
+
+class CausalBlockStrategy(BlockDenoisingStrategy):
+
+    def __init__(self, stage: Any) -> None:
+        self.stage = stage
+        self.num_transformer_blocks = 0
+        self.num_frames_per_block = 0
+        self.sliding_window_num_frames = 0
+        self.local_attn_size = -1
+        self.frame_seq_length = 0
+
+    def _ensure_model_constants(self) -> None:
+        transformer = self.stage.transformer
+        self.num_transformer_blocks = len(transformer.blocks)
+        arch_config = transformer.config.arch_config
+        self.num_frames_per_block = arch_config.num_frames_per_block
+        self.sliding_window_num_frames = arch_config.sliding_window_num_frames
+        try:
+            self.local_attn_size = getattr(transformer.model, "local_attn_size",
+                                           -1)
+        except Exception:
+            self.local_attn_size = -1
+
+    def _initialize_kv_cache(self, batch_size, dtype, device) -> list[dict]:
+        kv_cache1 = []
+        num_attention_heads = self.stage.transformer.num_attention_heads
+        attention_head_dim = self.stage.transformer.attention_head_dim
+        if self.local_attn_size != -1:
+            kv_cache_size = self.local_attn_size * self.frame_seq_length
+        else:
+            kv_cache_size = self.frame_seq_length * self.sliding_window_num_frames
+
+        for _ in range(self.num_transformer_blocks):
+            kv_cache1.append({
+                "k":
+                torch.zeros([
+                    batch_size, kv_cache_size, num_attention_heads,
+                    attention_head_dim
+                ],
+                            dtype=dtype,
+                            device=device),
+                "v":
+                torch.zeros([
+                    batch_size, kv_cache_size, num_attention_heads,
+                    attention_head_dim
+                ],
+                            dtype=dtype,
+                            device=device),
+                "global_end_index":
+                torch.tensor([0], dtype=torch.long, device=device),
+                "local_end_index":
+                torch.tensor([0], dtype=torch.long, device=device),
+            })
+
+        return kv_cache1
+
+    def _initialize_crossattn_cache(self, batch_size, max_text_len, dtype,
+                                    device) -> list[dict]:
+        crossattn_cache = []
+        num_attention_heads = self.stage.transformer.num_attention_heads
+        attention_head_dim = self.stage.transformer.attention_head_dim
+        for _ in range(self.num_transformer_blocks):
+            crossattn_cache.append({
+                "k":
+                torch.zeros([
+                    batch_size, max_text_len, num_attention_heads,
+                    attention_head_dim
+                ],
+                            dtype=dtype,
+                            device=device),
+                "v":
+                torch.zeros([
+                    batch_size, max_text_len, num_attention_heads,
+                    attention_head_dim
+                ],
+                            dtype=dtype,
+                            device=device),
+                "is_init":
+                False,
+            })
+        return crossattn_cache
+
+    def prepare(self, batch: ForwardBatch,
+                fastvideo_args: FastVideoArgs) -> StrategyState:
+        target_dtype = torch.bfloat16
+        autocast_enabled = (target_dtype != torch.float32
+                            ) and not fastvideo_args.disable_autocast
+
+        self._ensure_model_constants()
+        latents = batch.latents
+        if latents is None:
+            raise ValueError("latents must be provided")
+
+        latent_seq_length = latents.shape[-1] * latents.shape[-2]
+        patch_ratio = (
+            self.stage.transformer.config.arch_config.patch_size[-1] *
+            self.stage.transformer.config.arch_config.patch_size[-2])
+        self.frame_seq_length = latent_seq_length // patch_ratio
+
+        independent_first_frame = getattr(self.stage.transformer,
+                                          "independent_first_frame", False)
+
+        timesteps = torch.tensor(
+            fastvideo_args.pipeline_config.dmd_denoising_steps,
+            dtype=torch.long).cpu()
+        if getattr(fastvideo_args.pipeline_config, "warp_denoising_step",
+                   False):
+            scheduler_timesteps = torch.cat((
+                self.stage.scheduler.timesteps.cpu(),
+                torch.tensor([0], dtype=torch.float32),
+            ))
+            timesteps = scheduler_timesteps[1000 - timesteps]
+        timesteps = timesteps.to(get_local_torch_device())
+
+        boundary_ratio = fastvideo_args.pipeline_config.dit_config.boundary_ratio
+        if boundary_ratio is not None:
+            boundary_timestep = (boundary_ratio *
+                                 self.stage.scheduler.num_train_timesteps)
+            high_noise_timesteps = timesteps[timesteps >= boundary_timestep]
+        else:
+            boundary_timestep = None
+            high_noise_timesteps = None
+
+        image_kwargs: dict[str, Any] = {}
+        pos_cond_kwargs = self.stage.prepare_extra_func_kwargs(
+            self.stage.transformer.forward,
+            {
+                "encoder_attention_mask": batch.prompt_attention_mask,
+            },
+        )
+
+        if (st_attn_available
+                and self.stage.attn_backend == SlidingTileAttentionBackend):
+            self.stage.prepare_sta_param(batch, fastvideo_args)
+
+        prompt_embeds = batch.prompt_embeds
+        assert torch.isnan(prompt_embeds[0]).sum() == 0
+
+        kv_cache1 = self._initialize_kv_cache(batch_size=latents.shape[0],
+                                              dtype=target_dtype,
+                                              device=latents.device)
+        kv_cache2 = None
+        if boundary_timestep is not None:
+            kv_cache2 = self._initialize_kv_cache(
+                batch_size=latents.shape[0],
+                dtype=target_dtype,
+                device=latents.device,
+            )
+
+        text_len = None
+        if fastvideo_args.pipeline_config.text_encoder_configs:
+            text_len = getattr(
+                fastvideo_args.pipeline_config.text_encoder_configs[0].
+                arch_config, "text_len", None)
+        if not text_len:
+            if batch.prompt_attention_mask:
+                text_len = batch.prompt_attention_mask[0].shape[-1]
+            elif batch.prompt_embeds:
+                text_len = batch.prompt_embeds[0].shape[1]
+            else:
+                text_len = 0
+
+        crossattn_cache = self._initialize_crossattn_cache(
+            batch_size=latents.shape[0],
+            max_text_len=text_len,
+            dtype=target_dtype,
+            device=latents.device,
+        )
+
+        num_frames = latents.shape[2]
+        if num_frames % self.num_frames_per_block != 0:
+            raise ValueError(
+                "num_frames must be divisible by num_frames_per_block for "
+                "causal DMD denoising")
+        num_blocks = num_frames // self.num_frames_per_block
+        block_sizes = [self.num_frames_per_block] * num_blocks
+        start_index = 0
+
+        if boundary_timestep is not None:
+            block_sizes[0] = 1
+
+        pos_start_base = 0
+        if batch.pil_image is not None:
+            assert self.stage.vae is not None, (
+                "VAE is not provided for causal video gen task")
+            self.stage.vae = self.stage.vae.to(get_local_torch_device())
+            first_frame_latent = self.stage.vae.encode(
+                batch.pil_image).mean.float()
+            if (hasattr(self.stage.vae, "shift_factor")
+                    and self.stage.vae.shift_factor is not None):
+                if isinstance(self.stage.vae.shift_factor, torch.Tensor):
+                    first_frame_latent -= self.stage.vae.shift_factor.to(
+                        first_frame_latent.device, first_frame_latent.dtype)
+                else:
+                    first_frame_latent -= self.stage.vae.shift_factor
+
+            if isinstance(self.stage.vae.scaling_factor, torch.Tensor):
+                first_frame_latent = (
+                    first_frame_latent * self.stage.vae.scaling_factor.to(
+                        first_frame_latent.device, first_frame_latent.dtype))
+            else:
+                first_frame_latent = (first_frame_latent *
+                                      self.stage.vae.scaling_factor)
+
+            if fastvideo_args.vae_cpu_offload:
+                self.stage.vae = self.stage.vae.to("cpu")
+
+            t_zero = torch.zeros([latents.shape[0], 1],
+                                 device=latents.device,
+                                 dtype=torch.long)
+            with torch.autocast(device_type="cuda",
+                                dtype=target_dtype,
+                                enabled=autocast_enabled), \
+                set_forward_context(current_timestep=0,
+                                    attn_metadata=None,
+                                    forward_batch=batch):
+                self.stage.transformer(
+                    first_frame_latent.to(target_dtype),
+                    prompt_embeds,
+                    t_zero,
+                    kv_cache=kv_cache1,
+                    crossattn_cache=crossattn_cache,
+                    current_start=(pos_start_base + start_index) *
+                    self.frame_seq_length,
+                    start_frame=start_index,
+                    **image_kwargs,
+                    **pos_cond_kwargs,
+                )
+                if boundary_timestep is not None:
+                    self.stage.transformer_2(
+                        first_frame_latent.to(target_dtype),
+                        prompt_embeds,
+                        t_zero,
+                        kv_cache=kv_cache2,
+                        crossattn_cache=crossattn_cache,
+                        current_start=(pos_start_base + start_index) *
+                        self.frame_seq_length,
+                        start_frame=start_index,
+                        **image_kwargs,
+                        **pos_cond_kwargs,
+                    )
+
+            start_index += 1
+            block_sizes.pop(0)
+            latents[:, :, :1, :, :] = first_frame_latent
+
+        progress_bar = self.stage.progress_bar(total=len(block_sizes) *
+                                               len(timesteps))
+
+        extra: dict[str, Any] = {
+            "batch":
+            batch,
+            "fastvideo_args":
+            fastvideo_args,
+            "target_dtype":
+            target_dtype,
+            "autocast_enabled":
+            autocast_enabled,
+            "boundary_timestep":
+            boundary_timestep,
+            "high_noise_timesteps":
+            high_noise_timesteps,
+            "image_kwargs":
+            image_kwargs,
+            "pos_cond_kwargs":
+            pos_cond_kwargs,
+            "kv_cache1":
+            kv_cache1,
+            "kv_cache2":
+            kv_cache2,
+            "crossattn_cache":
+            crossattn_cache,
+            "block_sizes":
+            block_sizes,
+            "start_index":
+            start_index,
+            "progress_bar":
+            progress_bar,
+            "independent_first_frame":
+            independent_first_frame,
+            "context_noise":
+            getattr(fastvideo_args.pipeline_config, "context_noise", 0),
+            "pos_start_base":
+            pos_start_base,
+        }
+
+        return StrategyState(
+            latents=latents,
+            timesteps=timesteps,
+            num_inference_steps=len(timesteps),
+            prompt_embeds=batch.prompt_embeds,
+            negative_prompt_embeds=batch.negative_prompt_embeds,
+            prompt_attention_mask=batch.prompt_attention_mask,
+            negative_attention_mask=batch.negative_attention_mask,
+            image_embeds=batch.image_embeds,
+            guidance_scale=batch.guidance_scale,
+            guidance_scale_2=batch.guidance_scale_2,
+            guidance_rescale=batch.guidance_rescale,
+            do_cfg=batch.do_classifier_free_guidance,
+            extra=extra,
+        )
+
+    def block_plan(self, state: StrategyState) -> BlockPlan:
+        block_sizes = state.extra["block_sizes"]
+        start_index = state.extra["start_index"]
+        items: list[BlockPlanItem] = []
+        for block_size in block_sizes:
+            items.append(
+                BlockPlanItem(
+                    start_index=start_index,
+                    num_frames=block_size,
+                    use_kv_cache=True,
+                    model_selector="default",
+                ))
+            start_index += block_size
+        return BlockPlan(items=items)
+
+    def init_block_context(self, state: StrategyState,
+                           block_item: BlockPlanItem,
+                           block_idx: int) -> BlockContext:
+        return BlockContext(
+            kv_cache=state.extra["kv_cache1"],
+            kv_cache_2=state.extra["kv_cache2"],
+            crossattn_cache=state.extra["crossattn_cache"],
+            action_cache=None,
+            extra={
+                "block_idx": block_idx,
+                "start_index": block_item.start_index,
+                "num_frames": block_item.num_frames,
+            },
+        )
+
+    def process_block(self, state: StrategyState, block_ctx: BlockContext,
+                      block_item: BlockPlanItem) -> None:
+        batch = state.extra["batch"]
+        fastvideo_args = state.extra["fastvideo_args"]
+        target_dtype = state.extra["target_dtype"]
+        autocast_enabled = state.extra["autocast_enabled"]
+        boundary_timestep = state.extra["boundary_timestep"]
+        high_noise_timesteps = state.extra["high_noise_timesteps"]
+        image_kwargs = state.extra["image_kwargs"]
+        pos_cond_kwargs = state.extra["pos_cond_kwargs"]
+        progress_bar = state.extra["progress_bar"]
+        independent_first_frame = state.extra["independent_first_frame"]
+
+        start_index = block_item.start_index
+        current_num_frames = block_item.num_frames
+
+        kv_cache1 = block_ctx.kv_cache
+        kv_cache2 = block_ctx.kv_cache_2
+        crossattn_cache = block_ctx.crossattn_cache
+
+        def _get_kv_cache(timestep_val: float) -> list[dict]:
+            if boundary_timestep is not None:
+                if timestep_val >= boundary_timestep:
+                    return kv_cache1
+                if kv_cache2 is None:
+                    raise ValueError("kv_cache2 is not initialized")
+                return kv_cache2
+            return kv_cache1
+
+        current_latents = state.latents[:, :, start_index:start_index +
+                                        current_num_frames, :, :]
+        noise_latents_btchw = current_latents.permute(0, 2, 1, 3, 4)
+        video_raw_latent_shape = noise_latents_btchw.shape
+        h, w = current_latents.shape[-2:]
+
+        attn_metadata = None
+        for i, t_cur in enumerate(state.timesteps):
+            if boundary_timestep is not None and t_cur < boundary_timestep:
+                current_model = self.stage.transformer_2
+            else:
+                current_model = self.stage.transformer
+            noise_latents = noise_latents_btchw.clone()
+            latent_model_input = current_latents.to(target_dtype)
+
+            if (batch.image_latent is not None and independent_first_frame
+                    and start_index == 0):
+                latent_model_input = torch.cat(
+                    [latent_model_input,
+                     batch.image_latent.to(target_dtype)],
+                    dim=2)
+
+            t_expand = t_cur.repeat(latent_model_input.shape[0])
+
+            if (vsa_available
+                    and self.stage.attn_backend == VideoSparseAttentionBackend):
+                self.attn_metadata_builder_cls = (
+                    self.stage.attn_backend.get_builder_cls())
+                if self.attn_metadata_builder_cls is not None:
+                    self.attn_metadata_builder = (
+                        self.attn_metadata_builder_cls())
+                    attn_metadata = self.attn_metadata_builder.build(  # type: ignore
+                        current_timestep=i,  # type: ignore
+                        raw_latent_shape=(current_num_frames, h,
+                                          w),  # type: ignore
+                        patch_size=fastvideo_args.pipeline_config.dit_config.
+                        patch_size,  # type: ignore
+                        STA_param=batch.STA_param,  # type: ignore
+                        VSA_sparsity=fastvideo_args.
+                        VSA_sparsity,  # type: ignore
+                        device=get_local_torch_device(),  # type: ignore
+                    )
+                    assert attn_metadata is not None, (
+                        "attn_metadata cannot be None")
+                else:
+                    attn_metadata = None
+            else:
+                attn_metadata = None
+
+            with torch.autocast(device_type="cuda",
+                                dtype=target_dtype,
+                                enabled=autocast_enabled), \
+                set_forward_context(current_timestep=i,
+                                    attn_metadata=attn_metadata,
+                                    forward_batch=batch):
+                t_expanded_noise = t_cur * torch.ones(
+                    (latent_model_input.shape[0], 1),
+                    device=latent_model_input.device,
+                    dtype=torch.long)
+                pred_noise_btchw = current_model(
+                    latent_model_input,
+                    batch.prompt_embeds,
+                    t_expanded_noise,
+                    kv_cache=_get_kv_cache(t_cur),
+                    crossattn_cache=crossattn_cache,
+                    current_start=start_index * self.frame_seq_length,
+                    start_frame=start_index,
+                    **image_kwargs,
+                    **pos_cond_kwargs,
+                ).permute(0, 2, 1, 3, 4)
+
+            if boundary_timestep is not None and t_cur >= boundary_timestep:
+                pred_video_btchw = pred_noise_to_x_bound(
+                    pred_noise=pred_noise_btchw.flatten(0, 1),
+                    noise_input_latent=noise_latents.flatten(0, 1),
+                    timestep=t_expand,
+                    boundary_timestep=torch.ones_like(t_expand) *
+                    boundary_timestep,
+                    scheduler=self.stage.scheduler,
+                ).unflatten(0, pred_noise_btchw.shape[:2])
+            else:
+                pred_video_btchw = pred_noise_to_pred_video(
+                    pred_noise=pred_noise_btchw.flatten(0, 1),
+                    noise_input_latent=noise_latents.flatten(0, 1),
+                    timestep=t_expand,
+                    scheduler=self.stage.scheduler,
+                ).unflatten(0, pred_noise_btchw.shape[:2])
+
+            if i < len(state.timesteps) - 1:
+                next_timestep = state.timesteps[i + 1] * torch.ones(
+                    [1],
+                    dtype=torch.long,
+                    device=pred_video_btchw.device,
+                )
+                noise = torch.randn(
+                    video_raw_latent_shape,
+                    dtype=pred_video_btchw.dtype,
+                    generator=(batch.generator[0] if isinstance(
+                        batch.generator, list) else batch.generator),
+                ).to(self.stage.device)
+                noise_btchw = noise
+                if (boundary_timestep is not None
+                        and high_noise_timesteps is not None
+                        and i < len(high_noise_timesteps) - 1):
+                    noise_latents_btchw = self.stage.scheduler.add_noise_high(
+                        pred_video_btchw.flatten(0, 1),
+                        noise_btchw.flatten(0, 1),
+                        next_timestep,
+                        torch.ones_like(next_timestep) * boundary_timestep,
+                    ).unflatten(0, pred_video_btchw.shape[:2])
+                elif (boundary_timestep is not None
+                      and high_noise_timesteps is not None
+                      and i == len(high_noise_timesteps) - 1):
+                    noise_latents_btchw = pred_video_btchw
+                else:
+                    noise_latents_btchw = self.stage.scheduler.add_noise(
+                        pred_video_btchw.flatten(0, 1),
+                        noise_btchw.flatten(0, 1),
+                        next_timestep,
+                    ).unflatten(0, pred_video_btchw.shape[:2])
+                current_latents = noise_latents_btchw.permute(0, 2, 1, 3, 4)
+            else:
+                current_latents = pred_video_btchw.permute(0, 2, 1, 3, 4)
+
+            if progress_bar is not None:
+                progress_bar.update()
+
+        block_ctx.extra["attn_metadata"] = attn_metadata
+        state.latents[:, :, start_index:start_index +
+                      current_num_frames, :, :] = (current_latents)
+
+    def update_context(self, state: StrategyState, block_ctx: BlockContext,
+                       block_item: BlockPlanItem) -> None:
+        batch = state.extra["batch"]
+        target_dtype = state.extra["target_dtype"]
+        autocast_enabled = state.extra["autocast_enabled"]
+        boundary_timestep = state.extra["boundary_timestep"]
+        kv_cache1 = state.extra["kv_cache1"]
+        kv_cache2 = state.extra["kv_cache2"]
+        crossattn_cache = state.extra["crossattn_cache"]
+        image_kwargs = state.extra["image_kwargs"]
+        pos_cond_kwargs = state.extra["pos_cond_kwargs"]
+        context_noise = state.extra["context_noise"]
+
+        start_index = block_item.start_index
+        current_num_frames = block_item.num_frames
+
+        current_latents = state.latents[:, :, start_index:start_index +
+                                        current_num_frames, :, :]
+        latents_device = current_latents.device
+
+        t_context = torch.ones([current_latents.shape[0]],
+                               device=latents_device,
+                               dtype=torch.long) * int(context_noise)
+        context_bcthw = current_latents.to(target_dtype)
+
+        attn_metadata = block_ctx.extra.get("attn_metadata")
+        with torch.autocast(device_type="cuda",
+                            dtype=target_dtype,
+                            enabled=autocast_enabled), \
+            set_forward_context(current_timestep=0,
+                                attn_metadata=attn_metadata,
+                                forward_batch=batch):
+            t_expanded_context = t_context.unsqueeze(1)
+
+            if boundary_timestep is not None:
+                self.stage.transformer_2(
+                    context_bcthw,
+                    batch.prompt_embeds,
+                    t_expanded_context,
+                    kv_cache=kv_cache2,
+                    crossattn_cache=crossattn_cache,
+                    current_start=start_index * self.frame_seq_length,
+                    start_frame=start_index,
+                    **image_kwargs,
+                    **pos_cond_kwargs,
+                )
+
+            self.stage.transformer(
+                context_bcthw,
+                batch.prompt_embeds,
+                t_expanded_context,
+                kv_cache=kv_cache1,
+                crossattn_cache=crossattn_cache,
+                current_start=start_index * self.frame_seq_length,
+                start_frame=start_index,
+                **image_kwargs,
+                **pos_cond_kwargs,
+            )
+
+    def postprocess(self, state: StrategyState) -> ForwardBatch:
+        progress_bar = state.extra.get("progress_bar")
+        if progress_bar is not None:
+            progress_bar.close()
+
+        batch = state.extra["batch"]
+        boundary_timestep = state.extra["boundary_timestep"]
+
+        latents = state.latents
+        if boundary_timestep is not None:
+            num_frames_to_remove = self.num_frames_per_block - 1
+            latents = latents[:, :, :-num_frames_to_remove, :, :]
+
+        batch.latents = latents
+        return batch
+
+    def make_model_inputs(self, state: StrategyState, t: torch.Tensor,
+                          step_idx: int) -> ModelInputs:
+        raise NotImplementedError
+
+    def forward(self, state: StrategyState,
+                model_inputs: ModelInputs) -> torch.Tensor:
+        raise NotImplementedError
+
+    def cfg_combine(self, state: StrategyState,
+                    noise_pred: torch.Tensor) -> torch.Tensor:
+        raise NotImplementedError
+
+    def scheduler_step(self, state: StrategyState, noise_pred: torch.Tensor,
+                       t: torch.Tensor) -> torch.Tensor:
+        raise NotImplementedError

fastvideo/pipelines/stages/denoising_cosmos_strategy.py

@@ -0,0 +1,325 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+Cosmos denoising strategy using FlowMatchEulerDiscreteScheduler.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.forward_context import set_forward_context
+from fastvideo.logger import init_logger
+from fastvideo.models.loader.component_loader import TransformerLoader
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+from fastvideo.pipelines.stages.denoising_strategies import (
+    DenoisingStrategy,
+    ModelInputs,
+    StrategyState,
+)
+
+logger = init_logger(__name__)
+
+
+class CosmosStrategy(DenoisingStrategy):
+
+    def __init__(self, stage: Any) -> None:
+        self.stage = stage
+
+    def prepare(self, batch: ForwardBatch,
+                fastvideo_args: FastVideoArgs) -> StrategyState:
+        pipeline = self.stage.pipeline() if self.stage.pipeline else None
+        if not fastvideo_args.model_loaded["transformer"]:
+            loader = TransformerLoader()
+            self.stage.transformer = loader.load(
+                fastvideo_args.model_paths["transformer"], fastvideo_args)
+            if pipeline:
+                pipeline.add_module("transformer", self.stage.transformer)
+            fastvideo_args.model_loaded["transformer"] = True
+
+        extra_step_kwargs = self.stage.prepare_extra_func_kwargs(
+            self.stage.scheduler.step,
+            {
+                "generator": batch.generator,
+                "eta": batch.eta
+            },
+        )
+
+        if hasattr(self.stage.transformer, "module"):
+            transformer_dtype = next(
+                self.stage.transformer.module.parameters()).dtype
+        else:
+            transformer_dtype = next(self.stage.transformer.parameters()).dtype
+        target_dtype = transformer_dtype
+        autocast_enabled = (target_dtype != torch.float32
+                            ) and not fastvideo_args.disable_autocast
+
+        latents = batch.latents
+        num_inference_steps = batch.num_inference_steps
+
+        sigma_max = 80.0
+        sigma_min = 0.002
+        sigma_data = 1.0
+        final_sigmas_type = "sigma_min"
+
+        if self.stage.scheduler is not None:
+            self.stage.scheduler.register_to_config(
+                sigma_max=sigma_max,
+                sigma_min=sigma_min,
+                sigma_data=sigma_data,
+                final_sigmas_type=final_sigmas_type,
+            )
+
+        self.stage.scheduler.set_timesteps(num_inference_steps,
+                                           device=latents.device)
+        timesteps = self.stage.scheduler.timesteps
+
+        if (hasattr(self.stage.scheduler.config, "final_sigmas_type")
+                and self.stage.scheduler.config.final_sigmas_type == "sigma_min"
+                and len(self.stage.scheduler.sigmas) > 1):
+            self.stage.scheduler.sigmas[-1] = self.stage.scheduler.sigmas[-2]
+
+        conditioning_latents = getattr(batch, "conditioning_latents", None)
+        unconditioning_latents = conditioning_latents
+
+        progress_bar = self.stage.progress_bar(total=num_inference_steps)
+
+        extra: dict[str, Any] = {
+            "batch": batch,
+            "fastvideo_args": fastvideo_args,
+            "extra_step_kwargs": extra_step_kwargs,
+            "target_dtype": target_dtype,
+            "autocast_enabled": autocast_enabled,
+            "progress_bar": progress_bar,
+            "conditioning_latents": conditioning_latents,
+            "unconditioning_latents": unconditioning_latents,
+        }
+
+        return StrategyState(
+            latents=latents,
+            timesteps=timesteps,
+            num_inference_steps=num_inference_steps,
+            prompt_embeds=batch.prompt_embeds,
+            negative_prompt_embeds=batch.negative_prompt_embeds,
+            prompt_attention_mask=batch.prompt_attention_mask,
+            negative_attention_mask=batch.negative_attention_mask,
+            image_embeds=batch.image_embeds,
+            guidance_scale=batch.guidance_scale,
+            guidance_scale_2=batch.guidance_scale_2,
+            guidance_rescale=batch.guidance_rescale,
+            do_cfg=batch.do_classifier_free_guidance,
+            extra=extra,
+        )
+
+    def make_model_inputs(self, state: StrategyState, t: torch.Tensor,
+                          step_idx: int) -> ModelInputs:
+        state.extra["step_idx"] = step_idx
+        return ModelInputs(
+            latent_model_input=state.latents,
+            timestep=t,
+            prompt_embeds=state.prompt_embeds,
+            prompt_attention_mask=state.prompt_attention_mask,
+        )
+
+    def forward(self, state: StrategyState,
+                model_inputs: ModelInputs) -> torch.Tensor:
+        batch = state.extra["batch"]
+        target_dtype = state.extra["target_dtype"]
+        autocast_enabled = state.extra["autocast_enabled"]
+        conditioning_latents = state.extra["conditioning_latents"]
+        unconditioning_latents = state.extra["unconditioning_latents"]
+        step_idx = state.extra["step_idx"]
+
+        if getattr(self.stage, "interrupt", False):
+            return state.latents
+
+        current_sigma = self.stage.scheduler.sigmas[step_idx]
+        current_t = current_sigma / (current_sigma + 1)
+        c_in = 1 - current_t
+        c_skip = 1 - current_t
+        c_out = -current_t
+
+        timestep = current_t.view(1, 1, 1, 1,
+                                  1).expand(state.latents.size(0), -1,
+                                            state.latents.size(2), -1, -1)
+
+        with torch.autocast(device_type="cuda",
+                            dtype=target_dtype,
+                            enabled=autocast_enabled):
+            cond_latent = state.latents * c_in
+
+            if (hasattr(batch, "cond_indicator")
+                    and batch.cond_indicator is not None
+                    and conditioning_latents is not None):
+                cond_latent = (batch.cond_indicator * conditioning_latents +
+                               (1 - batch.cond_indicator) * cond_latent)
+            else:
+                logger.warning(
+                    "Step %s: Missing conditioning data - "
+                    "cond_indicator: %s, conditioning_latents: %s", step_idx,
+                    hasattr(batch, "cond_indicator"), conditioning_latents
+                    is not None)
+
+            cond_latent = cond_latent.to(target_dtype)
+
+            cond_timestep = timestep
+            if hasattr(batch,
+                       "cond_indicator") and batch.cond_indicator is not None:
+                sigma_conditioning = 0.0001
+                t_conditioning = sigma_conditioning / (sigma_conditioning + 1)
+                cond_timestep = (batch.cond_indicator * t_conditioning +
+                                 (1 - batch.cond_indicator) * timestep)
+                cond_timestep = cond_timestep.to(target_dtype)
+
+            with set_forward_context(
+                    current_timestep=step_idx,
+                    attn_metadata=None,
+                    forward_batch=batch,
+            ):
+                condition_mask = (batch.cond_mask.to(target_dtype) if hasattr(
+                    batch, "cond_mask") else None)
+                padding_mask = torch.zeros(1,
+                                           1,
+                                           batch.height,
+                                           batch.width,
+                                           device=cond_latent.device,
+                                           dtype=target_dtype)
+
+                if condition_mask is None:
+                    batch_size, _, num_frames, height, width = cond_latent.shape
+                    condition_mask = torch.zeros(batch_size,
+                                                 1,
+                                                 num_frames,
+                                                 height,
+                                                 width,
+                                                 device=cond_latent.device,
+                                                 dtype=target_dtype)
+
+                noise_pred = self.stage.transformer(
+                    hidden_states=cond_latent,
+                    timestep=cond_timestep.to(target_dtype),
+                    encoder_hidden_states=batch.prompt_embeds[0].to(
+                        target_dtype),
+                    fps=24,
+                    condition_mask=condition_mask,
+                    padding_mask=padding_mask,
+                    return_dict=False,
+                )[0]
+
+            cond_pred = (c_skip * state.latents +
+                         c_out * noise_pred.float()).to(target_dtype)
+
+            if (hasattr(batch, "cond_indicator")
+                    and batch.cond_indicator is not None
+                    and conditioning_latents is not None):
+                cond_pred = (batch.cond_indicator * conditioning_latents +
+                             (1 - batch.cond_indicator) * cond_pred)
+
+            if (state.do_cfg and batch.negative_prompt_embeds is not None):
+                uncond_latent = state.latents * c_in
+
+                if (hasattr(batch, "uncond_indicator")
+                        and batch.uncond_indicator is not None
+                        and unconditioning_latents is not None):
+                    uncond_latent = (
+                        batch.uncond_indicator * unconditioning_latents +
+                        (1 - batch.uncond_indicator) * uncond_latent)
+
+                with set_forward_context(
+                        current_timestep=step_idx,
+                        attn_metadata=None,
+                        forward_batch=batch,
+                ):
+                    uncond_condition_mask = (
+                        batch.uncond_mask.to(target_dtype) if
+                        (hasattr(batch, "uncond_mask")
+                         and batch.uncond_mask is not None) else condition_mask)
+
+                    uncond_timestep = timestep
+                    if (hasattr(batch, "uncond_indicator")
+                            and batch.uncond_indicator is not None):
+                        sigma_conditioning = 0.0001
+                        t_conditioning = sigma_conditioning / (
+                            sigma_conditioning + 1)
+                        uncond_timestep = (
+                            batch.uncond_indicator * t_conditioning +
+                            (1 - batch.uncond_indicator) * timestep)
+                        uncond_timestep = uncond_timestep.to(target_dtype)
+
+                    noise_pred_uncond = self.stage.transformer(
+                        hidden_states=uncond_latent.to(target_dtype),
+                        timestep=uncond_timestep.to(target_dtype),
+                        encoder_hidden_states=batch.negative_prompt_embeds[0].
+                        to(target_dtype),
+                        fps=24,
+                        condition_mask=uncond_condition_mask,
+                        padding_mask=padding_mask,
+                        return_dict=False,
+                    )[0]
+
+                uncond_pred = (
+                    c_skip * state.latents +
+                    c_out * noise_pred_uncond.float()).to(target_dtype)
+
+                if (hasattr(batch, "uncond_indicator")
+                        and batch.uncond_indicator is not None
+                        and unconditioning_latents is not None):
+                    uncond_pred = (
+                        batch.uncond_indicator * unconditioning_latents +
+                        (1 - batch.uncond_indicator) * uncond_pred)
+
+                guidance_diff = cond_pred - uncond_pred
+                final_pred = cond_pred + state.guidance_scale * guidance_diff
+            else:
+                final_pred = cond_pred
+
+        if current_sigma > 1e-8:
+            noise_for_scheduler = (state.latents - final_pred) / current_sigma
+        else:
+            logger.warning(
+                "Step %s: current_sigma too small (%s), using final_pred directly",
+                step_idx, current_sigma)
+            noise_for_scheduler = final_pred
+
+        if torch.isnan(noise_for_scheduler).sum() > 0:
+            logger.error(
+                "Step %s: NaN detected in noise_for_scheduler, sum: %s",
+                step_idx,
+                noise_for_scheduler.float().sum().item())
+            logger.error(
+                "Step %s: latents sum: %s, final_pred sum: %s, current_sigma: %s",
+                step_idx,
+                state.latents.float().sum().item(),
+                final_pred.float().sum().item(), current_sigma)
+
+        return noise_for_scheduler
+
+    def cfg_combine(self, state: StrategyState,
+                    noise_pred: torch.Tensor) -> torch.Tensor:
+        return noise_pred
+
+    def scheduler_step(self, state: StrategyState, noise_pred: torch.Tensor,
+                       t: torch.Tensor) -> torch.Tensor:
+        latents = self.stage.scheduler.step(
+            noise_pred,
+            t,
+            state.latents,
+            **state.extra["extra_step_kwargs"],
+            return_dict=False,
+        )[0]
+
+        progress_bar = state.extra["progress_bar"]
+        if progress_bar is not None:
+            progress_bar.update()
+
+        return latents
+
+    def postprocess(self, state: StrategyState) -> ForwardBatch:
+        progress_bar = state.extra.get("progress_bar")
+        if progress_bar is not None:
+            progress_bar.close()
+        batch = state.extra["batch"]
+        batch.latents = state.latents
+        return batch

fastvideo/pipelines/stages/denoising_dmd_strategy.py

@@ -0,0 +1,269 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+DMD denoising strategy (FlowMatch-based).
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+
+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.models.utils import pred_noise_to_pred_video
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+from fastvideo.pipelines.stages.denoising_strategies import (
+    DenoisingStrategy,
+    ModelInputs,
+    StrategyState,
+)
+from fastvideo.utils import dict_to_3d_list
+
+try:
+    from fastvideo.attention.backends.sliding_tile_attn import (
+        SlidingTileAttentionBackend)
+    st_attn_available = True
+except ImportError:
+    st_attn_available = False
+    SlidingTileAttentionBackend = None  # type: ignore
+
+try:
+    from fastvideo.attention.backends.video_sparse_attn import (
+        VideoSparseAttentionBackend)
+    vsa_available = True
+except ImportError:
+    vsa_available = False
+    VideoSparseAttentionBackend = None  # type: ignore
+
+logger = init_logger(__name__)
+
+
+class DmdStrategy(DenoisingStrategy):
+
+    def __init__(self, stage: Any) -> None:
+        self.stage = stage
+
+    def prepare(self, batch: ForwardBatch,
+                fastvideo_args: FastVideoArgs) -> StrategyState:
+        target_dtype = torch.bfloat16
+        autocast_enabled = (target_dtype != torch.float32
+                            ) and not fastvideo_args.disable_autocast
+
+        timesteps = batch.timesteps
+        if timesteps is None:
+            raise ValueError("Timesteps must be provided")
+        num_inference_steps = batch.num_inference_steps
+        num_warmup_steps = len(
+            timesteps) - num_inference_steps * self.stage.scheduler.order
+
+        image_embeds = batch.image_embeds
+        if len(image_embeds) > 0:
+            assert torch.isnan(image_embeds[0]).sum() == 0
+            image_embeds = [
+                image_embed.to(target_dtype) for image_embed in image_embeds
+            ]
+
+        image_kwargs = self.stage.prepare_extra_func_kwargs(
+            self.stage.transformer.forward,
+            {
+                "encoder_hidden_states_image": image_embeds,
+                "mask_strategy": dict_to_3d_list(
+                    None, t_max=50, l_max=60, h_max=24)
+            },
+        )
+
+        pos_cond_kwargs = self.stage.prepare_extra_func_kwargs(
+            self.stage.transformer.forward,
+            {
+                "encoder_hidden_states_2": batch.clip_embedding_pos,
+                "encoder_attention_mask": batch.prompt_attention_mask,
+            },
+        )
+
+        if (st_attn_available
+                and self.stage.attn_backend == SlidingTileAttentionBackend):
+            self.stage.prepare_sta_param(batch, fastvideo_args)
+
+        assert batch.latents is not None, "latents must be provided"
+        latents = batch.latents
+
+        prompt_embeds = batch.prompt_embeds
+        assert not torch.isnan(
+            prompt_embeds[0]).any(), "prompt_embeds contains nan"
+
+        timesteps = torch.tensor(
+            fastvideo_args.pipeline_config.dmd_denoising_steps,
+            dtype=torch.long,
+            device=get_local_torch_device())
+
+        progress_bar = self.stage.progress_bar(total=len(timesteps))
+
+        extra: dict[str, Any] = {
+            "batch": batch,
+            "fastvideo_args": fastvideo_args,
+            "target_dtype": target_dtype,
+            "autocast_enabled": autocast_enabled,
+            "num_warmup_steps": num_warmup_steps,
+            "image_kwargs": image_kwargs,
+            "pos_cond_kwargs": pos_cond_kwargs,
+            "progress_bar": progress_bar,
+            "video_raw_latent_shape": latents.shape,
+        }
+
+        return StrategyState(
+            latents=latents,
+            timesteps=timesteps,
+            num_inference_steps=num_inference_steps,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=batch.negative_prompt_embeds,
+            prompt_attention_mask=batch.prompt_attention_mask,
+            negative_attention_mask=batch.negative_attention_mask,
+            image_embeds=image_embeds,
+            guidance_scale=batch.guidance_scale,
+            guidance_scale_2=batch.guidance_scale_2,
+            guidance_rescale=batch.guidance_rescale,
+            do_cfg=batch.do_classifier_free_guidance,
+            extra=extra,
+        )
+
+    def make_model_inputs(self, state: StrategyState, t: torch.Tensor,
+                          step_idx: int) -> ModelInputs:
+        state.extra["step_idx"] = step_idx
+        return ModelInputs(
+            latent_model_input=state.latents,
+            timestep=t,
+            prompt_embeds=state.prompt_embeds,
+            prompt_attention_mask=state.prompt_attention_mask,
+        )
+
+    def forward(self, state: StrategyState,
+                model_inputs: ModelInputs) -> torch.Tensor:
+        batch = state.extra["batch"]
+        fastvideo_args = state.extra["fastvideo_args"]
+        target_dtype = state.extra["target_dtype"]
+        autocast_enabled = state.extra["autocast_enabled"]
+        step_idx = state.extra["step_idx"]
+
+        if getattr(self.stage, "interrupt", False):
+            state.extra["pred_latents"] = state.latents
+            return state.latents
+
+        latents = state.latents
+        noise_latents = latents.clone()
+        latent_model_input = latents.to(target_dtype)
+
+        if batch.image_latent is not None:
+            latent_model_input = torch.cat(
+                [latent_model_input,
+                 batch.image_latent.permute(0, 2, 1, 3, 4)],
+                dim=2).to(target_dtype)
+
+        t_expand = model_inputs.timestep.repeat(latent_model_input.shape[0])
+        guidance_expand = None
+        if fastvideo_args.pipeline_config.embedded_cfg_scale is not None:
+            guidance_expand = (torch.tensor(
+                [fastvideo_args.pipeline_config.embedded_cfg_scale] *
+                latent_model_input.shape[0],
+                dtype=torch.float32,
+                device=get_local_torch_device(),
+            ).to(target_dtype) * 1000.0)
+
+        with torch.autocast(device_type="cuda",
+                            dtype=target_dtype,
+                            enabled=autocast_enabled):
+            if (vsa_available
+                    and self.stage.attn_backend == VideoSparseAttentionBackend):
+                self.attn_metadata_builder_cls = (
+                    self.stage.attn_backend.get_builder_cls())
+
+                if self.attn_metadata_builder_cls is not None:
+                    self.attn_metadata_builder = (
+                        self.attn_metadata_builder_cls())
+                    attn_metadata = self.attn_metadata_builder.build(  # type: ignore
+                        current_timestep=step_idx,  # type: ignore
+                        raw_latent_shape=batch.
+                        raw_latent_shape[2:5],  # type: ignore
+                        patch_size=fastvideo_args.pipeline_config.dit_config.
+                        patch_size,  # type: ignore
+                        STA_param=batch.STA_param,  # type: ignore
+                        VSA_sparsity=fastvideo_args.
+                        VSA_sparsity,  # type: ignore
+                        device=get_local_torch_device(),  # type: ignore
+                    )
+                    assert attn_metadata is not None, (
+                        "attn_metadata cannot be None")
+                else:
+                    attn_metadata = None
+            else:
+                attn_metadata = None
+
+            with set_forward_context(
+                    current_timestep=step_idx,
+                    attn_metadata=attn_metadata,
+                    forward_batch=batch,
+            ):
+                pred_noise = self.stage.transformer(
+                    latent_model_input.permute(0, 2, 1, 3, 4),
+                    state.prompt_embeds,
+                    t_expand,
+                    guidance=guidance_expand,
+                    **state.extra["image_kwargs"],
+                    **state.extra["pos_cond_kwargs"],
+                ).permute(0, 2, 1, 3, 4)
+
+        pred_video = pred_noise_to_pred_video(
+            pred_noise=pred_noise.flatten(0, 1),
+            noise_input_latent=noise_latents.flatten(0, 1),
+            timestep=t_expand,
+            scheduler=self.stage.scheduler).unflatten(0, pred_noise.shape[:2])
+
+        if step_idx < len(state.timesteps) - 1:
+            next_timestep = state.timesteps[step_idx + 1] * torch.ones(
+                [1], dtype=torch.long, device=pred_video.device)
+            generator = batch.generator
+            if isinstance(generator, list):
+                generator = generator[0] if generator else None
+            noise = torch.randn(state.extra["video_raw_latent_shape"],
+                                dtype=pred_video.dtype,
+                                generator=generator).to(self.stage.device)
+            latents = self.stage.scheduler.add_noise(pred_video.flatten(0, 1),
+                                                     noise.flatten(0, 1),
+                                                     next_timestep).unflatten(
+                                                         0,
+                                                         pred_video.shape[:2])
+        else:
+            latents = pred_video
+
+        state.extra["pred_latents"] = latents
+        return latents
+
+    def cfg_combine(self, state: StrategyState,
+                    noise_pred: torch.Tensor) -> torch.Tensor:
+        return noise_pred
+
+    def scheduler_step(self, state: StrategyState, noise_pred: torch.Tensor,
+                       t: torch.Tensor) -> torch.Tensor:
+        progress_bar = state.extra["progress_bar"]
+        step_idx = state.extra["step_idx"]
+        num_warmup_steps = state.extra["num_warmup_steps"]
+
+        if step_idx == len(state.timesteps) - 1 or (
+            (step_idx + 1) > num_warmup_steps and
+            (step_idx + 1) % self.stage.scheduler.order == 0
+                and progress_bar is not None):
+            progress_bar.update()
+
+        return state.extra.get("pred_latents", state.latents)
+
+    def postprocess(self, state: StrategyState) -> ForwardBatch:
+        progress_bar = state.extra.get("progress_bar")
+        if progress_bar is not None:
+            progress_bar.close()
+
+        batch = state.extra["batch"]
+        latents = state.extra["pred_latents"].permute(0, 2, 1, 3, 4)
+        batch.latents = latents
+        return batch

fastvideo/pipelines/stages/denoising_engine.py

@@ -0,0 +1,177 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+Scaffolding for a unified denoising engine with hook support.
+"""
+
+from __future__ import annotations
+
+from functools import lru_cache
+from typing import Protocol, TYPE_CHECKING
+from collections.abc import Sequence
+
+import torch
+
+from fastvideo.models.schedulers.adapter import SchedulerAdapter
+from fastvideo.pipelines.stages.denoising_strategies import (
+    BlockDenoisingStrategy,
+    DenoisingStrategy,
+    StrategyState,
+)
+
+if TYPE_CHECKING:
+    from fastvideo.fastvideo_args import FastVideoArgs
+    from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+
+
+class EngineHook(Protocol):
+
+    def on_init(self, engine: DenoisingEngine, batch: ForwardBatch,
+                args: FastVideoArgs) -> None:
+        ...
+
+    def pre_run(self, state: StrategyState) -> None:
+        ...
+
+    def pre_step(self, state: StrategyState, step_idx: int,
+                 t: torch.Tensor) -> None:
+        ...
+
+    def post_step(self, state: StrategyState, step_idx: int,
+                  t: torch.Tensor) -> None:
+        ...
+
+    def post_run(self, state: StrategyState, batch: ForwardBatch) -> None:
+        ...
+
+
+class BaseEngineHook:
+
+    def on_init(self, engine: DenoisingEngine, batch: ForwardBatch,
+                args: FastVideoArgs) -> None:
+        return None
+
+    def pre_run(self, state: StrategyState) -> None:
+        return None
+
+    def pre_step(self, state: StrategyState, step_idx: int,
+                 t: torch.Tensor) -> None:
+        return None
+
+    def post_step(self, state: StrategyState, step_idx: int,
+                  t: torch.Tensor) -> None:
+        return None
+
+    def post_run(self, state: StrategyState, batch: ForwardBatch) -> None:
+        return None
+
+
+class GuidanceCache:
+
+    def __init__(self, maxsize: int = 8) -> None:
+        self._build = lru_cache(maxsize=maxsize)(self._build_impl)
+
+    def _build_impl(self, batch_size: int, dtype: torch.dtype,
+                    device: torch.device, guidance_val: float) -> torch.Tensor:
+        return (torch.full(
+            (batch_size, ),
+            guidance_val,
+            dtype=torch.float32,
+            device=device,
+        ).to(dtype) * 1000.0)
+
+    def get(self, batch_size: int, dtype: torch.dtype, device: torch.device,
+            guidance_val: float | None) -> torch.Tensor | None:
+        if guidance_val is None:
+            return None
+        return self._build(batch_size, dtype, device, guidance_val)
+
+
+class DenoisingEngine:
+
+    def __init__(
+        self,
+        strategy: DenoisingStrategy,
+        *,
+        scheduler_adapter: SchedulerAdapter | None = None,
+        hooks: Sequence[EngineHook] | None = None,
+    ) -> None:
+        self.strategy = strategy
+        self.scheduler_adapter = scheduler_adapter
+        self.hooks = list(hooks) if hooks is not None else []
+
+    def run(self, batch: ForwardBatch, args: FastVideoArgs) -> ForwardBatch:
+        for hook in self.hooks:
+            hook.on_init(self, batch, args)
+
+        state = self.strategy.prepare(batch, args)
+        if self.scheduler_adapter is not None:
+            state.extra.setdefault("scheduler_adapter", self.scheduler_adapter)
+        for hook in self.hooks:
+            hook.pre_run(state)
+
+        if isinstance(self.strategy, BlockDenoisingStrategy):
+            self.run_blocks(state)
+        else:
+            timesteps = state.timesteps
+            for i, t in enumerate(timesteps):
+                for hook in self.hooks:
+                    hook.pre_step(state, i, t)
+
+                model_inputs = self.strategy.make_model_inputs(state, t, i)
+                noise_pred = self.strategy.forward(state, model_inputs)
+                noise_pred = self.strategy.cfg_combine(state, noise_pred)
+                state.latents = self.strategy.scheduler_step(
+                    state,
+                    noise_pred,
+                    t,
+                )
+
+                for hook in self.hooks:
+                    hook.post_step(state, i, t)
+
+        for hook in self.hooks:
+            hook.post_run(state, batch)
+
+        return self.strategy.postprocess(state)
+
+    def run_blocks(
+        self,
+        state: StrategyState,
+        *,
+        block_plan=None,
+        start_block: int = 0,
+        num_blocks: int | None = None,
+    ) -> None:
+        if not isinstance(self.strategy, BlockDenoisingStrategy):
+            raise TypeError("run_blocks requires a BlockDenoisingStrategy")
+
+        strategy = self.strategy
+        if block_plan is None:
+            block_plan = strategy.block_plan(state)
+
+        items = block_plan.items
+        end_block = len(items)
+        if num_blocks is not None:
+            end_block = min(end_block, start_block + num_blocks)
+
+        for block_idx in range(start_block, end_block):
+            block_item = items[block_idx]
+            t_hook = self._block_hook_t(state, block_idx)
+            for hook in self.hooks:
+                hook.pre_step(state, block_idx, t_hook)
+            block_ctx = strategy.init_block_context(
+                state,
+                block_item,
+                block_idx,
+            )
+            strategy.process_block(state, block_ctx, block_item)
+            strategy.update_context(state, block_ctx, block_item)
+            for hook in self.hooks:
+                hook.post_step(state, block_idx, t_hook)
+
+    def _block_hook_t(self, state: StrategyState,
+                      block_idx: int) -> torch.Tensor:
+        timesteps = state.timesteps
+        if timesteps.numel() > 0:
+            return timesteps[0]
+        return torch.tensor(block_idx, device=state.latents.device)

fastvideo/pipelines/stages/denoising_engine_hooks.py

@@ -0,0 +1,77 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+Engine hook implementations for denoising runs.
+"""
+
+from __future__ import annotations
+
+import time
+
+from fastvideo.logger import init_logger
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+from fastvideo.pipelines.stages.denoising_engine import DenoisingEngine
+from fastvideo.pipelines.stages.denoising_strategies import StrategyState
+
+logger = init_logger(__name__)
+
+
+class PerfLoggingHook:
+    """
+    Record per-step and total denoising times into batch logging info.
+    """
+
+    def __init__(self, stage_name: str = "DenoisingEngine") -> None:
+        self.stage_name = stage_name
+        self._batch: ForwardBatch | None = None
+        self._run_start = 0.0
+        self._step_starts: dict[int, float] = {}
+        self._step_times_ms: list[float] = []
+
+    def on_init(self, engine: DenoisingEngine, batch: ForwardBatch, args):
+        self._batch = batch
+
+    def pre_run(self, state: StrategyState) -> None:
+        self._run_start = time.perf_counter()
+        self._step_starts.clear()
+        self._step_times_ms.clear()
+
+    def pre_step(self, state: StrategyState, step_idx: int, t) -> None:
+        self._step_starts[step_idx] = time.perf_counter()
+
+    def post_step(self, state: StrategyState, step_idx: int, t) -> None:
+        start = self._step_starts.pop(step_idx, None)
+        if start is None:
+            return
+        self._step_times_ms.append((time.perf_counter() - start) * 1000.0)
+
+    def post_run(self, state: StrategyState, batch: ForwardBatch) -> None:
+        total_ms = (time.perf_counter() - self._run_start) * 1000.0
+        target_batch = self._batch or batch
+        if target_batch is None:
+            return
+        target_batch.logging_info.add_stage_metric(self.stage_name,
+                                                   "denoise_step_times_ms",
+                                                   list(self._step_times_ms))
+        target_batch.logging_info.add_stage_metric(self.stage_name,
+                                                   "denoise_total_ms", total_ms)
+        if not self._step_times_ms:
+            return
+        if not self._is_primary_rank():
+            return
+        mean_ms = sum(self._step_times_ms) / len(self._step_times_ms)
+        logger.info(
+            "[%s] denoise steps=%d total_ms=%.2f mean_step_ms=%.2f min=%.2f max=%.2f",
+            self.stage_name,
+            len(self._step_times_ms),
+            total_ms,
+            mean_ms,
+            min(self._step_times_ms),
+            max(self._step_times_ms),
+        )
+
+    def _is_primary_rank(self) -> bool:
+        try:
+            from fastvideo.distributed import get_world_group
+            return get_world_group().local_rank == 0
+        except Exception:
+            return True

fastvideo/pipelines/stages/denoising_longcat_strategy.py

@@ -0,0 +1,551 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+LongCat denoising strategies (base, I2V, VC).
+"""
+
+from __future__ import annotations
+
+import time
+from typing import Any
+
+import torch
+from tqdm import tqdm
+
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.forward_context import set_forward_context
+from fastvideo.logger import init_logger
+from fastvideo.models.loader.component_loader import TransformerLoader
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+from fastvideo.pipelines.stages.denoising_strategies import (
+    DenoisingStrategy,
+    ModelInputs,
+    StrategyState,
+)
+
+logger = init_logger(__name__)
+
+
+class _BaseLongCatStrategy(DenoisingStrategy):
+
+    def __init__(self, stage: Any) -> None:
+        self.stage = stage
+
+    def _load_transformer(self, batch: ForwardBatch,
+                          fastvideo_args: FastVideoArgs,
+                          attach_pipeline: bool) -> None:
+        if fastvideo_args.model_loaded["transformer"]:
+            return
+        loader = TransformerLoader()
+        self.stage.transformer = loader.load(
+            fastvideo_args.model_paths["transformer"], fastvideo_args)
+        if attach_pipeline:
+            pipeline = self.stage.pipeline() if self.stage.pipeline else None
+            if pipeline:
+                pipeline.add_module("transformer", self.stage.transformer)
+        fastvideo_args.model_loaded["transformer"] = True
+
+    def _build_prompt_inputs(self, batch: ForwardBatch):
+        prompt_embeds = batch.prompt_embeds[0]
+        prompt_attention_mask = (batch.prompt_attention_mask[0]
+                                 if batch.prompt_attention_mask else None)
+        do_cfg = batch.do_classifier_free_guidance
+
+        if do_cfg:
+            negative_prompt_embeds = batch.negative_prompt_embeds[0]
+            negative_prompt_attention_mask = (batch.negative_attention_mask[0]
+                                              if batch.negative_attention_mask
+                                              else None)
+            prompt_embeds_combined = torch.cat(
+                [negative_prompt_embeds, prompt_embeds], dim=0)
+            if prompt_attention_mask is not None:
+                prompt_attention_mask_combined = torch.cat(
+                    [negative_prompt_attention_mask, prompt_attention_mask],
+                    dim=0)
+            else:
+                prompt_attention_mask_combined = None
+        else:
+            prompt_embeds_combined = prompt_embeds
+            prompt_attention_mask_combined = prompt_attention_mask
+
+        return prompt_embeds, prompt_attention_mask, prompt_embeds_combined, \
+            prompt_attention_mask_combined
+
+    def optimized_scale(self, positive_flat: torch.Tensor,
+                        negative_flat: torch.Tensor) -> torch.Tensor:
+        """
+        Calculate optimized scale from CFG-zero paper.
+
+        st_star = (v_cond^T * v_uncond) / ||v_uncond||^2
+        """
+        dot_product = torch.sum(positive_flat * negative_flat,
+                                dim=1,
+                                keepdim=True)
+        squared_norm = torch.sum(negative_flat**2, dim=1, keepdim=True) + 1e-8
+        return dot_product / squared_norm
+
+    def cfg_combine(self, state: StrategyState,
+                    noise_pred: torch.Tensor) -> torch.Tensor:
+        return noise_pred
+
+
+class LongCatStrategy(_BaseLongCatStrategy):
+
+    def prepare(self, batch: ForwardBatch,
+                fastvideo_args: FastVideoArgs) -> StrategyState:
+        self._load_transformer(batch, fastvideo_args, attach_pipeline=True)
+
+        if hasattr(self.stage.transformer, "module"):
+            transformer_dtype = next(
+                self.stage.transformer.module.parameters()).dtype
+        else:
+            transformer_dtype = next(self.stage.transformer.parameters()).dtype
+
+        target_dtype = transformer_dtype
+        autocast_enabled = (target_dtype != torch.float32
+                            ) and not fastvideo_args.disable_autocast
+
+        prompt_embeds, prompt_attention_mask, prompt_embeds_combined, \
+            prompt_attention_mask_combined = self._build_prompt_inputs(batch)
+
+        timesteps = batch.timesteps
+        num_inference_steps = len(timesteps)
+        progress_bar = tqdm(total=num_inference_steps, desc="LongCat Denoising")
+
+        extra: dict[str, Any] = {
+            "batch": batch,
+            "target_dtype": target_dtype,
+            "autocast_enabled": autocast_enabled,
+            "prompt_embeds_combined": prompt_embeds_combined,
+            "prompt_attention_mask_combined": prompt_attention_mask_combined,
+            "progress_bar": progress_bar,
+        }
+
+        return StrategyState(
+            latents=batch.latents,
+            timesteps=timesteps,
+            num_inference_steps=num_inference_steps,
+            prompt_embeds=[prompt_embeds],
+            negative_prompt_embeds=batch.negative_prompt_embeds,
+            prompt_attention_mask=[prompt_attention_mask]
+            if prompt_attention_mask is not None else None,
+            negative_attention_mask=batch.negative_attention_mask,
+            image_embeds=batch.image_embeds,
+            guidance_scale=batch.guidance_scale,
+            guidance_scale_2=batch.guidance_scale_2,
+            guidance_rescale=batch.guidance_rescale,
+            do_cfg=batch.do_classifier_free_guidance,
+            extra=extra,
+        )
+
+    def make_model_inputs(self, state: StrategyState, t: torch.Tensor,
+                          step_idx: int) -> ModelInputs:
+        target_dtype = state.extra["target_dtype"]
+        latents = state.latents
+
+        if state.do_cfg:
+            latent_model_input = torch.cat([latents] * 2)
+        else:
+            latent_model_input = latents
+
+        latent_model_input = latent_model_input.to(target_dtype)
+        timestep = t.expand(latent_model_input.shape[0]).to(target_dtype)
+
+        state.extra["step_idx"] = step_idx
+        return ModelInputs(
+            latent_model_input=latent_model_input,
+            timestep=timestep,
+            prompt_embeds=state.prompt_embeds,
+            prompt_attention_mask=state.prompt_attention_mask,
+        )
+
+    def forward(self, state: StrategyState,
+                model_inputs: ModelInputs) -> torch.Tensor:
+        batch = state.extra["batch"]
+        target_dtype = state.extra["target_dtype"]
+        autocast_enabled = state.extra["autocast_enabled"]
+        prompt_embeds_combined = state.extra["prompt_embeds_combined"]
+        prompt_attention_mask_combined = state.extra[
+            "prompt_attention_mask_combined"]
+        step_idx = state.extra["step_idx"]
+
+        batch.is_cfg_negative = False
+        with set_forward_context(
+                current_timestep=step_idx,
+                attn_metadata=None,
+                forward_batch=batch,
+        ), torch.autocast(device_type="cuda",
+                          dtype=target_dtype,
+                          enabled=autocast_enabled):
+            noise_pred = self.stage.transformer(
+                hidden_states=model_inputs.latent_model_input,
+                encoder_hidden_states=prompt_embeds_combined,
+                timestep=model_inputs.timestep,
+                encoder_attention_mask=prompt_attention_mask_combined,
+            )
+
+        if state.do_cfg:
+            noise_pred_uncond, noise_pred_cond = noise_pred.chunk(2)
+            B = noise_pred_cond.shape[0]
+            positive = noise_pred_cond.reshape(B, -1)
+            negative = noise_pred_uncond.reshape(B, -1)
+            st_star = self.optimized_scale(positive, negative)
+            st_star = st_star.view(B, 1, 1, 1, 1)
+            noise_pred = (noise_pred_uncond * st_star + state.guidance_scale *
+                          (noise_pred_cond - noise_pred_uncond * st_star))
+
+        noise_pred = -noise_pred
+        return noise_pred
+
+    def scheduler_step(self, state: StrategyState, noise_pred: torch.Tensor,
+                       t: torch.Tensor) -> torch.Tensor:
+        latents = self.stage.scheduler.step(noise_pred,
+                                            t,
+                                            state.latents,
+                                            return_dict=False)[0]
+
+        progress_bar = state.extra["progress_bar"]
+        if progress_bar is not None:
+            progress_bar.update()
+        return latents
+
+    def postprocess(self, state: StrategyState) -> ForwardBatch:
+        progress_bar = state.extra.get("progress_bar")
+        if progress_bar is not None:
+            progress_bar.close()
+        batch = state.extra["batch"]
+        batch.latents = state.latents
+        return batch
+
+
+class LongCatI2VStrategy(_BaseLongCatStrategy):
+
+    def prepare(self, batch: ForwardBatch,
+                fastvideo_args: FastVideoArgs) -> StrategyState:
+        self._load_transformer(batch, fastvideo_args, attach_pipeline=False)
+
+        target_dtype = torch.bfloat16
+        autocast_enabled = (target_dtype != torch.float32
+                            ) and not fastvideo_args.disable_autocast
+
+        prompt_embeds, prompt_attention_mask, prompt_embeds_combined, \
+            prompt_attention_mask_combined = self._build_prompt_inputs(batch)
+
+        timesteps = batch.timesteps
+        num_inference_steps = len(timesteps)
+        progress_bar = tqdm(total=num_inference_steps, desc="I2V Denoising")
+
+        num_cond_latents = getattr(batch, "num_cond_latents", 0)
+        if num_cond_latents > 0:
+            logger.info("I2V Denoising: num_cond_latents=%s, latent_shape=%s",
+                        num_cond_latents, batch.latents.shape)
+
+        extra: dict[str, Any] = {
+            "batch": batch,
+            "target_dtype": target_dtype,
+            "autocast_enabled": autocast_enabled,
+            "prompt_embeds_combined": prompt_embeds_combined,
+            "prompt_attention_mask_combined": prompt_attention_mask_combined,
+            "num_cond_latents": num_cond_latents,
+            "progress_bar": progress_bar,
+        }
+
+        return StrategyState(
+            latents=batch.latents,
+            timesteps=timesteps,
+            num_inference_steps=num_inference_steps,
+            prompt_embeds=[prompt_embeds],
+            negative_prompt_embeds=batch.negative_prompt_embeds,
+            prompt_attention_mask=[prompt_attention_mask]
+            if prompt_attention_mask is not None else None,
+            negative_attention_mask=batch.negative_attention_mask,
+            image_embeds=batch.image_embeds,
+            guidance_scale=batch.guidance_scale,
+            guidance_scale_2=batch.guidance_scale_2,
+            guidance_rescale=batch.guidance_rescale,
+            do_cfg=batch.do_classifier_free_guidance,
+            extra=extra,
+        )
+
+    def make_model_inputs(self, state: StrategyState, t: torch.Tensor,
+                          step_idx: int) -> ModelInputs:
+        target_dtype = state.extra["target_dtype"]
+        num_cond_latents = state.extra["num_cond_latents"]
+        latents = state.latents
+
+        if state.do_cfg:
+            latent_model_input = torch.cat([latents] * 2)
+        else:
+            latent_model_input = latents
+
+        latent_model_input = latent_model_input.to(target_dtype)
+        timestep = t.expand(latent_model_input.shape[0]).to(target_dtype)
+        timestep = timestep.unsqueeze(-1).repeat(1, latent_model_input.shape[2])
+        if num_cond_latents > 0:
+            timestep[:, :num_cond_latents] = 0
+
+        state.extra["step_idx"] = step_idx
+        return ModelInputs(
+            latent_model_input=latent_model_input,
+            timestep=timestep,
+            prompt_embeds=state.prompt_embeds,
+            prompt_attention_mask=state.prompt_attention_mask,
+            extra_kwargs={"num_cond_latents": num_cond_latents},
+        )
+
+    def forward(self, state: StrategyState,
+                model_inputs: ModelInputs) -> torch.Tensor:
+        batch = state.extra["batch"]
+        target_dtype = state.extra["target_dtype"]
+        autocast_enabled = state.extra["autocast_enabled"]
+        prompt_embeds_combined = state.extra["prompt_embeds_combined"]
+        prompt_attention_mask_combined = state.extra[
+            "prompt_attention_mask_combined"]
+        step_idx = state.extra["step_idx"]
+        num_cond_latents = state.extra["num_cond_latents"]
+
+        batch.is_cfg_negative = False
+        with set_forward_context(
+                current_timestep=step_idx,
+                attn_metadata=None,
+                forward_batch=batch,
+        ), torch.autocast(device_type="cuda",
+                          dtype=target_dtype,
+                          enabled=autocast_enabled):
+            noise_pred = self.stage.transformer(
+                hidden_states=model_inputs.latent_model_input,
+                encoder_hidden_states=prompt_embeds_combined,
+                timestep=model_inputs.timestep,
+                encoder_attention_mask=prompt_attention_mask_combined,
+                num_cond_latents=num_cond_latents,
+            )
+
+        if state.do_cfg:
+            noise_pred_uncond, noise_pred_cond = noise_pred.chunk(2)
+            B = noise_pred_cond.shape[0]
+            positive = noise_pred_cond.reshape(B, -1)
+            negative = noise_pred_uncond.reshape(B, -1)
+            st_star = self.optimized_scale(positive, negative)
+            st_star = st_star.view(B, 1, 1, 1, 1)
+            noise_pred = (noise_pred_uncond * st_star + state.guidance_scale *
+                          (noise_pred_cond - noise_pred_uncond * st_star))
+
+        noise_pred = -noise_pred
+        return noise_pred
+
+    def scheduler_step(self, state: StrategyState, noise_pred: torch.Tensor,
+                       t: torch.Tensor) -> torch.Tensor:
+        num_cond_latents = state.extra["num_cond_latents"]
+        latents = state.latents
+        if num_cond_latents > 0:
+            latents[:, :, num_cond_latents:] = self.stage.scheduler.step(
+                noise_pred[:, :, num_cond_latents:],
+                t,
+                latents[:, :, num_cond_latents:],
+                return_dict=False)[0]
+        else:
+            latents = self.stage.scheduler.step(noise_pred,
+                                                t,
+                                                latents,
+                                                return_dict=False)[0]
+
+        progress_bar = state.extra["progress_bar"]
+        if progress_bar is not None:
+            progress_bar.update()
+        return latents
+
+    def postprocess(self, state: StrategyState) -> ForwardBatch:
+        progress_bar = state.extra.get("progress_bar")
+        if progress_bar is not None:
+            progress_bar.close()
+        batch = state.extra["batch"]
+        batch.latents = state.latents
+        return batch
+
+
+class LongCatVCStrategy(_BaseLongCatStrategy):
+
+    def prepare(self, batch: ForwardBatch,
+                fastvideo_args: FastVideoArgs) -> StrategyState:
+        self._load_transformer(batch, fastvideo_args, attach_pipeline=False)
+
+        target_dtype = torch.bfloat16
+        autocast_enabled = (target_dtype != torch.float32
+                            ) and not fastvideo_args.disable_autocast
+
+        prompt_embeds, prompt_attention_mask, prompt_embeds_combined, \
+            prompt_attention_mask_combined = self._build_prompt_inputs(batch)
+
+        timesteps = batch.timesteps
+        num_inference_steps = len(timesteps)
+        progress_bar = tqdm(total=num_inference_steps, desc="VC Denoising")
+
+        num_cond_latents = getattr(batch, "num_cond_latents", 0)
+        use_kv_cache = getattr(batch, "use_kv_cache", False)
+        kv_cache_dict = getattr(batch, "kv_cache_dict", {})
+
+        logger.info(
+            "VC Denoising: num_cond_latents=%d, use_kv_cache=%s, latent_shape=%s",
+            num_cond_latents, use_kv_cache, batch.latents.shape)
+
+        extra: dict[str, Any] = {
+            "batch": batch,
+            "target_dtype": target_dtype,
+            "autocast_enabled": autocast_enabled,
+            "prompt_embeds_combined": prompt_embeds_combined,
+            "prompt_attention_mask_combined": prompt_attention_mask_combined,
+            "num_cond_latents": num_cond_latents,
+            "use_kv_cache": use_kv_cache,
+            "kv_cache_dict": kv_cache_dict,
+            "progress_bar": progress_bar,
+            "step_times": [],
+        }
+
+        return StrategyState(
+            latents=batch.latents,
+            timesteps=timesteps,
+            num_inference_steps=num_inference_steps,
+            prompt_embeds=[prompt_embeds],
+            negative_prompt_embeds=batch.negative_prompt_embeds,
+            prompt_attention_mask=[prompt_attention_mask]
+            if prompt_attention_mask is not None else None,
+            negative_attention_mask=batch.negative_attention_mask,
+            image_embeds=batch.image_embeds,
+            guidance_scale=batch.guidance_scale,
+            guidance_scale_2=batch.guidance_scale_2,
+            guidance_rescale=batch.guidance_rescale,
+            do_cfg=batch.do_classifier_free_guidance,
+            extra=extra,
+        )
+
+    def make_model_inputs(self, state: StrategyState, t: torch.Tensor,
+                          step_idx: int) -> ModelInputs:
+        target_dtype = state.extra["target_dtype"]
+        num_cond_latents = state.extra["num_cond_latents"]
+        use_kv_cache = state.extra["use_kv_cache"]
+        latents = state.latents
+
+        if state.do_cfg:
+            latent_model_input = torch.cat([latents] * 2)
+        else:
+            latent_model_input = latents
+
+        latent_model_input = latent_model_input.to(target_dtype)
+        timestep = t.expand(latent_model_input.shape[0]).to(target_dtype)
+        timestep = timestep.unsqueeze(-1).repeat(1, latent_model_input.shape[2])
+        if not use_kv_cache and num_cond_latents > 0:
+            timestep[:, :num_cond_latents] = 0
+
+        state.extra["step_idx"] = step_idx
+        state.extra["step_start"] = time.time()
+
+        extra_kwargs = {"num_cond_latents": num_cond_latents}
+        if use_kv_cache:
+            extra_kwargs["kv_cache_dict"] = state.extra["kv_cache_dict"]
+
+        return ModelInputs(
+            latent_model_input=latent_model_input,
+            timestep=timestep,
+            prompt_embeds=state.prompt_embeds,
+            prompt_attention_mask=state.prompt_attention_mask,
+            extra_kwargs=extra_kwargs,
+        )
+
+    def forward(self, state: StrategyState,
+                model_inputs: ModelInputs) -> torch.Tensor:
+        batch = state.extra["batch"]
+        target_dtype = state.extra["target_dtype"]
+        autocast_enabled = state.extra["autocast_enabled"]
+        prompt_embeds_combined = state.extra["prompt_embeds_combined"]
+        prompt_attention_mask_combined = state.extra[
+            "prompt_attention_mask_combined"]
+        step_idx = state.extra["step_idx"]
+
+        batch.is_cfg_negative = False
+        with set_forward_context(
+                current_timestep=step_idx,
+                attn_metadata=None,
+                forward_batch=batch,
+        ), torch.autocast(device_type="cuda",
+                          dtype=target_dtype,
+                          enabled=autocast_enabled):
+            noise_pred = self.stage.transformer(
+                hidden_states=model_inputs.latent_model_input,
+                encoder_hidden_states=prompt_embeds_combined,
+                timestep=model_inputs.timestep,
+                encoder_attention_mask=prompt_attention_mask_combined,
+                **model_inputs.extra_kwargs,
+            )
+
+        if state.do_cfg:
+            noise_pred_uncond, noise_pred_cond = noise_pred.chunk(2)
+            B = noise_pred_cond.shape[0]
+            positive = noise_pred_cond.reshape(B, -1)
+            negative = noise_pred_uncond.reshape(B, -1)
+            st_star = self.optimized_scale(positive, negative)
+            st_star = st_star.view(B, 1, 1, 1, 1)
+            noise_pred = (noise_pred_uncond * st_star + state.guidance_scale *
+                          (noise_pred_cond - noise_pred_uncond * st_star))
+
+        noise_pred = -noise_pred
+        return noise_pred
+
+    def scheduler_step(self, state: StrategyState, noise_pred: torch.Tensor,
+                       t: torch.Tensor) -> torch.Tensor:
+        num_cond_latents = state.extra["num_cond_latents"]
+        use_kv_cache = state.extra["use_kv_cache"]
+        latents = state.latents
+
+        if use_kv_cache:
+            latents = self.stage.scheduler.step(noise_pred,
+                                                t,
+                                                latents,
+                                                return_dict=False)[0]
+        else:
+            if num_cond_latents > 0:
+                latents[:, :, num_cond_latents:] = self.stage.scheduler.step(
+                    noise_pred[:, :, num_cond_latents:],
+                    t,
+                    latents[:, :, num_cond_latents:],
+                    return_dict=False,
+                )[0]
+            else:
+                latents = self.stage.scheduler.step(noise_pred,
+                                                    t,
+                                                    latents,
+                                                    return_dict=False)[0]
+
+        step_time = time.time() - state.extra["step_start"]
+        state.extra["step_times"].append(step_time)
+        if state.extra["step_idx"] < 3:
+            logger.info("Step %d: %.2fs", state.extra["step_idx"], step_time)
+
+        progress_bar = state.extra["progress_bar"]
+        if progress_bar is not None:
+            progress_bar.update()
+
+        return latents
+
+    def postprocess(self, state: StrategyState) -> ForwardBatch:
+        progress_bar = state.extra.get("progress_bar")
+        if progress_bar is not None:
+            progress_bar.close()
+
+        batch = state.extra["batch"]
+        use_kv_cache = state.extra["use_kv_cache"]
+        step_times = state.extra["step_times"]
+
+        latents = state.latents
+        if use_kv_cache and hasattr(
+                batch, "cond_latents") and batch.cond_latents is not None:
+            latents = torch.cat([batch.cond_latents, latents], dim=2)
+            logger.info(
+                "Concatenated conditioning latents back, final shape: %s",
+                latents.shape)
+
+        if step_times:
+            avg_time = sum(step_times) / len(step_times)
+            logger.info("Average step time: %.2fs (total: %.1fs)", avg_time,
+                        sum(step_times))
+
+        batch.latents = latents
+        return batch

fastvideo/pipelines/stages/denoising_matrixgame_strategy.py

@@ -0,0 +1,324 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+MatrixGame causal block denoising strategy.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+
+from fastvideo.distributed import get_local_torch_device
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+from fastvideo.pipelines.stages.denoising_strategies import (
+    BlockContext,
+    BlockDenoisingStrategy,
+    BlockPlan,
+    BlockPlanItem,
+    ModelInputs,
+    StrategyState,
+)
+from fastvideo.pipelines.stages.matrixgame_denoising import BlockProcessingContext
+
+try:
+    from fastvideo.attention.backends.sliding_tile_attn import (
+        SlidingTileAttentionBackend)
+    st_attn_available = True
+except ImportError:
+    st_attn_available = False
+    SlidingTileAttentionBackend = None  # type: ignore
+
+
+class MatrixGameBlockStrategy(BlockDenoisingStrategy):
+
+    def __init__(self, stage: Any) -> None:
+        self.stage = stage
+
+    def prepare(self, batch: ForwardBatch,
+                fastvideo_args: FastVideoArgs) -> StrategyState:
+        target_dtype = torch.bfloat16
+        autocast_enabled = (target_dtype != torch.float32
+                            ) and not fastvideo_args.disable_autocast
+
+        latents = batch.latents
+        if latents is None:
+            raise ValueError("latents must be provided")
+
+        latent_seq_length = latents.shape[-1] * latents.shape[-2]
+        patch_size = self.stage.transformer.patch_size
+        patch_ratio = patch_size[-1] * patch_size[-2]
+        self.stage.frame_seq_length = latent_seq_length // patch_ratio
+
+        timesteps = torch.tensor(
+            fastvideo_args.pipeline_config.dmd_denoising_steps,
+            dtype=torch.long).cpu()
+        if getattr(fastvideo_args.pipeline_config, "warp_denoising_step",
+                   False):
+            scheduler_timesteps = torch.cat((
+                self.stage.scheduler.timesteps.cpu(),
+                torch.tensor([0], dtype=torch.float32),
+            ))
+            timesteps = scheduler_timesteps[1000 - timesteps]
+        timesteps = timesteps.to(get_local_torch_device())
+
+        boundary_ratio = getattr(fastvideo_args.pipeline_config.dit_config,
+                                 "boundary_ratio", None)
+        if boundary_ratio is not None:
+            boundary_timestep = (boundary_ratio *
+                                 self.stage.scheduler.num_train_timesteps)
+            high_noise_timesteps = timesteps[timesteps >= boundary_timestep]
+        else:
+            boundary_timestep = None
+            high_noise_timesteps = None
+
+        image_embeds = batch.image_embeds
+        if len(image_embeds) > 0:
+            assert torch.isnan(image_embeds[0]).sum() == 0
+            image_embeds = [
+                image_embed.to(target_dtype) for image_embed in image_embeds
+            ]
+
+        image_kwargs = {"encoder_hidden_states_image": image_embeds}
+        pos_cond_kwargs: dict[str, Any] = {}
+
+        if (st_attn_available
+                and self.stage.attn_backend == SlidingTileAttentionBackend):
+            self.stage.prepare_sta_param(batch, fastvideo_args)
+
+        prompt_embeds = batch.prompt_embeds
+        assert torch.isnan(prompt_embeds[0]).sum() == 0
+
+        kv_cache1 = self.stage._initialize_kv_cache(batch_size=latents.shape[0],
+                                                    dtype=target_dtype,
+                                                    device=latents.device)
+        kv_cache2 = None
+        if boundary_timestep is not None:
+            kv_cache2 = self.stage._initialize_kv_cache(
+                batch_size=latents.shape[0],
+                dtype=target_dtype,
+                device=latents.device,
+            )
+
+        kv_cache_mouse = None
+        kv_cache_keyboard = None
+        if self.stage.use_action_module:
+            kv_cache_mouse, kv_cache_keyboard = (
+                self.stage._initialize_action_kv_cache(
+                    batch_size=latents.shape[0],
+                    dtype=target_dtype,
+                    device=latents.device,
+                ))
+
+        crossattn_cache = self.stage._initialize_crossattn_cache(
+            batch_size=latents.shape[0],
+            max_text_len=257,
+            dtype=target_dtype,
+            device=latents.device,
+        )
+
+        num_frames = latents.shape[2]
+        if num_frames % self.stage.num_frame_per_block != 0:
+            raise ValueError(
+                "num_frames must be divisible by num_frame_per_block for "
+                "causal denoising")
+        num_blocks = num_frames // self.stage.num_frame_per_block
+        block_sizes = [self.stage.num_frame_per_block] * num_blocks
+        start_index = 0
+
+        if boundary_timestep is not None:
+            block_sizes[0] = 1
+
+        ctx = BlockProcessingContext(
+            batch=batch,
+            block_idx=0,
+            start_index=0,
+            kv_cache1=kv_cache1,
+            kv_cache2=kv_cache2,
+            kv_cache_mouse=kv_cache_mouse,
+            kv_cache_keyboard=kv_cache_keyboard,
+            crossattn_cache=crossattn_cache,
+            timesteps=timesteps,
+            block_sizes=block_sizes,
+            noise_pool=None,
+            fastvideo_args=fastvideo_args,
+            target_dtype=target_dtype,
+            autocast_enabled=autocast_enabled,
+            boundary_timestep=boundary_timestep,
+            high_noise_timesteps=high_noise_timesteps,
+            context_noise=getattr(fastvideo_args.pipeline_config,
+                                  "context_noise", 0),
+            image_kwargs=image_kwargs,
+            pos_cond_kwargs=pos_cond_kwargs,
+        )
+
+        progress_bar = self.stage.progress_bar(total=len(block_sizes) *
+                                               len(timesteps))
+
+        extra: dict[str, Any] = {
+            "batch": batch,
+            "fastvideo_args": fastvideo_args,
+            "ctx": ctx,
+            "block_sizes": block_sizes,
+            "start_index": start_index,
+            "progress_bar": progress_bar,
+            "boundary_timestep": boundary_timestep,
+        }
+
+        return StrategyState(
+            latents=latents,
+            timesteps=timesteps,
+            num_inference_steps=len(timesteps),
+            prompt_embeds=batch.prompt_embeds,
+            negative_prompt_embeds=batch.negative_prompt_embeds,
+            prompt_attention_mask=batch.prompt_attention_mask,
+            negative_attention_mask=batch.negative_attention_mask,
+            image_embeds=image_embeds,
+            guidance_scale=batch.guidance_scale,
+            guidance_scale_2=batch.guidance_scale_2,
+            guidance_rescale=batch.guidance_rescale,
+            do_cfg=batch.do_classifier_free_guidance,
+            extra=extra,
+        )
+
+    def block_plan(self, state: StrategyState) -> BlockPlan:
+        block_sizes = state.extra["block_sizes"]
+        start_index = state.extra["start_index"]
+        items: list[BlockPlanItem] = []
+        for block_size in block_sizes:
+            items.append(
+                BlockPlanItem(
+                    start_index=start_index,
+                    num_frames=block_size,
+                    use_kv_cache=True,
+                    model_selector="default",
+                ))
+            start_index += block_size
+        return BlockPlan(items=items)
+
+    def init_block_context(self, state: StrategyState,
+                           block_item: BlockPlanItem,
+                           block_idx: int) -> BlockContext:
+        ctx = state.extra["ctx"]
+        ctx.block_idx = block_idx
+        ctx.start_index = block_item.start_index
+
+        action_kwargs = self.stage._prepare_action_kwargs(
+            state.extra["batch"],
+            block_item.start_index,
+            block_item.num_frames,
+        )
+
+        return BlockContext(
+            kv_cache=ctx.kv_cache1,
+            kv_cache_2=ctx.kv_cache2,
+            crossattn_cache=ctx.crossattn_cache,
+            action_cache=None,
+            extra={
+                "ctx": ctx,
+                "action_kwargs": action_kwargs,
+                "start_index": block_item.start_index,
+                "num_frames": block_item.num_frames,
+            },
+        )
+
+    def process_block(self, state: StrategyState, block_ctx: BlockContext,
+                      block_item: BlockPlanItem) -> None:
+        ctx = block_ctx.extra["ctx"]
+        batch = state.extra["batch"]
+        progress_bar = state.extra["progress_bar"]
+
+        start_index = block_item.start_index
+        current_num_frames = block_item.num_frames
+        action_kwargs = block_ctx.extra["action_kwargs"]
+
+        current_latents = state.latents[:, :, start_index:start_index +
+                                        current_num_frames, :, :]
+
+        noise_generator = None
+        if ctx.noise_pool is not None:
+            latents_device = state.latents.device
+
+            def noise_generator(shape: tuple, dtype: torch.dtype,
+                                step_idx: int) -> torch.Tensor:
+                if step_idx < len(ctx.noise_pool):
+                    noise = ctx.noise_pool[step_idx]
+                    if noise.shape != shape:
+                        noise = noise[:, :shape[1], :, :, :]
+                    return noise.to(device=latents_device, dtype=dtype)
+
+                generator = batch.generator
+                if isinstance(generator, list):
+                    generator = generator[0] if generator else None
+                return torch.randn(shape, dtype=dtype,
+                                   generator=generator).to(latents_device)
+
+        current_latents = self.stage._process_single_block(
+            current_latents=current_latents,
+            batch=batch,
+            start_index=start_index,
+            current_num_frames=current_num_frames,
+            timesteps=state.timesteps,
+            ctx=ctx,
+            action_kwargs=action_kwargs,
+            progress_bar=progress_bar,
+            noise_generator=noise_generator,
+        )
+
+        state.latents[:, :, start_index:start_index +
+                      current_num_frames, :, :] = (current_latents)
+
+    def update_context(self, state: StrategyState, block_ctx: BlockContext,
+                       block_item: BlockPlanItem) -> None:
+        ctx = block_ctx.extra["ctx"]
+        batch = state.extra["batch"]
+        action_kwargs = block_ctx.extra["action_kwargs"]
+        start_index = block_item.start_index
+        current_num_frames = block_item.num_frames
+
+        current_latents = state.latents[:, :, start_index:start_index +
+                                        current_num_frames, :, :]
+
+        self.stage._update_context_cache(
+            current_latents=current_latents,
+            batch=batch,
+            start_index=start_index,
+            current_num_frames=current_num_frames,
+            ctx=ctx,
+            action_kwargs=action_kwargs,
+            context_noise=ctx.context_noise,
+        )
+
+    def postprocess(self, state: StrategyState) -> ForwardBatch:
+        progress_bar = state.extra.get("progress_bar")
+        if progress_bar is not None:
+            progress_bar.close()
+
+        batch = state.extra["batch"]
+        boundary_timestep = state.extra["boundary_timestep"]
+
+        latents = state.latents
+        if boundary_timestep is not None:
+            num_frames_to_remove = self.stage.num_frame_per_block - 1
+            if num_frames_to_remove > 0:
+                latents = latents[:, :, :-num_frames_to_remove, :, :]
+
+        batch.latents = latents
+        return batch
+
+    def make_model_inputs(self, state: StrategyState, t: torch.Tensor,
+                          step_idx: int) -> ModelInputs:
+        raise NotImplementedError
+
+    def forward(self, state: StrategyState,
+                model_inputs: ModelInputs) -> torch.Tensor:
+        raise NotImplementedError
+
+    def cfg_combine(self, state: StrategyState,
+                    noise_pred: torch.Tensor) -> torch.Tensor:
+        raise NotImplementedError
+
+    def scheduler_step(self, state: StrategyState, noise_pred: torch.Tensor,
+                       t: torch.Tensor) -> torch.Tensor:
+        raise NotImplementedError

fastvideo/pipelines/stages/denoising_standard_strategy.py

@@ -0,0 +1,552 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+Standard denoising strategy backed by the legacy DenoisingStage utilities.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+
+from fastvideo.configs.pipelines.base import STA_Mode
+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.models.loader.component_loader import TransformerLoader
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+from fastvideo.pipelines.stages.denoising_strategies import (
+    DenoisingStrategy,
+    ModelInputs,
+    StrategyState,
+)
+from fastvideo.utils import dict_to_3d_list, masks_like
+
+try:
+    from fastvideo.attention.backends.sliding_tile_attn import (
+        SlidingTileAttentionBackend)
+    st_attn_available = True
+except ImportError:
+    st_attn_available = False
+    SlidingTileAttentionBackend = None  # type: ignore
+
+try:
+    from fastvideo.attention.backends.vmoba import VMOBAAttentionBackend
+    from fastvideo.utils import is_vmoba_available
+    vmoba_attn_available = is_vmoba_available()
+except ImportError:
+    vmoba_attn_available = False
+    VMOBAAttentionBackend = None  # type: ignore
+
+try:
+    from fastvideo.attention.backends.video_sparse_attn import (
+        VideoSparseAttentionBackend)
+    vsa_available = True
+except ImportError:
+    vsa_available = False
+    VideoSparseAttentionBackend = None  # type: ignore
+
+logger = init_logger(__name__)
+
+
+class StandardStrategy(DenoisingStrategy):
+
+    def __init__(self, stage: Any) -> None:
+        self.stage = stage
+
+    def prepare(self, batch: ForwardBatch,
+                fastvideo_args: FastVideoArgs) -> StrategyState:
+        pipeline = self.stage.pipeline() if self.stage.pipeline else None
+        if not fastvideo_args.model_loaded["transformer"]:
+            loader = TransformerLoader()
+            self.stage.transformer = loader.load(
+                fastvideo_args.model_paths["transformer"], fastvideo_args)
+            if pipeline:
+                pipeline.add_module("transformer", self.stage.transformer)
+            fastvideo_args.model_loaded["transformer"] = True
+
+        extra_step_kwargs = self.stage.prepare_extra_func_kwargs(
+            self.stage.scheduler.step,
+            {
+                "generator": batch.generator,
+                "eta": batch.eta
+            },
+        )
+
+        target_dtype = torch.bfloat16
+        autocast_enabled = (target_dtype != torch.float32
+                            ) and not fastvideo_args.disable_autocast
+
+        timesteps = batch.timesteps
+        if timesteps is None:
+            raise ValueError("Timesteps must be provided")
+        num_inference_steps = batch.num_inference_steps
+        num_warmup_steps = len(
+            timesteps) - num_inference_steps * self.stage.scheduler.order
+
+        image_embeds = batch.image_embeds
+        if len(image_embeds) > 0:
+            assert not torch.isnan(
+                image_embeds[0]).any(), "image_embeds contains nan"
+            image_embeds = [
+                image_embed.to(target_dtype) for image_embed in image_embeds
+            ]
+
+        image_kwargs = self.stage.prepare_extra_func_kwargs(
+            self.stage.transformer.forward,
+            {
+                "encoder_hidden_states_image": image_embeds,
+                "mask_strategy": dict_to_3d_list(
+                    None, t_max=50, l_max=60, h_max=24)
+            },
+        )
+
+        pos_cond_kwargs = self.stage.prepare_extra_func_kwargs(
+            self.stage.transformer.forward,
+            {
+                "encoder_hidden_states_2": batch.clip_embedding_pos,
+                "encoder_attention_mask": batch.prompt_attention_mask,
+            },
+        )
+
+        neg_cond_kwargs = self.stage.prepare_extra_func_kwargs(
+            self.stage.transformer.forward,
+            {
+                "encoder_hidden_states_2": batch.clip_embedding_neg,
+                "encoder_attention_mask": batch.negative_attention_mask,
+            },
+        )
+
+        action_kwargs = self.stage.prepare_extra_func_kwargs(
+            self.stage.transformer.forward,
+            {
+                "mouse_cond": batch.mouse_cond,
+                "keyboard_cond": batch.keyboard_cond,
+            },
+        )
+
+        if (st_attn_available
+                and self.stage.attn_backend == SlidingTileAttentionBackend):
+            self.stage.prepare_sta_param(batch, fastvideo_args)
+
+        latents = batch.latents
+        prompt_embeds = batch.prompt_embeds
+        assert not torch.isnan(
+            prompt_embeds[0]).any(), "prompt_embeds contains nan"
+
+        neg_prompt_embeds = None
+        if batch.do_classifier_free_guidance:
+            neg_prompt_embeds = batch.negative_prompt_embeds
+            assert neg_prompt_embeds is not None
+            assert not torch.isnan(
+                neg_prompt_embeds[0]).any(), "neg_prompt_embeds contains nan"
+
+        boundary_ratio = (
+            fastvideo_args.pipeline_config.dit_config.boundary_ratio)
+        if batch.boundary_ratio is not None:
+            logger.info("Overriding boundary ratio from %s to %s",
+                        boundary_ratio, batch.boundary_ratio)
+            boundary_ratio = batch.boundary_ratio
+        if boundary_ratio is not None:
+            boundary_timestep = (boundary_ratio *
+                                 self.stage.scheduler.num_train_timesteps)
+        else:
+            boundary_timestep = None
+
+        latent_model_input = latents.to(target_dtype)
+        if latent_model_input.ndim == 5:
+            assert latent_model_input.shape[0] == 1, (
+                "only support batch size 1")
+
+        ti2v_mask = None
+        ti2v_z = None
+        ti2v_seq_len = None
+        if (fastvideo_args.pipeline_config.ti2v_task
+                and batch.pil_image is not None):
+            assert batch.image_latent is None, (
+                "TI2V task should not have image latents")
+            assert self.stage.vae is not None, (
+                "VAE is not provided for TI2V task")
+            z = self.stage.vae.encode(batch.pil_image).mean.float()
+            if (hasattr(self.stage.vae, "shift_factor")
+                    and self.stage.vae.shift_factor is not None):
+                if isinstance(self.stage.vae.shift_factor, torch.Tensor):
+                    z -= self.stage.vae.shift_factor.to(z.device, z.dtype)
+                else:
+                    z -= self.stage.vae.shift_factor
+
+            if isinstance(self.stage.vae.scaling_factor, torch.Tensor):
+                z = z * self.stage.vae.scaling_factor.to(z.device, z.dtype)
+            else:
+                z = z * self.stage.vae.scaling_factor
+
+            latent_model_input = latents.to(target_dtype).squeeze(0)
+            _, mask2 = masks_like([latent_model_input], zero=True)
+            latent_model_input = ((1. - mask2[0]) * z +
+                                  mask2[0] * latent_model_input)
+            latent_model_input = latent_model_input.to(get_local_torch_device())
+            latents = latent_model_input
+            F = batch.num_frames
+            temporal_scale = (fastvideo_args.pipeline_config.vae_config.
+                              arch_config.scale_factor_temporal)
+            spatial_scale = (fastvideo_args.pipeline_config.vae_config.
+                             arch_config.scale_factor_spatial)
+            patch_size = (fastvideo_args.pipeline_config.dit_config.arch_config.
+                          patch_size)
+            seq_len = ((F - 1) // temporal_scale +
+                       1) * (batch.height // spatial_scale) * (
+                           batch.width // spatial_scale) // (patch_size[1] *
+                                                             patch_size[2])
+            ti2v_mask = mask2[0]
+            ti2v_z = z
+            ti2v_seq_len = seq_len
+
+        trajectory_timesteps: list[torch.Tensor] | None = None
+        trajectory_latents: list[torch.Tensor] | None = None
+        if batch.return_trajectory_latents:
+            trajectory_timesteps = []
+            trajectory_latents = []
+
+        progress_bar = self.stage.progress_bar(total=num_inference_steps)
+
+        extra: dict[str, Any] = {
+            "batch": batch,
+            "fastvideo_args": fastvideo_args,
+            "extra_step_kwargs": extra_step_kwargs,
+            "target_dtype": target_dtype,
+            "autocast_enabled": autocast_enabled,
+            "num_warmup_steps": num_warmup_steps,
+            "image_kwargs": image_kwargs,
+            "pos_cond_kwargs": pos_cond_kwargs,
+            "neg_cond_kwargs": neg_cond_kwargs,
+            "action_kwargs": action_kwargs,
+            "boundary_timestep": boundary_timestep,
+            "progress_bar": progress_bar,
+            "trajectory_timesteps": trajectory_timesteps,
+            "trajectory_latents": trajectory_latents,
+            "ti2v_mask": ti2v_mask,
+            "ti2v_z": ti2v_z,
+            "ti2v_seq_len": ti2v_seq_len,
+        }
+
+        return StrategyState(
+            latents=latents,
+            timesteps=timesteps,
+            num_inference_steps=num_inference_steps,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=neg_prompt_embeds,
+            prompt_attention_mask=batch.prompt_attention_mask,
+            negative_attention_mask=batch.negative_attention_mask,
+            image_embeds=image_embeds,
+            guidance_scale=batch.guidance_scale,
+            guidance_scale_2=batch.guidance_scale_2,
+            guidance_rescale=batch.guidance_rescale,
+            do_cfg=batch.do_classifier_free_guidance,
+            extra=extra,
+        )
+
+    def make_model_inputs(self, state: StrategyState, t: torch.Tensor,
+                          step_idx: int) -> ModelInputs:
+        batch = state.extra["batch"]
+        fastvideo_args = state.extra["fastvideo_args"]
+        target_dtype = state.extra["target_dtype"]
+        boundary_timestep = state.extra["boundary_timestep"]
+
+        if getattr(self.stage, "interrupt", False):
+            state.extra["skip_step"] = True
+        else:
+            state.extra["skip_step"] = False
+
+        if boundary_timestep is None or t >= boundary_timestep:
+            if (fastvideo_args.dit_cpu_offload
+                    and not fastvideo_args.dit_layerwise_offload
+                    and self.stage.transformer_2 is not None
+                    and next(self.stage.transformer_2.parameters()).device.type
+                    == 'cuda'):
+                self.stage.transformer_2.to('cpu')
+            current_model = self.stage.transformer
+            if (fastvideo_args.dit_cpu_offload
+                    and not fastvideo_args.dit_layerwise_offload
+                    and not fastvideo_args.use_fsdp_inference
+                    and current_model is not None):
+                transformer_device = next(
+                    current_model.parameters()).device.type
+                if transformer_device == 'cpu':
+                    current_model.to(get_local_torch_device())
+            current_guidance_scale = batch.guidance_scale
+        else:
+            if (fastvideo_args.dit_cpu_offload
+                    and not fastvideo_args.dit_layerwise_offload
+                    and next(self.stage.transformer.parameters()).device.type
+                    == 'cuda'):
+                self.stage.transformer.to('cpu')
+            current_model = self.stage.transformer_2
+            if (fastvideo_args.dit_cpu_offload
+                    and not fastvideo_args.dit_layerwise_offload
+                    and not fastvideo_args.use_fsdp_inference
+                    and current_model is not None):
+                transformer_2_device = next(
+                    current_model.parameters()).device.type
+                if transformer_2_device == 'cpu':
+                    current_model.to(get_local_torch_device())
+            current_guidance_scale = batch.guidance_scale_2
+
+        assert current_model is not None, "current_model is None"
+        state.extra["current_model"] = current_model
+        state.extra["current_guidance_scale"] = current_guidance_scale
+        state.extra["step_idx"] = step_idx
+
+        latent_model_input = state.latents.to(target_dtype)
+        if batch.video_latent is not None:
+            latent_model_input = torch.cat([
+                latent_model_input, batch.video_latent,
+                torch.zeros_like(state.latents)
+            ],
+                                           dim=1).to(target_dtype)
+        elif batch.image_latent is not None:
+            assert not fastvideo_args.pipeline_config.ti2v_task, (
+                "image latents should not be provided for TI2V task")
+            latent_model_input = torch.cat(
+                [latent_model_input, batch.image_latent],
+                dim=1).to(target_dtype)
+
+        if (fastvideo_args.pipeline_config.ti2v_task
+                and batch.pil_image is not None):
+            timestep = torch.stack([t]).to(get_local_torch_device())
+            mask2 = state.extra["ti2v_mask"]
+            seq_len = state.extra["ti2v_seq_len"]
+            temp_ts = (mask2[0][:, ::2, ::2] * timestep).flatten()
+            temp_ts = torch.cat([
+                temp_ts,
+                temp_ts.new_ones(seq_len - temp_ts.size(0)) * timestep
+            ])
+            timestep = temp_ts.unsqueeze(0)
+            t_expand = timestep.repeat(latent_model_input.shape[0], 1)
+        else:
+            t_expand = t.repeat(latent_model_input.shape[0])
+
+        latent_model_input = self.stage.scheduler.scale_model_input(
+            latent_model_input, t)
+
+        guidance_expand = None
+        if fastvideo_args.pipeline_config.embedded_cfg_scale is not None:
+            guidance_expand = (torch.tensor(
+                [fastvideo_args.pipeline_config.embedded_cfg_scale] *
+                latent_model_input.shape[0],
+                dtype=torch.float32,
+                device=get_local_torch_device(),
+            ).to(target_dtype) * 1000.0)
+
+        state.extra["guidance_expand"] = guidance_expand
+        return ModelInputs(
+            latent_model_input=latent_model_input,
+            timestep=t_expand,
+            prompt_embeds=state.prompt_embeds,
+            prompt_attention_mask=state.prompt_attention_mask,
+        )
+
+    def forward(self, state: StrategyState,
+                model_inputs: ModelInputs) -> torch.Tensor:
+        if state.extra.get("skip_step", False):
+            return state.latents
+
+        batch = state.extra["batch"]
+        fastvideo_args = state.extra["fastvideo_args"]
+        target_dtype = state.extra["target_dtype"]
+        autocast_enabled = state.extra["autocast_enabled"]
+        current_model = state.extra["current_model"]
+        current_guidance_scale = state.extra["current_guidance_scale"]
+        step_idx = state.extra["step_idx"]
+        guidance_expand = state.extra["guidance_expand"]
+
+        image_kwargs = state.extra["image_kwargs"]
+        pos_cond_kwargs = state.extra["pos_cond_kwargs"]
+        neg_cond_kwargs = state.extra["neg_cond_kwargs"]
+        action_kwargs = state.extra["action_kwargs"]
+
+        if ((st_attn_available
+             and self.stage.attn_backend == SlidingTileAttentionBackend) or
+            (vsa_available
+             and self.stage.attn_backend == VideoSparseAttentionBackend)):
+            self.attn_metadata_builder_cls = (
+                self.stage.attn_backend.get_builder_cls())
+
+            if self.attn_metadata_builder_cls is not None:
+                self.attn_metadata_builder = self.attn_metadata_builder_cls()
+                attn_metadata = self.attn_metadata_builder.build(  # type: ignore
+                    current_timestep=step_idx,  # type: ignore
+                    raw_latent_shape=batch.
+                    raw_latent_shape[2:5],  # type: ignore
+                    patch_size=fastvideo_args.pipeline_config.dit_config.
+                    patch_size,  # type: ignore
+                    STA_param=batch.STA_param,  # type: ignore
+                    VSA_sparsity=fastvideo_args.VSA_sparsity,  # type: ignore
+                    device=get_local_torch_device(),
+                )
+                assert attn_metadata is not None, (
+                    "attn_metadata cannot be None")
+            else:
+                attn_metadata = None
+        elif (vmoba_attn_available
+              and self.stage.attn_backend == VMOBAAttentionBackend):
+            self.attn_metadata_builder_cls = (
+                self.stage.attn_backend.get_builder_cls())
+            if self.attn_metadata_builder_cls is not None:
+                self.attn_metadata_builder = self.attn_metadata_builder_cls()
+                moba_params = fastvideo_args.moba_config.copy()
+                moba_params.update({
+                    "current_timestep":
+                    step_idx,
+                    "raw_latent_shape":
+                    batch.raw_latent_shape[2:5],
+                    "patch_size":
+                    fastvideo_args.pipeline_config.dit_config.patch_size,
+                    "device":
+                    get_local_torch_device(),
+                })
+                attn_metadata = self.attn_metadata_builder.build(**moba_params)
+                assert attn_metadata is not None, (
+                    "attn_metadata cannot be None")
+            else:
+                attn_metadata = None
+        else:
+            attn_metadata = None
+
+        with torch.autocast(device_type="cuda",
+                            dtype=target_dtype,
+                            enabled=autocast_enabled):
+            batch.is_cfg_negative = False
+            with set_forward_context(
+                    current_timestep=step_idx,
+                    attn_metadata=attn_metadata,
+                    forward_batch=batch,
+            ):
+                noise_pred = current_model(
+                    model_inputs.latent_model_input,
+                    state.prompt_embeds,
+                    model_inputs.timestep,
+                    guidance=guidance_expand,
+                    **image_kwargs,
+                    **pos_cond_kwargs,
+                    **action_kwargs,
+                )
+
+            if state.do_cfg:
+                batch.is_cfg_negative = True
+                with set_forward_context(
+                        current_timestep=step_idx,
+                        attn_metadata=attn_metadata,
+                        forward_batch=batch,
+                ):
+                    noise_pred_uncond = current_model(
+                        model_inputs.latent_model_input,
+                        state.negative_prompt_embeds,
+                        model_inputs.timestep,
+                        guidance=guidance_expand,
+                        **image_kwargs,
+                        **neg_cond_kwargs,
+                        **action_kwargs,
+                    )
+
+                noise_pred_text = noise_pred
+                noise_pred = noise_pred_uncond + current_guidance_scale * (
+                    noise_pred_text - noise_pred_uncond)
+
+                if state.guidance_rescale > 0.0:
+                    noise_pred = self.stage.rescale_noise_cfg(
+                        noise_pred,
+                        noise_pred_text,
+                        guidance_rescale=state.guidance_rescale,
+                    )
+
+        return noise_pred
+
+    def cfg_combine(self, state: StrategyState,
+                    noise_pred: torch.Tensor) -> torch.Tensor:
+        return noise_pred
+
+    def scheduler_step(self, state: StrategyState, noise_pred: torch.Tensor,
+                       t: torch.Tensor) -> torch.Tensor:
+        if state.extra.get("skip_step", False):
+            return state.latents
+
+        batch = state.extra["batch"]
+        extra_step_kwargs = state.extra["extra_step_kwargs"]
+        latents = self.stage.scheduler.step(
+            noise_pred,
+            t,
+            state.latents,
+            **extra_step_kwargs,
+            return_dict=False,
+        )[0]
+
+        if (state.extra["ti2v_mask"] is not None
+                and batch.pil_image is not None):
+            mask2 = state.extra["ti2v_mask"]
+            z = state.extra["ti2v_z"]
+            latents = latents.squeeze(0)
+            latents = (1. - mask2) * z + mask2 * latents
+
+        if state.extra["trajectory_latents"] is not None:
+            state.extra["trajectory_timesteps"].append(t)
+            state.extra["trajectory_latents"].append(latents)
+
+        progress_bar = state.extra["progress_bar"]
+        step_idx = state.extra["step_idx"]
+        num_warmup_steps = state.extra["num_warmup_steps"]
+        timesteps = state.timesteps
+        if step_idx == len(timesteps) - 1 or (
+            (step_idx + 1) > num_warmup_steps and
+            (step_idx + 1) % self.stage.scheduler.order == 0
+                and progress_bar is not None):
+            progress_bar.update()
+
+        return latents
+
+    def postprocess(self, state: StrategyState) -> ForwardBatch:
+        batch = state.extra["batch"]
+        fastvideo_args = state.extra["fastvideo_args"]
+        progress_bar = state.extra.get("progress_bar")
+
+        if state.extra["trajectory_latents"]:
+            trajectory_tensor = torch.stack(state.extra["trajectory_latents"],
+                                            dim=1)
+            trajectory_timesteps_tensor = torch.stack(
+                state.extra["trajectory_timesteps"], dim=0)
+            batch.trajectory_timesteps = trajectory_timesteps_tensor.cpu()
+            batch.trajectory_latents = trajectory_tensor.cpu()
+
+        batch.latents = state.latents
+
+        if fastvideo_args.dit_layerwise_offload:
+            mgr = getattr(self.stage.transformer, "_layerwise_offload_manager",
+                          None)
+            if mgr is not None and getattr(mgr, "enabled", False):
+                mgr.release_all()
+            if self.stage.transformer_2 is not None:
+                mgr2 = getattr(self.stage.transformer_2,
+                               "_layerwise_offload_manager", None)
+                if mgr2 is not None and getattr(mgr2, "enabled", False):
+                    mgr2.release_all()
+
+        if (st_attn_available
+                and self.stage.attn_backend == SlidingTileAttentionBackend
+                and fastvideo_args.STA_mode == STA_Mode.STA_SEARCHING):
+            self.stage.save_sta_search_results(batch)
+
+        pipeline = self.stage.pipeline() if self.stage.pipeline else None
+        if torch.backends.mps.is_available():
+            logger.info("Memory before deallocating transformer: %s",
+                        torch.mps.current_allocated_memory())
+            del self.stage.transformer
+            if pipeline is not None and "transformer" in pipeline.modules:
+                del pipeline.modules["transformer"]
+            fastvideo_args.model_loaded["transformer"] = False
+            logger.info("Memory after deallocating transformer: %s",
+                        torch.mps.current_allocated_memory())
+
+        if progress_bar is not None:
+            progress_bar.close()
+
+        return batch

fastvideo/pipelines/stages/denoising_strategies.py

@@ -0,0 +1,106 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+Strategy interfaces and shared types for unified denoising.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any, Protocol, runtime_checkable
+
+import torch
+
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+
+
+@dataclass
+class StrategyState:
+    latents: torch.Tensor
+    timesteps: torch.Tensor
+    num_inference_steps: int
+    prompt_embeds: list[torch.Tensor]
+    negative_prompt_embeds: list[torch.Tensor] | None
+    prompt_attention_mask: list[torch.Tensor] | None
+    negative_attention_mask: list[torch.Tensor] | None
+    image_embeds: list[torch.Tensor]
+    guidance_scale: float
+    guidance_scale_2: float | None
+    guidance_rescale: float
+    do_cfg: bool
+    extra: dict[str, Any] = field(default_factory=dict)
+
+
+@dataclass
+class ModelInputs:
+    latent_model_input: torch.Tensor
+    timestep: torch.Tensor
+    prompt_embeds: torch.Tensor | list[torch.Tensor]
+    prompt_attention_mask: torch.Tensor | list[torch.Tensor] | None
+    extra_kwargs: dict[str, Any] = field(default_factory=dict)
+
+
+@dataclass
+class BlockPlanItem:
+    start_index: int
+    num_frames: int
+    use_kv_cache: bool
+    model_selector: str
+
+
+@dataclass
+class BlockPlan:
+    items: list[BlockPlanItem] = field(default_factory=list)
+
+
+@dataclass
+class BlockContext:
+    kv_cache: list[dict] | None
+    kv_cache_2: list[dict] | None
+    crossattn_cache: list[dict] | None
+    action_cache: dict[str, list[dict]] | None
+    extra: dict[str, Any] = field(default_factory=dict)
+
+
+class DenoisingStrategy(Protocol):
+
+    def prepare(self, batch: ForwardBatch, args) -> StrategyState:
+        ...
+
+    def make_model_inputs(self, state: StrategyState, t: torch.Tensor,
+                          step_idx: int) -> ModelInputs:
+        ...
+
+    def forward(self, state: StrategyState,
+                model_inputs: ModelInputs) -> torch.Tensor:
+        ...
+
+    def cfg_combine(self, state: StrategyState,
+                    noise_pred: torch.Tensor) -> torch.Tensor:
+        ...
+
+    def scheduler_step(self, state: StrategyState, noise_pred: torch.Tensor,
+                       t: torch.Tensor) -> torch.Tensor:
+        ...
+
+    def postprocess(self, state: StrategyState) -> ForwardBatch:
+        ...
+
+
+@runtime_checkable
+class BlockDenoisingStrategy(DenoisingStrategy, Protocol):
+
+    def block_plan(self, state: StrategyState) -> BlockPlan:
+        ...
+
+    def init_block_context(self, state: StrategyState,
+                           block_item: BlockPlanItem,
+                           block_idx: int) -> BlockContext:
+        ...
+
+    def process_block(self, state: StrategyState, block_ctx: BlockContext,
+                      block_item: BlockPlanItem) -> None:
+        ...
+
+    def update_context(self, state: StrategyState, block_ctx: BlockContext,
+                       block_item: BlockPlanItem) -> None:
+        ...

fastvideo/pipelines/stages/longcat_denoising.py

@@ -1,179 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""
-LongCat-specific denoising stage implementing CFG-zero optimized guidance.
-"""
-
-import torch
-from tqdm import tqdm
-
-from fastvideo.fastvideo_args import FastVideoArgs
-from fastvideo.logger import init_logger
-from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
-from fastvideo.pipelines.stages.denoising import DenoisingStage
-from fastvideo.forward_context import set_forward_context
-
-logger = init_logger(__name__)
-
-
-class LongCatDenoisingStage(DenoisingStage):
-    """
-    LongCat denoising stage with CFG-zero optimized guidance scale.
-    
-    Implements:
-    1. Optimized CFG scale from CFG-zero paper
-    2. Negation of noise prediction before scheduler step (flow matching convention)
-    3. Batched CFG computation (unlike standard FastVideo separate passes)
-    """
-
-    def optimized_scale(self, positive_flat, negative_flat) -> torch.Tensor:
-        """
-        Calculate optimized scale from CFG-zero paper.
-        
-        st_star = (v_cond^T * v_uncond) / ||v_uncond||^2
-        
-        Args:
-            positive_flat: Conditional prediction, flattened [B, -1]
-            negative_flat: Unconditional prediction, flattened [B, -1]
-        
-        Returns:
-            st_star: Optimized scale [B, 1]
-        """
-        # Calculate dot product
-        dot_product = torch.sum(positive_flat * negative_flat,
-                                dim=1,
-                                keepdim=True)
-        # Squared norm of uncondition
-        squared_norm = torch.sum(negative_flat**2, dim=1, keepdim=True) + 1e-8
-        # st_star = v_cond^T * v_uncond / ||v_uncond||^2
-        st_star = dot_product / squared_norm
-        return st_star
-
-    def forward(
-        self,
-        batch: ForwardBatch,
-        fastvideo_args: FastVideoArgs,
-    ) -> ForwardBatch:
-        """
-        Run LongCat denoising loop with optimized CFG.
-        
-        Args:
-            batch: The current batch information.
-            fastvideo_args: The inference arguments.
-            
-        Returns:
-            The batch with denoised latents.
-        """
-        if not fastvideo_args.model_loaded["transformer"]:
-            from fastvideo.models.model_loader import TransformerLoader
-            loader = TransformerLoader()
-            self.transformer = loader.load(
-                fastvideo_args.model_paths["transformer"], fastvideo_args)
-            pipeline = self.pipeline() if self.pipeline else None
-            if pipeline:
-                pipeline.add_module("transformer", self.transformer)
-            fastvideo_args.model_loaded["transformer"] = True
-
-        # Get transformer dtype
-        if hasattr(self.transformer, 'module'):
-            transformer_dtype = next(self.transformer.module.parameters()).dtype
-        else:
-            transformer_dtype = next(self.transformer.parameters()).dtype
-
-        target_dtype = transformer_dtype
-        autocast_enabled = (target_dtype != torch.float32
-                            ) and not fastvideo_args.disable_autocast
-
-        # Extract batch parameters
-        latents = batch.latents
-        timesteps = batch.timesteps
-        prompt_embeds = batch.prompt_embeds[0]  # LongCat uses single encoder
-        prompt_attention_mask = batch.prompt_attention_mask[
-            0] if batch.prompt_attention_mask else None
-        guidance_scale = batch.guidance_scale
-        do_classifier_free_guidance = batch.do_classifier_free_guidance
-
-        # Get negative prompts if doing CFG
-        if do_classifier_free_guidance:
-            negative_prompt_embeds = batch.negative_prompt_embeds[0]
-            negative_prompt_attention_mask = (batch.negative_attention_mask[0]
-                                              if batch.negative_attention_mask
-                                              else None)
-            # Concatenate for batched processing
-            prompt_embeds_combined = torch.cat(
-                [negative_prompt_embeds, prompt_embeds], dim=0)
-            if prompt_attention_mask is not None:
-                prompt_attention_mask_combined = torch.cat(
-                    [negative_prompt_attention_mask, prompt_attention_mask],
-                    dim=0)
-            else:
-                prompt_attention_mask_combined = None
-        else:
-            prompt_embeds_combined = prompt_embeds
-            prompt_attention_mask_combined = prompt_attention_mask
-
-        # Denoising loop
-        num_inference_steps = len(timesteps)
-        with tqdm(total=num_inference_steps,
-                  desc="LongCat Denoising") as progress_bar:
-            for i, t in enumerate(timesteps):
-                # Expand latents for CFG
-                if do_classifier_free_guidance:
-                    latent_model_input = torch.cat([latents] * 2)
-                else:
-                    latent_model_input = latents
-
-                latent_model_input = latent_model_input.to(target_dtype)
-
-                # Expand timestep to match batch size
-                timestep = t.expand(
-                    latent_model_input.shape[0]).to(target_dtype)
-
-                # Run transformer with context
-                batch.is_cfg_negative = False
-                with set_forward_context(
-                        current_timestep=i,
-                        attn_metadata=None,
-                        forward_batch=batch,
-                ), torch.autocast(device_type='cuda',
-                                  dtype=target_dtype,
-                                  enabled=autocast_enabled):
-                    noise_pred = self.transformer(
-                        hidden_states=latent_model_input,
-                        encoder_hidden_states=prompt_embeds_combined,
-                        timestep=timestep,
-                        encoder_attention_mask=prompt_attention_mask_combined,
-                    )
-
-                # Apply CFG with optimized scale
-                if do_classifier_free_guidance:
-                    noise_pred_uncond, noise_pred_cond = noise_pred.chunk(2)
-
-                    B = noise_pred_cond.shape[0]
-                    positive = noise_pred_cond.reshape(B, -1)
-                    negative = noise_pred_uncond.reshape(B, -1)
-
-                    # Calculate optimized scale (CFG-zero)
-                    st_star = self.optimized_scale(positive, negative)
-
-                    # Reshape for broadcasting
-                    st_star = st_star.view(B, 1, 1, 1, 1)
-
-                    # Apply optimized CFG formula
-                    noise_pred = (
-                        noise_pred_uncond * st_star + guidance_scale *
-                        (noise_pred_cond - noise_pred_uncond * st_star))
-
-                # CRITICAL: Negate noise prediction for flow matching scheduler
-                noise_pred = -noise_pred
-
-                # Compute previous noisy sample x_t -> x_t-1
-                latents = self.scheduler.step(noise_pred,
-                                              t,
-                                              latents,
-                                              return_dict=False)[0]
-
-                progress_bar.update()
-
-        # Update batch with denoised latents
-        batch.latents = latents
-        return batch

fastvideo/pipelines/stages/longcat_i2v_latent_preparation.py

@@ -0,0 +1,105 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+LongCat I2V Latent Preparation Stage.
+
+This stage prepares latents with image conditioning for the first frame.
+"""
+
+import torch
+
+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.latent_preparation import LatentPreparationStage
+
+logger = init_logger(__name__)
+
+
+class LongCatI2VLatentPreparationStage(LatentPreparationStage):
+    """
+    Prepare latents with image conditioning for first frame.
+    
+    This stage:
+    1. Generates random noise for all frames
+    2. Replaces first latent frame with encoded image latent
+    3. Marks conditioning information in batch
+    """
+
+    # Uses parent __init__ - no need for additional constructor
+
+    def forward(
+        self,
+        batch: ForwardBatch,
+        fastvideo_args: FastVideoArgs,
+    ) -> ForwardBatch:
+        """Prepare latents with I2V conditioning."""
+
+        # IMPORTANT: Skip if latents already prepared (e.g., by refinement init stage)
+        # The refine_init stage encodes stage1 video and mixes with noise - don't overwrite!
+        if batch.latents is not None:
+            logger.info(
+                "I2V Latent Prep: Skipping - latents already prepared "
+                "(shape=%s), likely from refinement stage", batch.latents.shape)
+            return batch
+
+        # 1. Calculate dimensions
+        num_frames = batch.num_frames
+        height = batch.height
+        width = batch.width
+
+        # Get VAE compression factors
+        # IMPORTANT: Use VAE's temporal compression (4), NOT transformer's patch_size[0] (1)
+        vae_temporal_scale = fastvideo_args.pipeline_config.vae_config.arch_config.scale_factor_temporal
+        vae_spatial_scale = fastvideo_args.pipeline_config.vae_config.arch_config.scale_factor_spatial
+
+        num_latent_frames = (num_frames - 1) // vae_temporal_scale + 1
+        latent_height = height // vae_spatial_scale
+        latent_width = width // vae_spatial_scale
+
+        num_channels = self.transformer.config.in_channels
+
+        logger.info(
+            "I2V Latent Prep: num_frames=%s, num_latent_frames=%s "
+            "(vae_temporal_scale=%s), latent_shape=(%s, %s)", num_frames,
+            num_latent_frames, vae_temporal_scale, latent_height, latent_width)
+
+        # 2. Generate random noise for all frames
+        # batch_size might not be set, default to 1
+        batch_size = batch.batch_size if batch.batch_size is not None else 1
+        shape = (batch_size, num_channels, num_latent_frames, latent_height,
+                 latent_width)
+
+        # Handle generator - may be a list for batch handling
+        generator = batch.generator
+        if isinstance(generator, list):
+            generator = generator[0] if generator else None
+
+        # torch.randn requires specific argument order: size, generator, dtype
+        latents = torch.randn(*shape,
+                              generator=generator).to(get_local_torch_device(),
+                                                      dtype=torch.float32)
+
+        # 3. Replace first frame with conditioned image latent
+        if batch.image_latent is not None:
+            num_cond_latents = batch.num_cond_latents
+            latents[:, :, :
+                    num_cond_latents] = batch.image_latent[:, :, :
+                                                           num_cond_latents]
+
+            logger.info(
+                "I2V: Replaced first %s latent frame(s) with image conditioning",
+                num_cond_latents)
+        else:
+            logger.warning(
+                "No image_latent found in batch, proceeding without conditioning"
+            )
+
+        # 4. Store in batch
+        batch.latents = latents
+
+        # Required by base class output validator
+        batch.raw_latent_shape = (num_latent_frames, latent_height,
+                                  latent_width)
+
+        return batch

fastvideo/pipelines/stages/longcat_image_vae_encoding.py

@@ -0,0 +1,162 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+LongCat Image VAE Encoding Stage for I2V generation.
+
+This stage handles encoding a single input image to latent space with
+LongCat-specific normalization for I2V conditioning.
+"""
+
+import PIL
+import torch
+
+from fastvideo.distributed import get_local_torch_device
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.logger import init_logger
+from fastvideo.models.vision_utils import (normalize, numpy_to_pt, pil_to_numpy,
+                                           resize)
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+from fastvideo.pipelines.stages.base import PipelineStage
+from fastvideo.utils import PRECISION_TO_TYPE
+
+logger = init_logger(__name__)
+
+
+class LongCatImageVAEEncodingStage(PipelineStage):
+    """
+    Encode input image to latent space for I2V conditioning.
+    
+    This stage:
+    1. Preprocesses image to match target dimensions
+    2. Encodes via VAE to latent space
+    3. Applies LongCat-specific normalization
+    4. Stores latent and calculates num_cond_latents
+    """
+
+    def __init__(self, vae):
+        super().__init__()
+        self.vae = vae
+
+    def forward(
+        self,
+        batch: ForwardBatch,
+        fastvideo_args: FastVideoArgs,
+    ) -> ForwardBatch:
+        """Encode image to latent for I2V conditioning."""
+
+        # Skip image encoding for refinement tasks - we're refining an existing video
+        if getattr(batch, 'stage1_video', None) is not None or getattr(
+                batch, 'refine_from', None) is not None:
+            logger.info(
+                "Skipping image encoding - refinement mode (using stage1_video)"
+            )
+            return batch
+
+        # 1. Get image from batch
+        image = batch.pil_image  # PIL.Image
+        if image is None:
+            raise ValueError("pil_image must be provided for I2V")
+
+        if not isinstance(image, PIL.Image.Image):
+            raise TypeError(f"pil_image must be PIL.Image, got {type(image)}")
+
+        # 2. Get target dimensions
+        height = batch.height
+        width = batch.width
+
+        if height is None or width is None:
+            raise ValueError("height and width must be set for I2V")
+
+        # 3. Preprocess image
+        image = resize(image, height, width, resize_mode="default")
+        image = pil_to_numpy(image)
+        image = numpy_to_pt(image)
+        image = normalize(image)  # to [-1, 1]
+
+        # 4. Add temporal dimension
+        # After numpy_to_pt: [1, C, H, W] (batch already added by pil_to_numpy)
+        # Add T dimension: [1, C, H, W] -> [1, C, 1, H, W] = [B, C, T, H, W]
+        image = image.unsqueeze(2)
+        image = image.to(get_local_torch_device(), dtype=torch.float32)
+
+        # 5. Encode via VAE
+        self.vae = self.vae.to(get_local_torch_device())
+
+        # Setup VAE precision
+        vae_dtype = PRECISION_TO_TYPE[
+            fastvideo_args.pipeline_config.vae_precision]
+        vae_autocast_enabled = (
+            vae_dtype != torch.float32) and not fastvideo_args.disable_autocast
+
+        with torch.autocast(device_type="cuda",
+                            dtype=vae_dtype,
+                            enabled=vae_autocast_enabled):
+            if fastvideo_args.pipeline_config.vae_tiling:
+                self.vae.enable_tiling()
+
+            if not vae_autocast_enabled:
+                image = image.to(vae_dtype)
+
+            with torch.no_grad():
+                encoder_output = self.vae.encode(image)
+                latent = self.retrieve_latents(encoder_output, batch.generator)
+
+        # 6. Apply LongCat-specific normalization
+        # Formula: (latents - mean) / std
+        latent = self.normalize_latents(latent)
+
+        # 7. Calculate num_cond_latents
+        # Formula: 1 + (num_cond_frames - 1) // vae_temporal_scale
+        # For single image (num_cond_frames=1): always 1 latent frame
+        num_cond_frames = 1  # Single image
+        vae_temporal_scale = self.vae.config.scale_factor_temporal
+        batch.num_cond_latents = 1 + (num_cond_frames - 1) // vae_temporal_scale
+
+        # 8. Store in batch
+        batch.image_latent = latent
+        batch.num_cond_frames = 1
+
+        logger.info(
+            "I2V: Encoded image to latent shape %s, num_cond_latents=%s",
+            latent.shape, batch.num_cond_latents)
+
+        # Offload VAE if needed
+        if fastvideo_args.vae_cpu_offload:
+            self.vae.to("cpu")
+
+        return batch
+
+    def retrieve_latents(self, encoder_output: object,
+                         generator: torch.Generator | None) -> torch.Tensor:
+        """Sample from VAE posterior."""
+        # WAN VAE returns an object with .sample() method
+        if hasattr(encoder_output, 'sample'):
+            return encoder_output.sample(generator)
+        elif hasattr(encoder_output, 'latent_dist'):
+            return encoder_output.latent_dist.sample(generator)
+        elif hasattr(encoder_output, 'latents'):
+            return encoder_output.latents
+        else:
+            raise AttributeError("Could not access latents from encoder output")
+
+    def normalize_latents(self, latents: torch.Tensor) -> torch.Tensor:
+        """
+        Apply LongCat-specific latent normalization.
+        
+        Formula: (latents - mean) / std
+        
+        This matches the original LongCat implementation and is DIFFERENT
+        from standard VAE scaling (which uses scaling_factor).
+        """
+        if not hasattr(self.vae.config, 'latents_mean') or not hasattr(
+                self.vae.config, 'latents_std'):
+            raise ValueError(
+                "VAE config must have 'latents_mean' and 'latents_std' "
+                "for LongCat normalization")
+
+        latents_mean = torch.tensor(self.vae.config.latents_mean).view(
+            1, self.vae.config.z_dim, 1, 1, 1).to(latents.device, latents.dtype)
+
+        latents_std = torch.tensor(self.vae.config.latents_std).view(
+            1, self.vae.config.z_dim, 1, 1, 1).to(latents.device, latents.dtype)
+
+        return (latents - latents_mean) / latents_std

fastvideo/pipelines/stages/longcat_kv_cache_init.py

@@ -0,0 +1,123 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+LongCat KV Cache Initialization Stage for Video Continuation (VC).
+
+This stage pre-computes K/V cache for conditioning frames.
+
+"""
+
+import torch
+
+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
+
+logger = init_logger(__name__)
+
+
+class LongCatKVCacheInitStage(PipelineStage):
+    """
+    Pre-compute KV cache for conditioning frames.
+
+    After this stage:
+    - batch.kv_cache_dict contains {block_idx: (k, v)}
+    - batch.cond_latents contains the conditioning latents
+    - batch.latents contains ONLY noise latents
+    """
+
+    def __init__(self, transformer):
+        super().__init__()
+        self.transformer = transformer
+
+    def forward(
+        self,
+        batch: ForwardBatch,
+        fastvideo_args: FastVideoArgs,
+    ) -> ForwardBatch:
+        """Initialize KV cache from conditioning latents."""
+
+        # Check if KV cache is enabled
+        use_kv_cache = getattr(fastvideo_args.pipeline_config, 'use_kv_cache',
+                               True)
+        if not use_kv_cache:
+            batch.kv_cache_dict = {}
+            batch.use_kv_cache = False
+            logger.info("KV cache disabled, skipping initialization")
+            return batch
+
+        batch.use_kv_cache = True
+        offload_kv_cache = getattr(fastvideo_args.pipeline_config,
+                                   'offload_kv_cache', False)
+
+        # Get conditioning latents
+        num_cond_latents = batch.num_cond_latents
+        if num_cond_latents <= 0:
+            batch.kv_cache_dict = {}
+            logger.warning("num_cond_latents <= 0, skipping KV cache init")
+            return batch
+
+        # Extract conditioning latents
+        cond_latents = batch.latents[:, :, :num_cond_latents].clone()
+
+        logger.info(
+            "Initializing KV cache for %d conditioning latents, shape: %s",
+            num_cond_latents, cond_latents.shape)
+
+        # Timestep = 0 for conditioning (they are "clean")
+        B = cond_latents.shape[0]
+        T_cond = cond_latents.shape[2]
+        timestep = torch.zeros(B,
+                               T_cond,
+                               device=cond_latents.device,
+                               dtype=cond_latents.dtype)
+
+        # Empty prompt embeddings (cross-attn will be skipped)
+        max_seq_len = 512
+        # Get caption dimension from transformer config
+        caption_dim = self.transformer.config.caption_channels
+        empty_embeds = torch.zeros(B,
+                                   max_seq_len,
+                                   caption_dim,
+                                   device=cond_latents.device,
+                                   dtype=cond_latents.dtype)
+
+        # Get transformer dtype
+        if hasattr(self.transformer, 'module'):
+            transformer_dtype = next(self.transformer.module.parameters()).dtype
+        else:
+            transformer_dtype = next(self.transformer.parameters()).dtype
+
+        # Run transformer with return_kv=True, skip_crs_attn=True
+        with (
+                torch.no_grad(),
+                set_forward_context(
+                    current_timestep=0,
+                    attn_metadata=None,
+                    forward_batch=batch,
+                ),
+                torch.autocast(device_type='cuda', dtype=transformer_dtype),
+        ):
+            _, kv_cache_dict = self.transformer(
+                hidden_states=cond_latents.to(transformer_dtype),
+                encoder_hidden_states=empty_embeds.to(transformer_dtype),
+                timestep=timestep.to(transformer_dtype),
+                return_kv=True,
+                skip_crs_attn=True,
+                offload_kv_cache=offload_kv_cache,
+            )
+
+        # Store cache and save cond_latents for later concatenation
+        batch.kv_cache_dict = kv_cache_dict
+        batch.cond_latents = cond_latents
+
+        # Remove conditioning latents from main latents
+        # After this, batch.latents contains ONLY noise frames
+        batch.latents = batch.latents[:, :, num_cond_latents:]
+
+        logger.info(
+            "KV cache initialized: %d blocks, offload=%s, remaining latents shape: %s",
+            len(kv_cache_dict), offload_kv_cache, batch.latents.shape)
+
+        return batch

fastvideo/pipelines/stages/longcat_refine_init.py

@@ -257,12 +257,16 @@ class LongCatRefineInitStage(PipelineStage):
                     num_cond_frames_added, num_noise_frames_added,
                     new_num_frames)
 
-        # VAE encode
-        logger.info("Encoding stage1 video with VAE...")
+        # VAE encode with tiling for memory efficiency
+        logger.info("Encoding stage1 video with VAE (tiling enabled)...")
         vae_dtype = next(self.vae.parameters()).dtype
         vae_device = next(self.vae.parameters()).device
         video_up = video_up.to(dtype=vae_dtype, device=vae_device)
 
+        # Enable tiling for large video encoding
+        if hasattr(self.vae, 'enable_tiling'):
+            self.vae.enable_tiling()
+
         with torch.no_grad():
             latent_dist = self.vae.encode(video_up)
             # Extract tensor from latent distribution
@@ -301,10 +305,14 @@ class LongCatRefineInitStage(PipelineStage):
 
         logger.info("Applied t_thresh=%s noise mixing", t_thresh)
 
-        # Store in batch
-        batch.latents = latent_up.to(dtype)
+        # Store in batch - ensure correct dtype and device
+        # The latents need to be on the same device as the transformer (CUDA)
+        target_device = batch.prompt_embeds[0].device
+        batch.latents = latent_up.to(device=target_device, dtype=dtype)
         batch.raw_latent_shape = latent_up.shape
 
+        logger.info("Latents device: %s, dtype: %s", batch.latents.device,
+                    batch.latents.dtype)
         logger.info("LongCat refinement initialization complete")
 
         return batch

fastvideo/pipelines/stages/longcat_video_vae_encoding.py

@@ -0,0 +1,180 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+LongCat Video VAE Encoding Stage for Video Continuation (VC) generation.
+
+This stage handles encoding multiple video frames to latent space with
+LongCat-specific normalization for VC conditioning.
+"""
+
+from typing import Any
+
+import PIL.Image
+import torch
+
+from fastvideo.distributed import get_local_torch_device
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.logger import init_logger
+from fastvideo.models.vision_utils import normalize, numpy_to_pt, pil_to_numpy, resize
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+from fastvideo.pipelines.stages.base import PipelineStage
+from fastvideo.utils import PRECISION_TO_TYPE
+
+logger = init_logger(__name__)
+
+
+class LongCatVideoVAEEncodingStage(PipelineStage):
+    """
+    Encode video frames to latent space for VC conditioning.
+    
+    This stage:
+    1. Loads video frames from path or uses provided frames
+    2. Takes the last num_cond_frames from the video
+    3. Preprocesses and stacks frames
+    4. Encodes via VAE to latent space
+    5. Applies LongCat-specific normalization
+    6. Calculates num_cond_latents
+    """
+
+    def __init__(self, vae):
+        super().__init__()
+        self.vae = vae
+
+    def forward(
+        self,
+        batch: ForwardBatch,
+        fastvideo_args: FastVideoArgs,
+    ) -> ForwardBatch:
+        """Encode video frames to latent for VC conditioning."""
+
+        # Get video from batch - can be path, list of PIL images, or already loaded
+        video = getattr(batch, 'video_frames', None) or getattr(
+            batch, 'video_path', None)
+        num_cond_frames = getattr(batch, 'num_cond_frames',
+                                  13)  # Default 13 for VC
+
+        if video is None:
+            raise ValueError(
+                "video_frames or video_path must be provided for VC")
+
+        # Load video if path
+        if isinstance(video, str):
+            from diffusers.utils import load_video
+            video = load_video(video)
+            logger.info("Loaded video from path: %d frames", len(video))
+
+        # Take last num_cond_frames
+        if len(video) > num_cond_frames:
+            video = video[-num_cond_frames:]
+            logger.info("Using last %d frames for conditioning",
+                        num_cond_frames)
+        elif len(video) < num_cond_frames:
+            logger.warning(
+                "Video has only %d frames, less than num_cond_frames=%d",
+                len(video), num_cond_frames)
+            num_cond_frames = len(video)
+
+        # Get target dimensions
+        height = batch.height
+        width = batch.width
+
+        if height is None or width is None:
+            raise ValueError("height and width must be set for VC")
+
+        # Preprocess and stack frames
+        processed_frames = []
+        for frame in video:
+            if not isinstance(frame, PIL.Image.Image):
+                raise TypeError(f"Frame must be PIL.Image, got {type(frame)}")
+
+            frame = resize(frame, height, width, resize_mode="default")
+            frame = pil_to_numpy(frame)  # Returns [1, H, W, C] then converted
+            frame = numpy_to_pt(frame)  # Returns [1, C, H, W]
+            frame = normalize(frame)  # to [-1, 1]
+            processed_frames.append(frame)
+
+        # Stack frames: [num_frames, C, H, W] -> [1, C, T, H, W]
+        video_tensor = torch.cat(processed_frames, dim=0)  # [T, C, H, W]
+        video_tensor = video_tensor.permute(1, 0, 2,
+                                            3).unsqueeze(0)  # [1, C, T, H, W]
+        video_tensor = video_tensor.to(get_local_torch_device(),
+                                       dtype=torch.float32)
+
+        logger.info("VC: Preprocessed video tensor shape: %s",
+                    video_tensor.shape)
+
+        # Encode via VAE
+        self.vae = self.vae.to(get_local_torch_device())
+
+        # Setup VAE precision
+        vae_dtype = PRECISION_TO_TYPE[
+            fastvideo_args.pipeline_config.vae_precision]
+        vae_autocast_enabled = (
+            vae_dtype != torch.float32) and not fastvideo_args.disable_autocast
+
+        with torch.autocast(device_type="cuda",
+                            dtype=vae_dtype,
+                            enabled=vae_autocast_enabled):
+            if fastvideo_args.pipeline_config.vae_tiling:
+                self.vae.enable_tiling()
+
+            if not vae_autocast_enabled:
+                video_tensor = video_tensor.to(vae_dtype)
+
+            with torch.no_grad():
+                encoder_output = self.vae.encode(video_tensor)
+                latent = self.retrieve_latents(encoder_output, batch.generator)
+
+        # Apply LongCat-specific normalization
+        latent = self.normalize_latents(latent)
+
+        # Calculate num_cond_latents
+        # Formula: 1 + (num_cond_frames - 1) // vae_temporal_scale
+        vae_temporal_scale = self.vae.config.scale_factor_temporal
+        num_cond_latents = 1 + (num_cond_frames - 1) // vae_temporal_scale
+
+        # Store in batch
+        batch.video_latent = latent
+        batch.num_cond_frames = num_cond_frames
+        batch.num_cond_latents = num_cond_latents
+
+        logger.info(
+            "VC: Encoded %d frames to latent shape %s, num_cond_latents=%d",
+            num_cond_frames, latent.shape, num_cond_latents)
+
+        # Offload VAE if needed
+        if fastvideo_args.vae_cpu_offload:
+            self.vae.to("cpu")
+
+        return batch
+
+    def retrieve_latents(self, encoder_output: Any,
+                         generator: torch.Generator | None) -> torch.Tensor:
+        """Sample from VAE posterior."""
+        if hasattr(encoder_output, 'sample'):
+            return encoder_output.sample(generator)
+        elif hasattr(encoder_output, 'latent_dist'):
+            return encoder_output.latent_dist.sample(generator)
+        elif hasattr(encoder_output, 'latents'):
+            return encoder_output.latents
+        else:
+            raise AttributeError("Could not access latents from encoder output")
+
+    def normalize_latents(self, latents: torch.Tensor) -> torch.Tensor:
+        """
+        Apply LongCat-specific latent normalization.
+        
+        Formula: (latents - mean) / std
+        """
+        if not hasattr(self.vae.config, 'latents_mean') or not hasattr(
+                self.vae.config, 'latents_std'):
+            raise ValueError(
+                "VAE config must have 'latents_mean' and 'latents_std' "
+                "for LongCat normalization")
+
+        latents_mean = torch.tensor(self.vae.config.latents_mean).view(
+            1, self.vae.config.z_dim, 1, 1, 1).to(latents.device, latents.dtype)
+
+        latents_std = torch.tensor(self.vae.config.latents_std).view(
+            1, self.vae.config.z_dim, 1, 1, 1).to(latents.device, latents.dtype)
+
+        return (latents - latents_mean) / latents_std

fastvideo/pipelines/stages/matrixgame_denoising.py

@@ -15,14 +15,6 @@ from fastvideo.pipelines.stages.denoising import DenoisingStage
 from fastvideo.pipelines.stages.validators import StageValidators as V
 from fastvideo.pipelines.stages.validators import VerificationResult
 
-try:
-    from fastvideo.attention.backends.sliding_tile_attn import (
-        SlidingTileAttentionBackend)
-    st_attn_available = True
-except ImportError:
-    st_attn_available = False
-    SlidingTileAttentionBackend = None  # type: ignore
-
 try:
     from fastvideo.attention.backends.video_sparse_attn import (
         VideoSparseAttentionBackend)
@@ -123,169 +115,22 @@ class MatrixGameCausalDenoisingStage(DenoisingStage):
 
         self._streaming_initialized: bool = False
         self._streaming_ctx: BlockProcessingContext | None = None
+        self._streaming_engine = None
+        self._streaming_state = None
+        self._streaming_block_plan = None
 
     def forward(
         self,
         batch: ForwardBatch,
         fastvideo_args: FastVideoArgs,
     ) -> ForwardBatch:
-        target_dtype = torch.bfloat16
-        autocast_enabled = (target_dtype != torch.float32
-                            ) and not fastvideo_args.disable_autocast
+        from fastvideo.pipelines.stages.denoising_engine import DenoisingEngine
+        from fastvideo.pipelines.stages.denoising_matrixgame_strategy import (
+            MatrixGameBlockStrategy)
 
-        latent_seq_length = batch.latents.shape[-1] * batch.latents.shape[-2]
-        patch_size = self.transformer.patch_size
-        patch_ratio = patch_size[-1] * patch_size[-2]
-        self.frame_seq_length = latent_seq_length // patch_ratio
-
-        timesteps = torch.tensor(
-            fastvideo_args.pipeline_config.dmd_denoising_steps,
-            dtype=torch.long).cpu()
-        if fastvideo_args.pipeline_config.warp_denoising_step:
-            scheduler_timesteps = torch.cat((self.scheduler.timesteps.cpu(),
-                                             torch.tensor([0],
-                                                          dtype=torch.float32)))
-            timesteps = scheduler_timesteps[1000 - timesteps]
-        timesteps = timesteps.to(get_local_torch_device())
-
-        boundary_ratio = getattr(fastvideo_args.pipeline_config.dit_config,
-                                 'boundary_ratio', None)
-        if boundary_ratio is not None:
-            boundary_timestep = boundary_ratio * self.scheduler.num_train_timesteps
-            high_noise_timesteps = timesteps[timesteps >= boundary_timestep]
-        else:
-            boundary_timestep = None
-            high_noise_timesteps = None
-
-        image_embeds = batch.image_embeds
-        if len(image_embeds) > 0:
-            assert torch.isnan(image_embeds[0]).sum() == 0
-            image_embeds = [
-                image_embed.to(target_dtype) for image_embed in image_embeds
-            ]
-
-        # directly set the kwarg.
-        image_kwargs = {"encoder_hidden_states_image": image_embeds}
-        pos_cond_kwargs: dict[str, Any] = {}
-
-        if st_attn_available and self.attn_backend == SlidingTileAttentionBackend:
-            self.prepare_sta_param(batch, fastvideo_args)
-
-        assert batch.latents is not None, "latents must be provided"
-        latents = batch.latents
-        b, c, t, h, w = latents.shape
-        prompt_embeds = batch.prompt_embeds
-        assert torch.isnan(prompt_embeds[0]).sum() == 0
-
-        kv_cache1 = self._initialize_kv_cache(batch_size=latents.shape[0],
-                                              dtype=target_dtype,
-                                              device=latents.device)
-        kv_cache2 = None
-        if boundary_timestep is not None:
-            kv_cache2 = self._initialize_kv_cache(batch_size=latents.shape[0],
-                                                  dtype=target_dtype,
-                                                  device=latents.device)
-
-        kv_cache_mouse = None
-        kv_cache_keyboard = None
-        if self.use_action_module:
-            kv_cache_mouse, kv_cache_keyboard = self._initialize_action_kv_cache(
-                batch_size=latents.shape[0],
-                dtype=target_dtype,
-                device=latents.device)
-
-        crossattn_cache = self._initialize_crossattn_cache(
-            batch_size=latents.shape[0],
-            max_text_len=257,  # 1 CLS + 256 patch tokens
-            dtype=target_dtype,
-            device=latents.device)
-
-        if t % self.num_frame_per_block != 0:
-            raise ValueError(
-                "num_frames must be divisible by num_frame_per_block for causal denoising"
-            )
-        num_blocks = t // self.num_frame_per_block
-        block_sizes = [self.num_frame_per_block] * num_blocks
-        start_index = 0
-
-        if boundary_timestep is not None:
-            block_sizes[0] = 1
-
-        # NOTE: MatrixGame does NOT process the first frame separately.
-        # The first frame information is already encoded in batch.image_latent (cond_concat)
-        # and will be used by the model via channel concatenation: torch.cat([x, cond_concat], dim=1)
-
-        ctx = BlockProcessingContext(
-            batch=batch,
-            block_idx=0,
-            start_index=0,
-            kv_cache1=kv_cache1,
-            kv_cache2=kv_cache2,
-            kv_cache_mouse=kv_cache_mouse,
-            kv_cache_keyboard=kv_cache_keyboard,
-            crossattn_cache=crossattn_cache,
-            timesteps=timesteps,
-            block_sizes=block_sizes,
-            noise_pool=None,
-            fastvideo_args=fastvideo_args,
-            target_dtype=target_dtype,
-            autocast_enabled=autocast_enabled,
-            boundary_timestep=boundary_timestep,
-            high_noise_timesteps=high_noise_timesteps,
-            context_noise=getattr(fastvideo_args.pipeline_config,
-                                  "context_noise", 0),
-            image_kwargs=image_kwargs,
-            pos_cond_kwargs=pos_cond_kwargs,
-        )
-
-        context_noise = getattr(fastvideo_args.pipeline_config, "context_noise",
-                                0)
-
-        with self.progress_bar(total=len(block_sizes) *
-                               len(timesteps)) as progress_bar:
-            for block_idx, current_num_frames in enumerate(block_sizes):
-                ctx.block_idx = block_idx
-                ctx.start_index = start_index
-                current_latents = latents[:, :, start_index:start_index +
-                                          current_num_frames, :, :]
-
-                action_kwargs = self._prepare_action_kwargs(
-                    batch, start_index, current_num_frames)
-
-                current_latents = self._process_single_block(
-                    current_latents=current_latents,
-                    batch=batch,
-                    start_index=start_index,
-                    current_num_frames=current_num_frames,
-                    timesteps=timesteps,
-                    ctx=ctx,
-                    action_kwargs=action_kwargs,
-                    progress_bar=progress_bar,
-                )
-
-                latents[:, :, start_index:start_index +
-                        current_num_frames, :, :] = current_latents
-
-                # Update KV caches with clean context
-                self._update_context_cache(
-                    current_latents=current_latents,
-                    batch=batch,
-                    start_index=start_index,
-                    current_num_frames=current_num_frames,
-                    ctx=ctx,
-                    action_kwargs=action_kwargs,
-                    context_noise=context_noise,
-                )
-
-                start_index += current_num_frames
-
-        if boundary_timestep is not None:
-            num_frames_to_remove = self.num_frame_per_block - 1
-            if num_frames_to_remove > 0:
-                latents = latents[:, :, :-num_frames_to_remove, :, :]
-
-        batch.latents = latents
-        return batch
+        engine = DenoisingEngine(MatrixGameBlockStrategy(self),
+                                 hooks=self._build_engine_hooks())
+        return engine.run(batch, fastvideo_args)
 
     def _prepare_action_kwargs(self, batch: ForwardBatch, start_index: int,
                                num_frames: int) -> dict[str, Any]:
@@ -652,123 +497,44 @@ class MatrixGameCausalDenoisingStage(DenoisingStage):
 
     def streaming_reset(self, batch: ForwardBatch,
                         fastvideo_args: FastVideoArgs) -> ForwardBatch:
-        target_dtype = torch.bfloat16
-        autocast_enabled = (target_dtype != torch.float32
-                            ) and not fastvideo_args.disable_autocast
+        from fastvideo.pipelines.stages.denoising_engine import DenoisingEngine
+        from fastvideo.pipelines.stages.denoising_matrixgame_strategy import (
+            MatrixGameBlockStrategy)
 
-        latent_seq_length = batch.latents.shape[-1] * batch.latents.shape[-2]
-        patch_size = self.transformer.patch_size
-        patch_ratio = patch_size[-1] * patch_size[-2]
-        self.frame_seq_length = latent_seq_length // patch_ratio
+        strategy = MatrixGameBlockStrategy(self)
+        engine = DenoisingEngine(strategy, hooks=self._build_engine_hooks())
+        state = strategy.prepare(batch, fastvideo_args)
+        for hook in engine.hooks:
+            hook.on_init(engine, batch, fastvideo_args)
+        for hook in engine.hooks:
+            hook.pre_run(state)
+        block_plan = strategy.block_plan(state)
 
-        timesteps = torch.tensor(
-            fastvideo_args.pipeline_config.dmd_denoising_steps,
-            dtype=torch.long).cpu()
-        if fastvideo_args.pipeline_config.warp_denoising_step:
-            scheduler_timesteps = torch.cat((self.scheduler.timesteps.cpu(),
-                                             torch.tensor([0],
-                                                          dtype=torch.float32)))
-            timesteps = scheduler_timesteps[1000 - timesteps]
-        timesteps = timesteps.to(get_local_torch_device())
+        progress_bar = state.extra.get("progress_bar")
+        if progress_bar is not None:
+            progress_bar.close()
+        state.extra["progress_bar"] = None
 
-        boundary_ratio = getattr(fastvideo_args.pipeline_config.dit_config,
-                                 'boundary_ratio', None)
-        if boundary_ratio is not None:
-            boundary_timestep = boundary_ratio * self.scheduler.num_train_timesteps
-            high_noise_timesteps = timesteps[timesteps >= boundary_timestep]
-        else:
-            boundary_timestep = None
-            high_noise_timesteps = None
-
-        image_embeds = batch.image_embeds
-        if len(image_embeds) > 0:
-            assert torch.isnan(image_embeds[0]).sum() == 0
-            image_embeds = [
-                image_embed.to(target_dtype) for image_embed in image_embeds
-            ]
-
-        # directly set the kwarg.
-        image_kwargs = {"encoder_hidden_states_image": image_embeds}
-        pos_cond_kwargs: dict[str, Any] = {}
-
-        if st_attn_available and self.attn_backend == SlidingTileAttentionBackend:
-            self.prepare_sta_param(batch, fastvideo_args)
-
-        assert batch.latents is not None, "latents must be provided"
-        latents = batch.latents
+        ctx = state.extra["ctx"]
+        latents = state.latents
+        assert latents is not None, "latents must be provided"
         b, c, t, h, w = latents.shape
-        prompt_embeds = batch.prompt_embeds
-        assert torch.isnan(prompt_embeds[0]).sum() == 0
 
-        # Initialize caches
-        kv_cache1 = self._initialize_kv_cache(batch_size=latents.shape[0],
-                                              dtype=target_dtype,
-                                              device=latents.device)
-        kv_cache2 = None
-        if boundary_timestep is not None:
-            kv_cache2 = self._initialize_kv_cache(batch_size=latents.shape[0],
-                                                  dtype=target_dtype,
-                                                  device=latents.device)
-
-        kv_cache_mouse = None
-        kv_cache_keyboard = None
-        if self.use_action_module:
-            kv_cache_mouse, kv_cache_keyboard = self._initialize_action_kv_cache(
-                batch_size=latents.shape[0],
-                dtype=target_dtype,
-                device=latents.device)
-
-        crossattn_cache = self._initialize_crossattn_cache(
-            batch_size=latents.shape[0],
-            max_text_len=257,  # 1 CLS + 256 patch tokens
-            dtype=target_dtype,
-            device=latents.device)
-
-        # Calculate block sizes
-        if t % self.num_frame_per_block != 0:
-            raise ValueError(
-                "num_frames must be divisible by num_frame_per_block for causal denoising"
-            )
-        num_blocks = t // self.num_frame_per_block
-        block_sizes = [self.num_frame_per_block] * num_blocks
-        if boundary_timestep is not None:
-            block_sizes[0] = 1
-
-        # Pre-allocate noise pool
-        num_denoising_steps = len(timesteps)
+        num_denoising_steps = len(state.timesteps)
         noise_shape = (b, self.num_frame_per_block, c, h, w)
         noise_pool = [
             torch.randn(
                 noise_shape,
-                dtype=target_dtype,
+                dtype=ctx.target_dtype,
                 device=latents.device,
-            ) for _ in range(num_denoising_steps - 1)
+            ) for _ in range(max(num_denoising_steps - 1, 0))
         ]
+        ctx.noise_pool = noise_pool
 
-        # Create and store context
-        self._streaming_ctx = BlockProcessingContext(
-            batch=batch,
-            block_idx=0,
-            start_index=0,
-            kv_cache1=kv_cache1,
-            kv_cache2=kv_cache2,
-            kv_cache_mouse=kv_cache_mouse,
-            kv_cache_keyboard=kv_cache_keyboard,
-            crossattn_cache=crossattn_cache,
-            timesteps=timesteps,
-            block_sizes=block_sizes,
-            noise_pool=noise_pool,
-            fastvideo_args=fastvideo_args,
-            target_dtype=target_dtype,
-            autocast_enabled=autocast_enabled,
-            boundary_timestep=boundary_timestep,
-            high_noise_timesteps=high_noise_timesteps,
-            context_noise=getattr(fastvideo_args.pipeline_config,
-                                  "context_noise", 0),
-            image_kwargs=image_kwargs,
-            pos_cond_kwargs=pos_cond_kwargs,
-        )
-
+        self._streaming_engine = engine
+        self._streaming_state = state
+        self._streaming_block_plan = block_plan
+        self._streaming_ctx = ctx
         self._streaming_initialized = True
         return batch
 
@@ -776,23 +542,25 @@ class MatrixGameCausalDenoisingStage(DenoisingStage):
             self,
             keyboard_action: torch.Tensor | None = None,
             mouse_action: torch.Tensor | None = None) -> ForwardBatch:
-        if not self._streaming_initialized or self._streaming_ctx is None:
+        if (not self._streaming_initialized or self._streaming_ctx is None
+                or self._streaming_engine is None
+                or self._streaming_state is None
+                or self._streaming_block_plan is None):
             raise RuntimeError(
                 "Streaming not initialized! Call streaming_reset first.")
 
         ctx = self._streaming_ctx
-        if ctx.block_idx >= len(ctx.block_sizes):
+        block_plan = self._streaming_block_plan
+        if ctx.block_idx >= len(block_plan.items):
             return ctx.batch
 
         batch = ctx.batch
         latents = batch.latents
         assert latents is not None, "latents must be set in batch"
 
-        current_num_frames = ctx.block_sizes[ctx.block_idx]
-        start_index = ctx.start_index
-
-        current_latents = latents[:, :, start_index:start_index +
-                                  current_num_frames, :, :]
+        block_item = block_plan.items[ctx.block_idx]
+        start_index = block_item.start_index
+        current_num_frames = block_item.num_frames
 
         # Update batch with new actions for this block
         if keyboard_action is not None or mouse_action is not None:
@@ -810,58 +578,32 @@ class MatrixGameCausalDenoisingStage(DenoisingStage):
                 batch.mouse_cond[:, start_frame:start_frame +
                                  n] = mouse_action.to(batch.mouse_cond.device)
 
-        action_kwargs = self._prepare_action_kwargs(batch, start_index,
-                                                    current_num_frames)
-
-        # Create noise generator that uses pre-allocated noise pool
-        def streaming_noise_generator(shape: tuple, dtype: torch.dtype,
-                                      step_idx: int) -> torch.Tensor:
-            if ctx.noise_pool is not None and step_idx < len(ctx.noise_pool):
-                return ctx.noise_pool[step_idx][:, :shape[1], :, :, :].to(
-                    latents.device)
-            else:
-                # Fallback to dynamic allocation if pool not available
-                return torch.randn(
-                    shape,
-                    dtype=dtype,
-                    generator=(batch.generator[0] if isinstance(
-                        batch.generator, list) else batch.generator)).to(
-                            latents.device)
-
-        current_latents = self._process_single_block(
-            current_latents=current_latents,
-            batch=batch,
-            start_index=start_index,
-            current_num_frames=current_num_frames,
-            timesteps=ctx.timesteps,
-            ctx=ctx,
-            action_kwargs=action_kwargs,
-            noise_generator=streaming_noise_generator,
-        )
-
-        latents[:, :, start_index:start_index +
-                current_num_frames, :, :] = current_latents
-
-        # Update KV caches with clean context
-        self._update_context_cache(
-            current_latents=current_latents,
-            batch=batch,
-            start_index=start_index,
-            current_num_frames=current_num_frames,
-            ctx=ctx,
-            action_kwargs=action_kwargs,
-            context_noise=ctx.context_noise,
+        self._streaming_engine.run_blocks(
+            self._streaming_state,
+            block_plan=block_plan,
+            start_block=ctx.block_idx,
+            num_blocks=1,
         )
 
         # Advance streaming state
-        ctx.start_index += current_num_frames
+        ctx.start_index = start_index + current_num_frames
         ctx.block_idx += 1
 
         return batch
 
     def streaming_clear(self) -> None:
+        if (self._streaming_engine is not None
+                and self._streaming_state is not None):
+            batch = (self._streaming_ctx.batch if self._streaming_ctx
+                     is not None else self._streaming_state.extra.get("batch"))
+            if batch is not None:
+                for hook in self._streaming_engine.hooks:
+                    hook.post_run(self._streaming_state, batch)
         self._streaming_initialized = False
         self._streaming_ctx = None
+        self._streaming_engine = None
+        self._streaming_state = None
+        self._streaming_block_plan = None
 
     def verify_input(self, batch: ForwardBatch,
                      fastvideo_args: FastVideoArgs) -> VerificationResult:

fastvideo/tests/conftest.py

@@ -21,3 +21,8 @@ def distributed_setup():
     yield
 
     cleanup_dist_env_and_memory()
+
+
+def pytest_configure(config):
+    config.addinivalue_line("markers",
+                            "gpu: requires CUDA with BF16 support")

fastvideo/tests/ssim/L40S_reference_videos/TurboWan2.2-I2V-A14B-Diffusers/SLA_ATTN/steps4_An astronaut hatching from an egg, on the surface of the moon, the darkness and depth of space reali_ssim.json

@@ -0,0 +1,11 @@
+{
+  "mean_ssim": 0.9805062346988254,
+  "min_ssim": 0.9636241793632507,
+  "max_ssim": 0.9951534271240234,
+  "reference_video": "/mnt/fast-disks/hao_lab/loay/FastVideo/fastvideo/tests/ssim/L40S_reference_videos/TurboWan2.2-I2V-A14B-Diffusers/SLA_ATTN/An astronaut hatching from an egg, on the surface of the moon, the darkness and depth of space reali.mp4",
+  "generated_video": "/mnt/fast-disks/hao_lab/loay/FastVideo/fastvideo/tests/ssim/generated_videos/TurboWan2.2-I2V-A14B-Diffusers/SLA_ATTN/An astronaut hatching from an egg, on the surface of the moon, the darkness and depth of space reali.mp4",
+  "parameters": {
+    "num_inference_steps": 4,
+    "prompt": "An astronaut hatching from an egg, on the surface of the moon, the darkness and depth of space realised in the background. High quality, ultrarealistic detail and breath-taking movie-like camera shot."
+  }
+}

fastvideo/tests/ssim/generated_videos/TurboWan2.1-T2V-1.3B-Diffusers/SLA_ATTN/steps4_Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of _ssim.json

@@ -1,11 +0,0 @@
-{
-  "mean_ssim": 0.9917645101194028,
-  "min_ssim": 0.9908460974693298,
-  "max_ssim": 0.9925968050956726,
-  "reference_video": "/FastVideo/fastvideo/tests/ssim/L40S_reference_videos/TurboWan2.1-T2V-1.3B-Diffusers/SLA_ATTN/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of.mp4",
-  "generated_video": "/FastVideo/fastvideo/tests/ssim/generated_videos/TurboWan2.1-T2V-1.3B-Diffusers/SLA_ATTN/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of.mp4",
-  "parameters": {
-    "num_inference_steps": 4,
-    "prompt": "Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of a bustling street food market. The scene captures a mix of humor and authenticity. Mid-shot framing, vibrant lighting."
-  }
-}

fastvideo/tests/ssim/test_inference_similarity.py

@@ -19,6 +19,8 @@ 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 "A100" in device_name:
+    device_reference_folder = "A100" + device_reference_folder_suffix
 else:
     # device_reference_folder = "L40S" + device_reference_folder_suffix
     logger.warning(f"Unsupported device for ssim tests: {device_name}")

fastvideo/tests/ssim/test_matrixgame_similarity.py

@@ -22,6 +22,8 @@ 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 "A100" in device_name:
+    device_reference_folder = "A100" + device_reference_folder_suffix
 else:
     # device_reference_folder = "L40S" + device_reference_folder_suffix
     logger.warning(f"Unsupported device for ssim tests: {device_name}")

fastvideo/tests/ssim/test_turbodiffusion_similarity.py

@@ -24,6 +24,8 @@ 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 "A100" in device_name:
+    device_reference_folder = "A100" + device_reference_folder_suffix
 else:
     # device_reference_folder = "L40S" + device_reference_folder_suffix
     logger.warning(f"Unsupported device for ssim tests: {device_name}, using L40S references")
@@ -152,7 +154,147 @@ def test_turbodiffusion_inference_similarity(prompt, model_id):
         logger.error("Failed to write SSIM results to file")
 
     # TurboDiffusion uses fewer steps, may have slightly lower SSIM
-    min_acceptable_ssim = 0.90
+    min_acceptable_ssim = 0.95
+    assert mean_ssim >= min_acceptable_ssim, (
+        f"SSIM value {mean_ssim} is below threshold {min_acceptable_ssim} "
+        f"for {model_id} with backend {ATTENTION_BACKEND}"
+    )
+
+
+# TurboDiffusion I2V parameters (dual-model with RCM scheduler + SLA attention)
+TURBODIFFUSION_I2V_PARAMS = {
+    "num_gpus": 4,
+    "model_path": "loayrashid/TurboWan2.2-I2V-A14B-Diffusers",
+    "height": 480,
+    "width": 832,
+    "num_frames": 45,
+    "num_inference_steps": 4,  # TurboDiffusion uses 1-4 steps
+    "guidance_scale": 1.0,  # No CFG for TurboDiffusion
+    "seed": 42,
+    "sp_size": 4,
+    "tp_size": 1,
+    "fps": 24,
+}
+
+TURBODIFFUSION_I2V_MODEL_TO_PARAMS = {
+    "TurboWan2.2-I2V-A14B-Diffusers": TURBODIFFUSION_I2V_PARAMS,
+}
+
+TURBODIFFUSION_I2V_TEST_PROMPTS = [
+    "An astronaut hatching from an egg, on the surface of the moon, the darkness and depth of space realised in the background. High quality, ultrarealistic detail and breath-taking movie-like camera shot.",
+]
+
+TURBODIFFUSION_I2V_IMAGE_PATHS = [
+    "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg",
+]
+
+
+@pytest.mark.parametrize("prompt", TURBODIFFUSION_I2V_TEST_PROMPTS)
+@pytest.mark.parametrize("model_id", list(TURBODIFFUSION_I2V_MODEL_TO_PARAMS.keys()))
+def test_turbodiffusion_i2v_inference_similarity(prompt, model_id):
+    """
+    Test that runs TurboDiffusion I2V inference with dual-model switching,
+    then compares the output to reference videos using SSIM.
+    """
+    # TurboDiffusion requires SLA attention backend
+    ATTENTION_BACKEND = "SLA_ATTN"
+    os.environ["FASTVIDEO_ATTENTION_BACKEND"] = ATTENTION_BACKEND
+
+    assert len(TURBODIFFUSION_I2V_TEST_PROMPTS) == len(TURBODIFFUSION_I2V_IMAGE_PATHS), \
+        "Expect number of prompts equal to number of images"
+
+    script_dir = os.path.dirname(os.path.abspath(__file__))
+
+    base_output_dir = os.path.join(script_dir, 'generated_videos', model_id)
+    output_dir = os.path.join(base_output_dir, ATTENTION_BACKEND)
+    output_video_name = f"{prompt[:100].strip()}.mp4"
+
+    os.makedirs(output_dir, exist_ok=True)
+
+    BASE_PARAMS = TURBODIFFUSION_I2V_MODEL_TO_PARAMS[model_id]
+    num_inference_steps = BASE_PARAMS["num_inference_steps"]
+    image_path = TURBODIFFUSION_I2V_IMAGE_PATHS[TURBODIFFUSION_I2V_TEST_PROMPTS.index(prompt)]
+
+    init_kwargs = {
+        "num_gpus": BASE_PARAMS["num_gpus"],
+        "sp_size": BASE_PARAMS["sp_size"],
+        "tp_size": BASE_PARAMS["tp_size"],
+        "override_pipeline_cls_name": "TurboDiffusionI2VPipeline",
+        # Keep both transformers in VRAM - avoids CPU RAM bottleneck
+        "dit_cpu_offload": False,
+    }
+
+    generation_kwargs = {
+        "num_inference_steps": num_inference_steps,
+        "output_path": output_dir,
+        "image_path": image_path,
+        "height": BASE_PARAMS["height"],
+        "width": BASE_PARAMS["width"],
+        "num_frames": BASE_PARAMS["num_frames"],
+        "guidance_scale": BASE_PARAMS["guidance_scale"],
+        "seed": BASE_PARAMS["seed"],
+        "fps": BASE_PARAMS["fps"],
+    }
+
+    generator = VideoGenerator.from_pretrained(
+        model_path=BASE_PARAMS["model_path"],
+        **init_kwargs
+    )
+    generator.generate_video(prompt, **generation_kwargs)
+
+    if isinstance(generator.executor, MultiprocExecutor):
+        generator.executor.shutdown()
+
+    assert os.path.exists(output_dir), f"Output video was not generated at {output_dir}"
+
+    reference_folder = os.path.join(
+        script_dir, device_reference_folder, model_id, ATTENTION_BACKEND
+    )
+
+    if not os.path.exists(reference_folder):
+        logger.error("Reference folder missing")
+        raise FileNotFoundError(
+            f"Reference video folder does not exist: {reference_folder}"
+        )
+
+    # Find the matching reference video based on the prompt
+    reference_video_name = None
+
+    for filename in os.listdir(reference_folder):
+        if filename.endswith('.mp4') and prompt[:100].strip() in filename:
+            reference_video_name = filename
+            break
+
+    if not reference_video_name:
+        logger.error(
+            f"Reference video not found for prompt: {prompt} with backend: {ATTENTION_BACKEND}"
+        )
+        raise FileNotFoundError(f"Reference video missing")
+
+    reference_video_path = os.path.join(reference_folder, reference_video_name)
+    generated_video_path = os.path.join(output_dir, output_video_name)
+
+    logger.info(
+        f"Computing SSIM between {reference_video_path} and {generated_video_path}"
+    )
+    ssim_values = compute_video_ssim_torchvision(
+        reference_video_path, generated_video_path, use_ms_ssim=True
+    )
+
+    mean_ssim = ssim_values[0]
+    logger.info(f"SSIM mean value: {mean_ssim}")
+    logger.info(f"Writing SSIM results to directory: {output_dir}")
+
+    success = write_ssim_results(
+        output_dir, ssim_values, reference_video_path,
+        generated_video_path, num_inference_steps, prompt
+    )
+
+    if not success:
+        logger.error("Failed to write SSIM results to file")
+
+    # TurboDiffusion I2V uses fewer steps, may have slightly lower SSIM
+    min_acceptable_ssim = 0.95
     assert mean_ssim >= min_acceptable_ssim, (
         f"SSIM value {mean_ssim} is below threshold {min_acceptable_ssim} "
         f"for {model_id} with backend {ATTENTION_BACKEND}"

fastvideo/tests/stages/test_denoising_engine_hooks.py

@@ -0,0 +1,73 @@
+# SPDX-License-Identifier: Apache-2.0
+import torch
+
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+from fastvideo.pipelines.stages.denoising_engine import DenoisingEngine
+from fastvideo.pipelines.stages.denoising_engine_hooks import PerfLoggingHook
+from fastvideo.pipelines.stages.denoising_strategies import (ModelInputs,
+                                                             StrategyState)
+
+
+class DummyStrategy:
+    def prepare(self, batch: ForwardBatch,
+                args: FastVideoArgs) -> StrategyState:
+        latents = torch.zeros(1, 1, 1, 1, 1)
+        timesteps = torch.tensor([2, 1], dtype=torch.long)
+        return StrategyState(
+            latents=latents,
+            timesteps=timesteps,
+            num_inference_steps=len(timesteps),
+            prompt_embeds=[torch.zeros(1, 1, 1)],
+            negative_prompt_embeds=None,
+            prompt_attention_mask=None,
+            negative_attention_mask=None,
+            image_embeds=[],
+            guidance_scale=1.0,
+            guidance_scale_2=None,
+            guidance_rescale=0.0,
+            do_cfg=False,
+            extra={"batch": batch},
+        )
+
+    def make_model_inputs(self, state: StrategyState, t: torch.Tensor,
+                          step_idx: int) -> ModelInputs:
+        return ModelInputs(
+            latent_model_input=state.latents,
+            timestep=t,
+            prompt_embeds=state.prompt_embeds,
+            prompt_attention_mask=state.prompt_attention_mask,
+        )
+
+    def forward(self, state: StrategyState,
+                model_inputs: ModelInputs) -> torch.Tensor:
+        return state.latents
+
+    def cfg_combine(self, state: StrategyState,
+                    noise_pred: torch.Tensor) -> torch.Tensor:
+        return noise_pred
+
+    def scheduler_step(self, state: StrategyState, noise_pred: torch.Tensor,
+                       t: torch.Tensor) -> torch.Tensor:
+        return state.latents
+
+    def postprocess(self, state: StrategyState) -> ForwardBatch:
+        return state.extra["batch"]
+
+
+def test_perf_logging_hook_records_steps():
+    batch = ForwardBatch(data_type="test")
+    args = FastVideoArgs(model_path="dummy")
+    hook = PerfLoggingHook()
+    engine = DenoisingEngine(DummyStrategy(), hooks=[hook])
+
+    engine.run(batch, args)
+
+    info = batch.logging_info.get_stage_info("DenoisingEngine")
+    step_times = info.get("denoise_step_times_ms")
+    total_ms = info.get("denoise_total_ms")
+
+    assert isinstance(step_times, list)
+    assert len(step_times) == 2
+    assert total_ms is not None
+    assert total_ms >= 0.0

fastvideo/tests/stages/test_denoising_engine_parity.py

@@ -0,0 +1,811 @@
+# SPDX-License-Identifier: Apache-2.0
+import contextlib
+from types import SimpleNamespace
+
+import torch
+import pytest
+
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+from fastvideo.pipelines.stages.denoising_engine import DenoisingEngine
+from fastvideo.pipelines.stages.denoising_strategies import (
+    BlockDenoisingStrategy)
+from fastvideo.pipelines.stages.denoising_standard_strategy import (
+    StandardStrategy)
+from fastvideo.pipelines.stages.denoising_cosmos_strategy import (
+    CosmosStrategy)
+from fastvideo.pipelines.stages.denoising_dmd_strategy import DmdStrategy
+from fastvideo.pipelines.stages.denoising_longcat_strategy import (
+    LongCatStrategy,
+    LongCatI2VStrategy,
+    LongCatVCStrategy,
+)
+from fastvideo.pipelines.stages.denoising_causal_strategy import (
+    CausalBlockStrategy)
+from fastvideo.pipelines.stages.denoising_matrixgame_strategy import (
+    MatrixGameBlockStrategy)
+
+
+class DummyProgressBar:
+    def __init__(self):
+        self.updates = 0
+
+    def update(self):
+        self.updates += 1
+
+    def close(self):
+        return None
+
+
+class DummyScheduler:
+    def __init__(self, timesteps: torch.Tensor):
+        self.order = 1
+        self.timesteps = timesteps
+        self.sigmas = self._make_sigmas(timesteps)
+        self.num_train_timesteps = int(timesteps.max().item()
+                                       ) if timesteps.numel() else 1
+        self.init_noise_sigma = 1.0
+        self.config = SimpleNamespace(final_sigmas_type=None)
+
+    def _make_sigmas(self, timesteps: torch.Tensor) -> torch.Tensor:
+        if timesteps.numel() == 0:
+            return timesteps
+        max_val = float(timesteps.max().item()) or 1.0
+        return timesteps.float() / max_val
+
+    def register_to_config(self, **kwargs):
+        for key, value in kwargs.items():
+            setattr(self.config, key, value)
+
+    def set_timesteps(self, num_steps: int, device=None):
+        self.timesteps = torch.linspace(1.0, 0.0, num_steps, device=device)
+        self.sigmas = self._make_sigmas(self.timesteps)
+        return self.timesteps
+
+    def scale_model_input(self, latents: torch.Tensor,
+                          t: torch.Tensor) -> torch.Tensor:
+        return latents
+
+    def step(self, noise_pred: torch.Tensor, t: torch.Tensor,
+             latents: torch.Tensor, **kwargs):
+        return (latents - noise_pred * 0.1, )
+
+    def add_noise(self, latents: torch.Tensor, noise: torch.Tensor,
+                  t: torch.Tensor) -> torch.Tensor:
+        return latents + noise * 0.01
+
+    def add_noise_high(self, latents: torch.Tensor, noise: torch.Tensor,
+                       t: torch.Tensor, boundary_timestep: torch.Tensor
+                       ) -> torch.Tensor:
+        return latents + noise * 0.01
+
+
+class DummyTransformer(torch.nn.Module):
+    def __init__(self,
+                 hidden_size: int = 8,
+                 num_attention_heads: int = 2,
+                 patch_size: tuple[int, int, int] = (1, 1, 1),
+                 num_frames_per_block: int = 1,
+                 sliding_window_num_frames: int = 1,
+                 local_attn_size: int = 4) -> None:
+        super().__init__()
+        self.dummy = torch.nn.Parameter(torch.zeros(1))
+        self.hidden_size = hidden_size
+        self.num_attention_heads = num_attention_heads
+        self.attention_head_dim = hidden_size // num_attention_heads
+        self.blocks = [object(), object()]
+        self.config = SimpleNamespace(
+            arch_config=SimpleNamespace(
+                patch_size=patch_size,
+                num_frames_per_block=num_frames_per_block,
+                sliding_window_num_frames=sliding_window_num_frames,
+            ))
+        self.patch_size = patch_size
+        self.local_attn_size = local_attn_size
+        self.model = SimpleNamespace(local_attn_size=local_attn_size)
+        self.independent_first_frame = False
+
+    def forward(
+        self,
+        hidden_states=None,
+        encoder_hidden_states=None,
+        timestep=None,
+        guidance=None,
+        encoder_hidden_states_image=None,
+        encoder_hidden_states_2=None,
+        encoder_attention_mask=None,
+        mask_strategy=None,
+        mouse_cond=None,
+        keyboard_cond=None,
+        kv_cache=None,
+        crossattn_cache=None,
+        current_start=None,
+        start_frame=None,
+        num_cond_latents=None,
+        *args,
+        **kwargs,
+    ):
+        latent = hidden_states
+        if latent is None and args:
+            latent = args[0]
+        if latent is None:
+            raise ValueError("Missing latent input")
+        t_val = timestep
+        if t_val is None and len(args) > 2:
+            t_val = args[2]
+        scale = float(torch.as_tensor(t_val).float().mean().item()
+                      ) if t_val is not None else 0.0
+        out = latent * 0.01 + scale
+        out_channels = getattr(self, "out_channels", None)
+        if out_channels is not None and out.shape[1] != out_channels:
+            out = out[:, :out_channels, ...]
+        if kwargs.get("return_dict") is False:
+            return (out, )
+        return out
+
+
+class DummyStandardStage:
+    def __init__(self, transformer, scheduler):
+        self.transformer = transformer
+        self.transformer_2 = None
+        self.scheduler = scheduler
+        self.vae = None
+        self.pipeline = None
+        self.attn_backend = object()
+
+    def prepare_extra_func_kwargs(self, func, kwargs):
+        extra_step_kwargs = {}
+        for key, value in kwargs.items():
+            if key in set(func.__code__.co_varnames):
+                extra_step_kwargs[key] = value
+        return extra_step_kwargs
+
+    def progress_bar(self, iterable=None, total=None):
+        return DummyProgressBar()
+
+    def rescale_noise_cfg(self, noise_cfg, noise_pred_text,
+                          guidance_rescale=0.0) -> torch.Tensor:
+        std_text = noise_pred_text.std(dim=list(range(1, noise_pred_text.ndim)),
+                                       keepdim=True)
+        std_cfg = noise_cfg.std(dim=list(range(1, noise_cfg.ndim)),
+                                keepdim=True)
+        noise_pred_rescaled = noise_cfg * (std_text / std_cfg)
+        return (guidance_rescale * noise_pred_rescaled +
+                (1 - guidance_rescale) * noise_cfg)
+
+    def prepare_sta_param(self, batch, fastvideo_args):
+        return None
+
+    def save_sta_search_results(self, batch):
+        return None
+
+
+class DummyLongCatStage:
+    def __init__(self, transformer, scheduler):
+        self.transformer = transformer
+        self.scheduler = scheduler
+
+    def optimized_scale(self, positive_flat, negative_flat) -> torch.Tensor:
+        dot_product = torch.sum(positive_flat * negative_flat,
+                                dim=1,
+                                keepdim=True)
+        squared_norm = torch.sum(negative_flat**2, dim=1, keepdim=True) + 1e-8
+        return dot_product / squared_norm
+
+
+class DummyDmdStage(DummyStandardStage):
+    @property
+    def device(self) -> torch.device:
+        return torch.device("cpu")
+
+
+class DummyCausalStage:
+    def __init__(self, transformer, scheduler):
+        self.transformer = transformer
+        self.transformer_2 = None
+        self.scheduler = scheduler
+        self.vae = None
+        self.num_transformer_blocks = len(transformer.blocks)
+        self.num_frames_per_block = transformer.config.arch_config.num_frames_per_block
+        self.sliding_window_num_frames = (
+            transformer.config.arch_config.sliding_window_num_frames)
+        self.local_attn_size = transformer.model.local_attn_size
+        self.attn_backend = object()
+        self.frame_seq_length = 1
+
+    @property
+    def device(self) -> torch.device:
+        return next(self.transformer.parameters()).device
+
+    def prepare_extra_func_kwargs(self, func, kwargs):
+        extra_step_kwargs = {}
+        for key, value in kwargs.items():
+            if key in set(func.__code__.co_varnames):
+                extra_step_kwargs[key] = value
+        return extra_step_kwargs
+
+    def prepare_sta_param(self, batch, fastvideo_args):
+        return None
+
+    def progress_bar(self, iterable=None, total=None):
+        return DummyProgressBar()
+
+    def _initialize_kv_cache(self, batch_size, dtype, device):
+        return [{} for _ in range(self.num_transformer_blocks)]
+
+    def _initialize_crossattn_cache(self, batch_size, max_text_len, dtype,
+                                    device):
+        return [{} for _ in range(self.num_transformer_blocks)]
+
+
+class DummyMatrixGameStage:
+    def __init__(self, transformer, scheduler):
+        self.transformer = transformer
+        self.transformer_2 = None
+        self.scheduler = scheduler
+        self.vae = None
+        self.num_transformer_blocks = len(transformer.blocks)
+        self.num_frame_per_block = transformer.config.arch_config.num_frames_per_block
+        self.sliding_window_num_frames = (
+            transformer.config.arch_config.sliding_window_num_frames)
+        self.local_attn_size = transformer.local_attn_size
+        self.use_action_module = False
+        self.attn_backend = object()
+        self.frame_seq_length = 1
+
+    @property
+    def device(self) -> torch.device:
+        return torch.device("cpu")
+
+    def progress_bar(self, iterable=None, total=None):
+        return DummyProgressBar()
+
+    def prepare_sta_param(self, batch, fastvideo_args):
+        return None
+
+    def _initialize_kv_cache(self, batch_size, dtype, device):
+        return [{} for _ in range(self.num_transformer_blocks)]
+
+    def _initialize_action_kv_cache(self, batch_size, dtype, device):
+        return None, None
+
+    def _initialize_crossattn_cache(self, batch_size, max_text_len, dtype,
+                                    device):
+        return [{} for _ in range(self.num_transformer_blocks)]
+
+    def _prepare_action_kwargs(self, batch: ForwardBatch, start_index: int,
+                               num_frames: int) -> dict:
+        return {}
+
+    def _process_single_block(self, current_latents: torch.Tensor,
+                              batch: ForwardBatch, start_index: int,
+                              current_num_frames: int, timesteps: torch.Tensor,
+                              ctx, action_kwargs, noise_generator=None,
+                              progress_bar=None):
+        latents = current_latents
+        for t_cur in timesteps:
+            noise_pred = self.transformer(latents, batch.prompt_embeds, t_cur)
+            latents = self.scheduler.step(noise_pred,
+                                          t_cur,
+                                          latents,
+                                          return_dict=False)[0]
+            if progress_bar is not None:
+                progress_bar.update()
+        return latents
+
+    def _update_context_cache(self, current_latents: torch.Tensor,
+                              batch: ForwardBatch, start_index: int,
+                              current_num_frames: int, ctx, action_kwargs,
+                              context_noise: float):
+        return None
+
+
+@pytest.fixture(autouse=True)
+def _patch_autocast(monkeypatch, request):
+    if request.node.get_closest_marker("gpu"):
+        return
+
+    @contextlib.contextmanager
+    def _autocast(*args, **kwargs):
+        yield
+
+    monkeypatch.setattr(torch, "autocast", _autocast)
+
+
+@pytest.fixture(autouse=True)
+def _patch_local_device(monkeypatch, request):
+    if request.node.get_closest_marker("gpu"):
+        return
+
+    def _cpu_device():
+        return torch.device("cpu")
+
+    import fastvideo.distributed as dist
+    import fastvideo.pipelines.stages.denoising_standard_strategy as standard
+    import fastvideo.pipelines.stages.denoising_dmd_strategy as dmd
+    import fastvideo.pipelines.stages.denoising_causal_strategy as causal
+    import fastvideo.pipelines.stages.denoising_matrixgame_strategy as matrixgame
+
+    monkeypatch.setattr(dist, "get_local_torch_device", _cpu_device)
+    monkeypatch.setattr(standard, "get_local_torch_device", _cpu_device)
+    monkeypatch.setattr(dmd, "get_local_torch_device", _cpu_device)
+    monkeypatch.setattr(causal, "get_local_torch_device", _cpu_device)
+    monkeypatch.setattr(matrixgame, "get_local_torch_device", _cpu_device)
+
+
+def _make_args() -> FastVideoArgs:
+    args = FastVideoArgs(model_path="dummy")
+    args.pipeline_config.dit_config.arch_config.patch_size = (1, 1, 1)
+    args.pipeline_config.dit_config.boundary_ratio = None
+    args.pipeline_config.ti2v_task = False
+    args.pipeline_config.embedded_cfg_scale = None
+    return args
+
+
+def _make_batch(latents: torch.Tensor,
+                timesteps: torch.Tensor,
+                guidance_scale: float,
+                seed: int,
+                device: torch.device | str | None = None,
+                generator_device: torch.device | str = "cpu") -> ForwardBatch:
+    if device is None:
+        device = latents.device
+    embed_generator = torch.Generator(device=device).manual_seed(seed)
+    batch_generator = torch.Generator(device=generator_device).manual_seed(seed)
+    batch = ForwardBatch(data_type="test")
+    batch.latents = latents.clone()
+    batch.timesteps = timesteps.clone()
+    batch.num_inference_steps = len(timesteps)
+    batch.num_frames = latents.shape[2]
+    batch.height = latents.shape[-2]
+    batch.width = latents.shape[-1]
+    batch.raw_latent_shape = latents.shape
+    batch.prompt_embeds = [
+        torch.randn(1, 4, 8, generator=embed_generator, device=device)
+    ]
+    batch.negative_prompt_embeds = [
+        torch.randn(1, 4, 8, generator=embed_generator, device=device)
+    ]
+    batch.prompt_attention_mask = [
+        torch.ones(1, 4, dtype=torch.long, device=device)
+    ]
+    batch.negative_attention_mask = [
+        torch.ones(1, 4, dtype=torch.long, device=device)
+    ]
+    batch.guidance_scale = guidance_scale
+    batch.guidance_scale_2 = guidance_scale
+    batch.do_classifier_free_guidance = guidance_scale > 1.0
+    batch.guidance_rescale = 0.0
+    batch.generator = batch_generator
+    batch.image_embeds = []
+    batch.clip_embedding_pos = [
+        torch.randn(1, 4, 8, generator=embed_generator, device=device)
+    ]
+    batch.clip_embedding_neg = [
+        torch.randn(1, 4, 8, generator=embed_generator, device=device)
+    ]
+    return batch
+
+
+def _run_manual(strategy, batch, args):
+    state = strategy.prepare(batch, args)
+    if isinstance(strategy, BlockDenoisingStrategy):
+        block_plan = strategy.block_plan(state)
+        for block_idx, block_item in enumerate(block_plan.items):
+            block_ctx = strategy.init_block_context(state, block_item, block_idx)
+            strategy.process_block(state, block_ctx, block_item)
+            strategy.update_context(state, block_ctx, block_item)
+    else:
+        for i, t in enumerate(state.timesteps):
+            model_inputs = strategy.make_model_inputs(state, t, i)
+            noise_pred = strategy.forward(state, model_inputs)
+            noise_pred = strategy.cfg_combine(state, noise_pred)
+            state.latents = strategy.scheduler_step(state, noise_pred, t)
+    return strategy.postprocess(state)
+
+
+def _cuda_bf16_available() -> bool:
+    if not torch.cuda.is_available():
+        return False
+    is_bf16_supported = getattr(torch.cuda, "is_bf16_supported", None)
+    if is_bf16_supported is None:
+        return False
+    return torch.cuda.is_bf16_supported()
+
+
+def test_standard_strategy_parity_cpu():
+    timesteps = torch.tensor([2, 1], dtype=torch.long)
+    latents = torch.randn(1, 2, 2, 2, 2)
+    args = _make_args()
+    scheduler = DummyScheduler(timesteps)
+    stage_a = DummyStandardStage(DummyTransformer(), scheduler)
+    stage_b = DummyStandardStage(DummyTransformer(), DummyScheduler(timesteps))
+
+    batch_a = _make_batch(latents, timesteps, guidance_scale=2.0, seed=0)
+    batch_b = _make_batch(latents, timesteps, guidance_scale=2.0, seed=0)
+
+    engine_out = DenoisingEngine(StandardStrategy(stage_a)).run(batch_a, args)
+    manual_out = _run_manual(StandardStrategy(stage_b), batch_b, args)
+
+    assert torch.allclose(engine_out.latents, manual_out.latents)
+
+
+def test_standard_strategy_parity_i2v_cpu():
+    timesteps = torch.tensor([2, 1], dtype=torch.long)
+    latents = torch.randn(1, 2, 2, 2, 2)
+    args = _make_args()
+    scheduler = DummyScheduler(timesteps)
+    stage_a = DummyStandardStage(DummyTransformer(), scheduler)
+    stage_b = DummyStandardStage(DummyTransformer(), DummyScheduler(timesteps))
+    stage_a.transformer.out_channels = latents.shape[1]
+    stage_b.transformer.out_channels = latents.shape[1]
+
+    batch_a = _make_batch(latents, timesteps, guidance_scale=2.0, seed=5)
+    batch_b = _make_batch(latents, timesteps, guidance_scale=2.0, seed=5)
+
+    cond_gen = torch.Generator().manual_seed(123)
+    image_latent = torch.randn(latents.shape,
+                               generator=cond_gen,
+                               device=latents.device,
+                               dtype=latents.dtype)
+    image_embeds = [
+        torch.randn(1,
+                    4,
+                    8,
+                    generator=cond_gen,
+                    device=latents.device,
+                    dtype=latents.dtype)
+    ]
+
+    batch_a.image_latent = image_latent
+    batch_b.image_latent = image_latent.clone()
+    batch_a.image_embeds = image_embeds
+    batch_b.image_embeds = [image_embeds[0].clone()]
+
+    engine_out = DenoisingEngine(StandardStrategy(stage_a)).run(batch_a, args)
+    manual_out = _run_manual(StandardStrategy(stage_b), batch_b, args)
+
+    assert torch.allclose(engine_out.latents, manual_out.latents)
+
+
+def test_standard_strategy_parity_v2v_cpu():
+    timesteps = torch.tensor([2, 1], dtype=torch.long)
+    latents = torch.randn(1, 2, 2, 2, 2)
+    args = _make_args()
+    scheduler = DummyScheduler(timesteps)
+    stage_a = DummyStandardStage(DummyTransformer(), scheduler)
+    stage_b = DummyStandardStage(DummyTransformer(), DummyScheduler(timesteps))
+    stage_a.transformer.out_channels = latents.shape[1]
+    stage_b.transformer.out_channels = latents.shape[1]
+
+    batch_a = _make_batch(latents, timesteps, guidance_scale=2.0, seed=6)
+    batch_b = _make_batch(latents, timesteps, guidance_scale=2.0, seed=6)
+
+    cond_gen = torch.Generator().manual_seed(456)
+    video_latent = torch.randn(latents.shape,
+                               generator=cond_gen,
+                               device=latents.device,
+                               dtype=latents.dtype)
+
+    batch_a.video_latent = video_latent
+    batch_b.video_latent = video_latent.clone()
+
+    engine_out = DenoisingEngine(StandardStrategy(stage_a)).run(batch_a, args)
+    manual_out = _run_manual(StandardStrategy(stage_b), batch_b, args)
+
+    assert torch.allclose(engine_out.latents, manual_out.latents)
+
+
+def test_standard_strategy_parity_boundary_cpu():
+    timesteps = torch.tensor([2, 1], dtype=torch.long)
+    latents = torch.randn(1, 2, 2, 2, 2)
+    args = _make_args()
+    args.pipeline_config.dit_config.boundary_ratio = 0.75
+    scheduler = DummyScheduler(timesteps)
+    stage_a = DummyStandardStage(DummyTransformer(), scheduler)
+    stage_b = DummyStandardStage(DummyTransformer(), DummyScheduler(timesteps))
+    stage_a.transformer_2 = DummyTransformer()
+    stage_b.transformer_2 = DummyTransformer()
+
+    batch_a = _make_batch(latents, timesteps, guidance_scale=2.0, seed=7)
+    batch_b = _make_batch(latents, timesteps, guidance_scale=2.0, seed=7)
+    batch_a.guidance_scale_2 = 1.5
+    batch_b.guidance_scale_2 = 1.5
+
+    engine_out = DenoisingEngine(StandardStrategy(stage_a)).run(batch_a, args)
+    manual_out = _run_manual(StandardStrategy(stage_b), batch_b, args)
+
+    assert torch.allclose(engine_out.latents, manual_out.latents)
+
+
+def test_cosmos_strategy_parity_cpu():
+    timesteps = torch.tensor([2, 1], dtype=torch.long)
+    latents = torch.randn(1, 2, 2, 2, 2)
+    args = _make_args()
+    scheduler = DummyScheduler(timesteps)
+    stage_a = DummyStandardStage(DummyTransformer(), scheduler)
+    stage_b = DummyStandardStage(DummyTransformer(), DummyScheduler(timesteps))
+
+    batch_a = _make_batch(latents, timesteps, guidance_scale=2.0, seed=1)
+    batch_b = _make_batch(latents, timesteps, guidance_scale=2.0, seed=1)
+
+    engine_out = DenoisingEngine(CosmosStrategy(stage_a)).run(batch_a, args)
+    manual_out = _run_manual(CosmosStrategy(stage_b), batch_b, args)
+
+    assert torch.allclose(engine_out.latents, manual_out.latents)
+
+
+def test_dmd_strategy_parity_cpu():
+    timesteps = torch.tensor([4, 2], dtype=torch.long)
+    latents = torch.randn(1, 2, 2, 2, 2)
+    args = _make_args()
+    args.pipeline_config.dmd_denoising_steps = [4, 2]
+    scheduler = DummyScheduler(timesteps)
+    scheduler.timesteps = timesteps
+    scheduler.sigmas = scheduler._make_sigmas(timesteps)
+    stage_a = DummyDmdStage(DummyTransformer(), scheduler)
+    stage_b = DummyDmdStage(DummyTransformer(),
+                            DummyScheduler(timesteps))
+
+    batch_a = _make_batch(latents, timesteps, guidance_scale=1.0, seed=2)
+    batch_b = _make_batch(latents, timesteps, guidance_scale=1.0, seed=2)
+
+    engine_out = DenoisingEngine(DmdStrategy(stage_a)).run(batch_a, args)
+    manual_out = _run_manual(DmdStrategy(stage_b), batch_b, args)
+
+    assert torch.allclose(engine_out.latents, manual_out.latents)
+
+
+def test_dmd_strategy_parity_i2v_cpu():
+    timesteps = torch.tensor([4, 2], dtype=torch.long)
+    latents = torch.randn(1, 2, 2, 2, 2)
+    args = _make_args()
+    args.pipeline_config.dmd_denoising_steps = [4, 2]
+    scheduler = DummyScheduler(timesteps)
+    scheduler.timesteps = timesteps
+    scheduler.sigmas = scheduler._make_sigmas(timesteps)
+    stage_a = DummyDmdStage(DummyTransformer(), scheduler)
+    stage_b = DummyDmdStage(DummyTransformer(), DummyScheduler(timesteps))
+    stage_a.transformer.out_channels = latents.shape[1]
+    stage_b.transformer.out_channels = latents.shape[1]
+
+    batch_a = _make_batch(latents, timesteps, guidance_scale=1.0, seed=8)
+    batch_b = _make_batch(latents, timesteps, guidance_scale=1.0, seed=8)
+
+    cond_gen = torch.Generator().manual_seed(789)
+    image_latent = torch.randn(latents.shape,
+                               generator=cond_gen,
+                               device=latents.device,
+                               dtype=latents.dtype)
+    image_embeds = [
+        torch.randn(1,
+                    4,
+                    8,
+                    generator=cond_gen,
+                    device=latents.device,
+                    dtype=latents.dtype)
+    ]
+
+    batch_a.image_latent = image_latent
+    batch_b.image_latent = image_latent.clone()
+    batch_a.image_embeds = image_embeds
+    batch_b.image_embeds = [image_embeds[0].clone()]
+
+    engine_out = DenoisingEngine(DmdStrategy(stage_a)).run(batch_a, args)
+    manual_out = _run_manual(DmdStrategy(stage_b), batch_b, args)
+
+    assert torch.allclose(engine_out.latents, manual_out.latents)
+
+
+def test_longcat_strategy_parity_cpu():
+    timesteps = torch.tensor([2, 1], dtype=torch.long)
+    latents = torch.randn(1, 2, 2, 2, 2)
+    args = _make_args()
+    scheduler = DummyScheduler(timesteps)
+    stage_a = DummyLongCatStage(DummyTransformer(), scheduler)
+    stage_b = DummyLongCatStage(DummyTransformer(), DummyScheduler(timesteps))
+
+    batch_a = _make_batch(latents, timesteps, guidance_scale=2.0, seed=3)
+    batch_b = _make_batch(latents, timesteps, guidance_scale=2.0, seed=3)
+
+    engine_out = DenoisingEngine(LongCatStrategy(stage_a)).run(batch_a, args)
+    manual_out = _run_manual(LongCatStrategy(stage_b), batch_b, args)
+
+    assert torch.allclose(engine_out.latents, manual_out.latents)
+
+
+def test_longcat_i2v_strategy_parity_cpu():
+    timesteps = torch.tensor([2, 1], dtype=torch.long)
+    latents = torch.randn(1, 2, 3, 2, 2)
+    args = _make_args()
+    scheduler = DummyScheduler(timesteps)
+    stage_a = DummyLongCatStage(DummyTransformer(), scheduler)
+    stage_b = DummyLongCatStage(DummyTransformer(), DummyScheduler(timesteps))
+
+    batch_a = _make_batch(latents, timesteps, guidance_scale=2.0, seed=4)
+    batch_b = _make_batch(latents, timesteps, guidance_scale=2.0, seed=4)
+    batch_a.num_cond_latents = 1
+    batch_b.num_cond_latents = 1
+
+    engine_out = DenoisingEngine(LongCatI2VStrategy(stage_a)).run(batch_a,
+                                                                  args)
+    manual_out = _run_manual(LongCatI2VStrategy(stage_b), batch_b, args)
+
+    assert torch.allclose(engine_out.latents, manual_out.latents)
+
+
+def test_longcat_vc_strategy_parity_cpu():
+    timesteps = torch.tensor([2, 1], dtype=torch.long)
+    latents = torch.randn(1, 2, 3, 2, 2)
+    args = _make_args()
+    scheduler = DummyScheduler(timesteps)
+    stage_a = DummyLongCatStage(DummyTransformer(), scheduler)
+    stage_b = DummyLongCatStage(DummyTransformer(), DummyScheduler(timesteps))
+
+    batch_a = _make_batch(latents, timesteps, guidance_scale=2.0, seed=5)
+    batch_b = _make_batch(latents, timesteps, guidance_scale=2.0, seed=5)
+    batch_a.num_cond_latents = 1
+    batch_b.num_cond_latents = 1
+    batch_a.use_kv_cache = True
+    batch_b.use_kv_cache = True
+    batch_a.kv_cache_dict = {}
+    batch_b.kv_cache_dict = {}
+    batch_a.cond_latents = torch.randn(1, 2, 1, 2, 2)
+    batch_b.cond_latents = batch_a.cond_latents.clone()
+
+    engine_out = DenoisingEngine(LongCatVCStrategy(stage_a)).run(batch_a,
+                                                                 args)
+    manual_out = _run_manual(LongCatVCStrategy(stage_b), batch_b, args)
+
+    assert torch.allclose(engine_out.latents, manual_out.latents)
+
+
+def test_causal_block_strategy_parity_cpu():
+    timesteps = torch.tensor([4, 2], dtype=torch.long)
+    latents = torch.randn(1, 2, 2, 2, 2)
+    args = _make_args()
+    args.pipeline_config.dmd_denoising_steps = [4, 2]
+    scheduler = DummyScheduler(timesteps)
+    scheduler.timesteps = timesteps
+    scheduler.sigmas = scheduler._make_sigmas(timesteps)
+    stage_a = DummyCausalStage(DummyTransformer(num_frames_per_block=1),
+                               scheduler)
+    stage_b = DummyCausalStage(DummyTransformer(num_frames_per_block=1),
+                               DummyScheduler(timesteps))
+
+    batch_a = _make_batch(latents, timesteps, guidance_scale=1.0, seed=6)
+    batch_b = _make_batch(latents, timesteps, guidance_scale=1.0, seed=6)
+
+    engine_out = DenoisingEngine(CausalBlockStrategy(stage_a)).run(batch_a,
+                                                                   args)
+    manual_out = _run_manual(CausalBlockStrategy(stage_b), batch_b, args)
+
+    assert torch.allclose(engine_out.latents, manual_out.latents)
+
+
+def test_matrixgame_block_strategy_parity_cpu():
+    timesteps = torch.tensor([3, 1], dtype=torch.long)
+    latents = torch.randn(1, 2, 2, 2, 2)
+    args = _make_args()
+    args.pipeline_config.dmd_denoising_steps = [3, 1]
+    scheduler = DummyScheduler(timesteps)
+    scheduler.timesteps = timesteps
+    scheduler.sigmas = scheduler._make_sigmas(timesteps)
+    stage_a = DummyMatrixGameStage(DummyTransformer(num_frames_per_block=1),
+                                   scheduler)
+    stage_b = DummyMatrixGameStage(DummyTransformer(num_frames_per_block=1),
+                                   DummyScheduler(timesteps))
+
+    batch_a = _make_batch(latents, timesteps, guidance_scale=1.0, seed=7)
+    batch_b = _make_batch(latents, timesteps, guidance_scale=1.0, seed=7)
+
+    engine_out = DenoisingEngine(MatrixGameBlockStrategy(stage_a)).run(
+        batch_a, args)
+    manual_out = _run_manual(MatrixGameBlockStrategy(stage_b), batch_b, args)
+
+    assert torch.allclose(engine_out.latents, manual_out.latents)
+
+
+def test_matrixgame_block_strategy_streaming_parity_cpu():
+    timesteps = torch.tensor([3, 1], dtype=torch.long)
+    latents = torch.randn(1, 2, 2, 2, 2)
+    args = _make_args()
+    args.pipeline_config.dmd_denoising_steps = [3, 1]
+    scheduler = DummyScheduler(timesteps)
+    scheduler.timesteps = timesteps
+    scheduler.sigmas = scheduler._make_sigmas(timesteps)
+
+    stage_full = DummyMatrixGameStage(DummyTransformer(num_frames_per_block=1),
+                                      scheduler)
+    stage_stream = DummyMatrixGameStage(
+        DummyTransformer(num_frames_per_block=1), DummyScheduler(timesteps))
+
+    batch_full = _make_batch(latents, timesteps, guidance_scale=1.0, seed=10)
+    batch_stream = _make_batch(latents, timesteps, guidance_scale=1.0, seed=10)
+    batch_full.generator = None
+    batch_stream.generator = None
+
+    torch.manual_seed(1234)
+    noise_shape = (latents.shape[0], stage_full.num_frame_per_block,
+                   latents.shape[1], latents.shape[3], latents.shape[4])
+    noise_pool = [
+        torch.randn(noise_shape, dtype=latents.dtype)
+        for _ in range(max(len(timesteps) - 1, 0))
+    ]
+
+    strategy_full = MatrixGameBlockStrategy(stage_full)
+    state_full = strategy_full.prepare(batch_full, args)
+    state_full.extra["ctx"].noise_pool = [t.clone() for t in noise_pool]
+    block_plan_full = strategy_full.block_plan(state_full)
+
+    strategy_stream = MatrixGameBlockStrategy(stage_stream)
+    state_stream = strategy_stream.prepare(batch_stream, args)
+    state_stream.extra["ctx"].noise_pool = [t.clone() for t in noise_pool]
+    block_plan_stream = strategy_stream.block_plan(state_stream)
+
+    engine_full = DenoisingEngine(strategy_full)
+    engine_full.run_blocks(state_full, block_plan=block_plan_full)
+
+    engine_stream = DenoisingEngine(strategy_stream)
+    for block_idx in range(len(block_plan_stream.items)):
+        engine_stream.run_blocks(
+            state_stream,
+            block_plan=block_plan_stream,
+            start_block=block_idx,
+            num_blocks=1,
+        )
+
+    batch_full = strategy_full.postprocess(state_full)
+    batch_stream = strategy_stream.postprocess(state_stream)
+
+    assert torch.allclose(batch_full.latents, batch_stream.latents)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not _cuda_bf16_available(),
+                    reason="CUDA BF16 not available")
+def test_standard_strategy_parity_cuda():
+    device = torch.device("cuda")
+    timesteps = torch.tensor([2, 1], dtype=torch.long, device=device)
+    latents = torch.randn(1, 2, 2, 2, 2, device=device)
+    args = _make_args()
+    scheduler = DummyScheduler(timesteps)
+    stage_a = DummyStandardStage(DummyTransformer().to(device), scheduler)
+    stage_b = DummyStandardStage(DummyTransformer().to(device),
+                                 DummyScheduler(timesteps))
+
+    batch_a = _make_batch(latents, timesteps, guidance_scale=2.0, seed=8)
+    batch_b = _make_batch(latents, timesteps, guidance_scale=2.0, seed=8)
+
+    engine_out = DenoisingEngine(StandardStrategy(stage_a)).run(batch_a, args)
+    manual_out = _run_manual(StandardStrategy(stage_b), batch_b, args)
+
+    assert torch.allclose(engine_out.latents, manual_out.latents)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not _cuda_bf16_available(),
+                    reason="CUDA BF16 not available")
+def test_causal_block_strategy_parity_cuda():
+    device = torch.device("cuda")
+    timesteps = torch.tensor([4, 2], dtype=torch.long, device=device)
+    latents = torch.randn(1, 2, 2, 2, 2, device=device)
+    args = _make_args()
+    args.pipeline_config.dmd_denoising_steps = [4, 2]
+    scheduler = DummyScheduler(timesteps)
+    scheduler.timesteps = timesteps
+    scheduler.sigmas = scheduler._make_sigmas(timesteps)
+    stage_a = DummyCausalStage(
+        DummyTransformer(num_frames_per_block=1).to(device), scheduler)
+    stage_b = DummyCausalStage(
+        DummyTransformer(num_frames_per_block=1).to(device),
+        DummyScheduler(timesteps))
+
+    batch_a = _make_batch(latents, timesteps, guidance_scale=1.0, seed=9)
+    batch_b = _make_batch(latents, timesteps, guidance_scale=1.0, seed=9)
+
+    engine_out = DenoisingEngine(CausalBlockStrategy(stage_a)).run(batch_a,
+                                                                   args)
+    manual_out = _run_manual(CausalBlockStrategy(stage_b), batch_b, args)
+
+    assert torch.allclose(engine_out.latents, manual_out.latents)

fastvideo/tests/stages/test_scheduler_adapter.py

@@ -0,0 +1,52 @@
+# SPDX-License-Identifier: Apache-2.0
+import torch
+
+from fastvideo.models.schedulers.adapter import DefaultSchedulerAdapter
+
+
+class DummyScheduler:
+    def __init__(self) -> None:
+        self.calls: dict[str, tuple] = {}
+
+    def scale_model_input(self, latents: torch.Tensor,
+                          t: torch.Tensor) -> torch.Tensor:
+        self.calls["scale_model_input"] = (latents, t)
+        return latents + 1
+
+    def step(self, noise_pred: torch.Tensor, t: torch.Tensor,
+             latents: torch.Tensor, **kwargs):
+        self.calls["step"] = (noise_pred, t, latents, kwargs)
+        return "step-result"
+
+    def add_noise(self, latents: torch.Tensor, noise: torch.Tensor,
+                  t: torch.Tensor) -> torch.Tensor:
+        self.calls["add_noise"] = (latents, noise, t)
+        return latents + noise
+
+    def set_timesteps(self, num_steps: int, device=None, **kwargs):
+        self.calls["set_timesteps"] = (num_steps, device, kwargs)
+        return ["ok"]
+
+
+def test_default_scheduler_adapter_delegates():
+    scheduler = DummyScheduler()
+    adapter = DefaultSchedulerAdapter(scheduler)
+    latents = torch.zeros(2, 2)
+    noise = torch.ones(2, 2)
+    t = torch.tensor(1)
+
+    scaled = adapter.scale_model_input(latents, t)
+    assert torch.allclose(scaled, latents + 1)
+    assert scheduler.calls["scale_model_input"] == (latents, t)
+
+    step_out = adapter.step(torch.ones(1), t, latents, foo="bar")
+    assert step_out == "step-result"
+    assert scheduler.calls["step"][3]["foo"] == "bar"
+
+    noised = adapter.add_noise(latents, noise, t)
+    assert torch.allclose(noised, latents + noise)
+    assert scheduler.calls["add_noise"] == (latents, noise, t)
+
+    set_out = adapter.set_timesteps(4, device=torch.device("cpu"), baz=3)
+    assert set_out == ["ok"]
+    assert scheduler.calls["set_timesteps"][0] == 4

fastvideo/version.py

@@ -1 +1 @@
-__version__ = "0.1.6"
+__version__ = "0.1.7"

mkdocs.yml

@@ -153,6 +153,8 @@ nav:
     - Video Sparse Attention: attention/vsa/index.md
     - Sliding Tile Attention: attention/sta/index.md
     - Adding a New Attention Backend: attention/developer/index.md
+  - Utilities:
+    - LoRA: utilities/lora.md
   - Design:
     - Overview: design/overview.md
   - Developer Guide:

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "fastvideo"
-version = "0.1.6"
+version = "0.1.7"
 description = "FastVideo"
 readme = "README.md"
 requires-python = ">=3.10"
@@ -63,6 +63,7 @@ dependencies = [
     "remote-pdb",
 
     # Kernel & Packaging
+    "fastvideo-kernel==0.2.2",
     "wheel",
 
     # Training Dependencies
@@ -110,7 +111,7 @@ lint = [
 ]
 
 test = [
-    "av==14.3.0",
+    "av",
     "pytorch-msssim==1.0.0",
     "pytest",
 ]

pyproject_other.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "fastvideo"
-version = "0.1.6"
+version = "0.1.7"
 description = "FastVideo"
 readme = "README.md"
 requires-python = ">=3.10"
@@ -63,6 +63,7 @@ dependencies = [
     "remote-pdb",
 
     # Kernel & Packaging
+    "fastvideo-kernel==0.2.2",
     "wheel",
 
     # Training Dependencies
@@ -89,7 +90,7 @@ lint = [
 ]
 
 test = [
-    "av==14.3.0",
+    "av",
     "pytorch-msssim==1.0.0",
     "pytest",
 ]

scripts/checkpoint_conversion/convert_turbodiffusion_i2v_to_diffusers.py

@@ -0,0 +1,193 @@
+#!/usr/bin/env python3
+"""
+Convert TurboDiffusion I2V .pth checkpoint to Diffusers safetensors format.
+
+TurboDiffusion I2V uses two models: high-noise and low-noise.
+This script converts both checkpoints to Diffusers format.
+
+Usage:
+    python convert_turbodiffusion_i2v_to_diffusers.py \
+        --high_noise_path /path/to/TurboWan2.2-I2V-A14B-high-720P.pth \
+        --low_noise_path /path/to/TurboWan2.2-I2V-A14B-low-720P.pth \
+        --output_dir /path/to/output \
+        --reference_repo Wan-AI/Wan2.1-I2V-14B-720P-Diffusers
+"""
+
+import argparse
+import os
+import re
+import json
+import torch
+import shutil
+import glob
+from safetensors import safe_open
+from safetensors.torch import save_file
+from huggingface_hub import snapshot_download
+
+# Weight mapping from TurboDiffusion -> Diffusers/FastVideo format
+# Same as T2V but may need additional I2V-specific mappings
+TURBODIFFUSION_WEIGHT_MAPPING = {
+    # Self attention
+    r"^blocks\.(\d+)\.self_attn\.q\.(.*)$": r"blocks.\1.to_q.\2",
+    r"^blocks\.(\d+)\.self_attn\.k\.(.*)$": r"blocks.\1.to_k.\2",
+    r"^blocks\.(\d+)\.self_attn\.v\.(.*)$": r"blocks.\1.to_v.\2",
+    r"^blocks\.(\d+)\.self_attn\.o\.(.*)$": r"blocks.\1.to_out.\2",
+    r"^blocks\.(\d+)\.self_attn\.norm_q\.(.*)$": r"blocks.\1.norm_q.\2",
+    r"^blocks\.(\d+)\.self_attn\.norm_k\.(.*)$": r"blocks.\1.norm_k.\2",
+    # Cross attention
+    r"^blocks\.(\d+)\.cross_attn\.q\.(.*)$": r"blocks.\1.attn2.to_q.\2",
+    r"^blocks\.(\d+)\.cross_attn\.k\.(.*)$": r"blocks.\1.attn2.to_k.\2",
+    r"^blocks\.(\d+)\.cross_attn\.v\.(.*)$": r"blocks.\1.attn2.to_v.\2",
+    r"^blocks\.(\d+)\.cross_attn\.o\.(.*)$": r"blocks.\1.attn2.to_out.\2",
+    r"^blocks\.(\d+)\.cross_attn\.norm_q\.(.*)$": r"blocks.\1.attn2.norm_q.\2",
+    r"^blocks\.(\d+)\.cross_attn\.norm_k\.(.*)$": r"blocks.\1.attn2.norm_k.\2",
+    # I2V-specific cross attention (add_k/v_proj for image context)
+    r"^blocks\.(\d+)\.cross_attn\.add_k\.(.*)$": r"blocks.\1.attn2.add_k_proj.\2",
+    r"^blocks\.(\d+)\.cross_attn\.add_v\.(.*)$": r"blocks.\1.attn2.add_v_proj.\2",
+    r"^blocks\.(\d+)\.cross_attn\.norm_add_k\.(.*)$": r"blocks.\1.attn2.norm_added_k.\2",
+    r"^blocks\.(\d+)\.cross_attn\.norm_add_q\.(.*)$": r"blocks.\1.attn2.norm_added_q.\2",
+    # Norms and FFN
+    r"^blocks\.(\d+)\.norm1\.(.*)$": r"blocks.\1.norm1.\2",
+    r"^blocks\.(\d+)\.norm3\.(.*)$": r"blocks.\1.self_attn_residual_norm.norm.\2",
+    r"^blocks\.(\d+)\.norm2\.(.*)$": r"blocks.\1.norm3.\2",
+    r"^blocks\.(\d+)\.ffn\.0\.(.*)$": r"blocks.\1.ffn.fc_in.\2",
+    r"^blocks\.(\d+)\.ffn\.2\.(.*)$": r"blocks.\1.ffn.fc_out.\2",
+    r"^blocks\.(\d+)\.modulation$": r"blocks.\1.scale_shift_table",
+    # Embeddings
+    r"^text_embedding\.0\.(.*)$": r"condition_embedder.text_embedder.fc_in.\1",
+    r"^text_embedding\.2\.(.*)$": r"condition_embedder.text_embedder.fc_out.\1",
+    r"^time_embedding\.0\.(.*)$": r"condition_embedder.time_embedder.mlp.fc_in.\1",
+    r"^time_embedding\.2\.(.*)$": r"condition_embedder.time_embedder.mlp.fc_out.\1",
+    r"^time_projection\.1\.(.*)$": r"condition_embedder.time_modulation.linear.\1",
+    # Head
+    r"^head\.head\.(.*)$": r"proj_out.\1",
+    r"^head\.norm\.(.*)$": r"norm_out.\1",
+    r"^head\.modulation$": r"scale_shift_table",
+    # SLA proj_l weights
+    r"^blocks\.(\d+)\.self_attn\.attn_op\.local_attn\.proj_l\.(.*)$": r"blocks.\1.attn1.attn_impl.proj_l.\2",
+}
+
+SKIP_PATTERNS = []
+
+
+def should_skip_key(key: str) -> bool:
+    for pattern in SKIP_PATTERNS:
+        if re.match(pattern, key):
+            return True
+    return False
+
+
+def convert_key(turbo_key: str) -> str:
+    for pattern, replacement in TURBODIFFUSION_WEIGHT_MAPPING.items():
+        if re.match(pattern, turbo_key):
+            return re.sub(pattern, replacement, turbo_key)
+    return turbo_key
+
+
+def reshape_patch_embedding(tensor: torch.Tensor, target_shape: tuple) -> torch.Tensor:
+    if len(tensor.shape) == 2 and len(target_shape) == 5:
+        return tensor.view(target_shape)
+    return tensor
+
+
+def get_reference_shapes(reference_repo: str) -> dict:
+    print(f"Downloading reference model shapes from {reference_repo}...")
+    
+    local_dir = snapshot_download(
+        repo_id=reference_repo,
+        allow_patterns=["transformer/config.json", "transformer/diffusion_pytorch_model*.safetensors"],
+        local_dir_use_symlinks=False
+    )
+    
+    weight_files = glob.glob(os.path.join(local_dir, "transformer", "*.safetensors"))
+    shapes = {}
+    
+    for wf in weight_files:
+        with safe_open(wf, framework="pt") as f:
+            for key in f.keys():
+                shapes[key] = f.get_tensor(key).shape
+    
+    return shapes, local_dir
+
+
+def convert_checkpoint(input_path: str, output_dir: str, ref_shapes: dict, model_name: str) -> None:
+    """Convert a single TurboDiffusion checkpoint to Diffusers format."""
+    print(f"\n{'='*60}")
+    print(f"Converting {model_name}: {input_path}")
+    print(f"{'='*60}")
+    
+    turbo_state_dict = torch.load(input_path, map_location="cpu", weights_only=True)
+    print(f"Loaded {len(turbo_state_dict)} keys")
+    
+    converted_state_dict = {}
+    skipped_keys = []
+    
+    for turbo_key, tensor in turbo_state_dict.items():
+        if should_skip_key(turbo_key):
+            skipped_keys.append(turbo_key)
+            continue
+        
+        new_key = convert_key(turbo_key)
+        
+        if "patch_embedding" in new_key and new_key in ref_shapes:
+            target_shape = ref_shapes[new_key]
+            if tensor.shape != target_shape:
+                print(f"Reshaping {new_key}: {tensor.shape} -> {target_shape}")
+                tensor = reshape_patch_embedding(tensor, target_shape)
+        
+        if new_key in ref_shapes:
+            if tensor.shape != ref_shapes[new_key]:
+                print(f"WARNING: Shape mismatch for {new_key}: got {tensor.shape}, expected {ref_shapes[new_key]}")
+        
+        converted_state_dict[new_key] = tensor
+    
+    print(f"Converted: {len(converted_state_dict)} keys, Skipped: {len(skipped_keys)} keys")
+    
+    os.makedirs(output_dir, exist_ok=True)
+    output_path = os.path.join(output_dir, "diffusion_pytorch_model.safetensors")
+    print(f"Saving to {output_path}...")
+    save_file(converted_state_dict, output_path)
+    
+    return converted_state_dict
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Convert TurboDiffusion I2V checkpoints to Diffusers format")
+    parser.add_argument("--high_noise_path", type=str, required=True,
+                        help="Path to high-noise TurboDiffusion .pth checkpoint")
+    parser.add_argument("--low_noise_path", type=str, required=True,
+                        help="Path to low-noise TurboDiffusion .pth checkpoint")
+    parser.add_argument("--output_dir", type=str, required=True,
+                        help="Output directory for converted safetensors")
+    parser.add_argument("--reference_repo", type=str, default="Wan-AI/Wan2.2-I2V-A14B-Diffusers",
+                        help="Reference HF repo to get expected tensor shapes")
+    
+    args = parser.parse_args()
+    
+    # Get reference shapes
+    ref_shapes, ref_local_dir = get_reference_shapes(args.reference_repo)
+    print(f"Got {len(ref_shapes)} reference shapes")
+    
+    # Convert high-noise model
+    high_noise_output = os.path.join(args.output_dir, "transformer_high")
+    convert_checkpoint(args.high_noise_path, high_noise_output, ref_shapes, "high-noise")
+    
+    # Convert low-noise model
+    low_noise_output = os.path.join(args.output_dir, "transformer_low")
+    convert_checkpoint(args.low_noise_path, low_noise_output, ref_shapes, "low-noise")
+    
+    # Copy config.json to both
+    src_config = os.path.join(ref_local_dir, "transformer", "config.json")
+    shutil.copy(src_config, os.path.join(high_noise_output, "config.json"))
+    shutil.copy(src_config, os.path.join(low_noise_output, "config.json"))
+    print("Copied config.json to both transformer directories")
+    
+    print(f"\n{'='*60}")
+    print(f"Conversion complete!")
+    print(f"High-noise model: {high_noise_output}")
+    print(f"Low-noise model: {low_noise_output}")
+    print(f"{'='*60}")
+
+
+if __name__ == "__main__":
+    main()

scripts/inference/v1_inference_longcat.sh

@@ -1,14 +1,30 @@
 #!/bin/bash
 
-num_gpus=2
+# LongCat Text-to-Video (T2V) Inference Script
+# 
+# This script runs LongCat T2V inference using the fastvideo CLI.
+#
+# Usage:
+#   bash scripts/inference/v1_inference_longcat.sh
+#
+# Prerequisites:
+#   - Install fastvideo: pip install -e .
+#   - The model weights will be auto-downloaded from HuggingFace
+
+num_gpus=1
 
 export FASTVIDEO_ATTENTION_BACKEND=
 
-# For longcat, we must first convert the official weights to FastVideo native format
+# Model path options:
+# Option 1: HuggingFace model (auto-downloaded)
+export MODEL_BASE=FastVideo/LongCat-Video-T2V-Diffusers
+
+# Option 2: Local weights (uncomment if you have local weights)
+# For local weights, convert the official weights to FastVideo native format
 # conversion method: python scripts/checkpoint_conversion/longcat_to_fastvideo.py
 # --source /path/to/LongCat-Video/weights/LongCat-Video
 # --output weights/longcat-native
-export MODEL_BASE=weights/longcat-native
+# export MODEL_BASE=weights/longcat-native
 
 fastvideo generate \
     --model-path $MODEL_BASE \
@@ -26,7 +42,7 @@ fastvideo generate \
     --num-inference-steps 50 \
     --fps 15 \
     --guidance-scale 4.0 \
-    --prompt-txt assets/prompt.txt \
+    --prompt "In a realistic photography style, a white boy around seven or eight years old sits on a park bench, wearing a light blue T-shirt, denim shorts, and white sneakers. He holds an ice cream cone with vanilla and chocolate flavors, and beside him is a medium-sized golden Labrador. Smiling, the boy offers the ice cream to the dog, who eagerly licks it with its tongue. The sun is shining brightly, and the background features a green lawn and several tall trees, creating a warm and loving scene." \
     --negative-prompt "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards" \
     --seed 42 \
-    --output-path outputs_video/longcat_480p
+    --output-path outputs_video/longcat_t2v

scripts/inference/v1_inference_longcat_distill.sh

@@ -1,12 +1,23 @@
 #!/bin/bash
 
+# LongCat T2V Distilled Inference Script
+# 
+# This script runs LongCat T2V with distillation LoRA (16 steps instead of 50).
+# Uses the distilled LoRA for faster generation.
+#
+# Usage:
+#   bash scripts/inference/v1_inference_longcat_distill.sh
+#
+# Prerequisites:
+#   - Install fastvideo: pip install -e .
+#   - The model weights will be auto-downloaded from HuggingFace
+
 num_gpus=1
+
 export FASTVIDEO_ATTENTION_BACKEND=
-# For longcat, we must first convert the official weights to FastVideo native format
-# conversion method: python scripts/checkpoint_conversion/longcat_to_fastvideo.py
-# --source /path/to/LongCat-Video/weights/LongCat-Video
-# --output weights/longcat-native
-export MODEL_BASE=weights/longcat-native
+
+# Model path - HuggingFace model (auto-downloaded)
+export MODEL_BASE=FastVideo/LongCat-Video-T2V-Diffusers
 
 fastvideo generate \
     --model-path $MODEL_BASE \
@@ -14,11 +25,11 @@ fastvideo generate \
     --tp-size 1 \
     --num-gpus $num_gpus \
     --dit-cpu-offload False \
-    --vae-cpu-offload False \
-    --text-encoder-cpu-offload False \
+    --vae-cpu-offload True \
+    --text-encoder-cpu-offload True \
     --pin-cpu-memory False \
     --enable-bsa False \
-    --lora-path "$MODEL_BASE/lora/distilled" \
+    --lora-path "FastVideo/LongCat-Video-T2V-Distilled-LoRA" \
     --lora-nickname "distilled" \
     --height 480 \
     --width 832 \
@@ -26,6 +37,7 @@ fastvideo generate \
     --num-inference-steps 16 \
     --fps 15 \
     --guidance-scale 1.0 \
-    --prompt "In a realistic photography style, an asian boy around seven or eight years old sits on a park bench, wearing a light yellow T-shirt, denim shorts, and white sneakers. He holds an ice cream cone with vanilla and chocolate flavors, and beside him is a medium-sized golden Labrador. Smiling, the boy offers the ice cream to the dog, who eagerly licks it with its tongue. The sun is shining brightly, and the background features a green lawn and several tall trees, creating a warm and loving scene." \
+    --prompt "In a realistic photography style, a white boy around seven or eight years old sits on a park bench, wearing a light blue T-shirt, denim shorts, and white sneakers. He holds an ice cream cone with vanilla and chocolate flavors, and beside him is a medium-sized golden Labrador. Smiling, the boy offers the ice cream to the dog, who eagerly licks it with its tongue. The sun is shining brightly, and the background features a green lawn and several tall trees, creating a warm and loving scene." \
+    --negative-prompt "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards" \
     --seed 42 \
     --output-path outputs_video/longcat_distill

scripts/inference/v1_inference_longcat_i2v.sh

@@ -0,0 +1,60 @@
+#!/bin/bash
+
+# LongCat Image-to-Video (I2V) Inference Script
+# 
+# This script runs LongCat I2V inference using the fastvideo CLI.
+# LongCat I2V takes an input image and generates a video from it.
+#
+# Usage:
+#   bash scripts/inference/v1_inference_longcat_i2v.sh
+#
+# Prerequisites:
+#   - Install fastvideo: pip install -e .
+#   - The model weights will be auto-downloaded from HuggingFace
+#   - Or use local weights if you have them
+
+num_gpus=1
+
+export FASTVIDEO_ATTENTION_BACKEND=
+
+# Model path options:
+# Option 1: HuggingFace model (auto-downloaded)
+export MODEL_BASE=FastVideo/LongCat-Video-I2V-Diffusers
+
+# Option 2: Local weights (uncomment if you have local weights)
+# export MODEL_BASE=weights/longcat-for-i2v
+
+# Input image path (must be square for LongCat I2V)
+IMAGE_PATH="assets/girl.png"
+
+# Check if image exists
+if [ ! -f "$IMAGE_PATH" ]; then
+    echo "Error: Image not found at $IMAGE_PATH"
+    echo "Please provide a valid image path"
+    exit 1
+fi
+
+fastvideo generate \
+    --model-path $MODEL_BASE \
+    --sp-size $num_gpus \
+    --tp-size 1 \
+    --num-gpus $num_gpus \
+    --dit-cpu-offload False \
+    --vae-cpu-offload True \
+    --text-encoder-cpu-offload True \
+    --pin-cpu-memory False \
+    --enable-bsa False \
+    --image-path "$IMAGE_PATH" \
+    --height 480 \
+    --width 480 \
+    --num-frames 93 \
+    --num-inference-steps 50 \
+    --fps 15 \
+    --guidance-scale 4.0 \
+    --prompt "A woman sits at a wooden table by the window in a cozy café. She reaches out with her right hand, picks up the white coffee cup from the saucer, and gently brings it to her lips to take a sip. After drinking, she places the cup back on the table and looks out the window, enjoying the peaceful atmosphere." \
+    --negative-prompt "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards" \
+    --seed 42 \
+    --output-path outputs_video/longcat_i2v
+
+
+

scripts/inference/v1_inference_longcat_refine_fromvideo.sh

@@ -1,18 +1,30 @@
 #!/bin/bash
 
-num_gpus=1
-export FASTVIDEO_ATTENTION_BACKEND=
-# For longcat, we must first convert the official weights to FastVideo native format
-# conversion method: python scripts/checkpoint_conversion/longcat_to_fastvideo.py
-# --source /path/to/LongCat-Video/weights/LongCat-Video
-# --output weights/longcat-native
-export MODEL_BASE=weights/longcat-native
+# LongCat T2V Refinement Script (480p -> 720p)
+# 
+# This script refines a 480p distilled video to 720p using the refinement LoRA.
+# Run v1_inference_longcat_distill.sh first to generate the 480p video.
+#
+# Usage:
+#   bash scripts/inference/v1_inference_longcat_refine_fromvideo.sh
+#
+# Prerequisites:
+#   - Install fastvideo: pip install -e .
+#   - The model weights will be auto-downloaded from HuggingFace
+#   - Run v1_inference_longcat_distill.sh first to generate input video
 
-INPUT_VIDEO="outputs_video/longcat_distill/In a realistic photography style, an asian boy around seven or eight years old sits on a park bench,.mp4"
+num_gpus=1
+
+export FASTVIDEO_ATTENTION_BACKEND=
+
+# Model path - HuggingFace model (auto-downloaded)
+export MODEL_BASE=FastVideo/LongCat-Video-T2V-Diffusers
+
+INPUT_VIDEO="outputs_video/longcat_distill/In a realistic photography style, a white boy around seven or eight years old sits on a park bench,.mp4"
 REFINE_OUTPUT="outputs_video/longcat_refine_720p"
 
-# Prompt used for base generation
-PROMPT="In a realistic photography style, an asian boy around seven or eight years old sits on a park bench, wearing a light yellow T-shirt, denim shorts, and white sneakers. He holds an ice cream cone with vanilla and chocolate flavors, and beside him is a medium-sized golden Labrador. Smiling, the boy offers the ice cream to the dog, who eagerly licks it with its tongue. The sun is shining brightly, and the background features a green lawn and several tall trees, creating a warm and loving scene."
+# Prompt used for base generation (must match distill script)
+PROMPT="In a realistic photography style, a white boy around seven or eight years old sits on a park bench, wearing a light blue T-shirt, denim shorts, and white sneakers. He holds an ice cream cone with vanilla and chocolate flavors, and beside him is a medium-sized golden Labrador. Smiling, the boy offers the ice cream to the dog, who eagerly licks it with its tongue. The sun is shining brightly, and the background features a green lawn and several tall trees, creating a warm and loving scene."
 
 echo "=========================================="
 echo "LongCat 480p -> 720p Refinement"
@@ -25,14 +37,14 @@ echo ""
 # Check if input video exists
 if [ ! -f "$INPUT_VIDEO" ]; then
     echo "Error: Input video not found: $INPUT_VIDEO"
-    echo "Please set INPUT_VIDEO to your 480p video path"
+    echo "Please run v1_inference_longcat_distill.sh first to generate the 480p video"
     exit 1
 fi
 
-echo "🔧 Configuring refinement (BSA enabled, refinement LoRA)..."
-echo "✅ Input video: $INPUT_VIDEO"
-echo "✅ BSA enabled with sparsity=0.875"
-echo "✅ Refinement LoRA loaded"
+echo "Configuring refinement (BSA enabled, refinement LoRA)..."
+echo "Input video: $INPUT_VIDEO"
+echo "BSA enabled with sparsity=0.875"
+echo "Refinement LoRA loaded"
 echo ""
 
 fastvideo generate \
@@ -41,14 +53,14 @@ fastvideo generate \
     --tp-size 1 \
     --num-gpus $num_gpus \
     --dit-cpu-offload True \
-    --vae-cpu-offload False \
+    --vae-cpu-offload True \
     --text-encoder-cpu-offload True \
     --pin-cpu-memory False \
     --enable-bsa True \
     --bsa-sparsity 0.875 \
     --bsa-chunk-q 4 4 8 \
     --bsa-chunk-k 4 4 8 \
-    --lora-path "$MODEL_BASE/lora/refinement" \
+    --lora-path "FastVideo/LongCat-Video-T2V-Refinement-LoRA" \
     --lora-nickname "refinement" \
     --refine-from "$INPUT_VIDEO" \
     --t-thresh 0.5 \
@@ -60,12 +72,13 @@ fastvideo generate \
     --fps 30 \
     --guidance-scale 1.0 \
     --prompt "$PROMPT" \
+    --negative-prompt "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards" \
     --seed 42 \
     --output-path "$REFINE_OUTPUT"
 
 echo ""
 echo "=========================================="
-echo "✓ Refinement Complete!"
+echo "Refinement Complete!"
 echo "=========================================="
 echo ""
 echo "Output directory: $REFINE_OUTPUT"

scripts/inference/v1_inference_longcat_vc.sh

@@ -0,0 +1,61 @@
+#!/bin/bash
+
+# LongCat Video Continuation (VC) Inference Script
+# 
+# This script runs LongCat VC inference using the fastvideo CLI.
+# LongCat VC takes an input video and generates a continuation of it.
+#
+# Usage:
+#   bash scripts/inference/v1_inference_longcat_vc.sh
+#
+# Prerequisites:
+#   - Install fastvideo: pip install -e .
+#   - The model weights will be auto-downloaded from HuggingFace
+#   - Or use local weights if you have them
+
+num_gpus=1
+
+export FASTVIDEO_ATTENTION_BACKEND=
+
+# Model path options:
+# Option 1: HuggingFace model (auto-downloaded)
+export MODEL_BASE=FastVideo/LongCat-Video-VC-Diffusers
+
+# Option 2: Local weights (uncomment if you have local weights)
+# export MODEL_BASE=weights/longcat-vc-upload
+
+# Input video path
+VIDEO_PATH="assets/motorcycle.mp4"
+
+# Check if video exists
+if [ ! -f "$VIDEO_PATH" ]; then
+    echo "Error: Video not found at $VIDEO_PATH"
+    echo "Please provide a valid video path"
+    exit 1
+fi
+
+fastvideo generate \
+    --model-path $MODEL_BASE \
+    --sp-size $num_gpus \
+    --tp-size 1 \
+    --num-gpus $num_gpus \
+    --dit-cpu-offload False \
+    --vae-cpu-offload True \
+    --text-encoder-cpu-offload True \
+    --pin-cpu-memory False \
+    --enable-bsa False \
+    --video-path "$VIDEO_PATH" \
+    --num-cond-frames 13 \
+    --height 480 \
+    --width 832 \
+    --num-frames 93 \
+    --num-inference-steps 50 \
+    --fps 15 \
+    --guidance-scale 4.0 \
+    --prompt "A person rides a motorcycle along a long, straight road that stretches between a body of water and a forested hillside. The rider steadily accelerates, keeping the motorcycle centered between the guardrails, while the scenery passes by on both sides. The video captures the journey from the rider's perspective, emphasizing the sense of motion and adventure." \
+    --negative-prompt "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards" \
+    --seed 42 \
+    --output-path outputs_video/longcat_vc
+
+
+