separate critic lr

update
add noisy ground truth for teacher
2025-07-21 19:26:04 +00:00 · 2025-07-18 23:20:30 +00:00 · 2025-07-18 22:28:15 +00:00 · 2025-07-18 22:21:41 +00:00 · 2025-07-18 21:12:56 +00:00 · 2025-07-18 21:11:25 +00:00
205 changed files with 8060 additions and 2015 deletions
@@ -1,13 +1,12 @@
 env:
  IMAGE_VERSION: "py3.12-latest"
+  BUILDKITE_CLEAN_CHECKOUT: true

 steps:
  - label: "pre-commit"
    command: ".buildkite/scripts/pre_commit.sh"
    agents:
      queue: "default"
-    env:
-      - BUILDKITE_CLEAN_CHECKOUT=true

  - wait

@@ -23,10 +22,9 @@ steps:
                - "pyproject.toml"
                - "docker/Dockerfile.python3.12"
              config:
-                command: "timeout 30m .buildkite/scripts/pr_test.sh"
+                command: "timeout 15m .buildkite/scripts/pr_test.sh"
                label: "Encoder Tests"
                env:
-                  - BUILDKITE_CLEAN_CHECKOUT=true
                  - TEST_TYPE=encoder
                agents:
                  queue: "default"
@@ -37,10 +35,9 @@ steps:
                - "pyproject.toml"
                - "docker/Dockerfile.python3.12"
              config:
-                command: "timeout 30m .buildkite/scripts/pr_test.sh"
+                command: "timeout 15m .buildkite/scripts/pr_test.sh"
                label: "VAE Tests"
                env:
-                  - BUILDKITE_CLEAN_CHECKOUT=true
                  - TEST_TYPE=vae
                agents:
                  queue: "default"
@@ -53,20 +50,18 @@ steps:
                - "pyproject.toml"
                - "docker/Dockerfile.python3.12"
              config:
-                command: "timeout 30m .buildkite/scripts/pr_test.sh"
+                command: "timeout 15m .buildkite/scripts/pr_test.sh"
                label: "Transformer Tests"
                env:
-                  - BUILDKITE_CLEAN_CHECKOUT=true
                  - TEST_TYPE=transformer
                agents:
                  queue: "default"
            - path: 
                - "fastvideo/v1/**/*.py"
              config:
-                command: "timeout 60m .buildkite/scripts/pr_test.sh"
+                command: "timeout 30m .buildkite/scripts/pr_test.sh"
                label: "SSIM Tests"
                env:
-                  - BUILDKITE_CLEAN_CHECKOUT=true
                  - TEST_TYPE=ssim
                agents:
                  queue: "default"
@@ -75,10 +70,9 @@ steps:
                - "pyproject.toml"
                - "docker/Dockerfile.python3.12"
              config:
-                command: "timeout 30m .buildkite/scripts/pr_test.sh"
+                command: "timeout 15m .buildkite/scripts/pr_test.sh"
                label: "Training Tests"
                env:
-                  - BUILDKITE_CLEAN_CHECKOUT=true
                  - TEST_TYPE=training
                agents:
                  queue: "default"
@@ -92,10 +86,9 @@ steps:
                - "pyproject.toml"
                - "docker/Dockerfile.python3.12"
              config:
-                command: "timeout 30m .buildkite/scripts/pr_test.sh"
+                command: "timeout 15m .buildkite/scripts/pr_test.sh"
                label: "Training Tests VSA"
                env:
-                  - BUILDKITE_CLEAN_CHECKOUT=true
                  - TEST_TYPE=training_vsa
                agents:
                  queue: "default"
@@ -108,10 +101,9 @@ steps:
                - "pyproject.toml"
                - "docker/Dockerfile.python3.12"
              config:
-                command: "timeout 30m .buildkite/scripts/pr_test.sh"
+                command: "timeout 15m .buildkite/scripts/pr_test.sh"
                label: "Inference Tests STA"
                env:
-                  - BUILDKITE_CLEAN_CHECKOUT=true
                  - TEST_TYPE=inference_sta
                agents:
                  queue: "default"
@@ -123,10 +115,9 @@ steps:
                - "pyproject.toml"
                - "docker/Dockerfile.python3.12"
              config:
-                command: "timeout 30m .buildkite/scripts/pr_test.sh"
+                command: "timeout 15m .buildkite/scripts/pr_test.sh"
                label: "Precision Tests STA"
                env:
-                  - BUILDKITE_CLEAN_CHECKOUT=true
                  - TEST_TYPE=precision_sta
                agents:
                  queue: "default"
@@ -139,10 +130,9 @@ steps:
                - "pyproject.toml"
                - "docker/Dockerfile.python3.12"
              config:
-                command: "timeout 30m .buildkite/scripts/pr_test.sh"
+                command: "timeout 15m .buildkite/scripts/pr_test.sh"
                label: "Precision Tests VSA"
                env:
-                  - BUILDKITE_CLEAN_CHECKOUT=true
                  - TEST_TYPE=precision_vsa
                agents:
                  queue: "default"
@@ -58,7 +58,7 @@ if [ -z "${TEST_TYPE:-}" ]; then
 fi
 log "Test type: $TEST_TYPE"

-MODAL_ENV="BUILDKITE_REPO=$BUILDKITE_REPO BUILDKITE_COMMIT=$BUILDKITE_COMMIT IMAGE_VERSION=$IMAGE_VERSION"
+MODAL_ENV="BUILDKITE_REPO=$BUILDKITE_REPO BUILDKITE_COMMIT=$BUILDKITE_COMMIT BUILDKITE_PULL_REQUEST=$BUILDKITE_PULL_REQUEST IMAGE_VERSION=$IMAGE_VERSION"

 case "$TEST_TYPE" in
    "encoder")
@@ -122,17 +122,17 @@ jobs:
            # Actual tests
            encoder-test:
              - 'fastvideo/v1/models/encoders/**'
-              - 'fastvideo/v1/models/loaders/**'
+              - 'fastvideo/v1/models/loader/**'
              - 'fastvideo/v1/tests/encoders/**'
              - *common-paths
            vae-test:
              - 'fastvideo/v1/models/vaes/**'
-              - 'fastvideo/v1/models/loaders/**'
+              - 'fastvideo/v1/models/loader/**'
              - 'fastvideo/v1/tests/vaes/**'
              - *common-paths
            transformer-test:
              - 'fastvideo/v1/models/dits/**'
-              - 'fastvideo/v1/models/loaders/**'
+              - 'fastvideo/v1/models/loader/**'
              - 'fastvideo/v1/tests/transformers/**'
              - 'fastvideo/v1/layers/**'
              - 'fastvideo/v1/attention/**'
@@ -59,3 +59,8 @@ docs/source/inference/examples/

 # Static images
 !docs/source/_static/images/**/*.png
+
+# Local scripts (keep local but don't track in git)
+local_scripts/
+!comfyui/assets/**/*.png
+!comfyui/assets/**/*.gif
@@ -8,7 +8,7 @@ It features a clean, consistent API that works across popular video models, maki
 With FastVideo's optimizations, you can achieve more than 3x inference improvement compared to other systems.

 <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.html"><b> Quick Start</b></a> | 🤗 <a href="https://huggingface.co/FastVideo/FastHunyuan"  target="_blank"><b>FastHunyuan</b></a>  | 🤗 <a href="https://huggingface.co/FastVideo/FastMochi-diffusers" target="_blank"><b>FastMochi</b></a> | 🟣💬 <a href="https://join.slack.com/t/fastvideo/shared_invite/zt-2zf6ru791-sRwI9lPIUJQq1mIeB_yjJg" target="_blank"> <b>Slack</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.html"><b> Quick Start</b></a> | 🤗 <a href="https://huggingface.co/FastVideo/FastHunyuan"  target="_blank"><b>FastHunyuan</b></a>  | 🤗 <a href="https://huggingface.co/FastVideo/FastMochi-diffusers" target="_blank"><b>FastMochi</b></a> | 🟣💬 <a href="https://join.slack.com/t/fastvideo/shared_invite/zt-38u6p1jqe-yDI1QJOCEnbtkLoaI5bjZQ" target="_blank"> <b>Slack</b> </a> |
 </p>

 <div align="center">
@@ -0,0 +1,138 @@
+# ComfyUI-FastVideo
+
+A custom node suite for ComfyUI that provides accelerated video generation using [FastVideo](https://github.com/hao-ai-labs/FastVideo). See the [blog post](https://hao-ai-lab.github.io/blogs/fastvideo/) about FastVideo V1 to learn more.
+
+## Multi-GPU Parallel Inference
+
+One of the key features ComfyUI-FastVideo brings to ComfyUI is its ability to distribute the generation workload across multiple GPUs, resulting in significantly faster inference times.
+
+![Wan2.1-I2V-14B-480P-Diffusers](./assets/wani2v.gif)
+
+Example of Wan2.1-I2V-14B-480P-Diffusers model running on 4 GPUs.
+## Features
+
+- Generate high-quality videos from text prompts and images
+- Configurable video parameters (prompt, resolution, frame count, FPS)
+- Support for multiple GPUs with tensor and sequence parallelism
+- Advanced configuration options for VAE, Text Encoder, and DIT components
+- Interruption/cancellation support for long-running generations
+
+## Installation
+
+### Requirements
+
+- [ComfyUI](https://github.com/comfyanonymous/ComfyUI)
+- CUDA-capable GPU(s) with sufficient VRAM
+
+### Install using ComfyUI Manager
+
+Coming soon!
+
+### Manual Installation
+
+#### Copy the FastVideo `comfyui` directory into your ComfyUI custom_nodes directory:
+
+```bash
+cp -r /path/to/FastVideo/comfyui /path/to/ComfyUI/custom_nodes/FastVideo
+```
+
+#### Install dependencies:
+
+Currently, the only dependency is `fastvideo`, which can be installed using pip.
+
+```bash
+pip install fastvideo
+```
+
+#### Install missing custom nodes:
+
+`ComfyUI-VideoHelperSuite`:
+
+```bash
+cd /path/to/ComfyUI/custom_nodes
+git clone https://github.com/Kosinkadink/ComfyUI-VideoHelperSuite.git
+```
+
+If you're seeing `ImportError: libGL.so.1: cannot open shared object file: No such file or directory`,
+you may need to install ffmpeg
+
+```bash
+apt-get update && apt-get install ffmpeg
+```
+
+## Usage
+
+After installation, the following nodes will be available in the ComfyUI interface under the "fastvideo" category:
+
+- **Video Generator**: The main node for generating videos from prompts
+- **Inference Args**: Configure video generation parameters
+- **VAE Config**
+- **Text Encoder Config**
+- **DIT Config**
+- **Load Image Path**: Load images for potential conditioning
+
+You may have noticed many arguments on the nodes have 'auto' as the default value. This is because FastVideo will automatically detect the best values for these parameters based on the model and the hardware. However, you can also manually configure these parameters to get the best performance for your specific use case. We plan on releasing more optimized workflow files for different models and hardware configurations in the future.
+
+You can see what some of the default configurations are by looking at the FastVideo repo:
+- [Wan2.1-I2V-14B-480P-Diffusers](https://github.com/hao-ai-lab/FastVideo/blob/main/fastvideo/v1/configs/wan_14B_i2v_480p_pipeline.json)
+- [FastHunyuan-diffusers](https://github.com/hao-ai-lab/FastVideo/blob/main/fastvideo/v1/configs/fasthunyuan_t2v.json)
+
+### Node Configuration
+
+#### Video Generator
+
+- **prompt**: Text description of the video to generate
+- **output_path**: Directory where generated videos will be saved
+- **num_gpus**: Number of GPUs to use for generation
+- **model_path**: Path to the FastVideo model
+- **embedded_cfg_scale**: Classifier-free guidance scale
+- **sp_size**: Sequence parallelism size (usually should match num_gpus)
+- **tp_size**: Tensor parallelism size (usually should match num_gpus)
+- **precision**: Model precision (fp16 or bf16)
+
+`model_path takes either a model id from huggingface or a local path to a model. Models by default will be downloaded to ~/.cache/huggingface/hub/ and cached for subsequent runs.`
+
+#### Inference Args
+
+- **height/width**: Resolution of the output video
+- **num_frames**: Number of frames to generate
+- **num_inference_steps**: Number of diffusion steps per frame
+- **guidance_scale**: Classifier-free guidance scale
+- **flow_shift**: Frame flow shift parameter
+- **seed**: Random seed for reproducible generation
+- **fps**: Frames per second of the output video
+- **image_path**: Optional path to input image for conditioning (for i2v models)
+
+## Memory Management
+
+Models will remain loaded in GPU memory between runs when you only change inference arguments (such as prompt, resolution, frame count, FPS, guidance scale, etc.) or the prompt text. This allows for faster subsequent generations since the model doesn't need to be reloaded.
+
+However, if you need to change the following parameters, you will need to restart the ComfyUI server:
+- **Number of GPUs** (`num_gpus`)
+- **Model path** (`model_path`)
+- **Tensor parallelism size** (`tp_size`)
+- **Sequence parallelism size** (`sp_size`)
+
+These parameters affect the model's distribution across GPUs and require a complete reinitialization of the model pipeline.
+
+## Example workflows
+
+### Text to Video
+
+FastVideo-FastHunyuan-diffusers
+
+![FastVideo-FastHunyuan-diffusers](./assets/fasthunyuan.png)
+
+- [FastHunyuan-diffusers.json](./examples/FastHunyuan-diffusers.json)
+
+### Image to Video
+
+Wan2.1-I2V-14B-480P-Diffusers
+
+![Wan2.1-I2V-14B-480P-Diffusers](./assets/wani2v.png)
+
+- [Wan2.1-I2V-14B-480P-Diffusers.json](./examples/Wan2.1-I2V-14B-480P-Diffusers.json)
+
+## License
+
+This project is licensed under Apache 2.0.
@@ -0,0 +1,5 @@
+from .video_generator.nodes import (NODE_CLASS_MAPPINGS,
+                                    NODE_DISPLAY_NAME_MAPPINGS)
+
+WEB_DIRECTORY = "./web"
+__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS', 'WEB_DIRECTORY']
@@ -0,0 +1,645 @@
+{
+  "id": "23a1f065-bbba-4a8f-b144-944e1318fcbf",
+  "revision": 0,
+  "last_node_id": 8,
+  "last_link_id": 7,
+  "nodes": [
+    {
+      "id": 4,
+      "type": "VAEConfig",
+      "pos": [
+        374.2159423828125,
+        554.85888671875
+      ],
+      "size": [
+        334.080078125,
+        322
+      ],
+      "flags": {},
+      "order": 0,
+      "mode": 0,
+      "inputs": [],
+      "outputs": [
+        {
+          "name": "vae_config",
+          "type": "VAE_CONFIG",
+          "links": [
+            2
+          ]
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "VAEConfig"
+      },
+      "widgets_values": [
+        -99999,
+        -99999,
+        -99999,
+        -99999,
+        -99999,
+        -99999,
+        -99999,
+        -99999,
+        -99999,
+        -99999,
+        -99999,
+        -99999
+      ],
+      "auto_widget_states": {
+        "load_encoder": {
+          "isAuto": true,
+          "value": -99999,
+          "cachedValue": true
+        },
+        "load_decoder": {
+          "isAuto": true,
+          "value": -99999,
+          "cachedValue": true
+        },
+        "tile_sample_min_height": {
+          "isAuto": true,
+          "value": -99999,
+          "cachedValue": 256
+        },
+        "tile_sample_min_width": {
+          "isAuto": true,
+          "value": -99999,
+          "cachedValue": 256
+        },
+        "tile_sample_min_num_frames": {
+          "isAuto": true,
+          "value": -99999,
+          "cachedValue": 16
+        },
+        "tile_sample_stride_height": {
+          "isAuto": true,
+          "value": -99999,
+          "cachedValue": 192
+        },
+        "tile_sample_stride_width": {
+          "isAuto": true,
+          "value": -99999,
+          "cachedValue": 192
+        },
+        "tile_sample_stride_num_frames": {
+          "isAuto": true,
+          "value": -99999,
+          "cachedValue": 12
+        },
+        "blend_num_frames": {
+          "isAuto": true,
+          "value": -99999,
+          "cachedValue": 0
+        },
+        "use_tiling": {
+          "isAuto": true,
+          "value": -99999,
+          "cachedValue": true
+        },
+        "use_temporal_tiling": {
+          "isAuto": true,
+          "value": -99999,
+          "cachedValue": true
+        },
+        "use_parallel_tiling": {
+          "isAuto": true,
+          "value": -99999,
+          "cachedValue": true
+        }
+      }
+    },
+    {
+      "id": 5,
+      "type": "TextEncoderConfig",
+      "pos": [
+        416.4937744140625,
+        953.6171875
+      ],
+      "size": [
+        270,
+        106
+      ],
+      "flags": {},
+      "order": 1,
+      "mode": 0,
+      "inputs": [],
+      "outputs": [
+        {
+          "name": "text_encoder_config",
+          "type": "TEXT_ENCODER_CONFIG",
+          "links": [
+            7
+          ]
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "TextEncoderConfig"
+      },
+      "widgets_values": [
+        -99999,
+        -99999,
+        -99999
+      ],
+      "auto_widget_states": {
+        "prefix": {
+          "isAuto": true,
+          "value": -99999,
+          "cachedValue": ""
+        },
+        "quant_config": {
+          "isAuto": true,
+          "value": -99999,
+          "cachedValue": ""
+        },
+        "lora_config": {
+          "isAuto": true,
+          "value": -99999,
+          "cachedValue": ""
+        }
+      }
+    },
+    {
+      "id": 6,
+      "type": "DITConfig",
+      "pos": [
+        415.1928405761719,
+        1154.1573486328125
+      ],
+      "size": [
+        270,
+        82
+      ],
+      "flags": {},
+      "order": 2,
+      "mode": 0,
+      "inputs": [],
+      "outputs": [
+        {
+          "name": "dit_config",
+          "type": "DIT_CONFIG",
+          "links": [
+            6
+          ]
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "DITConfig"
+      },
+      "widgets_values": [
+        -99999,
+        -99999
+      ],
+      "auto_widget_states": {
+        "prefix": {
+          "isAuto": true,
+          "value": -99999,
+          "cachedValue": ""
+        },
+        "quant_config": {
+          "isAuto": true,
+          "value": -99999,
+          "cachedValue": ""
+        }
+      }
+    },
+    {
+      "id": 1,
+      "type": "VideoGenerator",
+      "pos": [
+        818.804931640625,
+        348.9299621582031
+      ],
+      "size": [
+        400,
+        436
+      ],
+      "flags": {},
+      "order": 4,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "inference_args",
+          "shape": 7,
+          "type": "INFERENCE_ARGS",
+          "link": 3
+        },
+        {
+          "name": "vae_config",
+          "shape": 7,
+          "type": "VAE_CONFIG",
+          "link": 2
+        },
+        {
+          "name": "text_encoder_config",
+          "shape": 7,
+          "type": "TEXT_ENCODER_CONFIG",
+          "link": 7
+        },
+        {
+          "name": "dit_config",
+          "shape": 7,
+          "type": "DIT_CONFIG",
+          "link": 6
+        }
+      ],
+      "outputs": [
+        {
+          "name": "video_path",
+          "type": "STRING",
+          "links": [
+            4
+          ]
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "VideoGenerator"
+      },
+      "widgets_values": [
+        "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 natural light filtering through the petals. Mid-shot, warm and cheerful tones.",
+        "/workspace/ComfyUI/outputs_video/",
+        2,
+        "FastVideo/FastHunyuan-diffusers",
+        -99999,
+        -99999,
+        -99999,
+        -99999,
+        -99999,
+        -99999,
+        -99999,
+        -99999,
+        -99999
+      ],
+      "auto_widget_states": {
+        "embedded_cfg_scale": {
+          "isAuto": true,
+          "value": -99999,
+          "cachedValue": 6
+        },
+        "sp_size": {
+          "isAuto": true,
+          "value": -99999,
+          "cachedValue": 2
+        },
+        "tp_size": {
+          "isAuto": true,
+          "value": -99999,
+          "cachedValue": 2
+        },
+        "vae_precision": {
+          "isAuto": true,
+          "value": -99999,
+          "cachedValue": "fp16"
+        },
+        "vae_tiling": {
+          "isAuto": true,
+          "value": -99999,
+          "cachedValue": true
+        },
+        "vae_sp": {
+          "isAuto": true,
+          "value": -99999,
+          "cachedValue": true
+        },
+        "text_encoder_precision": {
+          "isAuto": true,
+          "value": -99999,
+          "cachedValue": "fp16"
+        },
+        "precision": {
+          "isAuto": true,
+          "value": -99999,
+          "cachedValue": "fp16"
+        },
+        "use_cpu_offload": {
+          "isAuto": true,
+          "value": -99999,
+          "cachedValue": true
+        }
+      }
+    },
+    {
+      "id": 3,
+      "type": "VHS_LoadVideoPath",
+      "pos": [
+        1350.136962890625,
+        331.20361328125
+      ],
+      "size": [
+        231.8896484375,
+        286
+      ],
+      "flags": {},
+      "order": 5,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "meta_batch",
+          "shape": 7,
+          "type": "VHS_BatchManager",
+          "link": null
+        },
+        {
+          "name": "vae",
+          "shape": 7,
+          "type": "VAE",
+          "link": null
+        },
+        {
+          "name": "video",
+          "type": "STRING",
+          "widget": {
+            "name": "video"
+          },
+          "link": 4
+        }
+      ],
+      "outputs": [
+        {
+          "name": "IMAGE",
+          "type": "IMAGE",
+          "links": [
+            5
+          ]
+        },
+        {
+          "name": "frame_count",
+          "type": "INT",
+          "links": null
+        },
+        {
+          "name": "audio",
+          "type": "AUDIO",
+          "links": null
+        },
+        {
+          "name": "video_info",
+          "type": "VHS_VIDEOINFO",
+          "links": null
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "VHS_LoadVideoPath"
+      },
+      "widgets_values": {
+        "video": "",
+        "force_rate": 0,
+        "custom_width": 0,
+        "custom_height": 0,
+        "frame_load_cap": 0,
+        "skip_first_frames": 0,
+        "select_every_nth": 1,
+        "format": "Wan",
+        "videopreview": {
+          "hidden": false,
+          "paused": false,
+          "params": {
+            "filename": "",
+            "type": "path",
+            "format": "video/",
+            "force_rate": 0,
+            "custom_width": 0,
+            "custom_height": 0,
+            "frame_load_cap": 0,
+            "skip_first_frames": 0,
+            "select_every_nth": 1
+          }
+        }
+      }
+    },
+    {
+      "id": 2,
+      "type": "InferenceArgs",
+      "pos": [
+        411.46307373046875,
+        178.18182373046875
+      ],
+      "size": [
+        278.73828125,
+        298
+      ],
+      "flags": {},
+      "order": 3,
+      "mode": 0,
+      "inputs": [],
+      "outputs": [
+        {
+          "name": "inference_args",
+          "type": "INFERENCE_ARGS",
+          "links": [
+            3
+          ]
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "InferenceArgs"
+      },
+      "widgets_values": [
+        720,
+        1280,
+        45,
+        6,
+        -99999,
+        -99999,
+        1025,
+        "fixed",
+        24,
+        -99999,
+        -99999
+      ],
+      "auto_widget_states": {
+        "height": {
+          "isAuto": false,
+          "value": 720,
+          "cachedValue": 720
+        },
+        "width": {
+          "isAuto": false,
+          "value": 1280,
+          "cachedValue": 1280
+        },
+        "num_frames": {
+          "isAuto": false,
+          "value": 45,
+          "cachedValue": 45
+        },
+        "num_inference_steps": {
+          "isAuto": false,
+          "value": 6,
+          "cachedValue": 6
+        },
+        "guidance_scale": {
+          "isAuto": true,
+          "value": -99999,
+          "cachedValue": 1
+        },
+        "flow_shift": {
+          "isAuto": true,
+          "value": -99999,
+          "cachedValue": 17
+        },
+        "seed": {
+          "isAuto": false,
+          "value": 1025,
+          "cachedValue": 1024
+        },
+        "fps": {
+          "isAuto": false,
+          "value": 24,
+          "cachedValue": 24
+        },
+        "image_path": {
+          "isAuto": true,
+          "value": -99999,
+          "cachedValue": "X://insert/path/here.mp4"
+        },
+        "enable_teacache": {
+          "isAuto": true,
+          "value": -99999,
+          "cachedValue": true
+        }
+      }
+    },
+    {
+      "id": 8,
+      "type": "VHS_VideoCombine",
+      "pos": [
+        1668.3499755859375,
+        328.22625732421875
+      ],
+      "size": [
+        507.507080078125,
+        622.2227172851562
+      ],
+      "flags": {},
+      "order": 6,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "images",
+          "type": "IMAGE",
+          "link": 5
+        },
+        {
+          "name": "audio",
+          "shape": 7,
+          "type": "AUDIO",
+          "link": null
+        },
+        {
+          "name": "meta_batch",
+          "shape": 7,
+          "type": "VHS_BatchManager",
+          "link": null
+        },
+        {
+          "name": "vae",
+          "shape": 7,
+          "type": "VAE",
+          "link": null
+        }
+      ],
+      "outputs": [
+        {
+          "name": "Filenames",
+          "type": "VHS_FILENAMES",
+          "links": null
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "VHS_VideoCombine"
+      },
+      "widgets_values": {
+        "frame_rate": 24,
+        "loop_count": 0,
+        "filename_prefix": "",
+        "format": "video/h264-mp4",
+        "pix_fmt": "yuv420p",
+        "crf": 19,
+        "save_metadata": true,
+        "trim_to_audio": false,
+        "pingpong": false,
+        "save_output": false,
+        "videopreview": {
+          "hidden": false,
+          "paused": false,
+          "params": {
+            "filename": "._00003.mp4",
+            "subfolder": "",
+            "type": "temp",
+            "format": "video/h264-mp4",
+            "frame_rate": 24,
+            "workflow": "._00003.png",
+            "fullpath": "/workspace/ComfyUI/temp/._00003.mp4"
+          }
+        }
+      }
+    }
+  ],
+  "links": [
+    [
+      2,
+      4,
+      0,
+      1,
+      1,
+      "VAE_CONFIG"
+    ],
+    [
+      3,
+      2,
+      0,
+      1,
+      0,
+      "INFERENCE_ARGS"
+    ],
+    [
+      4,
+      1,
+      0,
+      3,
+      2,
+      "STRING"
+    ],
+    [
+      5,
+      3,
+      0,
+      8,
+      0,
+      "IMAGE"
+    ],
+    [
+      6,
+      6,
+      0,
+      1,
+      3,
+      "DIT_CONFIG"
+    ],
+    [
+      7,
+      5,
+      0,
+      1,
+      2,
+      "TEXT_ENCODER_CONFIG"
+    ]
+  ],
+  "groups": [],
+  "config": {},
+  "extra": {
+    "ds": {
+      "scale": 0.9090909090909091,
+      "offset": [
+        112.86678372727341,
+        -71.45635903989245
+      ]
+    },
+    "frontendVersion": "1.20.4",
+    "VHS_latentpreview": false,
+    "VHS_latentpreviewrate": 0,
+    "VHS_MetadataImage": true,
+    "VHS_KeepIntermediate": true
+  },
+  "version": 0.4
+}
@@ -0,0 +1,697 @@
+{
+  "id": "23a1f065-bbba-4a8f-b144-944e1318fcbf",
+  "revision": 0,
+  "last_node_id": 8,
+  "last_link_id": 7,
+  "nodes": [
+    {
+      "id": 7,
+      "type": "LoadImagePath",
+      "pos": [
+        33.15385437011719,
+        191.2037353515625
+      ],
+      "size": [
+        270,
+        334
+      ],
+      "flags": {},
+      "order": 0,
+      "mode": 0,
+      "inputs": [],
+      "outputs": [
+        {
+          "name": "image_path",
+          "type": "STRING",
+          "links": [
+            1
+          ]
+        },
+        {
+          "name": "IMAGE",
+          "type": "IMAGE",
+          "links": null
+        },
+        {
+          "name": "MASK",
+          "type": "MASK",
+          "links": null
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "LoadImagePath"
+      },
+      "widgets_values": [
+        "woman.jpg",
+        "image"
+      ]
+    },
+    {
+      "id": 3,
+      "type": "VHS_LoadVideoPath",
+      "pos": [
+        1350.136962890625,
+        331.20361328125
+      ],
+      "size": [
+        231.8896484375,
+        286
+      ],
+      "flags": {},
+      "order": 6,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "meta_batch",
+          "shape": 7,
+          "type": "VHS_BatchManager",
+          "link": null
+        },
+        {
+          "name": "vae",
+          "shape": 7,
+          "type": "VAE",
+          "link": null
+        },
+        {
+          "name": "video",
+          "type": "STRING",
+          "widget": {
+            "name": "video"
+          },
+          "link": 4
+        }
+      ],
+      "outputs": [
+        {
+          "name": "IMAGE",
+          "type": "IMAGE",
+          "links": [
+            5
+          ]
+        },
+        {
+          "name": "frame_count",
+          "type": "INT",
+          "links": null
+        },
+        {
+          "name": "audio",
+          "type": "AUDIO",
+          "links": null
+        },
+        {
+          "name": "video_info",
+          "type": "VHS_VIDEOINFO",
+          "links": null
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "VHS_LoadVideoPath"
+      },
+      "widgets_values": {
+        "video": "",
+        "force_rate": 0,
+        "custom_width": 0,
+        "custom_height": 0,
+        "frame_load_cap": 0,
+        "skip_first_frames": 0,
+        "select_every_nth": 1,
+        "format": "Wan",
+        "videopreview": {
+          "hidden": false,
+          "paused": false,
+          "params": {
+            "filename": "",
+            "type": "path",
+            "format": "video/",
+            "force_rate": 0,
+            "custom_width": 0,
+            "custom_height": 0,
+            "frame_load_cap": 0,
+            "skip_first_frames": 0,
+            "select_every_nth": 1
+          }
+        }
+      }
+    },
+    {
+      "id": 8,
+      "type": "VHS_VideoCombine",
+      "pos": [
+        1668.3499755859375,
+        328.22625732421875
+      ],
+      "size": [
+        214.7587890625,
+        334
+      ],
+      "flags": {},
+      "order": 7,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "images",
+          "type": "IMAGE",
+          "link": 5
+        },
+        {
+          "name": "audio",
+          "shape": 7,
+          "type": "AUDIO",
+          "link": null
+        },
+        {
+          "name": "meta_batch",
+          "shape": 7,
+          "type": "VHS_BatchManager",
+          "link": null
+        },
+        {
+          "name": "vae",
+          "shape": 7,
+          "type": "VAE",
+          "link": null
+        }
+      ],
+      "outputs": [
+        {
+          "name": "Filenames",
+          "type": "VHS_FILENAMES",
+          "links": null
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "VHS_VideoCombine"
+      },
+      "widgets_values": {
+        "frame_rate": 24,
+        "loop_count": 0,
+        "filename_prefix": "",
+        "format": "video/h264-mp4",
+        "pix_fmt": "yuv420p",
+        "crf": 19,
+        "save_metadata": true,
+        "trim_to_audio": false,
+        "pingpong": false,
+        "save_output": false,
+        "videopreview": {
+          "hidden": false,
+          "paused": false,
+          "params": {}
+        }
+      }
+    },
+    {
+      "id": 4,
+      "type": "VAEConfig",
+      "pos": [
+        374.2159423828125,
+        554.85888671875
+      ],
+      "size": [
+        334.080078125,
+        322
+      ],
+      "flags": {},
+      "order": 1,
+      "mode": 0,
+      "inputs": [],
+      "outputs": [
+        {
+          "name": "vae_config",
+          "type": "VAE_CONFIG",
+          "links": [
+            2
+          ]
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "VAEConfig"
+      },
+      "widgets_values": [
+        true,
+        true,
+        256,
+        256,
+        16,
+        192,
+        192,
+        12,
+        0,
+        true,
+        true,
+        true
+      ],
+      "auto_widget_states": {
+        "load_encoder": {
+          "isAuto": true,
+          "value": true,
+          "cachedValue": true
+        },
+        "load_decoder": {
+          "isAuto": true,
+          "value": true,
+          "cachedValue": true
+        },
+        "tile_sample_min_height": {
+          "isAuto": true,
+          "value": 256,
+          "cachedValue": 256
+        },
+        "tile_sample_min_width": {
+          "isAuto": true,
+          "value": 256,
+          "cachedValue": 256
+        },
+        "tile_sample_min_num_frames": {
+          "isAuto": true,
+          "value": 16,
+          "cachedValue": 16
+        },
+        "tile_sample_stride_height": {
+          "isAuto": true,
+          "value": 192,
+          "cachedValue": 192
+        },
+        "tile_sample_stride_width": {
+          "isAuto": true,
+          "value": 192,
+          "cachedValue": 192
+        },
+        "tile_sample_stride_num_frames": {
+          "isAuto": true,
+          "value": 12,
+          "cachedValue": 12
+        },
+        "blend_num_frames": {
+          "isAuto": true,
+          "value": 0,
+          "cachedValue": 0
+        },
+        "use_tiling": {
+          "isAuto": true,
+          "value": true,
+          "cachedValue": true
+        },
+        "use_temporal_tiling": {
+          "isAuto": true,
+          "value": true,
+          "cachedValue": true
+        },
+        "use_parallel_tiling": {
+          "isAuto": true,
+          "value": true,
+          "cachedValue": true
+        }
+      }
+    },
+    {
+      "id": 2,
+      "type": "InferenceArgs",
+      "pos": [
+        411.46307373046875,
+        178.18182373046875
+      ],
+      "size": [
+        278.73828125,
+        298
+      ],
+      "flags": {},
+      "order": 4,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "image_path",
+          "shape": 7,
+          "type": "STRING",
+          "widget": {
+            "name": "image_path"
+          },
+          "link": 1
+        }
+      ],
+      "outputs": [
+        {
+          "name": "inference_args",
+          "type": "INFERENCE_ARGS",
+          "links": [
+            3
+          ]
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "InferenceArgs"
+      },
+      "widgets_values": [
+        832,
+        480,
+        45,
+        20,
+        1,
+        17,
+        1024,
+        "fixed",
+        24,
+        "X://insert/path/here.mp4",
+        true
+      ],
+      "auto_widget_states": {
+        "height": {
+          "isAuto": false,
+          "value": 832,
+          "cachedValue": 720
+        },
+        "width": {
+          "isAuto": false,
+          "value": 480,
+          "cachedValue": 1280
+        },
+        "num_frames": {
+          "isAuto": false,
+          "value": 45,
+          "cachedValue": 45
+        },
+        "num_inference_steps": {
+          "isAuto": false,
+          "value": 20,
+          "cachedValue": 6
+        },
+        "guidance_scale": {
+          "isAuto": true,
+          "value": 1,
+          "cachedValue": 1
+        },
+        "flow_shift": {
+          "isAuto": true,
+          "value": 17,
+          "cachedValue": 17
+        },
+        "seed": {
+          "isAuto": false,
+          "value": 1024,
+          "cachedValue": 1024
+        },
+        "fps": {
+          "isAuto": false,
+          "value": 24,
+          "cachedValue": 24
+        },
+        "image_path": {
+          "isAuto": true,
+          "value": "X://insert/path/here.mp4",
+          "cachedValue": "X://insert/path/here.mp4"
+        },
+        "enable_teacache": {
+          "isAuto": true,
+          "value": true,
+          "cachedValue": true
+        }
+      }
+    },
+    {
+      "id": 1,
+      "type": "VideoGenerator",
+      "pos": [
+        818.804931640625,
+        348.9299621582031
+      ],
+      "size": [
+        400,
+        436
+      ],
+      "flags": {},
+      "order": 5,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "inference_args",
+          "shape": 7,
+          "type": "INFERENCE_ARGS",
+          "link": 3
+        },
+        {
+          "name": "vae_config",
+          "shape": 7,
+          "type": "VAE_CONFIG",
+          "link": 2
+        },
+        {
+          "name": "text_encoder_config",
+          "shape": 7,
+          "type": "TEXT_ENCODER_CONFIG",
+          "link": 7
+        },
+        {
+          "name": "dit_config",
+          "shape": 7,
+          "type": "DIT_CONFIG",
+          "link": 6
+        }
+      ],
+      "outputs": [
+        {
+          "name": "video_path",
+          "type": "STRING",
+          "links": [
+            4
+          ]
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "VideoGenerator"
+      },
+      "widgets_values": [
+        "A woman crying from laughter.",
+        "/workspace/ComfyUI/outputs_video/",
+        4,
+        "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers",
+        6,
+        2,
+        2,
+        "fp16",
+        true,
+        true,
+        "fp16",
+        "fp16",
+        true
+      ],
+      "auto_widget_states": {
+        "embedded_cfg_scale": {
+          "isAuto": true,
+          "value": 6,
+          "cachedValue": 6
+        },
+        "sp_size": {
+          "isAuto": true,
+          "value": 2,
+          "cachedValue": 2
+        },
+        "tp_size": {
+          "isAuto": true,
+          "value": 2,
+          "cachedValue": 2
+        },
+        "vae_precision": {
+          "isAuto": true,
+          "value": "fp16",
+          "cachedValue": "fp16"
+        },
+        "vae_tiling": {
+          "isAuto": true,
+          "value": true,
+          "cachedValue": true
+        },
+        "vae_sp": {
+          "isAuto": true,
+          "value": true,
+          "cachedValue": true
+        },
+        "text_encoder_precision": {
+          "isAuto": true,
+          "value": "fp16",
+          "cachedValue": "fp16"
+        },
+        "precision": {
+          "isAuto": true,
+          "value": "fp16",
+          "cachedValue": "fp16"
+        },
+        "use_cpu_offload": {
+          "isAuto": true,
+          "value": true,
+          "cachedValue": true
+        }
+      }
+    },
+    {
+      "id": 5,
+      "type": "TextEncoderConfig",
+      "pos": [
+        416.4937744140625,
+        953.6171875
+      ],
+      "size": [
+        270,
+        106
+      ],
+      "flags": {},
+      "order": 2,
+      "mode": 0,
+      "inputs": [],
+      "outputs": [
+        {
+          "name": "text_encoder_config",
+          "type": "TEXT_ENCODER_CONFIG",
+          "links": [
+            7
+          ]
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "TextEncoderConfig"
+      },
+      "widgets_values": [
+        "",
+        "",
+        ""
+      ],
+      "auto_widget_states": {
+        "prefix": {
+          "isAuto": true,
+          "value": "",
+          "cachedValue": ""
+        },
+        "quant_config": {
+          "isAuto": true,
+          "value": "",
+          "cachedValue": ""
+        },
+        "lora_config": {
+          "isAuto": true,
+          "value": "",
+          "cachedValue": ""
+        }
+      }
+    },
+    {
+      "id": 6,
+      "type": "DITConfig",
+      "pos": [
+        415.1928405761719,
+        1154.1573486328125
+      ],
+      "size": [
+        270,
+        82
+      ],
+      "flags": {},
+      "order": 3,
+      "mode": 0,
+      "inputs": [],
+      "outputs": [
+        {
+          "name": "dit_config",
+          "type": "DIT_CONFIG",
+          "links": [
+            6
+          ]
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "DITConfig"
+      },
+      "widgets_values": [
+        "",
+        ""
+      ],
+      "auto_widget_states": {
+        "prefix": {
+          "isAuto": true,
+          "value": "",
+          "cachedValue": ""
+        },
+        "quant_config": {
+          "isAuto": true,
+          "value": "",
+          "cachedValue": ""
+        }
+      }
+    }
+  ],
+  "links": [
+    [
+      1,
+      7,
+      0,
+      2,
+      0,
+      "STRING"
+    ],
+    [
+      2,
+      4,
+      0,
+      1,
+      1,
+      "VAE_CONFIG"
+    ],
+    [
+      3,
+      2,
+      0,
+      1,
+      0,
+      "INFERENCE_ARGS"
+    ],
+    [
+      4,
+      1,
+      0,
+      3,
+      2,
+      "STRING"
+    ],
+    [
+      5,
+      3,
+      0,
+      8,
+      0,
+      "IMAGE"
+    ],
+    [
+      6,
+      6,
+      0,
+      1,
+      3,
+      "DIT_CONFIG"
+    ],
+    [
+      7,
+      5,
+      0,
+      1,
+      2,
+      "TEXT_ENCODER_CONFIG"
+    ]
+  ],
+  "groups": [],
+  "config": {},
+  "extra": {
+    "ds": {
+      "scale": 0.8264462809917354,
+      "offset": [
+        646.7950212991898,
+        66.17259910028655
+      ]
+    },
+    "frontendVersion": "1.20.4",
+    "VHS_latentpreview": false,
+    "VHS_latentpreviewrate": 0,
+    "VHS_MetadataImage": true,
+    "VHS_KeepIntermediate": true
+  },
+  "version": 0.4
+}
@@ -0,0 +1,31 @@
+class DITConfig:
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "optional": {
+                "prefix": ("STRING", {
+                    "default": ""
+                }),
+                "quant_config": ("STRING", {
+                    "default": ""
+                }),
+            }
+        }
+
+    @classmethod
+    def VALIDATE_INPUTS(cls, **kwargs):
+        return True
+
+    RETURN_TYPES = ("DIT_CONFIG", )
+    RETURN_NAMES = ("dit_config", )
+    FUNCTION = "set_args"
+    CATEGORY = "fastvideo"
+
+    def set_args(self, prefix, quant_config):
+        raw_args = {"prefix": prefix, "quant_config": quant_config}
+
+        # Filter out keys where value is -99999
+        args = {k: v for k, v in raw_args.items() if str(int(v)) != str(-99999)}
+
+        return (args, )
@@ -0,0 +1,89 @@
+class InferenceArgs:
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "optional": {
+                "height": ("INT", {
+                    "default": 720
+                }),
+                "width": ("INT", {
+                    "default": 1280
+                }),
+                "num_frames": ("INT", {
+                    "default": 45
+                }),
+                "num_inference_steps": ("INT", {
+                    "default": 6
+                }),
+                "guidance_scale": ("FLOAT", {
+                    "default": 1.0
+                }),
+                "flow_shift": ("INT", {
+                    "default": 17
+                }),
+                "seed": ("INT", {
+                    "default": 1024
+                }),
+                "fps": ("INT", {
+                    "default": 24
+                }),
+                "image_path": ("STRING", {
+                    "default": "X://insert/path/here.mp4"
+                }),
+                "enable_teacache": ([True, False], {
+                    "default": False
+                }),
+            }
+        }
+
+    @classmethod
+    def VALIDATE_INPUTS(cls, **kwargs):
+        return True
+
+    RETURN_TYPES = ("INFERENCE_ARGS", )
+    RETURN_NAMES = ("inference_args", )
+    FUNCTION = "set_args"
+    CATEGORY = "fastvideo"
+
+    def set_args(
+        self,
+        height,
+        width,
+        num_frames,
+        num_inference_steps,
+        guidance_scale,
+        flow_shift,
+        seed,
+        fps,
+        image_path,
+        enable_teacache,
+    ):
+        raw_args = {
+            "height": height,
+            "width": width,
+            "num_frames": num_frames,
+            "num_inference_steps": num_inference_steps,
+            "guidance_scale": guidance_scale,
+            "flow_shift": flow_shift,
+            "seed": seed,
+            "fps": fps,
+            "image_path": image_path,
+            "enable_teacache": enable_teacache,
+        }
+
+        # Filter out keys where value is -99999, handling different types properly
+        args = {}
+        for k, v in raw_args.items():
+            try:
+                if isinstance(v, str):
+                    if v != "-99999":
+                        args[k] = v
+                elif v != -99999:
+                    # If it's not a string, compare directly
+                    args[k] = v
+            except (ValueError, TypeError):
+                # Include any value that causes an error in comparison
+                args[k] = v
+
+        return (args, )
@@ -0,0 +1,103 @@
+import hashlib
+import os
+
+import folder_paths
+import numpy as np
+import torch
+from PIL import Image, ImageOps, ImageSequence
+
+from .node_helpers import pillow
+
+
+class LoadImagePath:
+
+    @classmethod
+    def INPUT_TYPES(s):
+        input_dir = folder_paths.get_input_directory()
+        files = [
+            f for f in os.listdir(input_dir)
+            if os.path.isfile(os.path.join(input_dir, f))
+        ]
+        files = folder_paths.filter_files_content_types(files, ["image"])
+        return {
+            "required": {
+                "image": (sorted(files), {
+                    "image_upload": True
+                })
+            },
+        }
+
+    CATEGORY = "fastvideo"
+
+    RETURN_TYPES = ("STRING", "IMAGE", "MASK")
+    RETURN_NAMES = ("image_path", "IMAGE", "MASK")
+    FUNCTION = "load_image"
+
+    def load_image(self, image):
+        image_path = folder_paths.get_annotated_filepath(image)
+
+        img = pillow(Image.open, image_path)
+
+        output_images: list[torch.Tensor] = []
+        output_masks: list[torch.Tensor] = []
+        w, h = None, None
+
+        excluded_formats = ['MPO']
+
+        for i in ImageSequence.Iterator(img):
+            processed_image = pillow(ImageOps.exif_transpose, i)
+            if processed_image is None:
+                continue
+
+            if processed_image.mode == 'I':
+                processed_image = processed_image.point(lambda i: i * (1 / 255))
+            image = processed_image.convert("RGB")
+
+            if len(output_images) == 0:
+                w = image.size[0]
+                h = image.size[1]
+
+            if image.size[0] != w or image.size[1] != h:
+                continue
+
+            image = np.array(image).astype(np.float32) / 255.0
+            image = torch.from_numpy(image)[
+                None,
+            ]
+            if 'A' in processed_image.getbands():
+                mask = np.array(processed_image.getchannel('A')).astype(
+                    np.float32) / 255.0
+                mask = 1. - torch.from_numpy(mask)
+            elif processed_image.mode == 'P' and 'transparency' in processed_image.info:
+                mask = np.array(
+                    processed_image.convert('RGBA').getchannel('A')).astype(
+                        np.float32) / 255.0
+                mask = 1. - torch.from_numpy(mask)
+            else:
+                mask = torch.zeros((64, 64), dtype=torch.float32, device="cpu")
+            output_images.append(image)
+            output_masks.append(mask.unsqueeze(0))
+
+        if len(output_images) > 1 and img.format not in excluded_formats:
+            output_image = torch.cat(output_images, dim=0)
+            output_mask = torch.cat(output_masks, dim=0)
+        else:
+            output_image = output_images[0]
+            output_mask = output_masks[0]
+
+        return (image_path, output_image, output_mask)
+
+    @classmethod
+    def IS_CHANGED(s, image):
+        image_path = folder_paths.get_annotated_filepath(image)
+        m = hashlib.sha256()
+        with open(image_path, 'rb') as f:
+            m.update(f.read())
+        return m.digest().hex()
+
+    @classmethod
+    def VALIDATE_INPUTS(s, image):
+        if not folder_paths.exists_annotated_filepath(image):
+            return "Invalid image file: {}".format(image)
+
+        return True
@@ -0,0 +1,68 @@
+import hashlib
+from collections.abc import Callable
+from typing import Any, TypeVar
+
+import torch
+from comfy.cli_args import args
+from PIL import ImageFile, UnidentifiedImageError
+
+T = TypeVar('T')
+
+
+def conditioning_set_values(conditioning: list[Any],
+                            values: dict[str, Any] | None = None) -> list[Any]:
+    if values is None:
+        values = {}
+    c = []
+    for t in conditioning:
+        n = [t[0], t[1].copy()]
+        for k in values:
+            n[1][k] = values[k]
+        c.append(n)
+
+    return c
+
+
+def pillow(fn: Callable[[Any], T], arg: Any) -> T:
+    prev_value = None
+    try:
+        x = fn(arg)
+    except (OSError, UnidentifiedImageError, ValueError
+            ):  #PIL issues #4472 and #2445, also fixes ComfyUI issue #3416
+        prev_value = ImageFile.LOAD_TRUNCATED_IMAGES
+        ImageFile.LOAD_TRUNCATED_IMAGES = True
+        x = fn(arg)
+    finally:
+        if prev_value is not None:
+            ImageFile.LOAD_TRUNCATED_IMAGES = prev_value
+    return x
+
+
+def hasher() -> Callable[[], Any]:
+    hashfuncs = {
+        "md5": hashlib.md5,
+        "sha1": hashlib.sha1,
+        "sha256": hashlib.sha256,
+        "sha512": hashlib.sha512
+    }
+    return hashfuncs[args.default_hashing_function]
+
+
+def string_to_torch_dtype(string: str) -> torch.dtype | None:
+    if string == "fp32":
+        return torch.float32
+    if string == "fp16":
+        return torch.float16
+    if string == "bf16":
+        return torch.bfloat16
+    return None
+
+
+def image_alpha_fix(destination: torch.Tensor,
+                    source: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    if destination.shape[-1] < source.shape[-1]:
+        source = source[..., :destination.shape[-1]]
+    elif destination.shape[-1] > source.shape[-1]:
+        destination = torch.nn.functional.pad(destination, (0, 1))
+        destination[..., -1] = 1.0
+    return destination, source
@@ -0,0 +1,24 @@
+from .dit_config import DITConfig
+from .inference_args import InferenceArgs
+from .load_image import LoadImagePath
+from .text_encoder_config import TextEncoderConfig
+from .vae_config import VAEConfig
+from .video_generator import VideoGenerator
+
+NODE_CLASS_MAPPINGS = {
+    "VideoGenerator": VideoGenerator,
+    "InferenceArgs": InferenceArgs,
+    "VAEConfig": VAEConfig,
+    "TextEncoderConfig": TextEncoderConfig,
+    "DITConfig": DITConfig,
+    "LoadImagePath": LoadImagePath
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "VideoGenerator": "Video Generator",
+    "InferenceArgs": "Inference Args",
+    "VAEConfig": "VAE Config",
+    "TextEncoderConfig": "Text Encoder Config",
+    "DITConfig": "DIT Config",
+    "LoadImagePath": "Load Image Path"
+}
@@ -0,0 +1,38 @@
+class TextEncoderConfig:
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "optional": {
+                "prefix": ("STRING", {
+                    "default": ""
+                }),
+                "quant_config": ("STRING", {
+                    "default": ""
+                }),
+                "lora_config": ("STRING", {
+                    "default": ""
+                }),
+            }
+        }
+
+    @classmethod
+    def VALIDATE_INPUTS(cls, **kwargs):
+        return True
+
+    RETURN_TYPES = ("TEXT_ENCODER_CONFIG", )
+    RETURN_NAMES = ("text_encoder_config", )
+    FUNCTION = "set_args"
+    CATEGORY = "fastvideo"
+
+    def set_args(self, prefix, quant_config, lora_config):
+        raw_args = {
+            "prefix": prefix,
+            "quant_config": quant_config,
+            "lora_config": lora_config
+        }
+
+        # Filter out keys where value is -99999
+        args = {k: v for k, v in raw_args.items() if str(int(v)) != str(-99999)}
+
+        return (args, )
@@ -0,0 +1,88 @@
+class VAEConfig:
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "optional": {
+                "load_encoder": ([True, False], {
+                    "default": True
+                }),
+                "load_decoder": ([True, False], {
+                    "default": True
+                }),
+                "tile_sample_min_height": ("INT", {
+                    "default": 256
+                }),
+                "tile_sample_min_width": ("INT", {
+                    "default": 256
+                }),
+                "tile_sample_min_num_frames": ("INT", {
+                    "default": 16
+                }),
+                "tile_sample_stride_height": ("INT", {
+                    "default": 192
+                }),
+                "tile_sample_stride_width": ("INT", {
+                    "default": 192
+                }),
+                "tile_sample_stride_num_frames": ("INT", {
+                    "default": 12
+                }),
+                "blend_num_frames": ("INT", {
+                    "default": 0
+                }),
+                "use_tiling": ([True, False], {
+                    "default": True
+                }),
+                "use_temporal_tiling": ([True, False], {
+                    "default": True
+                }),
+                "use_parallel_tiling": ([True, False], {
+                    "default": True
+                }),
+            }
+        }
+
+    @classmethod
+    def VALIDATE_INPUTS(cls, **kwargs):
+        return True
+
+    RETURN_TYPES = ("VAE_CONFIG", )
+    RETURN_NAMES = ("vae_config", )
+    FUNCTION = "set_args"
+    CATEGORY = "fastvideo"
+
+    def set_args(
+        self,
+        load_encoder,
+        load_decoder,
+        tile_sample_min_height,
+        tile_sample_min_width,
+        tile_sample_min_num_frames,
+        tile_sample_stride_height,
+        tile_sample_stride_width,
+        tile_sample_stride_num_frames,
+        blend_num_frames,
+        use_tiling,
+        use_temporal_tiling,
+        use_parallel_tiling,
+    ):
+        raw_args = {
+            "load_encoder": load_encoder,
+            "load_decoder": load_decoder,
+            "tile_sample_min_height": tile_sample_min_height,
+            "tile_sample_min_width": tile_sample_min_width,
+            "tile_sample_min_num_frames": tile_sample_min_num_frames,
+            "tile_sample_stride_height": tile_sample_stride_height,
+            "tile_sample_stride_width": tile_sample_stride_width,
+            "tile_sample_stride_num_frames": tile_sample_stride_num_frames,
+            "blend_num_frames": blend_num_frames,
+            "use_tiling": use_tiling,
+            "use_temporal_tiling": use_temporal_tiling,
+            "use_parallel_tiling": use_parallel_tiling,
+        }
+
+        # Filter out any value explicitly set to -99999
+        args = {k: v for k, v in raw_args.items() if str(int(v)) != str(-99999)}
+
+        return (args, )
@@ -0,0 +1,315 @@
+from __future__ import annotations
+
+import glob
+import os
+import signal
+import sys
+import threading
+import time
+from typing import Any
+
+from comfy.model_management import processing_interrupted
+
+from fastvideo import PipelineConfig
+from fastvideo import VideoGenerator as FastVideoGenerator
+
+sys.path.insert(
+    0,
+    os.path.dirname(
+        os.path.dirname(
+            os.path.dirname(os.path.dirname(os.path.abspath(__file__))))))
+
+
+# Custom exception for interruption
+class GenerationInterruptedException(Exception):
+    pass
+
+
+# Custom exception for interruption that ComfyUI will recognize
+class GenerationCancelledException(Exception):
+
+    def __init__(self,
+                 message: str = "Generation was cancelled by user") -> None:
+        self.message = message
+        super().__init__(self.message)
+
+
+def update_config_from_args(config: Any, args_dict: dict[str, Any]) -> None:
+    """
+    Update configuration object from arguments dictionary.
+    
+    Args:
+        config: The configuration object to update
+        args_dict: Dictionary containing arguments
+    """
+    for key, value in args_dict.items():
+        if hasattr(config, key) and value is not None:
+            if key == "text_encoder_precisions" and isinstance(value, list):
+                setattr(config, key, tuple(value))
+            else:
+                setattr(config, key, value)
+
+
+class VideoGenerator:
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "prompt": ("STRING", {
+                    "multiline":
+                    True,
+                    "default":
+                    "A ripe orange tumbles gently from a tree and lands on the head of a lounging capybara, "
+                    "who blinks slowly in response. The moment is quietly humorous and oddly serene, framed by "
+                    "lush green foliage and dappled sunlight. Mid-shot, warm and whimsical tones."
+                }),
+                "output_path": ("STRING", {
+                    "default": "/workspace/ComfyUI/outputs_video/"
+                }),
+                "num_gpus": ("INT", {
+                    "default": 2,
+                    "min": 1,
+                    "max": 16
+                }),
+                "model_path": ("STRING", {
+                    "default": "FastVideo/FastHunyuan-diffusers"
+                })
+            },
+            "optional": {
+                "inference_args": ("INFERENCE_ARGS", ),
+                "embedded_cfg_scale": ("FLOAT", {
+                    "default": 6.0
+                }),
+                "sp_size": ("INT", {
+                    "default": 2
+                }),
+                "tp_size": ("INT", {
+                    "default": 2
+                }),
+                "vae_config": ("VAE_CONFIG", ),
+                "vae_precision": (["fp16", "bf16"], {
+                    "default": "fp16"
+                }),
+                "vae_tiling": ([True, False], {
+                    "default": True
+                }),
+                "vae_sp": ([True, False], {
+                    "default": False
+                }),
+                "text_encoder_config": ("TEXT_ENCODER_CONFIG", ),
+                "text_encoder_precision": (["fp16", "bf16"], {
+                    "default": "fp16"
+                }),
+                "dit_config": ("DIT_CONFIG", ),
+                "precision": (["fp16", "bf16"], {
+                    "default": "fp16"
+                }),
+                "use_cpu_offload": ([True, False], {
+                    "default": False
+                }),
+            }
+        }
+
+    @classmethod
+    def VALIDATE_INPUTS(cls, **kwargs):
+        return True
+
+    RETURN_TYPES = ("STRING", )
+    RETURN_NAMES = ("video_path", )
+    FUNCTION = "launch_inference"
+    CATEGORY = "fastvideo"
+
+    generator: FastVideoGenerator | None = None
+    _interrupt_thread: threading.Thread | None = None
+    _generation_active: bool = False
+    _generation_interrupted: bool = False
+    _interrupt_event: threading.Event = threading.Event()
+    _generation_thread: threading.Thread | None = None
+    _generation_result: str | None = None
+    _generation_exception: Exception | None = None
+
+    def _monitor_for_interruption(self):
+        """Background thread that monitors for interruption requests"""
+        time.sleep(2)  # Give the generation thread time to send execute_forward
+
+        while self._generation_active and not self._interrupt_event.is_set():
+            if processing_interrupted():
+                print("Video generation interrupted by user")
+                self._generation_interrupted = True
+
+                # Try to send interrupt signal to worker processes
+                if self.generator is not None and hasattr(
+                        self.generator, 'executor'):
+                    try:
+                        # The MultiprocExecutor has a workers attribute
+                        if hasattr(self.generator.executor, 'workers'):
+                            for worker in self.generator.executor.workers:
+                                if worker.is_alive():
+                                    os.kill(worker.pid, signal.SIGINT)
+                            print("Interrupt signal sent to worker processes")
+                    except Exception as e:
+                        print(f"Error sending interrupt signal: {e}")
+
+                # Set the interrupt event to notify other threads
+                self._interrupt_event.set()
+                break
+            time.sleep(0.5)
+
+    def _run_generation(self, prompt: str, output_path: str,
+                        inference_args: dict[str, Any]) -> None:
+        """Thread function to run the generation"""
+        try:
+            if self.generator is not None:
+                self.generator.generate_video(prompt=prompt,
+                                              output_path=output_path,
+                                              **inference_args)
+                self._generation_result = os.path.join(output_path,
+                                                       f"{prompt[:100]}.mp4")
+            else:
+                raise RuntimeError("Generator is not initialized")
+        except Exception as e:
+            self._generation_exception = e
+            self._interrupt_event.set()
+
+    def load_output_video(self, output_dir):
+        video_extensions = ["*.mp4", "*.avi", "*.mov", "*.mkv"]
+        video_files = []
+
+        for ext in video_extensions:
+            video_files.extend(glob.glob(os.path.join(output_dir, ext)))
+
+        if not video_files:
+            print("No video files found in output directory: %s", output_dir)
+            return ""
+
+        video_files.sort()
+        return video_files[0]
+
+    def launch_inference(
+        self,
+        prompt,
+        output_path,
+        num_gpus,
+        model_path,
+        embedded_cfg_scale,
+        sp_size,
+        tp_size,
+        vae_precision,
+        vae_tiling,
+        vae_sp,
+        text_encoder_precision,
+        precision,
+        inference_args=None,
+        vae_config=None,
+        text_encoder_config=None,
+        dit_config=None,
+        use_cpu_offload=None,
+    ):
+        print('Running FastVideo inference')
+
+        # Reset interruption flag and event
+        self._generation_interrupted = False
+        self._interrupt_event.clear()
+        self._generation_result = None
+        self._generation_exception = None
+
+        # Load pipeline config from model path
+        pipeline_config = PipelineConfig.from_pretrained(model_path)
+        print('pipeline_config', pipeline_config)
+
+        # Update configs with provided config dictionaries
+        if dit_config is not None:
+            update_config_from_args(pipeline_config.dit_config, dit_config)
+
+        if vae_config is not None:
+            update_config_from_args(pipeline_config.vae_config, vae_config)
+
+        if text_encoder_config is not None:
+            update_config_from_args(pipeline_config.text_encoder_configs,
+                                    text_encoder_config)
+
+        # Update top-level pipeline config with remaining arguments
+        raw_pipeline_args = {}
+        if embedded_cfg_scale is not None:
+            raw_pipeline_args['embedded_cfg_scale'] = embedded_cfg_scale
+        if precision is not None:
+            raw_pipeline_args['precision'] = precision
+        if vae_precision is not None:
+            raw_pipeline_args['vae_precision'] = vae_precision
+        if vae_tiling is not None:
+            raw_pipeline_args['vae_tiling'] = vae_tiling
+        if vae_sp is not None:
+            raw_pipeline_args['vae_sp'] = vae_sp
+        if text_encoder_precision is not None:
+            raw_pipeline_args['text_encoder_precision'] = text_encoder_precision
+
+        # Filter out any value explicitly set to -99999 (auto values)
+        pipeline_args = {
+            k: v
+            for k, v in raw_pipeline_args.items() if str(int(v)) != str(-99999)
+        }
+
+        update_config_from_args(pipeline_config, pipeline_args)
+
+        raw_generation_args = {}
+        if num_gpus is not None:
+            raw_generation_args['num_gpus'] = num_gpus
+        if tp_size is not None:
+            raw_generation_args['tp_size'] = tp_size
+        if sp_size is not None:
+            raw_generation_args['sp_size'] = sp_size
+        if use_cpu_offload is not None:
+            raw_generation_args['use_cpu_offload'] = use_cpu_offload
+
+        generation_args = {
+            k: v
+            for k, v in raw_generation_args.items()
+            if str(int(v)) != str(-99999)
+        }
+
+        if self.generator is None:
+            print('generation_args', generation_args)
+            print('pipeline_config', pipeline_config)
+            self.generator = FastVideoGenerator.from_pretrained(
+                model_path=model_path,
+                **generation_args,
+                pipeline_config=pipeline_config)
+
+        print('inference_args', inference_args)
+
+        # Start a thread to run the generation
+        self._generation_thread = threading.Thread(target=self._run_generation,
+                                                   args=(prompt, output_path,
+                                                         inference_args),
+                                                   daemon=True)
+        self._generation_thread.start()
+
+        # Start a background thread to monitor for interruptions
+        self._generation_active = True
+        self._interrupt_thread = threading.Thread(
+            target=self._monitor_for_interruption, daemon=True)
+        self._interrupt_thread.start()
+
+        # Wait for either completion or interruption
+        while self._generation_thread.is_alive(
+        ) and not self._interrupt_event.is_set():
+            self._generation_thread.join(timeout=0.5)
+
+        self._generation_active = False
+        if self._interrupt_thread:
+            self._interrupt_thread.join(timeout=1.0)
+            self._interrupt_thread = None
+
+        if self._generation_interrupted:
+            print("Video generation was cancelled by user")
+            raise GenerationCancelledException()
+        elif self._generation_exception:
+            # Re-raise the exception from the generation thread
+            raise self._generation_exception
+        elif self._generation_result:
+            return (self._generation_result, )
+        else:
+            # This shouldn't happen, but just in case
+            print("Generation completed but no result was produced")
+            raise Exception("Generation failed to produce a result")
@@ -0,0 +1,593 @@
+import { app } from '../../../scripts/app.js'
+
+function chainCallback(object, property, callback) {
+    if (object == undefined) {
+        console.error("Tried to add callback to non-existent object");
+        return;
+    }
+    if (property in object && object[property]) {
+        const callback_orig = object[property];
+        object[property] = function () {
+            const r = callback_orig.apply(this, arguments);
+            return callback.apply(this, arguments) ?? r;
+        };
+    } else {
+        object[property] = callback;
+    }
+}
+
+function drawAutoAnnotated(ctx, node, widget_width, y, H) {
+    const litegraph_base = LiteGraph;
+    const show_text = app.canvas.ds.scale >= 0.5;
+    const margin = 15;
+
+    const autoTextWidth = 30;
+    const autoTextRightMargin = 5;
+
+    ctx.textAlign = 'left';
+    ctx.strokeStyle = litegraph_base.WIDGET_OUTLINE_COLOR;
+    ctx.fillStyle = litegraph_base.WIDGET_BGCOLOR;
+
+    ctx.beginPath();
+    if (show_text && ctx.roundRect) {
+        ctx.roundRect(margin, y, widget_width - margin * 2, H, [H * 0.5]);
+    } else {
+        ctx.rect(margin, y, widget_width - margin * 2, H);
+    }
+    ctx.fill();
+
+    if (show_text) {
+        if (!this.disabled) ctx.stroke();
+        const isAuto = this.isAuto === true;
+
+        ctx.save();
+        if (isAuto) {
+            ctx.fillStyle = litegraph_base.WIDGET_TEXT_COLOR;
+            ctx.strokeStyle = litegraph_base.WIDGET_TEXT_COLOR;
+        } else {
+            ctx.fillStyle = litegraph_base.WIDGET_SECONDARY_TEXT_COLOR;
+            ctx.strokeStyle = litegraph_base.WIDGET_SECONDARY_TEXT_COLOR;
+        }
+
+        // Position for the cog
+        const cogX = widget_width - autoTextRightMargin - autoTextWidth - 6;
+        const cogY = y + H * 0.5;
+        const cogRadius = 6;
+        const toothLength = 2;
+        const numTeeth = 8;
+        const holeRadius = 2; // Radius of the center hole
+
+        // Draw the cog
+        ctx.beginPath();
+        ctx.arc(cogX, cogY, cogRadius - toothLength, 0, Math.PI * 2);
+        ctx.fill();
+
+        // Draw the center hole (by clearing it)
+        ctx.beginPath();
+        ctx.arc(cogX, cogY, holeRadius, 0, Math.PI * 2);
+        ctx.fillStyle = litegraph_base.WIDGET_BGCOLOR;
+        ctx.fill();
+
+        // Reset fill style for the teeth
+        if (isAuto) {
+            ctx.fillStyle = litegraph_base.WIDGET_TEXT_COLOR;
+        } else {
+            ctx.fillStyle = litegraph_base.WIDGET_SECONDARY_TEXT_COLOR;
+        }
+
+        // Draw teeth
+        ctx.beginPath();
+        for (let i = 0; i < numTeeth; i++) {
+            const angle = (i / numTeeth) * Math.PI * 2;
+            const innerX = cogX + (cogRadius - toothLength) * Math.cos(angle);
+            const innerY = cogY + (cogRadius - toothLength) * Math.sin(angle);
+            const outerX = cogX + cogRadius * Math.cos(angle);
+            const outerY = cogY + cogRadius * Math.sin(angle);
+
+            ctx.moveTo(innerX, innerY);
+            ctx.lineTo(outerX, outerY);
+        }
+        ctx.lineWidth = 2;
+        ctx.stroke();
+
+        ctx.restore();
+
+        // Draw label
+        ctx.fillStyle = litegraph_base.WIDGET_SECONDARY_TEXT_COLOR;
+        const label = this.label || this.name;
+        if (label != null) {
+            ctx.fillText(label, margin * 2 + 5, y + H * 0.7);
+        }
+
+        // Draw value
+        ctx.textAlign = 'right';
+        const text = isAuto ? "auto" : this.displayValue();
+        ctx.fillStyle = isAuto ? litegraph_base.WIDGET_SECONDARY_TEXT_COLOR : litegraph_base.WIDGET_TEXT_COLOR;
+        ctx.fillText(text, widget_width - autoTextRightMargin - autoTextWidth - 15, y + H * 0.7);
+
+        // Draw increment/decrement buttons if not in AUTO mode and not a string widget
+        if (!isAuto && !this.disabled && this.config[0] !== "FVAUTOSTRING") {
+            // Draw decrement button (left triangle)
+            ctx.fillStyle = litegraph_base.WIDGET_TEXT_COLOR;
+            ctx.beginPath();
+            ctx.moveTo(margin + 16, y + 5);
+            ctx.lineTo(margin + 6, y + H * 0.5);
+            ctx.lineTo(margin + 16, y + H - 5);
+            ctx.fill();
+
+            // Draw increment button (right triangle)
+            ctx.beginPath();
+            ctx.moveTo(widget_width - margin - 16, y + 5);
+            ctx.lineTo(widget_width - margin - 6, y + H * 0.5);
+            ctx.lineTo(widget_width - margin - 16, y + H - 5);
+            ctx.fill();
+        }
+    }
+}
+
+function mouseAutoAnnotated(event, [x, y], node) {
+    const widget_width = node.size[0];
+    const margin = 15;
+    const H = 20; // Widget height
+
+    const autoTextWidth = 30;
+    const autoTextRightMargin = 5;
+
+    const cogRadius = 6;
+
+    if (this.isAuto) {
+        if (event.type === "pointerup" || event.type === "mouseup") {
+            const cogX = widget_width - autoTextRightMargin - autoTextWidth - 6;
+            const cogLeftEdge = cogX - cogRadius;
+            const cogRightEdge = cogX + cogRadius;
+
+            if (x > cogLeftEdge && x < cogRightEdge) {
+                this.isAuto = false;
+                this.value = this.cachedValue !== undefined ? this.cachedValue : (this.options.default || 0);
+
+                if (this.callback) {
+                    this.callback(this.value);
+                }
+                node.graph.setDirtyCanvas(true, false);
+            }
+        }
+
+        // Block ALL events in auto mode except cog clicks
+        event.preventDefault?.();
+        event.stopPropagation?.();
+        event.stopImmediatePropagation?.();
+        return true; // Always return true to indicate event was handled
+    }
+
+    // Determine if clicking on increment/decrement buttons
+    const delta = this.config[0] === "FVAUTOSTRING" ? 0 :
+        (x < 40 ? -1 : x > widget_width - 48 ? 1 : 0);
+
+    if (event.type === "pointerdown" || event.type === "mousedown") {
+        // ComfyUI appears to intercept pointerdown events, so this code path is never reached
+        console.log("pointerdown received (unexpected)");
+        return false;
+    } else if (event.type === "pointerup" || event.type === "mouseup") {
+        // Stop event propagation to prevent double handling
+        event.preventDefault?.();
+        event.stopPropagation?.();
+        event.stopImmediatePropagation?.();
+
+        const cogX = widget_width - autoTextRightMargin - autoTextWidth - 6;
+        const cogLeftEdge = widget_width - autoTextRightMargin - autoTextWidth - 6 - cogRadius;
+        const cogRightEdge = widget_width - autoTextRightMargin - autoTextWidth - 6 + cogRadius;
+
+        if (x > cogLeftEdge && x < cogRightEdge) {
+            this.isAuto = !this.isAuto;
+
+            if (this.isAuto) {
+                this.cachedValue = this.value;
+                this.value = -99999;
+            } else {
+                this.value = this.cachedValue !== undefined ? this.cachedValue : (this.options.default || 0);
+            }
+
+            if (this.callback) {
+                this.callback(this.value);
+            }
+
+            node.graph.setDirtyCanvas(true, false);
+            return true;
+        }
+
+        // If in auto mode and NOT clicking the cog, block all other interactions
+        if (this.isAuto) {
+            return true;
+        }
+
+        // Handle increment/decrement buttons if not in auto mode
+        if (delta !== 0 && !this.isAuto) {
+            if (this.config[0] === "FVAUTOCOMBO") {
+                const options = this.options.values || [];
+                if (options.length === 0) return true;
+
+                let currentIndex = -1;
+                for (let i = 0; i < options.length; i++) {
+                    const optValue = typeof options[i] === 'object' ? options[i].value : options[i];
+                    if (optValue == this.value || String(optValue) === String(this.value)) {
+                        currentIndex = i;
+                        break;
+                    }
+                }
+
+                if (currentIndex === -1) {
+                    currentIndex = 0;
+                }
+
+                let newIndex = currentIndex + delta;
+                if (newIndex < 0) {
+                    newIndex = options.length - 1;
+                } else if (newIndex >= options.length) {
+                    newIndex = 0;
+                }
+
+                const newOption = options[newIndex];
+                this.value = typeof newOption === 'object' ? newOption.value : newOption;
+
+                if (this.callback) {
+                    this.callback(this.value);
+                }
+
+                node.graph.setDirtyCanvas(true, false);
+                return true;
+            } else {
+                let v = parseFloat(this.value);
+                const increment = delta * 0.1 * (this.options.step || 1);
+
+                v += increment;
+
+                // Apply min/max constraints
+                if (this.options.min != null) {
+                    v = Math.max(this.options.min, v);
+                }
+                if (this.options.max != null) {
+                    v = Math.min(this.options.max, v);
+                }
+
+                // Round to precision or to integer
+                if (this.config[0] === "FVAUTOINT") {
+                    v = Math.round(v);
+                } else if (this.options.precision !== undefined) {
+                    const precision = Math.pow(10, this.options.precision);
+                    v = Math.round(v * precision) / precision;
+                }
+
+                this.value = v;
+
+                if (this.callback) {
+                    this.callback(this.value);
+                }
+
+                node.graph.setDirtyCanvas(true, false);
+                return true;
+            }
+        }
+
+        if (delta === 0 && !this.isAuto) {
+            if (this.config[0] === "FVAUTOCOMBO") {
+                const options = this.options.values || [];
+
+                // Create menu items
+                const menuItems = options.map(opt => {
+                    const value = typeof opt === 'object' ? opt.value : opt;
+                    const label = typeof opt === 'object' ? opt.label : opt.toString();
+
+                    return {
+                        content: label,
+                        callback: () => {
+                            this.value = value;
+
+                            if (this.callback) {
+                                this.callback(this.value);
+                            }
+
+                            node.graph.setDirtyCanvas(true, false);
+                        }
+                    };
+                });
+
+                new LiteGraph.ContextMenu(menuItems, {
+                    event: event,
+                    title: null,
+                    callback: null,
+                    extra: node
+                });
+
+                return true;
+            } else if (this.config[0] === "FVAUTOSTRING") {
+                const d_callback = (v) => {
+                    this.value = v;
+
+                    if (this.callback) {
+                        this.callback(this.value);
+                    }
+
+                    node.graph.setDirtyCanvas(true, false);
+                };
+
+                const dialog = app.canvas.prompt(
+                    'Value',
+                    this.value,
+                    d_callback,
+                    event
+                );
+
+                return true;
+            } else {
+                // For numeric widgets, show input dialog
+                const d_callback = (v) => {
+                    this.value = this.parseValue?.(v) ?? Number(v);
+
+                    // Apply min/max constraints
+                    if (this.options.min != null) {
+                        this.value = Math.max(this.options.min, this.value);
+                    }
+                    if (this.options.max != null) {
+                        this.value = Math.min(this.options.max, this.value);
+                    }
+
+                    // Round to precision or to integer
+                    if (this.config[0] === "FVAUTOINT") {
+                        this.value = Math.round(this.value);
+                    } else if (this.options.precision !== undefined) {
+                        const precision = Math.pow(10, this.options.precision);
+                        this.value = Math.round(this.value * precision) / precision;
+                    }
+
+                    if (this.callback) {
+                        this.callback(this.value);
+                    }
+
+                    node.graph.setDirtyCanvas(true, false);
+                };
+
+                const dialog = app.canvas.prompt(
+                    'Value',
+                    this.value,
+                    d_callback,
+                    event
+                );
+
+                return true;
+            }
+        }
+
+        return true;
+    }
+
+    return false;
+}
+
+function makeAutoAnnotated(widget, inputData) {
+    const original = {
+        callback: widget.callback,
+        type: widget.type,
+        value: widget.value
+    };
+
+    // Add AUTO properties to the widget
+    Object.assign(widget, {
+        type: "BOOLEAN",
+        draw: drawAutoAnnotated,
+        mouse: mouseAutoAnnotated,
+        onMouse: null, // Explicitly disable original onMouse handler
+
+        isAuto: true,
+        cachedValue: widget.value,
+        config: inputData,
+        options: Object.assign({}, inputData[1], widget.options),
+        original: original,  // Store original properties for reference
+
+        // Disable other potential mouse handlers with no-op functions
+        onClick: function () {
+            return false;
+        },
+        onPointerUp: function () {
+            return false;
+        },
+        onPointerDown: function () {
+            return false;
+        },
+        onMouseUp: function () {
+            return false;
+        },
+        onMouseDown: function () {
+            return false;
+        },
+
+        computeSize(width) {
+            return [width, 20];
+        },
+        displayValue: function () {
+            if (this.config[0] === "FVAUTOINT") {
+                return Math.round(this.value).toString();
+            }
+            if (this.config[0] === "FVAUTOCOMBO") {
+                return this.value;
+            }
+            if (this.config[0] === "FVAUTOSTRING") {
+                return this.value;
+            }
+            // For FLOAT values, check if it's actually an integer
+            if (Number.isInteger(this.value)) {
+                return this.value.toString();
+            }
+            return this.value.toFixed(this.options.precision || 2);
+        },
+        parseValue: function (v) {
+            if (this.config[0] === "FVAUTOSTRING") {
+                return v;
+            }
+            if (typeof v === "string") {
+                return parseFloat(v);
+            }
+            return v;
+        },
+        serializeValue: function () {
+            // Return special value for AUTO mode
+            return this.isAuto ? -99999 : this.value;
+        },
+        deserializeValue: function (data) {
+            if (data === -99999) {
+                this.isAuto = true;
+                this.value = -99999;
+            } else {
+                this.isAuto = false;
+                this.value = data;
+                this.cachedValue = data;
+            }
+        }
+    });
+
+    // Override callback to handle AUTO mode
+    widget.callback = function (v) {
+        if (this.isAuto) {
+            return; // Don't call the original callback in AUTO mode
+        }
+        const result = original.callback?.call(this, v);
+        return result;
+    };
+
+    // Override any potential click handlers
+    const originalOnClick = widget.onClick;
+    if (originalOnClick) {
+        widget.onClick = function (...args) {
+            if (this.isAuto) {
+                return false;
+            }
+            return originalOnClick.call(this, ...args);
+        };
+    }
+
+    return widget;
+}
+
+app.registerExtension({
+    name: "FastVideo.AutoWidgets",
+
+    async beforeRegisterNodeDef(nodeType, nodeData, app) {
+        if (nodeData?.name == "VideoGenerator" || nodeData?.name === "InferenceArgs" || nodeData?.name === "VAEConfig" ||
+            nodeData?.name === "TextEncoderConfig" || nodeData?.name === "DITConfig") {
+            // Add serialization support
+            chainCallback(nodeType.prototype, "onSerialize", function (info) {
+                if (!this.widgets) {
+                    return;
+                }
+                // Ensure widgets_values exists
+                if (!info.widgets_values) {
+                    info.widgets_values = {};
+                }
+
+                // Store AUTO widget states in a separate property
+                if (!info.auto_widget_states) {
+                    info.auto_widget_states = {};
+                }
+
+                // Handle AUTO widgets specially
+                for (const w of this.widgets) {
+                    if (w.type === "BOOLEAN" && w.isAuto !== undefined) {
+                        // Store the serialized value (for Python node)
+                        info.widgets_values[w.name] = w.serializeValue();
+
+                        // Store the full state (for UI restoration)
+                        info.auto_widget_states[w.name] = {
+                            isAuto: w.isAuto,
+                            value: w.value,
+                            cachedValue: w.cachedValue
+                        };
+                    }
+                }
+            });
+
+            // Add deserialization support
+            chainCallback(nodeType.prototype, "onConfigure", function (info) {
+                if (!this.widgets) {
+                    return;
+                }
+
+                // First, restore from widgets_values (for backward compatibility)
+                if (info.widgets_values && Array.isArray(info.widgets_values)) {
+                    for (let i = 0; i < this.widgets.length && i < info.widgets_values.length; i++) {
+                        const w = this.widgets[i];
+                        const value = info.widgets_values[i];
+
+                        if (w.type === "BOOLEAN" && w.isAuto !== undefined) {
+                            w.deserializeValue(value);
+                        }
+                    }
+                }
+
+                // Then, restore full state if available
+                if (info.auto_widget_states) {
+                    for (const w of this.widgets) {
+                        if (w.type === "BOOLEAN" && w.isAuto !== undefined && w.name in info.auto_widget_states) {
+                            const state = info.auto_widget_states[w.name];
+
+                            w.isAuto = state.isAuto;
+                            w.cachedValue = state.cachedValue;
+                            w.value = state.isAuto ? -99999 : state.value;
+
+                            w.callback?.(w.value);
+                        }
+                    }
+                }
+
+                // Force a redraw
+                this.graph?.setDirtyCanvas(true, true);
+            });
+
+
+            // Override addInput to handle AUTO widgets
+            chainCallback(nodeType.prototype, "onNodeCreated", function () {
+                // Convert any existing widgets to AUTO widgets if needed
+                let new_widgets = [];
+                const intWidgetNames = ["sp_size", "tp_size", "height", "width", "num_frames", "num_inference_steps", "flow_shift", "seed", "fps", "scale_factor",
+                    "tile_sample_min_height", "tile_sample_min_width", "tile_sample_min_num_frames", "tile_sample_stride_height", "tile_sample_stride_width",
+                    "tile_sample_stride_num_frames", "blend_num_frames"
+                ]
+                const floatWidgetNames = ["embedded_cfg_scale", "guidance_scale"]
+                const comboWidgetNames = ["vae_tiling", "vae_precision", "vae_sp", "text_encoder_precision", "precision",
+                    "load_encoder", "load_decoder", "use_tiling", "use_temporal_tiling", "use_parallel_tiling", "use_cpu_offload", "enable_teacache"
+                ]
+                const stringWidgetNames = ["prefix", "quant_config", "lora_config", "image_path"]
+
+                if (this.widgets) {
+                    for (let w of this.widgets) {
+                        if (intWidgetNames.includes(w.name)) {
+                            new_widgets.push(makeAutoAnnotated(w, ["FVAUTOINT", { "default": 0 }]));
+                        } else if (floatWidgetNames.includes(w.name)) {
+                            new_widgets.push(makeAutoAnnotated(w, ["FVAUTOFLOAT", { "default": 0 }]));
+                        } else if (comboWidgetNames.includes(w.name)) {
+                            new_widgets.push(makeAutoAnnotated(w, ["FVAUTOCOMBO", { "default": 0 }]));
+                        } else if (stringWidgetNames.includes(w.name)) {
+                            new_widgets.push(makeAutoAnnotated(w, ["FVAUTOSTRING", { "default": "" }]));
+                        } else {
+                            new_widgets.push(w);
+                        }
+                    }
+                    this.widgets = new_widgets;
+
+                    const autoWidgets = this.widgets.filter(w => w.type === "BOOLEAN" && w.isAuto !== undefined);
+                }
+
+                this.graph?.setDirtyCanvas(true, true);
+            });
+        }
+    },
+
+    async init() {
+        // Force a redraw of all nodes when the extension initializes
+        if (app.graph) {
+            setTimeout(() => {
+                app.graph.setDirtyCanvas(true, true);
+            }, 1000);
+        }
+    }
+});
+
+console.log("FastVideo.core.js loaded");
@@ -279,11 +279,11 @@ def main(args):

 if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Benchmark Block Sparse Attention')
-    parser.add_argument('--batch_size', type=int, default=4, help='Batch size')
-    parser.add_argument('--num_heads', type=int, default=6, help='Number of heads')
-    parser.add_argument('--head_dim', type=int, default=128, help='Head dimension')
-    parser.add_argument('--topk', type=int, default=64, help='Number of kv blocks each q block attends to')
+    parser.add_argument('--batch_size', type=int, default=2, help='Batch size')
+    parser.add_argument('--num_heads', type=int, default=12, help='Number of heads')
+    parser.add_argument('--head_dim', type=int, default=64, help='Head dimension')
+    parser.add_argument('--topk', type=int, default=32, help='Number of kv blocks each q block attends to')
    parser.add_argument('--seq_lengths', type=int, nargs='+', default=[29120], help='Sequence lengths to benchmark')
-    parser.add_argument('--num_iterations', type=int, default=50, help='Number of test iterations to run')
+    parser.add_argument('--num_iterations', type=int, default=100, help='Number of test iterations to run')
    args = parser.parse_args()
    main(args)
@@ -18,7 +18,6 @@ import os
 import re
 import sys
 from pathlib import Path
-from typing import Optional

 import requests

@@ -168,8 +167,7 @@ _cached_base: str = ""
 _cached_branch: str = ""


-def get_repo_base_and_branch(
-        pr_number: str) -> tuple[Optional[str], Optional[str]]:
+def get_repo_base_and_branch(pr_number: str) -> tuple[str | None, str | None]:
    global _cached_base, _cached_branch
    if _cached_base and _cached_branch:
        return _cached_base, _cached_branch
@@ -8,12 +8,14 @@ You can easily use the FastVideo Docker image as a custom container on [RunPod](

 Choose a GPU that supports CUDA 12.4

+Pick 1 or 2 L40S GPU(s)
+
 ![RunPod CUDA selection](../../_static/images/runpod_cuda.png)

 When creating your pod template, use this image:

 ```
-ghcr.io/hao-ai-lab/fastvideo/fastvideo-dev:latest
+ghcr.io/hao-ai-lab/fastvideo/fastvideo-dev:py3.12-latest
 ```

 Paste Container Start Command to support SSH ([RunPod Docs](https://docs.runpod.io/pods/configuration/use-ssh)):
@@ -5,7 +5,6 @@ import itertools
 import re
 from dataclasses import dataclass, field
 from pathlib import Path
-from typing import Optional

 ROOT_DIR = Path(__file__).parent.parent.parent.resolve()
 ROOT_DIR_RELATIVE = '../../../..'
@@ -89,7 +88,7 @@ class Example:
        generate() -> str: Generates the documentation content.
    """ # noqa: E501
    path: Path
-    category: Optional[str] = None
+    category: str | None = None
    main_file: Path = field(init=False)
    other_files: list[Path] = field(init=False)
    title: str = field(init=False)
@@ -117,4 +117,4 @@ If you're planning to contribute to FastVideo please see the following page:

 If you encounter any issues during installation, please open an issue on our [GitHub repository](https://github.com/hao-ai-lab/FastVideo).

-You can also join our [Slack community](https://join.slack.com/t/fastvideo/shared_invite/zt-2zf6ru791-sRwI9lPIUJQq1mIeB_yjJg) for additional support.
+You can also join our [Slack community](https://join.slack.com/t/fastvideo/shared_invite/zt-38u6p1jqe-yDI1QJOCEnbtkLoaI5bjZQ) for additional support.
@@ -12,7 +12,7 @@ This guide explains how to implement a custom diffusion pipeline in FastVideo, l
 4. **Register Your Pipeline** - Make it discoverable by the framework
 5. **Configure Your Pipeline** - (Coming soon)

-Need help? Join our [Slack community](https://join.slack.com/t/fastvideo/shared_invite/zt-2zf6ru791-sRwI9lPIUJQq1mIeB_yjJg).
+Need help? Join our [Slack community](https://join.slack.com/t/fastvideo/shared_invite/zt-38u6p1jqe-yDI1QJOCEnbtkLoaI5bjZQ).

 ## Step 1: Pipeline Modules

@@ -27,7 +27,7 @@ fastvideo generate --help
 ### Hardware Configuration

 - `--num-gpus {NUM_GPUS}`: Number of GPUs to use
- `--tp-size {TP_SIZE}`: Tensor parallelism size (Typically should match the number of GPUs)
+- `--tp-size {TP_SIZE}`: Tensor parallelism size (only for the encoder, should not be larger than 1 if text encoder offload is enabled, as layerwise offload + prefetch is faster)
 - `--sp-size {SP_SIZE}`: Sequence parallelism size (Typically should match the number of GPUs)

 #### Video Configuration
@@ -68,7 +68,7 @@ Example configuration file (config.json):
    "output_path": "outputs/",
    "num_gpus": 2,
    "sp_size": 2,
-    "tp_size": 2,
+    "tp_size": 1,
    "num_frames": 45,
    "height": 720,
    "width": 1280,
@@ -102,7 +102,7 @@ prompt: "A beautiful woman in a red dress walking down a street"
 output_path: "outputs/"
 num_gpus: 2
 sp_size: 2
-tp_size: 2
+tp_size: 1
 num_frames: 45
 height: 720
 width: 1280
@@ -121,4 +121,4 @@ If the generated video doesn't match your prompt:
 - Learn about using [Optimizations](#inference-optimizations)
 - See [Examples](../examples/examples_inference_index.md) for more usage scenarios
 - Join our [Community Discord](https://discord.gg/JA7cksDz86).
- Join our [Community Slack](https://join.slack.com/t/fastvideo/shared_invite/zt-2zf6ru791-sRwI9lPIUJQq1mIeB_yjJg).
+- Join our [Community Slack](https://join.slack.com/t/fastvideo/shared_invite/zt-38u6p1jqe-yDI1QJOCEnbtkLoaI5bjZQ).
@@ -10,7 +10,7 @@ def main():
    # attempt to identify the optimal arguments.
    generator = VideoGenerator.from_pretrained(
        "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
-        # if num_gpus > 1, FastVideo will automatically handle distributed setup
+        # FastVideo will automatically handle distributed setup
        num_gpus=2,
        use_fsdp_inference=True,
        use_cpu_offload=False
@@ -0,0 +1,10 @@
+This directory contain e2e examples scripts for finetuning Wan2.1 I2V. 
+
+Execute the following commands from `FastVideo/` to run training:
+
+- Download crush-smol dataset:
+`bash examples/training/finetune/wan_i2v_14b_480p/crush_smol/download_dataset.sh`
+- Preprocess the videos and captions into latents:
+`bash examples/training/finetune/wan_i2v_14b_480p/crush_smol/preprocess_wan_data_i2v.sh`
+- Edit the following file and run finetuning:
+`bash examples/training/finetune/wan_i2v_14b_480p/crush_smol/finetune_i2v.sh`
@@ -0,0 +1,3 @@
+#!/bin/bash
+
+python scripts/huggingface/download_hf.py --repo_id "wlsaidhi/crush-smol-merged" --local_dir "data/crush-smol" --repo_type "dataset"
@@ -0,0 +1,93 @@
+#!/bin/bash
+
+export WANDB_BASE_URL="https://api.wandb.ai"
+export WANDB_MODE=online
+export TOKENIZERS_PARALLELISM=false
+# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
+
+MODEL_PATH="weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers"
+DATA_DIR="data/crush-smol_processed_i2v_1_3b_inp/combined_parquet_dataset/"
+VALIDATION_DATASET_FILE="examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/validation.json"
+NUM_GPUS=8
+# export CUDA_VISIBLE_DEVICES=4,5
+# IP=[MASTER NODE IP]
+
+# Training arguments
+training_args=(
+  --tracker_project_name "wan_i2v_finetune"
+  --output_dir "$DATA_DIR/outputs/wan_i2v_finetune"
+  --max_train_steps 2000
+  --train_batch_size 4
+  --train_sp_batch_size 1
+  --gradient_accumulation_steps 1
+  --num_latent_t 8
+  --num_height 480
+  --num_width 832
+  --num_frames 77
+)
+
+# Parallel arguments
+parallel_args=(
+  --num_gpus $NUM_GPUS
+  --sp_size 4
+  --tp_size 4
+  --hsdp_replicate_dim 2
+  --hsdp_shard_dim 4
+)
+
+# Model arguments
+model_args=(
+  --model_path $MODEL_PATH
+  --pretrained_model_name_or_path $MODEL_PATH
+)
+
+# Dataset arguments
+dataset_args=(
+  --data_path "$DATA_DIR"
+  --dataloader_num_workers 1
+)
+
+# Validation arguments
+validation_args=(
+  --log_validation
+  --validation_dataset_file "$VALIDATION_DATASET_FILE"
+  --validation_steps 100
+  --validation_sampling_steps "40"
+  --validation_guidance_scale "1.0"
+)
+
+# Optimizer arguments
+optimizer_args=(
+  --learning_rate 2e-5
+  --mixed_precision "bf16"
+  --checkpointing_steps 2000
+  --weight_decay 1e-4
+  --max_grad_norm 1.0
+)
+
+# Miscellaneous arguments
+miscellaneous_args=(
+  --inference_mode False
+  --allow_tf32
+  --checkpoints_total_limit 3
+  --training_cfg_rate 0.1
+  --multi_phased_distill_schedule "4000-1"
+  --not_apply_cfg_solver
+  --dit_precision "fp32"
+  --num_euler_timesteps 50
+  --ema_start_step 0
+  --enable_gradient_checkpointing_type "full"
+)
+
+# If you do not have 32 GPUs and to fit in memory, you can: 1. increase sp_size. 2. reduce num_latent_t
+torchrun \
+  --nnodes 1 \
+  --nproc_per_node $NUM_GPUS \
+    fastvideo/v1/training/wan_i2v_training_pipeline.py \
+    "${parallel_args[@]}" \
+    "${model_args[@]}" \
+    "${dataset_args[@]}" \
+    "${training_args[@]}" \
+    "${optimizer_args[@]}" \
+    "${validation_args[@]}" \
+    "${miscellaneous_args[@]}"
@@ -0,0 +1,131 @@
+#!/bin/bash
+#SBATCH --job-name=i2v
+#SBATCH --partition=main
+#SBATCH --qos=hao
+#SBATCH --nodes=4
+#SBATCH --ntasks=4
+#SBATCH --ntasks-per-node=1
+#SBATCH --gres=gpu:8
+#SBATCH --cpus-per-task=128
+#SBATCH --nodelist=fs-mbz-gpu-[100-850]
+#SBATCH --mem=1440G
+#SBATCH --output=i2v_output/i2v_%j.out
+#SBATCH --error=i2v_output/i2v_%j.err
+#SBATCH --exclusive
+set -e -x
+
+# Environment Setup
+source ~/conda/miniconda/bin/activate
+conda activate will-fv
+
+# Basic Info
+export WANDB_MODE="online"
+export NCCL_P2P_DISABLE=1
+export TORCH_NCCL_ENABLE_MONITORING=0
+# different cache dir for different processes
+export TRITON_CACHE_DIR=/tmp/triton_cache_${SLURM_PROCID}
+export MASTER_PORT=29500
+export NODE_RANK=$SLURM_PROCID
+nodes=( $(scontrol show hostnames $SLURM_JOB_NODELIST) )
+export MASTER_ADDR=${nodes[0]}
+export CUDA_VISIBLE_DEVICES=$SLURM_LOCALID
+export TOKENIZERS_PARALLELISM=false
+export WANDB_BASE_URL="https://api.wandb.ai"
+export WANDB_MODE=online
+# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
+
+echo "MASTER_ADDR: $MASTER_ADDR"
+echo "NODE_RANK: $NODE_RANK"
+
+# Configs
+NUM_GPUS=8
+
+MODEL_PATH="weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers"
+DATA_DIR="data/crush-smol_processed_i2v_1_3b_inp/combined_parquet_dataset/"
+VALIDATION_DATASET_FILE="examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/validation.json"
+# export CUDA_VISIBLE_DEVICES=4,5
+# IP=[MASTER NODE IP]
+
+# If you do not have 32 GPUs and to fit in memory, you can: 1. increase sp_size. 2. reduce num_latent_t
+
+# Training arguments
+training_args=(
+  --tracker_project_name wan_i2v_finetune
+  --output_dir="$DATA_DIR/outputs/wan_i2v_finetune_2n"
+  --max_train_steps=2000
+  --train_batch_size=2
+  --train_sp_batch_size 1
+  --gradient_accumulation_steps=1
+  --num_latent_t 8
+  --num_height 480
+  --num_width 832
+  --num_frames 77
+)
+
+# Parallel arguments
+parallel_args=(
+  --num_gpus $NUM_GPUS
+  --sp_size $NUM_GPUS
+  --tp_size $NUM_GPUS
+  --hsdp_replicate_dim $SLURM_JOB_NUM_NODES
+  --hsdp_shard_dim $NUM_GPUS
+)
+
+# Model arguments
+model_args=(
+  --model_path $MODEL_PATH
+  --pretrained_model_name_or_path $MODEL_PATH
+)
+
+# Dataset arguments
+dataset_args=(
+  --data_path "$DATA_DIR"
+  --dataloader_num_workers 10
+)
+
+# Validation arguments
+validation_args=(
+  --log_validation
+  --validation_dataset_file "$VALIDATION_DATASET_FILE"
+  --validation_steps 100
+  --validation_sampling_steps "40"
+  --validation_guidance_scale "1.0"
+)
+
+# Optimizer arguments
+optimizer_args=(
+  --learning_rate=1e-5
+  --mixed_precision="bf16"
+  --checkpointing_steps=1000
+  --weight_decay 1e-4
+  --max_grad_norm 1.0
+)
+
+# Miscellaneous arguments
+miscellaneous_args=(
+  --inference_mode False
+  --allow_tf32
+  --checkpoints_total_limit 3
+  --training_cfg_rate 0.1
+  --multi_phased_distill_schedule "4000-1"
+  --not_apply_cfg_solver
+  --dit_precision "fp32"
+  --num_euler_timesteps 50
+  --ema_start_step 0
+  --enable_gradient_checkpointing_type "full"
+)
+
+srun torchrun \
+--nnodes $SLURM_JOB_NUM_NODES \
+--nproc_per_node $NUM_GPUS \
+--node_rank $SLURM_PROCID \
+--rdzv_backend=c10d \
+--rdzv_endpoint="$MASTER_ADDR:$MASTER_PORT" \
+    fastvideo/v1/training/wan_i2v_training_pipeline.py \
+    "${parallel_args[@]}" \
+    "${model_args[@]}" \
+    "${dataset_args[@]}" \
+    "${training_args[@]}" \
+    "${optimizer_args[@]}" \
+    "${validation_args[@]}" \
+    "${miscellaneous_args[@]}"
@@ -0,0 +1,26 @@
+#!/bin/bash
+
+GPU_NUM=1 # 2,4,8
+MODEL_PATH="weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers"
+MODEL_TYPE="wan"
+DATA_MERGE_PATH="data/crush-smol/merge.txt"
+OUTPUT_DIR="data/crush-smol_processed_i2v_1_3b_inp/"
+VALIDATION_PATH="examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/validation.json"
+
+torchrun --nproc_per_node=$GPU_NUM \
+    fastvideo/v1/pipelines/preprocess/v1_preprocess.py \
+    --model_path $MODEL_PATH \
+    --data_merge_path $DATA_MERGE_PATH \
+    --preprocess_video_batch_size 8 \
+    --max_height 480 \
+    --max_width 832 \
+    --num_frames 77 \
+    --dataloader_num_workers 0 \
+    --output_dir=$OUTPUT_DIR \
+    --model_type $MODEL_TYPE \
+    --train_fps 16 \
+    --validation_dataset_file $VALIDATION_PATH \
+    --samples_per_file 8 \
+    --flush_frequency 8 \
+    --video_length_tolerance_range 5 \
+    --preprocess_task "i2v" 
@@ -0,0 +1,31 @@
+{
+  "data": [
+    {
+      "caption": "A large metal cylinder is seen pressing down on a pile of Oreo cookies, flattening them as if they were under a hydraulic press.",
+      "image_path": null,
+      "video_path": "validation_dataset/yYcK4nANZz4-Scene-034.mp4",
+      "num_inference_steps": 40,
+      "height": 480,
+      "width": 832,
+      "num_frames": 77
+    },
+    {
+      "caption": "A large metal cylinder is seen compressing colorful clay into a compact shape, demonstrating the power of a hydraulic press.",
+      "image_path": null,
+      "video_path": "validation_dataset/yYcK4nANZz4-Scene-027.mp4",
+      "num_inference_steps": 40,
+      "height": 480,
+      "width": 832,
+      "num_frames": 77
+    },
+    {
+      "caption": "A large metal cylinder is seen pressing down on a pile of colorful candies, flattening them as if they were under a hydraulic press. The candies are crushed and broken into small pieces, creating a mess on the table.",
+      "image_path": null,
+      "video_path": "validation_dataset/yYcK4nANZz4-Scene-030.mp4",
+      "num_inference_steps": 40,
+      "height": 480,
+      "width": 832,
+      "num_frames": 77
+    }
+  ]
+}
@@ -2,11 +2,12 @@

 export WANDB_BASE_URL="https://api.wandb.ai"
 export WANDB_MODE=online
+export TOKENIZERS_PARALLELISM=false
 # export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA

 MODEL_PATH="Wan-AI/Wan2.1-I2V-14B-480P-Diffusers"
 DATA_DIR="data/crush-smol_processed_i2v/combined_parquet_dataset/"
-VALIDATION_DIR="data/crush-smol_processed_i2v/validation_parquet_dataset/"
+VALIDATION_DATASET_FILE="examples/training/finetune/wan_i2v_14b_480p/crush_smol/validation.json"
 NUM_GPUS=8
 # export CUDA_VISIBLE_DEVICES=4,5
 # IP=[MASTER NODE IP]
@@ -29,7 +30,7 @@ training_args=(
 parallel_args=(
  --num_gpus $NUM_GPUS
  --sp_size 8
-  --tp_size 8
+  --tp_size 1
  --hsdp_replicate_dim 1
  --hsdp_shard_dim 8
 )
@@ -49,7 +50,7 @@ dataset_args=(
 # Validation arguments
 validation_args=(
  --log_validation
-  --validation_preprocessed_path "$VALIDATION_DIR"
+  --validation_dataset_file "$VALIDATION_DATASET_FILE"
  --validation_steps 100
  --validation_sampling_steps "40"
  --validation_guidance_scale "1.0"
@@ -75,6 +76,7 @@ miscellaneous_args=(
  --dit_precision "fp32"
  --num_euler_timesteps 50
  --ema_start_step 0
+  --enable_gradient_checkpointing_type "full"
 )

 # If you do not have 32 GPUs and to fit in memory, you can: 1. increase sp_size. 2. reduce num_latent_t
@@ -42,7 +42,7 @@ NUM_GPUS=8

 MODEL_PATH="Wan-AI/Wan2.1-I2V-14B-480P-Diffusers"
 DATA_DIR="data/crush-smol_processed_i2v/combined_parquet_dataset/"
-VALIDATION_DIR="data/crush-smol_processed_i2v/validation_parquet_dataset/"
+VALIDATION_DATASET_FILE="examples/training/finetune/wan_i2v_14b_480p/crush_smol/validation.json"
 # export CUDA_VISIBLE_DEVICES=4,5
 # IP=[MASTER NODE IP]

@@ -66,7 +66,7 @@ training_args=(
 parallel_args=(
  --num_gpus $NUM_GPUS
  --sp_size $NUM_GPUS
-  --tp_size $NUM_GPUS
+  --tp_size 1
  --hsdp_replicate_dim $SLURM_JOB_NUM_NODES
  --hsdp_shard_dim $NUM_GPUS
 )
@@ -86,7 +86,7 @@ dataset_args=(
 # Validation arguments
 validation_args=(
  --log_validation
-  --validation_preprocessed_path "$VALIDATION_DIR"
+  --validation_dataset_file "$VALIDATION_DATASET_FILE"
  --validation_steps 100
  --validation_sampling_steps "40"
  --validation_guidance_scale "1.0"
@@ -112,6 +112,7 @@ miscellaneous_args=(
  --dit_precision "fp32"
  --num_euler_timesteps 50
  --ema_start_step 0
+  --enable_gradient_checkpointing_type "full"
 )

 srun torchrun \
@@ -26,6 +26,24 @@
      "height": 480,
      "width": 832,
      "num_frames": 77
+    },
+    {
+      "caption": "The video shows a close-up of an orange being flattened as if it were under a hydraulic press, with the press moving down and compressing the fruit until it is completely flattened.",
+      "image_path": null,
+      "video_path": "validation_dataset/EJqsC21GSBY-Scene-059.mp4",
+      "num_inference_steps": 40,
+      "height": 480,
+      "width": 832,
+      "num_frames": 98
+    },
+    {
+      "caption": "The video shows a stack of colorful sponges being flattened as if they were under a hydraulic press. The sponges, which are pink, white, blue, and green, are compressed into a smaller size, demonstrating the press's power. The background features a green wall with a yellow and red sign, adding context to the setting.",
+      "image_path": null,
+      "video_path": "validation_dataset/EJqsC21GSBY-Scene-013.mp4",
+      "num_inference_steps": 40,
+      "height": 480,
+      "width": 832,
+      "num_frames": 148
    }
  ]
 }
@@ -2,11 +2,12 @@

 export WANDB_BASE_URL="https://api.wandb.ai"
 export WANDB_MODE=online
+export TOKENIZERS_PARALLELISM=false
 # export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA

 MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
 DATA_DIR="data/crush-smol_processed_t2v/combined_parquet_dataset/"
-VALIDATION_DIR="data/crush-smol_processed_t2v/validation_parquet_dataset/"
+VALIDATION_DATASET_FILE="examples/training/finetune/wan_t2v_1_3b/crush_smol/validation.json"
 NUM_GPUS=4
 # export CUDA_VISIBLE_DEVICES=4,5

@@ -29,7 +30,7 @@ training_args=(
 parallel_args=(
  --num_gpus $NUM_GPUS 
  --sp_size $NUM_GPUS 
-  --tp_size $NUM_GPUS
+  --tp_size 1
  --hsdp_replicate_dim 1
  --hsdp_shard_dim $NUM_GPUS
 )
@@ -49,7 +50,7 @@ dataset_args=(
 # Validation arguments
 validation_args=(
  --log_validation 
-  --validation_preprocessed_path $VALIDATION_DIR
+  --validation_dataset_file $VALIDATION_DATASET_FILE
  --validation_steps 50
  --validation_sampling_steps "50" 
  --validation_guidance_scale "1.0"
@@ -75,6 +76,7 @@ miscellaneous_args=(
  --dit_precision "fp32"
  --num_euler_timesteps 50
  --ema_start_step 0
+  --enable_gradient_checkpointing_type "full"
 )

 torchrun \
@@ -41,7 +41,7 @@ echo "NODE_RANK: $NODE_RANK"
 NUM_GPUS=8
 MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
 DATA_DIR="data/crush-smol_processed_t2v/combined_parquet_dataset/"
-VALIDATION_DIR="data/crush-smol_processed_t2v/validation_parquet_dataset/"
+VALIDATION_DATASET_FILE="examples/training/finetune/wan_t2v_1_3b/crush_smol/validation.json"
 # export CUDA_VISIBLE_DEVICES=4,5
 # IP=[MASTER NODE IP]

@@ -63,7 +63,7 @@ training_args=(
 parallel_args=(
  --num_gpus $NUM_GPUS
  --sp_size 4
-  --tp_size 4
+  --tp_size 1
  --hsdp_replicate_dim 2
  --hsdp_shard_dim 4
 )
@@ -83,7 +83,7 @@ dataset_args=(
 # Validation arguments
 validation_args=(
  --log_validation
-  --validation_preprocessed_path "$VALIDATION_DIR"
+  --validation_dataset_file "$VALIDATION_DATASET_FILE"
  --validation_steps 100
  --validation_sampling_steps "50"
  --validation_guidance_scale "1.0"
@@ -109,6 +109,7 @@ miscellaneous_args=(
  --dit_precision "fp32"
  --num_euler_timesteps 50
  --ema_start_step 0
+  --enable_gradient_checkpointing_type "full"
 )

 srun torchrun \
@@ -2,7 +2,7 @@
 import json
 import os
 from collections import defaultdict
-from typing import Any, Dict, List, Optional, Tuple
+from typing import Any

 import numpy as np

@@ -13,7 +13,7 @@ def configure_sta(mode: str = 'STA_searching',
                  layer_num: int = 40,
                  time_step_num: int = 50,
                  head_num: int = 40,
-                  **kwargs) -> List[List[List[Any]]]:
+                  **kwargs) -> list[list[list[Any]]]:
    """
    Configure Sliding Tile Attention (STA) parameters based on the specified mode.
    
@@ -53,22 +53,22 @@ def configure_sta(mode: str = 'STA_searching',

    if mode == 'STA_searching':
        # Get parameters with defaults
-        mask_candidates: Optional[List[str]] = kwargs.get('mask_candidates')
+        mask_candidates: list[str] | None = kwargs.get('mask_candidates')
        if mask_candidates is None:
            raise ValueError(
                "mask_candidates is required for STA_searching mode")
-        mask_selected: List[int] = kwargs.get('mask_selected',
+        mask_selected: list[int] = kwargs.get('mask_selected',
                                              list(range(len(mask_candidates))))

        # Parse selected masks
-        selected_masks: List[List[int]] = []
+        selected_masks: list[list[int]] = []
        for index in mask_selected:
            mask = mask_candidates[index]
            masks_list = [int(x) for x in mask.split(',')]
            selected_masks.append(masks_list)

        # Create 3D mask structure with fixed dimensions (t=50, l=60)
-        masks_3d: List[List[List[List[int]]]] = []
+        masks_3d: list[list[list[list[int]]]] = []
        for i in range(time_step_num):  # Fixed t dimension = 50
            row = []
            for j in range(layer_num):  # Fixed l dimension = 60
@@ -79,25 +79,23 @@ def configure_sta(mode: str = 'STA_searching',

    elif mode == 'STA_tuning':
        # Get required parameters
-        mask_search_files_path: Optional[str] = kwargs.get(
+        mask_search_files_path: str | None = kwargs.get(
            'mask_search_files_path')
        if not mask_search_files_path:
            raise ValueError(
                "mask_search_files_path is required for STA_tuning mode")

        # Get optional parameters with defaults
-        mask_candidates_tuning: Optional[List[str]] = kwargs.get(
-            'mask_candidates')
+        mask_candidates_tuning: list[str] | None = kwargs.get('mask_candidates')
        if mask_candidates_tuning is None:
            raise ValueError("mask_candidates is required for STA_tuning mode")
-        mask_selected_tuning: List[int] = kwargs.get(
+        mask_selected_tuning: list[int] = kwargs.get(
            'mask_selected', list(range(len(mask_candidates_tuning))))
-        skip_time_steps_tuning: Optional[int] = kwargs.get('skip_time_steps')
-        save_dir_tuning: Optional[str] = kwargs.get('save_dir',
-                                                    "mask_candidates")
+        skip_time_steps_tuning: int | None = kwargs.get('skip_time_steps')
+        save_dir_tuning: str | None = kwargs.get('save_dir', "mask_candidates")

        # Parse selected masks
-        selected_masks_tuning: List[List[int]] = []
+        selected_masks_tuning: list[list[int]] = []
        for index in mask_selected_tuning:
            mask = mask_candidates_tuning[index]
            masks_list = [int(x) for x in mask.split(',')]
@@ -108,7 +106,7 @@ def configure_sta(mode: str = 'STA_searching',
        averaged_results = average_head_losses(results, selected_masks_tuning)

        # Add full attention mask for specific cases
-        full_attention_mask_tuning: Optional[List[int]] = kwargs.get(
+        full_attention_mask_tuning: list[int] | None = kwargs.get(
            'full_attention_mask')
        if full_attention_mask_tuning is not None:
            selected_masks_tuning.append(full_attention_mask_tuning)
@@ -149,11 +147,11 @@ def configure_sta(mode: str = 'STA_searching',
        return mask_strategy_3d
    elif mode == 'STA_tuning_cfg':
        # Get required parameters for both positive and negative paths
-        mask_search_files_path_pos: Optional[str] = kwargs.get(
+        mask_search_files_path_pos: str | None = kwargs.get(
            'mask_search_files_path_pos')
-        mask_search_files_path_neg: Optional[str] = kwargs.get(
+        mask_search_files_path_neg: str | None = kwargs.get(
            'mask_search_files_path_neg')
-        save_dir_cfg: Optional[str] = kwargs.get('save_dir')
+        save_dir_cfg: str | None = kwargs.get('save_dir')

        if not mask_search_files_path_pos or not mask_search_files_path_neg or not save_dir_cfg:
            raise ValueError(
@@ -161,16 +159,16 @@ def configure_sta(mode: str = 'STA_searching',
            )

        # Get optional parameters with defaults
-        mask_candidates_cfg: Optional[List[str]] = kwargs.get('mask_candidates')
+        mask_candidates_cfg: list[str] | None = kwargs.get('mask_candidates')
        if mask_candidates_cfg is None:
            raise ValueError(
                "mask_candidates is required for STA_tuning_cfg mode")
-        mask_selected_cfg: List[int] = kwargs.get(
+        mask_selected_cfg: list[int] = kwargs.get(
            'mask_selected', list(range(len(mask_candidates_cfg))))
-        skip_time_steps_cfg: Optional[int] = kwargs.get('skip_time_steps')
+        skip_time_steps_cfg: int | None = kwargs.get('skip_time_steps')

        # Parse selected masks
-        selected_masks_cfg: List[List[int]] = []
+        selected_masks_cfg: list[list[int]] = []
        for index in mask_selected_cfg:
            mask = mask_candidates_cfg[index]
            masks_list = [int(x) for x in mask.split(',')]
@@ -187,7 +185,7 @@ def configure_sta(mode: str = 'STA_searching',
                                               selected_masks_cfg)

        # Add full attention mask for specific cases
-        full_attention_mask_cfg: Optional[List[int]] = kwargs.get(
+        full_attention_mask_cfg: list[int] | None = kwargs.get(
            'full_attention_mask')
        if full_attention_mask_cfg is not None:
            selected_masks_cfg.append(full_attention_mask_cfg)
@@ -227,7 +225,7 @@ def configure_sta(mode: str = 'STA_searching',

    else:  # STA_inference
        # Get parameters with defaults
-        load_path: Optional[str] = kwargs.get(
+        load_path: str | None = kwargs.get(
            'load_path', "mask_candidates/mask_strategy.json")
        if load_path is None:
            raise ValueError("load_path is required for STA_inference mode")
@@ -248,9 +246,9 @@ def configure_sta(mode: str = 'STA_searching',
 # Helper functions


-def read_specific_json_files(folder_path: str) -> List[Dict[str, Any]]:
+def read_specific_json_files(folder_path: str) -> list[dict[str, Any]]:
    """Read and parse JSON files containing mask search results."""
-    json_contents: List[Dict[str, Any]] = []
+    json_contents: list[dict[str, Any]] = []

    # List files only in the current directory (no walk)
    files = os.listdir(folder_path)
@@ -268,11 +266,11 @@ def read_specific_json_files(folder_path: str) -> List[Dict[str, Any]]:


 def average_head_losses(
-        results: List[Dict[str, Any]],
-        selected_masks: List[List[int]]) -> Dict[str, Dict[str, np.ndarray]]:
+        results: list[dict[str, Any]],
+        selected_masks: list[list[int]]) -> dict[str, dict[str, np.ndarray]]:
    """Average losses across all prompts for each mask strategy."""
    # Initialize a dictionary to store the averaged results
-    averaged_losses: Dict[str, Dict[str, np.ndarray]] = {}
+    averaged_losses: dict[str, dict[str, np.ndarray]] = {}
    loss_type = 'L2_loss'
    # Get all loss types (e.g., 'L2_loss')
    averaged_losses[loss_type] = {}
@@ -294,14 +292,14 @@ def average_head_losses(


 def select_best_mask_strategy(
-        averaged_results: Dict[str, Dict[str, np.ndarray]],
-        selected_masks: List[List[int]],
+        averaged_results: dict[str, dict[str, np.ndarray]],
+        selected_masks: list[list[int]],
        skip_time_steps: int = 12,
        timesteps: int = 50,
        head_num: int = 40
-) -> Tuple[Dict[str, List[int]], float, Dict[str, int]]:
+) -> tuple[dict[str, list[int]], float, dict[str, int]]:
    """Select the best mask strategy for each head based on loss minimization."""
-    best_mask_strategy: Dict[str, List[int]] = {}
+    best_mask_strategy: dict[str, list[int]] = {}
    loss_type = 'L2_loss'
    # Get the shape of time steps and layers
    layers = len(averaged_results[loss_type][str(selected_masks[0])][0])
@@ -310,7 +308,7 @@ def select_best_mask_strategy(
    total_tokens = 0  # total number of masked tokens
    total_length = 0  # total sequence length

-    strategy_counts: Dict[str, int] = {
+    strategy_counts: dict[str, int] = {
        str(strategy): 0
        for strategy in selected_masks
    }
@@ -352,22 +350,22 @@ def select_best_mask_strategy(


 def save_mask_search_results(
-        mask_search_final_result: List[Dict[str, List[float]]],
+        mask_search_final_result: list[dict[str, list[float]]],
        prompt: str,
-        mask_strategies: List[str],
-        output_dir: str = 'output/mask_search_result/') -> Optional[str]:
+        mask_strategies: list[str],
+        output_dir: str = 'output/mask_search_result/') -> str | None:
    if not mask_search_final_result:
        print("No mask search results to save")
        return None

    # Create result dictionary with defaultdict for nested lists
-    mask_search_dict: Dict[str, Dict[str, List[List[float]]]] = {
+    mask_search_dict: dict[str, dict[str, list[list[float]]]] = {
        "L2_loss": defaultdict(list),
        "L1_loss": defaultdict(list)
    }

    mask_selected = list(range(len(mask_strategies)))
-    selected_masks: List[List[int]] = []
+    selected_masks: list[list[int]] = []
    for index in mask_selected:
        mask = mask_strategies[index]
        masks_list = [int(x) for x in mask.split(',')]
@@ -377,7 +375,7 @@ def save_mask_search_results(
    for i, mask_strategy in enumerate(selected_masks):
        mask_strategy_str = str(mask_strategy)
        # Process L2 loss
-        step_results: List[List[float]] = []
+        step_results: list[list[float]] = []
        for step_data in mask_search_final_result:
            if isinstance(step_data, dict) and "L2_loss" in step_data:
                layer_losses = [float(loss) for loss in step_data["L2_loss"]]
@@ -3,8 +3,7 @@

 from abc import ABC, abstractmethod
 from dataclasses import dataclass, fields
-from typing import (TYPE_CHECKING, Any, Dict, Generic, Optional, Protocol, Set,
-                    Type, TypeVar)
+from typing import TYPE_CHECKING, Any, Generic, Protocol, TypeVar

 if TYPE_CHECKING:
    from fastvideo.v1.fastvideo_args import FastVideoArgs
@@ -27,12 +26,12 @@ class AttentionBackend(ABC):

    @staticmethod
    @abstractmethod
-    def get_impl_cls() -> Type["AttentionImpl"]:
+    def get_impl_cls() -> type["AttentionImpl"]:
        raise NotImplementedError

    @staticmethod
    @abstractmethod
-    def get_metadata_cls() -> Type["AttentionMetadata"]:
+    def get_metadata_cls() -> type["AttentionMetadata"]:
        raise NotImplementedError

    # @staticmethod
@@ -46,7 +45,7 @@ class AttentionBackend(ABC):

    @staticmethod
    @abstractmethod
-    def get_builder_cls() -> Type["AttentionMetadataBuilder"]:
+    def get_builder_cls() -> type["AttentionMetadataBuilder"]:
        raise NotImplementedError


@@ -57,8 +56,7 @@ class AttentionMetadata:
    current_timestep: int

    def asdict_zerocopy(self,
-                        skip_fields: Optional[Set[str]] = None
-                        ) -> Dict[str, Any]:
+                        skip_fields: set[str] | None = None) -> dict[str, Any]:
        """Similar to dataclasses.asdict, but avoids deepcopying."""
        if skip_fields is None:
            skip_fields = set()
@@ -124,7 +122,7 @@ class AttentionImpl(ABC, Generic[T]):
        head_size: int,
        softmax_scale: float,
        causal: bool = False,
-        num_kv_heads: Optional[int] = None,
+        num_kv_heads: int | None = None,
        prefix: str = "",
        **extra_impl_args,
    ) -> None:
@@ -1,7 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0

-from typing import List, Optional, Type
-
 import torch
 from flash_attn import flash_attn_func as flash_attn_2_func

@@ -28,7 +26,7 @@ class FlashAttentionBackend(AttentionBackend):
    accept_output_buffer: bool = True

    @staticmethod
-    def get_supported_head_sizes() -> List[int]:
+    def get_supported_head_sizes() -> list[int]:
        return [32, 64, 96, 128, 160, 192, 224, 256]

    @staticmethod
@@ -36,15 +34,15 @@ class FlashAttentionBackend(AttentionBackend):
        return "FLASH_ATTN"

    @staticmethod
-    def get_impl_cls() -> Type["FlashAttentionImpl"]:
+    def get_impl_cls() -> type["FlashAttentionImpl"]:
        return FlashAttentionImpl

    @staticmethod
-    def get_metadata_cls() -> Type["AttentionMetadata"]:
+    def get_metadata_cls() -> type["AttentionMetadata"]:
        raise NotImplementedError

    @staticmethod
-    def get_builder_cls() -> Type["AttentionMetadataBuilder"]:
+    def get_builder_cls() -> type["AttentionMetadataBuilder"]:
        raise NotImplementedError


@@ -56,7 +54,7 @@ class FlashAttentionImpl(AttentionImpl):
        head_size: int,
        causal: bool,
        softmax_scale: float,
-        num_kv_heads: Optional[int] = None,
+        num_kv_heads: int | None = None,
        prefix: str = "",
        **extra_impl_args,
    ) -> None:
@@ -1,5 +1,4 @@
 # SPDX-License-Identifier: Apache-2.0
-from typing import List, Optional, Type

 import torch
 from sageattention import sageattn
@@ -18,7 +17,7 @@ class SageAttentionBackend(AttentionBackend):
    accept_output_buffer: bool = True

    @staticmethod
-    def get_supported_head_sizes() -> List[int]:
+    def get_supported_head_sizes() -> list[int]:
        return [32, 64, 96, 128, 160, 192, 224, 256]

    @staticmethod
@@ -26,7 +25,7 @@ class SageAttentionBackend(AttentionBackend):
        return "SAGE_ATTN"

    @staticmethod
-    def get_impl_cls() -> Type["SageAttentionImpl"]:
+    def get_impl_cls() -> type["SageAttentionImpl"]:
        return SageAttentionImpl

    # @staticmethod
@@ -42,7 +41,7 @@ class SageAttentionImpl(AttentionImpl):
        head_size: int,
        causal: bool,
        softmax_scale: float,
-        num_kv_heads: Optional[int] = None,
+        num_kv_heads: int | None = None,
        prefix: str = "",
        **extra_impl_args,
    ) -> None:
@@ -1,5 +1,4 @@
 # SPDX-License-Identifier: Apache-2.0
-from typing import List, Optional, Type

 import torch

@@ -17,7 +16,7 @@ class SDPABackend(AttentionBackend):
    accept_output_buffer: bool = True

    @staticmethod
-    def get_supported_head_sizes() -> List[int]:
+    def get_supported_head_sizes() -> list[int]:
        return [32, 64, 96, 128, 160, 192, 224, 256]

    @staticmethod
@@ -25,7 +24,7 @@ class SDPABackend(AttentionBackend):
        return "SDPA"

    @staticmethod
-    def get_impl_cls() -> Type["SDPAImpl"]:
+    def get_impl_cls() -> type["SDPAImpl"]:
        return SDPAImpl

    # @staticmethod
@@ -41,7 +40,7 @@ class SDPAImpl(AttentionImpl):
        head_size: int,
        causal: bool,
        softmax_scale: float,
-        num_kv_heads: Optional[int] = None,
+        num_kv_heads: int | None = None,
        prefix: str = "",
        **extra_impl_args,
    ) -> None:
@@ -1,7 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 import json
 from dataclasses import dataclass
-from typing import Any, List, Optional, Type
+from typing import Any

 import torch
 from einops import rearrange
@@ -40,7 +40,7 @@ class SlidingTileAttentionBackend(AttentionBackend):
    accept_output_buffer: bool = True

    @staticmethod
-    def get_supported_head_sizes() -> List[int]:
+    def get_supported_head_sizes() -> list[int]:
        # TODO(will-refactor): check this
        return [32, 64, 96, 128, 160, 192, 224, 256]

@@ -49,22 +49,22 @@ class SlidingTileAttentionBackend(AttentionBackend):
        return "SLIDING_TILE_ATTN"

    @staticmethod
-    def get_impl_cls() -> Type["SlidingTileAttentionImpl"]:
+    def get_impl_cls() -> type["SlidingTileAttentionImpl"]:
        return SlidingTileAttentionImpl

    @staticmethod
-    def get_metadata_cls() -> Type["SlidingTileAttentionMetadata"]:
+    def get_metadata_cls() -> type["SlidingTileAttentionMetadata"]:
        return SlidingTileAttentionMetadata

    @staticmethod
-    def get_builder_cls() -> Type["SlidingTileAttentionMetadataBuilder"]:
+    def get_builder_cls() -> type["SlidingTileAttentionMetadataBuilder"]:
        return SlidingTileAttentionMetadataBuilder


@dataclass
 class SlidingTileAttentionMetadata(AttentionMetadata):
    current_timestep: int
-    STA_param: List[List[
+    STA_param: list[list[
        Any]]  # each timestep with one metadata, shape [num_layers, num_heads]


@@ -98,7 +98,7 @@ class SlidingTileAttentionImpl(AttentionImpl):
        head_size: int,
        causal: bool,
        softmax_scale: float,
-        num_kv_heads: Optional[int] = None,
+        num_kv_heads: int | None = None,
        prefix: str = "",
        **extra_impl_args,
    ) -> None:
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 import math
 from dataclasses import dataclass
-from typing import List, Optional, Type

 import torch
 from einops import rearrange
@@ -28,7 +27,7 @@ class VideoSparseAttentionBackend(AttentionBackend):
    accept_output_buffer: bool = True

    @staticmethod
-    def get_supported_head_sizes() -> List[int]:
+    def get_supported_head_sizes() -> list[int]:
        return [64, 128]

    @staticmethod
@@ -36,22 +35,22 @@ class VideoSparseAttentionBackend(AttentionBackend):
        return "VIDEO_SPARSE_ATTN"

    @staticmethod
-    def get_impl_cls() -> Type["VideoSparseAttentionImpl"]:
+    def get_impl_cls() -> type["VideoSparseAttentionImpl"]:
        return VideoSparseAttentionImpl

    @staticmethod
-    def get_metadata_cls() -> Type["VideoSparseAttentionMetadata"]:
+    def get_metadata_cls() -> type["VideoSparseAttentionMetadata"]:
        return VideoSparseAttentionMetadata

    @staticmethod
-    def get_builder_cls() -> Type["VideoSparseAttentionMetadataBuilder"]:
+    def get_builder_cls() -> type["VideoSparseAttentionMetadataBuilder"]:
        return VideoSparseAttentionMetadataBuilder


@dataclass
 class VideoSparseAttentionMetadata(AttentionMetadata):
    current_timestep: int
-    dit_seq_shape: List[int]
+    dit_seq_shape: list[int]
    VSA_sparsity: float


@@ -97,7 +96,7 @@ class VideoSparseAttentionImpl(AttentionImpl):
        head_size: int,
        causal: bool,
        softmax_scale: float,
-        num_kv_heads: Optional[int] = None,
+        num_kv_heads: int | None = None,
        prefix: str = "",
        **extra_impl_args,
    ) -> None:
@@ -105,8 +104,8 @@ class VideoSparseAttentionImpl(AttentionImpl):
        sp_group = get_sp_group()
        self.sp_size = sp_group.world_size
        self.VSA_base_tile_size = [4, 4, 4]
-        self.dit_seq_shape: List[int]
-        self.full_window_size: List[int]
+        self.dit_seq_shape: list[int]
+        self.full_window_size: list[int]
        self.img_seq_length: int

    def tile(self, x: torch.Tensor) -> torch.Tensor:
@@ -1,7 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0

-from typing import Optional, Tuple
-
 import torch
 import torch.nn as nn

@@ -23,11 +21,11 @@ class DistributedAttention(nn.Module):
    def __init__(self,
                 num_heads: int,
                 head_size: int,
-                 num_kv_heads: Optional[int] = None,
-                 softmax_scale: Optional[float] = None,
+                 num_kv_heads: int | None = None,
+                 softmax_scale: float | None = None,
                 causal: bool = False,
-                 supported_attention_backends: Optional[Tuple[
-                     AttentionBackendEnum, ...]] = None,
+                 supported_attention_backends: tuple[AttentionBackendEnum, ...]
+                 | None = None,
                 prefix: str = "",
                 **extra_impl_args) -> None:
        super().__init__()
@@ -63,10 +61,10 @@ class DistributedAttention(nn.Module):
        q: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
-        replicated_q: Optional[torch.Tensor] = None,
-        replicated_k: Optional[torch.Tensor] = None,
-        replicated_v: Optional[torch.Tensor] = None,
-    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+        replicated_q: torch.Tensor | None = None,
+        replicated_k: torch.Tensor | None = None,
+        replicated_v: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        """Forward pass for distributed attention.
        
        Args:
@@ -144,11 +142,11 @@ class DistributedAttention_VSA(DistributedAttention):
        q: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
-        replicated_q: Optional[torch.Tensor] = None,
-        replicated_k: Optional[torch.Tensor] = None,
-        replicated_v: Optional[torch.Tensor] = None,
-        gate_compress: Optional[torch.Tensor] = None,
-    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+        replicated_q: torch.Tensor | None = None,
+        replicated_k: torch.Tensor | None = None,
+        replicated_v: torch.Tensor | None = None,
+        gate_compress: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        """Forward pass for distributed attention.
        
        Args:
@@ -208,11 +206,11 @@ class LocalAttention(nn.Module):
    def __init__(self,
                 num_heads: int,
                 head_size: int,
-                 num_kv_heads: Optional[int] = None,
-                 softmax_scale: Optional[float] = None,
+                 num_kv_heads: int | None = None,
+                 softmax_scale: float | None = None,
                 causal: bool = False,
-                 supported_attention_backends: Optional[Tuple[
-                     AttentionBackendEnum, ...]] = None,
+                 supported_attention_backends: tuple[AttentionBackendEnum, ...]
+                 | None = None,
                 **extra_impl_args) -> None:
        super().__init__()
        if softmax_scale is None:
@@ -2,9 +2,10 @@
 # Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/attention/selector.py

 import os
+from collections.abc import Generator
 from contextlib import contextmanager
 from functools import cache
-from typing import Generator, Optional, Tuple, Type, cast
+from typing import cast

 import torch

@@ -17,7 +18,7 @@ from fastvideo.v1.utils import STR_BACKEND_ENV_VAR, resolve_obj_by_qualname
 logger = init_logger(__name__)


-def backend_name_to_enum(backend_name: str) -> Optional[AttentionBackendEnum]:
+def backend_name_to_enum(backend_name: str) -> AttentionBackendEnum | None:
    """
    Convert a string backend name to a _Backend enum value.

@@ -31,7 +32,7 @@ def backend_name_to_enum(backend_name: str) -> Optional[AttentionBackendEnum]:
          None


-def get_env_variable_attn_backend() -> Optional[AttentionBackendEnum]:
+def get_env_variable_attn_backend() -> AttentionBackendEnum | None:
    '''
    Get the backend override specified by the FastVideo attention
    backend environment variable, if one is specified.
@@ -53,11 +54,11 @@ def get_env_variable_attn_backend() -> Optional[AttentionBackendEnum]:
 #
 # THIS SELECTION TAKES PRECEDENCE OVER THE
 # FASTVIDEO ATTENTION BACKEND ENVIRONMENT VARIABLE
-forced_attn_backend: Optional[AttentionBackendEnum] = None
+forced_attn_backend: AttentionBackendEnum | None = None


 def global_force_attn_backend(
-        attn_backend: Optional[AttentionBackendEnum]) -> None:
+        attn_backend: AttentionBackendEnum | None) -> None:
    '''
    Force all attention operations to use a specified backend.

@@ -72,7 +73,7 @@ def global_force_attn_backend(
    forced_attn_backend = attn_backend


-def get_global_forced_attn_backend() -> Optional[AttentionBackendEnum]:
+def get_global_forced_attn_backend() -> AttentionBackendEnum | None:
    '''
    Get the currently-forced choice of attention backend,
    or None if auto-selection is currently enabled.
@@ -83,9 +84,9 @@ def get_global_forced_attn_backend() -> Optional[AttentionBackendEnum]:
 def get_attn_backend(
    head_size: int,
    dtype: torch.dtype,
-    supported_attention_backends: Optional[Tuple[AttentionBackendEnum,
-                                                 ...]] = None,
-) -> Type[AttentionBackend]:
+    supported_attention_backends: tuple[AttentionBackendEnum, ...]
+    | None = None,
+) -> type[AttentionBackend]:
    return _cached_get_attn_backend(head_size, dtype,
                                    supported_attention_backends)

@@ -94,9 +95,9 @@ def get_attn_backend(
 def _cached_get_attn_backend(
    head_size: int,
    dtype: torch.dtype,
-    supported_attention_backends: Optional[Tuple[AttentionBackendEnum,
-                                                 ...]] = None,
-) -> Type[AttentionBackend]:
+    supported_attention_backends: tuple[AttentionBackendEnum, ...]
+    | None = None,
+) -> type[AttentionBackend]:
    # Check whether a particular choice of backend was
    # previously forced.
    #
@@ -105,13 +106,13 @@ def _cached_get_attn_backend(
    if not supported_attention_backends:
        raise ValueError("supported_attention_backends is empty")
    selected_backend = None
-    backend_by_global_setting: Optional[AttentionBackendEnum] = (
+    backend_by_global_setting: AttentionBackendEnum | None = (
        get_global_forced_attn_backend())
    if backend_by_global_setting is not None:
        selected_backend = backend_by_global_setting
    else:
        # Check the environment variable and override if specified
-        backend_by_env_var: Optional[str] = envs.FASTVIDEO_ATTENTION_BACKEND
+        backend_by_env_var: str | None = envs.FASTVIDEO_ATTENTION_BACKEND
        if backend_by_env_var is not None:
            selected_backend = backend_name_to_enum(backend_by_env_var)

@@ -123,7 +124,7 @@ def _cached_get_attn_backend(
    if not attention_cls:
        raise ValueError(
            f"Invalid attention backend for {current_platform.device_name}")
-    return cast(Type[AttentionBackend], resolve_obj_by_qualname(attention_cls))
+    return cast(type[AttentionBackend], resolve_obj_by_qualname(attention_cls))


@contextmanager
@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass, field, fields
-from typing import Any, Dict
+from typing import Any

 from fastvideo.v1.logger import init_logger

@@ -12,7 +12,9 @@ logger = init_logger(__name__)
 # 3. Any field in ArchConfig is fixed upon initialization, and should be hidden away from users
@dataclass
 class ArchConfig:
-    pass
+    stacked_params_mapping: list[tuple[str, str, str]] = field(
+        default_factory=list
+    )  # mapping from huggingface weight names to custom names


@dataclass
@@ -42,7 +44,7 @@ class ModelConfig:
        self.__dict__.update(state)

    # This should be used only when loading from transformers/diffusers
-    def update_model_arch(self, source_model_dict: Dict[str, Any]) -> None:
+    def update_model_arch(self, source_model_dict: dict[str, Any]) -> None:
        arch_config = self.arch_config
        valid_fields = {f.name for f in fields(arch_config)}

@@ -56,7 +58,7 @@ class ModelConfig:
        if hasattr(arch_config, "__post_init__"):
            arch_config.__post_init__()

-    def update_model_config(self, source_model_dict: Dict[str, Any]) -> None:
+    def update_model_config(self, source_model_dict: dict[str, Any]) -> None:
        assert "arch_config" not in source_model_dict, "Source model config shouldn't contain arch_config."

        valid_fields = {f.name for f in fields(self)}
@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass, field
-from typing import Any, List, Optional, Tuple
+from typing import Any

 from fastvideo.v1.configs.models.base import ArchConfig, ModelConfig
 from fastvideo.v1.layers.quantization import QuantizationConfig
@@ -14,7 +14,7 @@ class DiTArchConfig(ArchConfig):
    _param_names_mapping: dict = field(default_factory=dict)
    _reverse_param_names_mapping: dict = field(default_factory=dict)
    _lora_param_names_mapping: dict = field(default_factory=dict)
-    _supported_attention_backends: Tuple[AttentionBackendEnum, ...] = (
+    _supported_attention_backends: tuple[AttentionBackendEnum, ...] = (
        AttentionBackendEnum.SLIDING_TILE_ATTN, AttentionBackendEnum.SAGE_ATTN,
        AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TORCH_SDPA,
        AttentionBackendEnum.VIDEO_SPARSE_ATTN)
@@ -22,7 +22,7 @@ class DiTArchConfig(ArchConfig):
    hidden_size: int = 0
    num_attention_heads: int = 0
    num_channels_latents: int = 0
-    exclude_lora_layers: List[str] = field(default_factory=list)
+    exclude_lora_layers: list[str] = field(default_factory=list)

    def __post_init__(self) -> None:
        if not self._compile_conditions:
@@ -35,7 +35,7 @@ class DiTConfig(ModelConfig):

    # FastVideoDiT-specific parameters
    prefix: str = ""
-    quant_config: Optional[QuantizationConfig] = None
+    quant_config: QuantizationConfig | None = None

    @staticmethod
    def add_cli_args(parser: Any, prefix: str = "dit-config") -> Any:
@@ -1,6 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass, field
-from typing import List, Optional, Tuple

 import torch

@@ -160,14 +159,14 @@ class HunyuanVideoArchConfig(DiTArchConfig):
    num_layers: int = 20
    num_single_layers: int = 40
    num_refiner_layers: int = 2
-    rope_axes_dim: Tuple[int, int, int] = (16, 56, 56)
+    rope_axes_dim: tuple[int, int, int] = (16, 56, 56)
    guidance_embeds: bool = False
-    dtype: Optional[torch.dtype] = None
+    dtype: torch.dtype | None = None
    text_embed_dim: int = 4096
    pooled_projection_dim: int = 768
    rope_theta: int = 256
    qk_norm: str = "rms_norm"
-    exclude_lora_layers: List[str] = field(
+    exclude_lora_layers: list[str] = field(
        default_factory=lambda: ["img_in", "txt_in", "time_in", "vector_in"])

    def __post_init__(self):
@@ -1,17 +1,14 @@
 # SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass, field
-from typing import List, Optional, Tuple, Union

 from fastvideo.v1.configs.models.dits.base import DiTArchConfig, DiTConfig


-def is_blocks(n: str, m) -> bool:
-    return "blocks" in n and str.isdigit(n.split(".")[-1])
-
-
@dataclass
 class StepVideoArchConfig(DiTArchConfig):
-    _fsdp_shard_conditions: list = field(default_factory=lambda: [is_blocks])
+    _fsdp_shard_conditions: list = field(
+        default_factory=lambda:
+        [lambda n, m: "transformer_blocks" in n and n.split(".")[-1].isdigit()])

    _param_names_mapping: dict = field(
        default_factory=lambda: {
@@ -41,18 +38,18 @@ class StepVideoArchConfig(DiTArchConfig):
    num_attention_heads: int = 48
    attention_head_dim: int = 128
    in_channels: int = 64
-    out_channels: Optional[int] = 64
+    out_channels: int | None = 64
    num_layers: int = 48
    dropout: float = 0.0
    patch_size: int = 1
    norm_type: str = "ada_norm_single"
    norm_elementwise_affine: bool = False
    norm_eps: float = 1e-6
-    caption_channels: Optional[Union[int, List[int], Tuple[int, ...]]] = field(
+    caption_channels: int | list[int] | tuple[int, ...] | None = field(
        default_factory=lambda: [6144, 1024])
-    attention_type: Optional[str] = "torch"
-    use_additional_conditions: Optional[bool] = False
-    exclude_lora_layers: List[str] = field(default_factory=lambda: [])
+    attention_type: str | None = "torch"
+    use_additional_conditions: bool | None = False
+    exclude_lora_layers: list[str] = field(default_factory=lambda: [])

    def __post_init__(self):
        self.hidden_size = self.num_attention_heads * self.attention_head_dim
@@ -1,6 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass, field
-from typing import List, Optional, Tuple

 from fastvideo.v1.configs.models.dits.base import DiTArchConfig, DiTConfig

@@ -74,7 +73,7 @@ class WanVideoArchConfig(DiTArchConfig):
            r"^blocks\.(\d+)\.ffn\.2\.(.*)$": r"blocks.\1.ffn.fc_out.\2",
        })

-    patch_size: Tuple[int, int, int] = (1, 2, 2)
+    patch_size: tuple[int, int, int] = (1, 2, 2)
    text_len = 512
    num_attention_heads: int = 40
    attention_head_dim: int = 128
@@ -87,10 +86,10 @@ class WanVideoArchConfig(DiTArchConfig):
    cross_attn_norm: bool = True
    qk_norm: str = "rms_norm_across_heads"
    eps: float = 1e-6
-    image_dim: Optional[int] = None
-    added_kv_proj_dim: Optional[int] = None
+    image_dim: int | None = None
+    added_kv_proj_dim: int | None = None
    rope_max_seq_len: int = 1024
-    exclude_lora_layers: List[str] = field(default_factory=lambda: ["embedder"])
+    exclude_lora_layers: list[str] = field(default_factory=lambda: ["embedder"])

    def __post_init__(self):
        super().__post_init__()
@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass, field
-from typing import Any, Dict, List, Optional, Tuple
+from typing import Any

 import torch

@@ -11,8 +11,8 @@ from fastvideo.v1.platforms import AttentionBackendEnum

@dataclass
 class EncoderArchConfig(ArchConfig):
-    architectures: List[str] = field(default_factory=lambda: [])
-    _supported_attention_backends: Tuple[AttentionBackendEnum, ...] = (
+    architectures: list[str] = field(default_factory=lambda: [])
+    _supported_attention_backends: tuple[AttentionBackendEnum, ...] = (
        AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TORCH_SDPA)
    output_hidden_states: bool = False
    use_return_dict: bool = True
@@ -32,8 +32,11 @@ class TextEncoderArchConfig(EncoderArchConfig):
    output_past: bool = True
    scalable_attention: bool = True
    tie_word_embeddings: bool = False
-
-    tokenizer_kwargs: Dict[str, Any] = field(default_factory=dict)
+    stacked_params_mapping: list[tuple[str, str, str]] = field(
+        default_factory=list
+    )  # mapping from huggingface weight names to custom names
+    tokenizer_kwargs: dict[str, Any] = field(default_factory=dict)
+    _fsdp_shard_conditions: list = field(default_factory=lambda: [])

    def __post_init__(self) -> None:
        self.tokenizer_kwargs = {
@@ -50,11 +53,11 @@ class ImageEncoderArchConfig(EncoderArchConfig):

@dataclass
 class BaseEncoderOutput:
-    last_hidden_state: Optional[torch.FloatTensor] = None
-    pooler_output: Optional[torch.FloatTensor] = None
-    hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
-    attentions: Optional[Tuple[torch.FloatTensor, ...]] = None
-    attention_mask: Optional[torch.Tensor] = None
+    last_hidden_state: torch.FloatTensor | None = None
+    pooler_output: torch.FloatTensor | None = None
+    hidden_states: tuple[torch.FloatTensor, ...] | None = None
+    attentions: tuple[torch.FloatTensor, ...] | None = None
+    attention_mask: torch.Tensor | None = None


@dataclass
@@ -62,8 +65,8 @@ class EncoderConfig(ModelConfig):
    arch_config: ArchConfig = field(default_factory=EncoderArchConfig)

    prefix: str = ""
-    quant_config: Optional[QuantizationConfig] = None
-    lora_config: Optional[Any] = None
+    quant_config: QuantizationConfig | None = None
+    lora_config: Any | None = None


@dataclass
@@ -1,6 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass, field
-from typing import Optional

 from fastvideo.v1.configs.models.encoders.base import (ImageEncoderArchConfig,
                                                       ImageEncoderConfig,
@@ -8,6 +7,14 @@ from fastvideo.v1.configs.models.encoders.base import (ImageEncoderArchConfig,
                                                       TextEncoderConfig)


+def _is_transformer_layer(n: str, m) -> bool:
+    return "layers" in n and str.isdigit(n.split(".")[-1])
+
+
+def _is_embeddings(n: str, m) -> bool:
+    return n.endswith("embeddings")
+
+
@dataclass
 class CLIPTextArchConfig(TextEncoderArchConfig):
    vocab_size: int = 49408
@@ -27,6 +34,15 @@ class CLIPTextArchConfig(TextEncoderArchConfig):
    bos_token_id: int = 49406
    eos_token_id: int = 49407
    text_len: int = 77
+    stacked_params_mapping: list[tuple[str, str,
+                                       str]] = field(default_factory=lambda: [
+                                           # (param_name, shard_name, shard_id)
+                                           ("qkv_proj", "q_proj", "q"),
+                                           ("qkv_proj", "k_proj", "k"),
+                                           ("qkv_proj", "v_proj", "v"),
+                                       ])
+    _fsdp_shard_conditions: list = field(
+        default_factory=lambda: [_is_transformer_layer, _is_embeddings])


@dataclass
@@ -45,6 +61,13 @@ class CLIPVisionArchConfig(ImageEncoderArchConfig):
    attention_dropout: float = 0.0
    initializer_range: float = 0.02
    initializer_factor: float = 1.0
+    stacked_params_mapping: list[tuple[str, str,
+                                       str]] = field(default_factory=lambda: [
+                                           # (param_name, shard_name, shard_id)
+                                           ("qkv_proj", "q_proj", "q"),
+                                           ("qkv_proj", "k_proj", "k"),
+                                           ("qkv_proj", "v_proj", "v"),
+                                       ])


@dataclass
@@ -52,8 +75,8 @@ class CLIPTextConfig(TextEncoderConfig):
    arch_config: TextEncoderArchConfig = field(
        default_factory=CLIPTextArchConfig)

-    num_hidden_layers_override: Optional[int] = None
-    require_post_norm: Optional[bool] = None
+    num_hidden_layers_override: int | None = None
+    require_post_norm: bool | None = None
    prefix: str = "clip"


@@ -62,6 +85,6 @@ class CLIPVisionConfig(ImageEncoderConfig):
    arch_config: ImageEncoderArchConfig = field(
        default_factory=CLIPVisionArchConfig)

-    num_hidden_layers_override: Optional[int] = None
-    require_post_norm: Optional[bool] = None
+    num_hidden_layers_override: int | None = None
+    require_post_norm: bool | None = None
    prefix: str = "clip"
@@ -1,11 +1,22 @@
 # SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass, field
-from typing import Optional

 from fastvideo.v1.configs.models.encoders.base import (TextEncoderArchConfig,
                                                       TextEncoderConfig)


+def _is_transformer_layer(n: str, m) -> bool:
+    return "layers" in n and str.isdigit(n.split(".")[-1])
+
+
+def _is_embeddings(n: str, m) -> bool:
+    return n.endswith("embed_tokens")
+
+
+def _is_final_norm(n: str, m) -> bool:
+    return n.endswith("norm")
+
+
@dataclass
 class LlamaArchConfig(TextEncoderArchConfig):
    vocab_size: int = 32000
@@ -13,7 +24,7 @@ class LlamaArchConfig(TextEncoderArchConfig):
    intermediate_size: int = 11008
    num_hidden_layers: int = 32
    num_attention_heads: int = 32
-    num_key_value_heads: Optional[int] = None
+    num_key_value_heads: int | None = None
    hidden_act: str = "silu"
    max_position_embeddings: int = 2048
    initializer_range: float = 0.02
@@ -25,13 +36,25 @@ class LlamaArchConfig(TextEncoderArchConfig):
    pretraining_tp: int = 1
    tie_word_embeddings: bool = False
    rope_theta: float = 10000.0
-    rope_scaling: Optional[float] = None
+    rope_scaling: float | None = None
    attention_bias: bool = False
    attention_dropout: float = 0.0
    mlp_bias: bool = False
-    head_dim: Optional[int] = None
+    head_dim: int | None = None
    hidden_state_skip_layer: int = 2
    text_len: int = 256
+    stacked_params_mapping: list[tuple[str, str, str]] = field(
+        default_factory=lambda: [
+            # (param_name, shard_name, shard_id)
+            (".qkv_proj", ".q_proj", "q"),
+            (".qkv_proj", ".k_proj", "k"),
+            (".qkv_proj", ".v_proj", "v"),
+            (".gate_up_proj", ".gate_proj", 0),  # type: ignore
+            (".gate_up_proj", ".up_proj", 1),  # type: ignore
+        ])
+    _fsdp_shard_conditions: list = field(
+        default_factory=lambda:
+        [_is_transformer_layer, _is_embeddings, _is_final_norm])


@dataclass
@@ -1,11 +1,22 @@
 # SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass, field
-from typing import Optional

 from fastvideo.v1.configs.models.encoders.base import (TextEncoderArchConfig,
                                                       TextEncoderConfig)


+def _is_transformer_layer(n: str, m) -> bool:
+    return "block" in n and str.isdigit(n.split(".")[-1])
+
+
+def _is_embeddings(n: str, m) -> bool:
+    return n.endswith("shared")
+
+
+def _is_final_layernorm(n: str, m) -> bool:
+    return n.endswith("final_layer_norm")
+
+
@dataclass
 class T5ArchConfig(TextEncoderArchConfig):
    vocab_size: int = 32128
@@ -13,7 +24,7 @@ class T5ArchConfig(TextEncoderArchConfig):
    d_kv: int = 64
    d_ff: int = 2048
    num_layers: int = 6
-    num_decoder_layers: Optional[int] = None
+    num_decoder_layers: int | None = None
    num_heads: int = 8
    relative_attention_num_buckets: int = 32
    relative_attention_max_distance: int = 128
@@ -29,6 +40,16 @@ class T5ArchConfig(TextEncoderArchConfig):
    eos_token_id: int = 1
    classifier_dropout: float = 0.0
    text_len: int = 512
+    stacked_params_mapping: list[tuple[str, str,
+                                       str]] = field(default_factory=lambda: [
+                                           # (param_name, shard_name, shard_id)
+                                           (".qkv_proj", ".q", "q"),
+                                           (".qkv_proj", ".k", "k"),
+                                           (".qkv_proj", ".v", "v"),
+                                       ])
+    _fsdp_shard_conditions: list = field(
+        default_factory=lambda:
+        [_is_transformer_layer, _is_embeddings, _is_final_layernorm])

    # Referenced from https://github.com/huggingface/transformers/blob/main/src/transformers/models/t5/configuration_t5.py
    def __post_init__(self):
@@ -2,7 +2,7 @@
 import argparse
 import dataclasses
 from dataclasses import dataclass, field
-from typing import Any, Union
+from typing import Any

 import torch

@@ -12,7 +12,7 @@ from fastvideo.v1.utils import StoreBoolean

@dataclass
 class VAEArchConfig(ArchConfig):
-    scaling_factor: Union[float, torch.tensor] = 0
+    scaling_factor: float | torch.Tensor = 0

    temporal_compression_ratio: int = 4
    spatial_compression_ratio: int = 8
@@ -1,6 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass, field
-from typing import Tuple

 from fastvideo.v1.configs.models.vaes.base import VAEArchConfig, VAEConfig

@@ -10,19 +9,19 @@ class HunyuanVAEArchConfig(VAEArchConfig):
    in_channels: int = 3
    out_channels: int = 3
    latent_channels: int = 16
-    down_block_types: Tuple[str, ...] = (
+    down_block_types: tuple[str, ...] = (
        "HunyuanVideoDownBlock3D",
        "HunyuanVideoDownBlock3D",
        "HunyuanVideoDownBlock3D",
        "HunyuanVideoDownBlock3D",
    )
-    up_block_types: Tuple[str, ...] = (
+    up_block_types: tuple[str, ...] = (
        "HunyuanVideoUpBlock3D",
        "HunyuanVideoUpBlock3D",
        "HunyuanVideoUpBlock3D",
        "HunyuanVideoUpBlock3D",
    )
-    block_out_channels: Tuple[int, ...] = (128, 256, 512, 512)
+    block_out_channels: tuple[int, ...] = (128, 256, 512, 512)
    layers_per_block: int = 2
    act_fn: str = "silu"
    norm_num_groups: int = 32
@@ -1,6 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass, field
-from typing import Tuple

 import torch

@@ -11,12 +10,12 @@ from fastvideo.v1.configs.models.vaes.base import VAEArchConfig, VAEConfig
 class WanVAEArchConfig(VAEArchConfig):
    base_dim: int = 96
    z_dim: int = 16
-    dim_mult: Tuple[int, ...] = (1, 2, 4, 4)
+    dim_mult: tuple[int, ...] = (1, 2, 4, 4)
    num_res_blocks: int = 2
-    attn_scales: Tuple[float, ...] = ()
-    temperal_downsample: Tuple[bool, ...] = (False, True, True)
+    attn_scales: tuple[float, ...] = ()
+    temperal_downsample: tuple[bool, ...] = (False, True, True)
    dropout: float = 0.0
-    latents_mean: Tuple[float, ...] = (
+    latents_mean: tuple[float, ...] = (
        -0.7571,
        -0.7089,
        -0.9113,
@@ -34,7 +33,7 @@ class WanVAEArchConfig(VAEArchConfig):
        0.2503,
        -0.2921,
    )
-    latents_std: Tuple[float, ...] = (
+    latents_std: tuple[float, ...] = (
        2.8184,
        1.4541,
        2.3275,
@@ -1,8 +1,9 @@
 # SPDX-License-Identifier: Apache-2.0
 import json
+from collections.abc import Callable
 from dataclasses import asdict, dataclass, field, fields
 from enum import Enum
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union, cast
+from typing import Any, cast

 import torch

@@ -39,11 +40,11 @@ def postprocess_text(output: BaseEncoderOutput) -> torch.tensor:
 class PipelineConfig:
    """Base configuration for all pipeline architectures."""
    model_path: str = ""
-    pipeline_config_path: Optional[str] = None
+    pipeline_config_path: str | None = None

    # Video generation parameters
    embedded_cfg_scale: float = 6.0
-    flow_shift: Optional[float] = None
+    flow_shift: float | None = None
    disable_autocast: bool = False

    # Model configuration
@@ -52,7 +53,7 @@ class PipelineConfig:

    # VAE configuration
    vae_config: VAEConfig = field(default_factory=VAEConfig)
-    vae_precision: str = "fp16"
+    vae_precision: str = "fp32"
    vae_tiling: bool = True
    vae_sp: bool = True

@@ -62,30 +63,29 @@ class PipelineConfig:

    # Text encoder configuration
    DEFAULT_TEXT_ENCODER_PRECISIONS = ("fp16", )
-    text_encoder_configs: Tuple[EncoderConfig, ...] = field(
+    text_encoder_configs: tuple[EncoderConfig, ...] = field(
        default_factory=lambda: (EncoderConfig(), ))
-    text_encoder_precisions: Tuple[str, ...] = field(
+    text_encoder_precisions: tuple[str, ...] = field(
        default_factory=lambda: ("fp16", ))
-    preprocess_text_funcs: Tuple[Callable[[str], str], ...] = field(
+    preprocess_text_funcs: tuple[Callable[[str], str], ...] = field(
        default_factory=lambda: (preprocess_text, ))
-    postprocess_text_funcs: Tuple[Callable[[BaseEncoderOutput], torch.tensor],
+    postprocess_text_funcs: tuple[Callable[[BaseEncoderOutput], torch.tensor],
                                  ...] = field(default_factory=lambda:
                                               (postprocess_text, ))

    # LoRA parameters
-    lora_path: Optional[str] = None
-    lora_nickname: Optional[
-        str] = "default"  # for swapping adapters in the pipeline
-    lora_target_names: Optional[List[
-        str]] = None  # can restrict list of layers to adapt, e.g. ["q_proj"]
+    lora_path: str | None = None
+    lora_nickname: str | None = "default"  # for swapping adapters in the pipeline
+    lora_target_names: list[
+        str] | None = None  # can restrict list of layers to adapt, e.g. ["q_proj"]

    # StepVideo specific parameters
-    pos_magic: Optional[str] = None
-    neg_magic: Optional[str] = None
-    timesteps_scale: Optional[bool] = None
+    pos_magic: str | None = None
+    neg_magic: str | None = None
+    timesteps_scale: bool | None = None

    # STA (Sliding Tile Attention) parameters
-    mask_strategy_file_path: Optional[str] = None
+    mask_strategy_file_path: str | None = None
    STA_mode: STA_Mode = STA_Mode.STA_INFERENCE
    skip_time_steps: int = 15

@@ -217,7 +217,7 @@ class PipelineConfig:
        return parser

    def update_config_from_dict(self,
-                                args: Dict[str, Any],
+                                args: dict[str, Any],
                                prefix: str = "") -> None:
        prefix_with_dot = f"{prefix}." if (prefix.strip() != "") else ""
        update_config_from_args(self, args, prefix, pop_args=True)
@@ -243,7 +243,7 @@ class PipelineConfig:

    @classmethod
    def from_kwargs(cls,
-                    kwargs: Dict[str, Any],
+                    kwargs: dict[str, Any],
                    config_cli_prefix: str = "") -> "PipelineConfig":
        """
        Load PipelineConfig from kwargs Dictionary.
@@ -255,11 +255,11 @@ class PipelineConfig:

        prefix_with_dot = f"{config_cli_prefix}." if (config_cli_prefix.strip()
                                                      != "") else ""
-        model_path: Optional[str] = kwargs.get(prefix_with_dot + 'model_path',
-                                               None) or kwargs.get('model_path')
-        pipeline_config_or_path: Optional[Union[str, PipelineConfig, Dict[
-            str, Any]]] = kwargs.get(prefix_with_dot + 'pipeline_config',
-                                     None) or kwargs.get('pipeline_config')
+        model_path: str | None = kwargs.get(prefix_with_dot + 'model_path',
+                                            None) or kwargs.get('model_path')
+        pipeline_config_or_path: str | PipelineConfig | dict[
+            str, Any] | None = kwargs.get(prefix_with_dot + 'pipeline_config',
+                                          None) or kwargs.get('pipeline_config')
        if model_path is None:
            raise ValueError("model_path is required in kwargs")

@@ -344,7 +344,7 @@ class PipelineConfig:
            input_pipeline_dict = json.load(f)
        self.update_pipeline_config(input_pipeline_dict)

-    def update_pipeline_config(self, source_pipeline_dict: Dict[str,
+    def update_pipeline_config(self, source_pipeline_dict: dict[str,
                                                                Any]) -> None:
        for f in fields(self):
            key = f.name
@@ -360,8 +360,9 @@ class PipelineConfig:
                    assert len(current_value) == len(
                        new_value
                    ), "Users shouldn't delete or add text encoder config objects in your json"
-                    for target_config, source_config in zip(
-                            current_value, new_value):
+                    for target_config, source_config in zip(current_value,
+                                                            new_value,
+                                                            strict=True):
                        target_config.update_model_config(source_config)
                else:
                    setattr(self, key, new_value)
@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
+from collections.abc import Callable
 from dataclasses import dataclass, field
-from typing import Callable, Tuple, TypedDict
+from typing import TypedDict

 import torch

@@ -39,7 +40,7 @@ def llama_preprocess_text(prompt: str) -> str:
 def llama_postprocess_text(outputs: BaseEncoderOutput) -> torch.tensor:
    hidden_state_skip_layer = 2
    assert outputs.hidden_states is not None
-    hidden_states: Tuple[torch.Tensor, ...] = outputs.hidden_states
+    hidden_states: tuple[torch.Tensor, ...] = outputs.hidden_states
    last_hidden_state: torch.tensor = hidden_states[-(hidden_state_skip_layer +
                                                      1)]
    crop_start = prompt_template_video.get("crop_start", -1)
@@ -70,11 +71,11 @@ class HunyuanConfig(PipelineConfig):
    flow_shift: int = 7

    # Text encoding stage
-    text_encoder_configs: Tuple[EncoderConfig, ...] = field(
+    text_encoder_configs: tuple[EncoderConfig, ...] = field(
        default_factory=lambda: (LlamaConfig(), CLIPTextConfig()))
-    preprocess_text_funcs: Tuple[Callable[[str], str], ...] = field(
+    preprocess_text_funcs: tuple[Callable[[str], str], ...] = field(
        default_factory=lambda: (llama_preprocess_text, clip_preprocess_text))
-    postprocess_text_funcs: Tuple[
+    postprocess_text_funcs: tuple[
        Callable[[BaseEncoderOutput], torch.tensor],
        ...] = field(default_factory=lambda:
                     (llama_postprocess_text, clip_postprocess_text))
@@ -82,7 +83,7 @@ class HunyuanConfig(PipelineConfig):
    # Precision for each component
    dit_precision: str = "bf16"
    vae_precision: str = "fp16"
-    text_encoder_precisions: Tuple[str, ...] = field(
+    text_encoder_precisions: tuple[str, ...] = field(
        default_factory=lambda: ("fp16", "fp16"))

    def __post_init__(self):
@@ -2,7 +2,7 @@
 """Registry for pipeline weight-specific configurations."""

 import os
-from typing import Callable, Dict, Optional, Type
+from collections.abc import Callable

 from fastvideo.v1.configs.pipelines.base import PipelineConfig
 from fastvideo.v1.configs.pipelines.hunyuan import (FastHunyuanConfig,
@@ -19,10 +19,11 @@ from fastvideo.v1.utils import (maybe_download_model_index,
 logger = init_logger(__name__)

 # Registry maps specific model weights to their config classes
-PIPE_NAME_TO_CONFIG: Dict[str, Type[PipelineConfig]] = {
+PIPE_NAME_TO_CONFIG: dict[str, type[PipelineConfig]] = {
    "FastVideo/FastHunyuan-diffusers": FastHunyuanConfig,
    "hunyuanvideo-community/HunyuanVideo": HunyuanConfig,
    "Wan-AI/Wan2.1-T2V-1.3B-Diffusers": WanT2V480PConfig,
+    "weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers": WanI2V480PConfig,
    "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers": WanI2V480PConfig,
    "Wan-AI/Wan2.1-I2V-14B-720P-Diffusers": WanI2V720PConfig,
    "Wan-AI/Wan2.1-T2V-14B-Diffusers": WanT2V720PConfig,
@@ -31,7 +32,7 @@ PIPE_NAME_TO_CONFIG: Dict[str, Type[PipelineConfig]] = {
 }

 # For determining pipeline type from model ID
-PIPELINE_DETECTOR: Dict[str, Callable[[str], bool]] = {
+PIPELINE_DETECTOR: dict[str, Callable[[str], bool]] = {
    "hunyuan": lambda id: "hunyuan" in id.lower(),
    "wanpipeline": lambda id: "wanpipeline" in id.lower(),
    "wanimagetovideo": lambda id: "wanimagetovideo" in id.lower(),
@@ -40,7 +41,7 @@ PIPELINE_DETECTOR: Dict[str, Callable[[str], bool]] = {
 }

 # Fallback configs when exact match isn't found but architecture is detected
-PIPELINE_FALLBACK_CONFIG: Dict[str, Type[PipelineConfig]] = {
+PIPELINE_FALLBACK_CONFIG: dict[str, type[PipelineConfig]] = {
    "hunyuan":
    HunyuanConfig,  # Base Hunyuan config as fallback for any Hunyuan variant
    "wanpipeline":
@@ -52,7 +53,7 @@ PIPELINE_FALLBACK_CONFIG: Dict[str, Type[PipelineConfig]] = {


 def get_pipeline_config_cls_from_name(
-        pipeline_name_or_path: str) -> Type[PipelineConfig]:
+        pipeline_name_or_path: str) -> type[PipelineConfig]:
    """Get the appropriate configuration class for a given pipeline name or path.

    This function implements a multi-step lookup process to find the most suitable
@@ -81,7 +82,7 @@ def get_pipeline_config_cls_from_name(
        - Warning messages are logged when falling back to less specific configurations
    """

-    pipeline_config_cls: Optional[Type[PipelineConfig]] = None
+    pipeline_config_cls: type[PipelineConfig] | None = None

    # First try exact match for specific weights
    if pipeline_name_or_path in PIPE_NAME_TO_CONFIG:
@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
+from collections.abc import Callable
 from dataclasses import dataclass, field
-from typing import Callable, Tuple

 import torch

@@ -17,7 +17,7 @@ def t5_postprocess_text(outputs: BaseEncoderOutput) -> torch.tensor:
    hidden_state: torch.tensor = outputs.last_hidden_state
    seq_lens = mask.gt(0).sum(dim=1).long()
    assert torch.isnan(hidden_state).sum() == 0
-    prompt_embeds = [u[:v] for u, v in zip(hidden_state, seq_lens)]
+    prompt_embeds = [u[:v] for u, v in zip(hidden_state, seq_lens, strict=True)]
    prompt_embeds_tensor: torch.tensor = torch.stack([
        torch.cat([u, u.new_zeros(512 - u.size(0), u.size(1))])
        for u in prompt_embeds
@@ -42,16 +42,16 @@ class WanT2V480PConfig(PipelineConfig):
    flow_shift: int = 3

    # Text encoding stage
-    text_encoder_configs: Tuple[EncoderConfig, ...] = field(
+    text_encoder_configs: tuple[EncoderConfig, ...] = field(
        default_factory=lambda: (T5Config(), ))
-    postprocess_text_funcs: Tuple[Callable[[BaseEncoderOutput], torch.tensor],
+    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 = "fp16"
-    text_encoder_precisions: Tuple[str, ...] = field(
+    vae_precision: str = "fp32"
+    text_encoder_precisions: tuple[str, ...] = field(
        default_factory=lambda: ("fp32", ))

    # WanConfig-specific added parameters
@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass
-from typing import Any, Dict, List, Optional, Union
+from typing import Any

 from fastvideo.v1.logger import init_logger

@@ -16,12 +16,12 @@ class SamplingParam:
    data_type: str = "video"

    # Image inputs
-    image_path: Optional[str] = None
+    image_path: str | None = None

    # Text inputs
-    prompt: Optional[Union[str, List[str]]] = None
-    negative_prompt: Optional[str] = None
-    prompt_path: Optional[str] = None
+    prompt: str | list[str] | None = None
+    negative_prompt: str | None = None
+    prompt_path: str | None = None
    output_path: str = "outputs/"

    # Batch info
@@ -54,7 +54,7 @@ class SamplingParam:
        if self.prompt_path and not self.prompt_path.endswith(".txt"):
            raise ValueError("prompt_path must be a txt file")

-    def update(self, source_dict: Dict[str, Any]) -> None:
+    def update(self, source_dict: dict[str, Any]) -> None:
        for key, value in source_dict.items():
            if hasattr(self, key):
                setattr(self, key, value)
@@ -1,11 +1,13 @@
 # SPDX-License-Identifier: Apache-2.0
 import os
-from typing import Any, Callable, Dict, Optional
+from collections.abc import Callable
+from typing import Any

 from fastvideo.v1.configs.sample.hunyuan import (FastHunyuanSamplingParam,
                                                 HunyuanSamplingParam)
 from fastvideo.v1.configs.sample.stepvideo import StepVideoT2VSamplingParam
-from fastvideo.v1.configs.sample.wan import (WanI2V_14B_480P_SamplingParam,
+from fastvideo.v1.configs.sample.wan import (Wan2_1_Fun_1_3B_InP_SamplingParam,
+                                             WanI2V_14B_480P_SamplingParam,
                                             WanI2V_14B_720P_SamplingParam,
                                             WanT2V_1_3B_SamplingParam,
                                             WanT2V_14B_SamplingParam)
@@ -15,19 +17,21 @@ from fastvideo.v1.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] = {
+SAMPLING_PARAM_REGISTRY: dict[str, Any] = {
    "FastVideo/FastHunyuan-diffusers": FastHunyuanSamplingParam,
    "hunyuanvideo-community/HunyuanVideo": HunyuanSamplingParam,
    "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,
    "FastVideo/stepvideo-t2v-diffusers": StepVideoT2VSamplingParam,
    # Add other specific weight variants
 }

 # For determining pipeline type from model ID
-SAMPLING_PARAM_DETECTOR: Dict[str, Callable[[str], bool]] = {
+SAMPLING_PARAM_DETECTOR: dict[str, Callable[[str], bool]] = {
    "hunyuan": lambda id: "hunyuan" in id.lower(),
    "wanpipeline": lambda id: "wanpipeline" in id.lower(),
    "wanimagetovideo": lambda id: "wanimagetovideo" in id.lower(),
@@ -36,7 +40,7 @@ SAMPLING_PARAM_DETECTOR: Dict[str, Callable[[str], bool]] = {
 }

 # Fallback configs when exact match isn't found but architecture is detected
-SAMPLING_FALLBACK_PARAM: Dict[str, Any] = {
+SAMPLING_FALLBACK_PARAM: dict[str, Any] = {
    "hunyuan":
    HunyuanSamplingParam,  # Base Hunyuan config as fallback for any Hunyuan variant
    "wanpipeline":
@@ -47,8 +51,7 @@ SAMPLING_FALLBACK_PARAM: Dict[str, Any] = {
 }


-def get_sampling_param_cls_for_name(
-        pipeline_name_or_path: str) -> Optional[Any]:
+def get_sampling_param_cls_for_name(pipeline_name_or_path: str) -> Any | None:
    """Get the appropriate sampling param for specific pretrained weights."""

    if os.path.exists(pipeline_name_or_path):
@@ -94,3 +94,20 @@ class WanI2V_14B_720P_SamplingParam(WanT2V_14B_SamplingParam):
                -5784.54975374, 5449.50911966, -1811.16591783, 256.27178429,
                -13.02252404
            ]))
+
+
+# =============================================
+# ============= Wan2.1 Fun Models =============
+# =============================================
+
+
+@dataclass
+class Wan2_1_Fun_1_3B_InP_SamplingParam(SamplingParam):
+    """Sampling parameters for Wan2.1 Fun 1.3B InP model."""
+    height: int = 480
+    width: int = 832
+    num_frames: int = 81
+    fps: int = 16
+    negative_prompt: str | None = "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走"
+    guidance_scale: float = 6.0
+    num_inference_steps: int = 50
@@ -1,8 +1,8 @@
-from typing import Any, Dict
+from typing import Any


 def update_config_from_args(config: Any,
-                            args_dict: Dict[str, Any],
+                            args_dict: dict[str, Any],
                            prefix: str = "",
                            pop_args: bool = False) -> None:
    """
@@ -4,7 +4,7 @@
  "use_cpu_offload": true,
  "disable_autocast": false,
  "precision": "bf16",
-  "vae_precision": "fp16",
+  "vae_precision": "fp32",
  "vae_tiling": false,
  "vae_sp": false,
  "vae_config": {
@@ -4,7 +4,7 @@
  "use_cpu_offload": true,
  "disable_autocast": false,
  "precision": "bf16",
-  "vae_precision": "fp16",
+  "vae_precision": "fp32",
  "vae_tiling": false,
  "vae_sp": false,
  "vae_config": {
@@ -12,7 +12,10 @@ from fastvideo.v1.dataset.validation_dataset import ValidationDataset


 def getdataset(args) -> VideoCaptionMergedDataset:
-    temporal_sample = TemporalRandomCrop(args.num_frames)  # 16 x
+    if args.do_temporal_sample:
+        temporal_sample = TemporalRandomCrop(args.num_frames)  # 16 x
+    else:
+        temporal_sample = None
    norm_fun = Lambda(lambda x: 2.0 * x - 1.0)
    resize_topcrop = [
        CenterCropResizeVideo((args.max_height, args.max_width), top_crop=True),
@@ -11,7 +11,7 @@ from fastvideo.v1.dataset.parquet_dataset_iterable_style import (
    build_parquet_iterable_style_dataloader)
 from fastvideo.v1.distributed import get_world_rank
 from fastvideo.v1.distributed.parallel_state import (
-    cleanup_dist_env_and_memory, get_torch_device,
+    cleanup_dist_env_and_memory, get_local_torch_device,
    maybe_init_distributed_environment_and_model_parallel)
 from fastvideo.v1.logger import init_logger

@@ -128,8 +128,8 @@ def main() -> None:

        # Verify all sums match
        if all(
-                abs(a - b) < 1e-6
-                for a, b in zip(first_pass_sums, second_pass_sums)):
+                abs(a - b) < 1e-6 for a, b in zip(
+                    first_pass_sums, second_pass_sums, strict=True)):
            logger.info(
                "All latent sums match between passes - resume verification successful!"
            )
@@ -148,8 +148,8 @@ def main() -> None:
                break

            # Move data to device
-            latents = latents.to(get_torch_device())
-            embeddings = embeddings.to(get_torch_device())
+            latents = latents.to(get_local_torch_device())
+            embeddings = embeddings.to(get_local_torch_device())

            # Calculate actual batch size
            batch_size = latents.size(0)
@@ -13,7 +13,7 @@ from fastvideo.v1.dataset.parquet_dataset_map_style import (
    build_parquet_map_style_dataloader)
 from fastvideo.v1.distributed import get_world_rank
 from fastvideo.v1.distributed.parallel_state import (
-    cleanup_dist_env_and_memory, get_torch_device,
+    cleanup_dist_env_and_memory, get_local_torch_device,
    maybe_init_distributed_environment_and_model_parallel)
 from fastvideo.v1.logger import init_logger

@@ -144,8 +144,8 @@ def main() -> None:

        # Verify all sums match
        if all(
-                abs(a - b) < 1e-6
-                for a, b in zip(first_pass_sums, second_pass_sums)):
+                abs(a - b) < 1e-6 for a, b in zip(
+                    first_pass_sums, second_pass_sums, strict=True)):
            logger.info(
                "All latent sums match between passes - resume verification successful!"
            )
@@ -165,8 +165,8 @@ def main() -> None:
                break

            # Move data to device
-            latents = latents.to(get_torch_device())
-            embeddings = embeddings.to(get_torch_device())
+            latents = latents.to(get_local_torch_device())
+            embeddings = embeddings.to(get_local_torch_device())

            # Calculate actual batch size
            batch_size = latents.size(0)
@@ -76,7 +76,7 @@ def latent_collate_function(batch):
    # latent_attn_mask: # b t h w
    # text_attn_mask: b 1 l
    # needs to check if the latent/prompt' size and apply padding & attn mask
-    latents, prompt_embeds, prompt_attention_masks = zip(*batch)
+    latents, prompt_embeds, prompt_attention_masks = zip(*batch, strict=True)
    # calculate max shape
    max_t = max([latent.shape[1] for latent in latents])
    max_h = max([latent.shape[2] for latent in latents])
@@ -1,19 +1,17 @@
 import os
 import pickle
 import random
-from typing import Dict, List, Tuple

 import numpy as np
 import pyarrow as pa
 import pyarrow.parquet as pq
-import torch
 import tqdm
 from torch.utils.data import IterableDataset, get_worker_info
 from torchdata.stateful_dataloader import StatefulDataLoader

 from fastvideo.v1.dataset.utils import collate_latents_embs_masks
-from fastvideo.v1.distributed import (get_sp_world_size, get_world_rank,
-                                      get_world_size)
+from fastvideo.v1.distributed import (get_sp_world_size, get_world_group,
+                                      get_world_rank, get_world_size)
 from fastvideo.v1.logger import init_logger

 logger = init_logger(__name__)
@@ -139,7 +137,7 @@ def shard_parquet_files_across_sp_groups_and_workers(
    num_sp_groups: int,
    num_workers: int,
    seed: int = 42,
-) -> Tuple[List[List[str]], List[int], List[Dict[str, int]]]:
+) -> tuple[list[list[str]], list[int], list[dict[str, int]]]:
    """
    Shard parquet files across SP groups and workers in a balanced way.
    
@@ -158,94 +156,123 @@ def shard_parquet_files_across_sp_groups_and_workers(
    # Check if sharding plan already exists
    sharding_info_dir = os.path.join(
        path, f"sharding_info_{num_sp_groups}_sp_groups_{num_workers}_workers")
-    if os.path.exists(sharding_info_dir):
-        logger.info("Sharding plan already exists")
-        logger.info("Loading sharding plan from %s", sharding_info_dir)
-        try:
+
+    # Only rank 0 handles cache checking and file scanning
+    if get_world_rank() == 0:
+        cache_loaded = False
+        shard_parquet_files = None
+        shard_total_samples = None
+        shard_parquet_lengths = None
+
+        # First try to load existing sharding plan
+        if os.path.exists(sharding_info_dir):
+            logger.info("Loading sharding plan from %s", sharding_info_dir)
+            try:
+                with open(
+                        os.path.join(sharding_info_dir,
+                                     "shard_parquet_files.pkl"), "rb") as f:
+                    shard_parquet_files = pickle.load(f)
+                with open(
+                        os.path.join(sharding_info_dir,
+                                     "shard_total_samples.pkl"), "rb") as f:
+                    shard_total_samples = pickle.load(f)
+                with open(
+                        os.path.join(sharding_info_dir,
+                                     "shard_parquet_lengths.pkl"), "rb") as f:
+                    shard_parquet_lengths = pickle.load(f)
+                cache_loaded = True
+                logger.info("Successfully loaded sharding plan")
+            except Exception as e:
+                logger.error("Error loading sharding plan: %s", str(e))
+                logger.info("Falling back to creating new sharding plan")
+                cache_loaded = False
+
+        # If cache not loaded (either doesn't exist or failed to load), create sharding plan
+        if not cache_loaded:
+            logger.info("Creating new sharding plan")
+            logger.info("Scanning for parquet files in %s", path)
+
+            # Find all parquet files
+            parquet_files = []
+
+            for root, _, files in os.walk(path):
+                for file in files:
+                    if file.endswith('.parquet'):
+                        parquet_files.append(os.path.join(root, file))
+
+            if not parquet_files:
+                raise ValueError("No parquet files found in %s", path)
+
+            # Calculate file lengths efficiently using a single pass
+            logger.info("Calculating file lengths...")
+            lengths = []
+            for file in tqdm.tqdm(parquet_files, desc="Reading parquet files"):
+                lengths.append(pq.ParquetFile(file).metadata.num_rows)
+
+            total_samples = sum(lengths)
+            logger.info("Found %d files with %d total samples",
+                        len(parquet_files), total_samples)
+
+            # Sort files by length for better balancing
+            sorted_indices = np.argsort(lengths)
+            sorted_files = [parquet_files[i] for i in sorted_indices]
+            sorted_lengths = [lengths[i] for i in sorted_indices]
+
+            # Create shards
+            num_shards = num_sp_groups * num_workers
+            shard_parquet_files = [[] for _ in range(num_shards)]
+            shard_total_samples = [0] * num_shards
+            shard_parquet_lengths = [{} for _ in range(num_shards)]
+
+            # Distribute files to shards using a greedy approach
+            logger.info("Distributing files to shards...")
+            for file, length in zip(reversed(sorted_files),
+                                    reversed(sorted_lengths),
+                                    strict=True):
+                # Find shard with minimum current length
+                target_shard = np.argmin(shard_total_samples)
+                shard_parquet_files[target_shard].append(file)
+                shard_total_samples[target_shard] += length
+                shard_parquet_lengths[target_shard][file] = length
+            #randomize each shard
+            for shard in shard_parquet_files:
+                random.seed(seed)
+                random.shuffle(shard)
+
+            # Save the sharding plan
+            os.makedirs(sharding_info_dir, exist_ok=True)
            with open(
                    os.path.join(sharding_info_dir, "shard_parquet_files.pkl"),
-                    "rb") as f:
-                shard_parquet_files = pickle.load(f)
+                    "wb") as f:
+                pickle.dump(shard_parquet_files, f)
            with open(
                    os.path.join(sharding_info_dir, "shard_total_samples.pkl"),
-                    "rb") as f:
-                shard_total_samples = pickle.load(f)
+                    "wb") as f:
+                pickle.dump(shard_total_samples, f)
            with open(
                    os.path.join(sharding_info_dir,
-                                 "shard_parquet_lengths.pkl"), "rb") as f:
-                shard_parquet_lengths = pickle.load(f)
-            return shard_parquet_files, shard_total_samples, shard_parquet_lengths
-        except Exception as e:
-            logger.error("Error loading sharding plan: %s", str(e))
-            logger.info("Falling back to creating new sharding plan")
+                                 "shard_parquet_lengths.pkl"), "wb") as f:
+                pickle.dump(shard_parquet_lengths, f)
+            logger.info("Saved sharding info to %s", sharding_info_dir)

-    if get_world_rank() == 0:
-        logger.info("Scanning for parquet files in %s", path)
+    # Wait for rank 0 to finish creating/loading sharding plan
+    world_group = get_world_group()
+    world_group.barrier()

-        # Find all parquet files
-        parquet_files = []
+    # Now all ranks load the sharding plan (it should exist and be valid now)
+    logger.info("Loading sharding plan from %s after barrier",
+                sharding_info_dir)
+    with open(os.path.join(sharding_info_dir, "shard_parquet_files.pkl"),
+              "rb") as f:
+        shard_parquet_files = pickle.load(f)
+    with open(os.path.join(sharding_info_dir, "shard_total_samples.pkl"),
+              "rb") as f:
+        shard_total_samples = pickle.load(f)
+    with open(os.path.join(sharding_info_dir, "shard_parquet_lengths.pkl"),
+              "rb") as f:
+        shard_parquet_lengths = pickle.load(f)

-        for root, _, files in os.walk(path):
-            for file in files:
-                if file.endswith('.parquet'):
-                    parquet_files.append(os.path.join(root, file))
-
-        if not parquet_files:
-            raise ValueError("No parquet files found in %s", path)
-
-        # Calculate file lengths efficiently using a single pass
-        logger.info("Calculating file lengths...")
-        lengths = []
-        for file in tqdm.tqdm(parquet_files, desc="Reading parquet files"):
-            lengths.append(pq.ParquetFile(file).metadata.num_rows)
-
-        total_samples = sum(lengths)
-        logger.info("Found %d files with %d total samples", len(parquet_files),
-                    total_samples)
-
-        # Sort files by length for better balancing
-        sorted_indices = np.argsort(lengths)
-        sorted_files = [parquet_files[i] for i in sorted_indices]
-        sorted_lengths = [lengths[i] for i in sorted_indices]
-
-        # Create shards
-        num_shards = num_sp_groups * num_workers
-        shard_parquet_files = [[] for _ in range(num_shards)]
-        shard_total_samples = [0] * num_shards
-        shard_parquet_lengths = [{} for _ in range(num_shards)]
-
-        # Distribute files to shards using a greedy approach
-        logger.info("Distributing files to shards...")
-        for file, length in zip(reversed(sorted_files),
-                                reversed(sorted_lengths)):
-            # Find shard with minimum current length
-            target_shard = np.argmin(shard_total_samples)
-            shard_parquet_files[target_shard].append(file)
-            shard_total_samples[target_shard] += length
-            shard_parquet_lengths[target_shard][file] = length
-        #randomize each shard
-        for shard in shard_parquet_files:
-            random.seed(seed)
-            random.shuffle(shard)
-
-        save_dir = os.path.join(
-            path,
-            f"sharding_info_{num_sp_groups}_sp_groups_{num_workers}_workers")
-        os.makedirs(save_dir, exist_ok=True)
-        with open(os.path.join(save_dir, "shard_parquet_files.pkl"), "wb") as f:
-            pickle.dump(shard_parquet_files, f)
-        with open(os.path.join(save_dir, "shard_total_samples.pkl"), "wb") as f:
-            pickle.dump(shard_total_samples, f)
-        with open(os.path.join(save_dir, "shard_parquet_lengths.pkl"),
-                  "wb") as f:
-            pickle.dump(shard_parquet_lengths, f)
-        logger.info("Saved sharding info to %s", save_dir)
-
-    # wait for all ranks to finish
-    torch.distributed.barrier()
-    # recursive call
-    return shard_parquet_files_across_sp_groups_and_workers(
-        path, num_sp_groups, num_workers, seed)
+    return shard_parquet_files, shard_total_samples, shard_parquet_lengths


 def build_parquet_iterable_style_dataloader(
@@ -257,7 +284,7 @@ def build_parquet_iterable_style_dataloader(
    text_padding_length: int = 512,
    seed: int = 42,
    read_batch_size: int = 32
-) -> Tuple[LatentsParquetIterStyleDataset, StatefulDataLoader]:
+) -> tuple[LatentsParquetIterStyleDataset, StatefulDataLoader]:
    """Build a dataloader for the LatentsParquetIterStyleDataset."""
    dataset = LatentsParquetIterStyleDataset(
        path=path,
@@ -1,7 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 import os
 import pickle
-from typing import Any, Dict, List, Tuple
+from typing import Any

 import pyarrow as pa
 import pyarrow.parquet as pq
@@ -13,14 +13,14 @@ from torch.utils.data import Dataset, Sampler
 from torchdata.stateful_dataloader import StatefulDataLoader

 from fastvideo.v1.dataset.utils import collate_rows_from_parquet_schema
-from fastvideo.v1.distributed import (get_sp_world_size, get_world_rank,
-                                      get_world_size)
+from fastvideo.v1.distributed import (get_sp_world_size, get_world_group,
+                                      get_world_rank, get_world_size)
 from fastvideo.v1.logger import init_logger

 logger = init_logger(__name__)


-class DP_SP_BatchSampler(Sampler[List[int]]):
+class DP_SP_BatchSampler(Sampler[list[int]]):
    """
    A simple sequential batch sampler that yields batches of indices.
    """
@@ -97,49 +97,68 @@ def get_parquet_files_and_length(path: str):
    cache_dir = os.path.join(path, "map_style_cache")
    cache_file = os.path.join(cache_dir, "file_info.pkl")

-    if os.path.exists(cache_file):
-        logger.info("Loading cached file info from %s", cache_file)
-        try:
-            with open(cache_file, "rb") as f:
-                file_names_sorted, lengths_sorted = pickle.load(f)
-            return file_names_sorted, lengths_sorted
-        except Exception as e:
-            logger.error("Error loading cached file info: %s", str(e))
-            logger.info("Falling back to scanning files")
-
-    # If no cache exists or loading failed, scan files
+    # Only rank 0 checks for cache and scans files if needed
    if get_world_rank() == 0:
-        lengths = []
-        file_names = []
-        for root, _, files in os.walk(path):
-            for file in sorted(files):
-                if file.endswith('.parquet'):
-                    file_path = os.path.join(root, file)
-                    file_names.append(file_path)
-        for file_path in tqdm.tqdm(file_names,
-                                   desc="Reading parquet files to get lengths"):
-            num_rows = pq.ParquetFile(file_path).metadata.num_rows
-            lengths.append(num_rows)
-        # sort according to file name to ensure all rank has the same order (in case os.walk is not sorted)
-        file_names_sorted, lengths_sorted = zip(
-            *sorted(zip(file_names, lengths), key=lambda x: x[0]))
-        assert len(
-            file_names_sorted) != 0, "No parquet files found in the dataset"
+        cache_loaded = False
+        file_names_sorted = None
+        lengths_sorted = None

-        os.makedirs(cache_dir, exist_ok=True)
-        with open(cache_file, "wb") as f:
-            pickle.dump((file_names_sorted, lengths_sorted), f)
-        logger.info("Saved file info to %s", cache_file)
+        # First try to load existing cache
+        if os.path.exists(cache_file):
+            logger.info("Loading cached file info from %s", cache_file)
+            try:
+                with open(cache_file, "rb") as f:
+                    file_names_sorted, lengths_sorted = pickle.load(f)
+                cache_loaded = True
+                logger.info("Successfully loaded cached file info")
+            except Exception as e:
+                logger.error("Error loading cached file info: %s", str(e))
+                logger.info("Falling back to scanning files")
+                cache_loaded = False

-    # Wait for rank 0 to finish saving
-    if get_world_size() > 1:
-        torch.distributed.barrier()
+        # If cache not loaded (either doesn't exist or failed to load), scan files
+        if not cache_loaded:
+            logger.info("Scanning parquet files to get lengths")
+            lengths = []
+            file_names = []
+            for root, _, files in os.walk(path):
+                for file in sorted(files):
+                    if file.endswith('.parquet'):
+                        file_path = os.path.join(root, file)
+                        file_names.append(file_path)
+            for file_path in tqdm.tqdm(
+                    file_names, desc="Reading parquet files to get lengths"):
+                num_rows = pq.ParquetFile(file_path).metadata.num_rows
+                lengths.append(num_rows)
+            # sort according to file name to ensure all rank has the same order
+            file_names_sorted, lengths_sorted = zip(*sorted(zip(file_names,
+                                                                lengths,
+                                                                strict=True),
+                                                            key=lambda x: x[0]),
+                                                    strict=True)
+            assert len(
+                file_names_sorted) != 0, "No parquet files found in the dataset"

-    return get_parquet_files_and_length(path)
+            # Save the cache
+            os.makedirs(cache_dir, exist_ok=True)
+            with open(cache_file, "wb") as f:
+                pickle.dump((file_names_sorted, lengths_sorted), f)
+            logger.info("Saved file info to %s", cache_file)
+
+    # Wait for rank 0 to finish creating/loading cache
+    world_group = get_world_group()
+    world_group.barrier()
+
+    # Now all ranks load the cache (it should exist and be valid now)
+    logger.info("Loading cached file info from %s after barrier", cache_file)
+    with open(cache_file, "rb") as f:
+        file_names_sorted, lengths_sorted = pickle.load(f)
+
+    return file_names_sorted, lengths_sorted


-def read_row_from_parquet_file(parquet_files: List[str], global_row_idx: int,
-                               lengths: List[int]) -> Dict[str, Any]:
+def read_row_from_parquet_file(parquet_files: list[str], global_row_idx: int,
+                               lengths: list[int]) -> dict[str, Any]:
    '''
    Read a row from a parquet file.
    Args:
@@ -150,17 +169,27 @@ def read_row_from_parquet_file(parquet_files: List[str], global_row_idx: int,
    '''
    # find the parquet file and local row index
    cumulative = 0
+    file_index = 0
+    local_row_idx = 0
+
    for file_index in range(len(lengths)):
        if cumulative + lengths[file_index] > global_row_idx:
            local_row_idx = global_row_idx - cumulative
            break
        cumulative += lengths[file_index]
+    else:
+        # If we reach here, global_row_idx is out of bounds
+        raise IndexError(
+            f"global_row_idx {global_row_idx} is out of bounds for dataset")

    parquet_file = pq.ParquetFile(parquet_files[file_index])

    # Calculate the row group to read into memory and the local idx
    # This way we can avoid reading in the entire parquet file
    cumulative = 0
+    row_group_index = 0
+    local_index = 0
+
    for i in range(parquet_file.num_row_groups):
        num_rows = parquet_file.metadata.row_group(i).num_rows
        if cumulative + num_rows > local_row_idx:
@@ -168,6 +197,11 @@ def read_row_from_parquet_file(parquet_files: List[str], global_row_idx: int,
            local_index = local_row_idx - cumulative
            break
        cumulative += num_rows
+    else:
+        # If we reach here, local_row_idx is out of bounds for this parquet file
+        raise IndexError(
+            f"local_row_idx {local_row_idx} is out of bounds for parquet file {parquet_files[file_index]}"
+        )

    row_group = parquet_file.read_row_group(row_group_index).to_pydict()
    row_dict = {k: v[local_index] for k, v in row_group.items()}
@@ -250,7 +284,7 @@ class LatentsParquetMapStyleDataset(Dataset):
        return negative_prompt_embedding, negative_prompt_attention_mask, negative_prompt

    # PyTorch calls this ONLY because the batch_sampler yields a list
-    def __getitems__(self, indices: List[int]) -> Dict[str, Any]:
+    def __getitems__(self, indices: list[int]) -> dict[str, Any]:
        """
        Batch fetch using read_row_from_parquet_file for each index.
        """
@@ -285,7 +319,7 @@ def build_parquet_map_style_dataloader(
        drop_last=True,
        drop_first_row=False,
        text_padding_length=512,
-        seed=42) -> Tuple[LatentsParquetMapStyleDataset, StatefulDataLoader]:
+        seed=42) -> tuple[LatentsParquetMapStyleDataset, StatefulDataLoader]:
    dataset = LatentsParquetMapStyleDataset(
        path,
        batch_size,
@@ -7,7 +7,7 @@ from abc import ABC, abstractmethod
 from collections import Counter
 from dataclasses import dataclass
 from os.path import join as opj
-from typing import Any, Dict, List, Optional, Union
+from typing import Any

 import numpy as np
 import torch
@@ -31,21 +31,21 @@ class PreprocessBatch:
    """
    # Raw metadata
    path: str
-    cap: Union[str, List[str]]
-    resolution: Optional[Dict] = None
-    fps: Optional[float] = None
-    duration: Optional[float] = None
+    cap: str | list[str]
+    resolution: dict | None = None
+    fps: float | None = None
+    duration: float | None = None

    # Processed metadata
-    num_frames: Optional[int] = None
-    sample_frame_index: Optional[List[int]] = None
-    sample_num_frames: Optional[int] = None
+    num_frames: int | None = None
+    sample_frame_index: list[int] | None = None
+    sample_num_frames: int | None = None

    # Processed data
-    pixel_values: Optional[torch.Tensor] = None
-    text: Optional[str] = None
-    input_ids: Optional[torch.Tensor] = None
-    cond_mask: Optional[torch.Tensor] = None
+    pixel_values: torch.Tensor | None = None
+    text: str | None = None
+    input_ids: torch.Tensor | None = None
+    cond_mask: torch.Tensor | None = None

    @property
    def is_video(self) -> bool:
@@ -281,10 +281,12 @@ class FrameSamplingStage(DatasetFilterStage):
                                  frame_interval).astype(int)

        # Temporal crop if too long
-        if len(frame_indices
-               ) > self.num_frames and temporal_sample_fn is not None:
-            begin_index, end_index = temporal_sample_fn(len(frame_indices))
-            frame_indices = frame_indices[begin_index:end_index]
+        if len(frame_indices) > self.num_frames:
+            if temporal_sample_fn is not None:
+                begin_index, end_index = temporal_sample_fn(len(frame_indices))
+                frame_indices = frame_indices[begin_index:end_index]
+            else:
+                frame_indices = frame_indices[:self.num_frames]

        batch.sample_frame_index = frame_indices.tolist()
        batch.sample_num_frames = len(frame_indices)
@@ -486,7 +488,7 @@ class VideoCaptionMergedDataset(torch.utils.data.IterableDataset,
            text_max_length=args.text_max_length,
            cfg_rate=args.training_cfg_rate)

-    def _load_raw_data(self) -> List[Dict]:
+    def _load_raw_data(self) -> list[dict]:
        """Load raw data from JSON files."""
        # Read folder-annotation pairs
        with open(self.data_merge_path) as f:
@@ -499,7 +501,7 @@ class VideoCaptionMergedDataset(torch.utils.data.IterableDataset,
                   ) == 2, "Folder-annotation pair should have two elements"
        folder, annotation_file = folder_anno_pairs[0]

-        data_items: List[Dict] = []
+        data_items: list[dict] = []
        with open(annotation_file) as f:
            data_items = json.load(f)

@@ -509,7 +511,7 @@ class VideoCaptionMergedDataset(torch.utils.data.IterableDataset,

        return data_items

-    def _process_metadata(self) -> List[PreprocessBatch]:
+    def _process_metadata(self) -> list[PreprocessBatch]:
        """Process the raw metadata through all filtering stages."""
        raw_data = self._load_raw_data()
        processed_batches = []
@@ -520,7 +522,7 @@ class VideoCaptionMergedDataset(torch.utils.data.IterableDataset,
            "resolution_failed": 0,
            "frame_sampling_failed": 0
        }
-        sample_num_frames: List[int] = []
+        sample_num_frames: list[int] = []

        for item in raw_data:
            batch = PreprocessBatch(path=item["path"],
@@ -546,7 +548,7 @@ class VideoCaptionMergedDataset(torch.utils.data.IterableDataset,
        return processed_batches

    def _apply_filter_stages(self, batch: PreprocessBatch,
-                             filter_counts: Dict[str, int]) -> bool:
+                             filter_counts: dict[str, int]) -> bool:
        """Apply all filter stages and update counters. Returns True if batch should be kept."""
        if not self.validation_stage.should_keep(batch):
            filter_counts["validation_failed"] += 1
@@ -562,8 +564,8 @@ class VideoCaptionMergedDataset(torch.utils.data.IterableDataset,

        return True

-    def _log_filtering_stats(self, filter_counts: Dict[str, int],
-                             sample_num_frames: List[int], before_count: int,
+    def _log_filtering_stats(self, filter_counts: dict[str, int],
+                             sample_num_frames: list[int], before_count: int,
                             after_count: int):
        """Log filtering statistics."""
        logger.info(
@@ -582,7 +584,7 @@ class VideoCaptionMergedDataset(torch.utils.data.IterableDataset,
    def __len__(self):
        return len(self.processed_batches)

-    def _get_item(self, idx: int) -> Dict:
+    def _get_item(self, idx: int) -> dict:
        """Get a single processed data item."""
        batch = self.processed_batches[idx]

@@ -606,10 +608,10 @@ class VideoCaptionMergedDataset(torch.utils.data.IterableDataset,

        return result

-    def state_dict(self) -> Dict[str, Any]:
+    def state_dict(self) -> dict[str, Any]:
        """Return state dict for checkpointing."""
        return {"processed_batches": self.processed_batches}

-    def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
+    def load_state_dict(self, state_dict: dict[str, Any]) -> None:
        """Load state dict from checkpoint."""
        self.processed_batches = state_dict["processed_batches"]
@@ -1,5 +1,5 @@
 import random
-from typing import Any, Dict, List, cast
+from typing import Any, cast

 import numpy as np
 import torch
@@ -21,7 +21,7 @@ def pad(t: torch.Tensor, padding_length: int) -> torch.Tensor:
        return t[:padding_length], torch.ones(padding_length)


-def get_torch_tensors_from_row_dict(row_dict, keys, cfg_rate) -> Dict[str, Any]:
+def get_torch_tensors_from_row_dict(row_dict, keys, cfg_rate) -> dict[str, Any]:
    """
    Get the latents and prompts from a row dictionary.
    """
@@ -61,7 +61,7 @@ def collate_latents_embs_masks(
        text_padding_length,
        keys,
        cfg_rate=0.0
-) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, List[str]]:
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[str]]:
    # Initialize tensors to hold padded embeddings and masks
    all_latents = []
    all_embs = []
@@ -95,7 +95,7 @@ def collate_latents_embs_masks(
 def collate_rows_from_parquet_schema(rows,
                                     parquet_schema,
                                     text_padding_length,
-                                     cfg_rate=0.0) -> Dict[str, Any]:
+                                     cfg_rate=0.0) -> dict[str, Any]:
    """
    Collate rows from parquet files based on the provided schema.
    Dynamically processes tensor fields based on schema and returns batched data.
@@ -108,10 +108,10 @@ def collate_rows_from_parquet_schema(rows,
        Dict containing batched tensors and metadata
    """
    if not rows:
-        return cast(Dict[str, Any], {})
+        return cast(dict[str, Any], {})

    # Initialize containers for different data types
-    batch_data: Dict[str, Any] = {}
+    batch_data: dict[str, Any] = {}

    # Get tensor and metadata field names from schema (fields ending with '_bytes')
    tensor_fields = []
@@ -4,15 +4,17 @@ import os
 import pathlib

 import datasets
-import torch
+from torch.utils.data import IterableDataset

+from fastvideo.v1.distributed import (get_sp_world_size, get_world_rank,
+                                      get_world_size)
 from fastvideo.v1.logger import init_logger
 from fastvideo.v1.models.vision_utils import load_image, load_video

 logger = init_logger(__name__)


-class ValidationDataset(torch.utils.data.IterableDataset):
+class ValidationDataset(IterableDataset):

    def __init__(self, filename: str):
        super().__init__()
@@ -48,10 +50,65 @@ class ValidationDataset(torch.utils.data.IterableDataset):
                f"Unsupported file format {self.filename.suffix} for validation dataset. Supported formats are: {_SUPPORTED_FILE_FORMATS}"
            )

-        self._data = data.to_iterable_dataset()
+        # Get distributed training info
+        self.global_rank = get_world_rank()
+        self.world_size = get_world_size()
+        self.sp_world_size = get_sp_world_size()
+        self.num_sp_groups = self.world_size // self.sp_world_size
+
+        # Convert to list to get total samples
+        self.all_samples = list(data)
+        self.original_total_samples = len(self.all_samples)
+
+        # Extend samples to be a multiple of DP degree (num_sp_groups)
+        remainder = self.original_total_samples % self.num_sp_groups
+        if remainder != 0:
+            samples_to_add = self.num_sp_groups - remainder
+
+            # Duplicate samples cyclically to reach the target
+            additional_samples = []
+            for i in range(samples_to_add):
+                additional_samples.append(
+                    self.all_samples[i % self.original_total_samples])
+
+            self.all_samples.extend(additional_samples)
+
+        self.total_samples = len(self.all_samples)
+
+        # Calculate which SP group this rank belongs to
+        self.sp_group_id = self.global_rank // self.sp_world_size
+
+        # Now all SP groups will have equal number of samples
+        self.samples_per_sp_group = self.total_samples // self.num_sp_groups
+
+        # Calculate start and end indices for this SP group
+        self.start_idx = self.sp_group_id * self.samples_per_sp_group
+        self.end_idx = self.start_idx + self.samples_per_sp_group
+
+        # Get samples for this SP group
+        self.sp_group_samples = self.all_samples[self.start_idx:self.end_idx]
+
+        logger.info(
+            "Rank %s (SP group %s): "
+            "Original samples: %s, "
+            "Extended samples: %s, "
+            "SP group samples: %s, "
+            "Range: [%s:%s]",
+            self.global_rank,
+            self.sp_group_id,
+            self.original_total_samples,
+            self.total_samples,
+            len(self.sp_group_samples),
+            self.start_idx,
+            self.end_idx,
+            local_main_process_only=False)
+
+    def __len__(self):
+        """Return the number of samples for this SP group."""
+        return len(self.sp_group_samples)

    def __iter__(self):
-        for sample in self._data:
+        for sample in self.sp_group_samples:
            # For consistency reasons, we mandate that "caption" is always present in the validation dataset.
            # However, since the model specifications use "prompt", we create an alias here.
            sample["prompt"] = sample["caption"]
@@ -64,6 +121,7 @@ class ValidationDataset(torch.utils.data.IterableDataset):
            if sample.get("image_path", None) is not None:
                image_path = sample["image_path"]
                image_path = os.path.join(self.dir, image_path)
+                sample["image_path"] = image_path
                if not pathlib.Path(image_path).is_file(
                ) and not image_path.startswith("http"):
                    logger.warning("Image file %s does not exist.", image_path)
@@ -73,6 +131,7 @@ class ValidationDataset(torch.utils.data.IterableDataset):
            if sample.get("video_path", None) is not None:
                video_path = sample["video_path"]
                video_path = os.path.join(self.dir, video_path)
+                sample["video_path"] = video_path
                if not pathlib.Path(video_path).is_file(
                ) and not video_path.startswith("http"):
                    logger.warning("Video file %s does not exist.", video_path)
@@ -82,6 +141,7 @@ class ValidationDataset(torch.utils.data.IterableDataset):
            if sample.get("control_image_path", None) is not None:
                control_image_path = sample["control_image_path"]
                control_image_path = os.path.join(self.dir, control_image_path)
+                sample["control_image_path"] = control_image_path
                if not pathlib.Path(control_image_path).is_file(
                ) and not control_image_path.startswith("http"):
                    logger.warning("Control Image file %s does not exist.",
@@ -92,6 +152,7 @@ class ValidationDataset(torch.utils.data.IterableDataset):
            if sample.get("control_video_path", None) is not None:
                control_video_path = sample["control_video_path"]
                control_video_path = os.path.join(self.dir, control_video_path)
+                sample["control_video_path"] = control_video_path
                if not pathlib.Path(control_video_path).is_file(
                ) and not control_video_path.startswith("http"):
                    logger.warning("Control Video file %s does not exist.",
@@ -3,10 +3,10 @@
 from fastvideo.v1.distributed.communication_op import *
 from fastvideo.v1.distributed.parallel_state import (
    cleanup_dist_env_and_memory, get_dp_group, get_dp_rank, get_dp_world_size,
-    get_sp_group, get_sp_parallel_rank, get_sp_world_size, get_torch_device,
-    get_tp_group, get_tp_rank, get_tp_world_size, get_world_group,
-    get_world_rank, get_world_size, init_distributed_environment,
-    initialize_model_parallel,
+    get_local_torch_device, get_sp_group, get_sp_parallel_rank,
+    get_sp_world_size, get_tp_group, get_tp_rank, get_tp_world_size,
+    get_world_group, get_world_rank, get_world_size,
+    init_distributed_environment, initialize_model_parallel,
    maybe_init_distributed_environment_and_model_parallel,
    model_parallel_is_initialized)
 from fastvideo.v1.distributed.utils import *
@@ -40,5 +40,5 @@ __all__ = [
    "get_tp_world_size",

    # Get torch device
-    "get_torch_device",
+    "get_local_torch_device",
 ]
@@ -1,7 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # Adapted from https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/distributed/device_communicators/base_device_communicator.py

-from typing import Any, Optional, Tuple
+from typing import Any

 import torch
 import torch.distributed as dist
@@ -28,7 +28,7 @@ class DistributedAutograd:
        def forward(ctx: Any,
                    group: ProcessGroup,
                    input_: Tensor,
-                    op: Optional[dist.ReduceOp] = None) -> Tensor:
+                    op: dist.ReduceOp | None = None) -> Tensor:
            ctx.group = group
            ctx.op = op
            output = input_.clone()
@@ -37,7 +37,7 @@ class DistributedAutograd:

        @staticmethod
        def backward(ctx: Any,
-                     grad_output: Tensor) -> Tuple[None, Tensor, None]:
+                     grad_output: Tensor) -> tuple[None, Tensor, None]:
            grad_output = grad_output.clone()
            dist.all_reduce(grad_output, group=ctx.group, op=ctx.op)
            return None, grad_output, None
@@ -75,7 +75,7 @@ class DistributedAutograd:

        @staticmethod
        def backward(ctx: Any,
-                     grad_output: Tensor) -> Tuple[None, Tensor, None, None]:
+                     grad_output: Tensor) -> tuple[None, Tensor, None, None]:
            # Split the gradient tensor along the gathered dimension
            dim_size = grad_output.size(ctx.dim) // ctx.world_size
            grad_chunks = grad_output.reshape(grad_output.shape[:ctx.dim] +
@@ -163,7 +163,7 @@ class DistributedAutograd:
        @staticmethod
        def backward(
                ctx: Any,
-                grad_output: Tensor) -> Tuple[None, Tensor, None, None, None]:
+                grad_output: Tensor) -> tuple[None, Tensor, None, None, None]:
            if ctx.world_size == 1:
                return None, grad_output, None, None, None

@@ -184,8 +184,8 @@ class DeviceCommunicatorBase:

    def __init__(self,
                 cpu_group: ProcessGroup,
-                 device: Optional[torch.device] = None,
-                 device_group: Optional[ProcessGroup] = None,
+                 device: torch.device | None = None,
+                 device_group: ProcessGroup | None = None,
                 unique_name: str = ""):
        self.device = device or torch.device("cpu")
        self.cpu_group = cpu_group
@@ -201,7 +201,7 @@ class DeviceCommunicatorBase:

    def all_reduce(self,
                   input_: torch.Tensor,
-                   op: Optional[dist.ReduceOp] = ReduceOp.SUM) -> torch.Tensor:
+                   op: dist.ReduceOp | None = ReduceOp.SUM) -> torch.Tensor:
        """Performs an all_reduce operation with gradient support."""
        return DistributedAutograd.AllReduce.apply(self.device_group, input_,
                                                   op)
@@ -225,7 +225,7 @@ class DeviceCommunicatorBase:
    def gather(self,
               input_: torch.Tensor,
               dst: int = 0,
-               dim: int = -1) -> Optional[torch.Tensor]:
+               dim: int = -1) -> torch.Tensor | None:
        """
        NOTE: We assume that the input tensor is on the same device across
        all the ranks.
@@ -254,7 +254,7 @@ class DeviceCommunicatorBase:
            output_tensor = None
        return output_tensor

-    def send(self, tensor: torch.Tensor, dst: Optional[int] = None) -> None:
+    def send(self, tensor: torch.Tensor, dst: int | None = None) -> None:
        """Sends a tensor to the destination rank in a non-blocking way"""
        """NOTE: `dst` is the local rank of the destination rank."""
        if dst is None:
@@ -264,7 +264,7 @@ class DeviceCommunicatorBase:
    def recv(self,
             size: torch.Size,
             dtype: torch.dtype,
-             src: Optional[int] = None) -> torch.Tensor:
+             src: int | None = None) -> torch.Tensor:
        """Receives a tensor from the source rank."""
        """NOTE: `src` is the local rank of the source rank."""
        if src is None:
@@ -1,8 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # Adapted from https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/distributed/device_communicators/cuda_communicator.py

-from typing import Optional
-
 import torch
 from torch.distributed import ProcessGroup

@@ -14,24 +12,22 @@ class CudaCommunicator(DeviceCommunicatorBase):

    def __init__(self,
                 cpu_group: ProcessGroup,
-                 device: Optional[torch.device] = None,
-                 device_group: Optional[ProcessGroup] = None,
+                 device: torch.device | None = None,
+                 device_group: ProcessGroup | None = None,
                 unique_name: str = ""):
        super().__init__(cpu_group, device, device_group, unique_name)

        from fastvideo.v1.distributed.device_communicators.pynccl import (
            PyNcclCommunicator)

-        self.pynccl_comm: Optional[PyNcclCommunicator] = None
+        self.pynccl_comm: PyNcclCommunicator | None = None
        if self.world_size > 1:
            self.pynccl_comm = PyNcclCommunicator(
                group=self.cpu_group,
                device=self.device,
            )

-    def all_reduce(self,
-                   input_,
-                   op: Optional[torch.distributed.ReduceOp] = None):
+    def all_reduce(self, input_, op: torch.distributed.ReduceOp | None = None):
        pynccl_comm = self.pynccl_comm
        assert pynccl_comm is not None
        out = pynccl_comm.all_reduce(input_, op=op)
@@ -44,7 +40,7 @@ class CudaCommunicator(DeviceCommunicatorBase):
            torch.distributed.all_reduce(out, group=self.device_group, op=op)
        return out

-    def send(self, tensor: torch.Tensor, dst: Optional[int] = None) -> None:
+    def send(self, tensor: torch.Tensor, dst: int | None = None) -> None:
        """Sends a tensor to the destination rank in a non-blocking way"""
        """NOTE: `dst` is the local rank of the destination rank."""
        if dst is None:
@@ -59,7 +55,7 @@ class CudaCommunicator(DeviceCommunicatorBase):
    def recv(self,
             size: torch.Size,
             dtype: torch.dtype,
-             src: Optional[int] = None) -> torch.Tensor:
+             src: int | None = None) -> torch.Tensor:
        """Receives a tensor from the source rank."""
        """NOTE: `src` is the local rank of the source rank."""
        if src is None:
@@ -1,8 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # Adapted from https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/distributed/device_communicators/pynccl.py

-from typing import Optional, Union
-
 # ===================== import region =====================
 import torch
 import torch.distributed as dist
@@ -22,9 +20,9 @@ class PyNcclCommunicator:

    def __init__(
        self,
-        group: Union[ProcessGroup, StatelessProcessGroup],
-        device: Union[int, str, torch.device],
-        library_path: Optional[str] = None,
+        group: ProcessGroup | StatelessProcessGroup,
+        device: int | str | torch.device,
+        library_path: str | None = None,
    ):
        """
        Args:
@@ -27,7 +27,7 @@
 import ctypes
 import platform
 from dataclasses import dataclass
-from typing import Any, Dict, List, Optional
+from typing import Any

 import torch
 from torch.distributed import ReduceOp
@@ -124,7 +124,7 @@ class ncclRedOpTypeEnum:
 class Function:
    name: str
    restype: Any
-    argtypes: List[Any]
+    argtypes: list[Any]


 class NCCLLibrary:
@@ -212,13 +212,13 @@ class NCCLLibrary:

    # class attribute to store the mapping from the path to the library
    # to avoid loading the same library multiple times
-    path_to_library_cache: Dict[str, Any] = {}
+    path_to_library_cache: dict[str, Any] = {}

    # class attribute to store the mapping from library path
    #  to the corresponding dictionary
-    path_to_dict_mapping: Dict[str, Dict[str, Any]] = {}
+    path_to_dict_mapping: dict[str, dict[str, Any]] = {}

-    def __init__(self, so_file: Optional[str] = None):
+    def __init__(self, so_file: str | None = None):

        so_file = so_file or find_nccl_library()

@@ -240,7 +240,7 @@ class NCCLLibrary:
            raise e

        if so_file not in NCCLLibrary.path_to_dict_mapping:
-            _funcs: Dict[str, Any] = {}
+            _funcs: dict[str, Any] = {}
            for func in NCCLLibrary.exported_functions:
                f = getattr(self.lib, func.name)
                f.restype = func.restype
@@ -23,15 +23,15 @@ If you only need to use the distributed environment without model parallelism,
 you can skip the model parallel initialization and destruction steps.
 """
 import contextlib
-import gc
 import os
 import pickle
 import weakref
 from collections import namedtuple
+from collections.abc import Callable
 from contextlib import contextmanager
 from dataclasses import dataclass
 from multiprocessing import shared_memory
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Optional
 from unittest.mock import patch

 import torch
@@ -58,15 +58,15 @@ TensorMetadata = namedtuple("TensorMetadata", ["device", "dtype", "size"])


 def _split_tensor_dict(
-    tensor_dict: Dict[str, Union[torch.Tensor, Any]]
-) -> Tuple[List[Tuple[str, Any]], List[torch.Tensor]]:
+    tensor_dict: dict[str, torch.Tensor | Any]
+) -> tuple[list[tuple[str, Any]], list[torch.Tensor]]:
    """Split the tensor dictionary into two parts:
    1. A list of (key, value) pairs. If the value is a tensor, it is replaced
         by its metadata.
    2. A list of tensors.
    """
-    metadata_list: List[Tuple[str, Any]] = []
-    tensor_list: List[torch.Tensor] = []
+    metadata_list: list[tuple[str, Any]] = []
+    tensor_list: list[torch.Tensor] = []
    for key, value in tensor_dict.items():
        if isinstance(value, torch.Tensor):
            # Note: we cannot use `value.device` here,
@@ -82,7 +82,7 @@ def _split_tensor_dict(
    return metadata_list, tensor_list


-_group_name_counter: Dict[str, int] = {}
+_group_name_counter: dict[str, int] = {}


 def _get_unique_name(name: str) -> str:
@@ -98,7 +98,7 @@ def _get_unique_name(name: str) -> str:
    return newname


-_groups: Dict[str, Callable[[], Optional["GroupCoordinator"]]] = {}
+_groups: dict[str, Callable[[], Optional["GroupCoordinator"]]] = {}


 def _register_group(group: "GroupCoordinator") -> None:
@@ -129,7 +129,7 @@ class GroupCoordinator:

    # available attributes:
    rank: int  # global rank
-    ranks: List[int]  # global ranks in the group
+    ranks: list[int]  # global ranks in the group
    world_size: int  # size of the group
    # difference between `local_rank` and `rank_in_group`:
    # if we have a group of size 4 across two nodes:
@@ -144,16 +144,16 @@ class GroupCoordinator:
    device_group: ProcessGroup  # group for device communication
    use_device_communicator: bool  # whether to use device communicator
    device_communicator: DeviceCommunicatorBase  # device communicator
-    mq_broadcaster: Optional[Any]  # shared memory broadcaster
+    mq_broadcaster: Any | None  # shared memory broadcaster

    def __init__(
        self,
-        group_ranks: List[List[int]],
+        group_ranks: list[list[int]],
        local_rank: int,
-        torch_distributed_backend: Union[str, Backend],
+        torch_distributed_backend: str | Backend,
        use_device_communicator: bool,
        use_message_queue_broadcaster: bool = False,
-        group_name: Optional[str] = None,
+        group_name: str | None = None,
    ):
        group_name = group_name or "anonymous"
        self.unique_name = _get_unique_name(group_name)
@@ -244,8 +244,8 @@ class GroupCoordinator:
        return self.ranks[(rank_in_group - 1) % world_size]

    @contextmanager
-    def graph_capture(
-            self, graph_capture_context: Optional[GraphCaptureContext] = None):
+    def graph_capture(self,
+                      graph_capture_context: GraphCaptureContext | None = None):
        if graph_capture_context is None:
            stream = torch.cuda.Stream()
            graph_capture_context = GraphCaptureContext(stream)
@@ -262,9 +262,9 @@ class GroupCoordinator:
            yield graph_capture_context

    def all_reduce(
-        self,
-        input_: torch.Tensor,
-        op: Optional[torch.distributed.ReduceOp] = ReduceOp.SUM
+            self,
+            input_: torch.Tensor,
+            op: torch.distributed.ReduceOp | None = ReduceOp.SUM
    ) -> torch.Tensor:
        """
        User-facing all-reduce function before we actually call the
@@ -291,9 +291,9 @@ class GroupCoordinator:
            return self._all_reduce_out_place(input_, op=op)

    def _all_reduce_out_place(
-        self,
-        input_: torch.Tensor,
-        op: Optional[torch.distributed.ReduceOp] = ReduceOp.SUM
+            self,
+            input_: torch.Tensor,
+            op: torch.distributed.ReduceOp | None = ReduceOp.SUM
    ) -> torch.Tensor:
        return self.device_communicator.all_reduce(input_, op=op)

@@ -310,7 +310,7 @@ class GroupCoordinator:
    def gather(self,
               input_: torch.Tensor,
               dst: int = 0,
-               dim: int = -1) -> Optional[torch.Tensor]:
+               dim: int = -1) -> torch.Tensor | None:
        """
        NOTE: We assume that the input tensor is on the same device across
        all the ranks.
@@ -346,7 +346,7 @@ class GroupCoordinator:
                                    group=self.device_group)
        return input_

-    def broadcast_object(self, obj: Optional[Any] = None, src: int = 0):
+    def broadcast_object(self, obj: Any | None = None, src: int = 0):
        """Broadcast the input object.
        NOTE: `src` is the local rank of the source rank.
        """
@@ -371,9 +371,9 @@ class GroupCoordinator:
            return recv[0]

    def broadcast_object_list(self,
-                              obj_list: List[Any],
+                              obj_list: list[Any],
                              src: int = 0,
-                              group: Optional[ProcessGroup] = None):
+                              group: ProcessGroup | None = None):
        """Broadcast the input object list.
        NOTE: `src` is the local rank of the source rank.
        """
@@ -453,11 +453,11 @@ class GroupCoordinator:

    def broadcast_tensor_dict(
        self,
-        tensor_dict: Optional[Dict[str, Union[torch.Tensor, Any]]] = None,
+        tensor_dict: dict[str, torch.Tensor | Any] | None = None,
        src: int = 0,
-        group: Optional[ProcessGroup] = None,
-        metadata_group: Optional[ProcessGroup] = None
-    ) -> Optional[Dict[str, Union[torch.Tensor, Any]]]:
+        group: ProcessGroup | None = None,
+        metadata_group: ProcessGroup | None = None
+    ) -> dict[str, torch.Tensor | Any] | None:
        """Broadcast the input tensor dictionary.
        NOTE: `src` is the local rank of the source rank.
        """
@@ -471,7 +471,7 @@ class GroupCoordinator:

        rank_in_group = self.rank_in_group
        if rank_in_group == src:
-            metadata_list: List[Tuple[Any, Any]] = []
+            metadata_list: list[tuple[Any, Any]] = []
            assert isinstance(
                tensor_dict,
                dict), (f"Expecting a dictionary, got {type(tensor_dict)}")
@@ -538,10 +538,10 @@ class GroupCoordinator:

    def send_tensor_dict(
        self,
-        tensor_dict: Dict[str, Union[torch.Tensor, Any]],
-        dst: Optional[int] = None,
+        tensor_dict: dict[str, torch.Tensor | Any],
+        dst: int | None = None,
        all_gather_group: Optional["GroupCoordinator"] = None,
-    ) -> Optional[Dict[str, Union[torch.Tensor, Any]]]:
+    ) -> dict[str, torch.Tensor | Any] | None:
        """Send the input tensor dictionary.
        NOTE: `dst` is the local rank of the source rank.
        """
@@ -561,7 +561,7 @@ class GroupCoordinator:
            dst = (self.rank_in_group + 1) % self.world_size
        assert dst < self.world_size, f"Invalid dst rank ({dst})"

-        metadata_list: List[Tuple[Any, Any]] = []
+        metadata_list: list[tuple[Any, Any]] = []
        assert isinstance(
            tensor_dict,
            dict), f"Expecting a dictionary, got {type(tensor_dict)}"
@@ -592,9 +592,9 @@ class GroupCoordinator:

    def recv_tensor_dict(
        self,
-        src: Optional[int] = None,
+        src: int | None = None,
        all_gather_group: Optional["GroupCoordinator"] = None,
-    ) -> Optional[Dict[str, Union[torch.Tensor, Any]]]:
+    ) -> dict[str, torch.Tensor | Any] | None:
        """Recv the input tensor dictionary.
        NOTE: `src` is the local rank of the source rank.
        """
@@ -615,7 +615,7 @@ class GroupCoordinator:
        assert src < self.world_size, f"Invalid src rank ({src})"

        recv_metadata_list = self.recv_object(src=src)
-        tensor_dict: Dict[str, Any] = {}
+        tensor_dict: dict[str, Any] = {}
        for key, value in recv_metadata_list:
            if isinstance(value, TensorMetadata):
                tensor = torch.empty(value.size,
@@ -665,7 +665,7 @@ class GroupCoordinator:
        """
        torch.distributed.barrier(group=self.cpu_group)

-    def send(self, tensor: torch.Tensor, dst: Optional[int] = None) -> None:
+    def send(self, tensor: torch.Tensor, dst: int | None = None) -> None:
        """Sends a tensor to the destination rank in a non-blocking way"""
        """NOTE: `dst` is the local rank of the destination rank."""
        self.device_communicator.send(tensor, dst)
@@ -673,7 +673,7 @@ class GroupCoordinator:
    def recv(self,
             size: torch.Size,
             dtype: torch.dtype,
-             src: Optional[int] = None) -> torch.Tensor:
+             src: int | None = None) -> torch.Tensor:
        """Receives a tensor from the source rank."""
        """NOTE: `src` is the local rank of the source rank."""
        return self.device_communicator.recv(size, dtype, src)
@@ -691,7 +691,7 @@ class GroupCoordinator:
            self.mq_broadcaster = None


-_WORLD: Optional[GroupCoordinator] = None
+_WORLD: GroupCoordinator | None = None


 def get_world_group() -> GroupCoordinator:
@@ -699,7 +699,7 @@ def get_world_group() -> GroupCoordinator:
    return _WORLD


-def init_world_group(ranks: List[int], local_rank: int,
+def init_world_group(ranks: list[int], local_rank: int,
                     backend: str) -> GroupCoordinator:
    return GroupCoordinator(
        group_ranks=[ranks],
@@ -711,11 +711,11 @@ def init_world_group(ranks: List[int], local_rank: int,


 def init_model_parallel_group(
-    group_ranks: List[List[int]],
+    group_ranks: list[list[int]],
    local_rank: int,
    backend: str,
    use_message_queue_broadcaster: bool = False,
-    group_name: Optional[str] = None,
+    group_name: str | None = None,
 ) -> GroupCoordinator:

    return GroupCoordinator(
@@ -728,7 +728,7 @@ def init_model_parallel_group(
    )


-_TP: Optional[GroupCoordinator] = None
+_TP: GroupCoordinator | None = None


 def get_tp_group() -> GroupCoordinator:
@@ -750,6 +750,7 @@ def init_distributed_environment(
    distributed_init_method: str = "env://",
    local_rank: int = 0,
    backend: str = "nccl",
+    device_id: torch.device | None = None,
 ):
    logger.debug(
        "world_size=%d rank=%d local_rank=%d "
@@ -764,7 +765,8 @@ def init_distributed_environment(
            backend=backend,
            init_method=distributed_init_method,
            world_size=world_size,
-            rank=rank)
+            rank=rank,
+            device_id=device_id)
    # set the local rank
    # local_rank is not available in torch ProcessGroup,
    # see https://github.com/pytorch/pytorch/issues/122816
@@ -784,7 +786,7 @@ def init_distributed_environment(
            "world group already initialized with a different world size")


-_SP: Optional[GroupCoordinator] = None
+_SP: GroupCoordinator | None = None


 def get_sp_group() -> GroupCoordinator:
@@ -792,7 +794,7 @@ def get_sp_group() -> GroupCoordinator:
    return _SP


-_DP: Optional[GroupCoordinator] = None
+_DP: GroupCoordinator | None = None


 def get_dp_group() -> GroupCoordinator:
@@ -803,7 +805,8 @@ def get_dp_group() -> GroupCoordinator:
 def initialize_model_parallel(
    tensor_model_parallel_size: int = 1,
    sequence_model_parallel_size: int = 1,
-    backend: Optional[str] = None,
+    data_parallel_size: int = 1,
+    backend: str | None = None,
 ) -> None:
    """
    Initialize model parallel groups.
@@ -904,7 +907,7 @@ def get_dp_rank() -> int:
    return get_dp_group().rank_in_group


-def get_torch_device() -> torch.device:
+def get_local_torch_device() -> torch.device:
    """Return the torch device for the current rank."""
    return torch.device(f"cuda:{envs.LOCAL_RANK}")

@@ -921,15 +924,17 @@ def maybe_init_distributed_environment_and_model_parallel(
    local_rank = int(os.environ.get("LOCAL_RANK", 0))
    world_size = int(os.environ.get("WORLD_SIZE", 1))
    rank = int(os.environ.get("RANK", 0))
+    device = torch.device(f"cuda:{local_rank}")

-    torch.cuda.set_device(local_rank)
    init_distributed_environment(
        world_size=world_size,
        rank=rank,
        local_rank=local_rank,
-        distributed_init_method=distributed_init_method)
+        distributed_init_method=distributed_init_method,
+        device_id=device)
    initialize_model_parallel(tensor_model_parallel_size=tp_size,
                              sequence_model_parallel_size=sp_size)
+    torch.cuda.set_device(device)


 def model_parallel_is_initialized() -> bool:
@@ -1010,19 +1015,10 @@ def cleanup_dist_env_and_memory(shutdown_ray: bool = False):
    if shutdown_ray:
        import ray  # Lazy import Ray
        ray.shutdown()
-    gc.collect()
-    from fastvideo.v1.platforms import current_platform
-    if not current_platform.is_cpu():
-        torch.cuda.empty_cache()
-    try:
-        torch._C._host_emptyCache()
-    except AttributeError:
-        logger.warning(
-            "torch._C._host_emptyCache() only available in Pytorch >=2.5")


-def in_the_same_node_as(pg: Union[ProcessGroup, StatelessProcessGroup],
-                        source_rank: int = 0) -> List[bool]:
+def in_the_same_node_as(pg: ProcessGroup | StatelessProcessGroup,
+                        source_rank: int = 0) -> list[bool]:
    """
    This is a collective operation that returns if each rank is in the same node
    as the source rank. It tests if processes are attached to the same
@@ -1108,7 +1104,7 @@ def in_the_same_node_as(pg: Union[ProcessGroup, StatelessProcessGroup],

 def initialize_tensor_parallel_group(
        tensor_model_parallel_size: int = 1,
-        backend: Optional[str] = None,
+        backend: str | None = None,
        group_name_suffix: str = "") -> GroupCoordinator:
    """Initialize a tensor parallel group for a specific model.
    
@@ -1172,7 +1168,7 @@ def initialize_tensor_parallel_group(

 def initialize_sequence_parallel_group(
        sequence_model_parallel_size: int = 1,
-        backend: Optional[str] = None,
+        backend: str | None = None,
        group_name_suffix: str = "") -> GroupCoordinator:
    """Initialize a sequence parallel group for a specific model.
    
@@ -9,7 +9,8 @@ import dataclasses
 import pickle
 import time
 from collections import deque
-from typing import Any, Deque, Dict, Optional, Sequence, Tuple
+from collections.abc import Sequence
+from typing import Any

 import torch
 from torch.distributed import TCPStore
@@ -72,15 +73,15 @@ class StatelessProcessGroup:
    data_expiration_seconds: int = 3600  # 1 hour

    # dst rank -> counter
-    send_dst_counter: Dict[int, int] = dataclasses.field(default_factory=dict)
+    send_dst_counter: dict[int, int] = dataclasses.field(default_factory=dict)
    # src rank -> counter
-    recv_src_counter: Dict[int, int] = dataclasses.field(default_factory=dict)
+    recv_src_counter: dict[int, int] = dataclasses.field(default_factory=dict)
    broadcast_send_counter: int = 0
-    broadcast_recv_src_counter: Dict[int, int] = dataclasses.field(
+    broadcast_recv_src_counter: dict[int, int] = dataclasses.field(
        default_factory=dict)

    # A deque to store the data entries, with key and timestamp.
-    entries: Deque[Tuple[str, float]] = dataclasses.field(default_factory=deque)
+    entries: deque[tuple[str, float]] = dataclasses.field(default_factory=deque)

    def __post_init__(self):
        assert self.rank < self.world_size
@@ -114,7 +115,7 @@ class StatelessProcessGroup:
        self.recv_src_counter[src] += 1
        return obj

-    def broadcast_obj(self, obj: Optional[Any], src: int) -> Any:
+    def broadcast_obj(self, obj: Any | None, src: int) -> Any:
        """Broadcast an object from a source rank to all other ranks.
        It does not clean up after all ranks have received the object.
        Use it for limited times, e.g., for initialization.
@@ -4,7 +4,7 @@
 import argparse
 import dataclasses
 import os
-from typing import List, cast
+from typing import cast

 from fastvideo import VideoGenerator
 from fastvideo.v1.configs.sample.base import SamplingParam
@@ -26,11 +26,11 @@ class GenerateSubcommand(CLISubcommand):
        self.init_arg_names = self._get_init_arg_names()
        self.generation_arg_names = self._get_generation_arg_names()

-    def _get_init_arg_names(self) -> List[str]:
+    def _get_init_arg_names(self) -> list[str]:
        """Get names of arguments for VideoGenerator initialization"""
        return ["num_gpus", "tp_size", "sp_size", "model_path"]

-    def _get_generation_arg_names(self) -> List[str]:
+    def _get_generation_arg_names(self) -> list[str]:
        """Get names of arguments for generate_video method"""
        return [field.name for field in dataclasses.fields(SamplingParam)]

@@ -120,5 +120,5 @@ class GenerateSubcommand(CLISubcommand):
        return cast(FlexibleArgumentParser, generate_parser)


-def cmd_init() -> List[CLISubcommand]:
+def cmd_init() -> list[CLISubcommand]:
    return [GenerateSubcommand()]
@@ -1,14 +1,12 @@
 # SPDX-License-Identifier: Apache-2.0
 # adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/entrypoints/cli/main.py

-from typing import List
-
 from fastvideo.v1.entrypoints.cli.cli_types import CLISubcommand
 from fastvideo.v1.entrypoints.cli.generate import cmd_init as generate_cmd_init
 from fastvideo.v1.utils import FlexibleArgumentParser


-def cmd_init() -> List[CLISubcommand]:
+def cmd_init() -> list[CLISubcommand]:
    """Initialize all commands from separate modules"""
    commands = []
    commands.extend(generate_cmd_init())
@@ -4,7 +4,6 @@ import argparse
 import os
 import subprocess
 import sys
-from typing import List, Optional

 from fastvideo.v1.logger import init_logger

@@ -19,8 +18,8 @@ class RaiseNotImplementedAction(argparse.Action):


 def launch_distributed(num_gpus: int,
-                       args: List[str],
-                       master_port: Optional[int] = None) -> int:
+                       args: list[str],
+                       master_port: int | None = None) -> int:
    """
    Launch a distributed job with the given arguments
    
@@ -6,11 +6,10 @@ This module provides a consolidated interface for generating videos using
 diffusion models.
 """

-import gc
 import math
 import os
 import time
-from typing import Any, Dict, List, Optional, Union
+from typing import Any

 import imageio
 import numpy as np
@@ -51,8 +50,8 @@ class VideoGenerator:
    @classmethod
    def from_pretrained(cls,
                        model_path: str,
-                        device: Optional[str] = None,
-                        torch_dtype: Optional[torch.dtype] = None,
+                        device: str | None = None,
+                        torch_dtype: torch.dtype | None = None,
                        **kwargs) -> "VideoGenerator":
        """
        Create a video generator from a pretrained model.
@@ -100,9 +99,9 @@ class VideoGenerator:
    def generate_video(
        self,
        prompt: str,
-        sampling_param: Optional[SamplingParam] = None,
+        sampling_param: SamplingParam | None = None,
        **kwargs,
-    ) -> Union[Dict[str, Any], List[np.ndarray]]:
+    ) -> dict[str, Any] | list[np.ndarray]:
        """
        Generate a video based on the given prompt.
        
@@ -277,5 +276,3 @@ class VideoGenerator:
        """
        self.executor.shutdown()
        del self.executor
-        gc.collect()
-        torch.cuda.empty_cache()
@@ -2,28 +2,29 @@
 # Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/envs.py

 import os
-from typing import TYPE_CHECKING, Any, Callable, Dict, Optional
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Any

 if TYPE_CHECKING:
    FASTVIDEO_RINGBUFFER_WARNING_INTERVAL: int = 60
-    FASTVIDEO_NCCL_SO_PATH: Optional[str] = None
-    LD_LIBRARY_PATH: Optional[str] = None
+    FASTVIDEO_NCCL_SO_PATH: str | None = None
+    LD_LIBRARY_PATH: str | None = None
    LOCAL_RANK: int = 0
-    CUDA_VISIBLE_DEVICES: Optional[str] = None
+    CUDA_VISIBLE_DEVICES: str | None = None
    FASTVIDEO_CACHE_ROOT: str = os.path.expanduser("~/.cache/fastvideo")
    FASTVIDEO_CONFIG_ROOT: str = os.path.expanduser("~/.config/fastvideo")
    FASTVIDEO_CONFIGURE_LOGGING: int = 1
    FASTVIDEO_LOGGING_LEVEL: str = "INFO"
    FASTVIDEO_LOGGING_PREFIX: str = ""
-    FASTVIDEO_LOGGING_CONFIG_PATH: Optional[str] = None
+    FASTVIDEO_LOGGING_CONFIG_PATH: str | None = None
    FASTVIDEO_TRACE_FUNCTION: int = 0
-    FASTVIDEO_ATTENTION_BACKEND: Optional[str] = None
-    FASTVIDEO_ATTENTION_CONFIG: Optional[str] = None
+    FASTVIDEO_ATTENTION_BACKEND: str | None = None
+    FASTVIDEO_ATTENTION_CONFIG: str | None = None
    FASTVIDEO_WORKER_MULTIPROC_METHOD: str = "fork"
    FASTVIDEO_TARGET_DEVICE: str = "cuda"
-    MAX_JOBS: Optional[str] = None
-    NVCC_THREADS: Optional[str] = None
-    CMAKE_BUILD_TYPE: Optional[str] = None
+    MAX_JOBS: str | None = None
+    NVCC_THREADS: str | None = None
+    CMAKE_BUILD_TYPE: str | None = None
    VERBOSE: bool = False
    FASTVIDEO_SERVER_DEV_MODE: bool = False

@@ -42,7 +43,7 @@ def get_default_config_root() -> str:
    )


-def maybe_convert_int(value: Optional[str]) -> Optional[int]:
+def maybe_convert_int(value: str | None) -> int | None:
    if value is None:
        return None
    return int(value)
@@ -53,7 +54,7 @@ def maybe_convert_int(value: Optional[str]) -> Optional[int]:

 # begin-env-vars-definition

-environment_variables: Dict[str, Callable[[], Any]] = {
+environment_variables: dict[str, Callable[[], Any]] = {

    # ================== Installation Time Env Vars ==================

@@ -6,7 +6,7 @@ import argparse
 import dataclasses
 from contextlib import contextmanager
 from dataclasses import field
-from typing import Any, Dict, List, Optional
+from typing import Any, List

 from fastvideo.v1.configs.pipelines.base import PipelineConfig, STA_Mode
 from fastvideo.v1.logger import init_logger
@@ -15,7 +15,7 @@ from fastvideo.v1.utils import FlexibleArgumentParser, StoreBoolean
 logger = init_logger(__name__)


-def clean_cli_args(args: argparse.Namespace) -> Dict[str, Any]:
+def clean_cli_args(args: argparse.Namespace) -> dict[str, Any]:
    """
    Clean the arguments by removing the ones that not explicitly provided by the user.
    """
@@ -44,7 +44,7 @@ class FastVideoArgs:

    # HuggingFace specific parameters
    trust_remote_code: bool = False
-    revision: Optional[str] = None
+    revision: str | None = None

    # Parallelism
    num_gpus: int = 1
@@ -52,17 +52,19 @@ class FastVideoArgs:
    sp_size: int = -1
    hsdp_replicate_dim: int = 1
    hsdp_shard_dim: int = -1
-    dist_timeout: Optional[int] = None  # timeout for torch.distributed
+    dist_timeout: int | None = None  # timeout for torch.distributed

    pipeline_config: PipelineConfig = field(default_factory=PipelineConfig)

    output_type: str = "pil"

-    use_cpu_offload: bool = True
+    use_cpu_offload: bool = True  # For DiT
    use_fsdp_inference: bool = True
+    text_encoder_offload: bool = True
+    pin_cpu_memory: bool = True

    # STA (Sliding Tile Attention) parameters
-    mask_strategy_file_path: Optional[str] = None
+    mask_strategy_file_path: str | None = None
    STA_mode: STA_Mode = STA_Mode.STA_INFERENCE
    skip_time_steps: int = 15

@@ -76,6 +78,8 @@ class FastVideoArgs:

    # Stage verification
    enable_stage_verification: bool = True
+    
+    denoising_step_list: List[int] | None = field(default=None)

    @property
    def training_mode(self) -> bool:
@@ -208,7 +212,7 @@ class FastVideoArgs:
            "--use-cpu-offload",
            action=StoreBoolean,
            help=
-            "Use CPU offload for model inference. Enable if run out of memory with FSDP.",
+            "Use CPU offload for DiT inference. Enable if run out of memory with FSDP.",
        )
        parser.add_argument(
            "--use-fsdp-inference",
@@ -216,7 +220,19 @@ class FastVideoArgs:
            help=
            "Use FSDP for inference by sharding the model weights. Latency is very low due to prefetch--enable if run out of memory.",
        )
-
+        parser.add_argument(
+            "--text-encoder-cpu-offload",
+            action=StoreBoolean,
+            help=
+            "Use CPU offload for text encoder. Enable if run out of memory.",
+        )
+        parser.add_argument(
+            "--pin-cpu-memory",
+            action=StoreBoolean,
+            help=
+            "Pin memory for CPU offload. Only added as a temp workaround if it throws \"CUDA error: invalid argument\". "
+            "Should be enabled in almost all cases",
+        )
        parser.add_argument(
            "--disable-autocast",
            action=StoreBoolean,
@@ -240,6 +256,12 @@ class FastVideoArgs:
            help="Enable input/output verification for pipeline stages",
        )

+        parser.add_argument("--denoising-step-list",
+            type=parse_int_list,
+            default=FastVideoArgs.denoising_step_list,
+            help="Comma-separated list of denoising steps (e.g., '1000,757,522')",
+        )
+        
        # Add pipeline configuration arguments
        PipelineConfig.add_cli_args(parser)

@@ -266,7 +288,7 @@ class FastVideoArgs:
        return cls(**kwargs)  # type: ignore

    @classmethod
-    def from_kwargs(cls, kwargs: Dict[str, Any]) -> "FastVideoArgs":
+    def from_kwargs(cls, kwargs: dict[str, Any]) -> "FastVideoArgs":
        kwargs['pipeline_config'] = PipelineConfig.from_kwargs(kwargs)
        return cls(**kwargs)

@@ -278,7 +300,7 @@ class FastVideoArgs:
            assert self.sp_size != -1, "sp_size must be set for training"

        if self.tp_size == -1:
-            self.tp_size = self.num_gpus
+            self.tp_size = 1
        if self.sp_size == -1:
            self.sp_size = self.num_gpus
        if self.hsdp_shard_dim == -1:
@@ -291,11 +313,6 @@ class FastVideoArgs:
        if self.num_gpus < max(self.tp_size, self.sp_size):
            self.num_gpus = max(self.tp_size, self.sp_size)

-        if self.tp_size != self.sp_size:
-            raise ValueError(
-                f"tp_size ({self.tp_size}) must be equal to sp_size ({self.sp_size})"
-            )
-
        if self.enable_torch_compile and self.num_gpus > 1:
            logger.warning(
                "Currently torch compile does not work with multi-gpu. Setting enable_torch_compile to False"
@@ -311,7 +328,7 @@ class FastVideoArgs:
 _current_fastvideo_args = None


-def prepare_fastvideo_args(argv: List[str]) -> FastVideoArgs:
+def prepare_fastvideo_args(argv: list[str]) -> FastVideoArgs:
    """
    Prepare the inference arguments from the command line arguments.

@@ -396,7 +413,7 @@ class TrainingArgs(FastVideoArgs):
    log_validation: bool = False
    tracker_project_name: str = ""
    wandb_run_name: str = ""
-    seed: Optional[int] = None
+    seed: int | None = None

    # output
    output_dir: str = ""
@@ -411,9 +428,10 @@ class TrainingArgs(FastVideoArgs):
    learning_rate: float = 0.0
    scale_lr: bool = False
    lr_scheduler: str = "constant"
+    lr_step_rules: str | None = None
    lr_warmup_steps: int = 0
    max_grad_norm: float = 0.0
-    enable_gradient_checkpointing_type: Optional[str] = None
+    enable_gradient_checkpointing_type: str | None = None
    selective_checkpointing: float = 0.0
    allow_tf32: bool = False
    mixed_precision: str = ""
@@ -446,6 +464,17 @@ class TrainingArgs(FastVideoArgs):
    # VSA training decay parameters
    VSA_decay_rate: float = 0.01  # decay rate -> 0.02
    VSA_decay_interval_steps: int = 1  # decay interval steps -> 50
+    
+    # distillation args
+    student_critic_update_ratio: int = 5
+    critic_learning_rate: float = 1e-5
+    critic_lr_scheduler: str = "constant"
+    critic_lr_step_rules: str | None = None
+    min_step_ratio: float = 0.2
+    max_step_ratio: float = 0.98
+    teacher_guidance_scale: float = 3.5
+    simulate_student_forward: bool = False
+    num_teacher_noisy_ground_truth_steps: int = 0

    @classmethod
    def from_cli_args(cls, args: argparse.Namespace) -> "TrainingArgs":
@@ -605,6 +634,9 @@ class TrainingArgs(FastVideoArgs):
                            type=str,
                            default="constant",
                            help="Learning rate scheduler type")
+        parser.add_argument("--lr-step-rules",
+                            type=str,
+                            help="Learning rate step rules")
        parser.add_argument("--lr-warmup-steps",
                            type=int,
                            default=10,
@@ -712,5 +744,48 @@ class TrainingArgs(FastVideoArgs):
            type=int,
            default=TrainingArgs.VSA_decay_interval_steps,
            help="VSA decay interval steps")
+        
+        # Distillation arguments
+        parser.add_argument("--student-critic-update-ratio",
+            type=int,
+            default=TrainingArgs.student_critic_update_ratio,
+            help="Ratio of student updates to critic updates.")
+        parser.add_argument("--critic-learning-rate",
+            type=float,
+            default=TrainingArgs.critic_learning_rate,
+            help="Learning rate for critic")
+        parser.add_argument("--critic-lr-scheduler",
+            type=str,
+            default=TrainingArgs.critic_lr_scheduler,
+            help="Learning rate scheduler type for critic")
+        parser.add_argument("--critic-lr-step-rules",
+            type=str,
+            help="Learning rate step rules for critic")
+        parser.add_argument("--min-step-ratio",
+            type=float,
+            default=TrainingArgs.min_step_ratio,
+            help="Minimum step ratio")
+        parser.add_argument("--max-step-ratio",
+            type=float,
+            default=TrainingArgs.max_step_ratio,
+            help="Maximum step ratio")
+        parser.add_argument("--teacher-guidance-scale",
+            type=float,
+            default=TrainingArgs.teacher_guidance_scale,
+            help="Teacher guidance scale")
+        parser.add_argument("--simulate-student-forward",
+            action=StoreBoolean,
+            default=TrainingArgs.simulate_student_forward,
+            help="Whether to simulate student forward")
+        parser.add_argument("--num-teacher-noisy-ground-truth-steps",
+            type=int,
+            default=TrainingArgs.num_teacher_noisy_ground_truth_steps,
+            help="Number of steps to use noisy ground truth for teacher")

        return parser
+
+def parse_int_list(value: str) -> List[int]:
+    """Parse a comma-separated string of integers into a list."""
+    if not value:
+        return []
+    return [int(x.strip()) for x in value.split(",")]
@@ -56,7 +56,7 @@ def get_forward_context() -> "ForwardContext":
 def set_forward_context(current_timestep,
                        attn_metadata,
                        forward_batch: Optional["ForwardBatch"] = None,
-                        fastvideo_args: Optional[FastVideoArgs] = None):
+                        fastvideo_args: FastVideoArgs | None = None):
    """A context manager that stores the current forward context,
    can be attention metadata, etc.
    Here we can inject common logic for every model forward pass.
@@ -1,7 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
 # Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/custom_op.py

-from typing import Any, Callable, Dict, Type
+from collections.abc import Callable
+from typing import Any

 import torch.nn as nn

@@ -81,7 +82,7 @@ class CustomOp(nn.Module):
    # Examples:
    # - MyOp.enabled()
    # - op_registry["my_op"].enabled()
-    op_registry: Dict[str, Type['CustomOp']] = {}
+    op_registry: dict[str, type['CustomOp']] = {}

    # Decorator to register custom ops.
    @classmethod
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/layers/layernorm.py
 """Custom normalization layers."""
-from typing import Optional, Tuple, Union

 import torch
 import torch.nn as nn
@@ -23,7 +22,7 @@ class RMSNorm(CustomOp):
        hidden_size: int,
        eps: float = 1e-6,
        dtype: torch.dtype = torch.float32,
-        var_hidden_size: Optional[int] = None,
+        var_hidden_size: int | None = None,
        has_weight: bool = True,
    ) -> None:
        super().__init__()
@@ -47,8 +46,8 @@ class RMSNorm(CustomOp):
    def forward_native(
        self,
        x: torch.Tensor,
-        residual: Optional[torch.Tensor] = None,
-    ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+        residual: torch.Tensor | None = None,
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
        """PyTorch-native implementation equivalent to forward()."""
        orig_dtype = x.dtype
        x = x.to(torch.float32)
@@ -159,7 +158,7 @@ class ScaleResidualLayerNormScaleShift(nn.Module):

    def forward(self, residual: torch.Tensor, x: torch.Tensor,
                gate: torch.Tensor, shift: torch.Tensor,
-                scale: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+                scale: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        """
        Apply gated residual connection, followed by layernorm and 
        scale/shift in a single fused operation.
--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
SolitaryThinker	4058e92fb4	separate critic lr	2025-07-21 19:26:04 +00:00
SolitaryThinker	de6f92b896	update	2025-07-18 23:20:30 +00:00
SolitaryThinker	aaeee0b3ac	add noisy ground truth for teacher	2025-07-18 22:28:15 +00:00
SolitaryThinker	ee2c06890b	multi-step simulation	2025-07-18 22:21:41 +00:00
SolitaryThinker	7186408be9	add sample param for wan2.1-fun inp	2025-07-18 21:12:56 +00:00
SolitaryThinker	878524bc28	add more samples to crush-smol val.json	2025-07-18 21:11:25 +00:00
SolitaryThinker	44da53f2a9	start from step 1	2025-07-18 20:58:05 +00:00
SolitaryThinker	a59b7b04d9	remove breakpoint	2025-07-18 20:31:09 +00:00
“BrianChen1129”	c8f4a6378f	update	2025-07-15 20:07:58 +00:00
“BrianChen1129”	a486a15bf4	update	2025-07-15 05:24:37 +00:00
“BrianChen1129”	207ce0e9b5	fix wandb step issue	2025-07-15 00:28:48 +00:00
“BrianChen1129”	bc818a0928	update i2v t2v distill scripts	2025-07-14 23:11:23 +00:00
“BrianChen1129”	8295fabda1	update	2025-07-14 23:09:47 +00:00
“BrianChen1129”	c81eabd50f	update new t2v dmd	2025-07-14 09:53:59 +00:00
“BrianChen1129”	81b404b17d	update	2025-07-13 22:03:51 +00:00
“BrianChen1129”	693a0f361d	update	2025-07-12 04:13:02 +00:00
“BrianChen1129”	c963128781	update	2025-07-12 02:30:29 +00:00
“BrianChen1129”	d7143bb00d	Remove local_scripts from git tracking and add to .gitignore	2025-07-11 20:48:28 +00:00
“BrianChen1129”	a652992506	update	2025-07-11 20:44:38 +00:00
“BrianChen1129”	8955d15f4e	bkp	2025-07-11 09:49:09 +00:00
“BrianChen1129”	411302db13	update	2025-07-10 08:06:37 +00:00
“BrianChen1129”	a3e53f0e97	update latents shape	2025-07-10 06:37:22 +00:00
“BrianChen1129”	7420c7d163	update	2025-07-10 02:43:45 +00:00
“BrianChen1129”	ca958fa4cc	fix seed diff	2025-07-09 22:45:58 +00:00
Wenxuan Tan	65ed588570	Set encoder TP size to 1 by default (#569 )	2025-07-09 17:44:24 -05:00
Wenxuan Tan	14adfe2edc	Remove all unnecessary torch.cuda.empty_cache (#606 )	2025-07-09 16:45:20 -05:00
“BrianChen1129”	1c99f157ec	add readme	2025-07-09 09:58:42 +00:00
“BrianChen1129”	3e3e0f6ccb	update	2025-07-09 09:34:09 +00:00
“BrianChen1129”	493d8c659a	training	2025-07-09 06:51:28 +00:00
“BrianChen1129”	b1f7acc4b5	update	2025-07-08 20:47:05 +00:00
“BrianChen1129”	3d6a69b154	update	2025-07-08 04:12:16 +00:00
William Lin	6198c6a640	[docs] update dev guide runpod image to py3.12 (#602 )	2025-07-07 19:56:07 -05:00
“BrianChen1129”	a3a3eea735	update	2025-07-08 00:19:51 +00:00
“BrianChen1129”	fb9e18ed20	update random seed and inference latent shape	2025-07-08 00:14:54 +00:00
William Lin	e6b71b531b	[docs] Update slack invite (#601 )	2025-07-07 14:20:42 -05:00
“BrianChen1129”	2bf1f5510b	update inference	2025-07-07 06:09:04 +00:00
“BrianChen1129”	5a9f2d2f05	update	2025-07-07 03:07:40 +00:00
William Lin	ae1d112c6a	[bugfix] [training] fix deadlock in latent datasets and init error in multi-node training (#598 )	2025-07-06 01:34:15 -05:00
“BrianChen1129”	b5bdf51697	update	2025-07-06 05:24:22 +00:00
“BrianChen1129”	5bad51acf2	update	2025-07-06 05:24:08 +00:00
“BrianChen1129”	6cdd5dca30	add wrapper	2025-07-05 09:59:55 +00:00
“BrianChen1129”	0c0fb63c56	update	2025-07-05 09:10:10 +00:00
“BrianChen1129”	91ef7567ac	update training	2025-07-05 04:13:22 +00:00
William Lin	bf4de1f38f	[chore] Upgrade min Python version from 3.8 to 3.10 (#597 )	2025-07-04 22:08:31 -05:00
William Lin	66fdcc8e76	[Training] Use inference pipeline for training validation (#585 )	2025-07-04 17:23:55 -05:00
“BrianChen1129”	09f2c64bbe	training val issue	2025-07-04 07:17:19 +00:00
Wenxuan Tan	ed1e8d6bad	[Feature] Offload all text encoders by default (#594 )	2025-07-03 19:30:14 -05:00
“BrianChen1129”	fc5e823875	update	2025-07-03 22:35:34 +00:00
“BrianChen1129”	0c910981e5	update	2025-07-03 21:37:45 +00:00
Kevin Lin	b9423ca3f8	Add ComfyUI custom node for inference (#596 )	2025-07-03 16:09:03 -05:00
“BrianChen1129”	5d809f19d1	update	2025-07-03 20:37:51 +00:00
“BrianChen1129”	8da919967f	update“	2025-07-03 04:16:08 +00:00
“BrianChen1129”	21ffa0d96a	update	2025-07-03 04:08:36 +00:00
“BrianChen1129”	d6a1e875f0	merge main	2025-07-03 03:43:32 +00:00
“BrianChen1129”	c025ea80e5	update	2025-07-03 03:40:28 +00:00
“BrianChen1129”	98653a2261	update validation	2025-07-03 02:28:29 +00:00
“BrianChen1129”	9589938583	update	2025-07-02 22:56:37 +00:00
Wenxuan Tan	ad16289871	[LoRA] Fix lora merge weights (#579 )	2025-07-02 00:12:52 -05:00
Wenxuan Tan	2a41da1e6b	Fix VAE precisions (#588 )	2025-07-01 14:49:04 -05:00
William Lin	32133171da	[chore] Release 0.1.1 (#592 )	2025-07-01 01:43:12 -05:00
Kevin Lin	19674c6f29	[CI] Fix fork builds (#590 )	2025-07-01 01:03:34 -05:00
“BrianChen1129”	13fa3d83c7	Merge branch 'main' of github.com:hao-ai-lab/FastVideo into yq/distill-dmd	2025-06-30 03:30:39 +00:00
“BrianChen1129”	db154cc313	update	2025-06-30 03:30:26 +00:00
“BrianChen1129”	8fb2708ff0	update	2025-06-29 23:46:31 +00:00
“BrianChen1129”	36dab19ff9	update	2025-06-29 08:45:56 +00:00
“BrianChen1129”	3c1179b4c2	update	2025-06-29 06:44:06 +00:00
“BrianChen1129”	8b60ac2964	Merge branch 'main' of github.com:hao-ai-lab/FastVideo into yq/distill-dmd	2025-06-29 01:57:31 +00:00
“BrianChen1129”	0b3de5e5eb	update	2025-06-29 00:38:52 +00:00
“BrianChen1129”	c627dc6421	Merge branch 'main' of github.com:hao-ai-lab/FastVideo into yq/distill-dmd	2025-06-29 00:03:26 +00:00
“BrianChen1129”	117d23fcc0	Merge branch 'main' of github.com:hao-ai-lab/FastVideo into yq/distill-dmd	2025-06-28 23:49:01 +00:00
“BrianChen1129”	4932793d00	update	2025-06-28 04:20:47 +00:00
“BrianChen1129”	152e6a6c77	update	2025-06-27 21:25:22 +00:00
“BrianChen1129”	3fd5531f4f	debugging	2025-06-27 04:30:35 +00:00
“BrianChen1129”	26ab6c84c3	update	2025-06-26 20:35:49 +00:00
“BrianChen1129”	65c0733a48	update	2025-06-25 23:36:05 +00:00
“BrianChen1129”	4f40eef184	e2e ready; fixing bugs	2025-06-25 07:07:56 +00:00
“BrianChen1129”	1d5bf58b34	debuging loss	2025-06-25 00:56:17 +00:00
“BrianChen1129”	4ad1281377	update	2025-06-24 23:13:05 +00:00
“BrianChen1129”	ee5da0838e	update	2025-06-24 03:29:09 +00:00
“BrianChen1129”	4006b3a5ee	update framework	2025-06-24 00:33:32 +00:00
“BrianChen1129”	6c731260f0	update	2025-06-23 22:42:23 +00:00
“BrianChen1129”	1de2eb2afd	update	2025-06-23 21:07:36 +00:00
“BrianChen1129”	4b9782e3f4	update	2025-06-23 07:10:57 +00:00