update pre commit

add DMD training pipeline
2025-07-27 05:56:13 +00:00 · 2025-07-27 05:54:53 +00:00 · 2025-07-27 05:08:10 +00:00 · 2025-07-26 19:46:26 -07:00 · 2025-07-25 19:02:52 -07:00 · 2025-07-25 18:39:05 -07:00
291 changed files with 3076 additions and 1189 deletions
@@ -16,9 +16,9 @@ steps:
          diff: 'git fetch origin "$BUILDKITE_PULL_REQUEST_BASE_BRANCH" && git diff --name-only origin/"$BUILDKITE_PULL_REQUEST_BASE_BRANCH"...HEAD'
          watch:
            - path: 
-                - "fastvideo/v1/models/encoders/**"
-                - "fastvideo/v1/models/loader/**"
-                - "fastvideo/v1/tests/encoders/**"
+                - "fastvideo/models/encoders/**"
+                - "fastvideo/models/loader/**"
+                - "fastvideo/tests/encoders/**"
                - "pyproject.toml"
                - "docker/Dockerfile.python3.12"
              config:
@@ -29,9 +29,9 @@ steps:
                agents:
                  queue: "default"
            - path:
-                - "fastvideo/v1/models/vaes/**"
-                - "fastvideo/v1/models/loader/**"
-                - "fastvideo/v1/tests/vaes/**"
+                - "fastvideo/models/vaes/**"
+                - "fastvideo/models/loader/**"
+                - "fastvideo/tests/vaes/**"
                - "pyproject.toml"
                - "docker/Dockerfile.python3.12"
              config:
@@ -42,11 +42,11 @@ steps:
                agents:
                  queue: "default"
            - path:
-                - "fastvideo/v1/models/dits/**"
-                - "fastvideo/v1/models/loader/**"
-                - "fastvideo/v1/tests/transformers/**"
-                - "fastvideo/v1/layers/**"
-                - "fastvideo/v1/attention/**"
+                - "fastvideo/models/dits/**"
+                - "fastvideo/models/loader/**"
+                - "fastvideo/tests/transformers/**"
+                - "fastvideo/layers/**"
+                - "fastvideo/attention/**"
                - "pyproject.toml"
                - "docker/Dockerfile.python3.12"
              config:
@@ -57,7 +57,7 @@ steps:
                agents:
                  queue: "default"
            - path: 
-                - "fastvideo/v1/**/*.py"
+                - "fastvideo/**/*.py"
                - "pyproject.toml"
                - "docker/Dockerfile.python3.12"
              config:
@@ -68,11 +68,11 @@ steps:
                agents:
                  queue: "default"
            - path:
-                - "fastvideo/v1/tests/lora/**"
-                - "fastvideo/v1/models/loader/**"
-                - "fastvideo/v1/tests/transformers/**"
-                - "fastvideo/v1/pipelines/**"
-                - "fastvideo/v1/layers/lora/**"
+                - "fastvideo/tests/lora/**"
+                - "fastvideo/models/loader/**"
+                - "fastvideo/tests/transformers/**"
+                - "fastvideo/pipelines/**"
+                - "fastvideo/layers/lora/**"
                - "pyproject.toml"
                - "docker/Dockerfile.python3.12"
              config:
@@ -83,7 +83,7 @@ steps:
                agents:
                  queue: "default"
            - path:
-                - "fastvideo/v1/**"
+                - "fastvideo/**"
                - "pyproject.toml"
                - "docker/Dockerfile.python3.12"
              config:
@@ -94,7 +94,7 @@ steps:
                agents:
                  queue: "default"
            - path:
-                - "fastvideo/v1/**"
+                - "fastvideo/**"
                - "pyproject.toml"
                - "docker/Dockerfile.python3.12"
              config:
@@ -105,7 +105,7 @@ steps:
                agents:
                  queue: "default"
            - path:
-                - "fastvideo/v1/**"
+                - "fastvideo/**"
                - "csrc/attn/vsa/**"
                - "csrc/attn/tk/**"
                - "csrc/attn/setup_vsa.py"
@@ -121,7 +121,7 @@ steps:
                agents:
                  queue: "default"
            - path:
-                - "fastvideo/v1/**"
+                - "fastvideo/**"
                - "csrc/attn/st_attn/**"
                - "csrc/attn/setup_sta.py"
                - "csrc/attn/config_sta.py"
@@ -50,7 +50,7 @@ else
    exit 1
 fi

-MODAL_TEST_FILE="fastvideo/v1/tests/modal/pr_test.py"
+MODAL_TEST_FILE="fastvideo/tests/modal/pr_test.py"

 if [ -z "${TEST_TYPE:-}" ]; then
    log "Error: TEST_TYPE environment variable is not set"
@@ -23,7 +23,7 @@ body:
  attributes:
    label: Environment
    description: |
-      Please share your environment with us. You can run the command **python fastvideo/utils/collect_env.py** and copy-paste its output below.
+      Please share your environment with us. You can run the command **python collect_env.py** and copy-paste its output below.
    placeholder: FastVideo version, platform, python version, cuda version...
  validations:
    required: true
@@ -8,14 +8,14 @@ on:
      - main
    paths:
      - "docs/**/*.md"
-      - "fastvideo/v1/examples/**/*.py"
+      - "fastvideo/examples/**/*.py"
  pull_request:
    branches: 
      - main
    types: [opened, ready_for_review, synchronize, reopened]
    paths:
      - "docs/**/*.md"
-      - "fastvideo/v1/examples/**/*.py"
+      - "fastvideo/examples/**/*.py"

  # Allows you to run this workflow manually from the Actions tab
  workflow_dispatch:
@@ -115,37 +115,37 @@ jobs:
              - 'csrc/attn/config_vsa.py'
              - 'csrc/attn/vsa.cpp'
            vsa-paths: &vsa-paths
-              - 'fastvideo/v1/**'
+              - 'fastvideo/**'
              - *common-paths
              - *vsa-kernel-paths
            
            # Actual tests
            encoder-test:
-              - 'fastvideo/v1/models/encoders/**'
-              - 'fastvideo/v1/models/loader/**'
-              - 'fastvideo/v1/tests/encoders/**'
+              - 'fastvideo/models/encoders/**'
+              - 'fastvideo/models/loader/**'
+              - 'fastvideo/tests/encoders/**'
              - *common-paths
            vae-test:
-              - 'fastvideo/v1/models/vaes/**'
-              - 'fastvideo/v1/models/loader/**'
-              - 'fastvideo/v1/tests/vaes/**'
+              - 'fastvideo/models/vaes/**'
+              - 'fastvideo/models/loader/**'
+              - 'fastvideo/tests/vaes/**'
              - *common-paths
            transformer-test:
-              - 'fastvideo/v1/models/dits/**'
-              - 'fastvideo/v1/models/loader/**'
-              - 'fastvideo/v1/tests/transformers/**'
-              - 'fastvideo/v1/layers/**'
-              - 'fastvideo/v1/attention/**'
+              - 'fastvideo/models/dits/**'
+              - 'fastvideo/models/loader/**'
+              - 'fastvideo/tests/transformers/**'
+              - 'fastvideo/layers/**'
+              - 'fastvideo/attention/**'
              - *common-paths
            training-test:
-              - 'fastvideo/v1/**'
+              - 'fastvideo/**'
              - *common-paths
            training-test-VSA:
-              - 'fastvideo/v1/**'
+              - 'fastvideo/**'
              - *common-paths
              - *vsa-kernel-paths
            inference-test-STA:
-              - 'fastvideo/v1/**'
+              - 'fastvideo/**'
              - *common-paths
              - *sta-kernel-paths
            precision-test-STA:
@@ -167,7 +167,7 @@ jobs:
      gpu_count: 1
      volume_size: 100
      image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest"
-      test_command: "uv pip install -e .[test] && pytest ./fastvideo/v1/tests/encoders -s"
+      test_command: "uv pip install -e .[test] && pytest ./fastvideo/tests/encoders -s"
      timeout_minutes: 30
    secrets:
      RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
@@ -185,7 +185,7 @@ jobs:
      gpu_count: 1
      volume_size: 100
      image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest"
-      test_command: "uv pip install -e .[test] && pytest ./fastvideo/v1/tests/vaes -s"
+      test_command: "uv pip install -e .[test] && pytest ./fastvideo/tests/vaes -s"
      timeout_minutes: 30
    secrets:
      RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
@@ -203,7 +203,7 @@ jobs:
      gpu_count: 1
      volume_size: 100
      image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest"
-      test_command: "uv pip install -e .[test] && pytest ./fastvideo/v1/tests/transformers -s"
+      test_command: "uv pip install -e .[test] && pytest ./fastvideo/tests/transformers -s"
      timeout_minutes: 30
    secrets:
      RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
@@ -229,7 +229,7 @@ jobs:
      volume_size: 200
      disk_size: 200
      image: "ghcr.io/${{ github.repository }}/fastvideo-dev:${{ matrix.python-version.tag }}"
-      test_command: "uv pip install -e .[test] && pytest ./fastvideo/v1/tests/ssim -vs"
+      test_command: "uv pip install -e .[test] && pytest ./fastvideo/tests/ssim -vs"
      timeout_minutes: 60
    secrets:
      RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
@@ -248,7 +248,7 @@ jobs:
      volume_size: 100
      disk_size: 100
      image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest"
-      test_command: "wandb login $WANDB_API_KEY && uv pip install -e .[test] && pytest ./fastvideo/v1/tests/training/Vanilla -srP"
+      test_command: "wandb login $WANDB_API_KEY && uv pip install -e .[test] && pytest ./fastvideo/tests/training/Vanilla -srP"
      timeout_minutes: 30
    secrets:
      RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
@@ -268,7 +268,7 @@ jobs:
      volume_size: 100
      disk_size: 100
      image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest"
-      test_command: "wandb login $WANDB_API_KEY && uv pip install -e .[test] && pytest ./fastvideo/v1/tests/training/VSA -srP"
+      test_command: "wandb login $WANDB_API_KEY && uv pip install -e .[test] && pytest ./fastvideo/tests/training/VSA -srP"
      timeout_minutes: 30
    secrets:
      RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
@@ -288,7 +288,7 @@ jobs:
      volume_size: 100
      disk_size: 100
      image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest"
-      test_command: "uv pip install -e .[test] && pytest ./fastvideo/v1/tests/inference/STA -srP"
+      test_command: "uv pip install -e .[test] && pytest ./fastvideo/tests/inference/STA -srP"
      timeout_minutes: 30
    secrets:
      RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
@@ -343,7 +343,7 @@ jobs:
      volume_size: 100
      disk_size: 100
      image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest"
-      test_command: "wandb login $WANDB_API_KEY && uv pip install -e .[test] && pytest ./fastvideo/v1/tests/nightly/test_e2e_overfit_single_sample.py -vs"
+      test_command: "wandb login $WANDB_API_KEY && uv pip install -e .[test] && pytest ./fastvideo/tests/nightly/test_e2e_overfit_single_sample.py -vs"
      timeout_minutes: 30
    secrets:
      RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
@@ -0,0 +1,28 @@
+name: Publish to Comfy registry
+on:
+  workflow_dispatch:
+  push:
+    branches:
+      - main
+      - master
+    paths:
+      - "pyproject.toml"
+
+permissions:
+  issues: write
+
+jobs:
+  publish-node:
+    name: Publish Custom Node to registry
+    runs-on: ubuntu-latest
+    if: ${{ github.repository_owner == 'hao-ai-lab' }}
+    steps:
+      - name: Check out code
+        uses: actions/checkout@v4
+        with:
+          submodules: true
+      - name: Publish Custom Node
+        uses: Comfy-Org/publish-node-action@v1
+        with:
+          ## Add your own personal access token to your Github Repository secrets and reference it here.
+          personal_access_token: ${{ secrets.REGISTRY_ACCESS_TOKEN }}
@@ -57,7 +57,7 @@ docs/source/training/examples/
 *.pkl

 # Reference videos
-!fastvideo/v1/tests/ssim/reference_videos/**/*.mp4
+!fastvideo/tests/ssim/reference_videos/**/*.mp4

 # Static images
 !docs/source/_static/images/**/*.png
@@ -3,7 +3,7 @@ default_stages:
  - manual # Run in CI
 exclude: |
  (?x)(
-    fastvideo/v1/third_party/.*|
+    fastvideo/third_party/.*|
    csrc/.*|
    assets/.*|
    tests/.*|
@@ -69,7 +69,7 @@ repos:
    entry: bash
    args:
      - -c
-      - 'git ls-files | grep -v "^fastvideo/v1/tests/ssim/" | grep -v "^fastvideo/v1/tests/inference/lora/L40S_reference_videos/" | grep " " && echo "Filenames should not contain spaces!" && exit 1 || exit 0'
+      - 'git ls-files | grep -v "^fastvideo/tests/ssim/" | grep -v "^fastvideo/tests/inference/lora/L40S_reference_videos/" | grep " " && echo "Filenames should not contain spaces!" && exit 1 || exit 0'
    language: system
    always_run: true
    pass_filenames: false
@@ -16,7 +16,7 @@ import sys
 # Run it with `python collect_env.py` or `python -m torch.utils.collect_env`
 from collections import namedtuple

-from fastvideo.v1.envs import environment_variables
+from fastvideo.envs import environment_variables

 try:
    import torch
@@ -81,6 +81,7 @@ DEFAULT_CONDA_PATTERNS = {
 DEFAULT_PIP_PATTERNS = {
    "torch",
    "numpy",
+    "mypy",
    "flake8",
    "triton",
    "optree",
@@ -74,8 +74,8 @@ After installation, the following nodes will be available in the ComfyUI interfa
 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)
+- [Wan2.1-I2V-14B-480P-Diffusers](https://github.com/hao-ai-lab/FastVideo/blob/main/fastvideo/configs/wan_14B_i2v_480p_pipeline.json)
+- [FastHunyuan-diffusers](https://github.com/hao-ai-lab/FastVideo/blob/main/fastvideo/configs/fasthunyuan_t2v.json)

 ### Node Configuration

@@ -310,7 +310,7 @@
          "value": -99999,
          "cachedValue": "fp16"
        },
-        "use_cpu_offload": {
+        "dit_cpu_offload": {
          "isAuto": true,
          "value": -99999,
          "cachedValue": true
@@ -517,7 +517,7 @@
          "value": "fp16",
          "cachedValue": "fp16"
        },
-        "use_cpu_offload": {
+        "dit_cpu_offload": {
          "isAuto": true,
          "value": true,
          "cachedValue": true
@@ -105,7 +105,7 @@ class VideoGenerator:
                "precision": (["fp16", "bf16"], {
                    "default": "fp16"
                }),
-                "use_cpu_offload": ([True, False], {
+                "dit_cpu_offload": ([True, False], {
                    "default": False
                }),
            }
@@ -204,7 +204,7 @@ class VideoGenerator:
        vae_config=None,
        text_encoder_config=None,
        dit_config=None,
-        use_cpu_offload=None,
+        dit_cpu_offload=None,
    ):
        print('Running FastVideo inference')

@@ -259,8 +259,8 @@ class VideoGenerator:
            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
+        if dit_cpu_offload is not None:
+            raw_generation_args['dit_cpu_offload'] = dit_cpu_offload

        generation_args = {
            k: v
@@ -552,7 +552,7 @@ app.registerExtension({
                ]
                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"
+                    "load_encoder", "load_decoder", "use_tiling", "use_temporal_tiling", "use_parallel_tiling", "dit_cpu_offload", "enable_teacache"
                ]
                const stringWidgetNames = ["prefix", "quant_config", "lora_config", "image_path"]

@@ -9,11 +9,11 @@
 ## Initialization Configuration

 ```{autodoc2-summary}
-    fastvideo.v1.configs.pipelines.PipelineConfig
+    fastvideo.configs.pipelines.PipelineConfig
 ```

 ## Sampling Configuration

 ```{autodoc2-summary}
-    fastvideo.v1.configs.sample.SamplingParam
+    fastvideo.configs.sample.SamplingParam
 ```
@@ -1,24 +1,24 @@
 # 🔍 FastVideo Overview

-This document outlines FastVideo's architecture for developers interested in framework internals or contributions. It serves as an onboarding guide for new contributors by providing an overview of the most important directories and files within the `fastvideo/v1/` codebase.
+This document outlines FastVideo's architecture for developers interested in framework internals or contributions. It serves as an onboarding guide for new contributors by providing an overview of the most important directories and files within the `fastvideo/` codebase.

-## Table of Contents - V1 Directory Structure and Files
+## Table of Contents - Directory Structure and Files

- [`fastvideo/v1/pipelines/`](#design-pipeline-system) - Core diffusion pipeline components
- [`fastvideo/v1/models/`](#design-model-components) - Model implementations
+- [`fastvideo/pipelines/`](#design-pipeline-system) - Core diffusion pipeline components
+- [`fastvideo/models/`](#design-model-components) - Model implementations
  - [`dits/`](#design-transformer-models) - Transformer-based diffusion models
  - [`vaes/`](#design-vae-variational-auto-encoder) - Variational autoencoders
  - [`encoders/`](#design-text-and-image-encoders) - Text and image encoders
  - [`schedulers/`](#design-schedulers) - Diffusion schedulers
- [`fastvideo/v1/attention/`](#design-optimized-attention) - Optimized attention implementations
- [`fastvideo/v1/distributed/`](#design-distributed-processing) - Distributed computing utilities
- [`fastvideo/v1/layers/`](#design-tensor-parallelism) - Custom neural network layers
- [`fastvideo/v1/platforms/`](#design-platforms) - Hardware platform abstractions
- [`fastvideo/v1/worker/`](#design-executor-and-worker-abstractions) - Multi-GPU process management
- [`fastvideo/v1/fastvideo_args.py`](#design-fastvideo-args) - Argument handling
- [`fastvideo/v1/forward_context.py`](#design-forwardcontext) - Forward pass context management
- `fastvideo/v1/utils.py` - Utility functions
- [`fastvideo/v1/logger.py`](#design-logger) - Logging infrastructure
+- [`fastvideo/attention/`](#design-optimized-attention) - Optimized attention implementations
+- [`fastvideo/distributed/`](#design-distributed-processing) - Distributed computing utilities
+- [`fastvideo/layers/`](#design-tensor-parallelism) - Custom neural network layers
+- [`fastvideo/platforms/`](#design-platforms) - Hardware platform abstractions
+- [`fastvideo/worker/`](#design-executor-and-worker-abstractions) - Multi-GPU process management
+- [`fastvideo/fastvideo_args.py`](#design-fastvideo-args) - Argument handling
+- [`fastvideo/forward_context.py`](#design-forwardcontext) - Forward pass context management
+- `fastvideo/utils.py` - Utility functions
+- [`fastvideo/logger.py`](#design-logger) - Logging infrastructure

 ## Core Architecture

@@ -32,7 +32,7 @@ FastVideo separates model components from execution logic with these principles:
 (design-fastvideo-args)=
 ## FastVideoArgs

-The `FastVideoArgs` class in `fastvideo/v1/fastvideo_args.py` serves as the central configuration system for FastVideo. It contains all parameters needed to control model loading, inference configuration, performance optimization settings, and more.
+The `FastVideoArgs` class in `fastvideo/fastvideo_args.py` serves as the central configuration system for FastVideo. It contains all parameters needed to control model loading, inference configuration, performance optimization settings, and more.

 Key features include:
 - **Command-line Interface**: Automatic conversion between CLI arguments and dataclass fields
@@ -111,7 +111,7 @@ def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> Forward
 (design-forwardbatch)=
 ### ForwardBatch

-Defined in `fastvideo/v1/pipelines/pipeline_batch_info.py`, `ForwardBatch` encapsulates the data payload passed between pipeline stages. It typically holds:
+Defined in `fastvideo/pipelines/pipeline_batch_info.py`, `ForwardBatch` encapsulates the data payload passed between pipeline stages. It typically holds:

 - **Input Data**: Prompts, images, generation parameters
 - **Intermediate State**: Embeddings, latents, timesteps, accumulated during stage execution
@@ -123,14 +123,14 @@ This structure facilitates clear state transitions between stages.
 (design-model-components)=
 ## Model Components

-The `fastvideo/v1/models/` directory contains implementations of the core neural network models used in video diffusion:
+The `fastvideo/models/` directory contains implementations of the core neural network models used in video diffusion:

 (design-transformer-models)=
 ### Transformer Models

 Transformer networks perform the actual denoising during diffusion:

- **Location**: `fastvideo/v1/models/dits/`
+- **Location**: `fastvideo/models/dits/`
 - **Examples**:
  - `WanTransformer3DModel`
  - `HunyuanVideoTransformer3DModel`
@@ -157,7 +157,7 @@ def forward(

 VAEs handle conversion between pixel space and latent space:

- **Location**: `fastvideo/v1/models/vaes/`
+- **Location**: `fastvideo/models/vaes/`
 - **Examples**:
  - `AutoencoderKLWan`
  - `AutoencoderKLHunyuanVideo`
@@ -175,7 +175,7 @@ FastVideo's VAE implementations include:

 Encoders process conditioning inputs into embeddings:

- **Location**: `fastvideo/v1/models/encoders/`
+- **Location**: `fastvideo/models/encoders/`
 - **Text Encoders**:
  - `CLIPTextModel`
  - `LlamaModel`
@@ -193,7 +193,7 @@ FastVideo implements optimizations such as:

 Schedulers manage the diffusion sampling process:

- **Location**: `fastvideo/v1/models/schedulers/`
+- **Location**: `fastvideo/models/schedulers/`
 - **Examples**:
  - `UniPCMultistepScheduler`
  - `FlowMatchEulerDiscreteScheduler`
@@ -219,7 +219,7 @@ def step(
 (design-optimized-attention)=
 ## Optimized Attention

-The `fastvideo/v1/attention/` directory contains optimized attention implementations crucial for efficient video diffusion:
+The `fastvideo/attention/` directory contains optimized attention implementations crucial for efficient video diffusion:

 ### Attention Backends
 Multiple implementations with automatic selection:
@@ -248,7 +248,7 @@ Supports various patterns with memory optimization techniques:
 (design-distributed-processing)=
 ## Distributed Processing

-The `fastvideo/v1/distributed/` directory contains implementations for distributed model execution:
+The `fastvideo/distributed/` directory contains implementations for distributed model execution:

 (design-tensor-parallelism)=
 ### Tensor Parallelism
@@ -260,7 +260,7 @@ Tensor parallelism splits model weights across devices:

 ```python
 # Tensor-parallel layers in a transformer block
-from fastvideo.v1.layers.linear import ColumnParallelLinear, RowParallelLinear
+from fastvideo.layers.linear import ColumnParallelLinear, RowParallelLinear

 # Split along output dimension
 self.qkv_proj = ColumnParallelLinear(
@@ -288,7 +288,7 @@ Sequence parallelism splits sequences across devices:

 ```python
 # Distributed attention for long sequences
-from fastvideo.v1.attention import DistributedAttention
+from fastvideo.attention import DistributedAttention

 self.attn = DistributedAttention(
    num_heads=num_heads,
@@ -312,7 +312,7 @@ Efficient communication primitives minimize distributed overhead:

 ### ForwardContext

-Defined in `fastvideo/v1/forward_context.py`, `ForwardContext` manages execution-specific state *within* a forward pass, particularly for low-level optimizations. It is accessed via `get_forward_context()`.
+Defined in `fastvideo/forward_context.py`, `ForwardContext` manages execution-specific state *within* a forward pass, particularly for low-level optimizations. It is accessed via `get_forward_context()`.

 - **Attention Metadata**: Configuration for optimized attention kernels (`attn_metadata`)
 - **Profiling Data**: Potential hooks for performance metrics collection
@@ -333,7 +333,7 @@ with set_forward_context(current_timestep, attn_metadata, fastvideo_args):
 (design-executor-and-worker-abstractions)=
 ## Executor and Worker System

-The `fastvideo/v1/worker/` directory contains the distributed execution framework:
+The `fastvideo/worker/` directory contains the distributed execution framework:

 ### Executor Abstraction

@@ -360,7 +360,7 @@ This design allows FastVideo to efficiently utilize multiple GPUs while providin
 (design-platforms)=
 ## Platforms

-The `fastvideo/v1/platforms/` directory provides hardware platform abstractions that enable FastVideo to run efficiently on different hardware configurations:
+The `fastvideo/platforms/` directory provides hardware platform abstractions that enable FastVideo to run efficiently on different hardware configurations:

 ### Platform Abstraction

@@ -377,7 +377,7 @@ The primary components include:
 Usage example:

 ```python
-from fastvideo.v1.platforms import current_platform, _Backend
+from fastvideo.platforms import current_platform, _Backend

 # Check hardware capabilities
 if current_platform.supports_backend(_Backend.FLASH_ATTN):
@@ -398,9 +398,9 @@ See [PR](https://github.com/hao-ai-lab/FastVideo/pull/356)

 If you're a new contributor, here are some common areas to explore:

-1. **Adding a new model**: Implement new model types in the appropriate subdirectory of `fastvideo/v1/models/`
+1. **Adding a new model**: Implement new model types in the appropriate subdirectory of `fastvideo/models/`
 2. **Optimizing performance**: Look at attention implementations or memory management
-3. **Adding a new pipeline**: Create a new pipeline subclass in `fastvideo/v1/pipelines/`
+3. **Adding a new pipeline**: Create a new pipeline subclass in `fastvideo/pipelines/`
 4. **Hardware support**: Extend the `platforms` module for new hardware targets

 When adding code, follow these practices:
@@ -10,20 +10,20 @@ This class will be the primary Python API for generating videos and images.
    fastvideo.VideoGenerator
 ```

-`````{py:class} VideoGenerator(fastvideo_args: fastvideo.v1.fastvideo_args.FastVideoArgs, executor_class: type[fastvideo.v1.worker.executor.Executor], log_stats: bool)
-:canonical: fastvideo.v1.entrypoints.video_generator.VideoGenerator
+`````{py:class} VideoGenerator(fastvideo_args: fastvideo.fastvideo_args.FastVideoArgs, executor_class: type[fastvideo.worker.executor.Executor], log_stats: bool)
+:canonical: fastvideo.entrypoints.video_generator.VideoGenerator

-```{autodoc2-docstring} fastvideo.v1.entrypoints.video_generator.VideoGenerator
+```{autodoc2-docstring} fastvideo.entrypoints.video_generator.VideoGenerator
 :parser: docs.source.autodoc2_docstring_parser
 ```

 `VideoGenerator.from_pretrained()` should be the primary way of creating a new video generator.

-````{py:method} from_pretrained(model_path: str, device: typing.Optional[str] = None, torch_dtype: typing.Optional[torch.dtype] = None, pipeline_config: typing.Optional[typing.Union[str | fastvideo.v1.configs.pipelines.PipelineConfig]] = None, **kwargs) -> fastvideo.v1.entrypoints.video_generator.VideoGenerator
-:canonical: fastvideo.v1.entrypoints.video_generator.VideoGenerator.from_pretrained
+````{py:method} from_pretrained(model_path: str, device: typing.Optional[str] = None, torch_dtype: typing.Optional[torch.dtype] = None, pipeline_config: typing.Optional[typing.Union[str | fastvideo.configs.pipelines.PipelineConfig]] = None, **kwargs) -> fastvideo.entrypoints.video_generator.VideoGenerator
+:canonical: fastvideo.entrypoints.video_generator.VideoGenerator.from_pretrained
 :classmethod:

-```{autodoc2-docstring} fastvideo.v1.entrypoints.video_generator.VideoGenerator.from_pretrained
+```{autodoc2-docstring} fastvideo.entrypoints.video_generator.VideoGenerator.from_pretrained
 :parser: docs.source.autodoc2_docstring_parser
 ```

@@ -38,25 +38,25 @@ The follow two classes `PipelineConfig` and `SamplingParam` are used to configur
 ```

 `````{py:class} PipelineConfig
-:canonical: fastvideo.v1.configs.pipelines.base.PipelineConfig
+:canonical: fastvideo.configs.pipelines.base.PipelineConfig

-```{autodoc2-docstring} fastvideo.v1.configs.pipelines.base.PipelineConfig
+```{autodoc2-docstring} fastvideo.configs.pipelines.base.PipelineConfig
 :parser: docs.source.autodoc2_docstring_parser
 ```

-````{py:method} from_pretrained(model_path: str) -> fastvideo.v1.configs.pipelines.base.PipelineConfig
-:canonical: fastvideo.v1.configs.pipelines.base.PipelineConfig.from_pretrained
+````{py:method} from_pretrained(model_path: str) -> fastvideo.configs.pipelines.base.PipelineConfig
+:canonical: fastvideo.configs.pipelines.base.PipelineConfig.from_pretrained
 :classmethod:

-```{autodoc2-docstring} fastvideo.v1.configs.pipelines.base.PipelineConfig.from_pretrained
+```{autodoc2-docstring} fastvideo.configs.pipelines.base.PipelineConfig.from_pretrained
 :parser: docs.source.autodoc2_docstring_parser
 ```


 ````{py:method} dump_to_json(file_path: str)
-:canonical: fastvideo.v1.configs.pipelines.base.PipelineConfig.dump_to_json
+:canonical: fastvideo.configs.pipelines.base.PipelineConfig.dump_to_json

-```{autodoc2-docstring} fastvideo.v1.configs.pipelines.base.PipelineConfig.dump_to_json
+```{autodoc2-docstring} fastvideo.configs.pipelines.base.PipelineConfig.dump_to_json
 :parser: docs.source.autodoc2_docstring_parser
 ```

@@ -68,16 +68,16 @@ The follow two classes `PipelineConfig` and `SamplingParam` are used to configur
 ```

 `````{py:class} SamplingParam
-:canonical: fastvideo.v1.configs.sample.base.SamplingParam
+:canonical: fastvideo.configs.sample.base.SamplingParam

-```{autodoc2-docstring} fastvideo.v1.configs.sample.base.SamplingParam
+```{autodoc2-docstring} fastvideo.configs.sample.base.SamplingParam
 :parser: docs.source.autodoc2_docstring_parser
 ```

-````{py:method} from_pretrained(model_path: str) -> fastvideo.v1.configs.sample.base.SamplingParam
-:canonical: fastvideo.v1.configs.sample.base.SamplingParam.from_pretrained
+````{py:method} from_pretrained(model_path: str) -> fastvideo.configs.sample.base.SamplingParam
+:canonical: fastvideo.configs.sample.base.SamplingParam.from_pretrained
 :classmethod:

-```{autodoc2-docstring} fastvideo.v1.configs.sample.base.SamplingParam.from_pretrained
+```{autodoc2-docstring} fastvideo.configs.sample.base.SamplingParam.from_pretrained
 :parser: docs.source.autodoc2_docstring_parser
 ```
@@ -46,25 +46,25 @@ FastVideo uses the Hugging Face Diffusers format for model organization:
 ### Implementing Modules

 Place new modules in the appropriate directories:
- Encoders: `fastvideo/v1/models/encoders/`
- VAEs: `fastvideo/v1/models/vaes/`
- Transformer models: `fastvideo/v1/models/dits/`
- Schedulers: `fastvideo/v1/models/schedulers/`
+- Encoders: `fastvideo/models/encoders/`
+- VAEs: `fastvideo/models/vaes/`
+- Transformer models: `fastvideo/models/dits/`
+- Schedulers: `fastvideo/models/schedulers/`

 ### Adapting Model Layers

 #### Layer Replacements
 Replace standard PyTorch layers with FastVideo optimized versions:
- nn.LayerNorm → fastvideo.v1.layers.layernorm.RMSNorm
- Embedding layers → fastvideo.v1.layers.vocab_parallel_embedding modules
- Activation functions → versions from fastvideo.v1.layers.activation
+- nn.LayerNorm → fastvideo.layers.layernorm.RMSNorm
+- Embedding layers → fastvideo.layers.vocab_parallel_embedding modules
+- Activation functions → versions from fastvideo.layers.activation

 #### Distributed Linear Layers
 Use appropriate parallel layers for distribution:

 ```python
 # Output dimension parallelism
-from fastvideo.v1.layers.linear import ColumnParallelLinear
+from fastvideo.layers.linear import ColumnParallelLinear
 self.q_proj = ColumnParallelLinear(
    input_size=hidden_size,
    output_size=head_size * num_heads,
@@ -73,7 +73,7 @@ self.q_proj = ColumnParallelLinear(
 )

 # Fused QKV projection
-from fastvideo.v1.layers.linear import QKVParallelLinear
+from fastvideo.layers.linear import QKVParallelLinear
 self.qkv_proj = QKVParallelLinear(
    hidden_size=hidden_size,
    head_size=attention_head_dim,
@@ -82,7 +82,7 @@ self.qkv_proj = QKVParallelLinear(
 )

 # Input dimension parallelism
-from fastvideo.v1.layers.linear import RowParallelLinear
+from fastvideo.layers.linear import RowParallelLinear
 self.out_proj = RowParallelLinear(
    input_size=head_size * num_heads,
    output_size=hidden_size,
@@ -96,8 +96,8 @@ Replace standard attention with FastVideo's optimized attention:

 ```python
 # Local attention patterns
-from fastvideo.v1.attention import LocalAttention
-from fastvideo.v1.attention.backends.abstract import _Backend
+from fastvideo.attention import LocalAttention
+from fastvideo.attention.backends.abstract import _Backend
 self.attn = LocalAttention(
    num_heads=num_heads,
    head_size=head_dim,
@@ -108,7 +108,7 @@ self.attn = LocalAttention(
 )

 # Distributed attention for long sequences
-from fastvideo.v1.attention import DistributedAttention
+from fastvideo.attention import DistributedAttention
 self.attn = DistributedAttention(
    num_heads=num_heads,
    head_size=head_dim,
@@ -130,7 +130,7 @@ self.attn = DistributedAttention(
 Register implemented modules in the model registry:

 ```python
-# In fastvideo/v1/models/registry.py
+# In fastvideo/models/registry.py
 _TEXT_TO_VIDEO_DIT_MODELS = {
    "YourTransformerModel": ("dits", "yourmodule", "YourTransformerClass"),
 }
@@ -145,7 +145,7 @@ _VAE_MODELS = {
 Create a new directory for your pipeline:

 ```
-fastvideo/v1/pipelines/
+fastvideo/pipelines/
 ├── your_pipeline/
 │   ├── __init__.py
 │   └── your_pipeline.py
@@ -167,13 +167,13 @@ Pipelines are composed of stages, each handling a specific part of the diffusion
 ### Creating Your Pipeline

 ```python
-from fastvideo.v1.pipelines.composed_pipeline_base import ComposedPipelineBase
-from fastvideo.v1.pipelines.stages import (
+from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase
+from fastvideo.pipelines.stages import (
    InputValidationStage, CLIPTextEncodingStage, TimestepPreparationStage,
    LatentPreparationStage, DenoisingStage, DecodingStage
 )
-from fastvideo.v1.fastvideo_args import FastVideoArgs
-from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
 import torch

 class MyCustomPipeline(ComposedPipelineBase):
@@ -246,7 +246,7 @@ EntryClass = MyCustomPipeline
 If existing stages don't meet your needs, create custom ones:

 ```python
-from fastvideo.v1.pipelines.stages.base import PipelineStage
+from fastvideo.pipelines.stages.base import PipelineStage

 class MyCustomStage(PipelineStage):
    """Custom processing stage for the pipeline."""
@@ -305,7 +305,7 @@ EntryClass = [MyCustomPipeline, MyOtherPipeline]
 ```

 The registry will automatically:
-1. Scan all packages under `fastvideo/v1/pipelines/`
+1. Scan all packages under `fastvideo/pipelines/`
 2. Look for `EntryClass` variables
 3. Register pipelines using their class names as identifiers

@@ -27,7 +27,7 @@ def main():
    model_name = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
    config = PipelineConfig.from_pretrained(model_name)
    config.vae_precision = "fp16"
-    config.use_cpu_offload = True
+    config.dit_cpu_offload = True

    # Create the generator
    generator = VideoGenerator.from_pretrained(
@@ -1,6 +1,6 @@
 from fastvideo import VideoGenerator

-# from fastvideo.v1.configs.sample import SamplingParam
+# from fastvideo.configs.sample import SamplingParam

 OUTPUT_PATH = "video_samples"
 def main():
@@ -13,7 +13,10 @@ def main():
        # FastVideo will automatically handle distributed setup
        num_gpus=2,
        use_fsdp_inference=True,
-        use_cpu_offload=False
+        dit_cpu_offload=False,
+        vae_cpu_offload=False,
+        text_encoder_cpu_offload=False,
+        # image_encoder_cpu_offload=False,
    )

    # sampling_param = SamplingParam.from_pretrained("Wan-AI/Wan2.1-T2V-1.3B-Diffusers")
@@ -2,7 +2,7 @@ import os
 import time
 from fastvideo import VideoGenerator

-from fastvideo.v1.configs.sample import SamplingParam
+from fastvideo.configs.sample import SamplingParam

 OUTPUT_PATH = "video_samples_dmd2"
 def main():
@@ -16,8 +16,9 @@ def main():
        num_gpus=1,
        use_fsdp_inference=True,
        # Adjust these offload parameters if you have < 32GB of VRAM
-        text_encoder_offload=False,
-        use_cpu_offload=False,
+        text_encoder_cpu_offload=False,
+        dit_cpu_offload=False,
+        vae_cpu_offload=False,
        VSA_sparsity=0.8,
    )
    load_end_time = time.perf_counter()
@@ -35,11 +36,24 @@ def main():
    video = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True, sampling_param=sampling_param)
    end_time = time.perf_counter()
    gen_time = end_time - start_time
-    e2e_gen_time = end_time - start_time
+
+    # Generate another video with a different prompt, without reloading the
+    # model!
+    prompt2 = (
+        "A majestic lion strides across the golden savanna, its powerful frame "
+        "glistening under the warm afternoon sun. The tall grass ripples gently in "
+        "the breeze, enhancing the lion's commanding presence. The tone is vibrant, "
+        "embodying the raw energy of the wild. Low angle, steady tracking shot, "
+        "cinematic.")
+    start_time = time.perf_counter()
+    video2 = generator.generate_video(prompt2, output_path=OUTPUT_PATH, save_video=True)
+    end_time = time.perf_counter()
+    gen_time2 = end_time - start_time
+

    print(f"Time taken to load model: {load_time} seconds")
    print(f"Time taken to generate video: {gen_time} seconds")
-    print(f"Time taken for e2e generation: {e2e_gen_time} seconds")
+    print(f"Time taken to generate video2: {gen_time2} seconds")


 if __name__ == "__main__":
@@ -1,5 +1,5 @@
 from fastvideo import VideoGenerator, PipelineConfig
-from fastvideo.v1.configs.sample import SamplingParam
+from fastvideo.configs.sample import SamplingParam

 def main():
    config = PipelineConfig.from_pretrained("Wan-AI/Wan2.1-T2V-1.3B-Diffusers")
@@ -9,8 +9,8 @@ def main():
        "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
        pipeline_config=config,
        use_fsdp_inference=False,      # Disable FSDP for MPS
-        use_cpu_offload=True,          
-        text_encoder_offload=True,    
+        dit_cpu_offload=True,          
+        text_encoder_cpu_offload=True,    
        pin_cpu_memory=True,           
        disable_autocast=False,        
        num_gpus=1,      
@@ -1,5 +1,5 @@
 from fastvideo import VideoGenerator
-from fastvideo.v1.configs.pipelines.base import PipelineConfig
+from fastvideo.configs.pipelines.base import PipelineConfig

 def main():

@@ -1,6 +1,6 @@
 from fastvideo import VideoGenerator, PipelineConfig, SamplingParam

-# from fastvideo.v1.configs.sample import SamplingParam
+# from fastvideo.configs.sample import SamplingParam

 OUTPUT_PATH = "video_samples_fp16"
 def main():
@@ -16,8 +16,8 @@ def main():
        # if num_gpus > 1, FastVideo will automatically handle distributed setup
        pipeline_config=pipeline_config,
        use_fsdp_inference=False,      # Disable FSDP for MPS
-        use_cpu_offload=True,          
-        text_encoder_offload=True,    
+        dit_cpu_offload=True,          
+        text_encoder_cpu_offload=True,    
        pin_cpu_memory=True,           
        disable_autocast=False,        
        num_gpus=1, 
@@ -6,7 +6,7 @@ import gradio as gr
 import torch

 from fastvideo import VideoGenerator
-from fastvideo.v1.configs.sample.base import SamplingParam
+from fastvideo.configs.sample.base import SamplingParam

 if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="FastVideo Gradio Demo")
@@ -1,5 +1,5 @@
 from fastvideo import VideoGenerator
-from fastvideo.v1.configs.sample import SamplingParam
+from fastvideo.configs.sample import SamplingParam

 OUTPUT_PATH = "./lora_out"
 def main():
@@ -2,7 +2,7 @@
 Inference using a LoRA checkpoint from FastVideo trainer.
 """
 from fastvideo import VideoGenerator
-from fastvideo.v1.configs.sample import SamplingParam
+from fastvideo.configs.sample import SamplingParam

 OUTPUT_PATH = "./lora_out"
 def main():
@@ -9,7 +9,7 @@ def main():
    gen = VideoGenerator.from_pretrained(
        model_path="Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
        num_gpus=1,
-        use_cpu_offload=False,
+        dit_cpu_offload=False,
    )
    load_time = time.perf_counter() - start_time
    print(f"Model loading time: {load_time:.2f} seconds")
@@ -84,7 +84,7 @@ miscellaneous_args=(
 torchrun \
  --nnodes 1 \
  --nproc_per_node $NUM_GPUS \
-    fastvideo/v1/training/wan_i2v_training_pipeline.py \
+    fastvideo/training/wan_i2v_training_pipeline.py \
    "${parallel_args[@]}" \
    "${model_args[@]}" \
    "${dataset_args[@]}" \
@@ -120,7 +120,7 @@ srun torchrun \
 --node_rank $SLURM_PROCID \
 --rdzv_backend=c10d \
 --rdzv_endpoint="$MASTER_ADDR:$MASTER_PORT" \
-    fastvideo/v1/training/wan_i2v_training_pipeline.py \
+    fastvideo/training/wan_i2v_training_pipeline.py \
    "${parallel_args[@]}" \
    "${model_args[@]}" \
    "${dataset_args[@]}" \
@@ -7,7 +7,7 @@ DATA_MERGE_PATH="data/crush-smol/merge.txt"
 OUTPUT_DIR="data/crush-smol_processed_i2v_1_3b_inp/"

 torchrun --nproc_per_node=$GPU_NUM \
-    fastvideo/v1/pipelines/preprocess/v1_preprocess.py \
+    fastvideo/pipelines/preprocess/v1_preprocess.py \
    --model_path $MODEL_PATH \
    --data_merge_path $DATA_MERGE_PATH \
    --preprocess_video_batch_size 8 \
@@ -84,7 +84,7 @@ miscellaneous_args=(
 torchrun \
  --nnodes 1 \
  --nproc_per_node $NUM_GPUS \
-    fastvideo/v1/training/wan_i2v_training_pipeline.py \
+    fastvideo/training/wan_i2v_training_pipeline.py \
    "${parallel_args[@]}" \
    "${model_args[@]}" \
    "${dataset_args[@]}" \
@@ -120,7 +120,7 @@ srun torchrun \
 --node_rank $SLURM_PROCID \
 --rdzv_backend=c10d \
 --rdzv_endpoint="$MASTER_ADDR:$MASTER_PORT" \
-    fastvideo/v1/training/wan_i2v_training_pipeline.py \
+    fastvideo/training/wan_i2v_training_pipeline.py \
    "${parallel_args[@]}" \
    "${model_args[@]}" \
    "${dataset_args[@]}" \
@@ -83,7 +83,7 @@ miscellaneous_args=(
 torchrun \
  --nnodes 1 \
  --nproc_per_node $NUM_GPUS \
-    fastvideo/v1/training/wan_i2v_training_pipeline.py \
+    fastvideo/training/wan_i2v_training_pipeline.py \
    "${parallel_args[@]}" \
    "${model_args[@]}" \
    "${dataset_args[@]}" \
@@ -7,7 +7,7 @@ DATA_MERGE_PATH="data/crush-smol/merge.txt"
 OUTPUT_DIR="data/crush-smol_processed_i2v/"

 torchrun --nproc_per_node=$GPU_NUM \
-    fastvideo/v1/pipelines/preprocess/v1_preprocess.py \
+    fastvideo/pipelines/preprocess/v1_preprocess.py \
    --model_path $MODEL_PATH \
    --data_merge_path $DATA_MERGE_PATH \
    --preprocess_video_batch_size 8 \
@@ -83,7 +83,7 @@ miscellaneous_args=(
 torchrun \
  --nnodes 1 \
  --nproc_per_node $NUM_GPUS \
-    fastvideo/v1/training/wan_training_pipeline.py \
+    fastvideo/training/wan_training_pipeline.py \
    "${parallel_args[@]}" \
    "${model_args[@]}" \
    "${dataset_args[@]}" \
@@ -117,7 +117,7 @@ srun torchrun \
 --node_rank $SLURM_PROCID \
 --rdzv_backend=c10d \
 --rdzv_endpoint="$MASTER_ADDR:$MASTER_PORT" \
-    fastvideo/v1/training/wan_training_pipeline.py \
+    fastvideo/training/wan_training_pipeline.py \
    "${parallel_args[@]}" \
    "${model_args[@]}" \
    "${dataset_args[@]}" \
@@ -83,7 +83,7 @@ torchrun \
  --nnodes 1 \
  --nproc_per_node $NUM_GPUS \
  --master_port 29501 \
-    fastvideo/v1/training/wan_training_pipeline.py \
+    fastvideo/training/wan_training_pipeline.py \
    "${parallel_args[@]}" \
    "${model_args[@]}" \
    "${dataset_args[@]}" \
@@ -7,7 +7,7 @@ DATA_MERGE_PATH="data/crush-smol/merge.txt"
 OUTPUT_DIR="data/crush-smol_processed_t2v/"

 torchrun --nproc_per_node=$GPU_NUM \
-    fastvideo/v1/pipelines/preprocess/v1_preprocess.py \
+    fastvideo/pipelines/preprocess/v1_preprocess.py \
    --model_path $MODEL_PATH \
    --data_merge_path $DATA_MERGE_PATH \
    --preprocess_video_batch_size 8 \
@@ -6,7 +6,7 @@ from typing import Any

 import numpy as np

-from fastvideo.v1.utils import dict_to_3d_list
+from fastvideo.utils import dict_to_3d_list


 def configure_sta(mode: str = 'STA_searching',
@@ -1,6 +1,6 @@
-from fastvideo.v1.configs.pipelines import PipelineConfig
-from fastvideo.v1.configs.sample import SamplingParam
-from fastvideo.v1.entrypoints.video_generator import VideoGenerator
+from fastvideo.configs.pipelines import PipelineConfig
+from fastvideo.configs.sample import SamplingParam
+from fastvideo.entrypoints.video_generator import VideoGenerator
 from fastvideo.version import __version__

-__all__ = ["VideoGenerator", "PipelineConfig", "SamplingParam", "__version__"]
+__all__ = ["VideoGenerator", "PipelineConfig", "SamplingParam", "__version__"]
@@ -0,0 +1,19 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from fastvideo.attention.backends.abstract import (AttentionBackend,
+                                                   AttentionMetadata,
+                                                   AttentionMetadataBuilder)
+from fastvideo.attention.layer import (DistributedAttention,
+                                       DistributedAttention_VSA, LocalAttention)
+from fastvideo.attention.selector import get_attn_backend
+
+__all__ = [
+    "DistributedAttention",
+    "LocalAttention",
+    "DistributedAttention_VSA",
+    "AttentionBackend",
+    "AttentionMetadata",
+    "AttentionMetadataBuilder",
+    # "AttentionState",
+    "get_attn_backend",
+]
@@ -6,8 +6,8 @@ from dataclasses import dataclass, fields
 from typing import TYPE_CHECKING, Any, Generic, Protocol, TypeVar

 if TYPE_CHECKING:
-    from fastvideo.v1.fastvideo_args import FastVideoArgs
-    from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
+    from fastvideo.fastvideo_args import FastVideoArgs
+    from fastvideo.pipelines.pipeline_batch_info import ForwardBatch

 import torch

@@ -12,11 +12,11 @@ try:
 except ImportError:
    flash_attn_func = flash_attn_2_func

-from fastvideo.v1.attention.backends.abstract import (AttentionBackend,
-                                                      AttentionImpl,
-                                                      AttentionMetadata,
-                                                      AttentionMetadataBuilder)
-from fastvideo.v1.logger import init_logger
+from fastvideo.attention.backends.abstract import (AttentionBackend,
+                                                   AttentionImpl,
+                                                   AttentionMetadata,
+                                                   AttentionMetadataBuilder)
+from fastvideo.logger import init_logger

 logger = init_logger(__name__)

@@ -3,11 +3,9 @@
 import torch
 from sageattention import sageattn

-from fastvideo.v1.attention.backends.abstract import (
-    AttentionBackend)  # FlashAttentionMetadata,
-from fastvideo.v1.attention.backends.abstract import (AttentionImpl,
-                                                      AttentionMetadata)
-from fastvideo.v1.logger import init_logger
+from fastvideo.attention.backends.abstract import (  # FlashAttentionMetadata,
+    AttentionBackend, AttentionImpl, AttentionMetadata)
+from fastvideo.logger import init_logger

 logger = init_logger(__name__)

@@ -2,11 +2,9 @@

 import torch

-from fastvideo.v1.attention.backends.abstract import (
-    AttentionBackend)  # FlashAttentionMetadata,
-from fastvideo.v1.attention.backends.abstract import (AttentionImpl,
-                                                      AttentionMetadata)
-from fastvideo.v1.logger import init_logger
+from fastvideo.attention.backends.abstract import (  # FlashAttentionMetadata,
+    AttentionBackend, AttentionImpl, AttentionMetadata)
+from fastvideo.logger import init_logger

 logger = init_logger(__name__)

@@ -7,17 +7,17 @@ import torch
 from einops import rearrange
 from st_attn import sliding_tile_attention

-import fastvideo.v1.envs as envs
-from fastvideo.v1.attention.backends.abstract import (AttentionBackend,
-                                                      AttentionImpl,
-                                                      AttentionMetadata,
-                                                      AttentionMetadataBuilder)
-from fastvideo.v1.distributed import get_sp_group
-from fastvideo.v1.fastvideo_args import FastVideoArgs
-from fastvideo.v1.forward_context import ForwardContext, get_forward_context
-from fastvideo.v1.logger import init_logger
-from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
-from fastvideo.v1.utils import dict_to_3d_list
+import fastvideo.envs as envs
+from fastvideo.attention.backends.abstract import (AttentionBackend,
+                                                   AttentionImpl,
+                                                   AttentionMetadata,
+                                                   AttentionMetadataBuilder)
+from fastvideo.distributed import get_sp_group
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.forward_context import ForwardContext, get_forward_context
+from fastvideo.logger import init_logger
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+from fastvideo.utils import dict_to_3d_list

 logger = init_logger(__name__)

@@ -10,14 +10,14 @@ try:
 except ImportError:
    video_sparse_attn = None

-from fastvideo.v1.attention.backends.abstract import (AttentionBackend,
-                                                      AttentionImpl,
-                                                      AttentionMetadata,
-                                                      AttentionMetadataBuilder)
-from fastvideo.v1.distributed import get_sp_group
-from fastvideo.v1.fastvideo_args import FastVideoArgs
-from fastvideo.v1.logger import init_logger
-from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
+from fastvideo.attention.backends.abstract import (AttentionBackend,
+                                                   AttentionImpl,
+                                                   AttentionMetadata,
+                                                   AttentionMetadataBuilder)
+from fastvideo.distributed import get_sp_group
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.logger import init_logger
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch

 logger = init_logger(__name__)

@@ -3,15 +3,14 @@
 import torch
 import torch.nn as nn

-from fastvideo.v1.attention.selector import (backend_name_to_enum,
-                                             get_attn_backend)
-from fastvideo.v1.distributed.communication_op import (
+from fastvideo.attention.selector import backend_name_to_enum, get_attn_backend
+from fastvideo.distributed.communication_op import (
    sequence_model_parallel_all_gather, sequence_model_parallel_all_to_all_4D)
-from fastvideo.v1.distributed.parallel_state import (get_sp_parallel_rank,
-                                                     get_sp_world_size)
-from fastvideo.v1.forward_context import ForwardContext, get_forward_context
-from fastvideo.v1.platforms import AttentionBackendEnum
-from fastvideo.v1.utils import get_compute_dtype
+from fastvideo.distributed.parallel_state import (get_sp_parallel_rank,
+                                                  get_sp_world_size)
+from fastvideo.forward_context import ForwardContext, get_forward_context
+from fastvideo.platforms import AttentionBackendEnum
+from fastvideo.utils import get_compute_dtype


 class DistributedAttention(nn.Module):
@@ -9,11 +9,11 @@ from typing import cast

 import torch

-import fastvideo.v1.envs as envs
-from fastvideo.v1.attention.backends.abstract import AttentionBackend
-from fastvideo.v1.logger import init_logger
-from fastvideo.v1.platforms import AttentionBackendEnum, current_platform
-from fastvideo.v1.utils import STR_BACKEND_ENV_VAR, resolve_obj_by_qualname
+import fastvideo.envs as envs
+from fastvideo.attention.backends.abstract import AttentionBackend
+from fastvideo.logger import init_logger
+from fastvideo.platforms import AttentionBackendEnum, current_platform
+from fastvideo.utils import STR_BACKEND_ENV_VAR, resolve_obj_by_qualname

 logger = init_logger(__name__)

@@ -0,0 +1,145 @@
+import dataclasses
+from typing import Any, Optional
+
+from fastvideo.configs.utils import update_config_from_args
+from fastvideo.utils import FlexibleArgumentParser, StoreBoolean
+
+
+@dataclasses.dataclass
+class PreprocessConfig:
+    """Configuration for preprocessing operations."""
+
+    # Model and dataset configuration
+    model_path: str = ""
+    dataset_path: str = ""
+    dataset_output_dir: str = "./output"
+
+    # Dataloader configuration
+    dataloader_num_workers: int = 1
+    preprocess_video_batch_size: int = 2
+
+    # Saver configuration
+    samples_per_file: int = 64
+    flush_frequency: int = 256
+
+    # Video processing parameters
+    max_height: int = 480
+    max_width: int = 848
+    num_frames: int = 163
+    video_length_tolerance_range: float = 2.0
+    train_fps: int = 30
+    speed_factor: float = 1.0
+    drop_short_ratio: float = 1.0
+    do_temporal_sample: bool = False
+
+    # Model configuration
+    training_cfg_rate: float = 0.0
+
+    @staticmethod
+    def add_cli_args(parser: FlexibleArgumentParser,
+                     prefix: str = "preprocess") -> FlexibleArgumentParser:
+        """Add preprocessing configuration arguments to the parser."""
+        prefix_with_dot = f"{prefix}." if (prefix.strip() != "") else ""
+
+        preprocess_args = parser.add_argument_group("Preprocessing Arguments")
+        # Model & Dataset
+        preprocess_args.add_argument(f"--{prefix_with_dot}model-path",
+                                     type=str,
+                                     default=PreprocessConfig.model_path,
+                                     help="Path to the model for preprocessing")
+        preprocess_args.add_argument(
+            f"--{prefix_with_dot}dataset-path",
+            type=str,
+            default=PreprocessConfig.dataset_path,
+            help="Path to the dataset directory for preprocessing")
+        preprocess_args.add_argument(
+            f"--{prefix_with_dot}dataset-output-dir",
+            type=str,
+            default=PreprocessConfig.dataset_output_dir,
+            help="The output directory where the dataset will be written.")
+
+        # Dataloader
+        preprocess_args.add_argument(
+            f"--{prefix_with_dot}dataloader-num-workers",
+            type=int,
+            default=PreprocessConfig.dataloader_num_workers,
+            help=
+            "Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process."
+        )
+        preprocess_args.add_argument(
+            f"--{prefix_with_dot}preprocess-video-batch-size",
+            type=int,
+            default=PreprocessConfig.preprocess_video_batch_size,
+            help="Batch size (per device) for the training dataloader.")
+
+        # Saver
+        preprocess_args.add_argument(f"--{prefix_with_dot}samples-per-file",
+                                     type=int,
+                                     default=PreprocessConfig.samples_per_file,
+                                     help="Number of samples per output file")
+        preprocess_args.add_argument(f"--{prefix_with_dot}flush-frequency",
+                                     type=int,
+                                     default=PreprocessConfig.flush_frequency,
+                                     help="How often to save to parquet files")
+
+        # Video processing parameters
+        preprocess_args.add_argument(f"--{prefix_with_dot}max-height",
+                                     type=int,
+                                     default=PreprocessConfig.max_height,
+                                     help="Maximum height for video processing")
+        preprocess_args.add_argument(f"--{prefix_with_dot}max-width",
+                                     type=int,
+                                     default=PreprocessConfig.max_width,
+                                     help="Maximum width for video processing")
+        preprocess_args.add_argument(f"--{prefix_with_dot}num-frames",
+                                     type=int,
+                                     default=PreprocessConfig.num_frames,
+                                     help="Number of frames to process")
+        preprocess_args.add_argument(
+            f"--{prefix_with_dot}video-length-tolerance-range",
+            type=float,
+            default=PreprocessConfig.video_length_tolerance_range,
+            help="Video length tolerance range")
+        preprocess_args.add_argument(f"--{prefix_with_dot}train-fps",
+                                     type=int,
+                                     default=PreprocessConfig.train_fps,
+                                     help="Training FPS")
+        preprocess_args.add_argument(f"--{prefix_with_dot}speed-factor",
+                                     type=float,
+                                     default=PreprocessConfig.speed_factor,
+                                     help="Speed factor for video processing")
+        preprocess_args.add_argument(f"--{prefix_with_dot}drop-short-ratio",
+                                     type=float,
+                                     default=PreprocessConfig.drop_short_ratio,
+                                     help="Ratio for dropping short videos")
+        preprocess_args.add_argument(
+            f"--{prefix_with_dot}do-temporal-sample",
+            action=StoreBoolean,
+            default=PreprocessConfig.do_temporal_sample,
+            help="Whether to do temporal sampling")
+
+        # Model Training configuration
+        preprocess_args.add_argument(f"--{prefix_with_dot}training-cfg-rate",
+                                     type=float,
+                                     default=PreprocessConfig.training_cfg_rate,
+                                     help="Training CFG rate")
+
+        return parser
+
+    @classmethod
+    def from_kwargs(cls, kwargs: dict[str,
+                                      Any]) -> Optional["PreprocessConfig"]:
+        """Create PreprocessConfig from keyword arguments."""
+        preprocess_config = cls()
+        if not update_config_from_args(
+                preprocess_config, kwargs, prefix="preprocess", pop_args=True):
+            return None
+        return preprocess_config
+
+    def check_preprocess_config(self) -> None:
+        if self.dataset_path == "":
+            raise ValueError("dataset_path must be set for preprocessing mode")
+        if self.samples_per_file <= 0:
+            raise ValueError("samples_per_file must be greater than 0")
+        if self.flush_frequency <= 0:
+            raise ValueError("flush_frequency must be greater than 0")
@@ -1,7 +1,7 @@
 {
  "embedded_cfg_scale": 6,
  "flow_shift": 17,
-  "use_cpu_offload": false,
+  "dit_cpu_offload": false,
  "disable_autocast": false,
  "precision": "bf16",
  "vae_precision": "fp32",
@@ -0,0 +1,6 @@
+from fastvideo.configs.models.base import ModelConfig
+from fastvideo.configs.models.dits.base import DiTConfig
+from fastvideo.configs.models.encoders.base import EncoderConfig
+from fastvideo.configs.models.vaes.base import VAEConfig
+
+__all__ = ["ModelConfig", "VAEConfig", "DiTConfig", "EncoderConfig"]
@@ -2,7 +2,7 @@
 from dataclasses import dataclass, field, fields
 from typing import Any

-from fastvideo.v1.logger import init_logger
+from fastvideo.logger import init_logger

 logger = init_logger(__name__)

@@ -0,0 +1,5 @@
+from fastvideo.configs.models.dits.hunyuanvideo import HunyuanVideoConfig
+from fastvideo.configs.models.dits.stepvideo import StepVideoConfig
+from fastvideo.configs.models.dits.wanvideo import WanVideoConfig
+
+__all__ = ["HunyuanVideoConfig", "WanVideoConfig", "StepVideoConfig"]
@@ -2,9 +2,9 @@
 from dataclasses import dataclass, field
 from typing import Any

-from fastvideo.v1.configs.models.base import ArchConfig, ModelConfig
-from fastvideo.v1.layers.quantization import QuantizationConfig
-from fastvideo.v1.platforms import AttentionBackendEnum
+from fastvideo.configs.models.base import ArchConfig, ModelConfig
+from fastvideo.layers.quantization import QuantizationConfig
+from fastvideo.platforms import AttentionBackendEnum


@dataclass
@@ -3,7 +3,7 @@ from dataclasses import dataclass, field

 import torch

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


 def is_double_block(n: str, m) -> bool:
@@ -1,7 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass, field

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


@dataclass
@@ -1,7 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass, field

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


 def is_blocks(n: str, m) -> bool:
@@ -0,0 +1,14 @@
+from fastvideo.configs.models.encoders.base import (BaseEncoderOutput,
+                                                    EncoderConfig,
+                                                    ImageEncoderConfig,
+                                                    TextEncoderConfig)
+from fastvideo.configs.models.encoders.clip import (CLIPTextConfig,
+                                                    CLIPVisionConfig)
+from fastvideo.configs.models.encoders.llama import LlamaConfig
+from fastvideo.configs.models.encoders.t5 import T5Config
+
+__all__ = [
+    "EncoderConfig", "TextEncoderConfig", "ImageEncoderConfig",
+    "BaseEncoderOutput", "CLIPTextConfig", "CLIPVisionConfig", "LlamaConfig",
+    "T5Config"
+]
@@ -4,9 +4,9 @@ from typing import Any

 import torch

-from fastvideo.v1.configs.models.base import ArchConfig, ModelConfig
-from fastvideo.v1.layers.quantization import QuantizationConfig
-from fastvideo.v1.platforms import AttentionBackendEnum
+from fastvideo.configs.models.base import ArchConfig, ModelConfig
+from fastvideo.layers.quantization import QuantizationConfig
+from fastvideo.platforms import AttentionBackendEnum


@dataclass
@@ -1,10 +1,10 @@
 # SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass, field

-from fastvideo.v1.configs.models.encoders.base import (ImageEncoderArchConfig,
-                                                       ImageEncoderConfig,
-                                                       TextEncoderArchConfig,
-                                                       TextEncoderConfig)
+from fastvideo.configs.models.encoders.base import (ImageEncoderArchConfig,
+                                                    ImageEncoderConfig,
+                                                    TextEncoderArchConfig,
+                                                    TextEncoderConfig)


 def _is_transformer_layer(n: str, m) -> bool:
@@ -1,8 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass, field

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


 def _is_transformer_layer(n: str, m) -> bool:
@@ -1,8 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass, field

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


 def _is_transformer_layer(n: str, m) -> bool:
@@ -0,0 +1,9 @@
+from fastvideo.configs.models.vaes.hunyuanvae import HunyuanVAEConfig
+from fastvideo.configs.models.vaes.stepvideovae import StepVideoVAEConfig
+from fastvideo.configs.models.vaes.wanvae import WanVAEConfig
+
+__all__ = [
+    "HunyuanVAEConfig",
+    "WanVAEConfig",
+    "StepVideoVAEConfig",
+]
@@ -6,8 +6,8 @@ from typing import Any

 import torch

-from fastvideo.v1.configs.models.base import ArchConfig, ModelConfig
-from fastvideo.v1.utils import StoreBoolean
+from fastvideo.configs.models.base import ArchConfig, ModelConfig
+from fastvideo.utils import StoreBoolean


@dataclass
@@ -1,7 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass, field

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


@dataclass
@@ -1,7 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass, field

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


@dataclass
@@ -3,7 +3,7 @@ from dataclasses import dataclass, field

 import torch

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


@dataclass
@@ -0,0 +1,15 @@
+from fastvideo.configs.pipelines.base import (PipelineConfig,
+                                              SlidingTileAttnConfig)
+from fastvideo.configs.pipelines.hunyuan import FastHunyuanConfig, HunyuanConfig
+from fastvideo.configs.pipelines.registry import (
+    get_pipeline_config_cls_from_name)
+from fastvideo.configs.pipelines.stepvideo import StepVideoT2VConfig
+from fastvideo.configs.pipelines.wan import (WanI2V480PConfig, WanI2V720PConfig,
+                                             WanT2V480PConfig, WanT2V720PConfig)
+
+__all__ = [
+    "HunyuanConfig", "FastHunyuanConfig", "PipelineConfig",
+    "SlidingTileAttnConfig", "WanT2V480PConfig", "WanI2V480PConfig",
+    "WanT2V720PConfig", "WanI2V720PConfig", "StepVideoT2VConfig",
+    "get_pipeline_config_cls_from_name"
+]
@@ -7,13 +7,12 @@ from typing import Any, cast

 import torch

-from fastvideo.v1.configs.models import (DiTConfig, EncoderConfig, ModelConfig,
-                                         VAEConfig)
-from fastvideo.v1.configs.models.encoders import BaseEncoderOutput
-from fastvideo.v1.configs.utils import update_config_from_args
-from fastvideo.v1.logger import init_logger
-from fastvideo.v1.utils import (FlexibleArgumentParser, StoreBoolean,
-                                shallow_asdict)
+from fastvideo.configs.models import (DiTConfig, EncoderConfig, ModelConfig,
+                                      VAEConfig)
+from fastvideo.configs.models.encoders import BaseEncoderOutput
+from fastvideo.configs.utils import update_config_from_args
+from fastvideo.logger import init_logger
+from fastvideo.utils import FlexibleArgumentParser, StoreBoolean, shallow_asdict

 logger = init_logger(__name__)

@@ -213,11 +212,11 @@ class PipelineConfig:
        )

        # Add VAE configuration arguments
-        from fastvideo.v1.configs.models.vaes.base import VAEConfig
+        from fastvideo.configs.models.vaes.base import VAEConfig
        VAEConfig.add_cli_args(parser, prefix=f"{prefix_with_dot}vae-config")

        # Add DiT configuration arguments
-        from fastvideo.v1.configs.models.dits.base import DiTConfig
+        from fastvideo.configs.models.dits.base import DiTConfig
        DiTConfig.add_cli_args(parser, prefix=f"{prefix_with_dot}dit-config")

        return parser
@@ -241,7 +240,7 @@ class PipelineConfig:
        """
        use the pipeline class setting from model_path to match the pipeline config
        """
-        from fastvideo.v1.configs.pipelines.registry import (
+        from fastvideo.configs.pipelines.registry import (
            get_pipeline_config_cls_from_name)
        pipeline_config_cls = get_pipeline_config_cls_from_name(model_path)

@@ -256,7 +255,7 @@ class PipelineConfig:
        kwargs: dictionary of kwargs
        config_cli_prefix: prefix of CLI arguments for this PipelineConfig instance
        """
-        from fastvideo.v1.configs.pipelines.registry import (
+        from fastvideo.configs.pipelines.registry import (
            get_pipeline_config_cls_from_name)

        prefix_with_dot = f"{config_cli_prefix}." if (config_cli_prefix.strip()
@@ -5,12 +5,12 @@ from typing import TypedDict

 import torch

-from fastvideo.v1.configs.models import DiTConfig, EncoderConfig, VAEConfig
-from fastvideo.v1.configs.models.dits import HunyuanVideoConfig
-from fastvideo.v1.configs.models.encoders import (BaseEncoderOutput,
-                                                  CLIPTextConfig, LlamaConfig)
-from fastvideo.v1.configs.models.vaes import HunyuanVAEConfig
-from fastvideo.v1.configs.pipelines.base import PipelineConfig
+from fastvideo.configs.models import DiTConfig, EncoderConfig, VAEConfig
+from fastvideo.configs.models.dits import HunyuanVideoConfig
+from fastvideo.configs.models.encoders import (BaseEncoderOutput,
+                                               CLIPTextConfig, LlamaConfig)
+from fastvideo.configs.models.vaes import HunyuanVAEConfig
+from fastvideo.configs.pipelines.base import PipelineConfig

 PROMPT_TEMPLATE_ENCODE_VIDEO = (
    "<|start_header_id|>system<|end_header_id|>\n\nDescribe the video by detailing the following aspects: "
@@ -4,18 +4,15 @@
 import os
 from collections.abc import Callable

-from fastvideo.v1.configs.pipelines.base import PipelineConfig
-from fastvideo.v1.configs.pipelines.hunyuan import (FastHunyuanConfig,
-                                                    HunyuanConfig)
-from fastvideo.v1.configs.pipelines.stepvideo import StepVideoT2VConfig
-from fastvideo.v1.configs.pipelines.wan import (FastWanT2V480PConfig,
-                                                WanI2V480PConfig,
-                                                WanI2V720PConfig,
-                                                WanT2V480PConfig,
-                                                WanT2V720PConfig)
-from fastvideo.v1.logger import init_logger
-from fastvideo.v1.utils import (maybe_download_model_index,
-                                verify_model_config_and_directory)
+from fastvideo.configs.pipelines.base import PipelineConfig
+from fastvideo.configs.pipelines.hunyuan import FastHunyuanConfig, HunyuanConfig
+from fastvideo.configs.pipelines.stepvideo import StepVideoT2VConfig
+from fastvideo.configs.pipelines.wan import (FastWanT2V480PConfig,
+                                             WanI2V480PConfig, WanI2V720PConfig,
+                                             WanT2V480PConfig, WanT2V720PConfig)
+from fastvideo.logger import init_logger
+from fastvideo.utils import (maybe_download_model_index,
+                             verify_model_config_and_directory)

 logger = init_logger(__name__)

@@ -1,10 +1,10 @@
 # SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass, field

-from fastvideo.v1.configs.models import DiTConfig, VAEConfig
-from fastvideo.v1.configs.models.dits import StepVideoConfig
-from fastvideo.v1.configs.models.vaes import StepVideoVAEConfig
-from fastvideo.v1.configs.pipelines.base import PipelineConfig
+from fastvideo.configs.models import DiTConfig, VAEConfig
+from fastvideo.configs.models.dits import StepVideoConfig
+from fastvideo.configs.models.vaes import StepVideoVAEConfig
+from fastvideo.configs.pipelines.base import PipelineConfig


@dataclass
@@ -4,12 +4,12 @@ from dataclasses import dataclass, field

 import torch

-from fastvideo.v1.configs.models import DiTConfig, EncoderConfig, VAEConfig
-from fastvideo.v1.configs.models.dits import WanVideoConfig
-from fastvideo.v1.configs.models.encoders import (BaseEncoderOutput,
-                                                  CLIPVisionConfig, T5Config)
-from fastvideo.v1.configs.models.vaes import WanVAEConfig
-from fastvideo.v1.configs.pipelines.base import PipelineConfig
+from fastvideo.configs.models import DiTConfig, EncoderConfig, VAEConfig
+from fastvideo.configs.models.dits import WanVideoConfig
+from fastvideo.configs.models.encoders import (BaseEncoderOutput,
+                                               CLIPVisionConfig, T5Config)
+from fastvideo.configs.models.vaes import WanVAEConfig
+from fastvideo.configs.pipelines.base import PipelineConfig


 def t5_postprocess_text(outputs: BaseEncoderOutput) -> torch.tensor:
@@ -0,0 +1,3 @@
+from fastvideo.configs.sample.base import SamplingParam
+
+__all__ = ["SamplingParam"]
@@ -2,7 +2,7 @@
 from dataclasses import dataclass
 from typing import Any

-from fastvideo.v1.logger import init_logger
+from fastvideo.logger import init_logger

 logger = init_logger(__name__)

@@ -23,6 +23,7 @@ class SamplingParam:
    negative_prompt: str | None = None
    prompt_path: str | None = None
    output_path: str = "outputs/"
+    output_video_name: str | None = None

    # Batch info
    num_videos_per_prompt: int = 1
@@ -66,7 +67,7 @@ class SamplingParam:

    @classmethod
    def from_pretrained(cls, model_path: str) -> "SamplingParam":
-        from fastvideo.v1.configs.sample.registry import (
+        from fastvideo.configs.sample.registry import (
            get_sampling_param_cls_for_name)
        sampling_cls = get_sampling_param_cls_for_name(model_path)
        if sampling_cls is not None:
@@ -106,6 +107,12 @@ class SamplingParam:
            default=SamplingParam.output_path,
            help="Path to save the generated video",
        )
+        parser.add_argument(
+            "--output-video-name",
+            type=str,
+            default=SamplingParam.output_video_name,
+            help="Name of the output video",
+        )
        parser.add_argument(
            "--num-videos-per-prompt",
            type=int,
@@ -1,8 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass, field

-from fastvideo.v1.configs.sample.base import SamplingParam
-from fastvideo.v1.configs.sample.teacache import TeaCacheParams
+from fastvideo.configs.sample.base import SamplingParam
+from fastvideo.configs.sample.teacache import TeaCacheParams


@dataclass
@@ -3,17 +3,17 @@ import os
 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 (FastWanT2V480PConfig,
-                                             WanI2V_14B_480P_SamplingParam,
-                                             WanI2V_14B_720P_SamplingParam,
-                                             WanT2V_1_3B_SamplingParam,
-                                             WanT2V_14B_SamplingParam)
-from fastvideo.v1.logger import init_logger
-from fastvideo.v1.utils import (maybe_download_model_index,
-                                verify_model_config_and_directory)
+from fastvideo.configs.sample.hunyuan import (FastHunyuanSamplingParam,
+                                              HunyuanSamplingParam)
+from fastvideo.configs.sample.stepvideo import StepVideoT2VSamplingParam
+from fastvideo.configs.sample.wan import (FastWanT2V480PConfig,
+                                          WanI2V_14B_480P_SamplingParam,
+                                          WanI2V_14B_720P_SamplingParam,
+                                          WanT2V_1_3B_SamplingParam,
+                                          WanT2V_14B_SamplingParam)
+from fastvideo.logger import init_logger
+from fastvideo.utils import (maybe_download_model_index,
+                             verify_model_config_and_directory)

 logger = init_logger(__name__)
 # Registry maps specific model weights to their config classes
@@ -1,7 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass

-from fastvideo.v1.configs.sample.base import SamplingParam
+from fastvideo.configs.sample.base import SamplingParam


@dataclass
@@ -1,7 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass, field

-from fastvideo.v1.configs.sample.base import CacheParams
+from fastvideo.configs.sample.base import CacheParams


@dataclass
@@ -1,8 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass, field

-from fastvideo.v1.configs.sample.base import SamplingParam
-from fastvideo.v1.configs.sample.teacache import WanTeaCacheParams
+from fastvideo.configs.sample.base import SamplingParam
+from fastvideo.configs.sample.teacache import WanTeaCacheParams


@dataclass
@@ -1,10 +1,11 @@
+import argparse
 from typing import Any


 def update_config_from_args(config: Any,
                            args_dict: dict[str, Any],
                            prefix: str = "",
-                            pop_args: bool = False) -> None:
+                            pop_args: bool = False) -> bool:
    """
    Update configuration object from arguments dictionary.
    
@@ -43,3 +44,18 @@ def update_config_from_args(config: Any,
        for key in args_to_remove:
            if key not in args_not_to_remove:
                args_dict.pop(key)
+
+    return len(args_to_remove) > 0
+
+
+def clean_cli_args(args: argparse.Namespace) -> dict[str, Any]:
+    """
+    Clean the arguments by removing the ones that not explicitly provided by the user.
+    """
+    provided_args = {}
+    for k, v in vars(args).items():
+        if (v is not None and hasattr(args, '_provided')
+                and k in args._provided):
+            provided_args[k] = v
+
+    return provided_args
@@ -1,7 +1,7 @@
 {
  "embedded_cfg_scale": 6.0,
  "flow_shift": 3,
-  "use_cpu_offload": true,
+  "dit_cpu_offload": true,
  "disable_autocast": false,
  "precision": "bf16",
  "vae_precision": "fp32",
@@ -1,7 +1,7 @@
 {
  "embedded_cfg_scale": 6.0,
  "flow_shift": 3,
-  "use_cpu_offload": true,
+  "dit_cpu_offload": true,
  "disable_autocast": false,
  "precision": "bf16",
  "vae_precision": "fp32",
@@ -2,13 +2,12 @@
 from torchvision import transforms
 from torchvision.transforms import Lambda

-from fastvideo.v1.dataset.parquet_dataset_map_style import (
+from fastvideo.dataset.parquet_dataset_map_style import (
    build_parquet_map_style_dataloader)
-from fastvideo.v1.dataset.preprocessing_datasets import (
-    VideoCaptionMergedDataset)
-from fastvideo.v1.dataset.transform import (CenterCropResizeVideo, Normalize255,
-                                            TemporalRandomCrop)
-from fastvideo.v1.dataset.validation_dataset import ValidationDataset
+from fastvideo.dataset.preprocessing_datasets import VideoCaptionMergedDataset
+from fastvideo.dataset.transform import (CenterCropResizeVideo, Normalize255,
+                                         TemporalRandomCrop)
+from fastvideo.dataset.validation_dataset import ValidationDataset


 def getdataset(args) -> VideoCaptionMergedDataset:
@@ -7,13 +7,13 @@ import time
 import torch.distributed as dist
 import torch.distributed.checkpoint as dist_cp

-from fastvideo.v1.dataset.parquet_dataset_iterable_style import (
+from fastvideo.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 (
+from fastvideo.distributed import get_world_rank
+from fastvideo.distributed.parallel_state import (
    cleanup_dist_env_and_memory, get_local_torch_device,
    maybe_init_distributed_environment_and_model_parallel)
-from fastvideo.v1.logger import init_logger
+from fastvideo.logger import init_logger

 logger = init_logger(__name__)

@@ -52,9 +52,9 @@ def main() -> None:
                        help='Path to save/load checkpoint')
    '''
    example launch command:
-    torchrun --nproc_per_node=1 --master_port=12358 fastvideo/v1/dataset/benchmarks/benchmark_parquet_dataset_iterable_style.py --path data/crush-smol/latents/combined_parquet_dataset --batch_size 2 --num_data_workers 2 --num_epoch 2 --num_batches_per_epoch 2 --verify_resume
-    torchrun --nproc_per_node=8 --master_port=12358 fastvideo/v1/dataset/benchmarks/benchmark_parquet_dataset_iterable_style.py --path data/crush-smol/latents/combined_parquet_dataset --batch_size 2 --num_data_workers 1 --num_epoch 2 --num_batches_per_epoch 5 --verify_resume
-    torchrun --nproc_per_node=8 --master_port=12358 fastvideo/v1/dataset/benchmarks/benchmark_parquet_dataset_iterable_style.py --path /mnt/sharefs/users/hao.zhang/Vchitect-2M/Wan-Syn/latents/ --batch_size 2 --num_data_workers 4 --num_epoch 2 --num_batches_per_epoch 100 
+    torchrun --nproc_per_node=1 --master_port=12358 fastvideo/dataset/benchmarks/benchmark_parquet_dataset_iterable_style.py --path data/crush-smol/latents/combined_parquet_dataset --batch_size 2 --num_data_workers 2 --num_epoch 2 --num_batches_per_epoch 2 --verify_resume
+    torchrun --nproc_per_node=8 --master_port=12358 fastvideo/dataset/benchmarks/benchmark_parquet_dataset_iterable_style.py --path data/crush-smol/latents/combined_parquet_dataset --batch_size 2 --num_data_workers 1 --num_epoch 2 --num_batches_per_epoch 5 --verify_resume
+    torchrun --nproc_per_node=8 --master_port=12358 fastvideo/dataset/benchmarks/benchmark_parquet_dataset_iterable_style.py --path /mnt/sharefs/users/hao.zhang/Vchitect-2M/Wan-Syn/latents/ --batch_size 2 --num_data_workers 4 --num_epoch 2 --num_batches_per_epoch 100 
    '''
    args = parser.parse_args()
    world_size = int(os.environ.get("WORLD_SIZE", 1))
@@ -8,14 +8,14 @@ import torch
 import torch.distributed as dist
 import torch.distributed.checkpoint as dist_cp

-from fastvideo.v1.dataset.dataloader.schema import pyarrow_schema_t2v
-from fastvideo.v1.dataset.parquet_dataset_map_style import (
+from fastvideo.dataset.dataloader.schema import pyarrow_schema_t2v
+from fastvideo.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 (
+from fastvideo.distributed import get_world_rank
+from fastvideo.distributed.parallel_state import (
    cleanup_dist_env_and_memory, get_local_torch_device,
    maybe_init_distributed_environment_and_model_parallel)
-from fastvideo.v1.logger import init_logger
+from fastvideo.logger import init_logger

 logger = init_logger(__name__)

@@ -55,9 +55,9 @@ def main() -> None:
                        help='Path to save/load checkpoint')
    '''
    example launch command:
-    torchrun --nproc_per_node=1 --master_port=12358 fastvideo/v1/dataset/benchmarks/benchmark_parquet_dataset_map_style.py --path data/crush-smol/latents/combined_parquet_dataset --batch_size 4 --num_data_workers 1 --num_epoch 2 --num_batches_per_epoch 3 --verify_resume
-    torchrun --nproc_per_node=8 --master_port=12358 fastvideo/v1/dataset/benchmarks/benchmark_parquet_dataset_map_style.py --path data/crush-smol/latents/combined_parquet_dataset --batch_size 2 --num_data_workers 1 --num_epoch 2 --num_batches_per_epoch 5 --verify_resume
-    torchrun --nproc_per_node=8 --master_port=12358 fastvideo/v1/dataset/benchmarks/benchmark_parquet_dataset_map_style.py --path /mnt/sharefs/users/hao.zhang/Vchitect-2M/Wan-Syn/latents/ --batch_size 2 --num_data_workers 4 --num_epoch 2 --num_batches_per_epoch 100 
+    torchrun --nproc_per_node=1 --master_port=12358 fastvideo/dataset/benchmarks/benchmark_parquet_dataset_map_style.py --path data/crush-smol/latents/combined_parquet_dataset --batch_size 4 --num_data_workers 1 --num_epoch 2 --num_batches_per_epoch 3 --verify_resume
+    torchrun --nproc_per_node=8 --master_port=12358 fastvideo/dataset/benchmarks/benchmark_parquet_dataset_map_style.py --path data/crush-smol/latents/combined_parquet_dataset --batch_size 2 --num_data_workers 1 --num_epoch 2 --num_batches_per_epoch 5 --verify_resume
+    torchrun --nproc_per_node=8 --master_port=12358 fastvideo/dataset/benchmarks/benchmark_parquet_dataset_map_style.py --path /mnt/sharefs/users/hao.zhang/Vchitect-2M/Wan-Syn/latents/ --batch_size 2 --num_data_workers 4 --num_epoch 2 --num_batches_per_epoch 100 
    '''
    args = parser.parse_args()
    world_size = int(os.environ.get("WORLD_SIZE", 1))
@@ -9,10 +9,10 @@ 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_group,
-                                      get_world_rank, get_world_size)
-from fastvideo.v1.logger import init_logger
+from fastvideo.dataset.utils import collate_latents_embs_masks
+from fastvideo.distributed import (get_sp_world_size, get_world_group,
+                                   get_world_rank, get_world_size)
+from fastvideo.logger import init_logger

 logger = init_logger(__name__)

@@ -13,10 +13,10 @@ import tqdm
 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_group,
-                                      get_world_rank, get_world_size)
-from fastvideo.v1.logger import init_logger
+from fastvideo.dataset.utils import collate_rows_from_parquet_schema
+from fastvideo.distributed import (get_sp_world_size, get_world_group,
+                                   get_world_rank, get_world_size)
+from fastvideo.logger import init_logger

 logger = init_logger(__name__)

@@ -16,7 +16,7 @@ from einops import rearrange
 from PIL import Image
 from transformers import AutoTokenizer

-from fastvideo.v1.logger import init_logger
+from fastvideo.logger import init_logger

 logger = init_logger(__name__)

@@ -6,10 +6,10 @@ import pathlib
 import datasets
 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
+from fastvideo.distributed import (get_sp_world_size, get_world_rank,
+                                   get_world_size)
+from fastvideo.logger import init_logger
+from fastvideo.models.vision_utils import load_image, load_video

 logger = init_logger(__name__)

--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
“BrianChen1129”	586d144a0f	update pre commit	2025-07-27 05:56:13 +00:00
“BrianChen1129”	3601e33c4a	update pre commit	2025-07-27 05:54:53 +00:00
“BrianChen1129”	6345b6bff7	add DMD training pipeline	2025-07-27 05:08:10 +00:00
William Lin	cb320dfa3a	[bugfix] VideoGenerator improperly extracts output_video_name (#649 )	2025-07-26 19:46:26 -07:00
Kevin Lin	d8bb5abc46	[CI] Fix ComfyUI publisher ID (#648 )	2025-07-25 19:02:52 -07:00
Kevin Lin	cc703eca51	[CI] Add publish workflow for ComfyUI (#647 )	2025-07-25 18:39:05 -07:00
William Lin	81c9df629c	[core] Add offloading for vae and image encoder and rename offloading args (#643 )	2025-07-25 17:55:03 -07:00
Yongqi Chen	d3c0c52208	[Feature] Add prompt_txt support for CLI inference; Add DMD CLI inference (#646 )	2025-07-25 19:44:10 -04:00
William Lin	744e0555c0	[misc] Use FASTVIDEO_STAGE_LOGGING for perf timing of stage (#644 )	2025-07-25 16:20:53 -07:00
Jinzhe Pan	3a38f7dfdc	[1/3][Preprocess] refactor preprocessing configs (#638 )	2025-07-25 14:27:12 -07:00
William Lin	f572319bd9	[Feature] Remove V1 folder (#642 )	2025-07-24 22:43:12 -07:00