Revert "[feat] Job Runner UI (#1172 )"

This reverts commit f72618dafb.
[feat] Job Runner UI (#1172 )
2026-03-29 16:42:04 +08:00 · 2026-03-29 16:17:56 +08:00 · 2026-03-28 10:28:50 +08:00 · 2026-03-26 16:38:31 -07:00 · 2026-03-26 08:07:11 +08:00 · 2026-03-24 16:49:51 -07:00
320 changed files with 9243 additions and 11132 deletions
@@ -3,19 +3,17 @@ env:
  BUILDKITE_CLEAN_CHECKOUT: true

 steps:
-  - label: "pre-commit"
-    command: ".buildkite/scripts/pre_commit.sh"
-    agents:
-      queue: "default"
-
-  - wait
-
-  - label: "Trigger Tests"
+  # ============================================================
+  # Fastcheck: Runs on every PR (~10-15 min parallel)
+  # Core component validation: encoders, VAEs, transformers,
+  # CUDA kernels, and unit tests.
+  # ============================================================
+  - label: "Trigger Fastcheck"
    plugins:
      - monorepo-diff#v1.4.0:
          diff: 'git fetch origin "$BUILDKITE_PULL_REQUEST_BASE_BRANCH" && git diff --name-only origin/"$BUILDKITE_PULL_REQUEST_BASE_BRANCH"...HEAD'
          watch:
-            - path: 
+            - path:
                - "fastvideo/models/encoders/**"
                - "fastvideo/models/loader/**"
                - "fastvideo/tests/encoders/**"
@@ -23,7 +21,7 @@ steps:
                - "docker/Dockerfile.python3.12"
              config:
                command: "timeout 20m .buildkite/scripts/pr_test.sh"
-                label: "Encoder Tests"
+                label: ":microscope: Encoder Tests"
                env:
                  - TEST_TYPE=encoder
                agents:
@@ -36,7 +34,7 @@ steps:
                - "docker/Dockerfile.python3.12"
              config:
                command: "timeout 20m .buildkite/scripts/pr_test.sh"
-                label: "VAE Tests"
+                label: ":microscope: VAE Tests"
                env:
                  - TEST_TYPE=vae
                agents:
@@ -51,18 +49,55 @@ steps:
                - "docker/Dockerfile.python3.12"
              config:
                command: "timeout 15m .buildkite/scripts/pr_test.sh"
-                label: "Transformer Tests"
+                label: ":microscope: Transformer Tests"
                env:
                  - TEST_TYPE=transformer
                agents:
                  queue: "default"
-            - path: 
+            - path:
+                - "fastvideo-kernel/**"
+                - "pyproject.toml"
+                - "docker/Dockerfile.python3.12"
+              config:
+                command: "timeout 15m .buildkite/scripts/pr_test.sh"
+                label: ":microscope: Kernel Tests"
+                env:
+                  - TEST_TYPE=kernel_tests
+                agents:
+                  queue: "default"
+            - path:
+                - "fastvideo/**"
+                - ".buildkite/**"
+                - ".github/**"
+                - "pyproject.toml"
+                - "docker/Dockerfile.python3.12"
+              config:
+                command: "timeout 15m .buildkite/scripts/pr_test.sh"
+                label: ":microscope: Unit Tests"
+                env:
+                  - TEST_TYPE=unit_test
+                agents:
+                  queue: "default"
+
+  # ============================================================
+  # Full Suite: Runs only when FULL_SUITE=true
+  # Triggered by adding the 'ready' label to a PR (via GitHub
+  # Actions → Buildkite API). Includes integration tests, SSIM
+  # regression, training pipelines, and performance benchmarks.
+  # ============================================================
+  - label: "Trigger Full Suite"
+    if: build.env("FULL_SUITE") == "true"
+    plugins:
+      - monorepo-diff#v1.4.0:
+          diff: 'git fetch origin "$BUILDKITE_PULL_REQUEST_BASE_BRANCH" && git diff --name-only origin/"$BUILDKITE_PULL_REQUEST_BASE_BRANCH"...HEAD'
+          watch:
+            - path:
                - "fastvideo/**/*.py"
                - "pyproject.toml"
                - "docker/Dockerfile.python3.12"
              config:
                command: "timeout 90m .buildkite/scripts/pr_test.sh"
-                label: "SSIM Tests"
+                label: ":bar_chart: SSIM Tests"
                env:
                  - TEST_TYPE=ssim
                agents:
@@ -77,7 +112,7 @@ steps:
                - "docker/Dockerfile.python3.12"
              config:
                command: "timeout 20m .buildkite/scripts/pr_test.sh"
-                label: "LoRA Inference Tests"
+                label: ":test_tube: LoRA Inference Tests"
                env:
                  - TEST_TYPE=inference_lora
                agents:
@@ -88,7 +123,7 @@ steps:
                - "docker/Dockerfile.python3.12"
              config:
                command: "timeout 15m .buildkite/scripts/pr_test.sh"
-                label: "Training Tests"
+                label: ":test_tube: Training Tests"
                env:
                  - TEST_TYPE=training
                agents:
@@ -99,7 +134,7 @@ steps:
                - "docker/Dockerfile.python3.12"
              config:
                command: "timeout 15m .buildkite/scripts/pr_test.sh"
-                label: "Distillation DMDTests"
+                label: ":test_tube: Distillation DMD Tests"
                env:
                  - TEST_TYPE=distillation_dmd
                agents:
@@ -111,7 +146,7 @@ steps:
                - "docker/Dockerfile.python3.12"
              config:
                command: "timeout 15m .buildkite/scripts/pr_test.sh"
-                label: "Self-Forcing Tests"
+                label: ":test_tube: Self-Forcing Tests"
                env:
                  - TEST_TYPE=self_forcing
                agents:
@@ -122,7 +157,7 @@ steps:
                - "docker/Dockerfile.python3.12"
              config:
                command: "timeout 15m .buildkite/scripts/pr_test.sh"
-                label: "LoRA Training Tests"
+                label: ":test_tube: LoRA Training Tests"
                env:
                  - TEST_TYPE=training_lora
                agents:
@@ -134,20 +169,11 @@ steps:
                - "docker/Dockerfile.python3.12"
              config:
                command: "timeout 15m .buildkite/scripts/pr_test.sh"
-                label: "Training Tests VSA"
+                label: ":test_tube: Training Tests VSA"
                env:
                  - TEST_TYPE=training_vsa
                agents:
                  queue: "default"
-            - path:
-                - "fastvideo-kernel/**"
-                - "pyproject.toml"
-                - "docker/Dockerfile.python3.12"
-              config:
-                command: "timeout 15m .buildkite/scripts/pr_test.sh"
-                label: "Kernel Tests"
-                env:
-                  - TEST_TYPE=kernel_tests
            - path:
                - "fastvideo-kernel/**"
                - "fastvideo/attention/backends/vmoba.py"
@@ -155,20 +181,9 @@ steps:
                - "docker/Dockerfile.python3.12"
              config:
                command: "timeout 15m .buildkite/scripts/pr_test.sh"
-                label: "Inference Tests VMoBA"
-                env:  
-                  - TEST_TYPE=inference_vmoba
-                agents:
-                  queue: "default"
-            - path:
-                - "fastvideo/**"
-                - "pyproject.toml"
-                - "docker/Dockerfile.python3.12"
-              config:
-                command: "timeout 15m .buildkite/scripts/pr_test.sh"
-                label: "Unit Tests"
+                label: ":test_tube: Inference Tests VMoBA"
                env:
-                  - TEST_TYPE=unit_test
+                  - TEST_TYPE=inference_vmoba
                agents:
                  queue: "default"
            - path:
@@ -184,7 +199,7 @@ steps:
                - "docker/Dockerfile.python3.12"
              config:
                command: "timeout 30m .buildkite/scripts/pr_test.sh"
-                label: "Performance Tests"
+                label: ":test_tube: Performance Tests"
                env:
                  - TEST_TYPE=performance
                agents:
@@ -197,19 +212,8 @@ steps:
                - "docker/Dockerfile.python3.12"
              config:
                command: "timeout 30m .buildkite/scripts/pr_test.sh"
-                label: "API Server Tests"
+                label: ":test_tube: API Server Tests"
                env:
                  - TEST_TYPE=api_server
                agents:
                  queue: "default"
-            # - path:
-            #     - "scripts/lora_extraction/**"
-            #     - "pyproject.toml"
-            #     - "docker/Dockerfile.python3.12"
-            #   config:
-            #     command: "timeout 90m .buildkite/scripts/pr_test.sh"
-            #     label: "LoRA Extraction Tests"
-            #     env:
-            #       - TEST_TYPE=lora_extraction
-            #     agents:
-            #       queue: "default"
@@ -0,0 +1,255 @@
+pull_request_rules:
+
+  # ============================================================
+  # Auto-labeling by file path
+  # ============================================================
+
+  - name: label attention changes
+    conditions:
+      - or:
+        - files~=^fastvideo/attention/
+        - files~=^csrc/attn/
+      - -closed
+      - label!=stale
+    actions:
+      label:
+        add: [attention]
+
+  - name: label encoder changes
+    conditions:
+      - files~=^fastvideo/models/encoders/
+      - -closed
+      - label!=stale
+    actions:
+      label:
+        add: [encoder]
+
+  - name: label vae changes
+    conditions:
+      - files~=^fastvideo/models/vaes/
+      - -closed
+      - label!=stale
+    actions:
+      label:
+        add: [vae]
+
+  - name: label transformer/DiT changes
+    conditions:
+      - files~=^fastvideo/models/dits/
+      - -closed
+      - label!=stale
+    actions:
+      label:
+        add: [transformer]
+
+  - name: label training changes
+    conditions:
+      - files~=^fastvideo/training/
+      - -closed
+      - label!=stale
+    actions:
+      label:
+        add: [training]
+
+  - name: label pipeline changes
+    conditions:
+      - files~=^fastvideo/pipelines/
+      - -closed
+      - label!=stale
+    actions:
+      label:
+        add: [pipeline]
+
+  - name: label distributed changes
+    conditions:
+      - files~=^fastvideo/distributed/
+      - -closed
+      - label!=stale
+    actions:
+      label:
+        add: [distributed]
+
+  - name: label kernel changes
+    conditions:
+      - or:
+        - files~=^csrc/
+        - files~=^fastvideo-kernel/
+      - -closed
+      - label!=stale
+    actions:
+      label:
+        add: [kernel]
+
+  - name: label data/dataset changes
+    conditions:
+      - or:
+        - files~=^fastvideo/dataset/
+        - files~=^fastvideo/data_preprocess/
+      - -closed
+      - label!=stale
+    actions:
+      label:
+        add: [data]
+
+  - name: label CLI/entrypoint changes
+    conditions:
+      - files~=^fastvideo/entrypoints/
+      - -closed
+      - label!=stale
+    actions:
+      label:
+        add: [cli]
+
+  - name: label documentation changes
+    conditions:
+      - or:
+        - files~=^docs/
+        - files~=\.md$
+      - -closed
+      - label!=stale
+    actions:
+      label:
+        add: [documentation]
+
+  - name: label CI changes
+    conditions:
+      - or:
+        - files~=^\.github/
+        - files~=^\.buildkite/
+        - files~=^docker/
+      - -closed
+      - label!=stale
+    actions:
+      label:
+        add: [ci]
+
+  - name: label test changes
+    conditions:
+      - or:
+        - files~=^fastvideo/tests/
+        - files~=^tests/
+      - -closed
+      - label!=stale
+    actions:
+      label:
+        add: [tests]
+
+  - name: label ComfyUI changes
+    conditions:
+      - files~=^comfyui/
+      - -closed
+      - label!=stale
+    actions:
+      label:
+        add: [comfyui]
+
+  - name: label layer changes
+    conditions:
+      - files~=^fastvideo/layers/
+      - -closed
+      - label!=stale
+    actions:
+      label:
+        add: [layers]
+
+  - name: label LoRA changes
+    conditions:
+      - or:
+        - files~=^fastvideo/layers/lora/
+        - files~=^fastvideo/tests/lora/
+        - title~=(?i)lora
+      - -closed
+      - label!=stale
+    actions:
+      label:
+        add: [lora]
+
+  # ============================================================
+  # Pre-commit failure help comment
+  # ============================================================
+
+  - name: comment on pre-commit failure
+    conditions:
+      - check-failure~=pre-commit
+      - -closed
+    actions:
+      comment:
+        message: |
+          ## Pre-commit checks failed
+
+          Hi @{{author}}, the pre-commit checks have failed. To fix them locally:
+
+          ```bash
+          # Install pre-commit if you haven't already
+          uv pip install pre-commit
+          pre-commit install
+
+          # Run all checks and auto-fix what's possible
+          pre-commit run --all-files
+          ```
+
+          Common fixes:
+          - **yapf**: `yapf -i <file>` (formatting)
+          - **ruff**: `ruff check --fix <file>` (linting)
+          - **codespell**: `codespell --write-changes <file>` (spelling)
+
+          After fixing, commit and push the changes. The checks will re-run automatically.
+
+          For future commits, `pre-commit` will run automatically on changed files before each commit.
+
+  # ============================================================
+  # Merge conflict detection
+  # ============================================================
+
+  - name: label conflicting PRs
+    conditions:
+      - conflict
+      - -closed
+      - label!=stale
+    actions:
+      label:
+        add: [needs-rebase]
+      comment:
+        message: |
+          This PR has merge conflicts with the base branch. Please rebase:
+
+          ```bash
+          git fetch origin main
+          git rebase origin/main
+          # Resolve any conflicts, then:
+          git push --force-with-lease
+          ```
+
+  - name: remove conflict label when resolved
+    conditions:
+      - -conflict
+      - -closed
+      - label=needs-rebase
+    actions:
+      label:
+        remove: [needs-rebase]
+
+  # ============================================================
+  # Auto-rebase stale PRs
+  # ============================================================
+
+  - name: auto-merge when ready and approved
+    conditions:
+      - label=ready
+      - "#approved-reviews-by>=1"
+      - check-success~=pre-commit
+      - -conflict
+      - -closed
+      - -draft
+    actions:
+      merge:
+        method: squash
+
+  - name: auto-rebase on ready label
+    conditions:
+      - label=ready
+      - "#approved-reviews-by>=1"
+      - -conflict
+      - -closed
+    actions:
+      rebase: {}
@@ -36,11 +36,11 @@ jobs:
          
          if [ "$NEW_VERSION" != "$OLD_VERSION" ]; then
            echo "Version changed from $OLD_VERSION to $NEW_VERSION"
-            echo "changed=true" >> $GITHUB_OUTPUT
-            echo "new-version=$NEW_VERSION" >> $GITHUB_OUTPUT
+            echo "changed=true" >> "$GITHUB_OUTPUT"
+            echo "new-version=$NEW_VERSION" >> "$GITHUB_OUTPUT"
          else
            echo "Version did not change"
-            echo "changed=false" >> $GITHUB_OUTPUT
+            echo "changed=false" >> "$GITHUB_OUTPUT"
          fi

  build_wheels:
@@ -0,0 +1,53 @@
+name: Auto-Label Issues
+
+on:
+  issues:
+    types: [opened, edited]
+
+permissions:
+  issues: write
+
+jobs:
+  label-issues:
+    if: github.repository == 'hao-ai-lab/FastVideo'
+    runs-on: ubuntu-latest
+    steps:
+      - name: Label by keywords
+        uses: actions/github-script@60a0d83039c74a4aee543508d2ffcb1c3799cdea # v7.0.1
+        with:
+          script: |
+            const title = context.payload.issue.title.toLowerCase();
+            const body = (context.payload.issue.body || '').toLowerCase();
+            const text = title + ' ' + body;
+            const labels = [];
+
+            const rules = [
+              { keywords: ['training', 'finetune', 'fine-tune', 'lora', 'fsdp', 'distill'], label: 'training' },
+              { keywords: ['inference', 'generate', 'pipeline', 'slow', 'latency'], label: 'inference' },
+              { keywords: ['attention', 'vsa', 'flash', 'sta', 'vmoba', 'sparse attn'], label: 'attention' },
+              { keywords: ['install', 'setup', 'pip', 'cuda', 'uv ', 'import error', 'modulenotfound'], label: 'installation' },
+              { keywords: ['docs', 'documentation', 'tutorial', 'example'], label: 'documentation' },
+              { keywords: ['memory', 'oom', 'out of memory', 'gpu memory', 'vram'], label: 'performance' },
+              { keywords: ['wan', 'hunyuan', 'mochi', 'ltx', 'cogvideo'], label: 'model-specific' },
+              { keywords: ['windows', 'macos', 'mac os', 'apple', 'mps'], label: 'platform' },
+              { keywords: ['comfyui', 'comfy'], label: 'comfyui' },
+              { keywords: ['kernel', 'csrc', 'cuda kernel', 'thunderkittens'], label: 'kernel' },
+            ];
+
+            for (const rule of rules) {
+              if (rule.keywords.some(kw => text.includes(kw))) {
+                labels.push(rule.label);
+              }
+            }
+
+            if (labels.length > 0) {
+              await github.rest.issues.addLabels({
+                owner: context.repo.owner,
+                repo: context.repo.repo,
+                issue_number: context.payload.issue.number,
+                labels: labels,
+              });
+              console.log(`Added labels: ${labels.join(', ')}`);
+            } else {
+              console.log('No keyword matches found');
+            }
@@ -0,0 +1,17 @@
+{
+  "problemMatcher": [
+    {
+      "owner": "ruff",
+      "pattern": [
+        {
+          "regexp": "^(.+):(\\d+):(\\d+): (\\w+) (.+)$",
+          "file": 1,
+          "line": 2,
+          "column": 3,
+          "code": 4,
+          "message": 5
+        }
+      ]
+    }
+  ]
+}
@@ -67,12 +67,8 @@ concurrency:
  cancel-in-progress: true

 jobs:
-  pre-commit:
-    uses: ./.github/workflows/pre-commit.yml
-
  change-filter:
    runs-on: ubuntu-latest
-    needs: pre-commit
    if: ${{ github.event.pull_request.draft == false || github.event_name == 'workflow_dispatch' }}
    outputs:
      encoder-test: ${{ steps.filter.outputs.encoder-test }}
@@ -1,7 +1,18 @@
 name: pre-commit

 on:
-    workflow_call:
+  pull_request:
+    branches: [main]
+  push:
+    branches: [main]
+  workflow_call:
+
+concurrency:
+  group: pre-commit-${{ github.ref }}
+  cancel-in-progress: ${{ github.event_name == 'pull_request' }}
+
+permissions:
+  contents: read

 jobs:
  pre-commit:
@@ -13,6 +24,7 @@ jobs:
        python-version: "3.12"
    - run: echo "::add-matcher::.github/workflows/matchers/actionlint.json"
    - run: echo "::add-matcher::.github/workflows/matchers/mypy.json"
+    - run: echo "::add-matcher::.github/workflows/matchers/ruff.json"
    - uses: pre-commit/action@v3.0.1
      with:
-        extra_args: --all-files --hook-stage manual
+        extra_args: --all-files --hook-stage manual
@@ -0,0 +1,51 @@
+name: Close Stale Issues and PRs
+
+on:
+  schedule:
+    # Daily at 1:30 AM UTC
+    - cron: '30 1 * * *'
+
+jobs:
+  stale:
+    if: github.repository == 'hao-ai-lab/FastVideo'
+    permissions:
+      issues: write
+      pull-requests: write
+      actions: write
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/stale@997185467fa4f803885201cee163a9f38240193d # v10.1.1
+        with:
+          operations-per-run: 500
+
+          exempt-draft-pr: true
+          exempt-issue-labels: 'keep-open,pinned,security,Bug,RFC'
+          exempt-pr-labels: 'keep-open,pinned'
+
+          labels-to-add-when-unstale: 'unstale'
+          labels-to-remove-when-stale: 'unstale'
+
+          days-before-issue-stale: 90
+          days-before-issue-close: 30
+          stale-issue-label: 'stale'
+          stale-issue-message: >
+            This issue has been automatically marked as stale because it has not
+            had any activity within 90 days. It will be automatically closed if
+            no further activity occurs within 30 days. Leave a comment if you
+            feel this issue should remain open. Thank you!
+          close-issue-message: >
+            This issue has been automatically closed due to inactivity. Please
+            feel free to reopen if you feel it is still relevant. Thank you!
+
+          days-before-pr-stale: 60
+          days-before-pr-close: 14
+          stale-pr-label: 'stale'
+          stale-pr-message: >
+            This pull request has been automatically marked as stale because it
+            has not had any activity within 60 days. It will be automatically
+            closed if no further activity occurs within 14 days. Leave a comment
+            if you feel this pull request should remain open. Thank you!
+          close-pr-message: >
+            This pull request has been automatically closed due to inactivity.
+            Please feel free to reopen if you intend to continue working on it.
+            Thank you!
@@ -0,0 +1,26 @@
+name: Trigger Full Test Suite
+
+on:
+  pull_request:
+    types: [labeled]
+
+jobs:
+  trigger-buildkite:
+    if: >-
+      github.repository == 'hao-ai-lab/FastVideo'
+      && github.event.label.name == 'ready'
+    runs-on: ubuntu-latest
+    permissions:
+      pull-requests: read
+    steps:
+      - name: Trigger Buildkite build with FULL_SUITE
+        uses: buildkite/trigger-pipeline-action@v2.4.1
+        with:
+          buildkite_api_access_token: ${{ secrets.BUILDKITE_API_TOKEN }}
+          pipeline: "${{ vars.BUILDKITE_ORG_SLUG }}/${{ vars.BUILDKITE_PIPELINE_SLUG }}"
+          branch: "${{ github.event.pull_request.head.ref }}"
+          commit: "${{ github.event.pull_request.head.sha }}"
+          message: "Full suite for PR #${{ github.event.pull_request.number }}"
+          send_pull_request: "true"
+          pull_request_base_branch: "${{ github.event.pull_request.base.ref }}"
+          build_env_vars: '{"FULL_SUITE": "true"}'
@@ -0,0 +1,56 @@
+name: Welcome First-Time Contributors
+
+on:
+  issues:
+    types: [opened]
+  pull_request_target:
+    types: [opened]
+
+permissions:
+  issues: write
+  pull-requests: write
+
+jobs:
+  welcome:
+    if: github.repository == 'hao-ai-lab/FastVideo'
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/first-interaction@34f15e814fe48ac9312ccf29db4e74fa767cbab7 # v1.3.0
+        with:
+          repo-token: ${{ secrets.GITHUB_TOKEN }}
+          issue-message: |
+            Welcome to FastVideo! Thanks for opening your first issue.
+
+            To help us investigate, please include:
+            - **FastVideo version**: `pip show fastvideo`
+            - **GPU**: `nvidia-smi` output (GPU model, driver, CUDA version)
+            - **Python version**: `python --version`
+            - **OS**: e.g., Ubuntu 22.04
+
+            If this is a bug, a minimal reproduction script helps us fix it faster.
+
+            Useful links:
+            - [Documentation](https://hao-ai-lab.github.io/FastVideo)
+            - [Contributing Guide](https://hao-ai-lab.github.io/FastVideo/contributing/overview/)
+            - [Slack](https://join.slack.com/t/fastvideo/shared_invite/zt-3f4lao1uq-u~Ipx6Lt4J27AlD2y~IdLQ)
+          pr-message: |
+            Welcome to FastVideo! Thanks for your first pull request.
+
+            **How our CI works:**
+
+            PRs run a two-tier CI system:
+            1. **Pre-commit** — formatting (yapf), linting (ruff), type checking (mypy). Runs immediately on every PR.
+            2. **Fastcheck** — core GPU tests (encoders, VAEs, transformers, kernels, unit tests). Runs automatically via Buildkite on relevant file changes (~10-15 min).
+            3. **Full Suite** — integration tests, training pipelines, SSIM regression. Runs only when a reviewer adds the `ready` label.
+
+            **Before your PR is reviewed:**
+            - [ ] `pre-commit run --all-files` passes locally
+            - [ ] You've added or updated tests for your changes
+            - [ ] The PR description explains what and why
+
+            If pre-commit fails, a bot comment will explain how to fix it. Fastcheck and Full Suite results appear in the Checks section below.
+
+            **Useful links:**
+            - [Contributing Guide](https://hao-ai-lab.github.io/FastVideo/contributing/overview/)
+            - [Development Roadmap](https://github.com/hao-ai-lab/FastVideo/issues/899)
+            - [Slack](https://join.slack.com/t/fastvideo/shared_invite/zt-3f4lao1uq-u~Ipx6Lt4J27AlD2y~IdLQ)
@@ -33,6 +33,7 @@ env
 **.txt
 *.log
 weights/
+logs/

 # SSIM test outputs
 fastvideo/tests/ssim/generated_videos/
@@ -6,16 +6,17 @@
     | <a href="https://hao-ai-lab.github.io/FastVideo"><b>Documentation</b></a> | <a href="https://hao-ai-lab.github.io/FastVideo/inference/inference_quick_start/"><b> Quick Start</b></a> | <a href="https://github.com/hao-ai-lab/FastVideo/discussions/982"  target="_blank"><b>Weekly Dev Meeting</b></a>  | 🟣💬 <a href="https://join.slack.com/t/fastvideo/shared_invite/zt-3f4lao1uq-u~Ipx6Lt4J27AlD2y~IdLQ" target="_blank"> <b>Slack</b> </a> |  🟣💬 <a href="https://github.com/hao-ai-lab/FastVideo/discussions/1097" target="_blank"> <b> WeChat </b> </a> |
 </p>

-**FastVideo is a unified post-training and inference framework for accelerated video generation.**
+**FastVideo is a unified post-training and real-time inference framework for accelerated video generation.**

 ## NEWS
-
- `2025/11/19`: Release [CausalWan2.2 I2V A14B Preview](https://huggingface.co/FastVideo/CausalWan2.2-I2V-A14B-Preview-Diffusers) models, [Blog](https://hao-ai-lab.github.io/blogs/fastvideo_causalwan_preview/) and [Inference Code!](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_self_forcing_causal_wan2_2_i2v.py)
+- `2026/03/17`: Release Live demo: [Into the Dreamverse: Vibe Directing in FastVideo](https://dreamverse.fastvideo.org/), check out the [Blog](https://haoailab.com/blogs/dreamverse/).
+- `2026/03/13`: Release Live demo: [Create a 5s 1080p Video in 4.5s with FastVideo on a Single GPU](https://1080p.fastvideo.org/), check out the [Blog](https://haoailab.com/blogs/fastvideo_realtime_1080p/).
+- `2025/11/19`: Release [CausalWan2.2 I2V A14B Preview](https://huggingface.co/FastVideo/CausalWan2.2-I2V-A14B-Preview-Diffusers) models, [Blog](https://hao-ai-lab.github.io/blogs/fastvideo_causalwan_preview/) and [Inference Code!](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_self_forcing_causal_wan2_2_i2v.py).
 - `2025/08/04`: Release [FastWan](https://hao-ai-lab.github.io/FastVideo/distillation/dmd) models and [Sparse-Distillation](https://hao-ai-lab.github.io/blogs/fastvideo_post_training/).

 ### More News

- `2025/06/14`: Release finetuning and inference code for [VSA](https://arxiv.org/pdf/2505.13389)
+- `2025/06/14`: Release finetuning and inference code for [VSA](https://arxiv.org/pdf/2505.13389).
 - `2025/04/24`: [FastVideo V1](https://hao-ai-lab.github.io/blogs/fastvideo/) is released!
 - `2025/02/18`: Release the inference code for [Sliding Tile Attention](https://hao-ai-lab.github.io/blogs/sta/).

@@ -1,15 +1,10 @@
 try:
-    from .comfyui.video_generator.nodes import (NODE_CLASS_MAPPINGS,
-                                                NODE_DISPLAY_NAME_MAPPINGS)
+    from .comfyui.video_generator.nodes import (NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS)
    WEB_DIRECTORY = "./web"
-    __all__ = [
-        'NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS', 'WEB_DIRECTORY'
-    ]
+    __all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS', 'WEB_DIRECTORY']
 except ImportError:
    # ComfyUI environment not available, skip comfyui imports
    NODE_CLASS_MAPPINGS = {}
    NODE_DISPLAY_NAME_MAPPINGS = {}
    WEB_DIRECTORY = "./web"
-    __all__ = [
-        'NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS', 'WEB_DIRECTORY'
-    ]
+    __all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS', 'WEB_DIRECTORY']
@@ -13,8 +13,7 @@ from .fvd import (
    compute_statistics,
    FVDConfig,
 )
-from .feature_extractors import (BaseFeatureExtractor, I3DFeatureExtractor,
-                                 load_extractor)
+from .feature_extractors import (BaseFeatureExtractor, I3DFeatureExtractor, load_extractor)
 from .video_utils import (
    load_video_auto,
    sample_clips_from_video,
@@ -5,18 +5,11 @@ from .fvd import compute_fvd_with_config, FVDConfig


 def main() -> int:
-    parser = argparse.ArgumentParser(
-        description='Compute Fréchet Video Distance (FVD)')
+    parser = argparse.ArgumentParser(description='Compute Fréchet Video Distance (FVD)')

    # Required arguments
-    parser.add_argument('--real-path',
-                        type=str,
-                        required=True,
-                        help='Path to real videos')
-    parser.add_argument('--gen-path',
-                        type=str,
-                        required=True,
-                        help='Path to generated videos')
+    parser.add_argument('--real-path', type=str, required=True, help='Path to real videos')
+    parser.add_argument('--gen-path', type=str, required=True, help='Path to generated videos')

    # Extractor selection
    parser.add_argument('--extractor',
@@ -26,42 +19,19 @@ def main() -> int:
                        help='Feature extractor model to use (default: i3d)')

    # Standard args
-    parser.add_argument('--seed',
-                        type=int,
-                        default=None,
-                        help='Random seed for reproducibility')
+    parser.add_argument('--seed', type=int, default=None, help='Random seed for reproducibility')
    parser.add_argument('--protocol',
                        type=str,
                        default=None,
                        choices=['fvd2048_16f', 'fvd2048_128f', 'quick_test'],
                        help='Use standard protocol (overrides other settings)')
-    parser.add_argument('--num-videos',
-                        type=int,
-                        default=2048,
-                        help='Number of videos to use')
-    parser.add_argument('--num-frames',
-                        type=int,
-                        default=16,
-                        help='Number of frames per clip')
-    parser.add_argument('--clip-strategy',
-                        type=str,
-                        default='beginning',
-                        help='Clip sampling strategy')
-    parser.add_argument('--batch-size',
-                        type=int,
-                        default=32,
-                        help='Batch size for feature extraction')
-    parser.add_argument('--device',
-                        type=str,
-                        default='cuda',
-                        help='Device to use (cuda or cpu)')
-    parser.add_argument('--cache-real-features',
-                        type=str,
-                        default=None,
-                        help='Path to cache real video features')
-    parser.add_argument('--quiet',
-                        action='store_true',
-                        help='Suppress progress output')
+    parser.add_argument('--num-videos', type=int, default=2048, help='Number of videos to use')
+    parser.add_argument('--num-frames', type=int, default=16, help='Number of frames per clip')
+    parser.add_argument('--clip-strategy', type=str, default='beginning', help='Clip sampling strategy')
+    parser.add_argument('--batch-size', type=int, default=32, help='Batch size for feature extraction')
+    parser.add_argument('--device', type=str, default='cuda', help='Device to use (cuda or cpu)')
+    parser.add_argument('--cache-real-features', type=str, default=None, help='Path to cache real video features')
+    parser.add_argument('--quiet', action='store_true', help='Suppress progress output')

    args = parser.parse_args()

@@ -22,8 +22,7 @@ class BaseFeatureExtractor(ABC, nn.Module):

    def __init__(self, device: str = 'cuda'):
        super().__init__()
-        self.device = torch.device(
-            device if torch.cuda.is_available() else 'cpu')
+        self.device = torch.device(device if torch.cuda.is_available() else 'cpu')

    @property
    @abstractmethod
@@ -53,10 +52,7 @@ class BaseFeatureExtractor(ABC, nn.Module):
        pass

    @torch.no_grad()
-    def extract_features(self,
-                         videos: torch.Tensor,
-                         batch_size: int = 32,
-                         verbose: bool = True) -> torch.Tensor:
+    def extract_features(self, videos: torch.Tensor, batch_size: int = 32, verbose: bool = True) -> torch.Tensor:
        """
        Extract features for a large tensor of videos by batching.
        """
@@ -65,9 +61,7 @@ class BaseFeatureExtractor(ABC, nn.Module):

        iterator = range(0, N, batch_size)
        if verbose:
-            iterator = tqdm(
-                iterator,
-                desc=f"Extracting features ({self.__class__.__name__})")
+            iterator = tqdm(iterator, desc=f"Extracting features ({self.__class__.__name__})")

        for i in iterator:
            batch = videos[i:i + batch_size].to(self.device)
@@ -95,9 +89,7 @@ class I3DFeatureExtractor(BaseFeatureExtractor):

    def _load_model(self) -> torch.nn.Module:
        try:
-            model_path = hf_hub_download(repo_id=self.REPO_ID,
-                                         filename=self.MODEL_FILENAME,
-                                         cache_dir=self.cache_dir)
+            model_path = hf_hub_download(repo_id=self.REPO_ID, filename=self.MODEL_FILENAME, cache_dir=self.cache_dir)
            return torch.jit.load(model_path, map_location=self.device)
        except Exception as e:
            raise RuntimeError(f"Failed to load I3D model: {e}") from e
@@ -119,10 +111,7 @@ class I3DFeatureExtractor(BaseFeatureExtractor):
        # Resize to 224x224
        if H != 224 or W != 224:
            videos = videos.reshape(B * T, C, H, W)
-            videos = F.interpolate(videos,
-                                   size=(224, 224),
-                                   mode='bilinear',
-                                   align_corners=False)
+            videos = F.interpolate(videos, size=(224, 224), mode='bilinear', align_corners=False)
            videos = videos.reshape(B, T, C, 224, 224)

        # [B, T, C, H, W] -> [B, C, T, H, W]
@@ -131,21 +120,15 @@ class I3DFeatureExtractor(BaseFeatureExtractor):
    def extract_features_batch(self, videos: torch.Tensor) -> torch.Tensor:
        batch = self.preprocess(videos)
        # TorchScript I3D returns raw logits when return_features=True
-        return self.model(batch,
-                          rescale=False,
-                          resize=False,
-                          return_features=True)
+        return self.model(batch, rescale=False, resize=False, return_features=True)


 # 2. CLIP Extractor (Semantic/Content Quality)
 class CLIPFeatureExtractor(BaseFeatureExtractor):

-    def __init__(self,
-                 device: str = 'cuda',
-                 model_name: str = "openai/clip-vit-base-patch32"):
+    def __init__(self, device: str = 'cuda', model_name: str = "openai/clip-vit-base-patch32"):
        if not TRANSFORMERS_AVAILABLE:
-            raise ImportError(
-                "Please install transformers: pip install transformers")
+            raise ImportError("Please install transformers: pip install transformers")
        super().__init__(device)
        self.processor = CLIPProcessor.from_pretrained(model_name)
        self.model = CLIPModel.from_pretrained(model_name).to(self.device)
@@ -172,9 +155,7 @@ class CLIPFeatureExtractor(BaseFeatureExtractor):
        images = videos.view(B * T, C, H, W)

        # HF Processor
-        inputs = self.processor(images=images,
-                                return_tensors="pt",
-                                padding=True)
+        inputs = self.processor(images=images, return_tensors="pt", padding=True)
        inputs = {k: v.to(self.device) for k, v in inputs.items()}

        # Extract features [B*T, Dim]
@@ -188,24 +169,15 @@ class CLIPFeatureExtractor(BaseFeatureExtractor):
 # 3. VideoMAE Extractor (Structure/Motion Quality)
 class VideoMAEFeatureExtractor(BaseFeatureExtractor):

-    def __init__(self,
-                 device: str = 'cuda',
-                 model_name: str = "MCG-NJU/videomae-base"):
+    def __init__(self, device: str = 'cuda', model_name: str = "MCG-NJU/videomae-base"):
        if not TRANSFORMERS_AVAILABLE:
-            raise ImportError(
-                "Please install transformers: pip install transformers")
+            raise ImportError("Please install transformers: pip install transformers")
        super().__init__(device)
        self.model = VideoMAEModel.from_pretrained(model_name).to(self.device)
        self.model.eval()

-        self.register_buffer(
-            'mean',
-            torch.tensor([0.485, 0.456, 0.406],
-                         device=self.device).view(1, 1, 3, 1, 1))
-        self.register_buffer(
-            'std',
-            torch.tensor([0.229, 0.224, 0.225],
-                         device=self.device).view(1, 1, 3, 1, 1))
+        self.register_buffer('mean', torch.tensor([0.485, 0.456, 0.406], device=self.device).view(1, 1, 3, 1, 1))
+        self.register_buffer('std', torch.tensor([0.229, 0.224, 0.225], device=self.device).view(1, 1, 3, 1, 1))

    @property
    def feature_dim(self) -> int:
@@ -221,10 +193,7 @@ class VideoMAEFeatureExtractor(BaseFeatureExtractor):
        # 1. Resize to 224x224
        if H != 224 or W != 224:
            videos = videos.view(B * T, C, H, W)
-            videos = F.interpolate(videos,
-                                   size=(224, 224),
-                                   mode='bilinear',
-                                   align_corners=False)
+            videos = F.interpolate(videos, size=(224, 224), mode='bilinear', align_corners=False)
            videos = videos.view(B, T, C, 224, 224)

        # 2. Normalize to [0, 1]
@@ -260,5 +229,4 @@ def load_extractor(name: str, device: str = 'cuda') -> BaseFeatureExtractor:
    elif name == 'videomae':
        return VideoMAEFeatureExtractor(device)
    else:
-        raise ValueError(
-            f"Unknown extractor: {name}. Options: i3d, clip, videomae")
+        raise ValueError(f"Unknown extractor: {name}. Options: i3d, clip, videomae")
@@ -36,8 +36,7 @@ def compute_frechet_distance(mu1: np.ndarray,

    if np.iscomplexobj(covmean):
        if not np.allclose(np.diagonal(covmean).imag, 0, atol=1e-3):
-            print(
-                f"Warning: Imaginary component: {np.max(np.abs(covmean.imag))}")
+            print(f"Warning: Imaginary component: {np.max(np.abs(covmean.imag))}")
        covmean = covmean.real

    trace_product = np.trace(covmean)
@@ -67,8 +66,7 @@ class FVDConfig:
    temporal_stride: int = 1  # For sliding window clips

    # Data processing
-    video_extensions: list[str] = field(
-        default_factory=lambda: ['.mp4', '.avi', '.mov', '.mkv'])
+    video_extensions: list[str] = field(default_factory=lambda: ['.mp4', '.avi', '.mov', '.mkv'])
    support_frame_dirs: bool = True

    # Computation
@@ -88,25 +86,17 @@ class FVDConfig:
    @classmethod
    def fvd2048_16f(cls) -> 'FVDConfig':
        """Standard FVD protocol: 2048 videos, 16 frames, beginning clip."""
-        return cls(num_videos=2048,
-                   num_frames_per_clip=16,
-                   clip_strategy='beginning',
-                   use_streaming=True)
+        return cls(num_videos=2048, num_frames_per_clip=16, clip_strategy='beginning', use_streaming=True)

    @classmethod
    def fvd2048_128f(cls) -> 'FVDConfig':
        """Long video protocol: 2048 videos, 128 frames."""
-        return cls(num_videos=2048,
-                   num_frames_per_clip=128,
-                   clip_strategy='beginning',
-                   use_streaming=True)
+        return cls(num_videos=2048, num_frames_per_clip=128, clip_strategy='beginning', use_streaming=True)

    @classmethod
    def quick_test(cls) -> 'FVDConfig':
        """Quick test config: 100 videos, 16 frames."""
-        return cls(num_videos=100,
-                   num_frames_per_clip=16,
-                   clip_strategy='beginning')
+        return cls(num_videos=100, num_frames_per_clip=16, clip_strategy='beginning')

    def to_dict(self) -> dict:
        """Export config to dict for logging"""
@@ -196,14 +186,10 @@ def load_or_compute_features(videos: str | Path | torch.Tensor,
            max_features = config.num_videos * config.num_clips_per_video

            if len(features) < max_features:
-                print(
-                    f"WARNING: Cache has {len(features)} features but need {max_features}"
-                )
+                print(f"WARNING: Cache has {len(features)} features but need {max_features}")
                print("Cached features insufficient - will recompute...")
            elif len(features) > max_features:
-                print(
-                    f"Using {max_features} features from cache (truncated from {len(features)})"
-                )
+                print(f"Using {max_features} features from cache (truncated from {len(features)})")
                features = features[:max_features]
                return features
            else:
@@ -215,17 +201,16 @@ def load_or_compute_features(videos: str | Path | torch.Tensor,
    if isinstance(videos, (str | Path)):
        target_size = (224, 224) if config.resize_before_extraction else None

-        video_generator = load_video_clips_streaming(
-            videos,
-            num_frames=config.num_frames_per_clip,
-            max_videos=config.num_videos,
-            clip_strategy=config.clip_strategy,
-            frame_stride=config.frame_stride,
-            num_clips_per_video=config.num_clips_per_video,
-            video_extensions=config.video_extensions,
-            support_frame_dirs=config.support_frame_dirs,
-            target_size=target_size,
-            verbose=True)
+        video_generator = load_video_clips_streaming(videos,
+                                                     num_frames=config.num_frames_per_clip,
+                                                     max_videos=config.num_videos,
+                                                     clip_strategy=config.clip_strategy,
+                                                     frame_stride=config.frame_stride,
+                                                     num_clips_per_video=config.num_clips_per_video,
+                                                     video_extensions=config.video_extensions,
+                                                     support_frame_dirs=config.support_frame_dirs,
+                                                     target_size=target_size,
+                                                     verbose=True)

        max_clips = config.num_videos * config.num_clips_per_video
        features = extract_features_streaming(video_generator,
@@ -235,17 +220,14 @@ def load_or_compute_features(videos: str | Path | torch.Tensor,
                                              verbose=True)
    else:
        print(f"Extracting features from {len(videos)} video tensors...")
-        features = extractor.extract_features(videos,
-                                              batch_size=config.batch_size,
-                                              verbose=True)
+        features = extractor.extract_features(videos, batch_size=config.batch_size, verbose=True)
        features = features.numpy()

    # Validate feature count
    expected_count = config.num_videos * config.num_clips_per_video
    if len(features) < expected_count:
-        raise ValueError(
-            f"ERROR: Only extracted {len(features)} features, but need {expected_count}!\n"
-            f"Found fewer videos than expected. Check your video directory.")
+        raise ValueError(f"ERROR: Only extracted {len(features)} features, but need {expected_count}!\n"
+                         f"Found fewer videos than expected. Check your video directory.")
    elif len(features) > expected_count:
        print(f"Truncating {len(features)} features to {expected_count}")
        features = features[:expected_count]
@@ -312,9 +294,7 @@ def compute_fvd_with_config(real_videos: str | Path | torch.Tensor,

    # Initialize Extractor using Factory
    if verbose:
-        print(
-            f"\nInitializing {config.extractor_model.upper()} model on {config.device}..."
-        )
+        print(f"\nInitializing {config.extractor_model.upper()} model on {config.device}...")

    extractor = load_extractor(config.extractor_model, device=config.device)

@@ -324,12 +304,11 @@ def compute_fvd_with_config(real_videos: str | Path | torch.Tensor,
        print("Extracting REAL video features...")
        print(f"{'='*70}")

-    real_features = load_or_compute_features(
-        videos=real_videos,
-        extractor=extractor,
-        config=config,
-        cache_path=config.cache_real_features,
-        cache_name="real_features")
+    real_features = load_or_compute_features(videos=real_videos,
+                                             extractor=extractor,
+                                             config=config,
+                                             cache_path=config.cache_real_features,
+                                             cache_name="real_features")

    if verbose:
        print(f"\n{'='*70}")
@@ -19,16 +19,12 @@ class I3DFeatureExtractor(nn.Module):
    REPO_ID = 'flateon/FVD-I3D-torchscript'
    MODEL_FILENAME = 'i3d_torchscript.pt'

-    def __init__(self,
-                 device: str = 'cuda',
-                 cache_dir: str | Path | None = None):
+    def __init__(self, device: str = 'cuda', cache_dir: str | Path | None = None):
        super().__init__()

        self.device_str = device
        if device == 'cuda' and not torch.cuda.is_available():
-            print(
-                "Warning: CUDA requested but not available – falling back to CPU"
-            )
+            print("Warning: CUDA requested but not available – falling back to CPU")
            self.device = torch.device('cpu')
        else:
            self.device = torch.device(device)
@@ -51,9 +47,7 @@ class I3DFeatureExtractor(nn.Module):

        try:
            # Download model from Hugging Face Hub
-            model_path = hf_hub_download(repo_id=self.REPO_ID,
-                                         filename=self.MODEL_FILENAME,
-                                         cache_dir=self.cache_dir)
+            model_path = hf_hub_download(repo_id=self.REPO_ID, filename=self.MODEL_FILENAME, cache_dir=self.cache_dir)

            # Load directly to chosen device
            model = torch.jit.load(model_path, map_location=self.device)
@@ -61,10 +55,9 @@ class I3DFeatureExtractor(nn.Module):
            return model

        except Exception as e:
-            raise RuntimeError(
-                f"Failed to load I3D model from Hugging Face Hub. Error: {e}\n"
-                f"Ensure you have internet connection and huggingface_hub installed:\n"
-                f"pip install huggingface_hub") from e
+            raise RuntimeError(f"Failed to load I3D model from Hugging Face Hub. Error: {e}\n"
+                               f"Ensure you have internet connection and huggingface_hub installed:\n"
+                               f"pip install huggingface_hub") from e

    def preprocess(self, videos: torch.Tensor) -> torch.Tensor:
        """
@@ -88,10 +81,7 @@ class I3DFeatureExtractor(nn.Module):
        # Resize to 224x224 if needed
        if H != 224 or W != 224:
            videos = videos.reshape(B * T, C, H, W)
-            videos = F.interpolate(videos,
-                                   size=(224, 224),
-                                   mode='bilinear',
-                                   align_corners=False)
+            videos = F.interpolate(videos, size=(224, 224), mode='bilinear', align_corners=False)
            videos = videos.reshape(B, T, C, 224, 224)

        # Convert to [B, C, T, H, W] format
@@ -100,10 +90,7 @@ class I3DFeatureExtractor(nn.Module):
        return videos

    @torch.no_grad()
-    def extract_features(self,
-                         videos: torch.Tensor,
-                         batch_size: int = 32,
-                         verbose: bool = True) -> torch.Tensor:
+    def extract_features(self, videos: torch.Tensor, batch_size: int = 32, verbose: bool = True) -> torch.Tensor:
        """
        Extract I3D features
        
@@ -127,16 +114,11 @@ class I3DFeatureExtractor(nn.Module):
            batch = self.preprocess(batch)  # Now returns [B, C, T, H, W]

            # Use the HF model without rescale/resize (we handle it in preprocess)
-            features = self.model(batch,
-                                  rescale=False,
-                                  resize=False,
-                                  return_features=True)
+            features = self.model(batch, rescale=False, resize=False, return_features=True)

            all_features.append(features.cpu())

        return torch.cat(all_features, dim=0)

-    def __call__(self,
-                 videos: torch.Tensor,
-                 batch_size: int = 32) -> torch.Tensor:
+    def __call__(self, videos: torch.Tensor, batch_size: int = 32) -> torch.Tensor:
        return self.extract_features(videos, batch_size=batch_size)
@@ -36,10 +36,7 @@ def main() -> None:
                cache_real_features=str(script_dir / f'fvd-cache/{model_name}'),
            )

-            results = compute_fvd_with_config(real_dir,
-                                              gen_dir,
-                                              cfg,
-                                              verbose=False)
+            results = compute_fvd_with_config(real_dir, gen_dir, cfg, verbose=False)
            print(f"FVD: {results['fvd']}\nModel: {results['model']}")

        except Exception as e:
@@ -59,10 +59,7 @@ def validate_fvd(subset_a: Path, subset_b: Path, num_videos: int):
    print("TEST 1: Identity Test")
    print("=" * 70)

-    result1 = compute_fvd_with_config(real_videos=str(subset_a),
-                                      gen_videos=str(subset_a),
-                                      config=config,
-                                      verbose=False)
+    result1 = compute_fvd_with_config(real_videos=str(subset_a), gen_videos=str(subset_a), config=config, verbose=False)
    fvd_identity = result1['fvd']
    print(f"\nIdentity FVD: {fvd_identity:.2f}")

@@ -70,10 +67,7 @@ def validate_fvd(subset_a: Path, subset_b: Path, num_videos: int):
    print("TEST 2: Real vs Real")
    print("=" * 70)

-    result2 = compute_fvd_with_config(real_videos=str(subset_a),
-                                      gen_videos=str(subset_b),
-                                      config=config,
-                                      verbose=False)
+    result2 = compute_fvd_with_config(real_videos=str(subset_a), gen_videos=str(subset_b), config=config, verbose=False)
    fvd_real = result2['fvd']
    print(f"\nReal vs Real FVD: {fvd_real:.2f}")

@@ -87,9 +81,7 @@ def validate_fvd(subset_a: Path, subset_b: Path, num_videos: int):
 def main() -> None:
    bair_dir = Path('benchmarks/data/bair_full_videos')

-    subset_a, subset_b, count = split_videos(bair_dir,
-                                             n_per_subset=128,
-                                             seed=42)
+    subset_a, subset_b, count = split_videos(bair_dir, n_per_subset=128, seed=42)
    validate_fvd(subset_a, subset_b, count)


@@ -54,8 +54,7 @@ def _load_video_cv2(video_path: str | Path,
            raise RuntimeError(f"Video has 0 frames: {video_path}")

        frames = np.stack(frames)  # [T, H, W, C]
-        frames = torch.from_numpy(frames).permute(0, 3, 1,
-                                                  2).float()  # [T, C, H, W]
+        frames = torch.from_numpy(frames).permute(0, 3, 1, 2).float()  # [T, C, H, W]
        return frames

    if total_frames == 0:
@@ -63,14 +62,11 @@ def _load_video_cv2(video_path: str | Path,

    # Determine frame indices for sampling
    if total_frames < num_frames:
-        frame_indices = list(range(
-            total_frames)) + [total_frames - 1] * (num_frames - total_frames)
+        frame_indices = list(range(total_frames)) + [total_frames - 1] * (num_frames - total_frames)
    elif sample_strategy == 'uniform':
-        frame_indices = np.linspace(0, total_frames - 1, num_frames,
-                                    dtype=int).tolist()
+        frame_indices = np.linspace(0, total_frames - 1, num_frames, dtype=int).tolist()
    elif sample_strategy == 'random':
-        frame_indices = sorted(
-            np.random.choice(total_frames, num_frames, replace=False))
+        frame_indices = sorted(np.random.choice(total_frames, num_frames, replace=False))
    else:
        raise ValueError(f"Unknown sample_strategy: {sample_strategy}")

@@ -94,17 +90,15 @@ def _load_video_cv2(video_path: str | Path,
    cap.release()

    frames = np.stack(frames)  # [T, H, W, C]
-    frames = torch.from_numpy(frames).permute(0, 3, 1,
-                                              2).float()  # [T, C, H, W]
+    frames = torch.from_numpy(frames).permute(0, 3, 1, 2).float()  # [T, C, H, W]

    return frames


-def _load_video_from_frames(
-        frame_dir: str | Path,
-        num_frames: int | None = 16,
-        sample_strategy: str = 'uniform',
-        frame_extensions: list[str] | None = None) -> torch.Tensor:
+def _load_video_from_frames(frame_dir: str | Path,
+                            num_frames: int | None = 16,
+                            sample_strategy: str = 'uniform',
+                            frame_extensions: list[str] | None = None) -> torch.Tensor:
    """
    Load video from directory of frame images.
    
@@ -131,9 +125,7 @@ def _load_video_from_frames(
        frame_files.extend(frame_dir.glob(f"*{ext}"))

    if len(frame_files) == 0:
-        raise ValueError(
-            f"No frames found in {frame_dir} with extensions {frame_extensions}"
-        )
+        raise ValueError(f"No frames found in {frame_dir} with extensions {frame_extensions}")

    frame_files = sorted(frame_files, key=lambda x: x.name)
    total_frames = len(frame_files)
@@ -143,16 +135,11 @@ def _load_video_from_frames(
        frame_indices = list(range(total_frames))
    else:
        if total_frames < num_frames:
-            frame_indices = list(range(total_frames)) + [total_frames - 1] * (
-                num_frames - total_frames)
+            frame_indices = list(range(total_frames)) + [total_frames - 1] * (num_frames - total_frames)
        elif sample_strategy == 'uniform':
-            frame_indices = np.linspace(0,
-                                        total_frames - 1,
-                                        num_frames,
-                                        dtype=int).tolist()
+            frame_indices = np.linspace(0, total_frames - 1, num_frames, dtype=int).tolist()
        elif sample_strategy == 'random':
-            frame_indices = sorted(
-                np.random.choice(total_frames, num_frames, replace=False))
+            frame_indices = sorted(np.random.choice(total_frames, num_frames, replace=False))
        else:
            raise ValueError(f"Unknown sample_strategy: {sample_strategy}")

@@ -170,8 +157,7 @@ def _load_video_from_frames(

    # Stack and convert to tensor
    frames = np.stack(frames)  # [T, H, W, C]
-    frames = torch.from_numpy(frames).permute(0, 3, 1,
-                                              2).float()  # [T, C, H, W]
+    frames = torch.from_numpy(frames).permute(0, 3, 1, 2).float()  # [T, C, H, W]

    return frames

@@ -226,13 +212,12 @@ def load_video_auto(video_path: str | Path,
        raise ValueError(f"Unknown video format at {video_path}")


-def sample_clips_from_video(
-        video: torch.Tensor,
-        num_frames_per_clip: int = 16,
-        num_clips: int = 1,
-        strategy: str | ClipSamplingStrategy = ClipSamplingStrategy.BEGINNING,
-        frame_stride: int = 1,
-        temporal_stride: int = 1) -> list[torch.Tensor]:
+def sample_clips_from_video(video: torch.Tensor,
+                            num_frames_per_clip: int = 16,
+                            num_clips: int = 1,
+                            strategy: str | ClipSamplingStrategy = ClipSamplingStrategy.BEGINNING,
+                            frame_stride: int = 1,
+                            temporal_stride: int = 1) -> list[torch.Tensor]:
    """
    Sample clips from a video with various strategies.
        
@@ -296,10 +281,7 @@ def sample_clips_from_video(
    elif strategy == ClipSamplingStrategy.RANDOM:
        # Sample N random clips
        for _ in range(num_clips):
-            if effective_clip_length == T:
-                start = 0
-            else:
-                start = np.random.randint(0, T - effective_clip_length + 1)
+            start = 0 if effective_clip_length == T else np.random.randint(0, T - effective_clip_length + 1)
            clip = video[start:start + effective_clip_length]
            clips.append(clip)

@@ -312,8 +294,7 @@ def sample_clips_from_video(
            clips.append(clip)
        else:
            # Multiple uniformly spaced clips
-            step = (T - effective_clip_length) / (num_clips -
-                                                  1) if num_clips > 1 else 0
+            step = (T - effective_clip_length) / (num_clips - 1) if num_clips > 1 else 0
            for i in range(num_clips):
                start = int(i * step)
                start = min(start, T - effective_clip_length)
@@ -422,15 +403,12 @@ def load_video_clips_streaming(directory: str | Path,
    failed_count = 0
    total_clips = 0

-    iterator = tqdm(video_paths,
-                    desc="Loading videos") if verbose else video_paths
+    iterator = tqdm(video_paths, desc="Loading videos") if verbose else video_paths

    for video_path in iterator:
        try:
            # Load full video
-            video = load_video_auto(video_path,
-                                    num_frames=None,
-                                    sample_strategy='uniform')
+            video = load_video_auto(video_path, num_frames=None, sample_strategy='uniform')

            # Sample clips from video
            clips = sample_clips_from_video(video,
@@ -445,18 +423,14 @@ def load_video_clips_streaming(directory: str | Path,
                    T, C, H, W = clip.shape
                    if target_size != (H, W):
                        # Resize to target size
-                        clip = clip.contiguous(
-                        )  # Fix non-contiguous tensors first
-                        clip_flat = clip.view(T * C, H,
-                                              W).unsqueeze(0)  # [1, T*C, H, W]
-                        clip_resized = torch.nn.functional.interpolate(
-                            clip_flat,
-                            size=target_size,
-                            mode='bilinear',
-                            align_corners=False)
-                        clip = clip_resized.squeeze(0).view(
-                            T, C, target_size[0],
-                            target_size[1])  # Back to [T, C, H, W]
+                        clip = clip.contiguous()  # Fix non-contiguous tensors first
+                        clip_flat = clip.view(T * C, H, W).unsqueeze(0)  # [1, T*C, H, W]
+                        clip_resized = torch.nn.functional.interpolate(clip_flat,
+                                                                       size=target_size,
+                                                                       mode='bilinear',
+                                                                       align_corners=False)
+                        clip = clip_resized.squeeze(0).view(T, C, target_size[0],
+                                                            target_size[1])  # Back to [T, C, H, W]
                    resized_clips.append(clip)
                clips = resized_clips

@@ -480,11 +454,7 @@ def load_video_clips_streaming(directory: str | Path,

    failure_rate = failed_count / len(video_paths)
    if failure_rate > 0.1:  # More than 10% failed
-        print(
-            f"\nWARNING: {failure_rate:.1%} of videos failed to load ({failed_count}/{len(video_paths)})"
-        )
+        print(f"\nWARNING: {failure_rate:.1%} of videos failed to load ({failed_count}/{len(video_paths)})")

    if verbose:
-        print(
-            f"\nSuccessfully loaded {total_clips} clips from {len(video_paths) - failed_count} videos"
-        )
+        print(f"\nSuccessfully loaded {total_clips} clips from {len(video_paths) - failed_count} videos")
@@ -100,10 +100,7 @@ DEFAULT_PIP_PATTERNS = {
 def run(command):
    """Return (return-code, stdout, stderr)."""
    shell = True if type(command) is str else False
-    p = subprocess.Popen(command,
-                         stdout=subprocess.PIPE,
-                         stderr=subprocess.PIPE,
-                         shell=shell)
+    p = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=shell)
    raw_output, raw_err = p.communicate()
    rc = p.returncode
    if get_platform() == 'win32':
@@ -157,8 +154,7 @@ def get_conda_packages(run_lambda, patterns=None):
        return out

    return "\n".join(line for line in out.splitlines()
-                     if not line.startswith("#") and any(name in line
-                                                         for name in patterns))
+                     if not line.startswith("#") and any(name in line for name in patterns))


 def get_gcc_version(run_lambda):
@@ -166,27 +162,24 @@ def get_gcc_version(run_lambda):


 def get_clang_version(run_lambda):
-    return run_and_parse_first_match(run_lambda, 'clang --version',
-                                     r'clang version (.*)')
+    return run_and_parse_first_match(run_lambda, 'clang --version', r'clang version (.*)')


 def get_cmake_version(run_lambda):
-    return run_and_parse_first_match(run_lambda, 'cmake --version',
-                                     r'cmake (.*)')
+    return run_and_parse_first_match(run_lambda, 'cmake --version', r'cmake (.*)')


 def get_nvidia_driver_version(run_lambda):
    if get_platform() == 'darwin':
        cmd = 'kextstat | grep -i cuda'
-        return run_and_parse_first_match(run_lambda, cmd,
-                                         r'com[.]nvidia[.]CUDA [(](.*?)[)]')
+        return run_and_parse_first_match(run_lambda, cmd, r'com[.]nvidia[.]CUDA [(](.*?)[)]')
    smi = get_nvidia_smi()
    return run_and_parse_first_match(run_lambda, smi, r'Driver Version: (.*?) ')


 def get_gpu_info(run_lambda):
-    if get_platform() == 'darwin' or (TORCH_AVAILABLE and hasattr(
-            torch.version, 'hip') and torch.version.hip is not None):
+    if get_platform() == 'darwin' or (TORCH_AVAILABLE and hasattr(torch.version, 'hip')
+                                      and torch.version.hip is not None):
        if TORCH_AVAILABLE and torch.cuda.is_available():
            if torch.version.hip is not None:
                prop = torch.cuda.get_device_properties(0)
@@ -208,8 +201,7 @@ def get_gpu_info(run_lambda):


 def get_running_cuda_version(run_lambda):
-    return run_and_parse_first_match(run_lambda, 'nvcc --version',
-                                     r'release .+ V(.*)')
+    return run_and_parse_first_match(run_lambda, 'nvcc --version', r'release .+ V(.*)')


 def get_cudnn_version(run_lambda):
@@ -255,8 +247,7 @@ def get_nvidia_smi():
    if get_platform() == 'win32':
        system_root = os.environ.get('SYSTEMROOT', 'C:\\Windows')
        program_files_root = os.environ.get('PROGRAMFILES', 'C:\\Program Files')
-        legacy_path = os.path.join(program_files_root, 'NVIDIA Corporation',
-                                   'NVSMI', smi)
+        legacy_path = os.path.join(program_files_root, 'NVIDIA Corporation', 'NVSMI', smi)
        new_path = os.path.join(system_root, 'System32', smi)
        smis = [new_path, legacy_path]
        for candidate_smi in smis:
@@ -388,10 +379,8 @@ def get_cpu_info(run_lambda):
    if get_platform() == 'linux':
        rc, out, err = run_lambda('lscpu')
    elif get_platform() == 'win32':
-        rc, out, err = run_lambda(
-            'wmic cpu get Name,Manufacturer,Family,Architecture,ProcessorType,DeviceID, \
-        CurrentClockSpeed,MaxClockSpeed,L2CacheSize,L2CacheSpeed,Revision /VALUE'
-        )
+        rc, out, err = run_lambda('wmic cpu get Name,Manufacturer,Family,Architecture,ProcessorType,DeviceID, \
+        CurrentClockSpeed,MaxClockSpeed,L2CacheSize,L2CacheSpeed,Revision /VALUE')
    elif get_platform() == 'darwin':
        rc, out, err = run_lambda("sysctl -n machdep.cpu.brand_string")
    cpu_info = 'None'
@@ -416,27 +405,22 @@ def get_platform():


 def get_mac_version(run_lambda):
-    return run_and_parse_first_match(run_lambda, 'sw_vers -productVersion',
-                                     r'(.*)')
+    return run_and_parse_first_match(run_lambda, 'sw_vers -productVersion', r'(.*)')


 def get_windows_version(run_lambda):
    system_root = os.environ.get('SYSTEMROOT', 'C:\\Windows')
    wmic_cmd = os.path.join(system_root, 'System32', 'Wbem', 'wmic')
    findstr_cmd = os.path.join(system_root, 'System32', 'findstr')
-    return run_and_read_all(
-        run_lambda,
-        '{} os get Caption | {} /v Caption'.format(wmic_cmd, findstr_cmd))
+    return run_and_read_all(run_lambda, '{} os get Caption | {} /v Caption'.format(wmic_cmd, findstr_cmd))


 def get_lsb_version(run_lambda):
-    return run_and_parse_first_match(run_lambda, 'lsb_release -a',
-                                     r'Description:\t(.*)')
+    return run_and_parse_first_match(run_lambda, 'lsb_release -a', r'Description:\t(.*)')


 def check_release_file(run_lambda):
-    return run_and_parse_first_match(run_lambda, 'cat /etc/*-release',
-                                     r'PRETTY_NAME="(.*)"')
+    return run_and_parse_first_match(run_lambda, 'cat /etc/*-release', r'PRETTY_NAME="(.*)"')


 def get_os(run_lambda):
@@ -499,13 +483,10 @@ def get_pip_packages(run_lambda, patterns=None):
        elif shutil.which("uv") is not None:
            cmd = ["uv", "pip", "list", "--format=freeze"]
        else:
-            raise RuntimeError(
-                "Could not collect pip list output (pip or uv module not available)"
-            )
+            raise RuntimeError("Could not collect pip list output (pip or uv module not available)")

        out = run_and_read_all(run_lambda, cmd)
-        return "\n".join(line for line in out.splitlines()
-                         if any(name in line for name in patterns))
+        return "\n".join(line for line in out.splitlines() if any(name in line for name in patterns))

    pip_version = 'pip3' if sys.version[0] == '3' else 'pip'
    out = run_with_pip()
@@ -537,8 +518,7 @@ def is_xnnpack_available():
 def get_env_vars():
    env_vars = ''
    secret_terms = ('secret', 'token', 'api', 'access', 'password')
-    report_prefix = ("TORCH", "NCCL", "PYTORCH", "CUDA", "CUBLAS", "CUDNN",
-                     "OMP_", "MKL_", "NVIDIA")
+    report_prefix = ("TORCH", "NCCL", "PYTORCH", "CUDA", "CUBLAS", "CUDNN", "OMP_", "MKL_", "NVIDIA")
    for k, v in os.environ.items():
        if any(term in k.lower() for term in secret_terms):
            continue
@@ -559,8 +539,7 @@ def get_env_info():
        debug_mode_str = str(torch.version.debug)
        cuda_available_str = str(torch.cuda.is_available())
        cuda_version_str = torch.version.cuda
-        if not hasattr(torch.version,
-                       'hip') or torch.version.hip is None:  # cuda version
+        if not hasattr(torch.version, 'hip') or torch.version.hip is None:  # cuda version
            hip_compiled_version = hip_runtime_version = miopen_runtime_version = 'N/A'
        else:  # HIP version

@@ -588,9 +567,8 @@ def get_env_info():
    return SystemEnv(
        torch_version=version_str,
        is_debug_build=debug_mode_str,
-        python_version='{} ({}-bit runtime)'.format(
-            sys_version,
-            sys.maxsize.bit_length() + 1),
+        python_version='{} ({}-bit runtime)'.format(sys_version,
+                                                    sys.maxsize.bit_length() + 1),
        python_platform=get_python_platform(),
        is_cuda_available=cuda_available_str,
        cuda_compiled_version=cuda_version_str,
@@ -714,10 +692,8 @@ def pretty_str(envinfo):
        'nvidia_driver_version',
    ]
    all_cuda_fields = dynamic_cuda_fields + ['cudnn_version']
-    all_dynamic_cuda_fields_missing = all(mutable_dict[field] is None
-                                          for field in dynamic_cuda_fields)
-    if TORCH_AVAILABLE and not torch.cuda.is_available(
-    ) and all_dynamic_cuda_fields_missing:
+    all_dynamic_cuda_fields_missing = all(mutable_dict[field] is None for field in dynamic_cuda_fields)
+    if TORCH_AVAILABLE and not torch.cuda.is_available() and all_dynamic_cuda_fields_missing:
        for field in all_cuda_fields:
            mutable_dict[field] = 'No CUDA'
        if envinfo.cuda_compiled_version is None:
@@ -730,19 +706,15 @@ def pretty_str(envinfo):
    mutable_dict = replace_nones(mutable_dict)

    # If either of these are '', replace with 'No relevant packages'
-    mutable_dict['pip_packages'] = replace_if_empty(
-        mutable_dict['pip_packages'])
-    mutable_dict['conda_packages'] = replace_if_empty(
-        mutable_dict['conda_packages'])
+    mutable_dict['pip_packages'] = replace_if_empty(mutable_dict['pip_packages'])
+    mutable_dict['conda_packages'] = replace_if_empty(mutable_dict['conda_packages'])

    # Tag conda and pip packages with a prefix
    # If they were previously None, they'll show up as ie '[conda] Could not collect'
    if mutable_dict['pip_packages']:
-        mutable_dict['pip_packages'] = prepend(
-            mutable_dict['pip_packages'], '[{}] '.format(envinfo.pip_version))
+        mutable_dict['pip_packages'] = prepend(mutable_dict['pip_packages'], '[{}] '.format(envinfo.pip_version))
    if mutable_dict['conda_packages']:
-        mutable_dict['conda_packages'] = prepend(mutable_dict['conda_packages'],
-                                                 '[conda] ')
+        mutable_dict['conda_packages'] = prepend(mutable_dict['conda_packages'], '[conda] ')
    mutable_dict['cpu_info'] = envinfo.cpu_info
    return env_info_fmt.format(**mutable_dict)

@@ -756,18 +728,13 @@ def main():
    output = get_pretty_env_info()
    print(output)

-    if TORCH_AVAILABLE and hasattr(torch, 'utils') and hasattr(
-            torch.utils, '_crash_handler'):
+    if TORCH_AVAILABLE and hasattr(torch, 'utils') and hasattr(torch.utils, '_crash_handler'):
        minidump_dir = torch.utils._crash_handler.DEFAULT_MINIDUMP_DIR
        if sys.platform == "linux" and os.path.exists(minidump_dir):
-            dumps = [
-                os.path.join(minidump_dir, dump)
-                for dump in os.listdir(minidump_dir)
-            ]
+            dumps = [os.path.join(minidump_dir, dump) for dump in os.listdir(minidump_dir)]
            latest = max(dumps, key=os.path.getctime)
            ctime = os.path.getctime(latest)
-            creation_time = datetime.datetime.fromtimestamp(ctime).strftime(
-                '%Y-%m-%d %H:%M:%S')
+            creation_time = datetime.datetime.fromtimestamp(ctime).strftime('%Y-%m-%d %H:%M:%S')
            msg = "\n*** Detected a minidump at {} created on {}, ".format(latest, creation_time) + \
                  "if this is related to your bug please include it when you file a report ***"
            print(msg, file=sys.stderr)
@@ -1,5 +1,4 @@
-from .video_generator.nodes import (NODE_CLASS_MAPPINGS,
-                                    NODE_DISPLAY_NAME_MAPPINGS)
+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']
+__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS', 'WEB_DIRECTORY']
@@ -14,10 +14,7 @@ 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 = [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": {
@@ -65,13 +62,10 @@ class LoadImagePath:
                None,
            ]
            if 'A' in processed_image.getbands():
-                mask = np.array(processed_image.getchannel('A')).astype(
-                    np.float32) / 255.0
+                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 = 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")
@@ -9,8 +9,7 @@ from PIL import ImageFile, UnidentifiedImageError
 T = TypeVar('T')


-def conditioning_set_values(conditioning: list[Any],
-                            values: dict[str, Any] | None = None) -> list[Any]:
+def conditioning_set_values(conditioning: list[Any], values: dict[str, Any] | None = None) -> list[Any]:
    if values is None:
        values = {}
    c = []
@@ -27,8 +26,7 @@ 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
+    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)
@@ -39,12 +37,7 @@ def pillow(fn: Callable[[Any], T], arg: Any) -> T:


 def hasher() -> Callable[[], Any]:
-    hashfuncs = {
-        "md5": hashlib.md5,
-        "sha1": hashlib.sha1,
-        "sha256": hashlib.sha256,
-        "sha512": hashlib.sha512
-    }
+    hashfuncs = {"md5": hashlib.md5, "sha1": hashlib.sha1, "sha256": hashlib.sha256, "sha512": hashlib.sha512}
    return hashfuncs[args.default_hashing_function]


@@ -58,8 +51,7 @@ def string_to_torch_dtype(string: str) -> torch.dtype | None:
    return None


-def image_alpha_fix(destination: torch.Tensor,
-                    source: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+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]:
@@ -26,11 +26,7 @@ class TextEncoderConfig:
    CATEGORY = "fastvideo"

    def set_args(self, prefix, quant_config, lora_config):
-        raw_args = {
-            "prefix": prefix,
-            "quant_config": quant_config,
-            "lora_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)}
@@ -13,11 +13,7 @@ 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__))))))
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))))


 # Custom exception for interruption
@@ -28,8 +24,7 @@ class GenerationInterruptedException(Exception):
 # Custom exception for interruption that ComfyUI will recognize
 class GenerationCancelledException(Exception):

-    def __init__(self,
-                 message: str = "Generation was cancelled by user") -> None:
+    def __init__(self, message: str = "Generation was cancelled by user") -> None:
        self.message = message
        super().__init__(self.message)

@@ -139,8 +134,7 @@ class VideoGenerator:
                self._generation_interrupted = True

                # Try to send interrupt signal to worker processes
-                if self.generator is not None and hasattr(
-                        self.generator, 'executor'):
+                if self.generator is not None and hasattr(self.generator, 'executor'):
                    try:
                        # The MultiprocExecutor has a workers attribute
                        if hasattr(self.generator.executor, 'workers'):
@@ -156,16 +150,12 @@ class VideoGenerator:
                break
            time.sleep(0.5)

-    def _run_generation(self, prompt: str, output_path: str,
-                        inference_args: dict[str, Any]) -> None:
+    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")
+                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:
@@ -226,8 +216,7 @@ class VideoGenerator:
            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_config_from_args(pipeline_config.text_encoder_configs, text_encoder_config)

        # Update top-level pipeline config with remaining arguments
        raw_pipeline_args = {}
@@ -245,10 +234,7 @@ class VideoGenerator:
            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)
-        }
+        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)

@@ -262,38 +248,30 @@ class VideoGenerator:
        if dit_cpu_offload is not None:
            raw_generation_args['dit_cpu_offload'] = dit_cpu_offload

-        generation_args = {
-            k: v
-            for k, v in raw_generation_args.items()
-            if str(int(v)) != str(-99999)
-        }
+        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)
+            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),
+                                                   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 = 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():
+        while self._generation_thread.is_alive() and not self._interrupt_event.is_set():
            self._generation_thread.join(timeout=0.5)

        self._generation_active = False
@@ -42,8 +42,7 @@ def fix_case(text: str) -> str:
        r"int\d+": lambda x: x.group(0).upper(),  # e.g. int8, int16
    }
    for pattern, repl in subs.items():
-        text = re.sub(rf'\b{pattern}\b', repl, text,
-                      flags=re.IGNORECASE)  # type: ignore[call-overload]
+        text = re.sub(rf'\b{pattern}\b', repl, text, flags=re.IGNORECASE)  # type: ignore[call-overload]
    return text


@@ -134,9 +133,7 @@ class Example:
        if not markdown_files:
            raise IndexError(f"No Markdown files found in {self.path}")

-        readme_files = [
-            f for f in markdown_files if f.name.lower() == "readme.md"
-        ]
+        readme_files = [f for f in markdown_files if f.name.lower() == "readme.md"]
        if readme_files:
            return readme_files[0]

@@ -156,8 +153,7 @@ class Example:
        if self.path.is_file():
            return []
        is_other_file = lambda file: file.is_file() and file != self.main_file
-        return [file for file in self.path.rglob("*")
-                if is_other_file(file)]  # type: ignore[no-untyped-call]
+        return [file for file in self.path.rglob("*") if is_other_file(file)]  # type: ignore[no-untyped-call]

    def determine_title(self) -> str:
        return fix_case(self.path.stem.replace("_", " ").title())
@@ -190,8 +186,8 @@ class Example:
        content += "## Additional Files\n\n"
        # Define binary/non-text file extensions to skip
        binary_extensions = {
-            '.mp4', '.avi', '.mov', '.mkv', '.gif', '.jpg', '.jpeg', '.png',
-            '.webp', '.bmp', '.pdf', '.zip', '.tar', '.gz', '.mp3', '.wav'
+            '.mp4', '.avi', '.mov', '.mkv', '.gif', '.jpg', '.jpeg', '.png', '.webp', '.bmp', '.pdf', '.zip', '.tar',
+            '.gz', '.mp3', '.wav'
        }

        for file in sorted(self.other_files):
@@ -259,8 +255,7 @@ def create_category_indices() -> dict[str, Index]:
    category_indices = {
        "inference":
        Index(
-            path=ROOT_DIR /
-            "docs/inference/examples/examples_inference_index.md",
+            path=ROOT_DIR / "docs/inference/examples/examples_inference_index.md",
            title="🚀 Examples",
            description=
            "Inference examples demonstrate how to use FastVideo inference. We recommend starting with [basic.md](basic.md).",
@@ -271,18 +266,15 @@ def create_category_indices() -> dict[str, Index]:
        Index(
            path=ROOT_DIR / "docs/training/examples/examples_training_index.md",
            title="🚀 Examples",
-            description=
-            "Training examples demonstrate how to use FastVideo training.",
+            description="Training examples demonstrate how to use FastVideo training.",
            caption="Examples",
            maxdepth=3,
        ),
        "distillation":
        Index(
-            path=ROOT_DIR /
-            "docs/distillation/examples/examples_distillation_index.md",
+            path=ROOT_DIR / "docs/distillation/examples/examples_distillation_index.md",
            title="🚀 Examples",
-            description=
-            "Distillation examples demonstrate how to use FastVideo distillation.",
+            description="Distillation examples demonstrate how to use FastVideo distillation.",
            caption="Examples",
            maxdepth=3,
        ),
@@ -296,8 +288,7 @@ def create_category_indices() -> dict[str, Index]:
    return category_indices


-def find_examples(category_indices: dict[str, Index],
-                  generate_main_index: bool) -> list[Example]:
+def find_examples(category_indices: dict[str, Index], generate_main_index: bool) -> list[Example]:
    """Find all examples from the examples directory."""
    examples = []
    glob_patterns = ["*.py", "*.md", "*.sh"]
@@ -339,13 +330,9 @@ def find_examples(category_indices: dict[str, Index],
    return examples


-def create_nested_structures(
-    examples: list[Example]
-) -> dict[str, dict[str, dict[str, dict[str, NestedStructure]]]]:
+def create_nested_structures(examples: list[Example]) -> dict[str, dict[str, dict[str, dict[str, NestedStructure]]]]:
    """Create nested structures for training and distillation categories."""
-    nested_structures: dict[str, dict[str, dict[str,
-                                                dict[str,
-                                                     NestedStructure]]]] = {}
+    nested_structures: dict[str, dict[str, dict[str, dict[str, NestedStructure]]]] = {}

    # Map category names to actual directory names
    category_dir_mapping = {
@@ -378,13 +365,11 @@ def create_nested_structures(
                    nested_structures[example.category][method][model] = {}

                # Store the nested structure
-                nested_structures[
-                    example.category][method][model][dataset] = NestedStructure(
-                        category=example.category,
-                        method=method,
-                        model=model,
-                        dataset=dataset,
-                        example=example)
+                nested_structures[example.category][method][model][dataset] = NestedStructure(category=example.category,
+                                                                                              method=method,
+                                                                                              model=model,
+                                                                                              dataset=dataset,
+                                                                                              example=example)

        elif example.category == "distillation" and len(path_parts) >= 2:
            # For distillation examples like Wan2.1-T2V/Wan-Syn-Data-480P
@@ -401,20 +386,16 @@ def create_nested_structures(
                nested_structures[example.category][method][model] = {}

            # Store the nested structure
-            nested_structures[
-                example.category][method][model][dataset] = NestedStructure(
-                    category=example.category,
-                    method=method,
-                    model=model,
-                    dataset=dataset,
-                    example=example)
+            nested_structures[example.category][method][model][dataset] = NestedStructure(category=example.category,
+                                                                                          method=method,
+                                                                                          model=model,
+                                                                                          dataset=dataset,
+                                                                                          example=example)

    return nested_structures


-def generate_flat_examples(examples: list[Example],
-                           category_indices: dict[str, Index],
-                           examples_index: Index | None,
+def generate_flat_examples(examples: list[Example], category_indices: dict[str, Index], examples_index: Index | None,
                           generate_main_index: bool) -> None:
    """Generate documentation for flat structure examples (inference, etc.)."""
    for example in examples:
@@ -437,9 +418,8 @@ def generate_flat_examples(examples: list[Example],
        index.documents.append(example.path.stem)


-def generate_nested_examples(nested_structures: dict[str, dict[str, dict[
-    str, dict[str, NestedStructure]]]], category_indices: dict[str,
-                                                               Index]) -> None:
+def generate_nested_examples(nested_structures: dict[str, dict[str, dict[str, dict[str, NestedStructure]]]],
+                             category_indices: dict[str, Index]) -> None:
    """Generate documentation for nested structure examples (training, distillation)."""
    for category_name in ["training", "distillation"]:
        if category_name not in category_indices or category_name not in nested_structures:
@@ -464,12 +444,11 @@ def generate_nested_examples(nested_structures: dict[str, dict[str, dict[
                        f.write(nested_struct.example.generate())

                # Create model-level index
-                model_index = Index(
-                    path=category_base_dir / f"{model}.md",
-                    title=fix_case(model.replace('_', ' ')),
-                    description=f"Examples for the {model} model.",
-                    caption=f"{fix_case(model.replace('_', ' '))} Datasets",
-                    maxdepth=1)
+                model_index = Index(path=category_base_dir / f"{model}.md",
+                                    title=fix_case(model.replace('_', ' ')),
+                                    description=f"Examples for the {model} model.",
+                                    caption=f"{fix_case(model.replace('_', ' '))} Datasets",
+                                    maxdepth=1)

                # Add dataset indices to model index
                for dataset, nested_struct in datasets.items():
@@ -508,8 +487,7 @@ def generate_examples(generate_main_index: bool = False) -> None:
        examples_index = Index(
            path=main_index_dir / "examples_index.md",
            title="💡 Examples",
-            description=
-            "A collection of examples demonstrating usage of FastVideo.\n\n"
+            description="A collection of examples demonstrating usage of FastVideo.\n\n"
            f"All documented examples are autogenerated using [generate_examples.py](https://github.com/{GITHUB_REPO}/blob/main/docs/generate_examples.py) "
            f"from examples found in the [examples](https://github.com/{GITHUB_REPO}/tree/main/examples) directory.",
            caption="Examples",
@@ -522,8 +500,7 @@ def generate_examples(generate_main_index: bool = False) -> None:
    nested_structures = create_nested_structures(examples)

    # Generate flat structure examples (inference, etc.)
-    generate_flat_examples(examples, category_indices, examples_index,
-                           generate_main_index)
+    generate_flat_examples(examples, category_indices, examples_index, generate_main_index)

    # Generate nested structure examples (training, distillation)
    generate_nested_examples(nested_structures, category_indices)
@@ -534,11 +511,8 @@ def generate_examples(generate_main_index: bool = False) -> None:
            # Add to main index if it exists
            if generate_main_index and examples_index:
                main_index_dir = examples_index.path.parent
-                rel_path = os.path.relpath(category_index.path,
-                                           start=main_index_dir)
-                examples_index.documents.insert(
-                    0,
-                    str(rel_path).replace("\\", "/").replace(".md", ""))
+                rel_path = os.path.relpath(category_index.path, start=main_index_dir)
+                examples_index.documents.insert(0, str(rel_path).replace("\\", "/").replace(".md", ""))

            # Write the category index file
            with open(category_index.path, "w+") as f:
@@ -0,0 +1,93 @@
+# Training
+
+## Quick Start
+
+### Single-node
+
+```bash
+bash examples/train/run.sh <config.yaml> [--dotted.key value ...]
+```
+
+The script auto-detects available GPUs. Override with `NUM_GPUS`:
+
+```bash
+NUM_GPUS=4 bash examples/train/run.sh \
+    examples/train/configs/fine_tuning/wan/t2v.yaml
+```
+
+### Multi-node (Slurm)
+
+```bash
+bash examples/train/run_slurm.sh <config.yaml> <num_nodes> [--dotted.key value ...]
+```
+
+```bash
+bash examples/train/run_slurm.sh \
+    examples/train/configs/fine_tuning/wan/t2v.yaml 4 \
+    --training.distributed.hsdp_shard_dim 8 \
+    --training.distributed.hsdp_replicate_dim 4
+```
+
+Slurm environment variables:
+
+| Variable | Default | Description |
+|---|---|---|
+| `PARTITION` | `main` | Slurm partition |
+| `NUM_GPUS` | `8` | GPUs per node |
+| `CPUS_PER_TASK` | `128` | CPUs per task |
+| `MEM` | `1440G` | Memory per node |
+| `EXCLUDE` | `""` | Nodes to exclude |
+
+### Overriding config values
+
+Any config field can be overridden from the command line using dotted keys:
+
+```bash
+bash examples/train/run.sh examples/train/configs/fine_tuning/wan/t2v.yaml \
+    --training.optimizer.learning_rate 1e-5 \
+    --training.loop.max_train_steps 1000 \
+    --training.checkpoint.output_dir outputs/my_experiment
+```
+
+### Resuming from a checkpoint
+
+```bash
+bash examples/train/run.sh examples/train/configs/fine_tuning/wan/t2v.yaml \
+    --training.checkpoint.resume_from_checkpoint outputs/my_experiment/checkpoint-500
+```
+
+## W&B Logging
+
+Training metrics and validation videos are logged to
+[Weights & Biases](https://wandb.ai). Set your API key before launching:
+
+```bash
+export WANDB_API_KEY=your_key_here
+```
+
+To disable W&B (e.g. for local debugging):
+
+```bash
+export WANDB_MODE=offline
+```
+
+The project name and run name are set in the config under `training.tracker`:
+
+```yaml
+training:
+  tracker:
+    project_name: my_project
+    run_name: my_run
+```
+
+## Directory Layout
+
+```
+examples/train/
+├── configs/      # Single-step training configs (by method and model)
+├── scenario/     # Multi-step end-to-end training pipelines
+├── run.sh        # Single-node launcher
+└── run_slurm.sh  # Multi-node Slurm launcher
+```
+
+See `configs/README.md` and `scenario/README.md` for details.
@@ -0,0 +1,21 @@
+# Training Configs
+
+Single-step training configurations organized by method and model.
+
+```
+configs/
+├── fine_tuning/          # Standard finetuning and DFSFT
+├── distribution_matching/  # DMD2 and Self-Forcing
+├── knowledge_distillation/ # KD from teacher to student
+└── example.yaml          # Annotated reference config with all fields
+```
+
+Each method directory contains per-model subdirectories (e.g. `wan/`, `hunyuan/`).
+
+Launch any config with:
+
+```bash
+bash examples/train/run.sh examples/train/configs/<method>/<model>/<config>.yaml
+```
+
+For multi-step training pipelines, see `examples/train/scenario/`.
@@ -0,0 +1,72 @@
+# HunyuanVideo T2V finetune config.
+#
+# Data must be preprocessed with Hunyuan VAE + dual text encoders
+# (LLaMA + CLIP) into parquet format before training.
+# See scripts/preprocess/ for preprocessing pipeline.
+
+models:
+  student:
+    _target_: fastvideo.train.models.hunyuan.HunyuanModel
+    init_from: hunyuanvideo-community/HunyuanVideo
+    trainable: true
+
+method:
+  _target_: fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod
+
+training:
+  distributed:
+    num_gpus: 8
+    sp_size: 1
+    tp_size: 1
+    hsdp_replicate_dim: 1
+    hsdp_shard_dim: 8
+
+  data:
+    data_path: data/hunyuan_training_data
+    dataloader_num_workers: 4
+    train_batch_size: 1
+    training_cfg_rate: 0.0
+    seed: 1000
+    num_latent_t: 32
+    num_height: 720
+    num_width: 1280
+    num_frames: 125
+
+  optimizer:
+    learning_rate: 1.0e-5
+    betas: [0.9, 0.999]
+    weight_decay: 0.01
+    lr_scheduler: constant
+    lr_warmup_steps: 0
+
+  loop:
+    max_train_steps: 5000
+    gradient_accumulation_steps: 1
+
+  checkpoint:
+    output_dir: outputs/hunyuan_finetune
+    training_state_checkpointing_steps: 500
+    checkpoints_total_limit: 3
+    resume_from_checkpoint: latest
+
+  tracker:
+    project_name: fastvideo_hunyuan
+    run_name: hunyuan_finetune
+
+  model:
+    enable_gradient_checkpointing_type: full
+
+callbacks:
+  grad_clip:
+    _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback
+    max_grad_norm: 1.0
+  validation:
+    _target_: fastvideo.train.callbacks.validation.ValidationCallback
+    pipeline_target: fastvideo.pipelines.basic.hunyuan.hunyuan_pipeline.HunyuanVideoPipeline
+    dataset_file: data/hunyuan_training_data/validation_prompts.json
+    every_steps: 100
+    sampling_steps: [50]
+    guidance_scale: 1.0
+
+pipeline:
+  flow_shift: 7
@@ -0,0 +1,68 @@
+# Wan 2.1 T2V 1.3B finetune.
+
+models:
+  student:
+    _target_: fastvideo.train.models.wan.WanModel
+    init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers
+    trainable: true
+
+method:
+  _target_: fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod
+
+training:
+  distributed:
+    num_gpus: 8
+    sp_size: 1
+    tp_size: 1
+    hsdp_replicate_dim: 8
+    hsdp_shard_dim: 1
+
+  data:
+    data_path: data/preprocessed
+    dataloader_num_workers: 4
+    train_batch_size: 1
+    training_cfg_rate: 0.0
+    seed: 1000
+    num_latent_t: 20
+    num_height: 448
+    num_width: 832
+    num_frames: 77
+
+  optimizer:
+    learning_rate: 1.0e-6
+    betas: [0.9, 0.999]
+    weight_decay: 0.01
+    lr_scheduler: constant
+    lr_warmup_steps: 0
+
+  loop:
+    max_train_steps: 4000
+    gradient_accumulation_steps: 1
+
+  checkpoint:
+    output_dir: outputs/wan_finetune
+    training_state_checkpointing_steps: 100
+    checkpoints_total_limit: 3
+    resume_from_checkpoint: latest
+
+  tracker:
+    project_name: fastvideo
+    run_name: wan_finetune
+
+  model:
+    enable_gradient_checkpointing_type: full
+
+callbacks:
+  grad_clip:
+    _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback
+    max_grad_norm: 1.0
+  validation:
+    _target_: fastvideo.train.callbacks.validation.ValidationCallback
+    pipeline_target: fastvideo.pipelines.basic.wan.wan_pipeline.WanPipeline
+    dataset_file: data/validation_prompts.json
+    every_steps: 100
+    sampling_steps: [50]
+    guidance_scale: 5.0
+
+pipeline:
+  flow_shift: 3
@@ -0,0 +1,89 @@
+# Knowledge Distillation: Wan 2.1 Causal T2V 1.3B
+#
+# Trains a causal student (1.3B) with per-frame block-quantized MSE loss
+# on teacher (14B) ODE trajectories. The resulting checkpoint serves as
+# the ode_init weight initialization for downstream Self-Forcing training.
+#
+# The same teacher_path_cache can be shared with kd_wan2.1_t2v_1.3B.yaml
+# since the cache stores teacher trajectories independent of the student.
+#
+# After training, export the ode_init checkpoint:
+#   python -m fastvideo.train.entrypoint.dcp_to_diffusers \
+#       --role student \
+#       --output_dir outputs/wan2.1_causal_kd_ode_init/diffusers \
+#       --checkpoint_dir outputs/wan2.1_causal_kd_ode_init/latest
+#
+# The exported model.safetensors can then be passed as:
+#   transformer_override_safetensor in a Self-Forcing config.
+#
+# Remove the 'teacher' block once the cache is complete to avoid
+# loading the 14B model on subsequent runs.
+
+models:
+  student:
+    _target_: fastvideo.train.models.wan.WanCausalModel
+    init_from: wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers
+    trainable: true
+  teacher:                          # remove once cache is complete
+    _target_: fastvideo.train.models.wan.WanModel
+    init_from: Wan-AI/Wan2.1-T2V-14B-Diffusers
+    trainable: false
+    disable_custom_init_weights: true
+
+method:
+  _target_: fastvideo.train.methods.knowledge_distillation.kd.KDCausalMethod
+  teacher_path_cache: data/kd_cache/wan14b_4step
+  t_list: [999, 937, 833, 624, 0]   # student inference schedule (num_steps+1)
+  student_sample_steps: 4
+  teacher_inference_steps: 48       # teacher ODE steps; nearest indices selected for t_list
+  teacher_guidance_scale: 3.5       # CFG scale during teacher path generation
+  num_frames_per_block: 3           # frames sharing the same noise level (per-frame KD)
+
+training:
+  distributed:
+    num_gpus: 8
+    sp_size: 1
+    tp_size: 1
+    hsdp_replicate_dim: 1
+    hsdp_shard_dim: 8
+
+  data:
+    data_path: data/Wan-Syn_77x448x832_600k
+    dataloader_num_workers: 4
+    train_batch_size: 2
+    training_cfg_rate: 0.0
+    seed: 1000
+    num_latent_t: 20
+    num_height: 448
+    num_width: 832
+    num_frames: 77
+
+  optimizer:
+    learning_rate: 7e-5
+    betas: [0.9, 0.999]
+    weight_decay: 0.0
+    lr_scheduler: constant
+    lr_warmup_steps: 0
+
+  loop:
+    max_train_steps: 10000
+    gradient_accumulation_steps: 1
+
+  checkpoint:
+    output_dir: outputs/wan2.1_causal_kd_ode_init
+    training_state_checkpointing_steps: 1000
+    checkpoints_total_limit: 3
+
+  tracker:
+    project_name: kd-wan
+    run_name: wan2.1_causal_kd_ode_init
+
+  model:
+    enable_gradient_checkpointing_type: full
+
+callbacks:
+  grad_clip:
+    max_grad_norm: 1.0
+
+pipeline:
+  flow_shift: 5
@@ -0,0 +1,85 @@
+# Knowledge Distillation: Wan 2.1 T2V 1.3B
+#
+# Trains student (1.3B) with MSE loss on teacher (14B) ODE trajectories.
+# The resulting checkpoint serves as the ode_init weight initialization
+# for downstream Self-Forcing training.
+#
+# After training, export the ode_init checkpoint:
+#   python -m fastvideo.train.entrypoint.dcp_to_diffusers \
+#       --role student \
+#       --output_dir outputs/wan2.1_kd_ode_init/diffusers \
+#       --checkpoint_dir outputs/wan2.1_kd_ode_init/latest
+#
+# The exported model.safetensors can then be passed as:
+#   transformer_override_safetensor in a Self-Forcing config.
+#
+# Remove the 'teacher' block once the cache is complete to avoid
+# loading the 14B model on subsequent runs.
+
+models:
+  student:
+    _target_: fastvideo.train.models.wan.WanModel
+    init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers
+    trainable: true
+  teacher:                          # remove once cache is complete
+    _target_: fastvideo.train.models.wan.WanModel
+    init_from: Wan-AI/Wan2.1-T2V-14B-Diffusers
+    trainable: false
+    disable_custom_init_weights: true
+
+method:
+  _target_: fastvideo.train.methods.knowledge_distillation.kd.KDMethod
+  teacher_path_cache: data/kd_cache/wan14b_4step
+  t_list: [999, 937, 833, 624, 0]   # student inference schedule (num_steps+1)
+  student_sample_steps: 4
+  teacher_inference_steps: 48       # teacher ODE steps; nearest indices selected for t_list
+  teacher_guidance_scale: 3.5       # CFG scale during teacher path generation
+
+training:
+  distributed:
+    num_gpus: 8
+    sp_size: 1
+    tp_size: 1
+    hsdp_replicate_dim: 1
+    hsdp_shard_dim: 8
+
+  data:
+    data_path: data/Wan-Syn_77x448x832_600k
+    dataloader_num_workers: 4
+    train_batch_size: 2
+    training_cfg_rate: 0.0
+    seed: 1000
+    num_latent_t: 20
+    num_height: 448
+    num_width: 832
+    num_frames: 77
+
+  optimizer:
+    learning_rate: 7e-5
+    betas: [0.9, 0.999]
+    weight_decay: 0.0
+    lr_scheduler: constant
+    lr_warmup_steps: 0
+
+  loop:
+    max_train_steps: 10000
+    gradient_accumulation_steps: 1
+
+  checkpoint:
+    output_dir: outputs/wan2.1_kd_ode_init
+    training_state_checkpointing_steps: 1000
+    checkpoints_total_limit: 3
+
+  tracker:
+    project_name: kd-wan
+    run_name: wan2.1_kd_ode_init
+
+  model:
+    enable_gradient_checkpointing_type: full
+
+callbacks:
+  grad_clip:
+    max_grad_norm: 1.0
+
+pipeline:
+  flow_shift: 5
@@ -5,12 +5,12 @@
 #   bash examples/train/run.sh <config.yaml> [--dotted.key value ...]
 #
 # Examples:
-#   bash examples/train/run.sh examples/train/finetune_wan2.1_t2v_1.3B_vsa_phase3.4_0.9sparsity.yaml
-#   bash examples/train/run.sh examples/train/distill_wan2.1_t2v_1.3B_dmd2.yaml --dry-run
-#   bash examples/train/run.sh examples/train/distill_wan2.1_t2v_1.3B_dmd2.yaml \
+#   bash examples/train/run.sh examples/train/configs/fine_tuning/wan/t2v.yaml
+#   bash examples/train/run.sh examples/train/configs/distribution_matching/wan/dmd2_t2v.yaml --dry-run
+#   bash examples/train/run.sh examples/train/configs/distribution_matching/wan/dmd2_t2v.yaml \
 #       --training.distributed.num_gpus 4 \
 #       --training.optimizer.learning_rate 1e-5
-#   bash examples/train/run.sh examples/train/distill_wan2.1_t2v_1.3B_dmd2.yaml \
+#   bash examples/train/run.sh examples/train/configs/distribution_matching/wan/dmd2_t2v.yaml \
 #       --training.checkpoint.resume_from_checkpoint outputs/my_run/checkpoint-1000
 #
 # Logs are written to logs/<config_name>_<timestamp>.log (and also printed to stdout).
@@ -6,7 +6,7 @@
 #
 # Examples:
 #   bash examples/train/run_slurm.sh examples/train/configs/example.yaml 2
-#   bash examples/train/run_slurm.sh examples/train/configs/distill_wan2.1_t2v_1.3B_dmd2.yaml 4 \
+#   bash examples/train/run_slurm.sh examples/train/configs/distribution_matching/wan/dmd2_t2v.yaml 4 \
 #       --training.optimizer.learning_rate 1e-5
 #   bash examples/train/run_slurm.sh examples/train/configs/example.yaml 8 \
 #       --training.checkpoint.resume_from_checkpoint outputs/my_run/checkpoint-1000
@@ -0,0 +1,13 @@
+# Training Scenarios
+
+End-to-end multi-step training pipelines. Each scenario directory contains
+all configs, scripts, and data needed to run a complete workflow.
+
+```
+scenario/
+└── ode_init_self_forcing_wan_causal/   # KD → export → Self-Forcing
+```
+
+See the `usage.md` inside each scenario for step-by-step instructions.
+
+For single-step configs, see `examples/train/configs/`.
@@ -0,0 +1,92 @@
+"""Decode 'real' latents from a KD cache to videos.
+
+Uses fastvideo's AutoencoderKLWan and its built-in denormalization
+(latent * std + mean, matching normalize_dit_input's inverse).
+
+Usage:
+    python examples/train/scenario/ode_init_self_forcing_wan_causal/decode_kd_cache.py \
+        --cache_dir data/kd_test_cache_small \
+        --out_dir   data/kd_test_cache_small/video \
+        --fps 16
+"""
+
+import argparse
+from pathlib import Path
+
+import imageio
+import numpy as np
+import torch
+from diffusers import AutoencoderKLWan
+
+MODEL_ID = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
+
+
+def denormalize(latents: torch.Tensor, vae: AutoencoderKLWan) -> torch.Tensor:
+    """Inverse of normalize_dit_input('wan', ...).
+
+    normalize_dit_input: normalized = (latent - mean) * (1/std)
+    inverse:             latent     = normalized * std + mean
+    """
+    cfg = vae.config
+    mean = torch.tensor(cfg.latents_mean, dtype=latents.dtype,
+                        device=latents.device).view(1, -1, 1, 1, 1)
+    std = torch.tensor(cfg.latents_std, dtype=latents.dtype,
+                       device=latents.device).view(1, -1, 1, 1, 1)
+    return latents * std + mean
+
+
+def decode_real(pt_path: Path, vae: AutoencoderKLWan,
+                device: torch.device) -> np.ndarray:
+    """Load one .pt cache file and decode its 'real' latent to uint8 [T,H,W,3]."""
+    d = torch.load(pt_path, weights_only=True, map_location="cpu")
+    real = d["real"].float()  # [T, C, H, W]  (normalized)
+
+    # [T, C, H, W] → [1, C, T, H, W]
+    latent = real.permute(1, 0, 2, 3).unsqueeze(0).to(device)
+    latent = denormalize(latent, vae)
+
+    with torch.no_grad():
+        video = vae.decode(latent).sample  # [1, 3, T_px, H_px, W_px]
+
+    video = video.squeeze(0).permute(1, 2, 3, 0)  # [T, H, W, 3]
+    video = ((video.clamp(-1, 1) + 1) / 2 * 255).cpu().float().numpy()
+    return video.astype(np.uint8)
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--cache_dir", default="data/kd_test_cache_small")
+    parser.add_argument("--out_dir", default="data/kd_test_cache_small/video")
+    parser.add_argument("--model_id", default=MODEL_ID)
+    parser.add_argument("--fps", type=int, default=16)
+    parser.add_argument("--overwrite", action="store_true")
+    args = parser.parse_args()
+
+    samples_dir = Path(args.cache_dir) / "samples"
+    out_dir = Path(args.out_dir)
+    out_dir.mkdir(parents=True, exist_ok=True)
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    print(f"Loading VAE from {args.model_id} ...")
+    vae = AutoencoderKLWan.from_pretrained(
+        args.model_id, subfolder="vae", torch_dtype=torch.float32)
+    vae.eval().to(device)
+
+    pts = sorted(samples_dir.glob("*.pt"))
+    print(f"Found {len(pts)} samples in {samples_dir}")
+
+    for pt in pts:
+        out_path = out_dir / (pt.stem + ".mp4")
+        if out_path.exists() and not args.overwrite:
+            print(f"  skip {pt.name} (exists)")
+            continue
+        video_np = decode_real(pt, vae, device)
+        imageio.mimsave(str(out_path), video_np, fps=args.fps)
+        print(f"  {pt.name} → {out_path.name}  {video_np.shape}")
+
+    print(f"Done. Videos in {out_dir}/")
+
+
+if __name__ == "__main__":
+    main()
@@ -0,0 +1,91 @@
+# KD Test: Causal Wan 1.3B ODE-init
+#
+# Smoke-test for KDCausalMethod on a tiny 16-sample dataset.
+# Uses 1.3B as teacher (instead of 14B) so everything fits on 1 GPU.
+#
+# Step 1 — Run KD training:
+#   torchrun --nproc_per_node=1 -m fastvideo.train.entrypoint.train \
+#       --config examples/train/scenario/ode_init_self_forcing_wan_causal/step1_kd.yaml
+#
+# Step 2 — Export ode_init checkpoint:
+#   python -m fastvideo.train.entrypoint.dcp_to_diffusers \
+#       --role student \
+#       --checkpoint_dir outputs/kdtest/kd_causal/latest \
+#       --output_dir outputs/kdtest/kd_ode_init
+#
+# Step 3 — Run Self-Forcing with the exported ode_init:
+#   torchrun --nproc_per_node=1 -m fastvideo.train.entrypoint.train \
+#       --config examples/train/scenario/ode_init_self_forcing_wan_causal/step3_self_forcing.yaml
+#
+# Success criteria:
+#   - Cache generates 16 .pt files under data/kd_test_cache/samples/
+#   - kd_loss decreases steadily and approaches ~0 by step 200
+
+models:
+  student:
+    _target_: fastvideo.train.models.wan.WanCausalModel
+    init_from: wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers
+    trainable: true
+  teacher:                          # using 1.3B for single-GPU test; swap to 14B for real runs
+    _target_: fastvideo.train.models.wan.WanModel
+    init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers
+    trainable: false
+    disable_custom_init_weights: true
+
+method:
+  _target_: fastvideo.train.methods.knowledge_distillation.kd.KDCausalMethod
+  teacher_path_cache: data/kd_test_cache_small
+  t_list: [995, 937, 833, 625, 0]
+  student_sample_steps: 4
+  teacher_inference_steps: 48
+  teacher_guidance_scale: 3.5
+  num_frames_per_block: 3
+
+training:
+  distributed:
+    num_gpus: 8
+    sp_size: 1
+    tp_size: 1
+    hsdp_replicate_dim: 1
+    hsdp_shard_dim: 8
+
+  data:
+    data_path: data/Wan-Syn_small
+    dataloader_num_workers: 0
+    train_batch_size: 1
+    training_cfg_rate: 0.0
+    seed: 42
+    num_latent_t: 20
+    num_height: 448
+    num_width: 832
+    num_frames: 77
+
+  optimizer:
+    learning_rate: 7e-5
+    betas: [0.9, 0.999]
+    weight_decay: 0.0
+    lr_scheduler: constant
+    lr_warmup_steps: 0
+
+  loop:
+    max_train_steps: 300
+    gradient_accumulation_steps: 1
+
+  checkpoint:
+    output_dir: outputs/kdtest/kd_causal
+    training_state_checkpointing_steps: 300
+    checkpoints_total_limit: 1
+
+  tracker:
+    project_name: kd-test
+    run_name: kd_causal_smoke
+
+  model:
+    enable_gradient_checkpointing_type: full
+
+callbacks:
+  grad_clip:
+    max_grad_norm: 1.0
+
+pipeline:
+  flow_shift: 5
@@ -0,0 +1,21 @@
+#!/usr/bin/env bash
+# Step 2 — Export ODE-init weights from KD checkpoint to diffusers format.
+#
+# Usage:
+#   bash examples/train/scenario/ode_init_self_forcing_wan_causal/step2_export.sh
+
+set -euo pipefail
+
+CHECKPOINT_DIR="${1:-outputs/kdtest/kd_causal/checkpoint-300}"
+OUTPUT_DIR="${2:-outputs/kdtest/kd_ode_init}"
+
+echo "Exporting ODE-init weights..."
+echo "  checkpoint: ${CHECKPOINT_DIR}"
+echo "  output:     ${OUTPUT_DIR}"
+
+python -m fastvideo.train.entrypoint.dcp_to_diffusers \
+    --role student \
+    --checkpoint "${CHECKPOINT_DIR}" \
+    --output-dir "${OUTPUT_DIR}"
+
+echo "Done. Use the exported weights in step3_self_forcing.yaml."
@@ -0,0 +1,111 @@
+# Self-Forcing Test: Causal Wan 1.3B
+#
+# Smoke-test for SelfForcingMethod on a tiny 16-sample dataset.
+# Run this AFTER step1_kd.yaml + dcp_to_diffusers export.
+#
+# Prerequisite — export ode_init from the KD run:
+#   python -m fastvideo.train.entrypoint.dcp_to_diffusers \
+#       --role student \
+#       --checkpoint_dir outputs/kdtest/kd_causal/latest \
+#       --output_dir outputs/kdtest/kd_ode_init
+#
+# Then run this config:
+#   torchrun --nproc_per_node=1 -m fastvideo.train.entrypoint.train \
+#       --config examples/train/scenario/ode_init_self_forcing_wan_causal/step3_self_forcing.yaml
+#
+# Success criteria:
+#   - generator_loss and critic_loss both finite at step 0
+#   - Validation videos logged to tracker every 20 steps
+#   - No crash through max_train_steps=100
+
+models:
+  student:
+    _target_: fastvideo.train.models.wan.WanCausalModel
+    init_from: wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers
+    trainable: true
+    transformer_override_safetensor: outputs/kdtest/kd_ode_init/transformer/model.safetensors
+  teacher:
+    _target_: fastvideo.train.models.wan.WanModel
+    init_from: Wan-AI/Wan2.1-T2V-14B-Diffusers
+    trainable: false
+    disable_custom_init_weights: true
+  critic:
+    _target_: fastvideo.train.models.wan.WanModel
+    init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers
+    trainable: true
+    disable_custom_init_weights: true
+
+method:
+  _target_: fastvideo.train.methods.distribution_matching.self_forcing.SelfForcingMethod
+  rollout_mode: simulate
+  generator_update_interval: 5
+  real_score_guidance_scale: 4.5
+  dmd_denoising_steps: [1000, 750, 500, 250]
+
+  chunk_size: 3
+  student_sample_type: sde
+  context_noise: 0.0
+  enable_gradient_in_rollout: true
+  start_gradient_frame: 0
+
+  # Critic optimizer
+  fake_score_learning_rate: 8.0e-6
+  fake_score_betas: [0.0, 0.999]
+  fake_score_lr_scheduler: constant
+
+training:
+  distributed:
+    num_gpus: 8
+    sp_size: 1
+    tp_size: 1
+    hsdp_replicate_dim: 1
+    hsdp_shard_dim: 8
+
+  data:
+    data_path: data/Wan-Syn_small
+    dataloader_num_workers: 0
+    train_batch_size: 1
+    training_cfg_rate: 0.0
+    seed: 42
+    num_latent_t: 20
+    num_height: 448
+    num_width: 832
+    num_frames: 77
+
+  optimizer:
+    learning_rate: 1e-5
+    betas: [0.0, 0.999]
+    weight_decay: 0.01
+    lr_scheduler: constant
+    lr_warmup_steps: 0
+
+  loop:
+    max_train_steps: 100
+    gradient_accumulation_steps: 1
+
+  checkpoint:
+    output_dir: outputs/kdtest/self_forcing_causal
+    training_state_checkpointing_steps: 100
+    checkpoints_total_limit: 1
+
+  tracker:
+    project_name: kd-test
+    run_name: self_forcing_causal_smoke
+
+  model:
+    enable_gradient_checkpointing_type: full
+
+callbacks:
+  grad_clip:
+    max_grad_norm: 1.0
+  validation:
+    pipeline_target: fastvideo.pipelines.basic.wan.wan_causal_dmd_pipeline.WanCausalDMDPipeline
+    dataset_file: examples/train/scenario/ode_init_self_forcing_wan_causal/validation.json
+    every_steps: 20
+    sampling_steps: [4]
+    sampling_timesteps: [1000, 750, 500, 250]
+    num_frames: 69
+    guidance_scale: 6.0
+
+pipeline:
+  flow_shift: 5
@@ -0,0 +1,19 @@
+# ODE-Init Self-Forcing: Wan 2.1 Causal 1.3B
+
+End-to-end scenario that trains a causal Wan 1.3B model via
+ODE-init knowledge distillation followed by self-forcing.
+
+## Steps
+
+```bash
+# Step 1 — KD training (produces ODE-init checkpoint)
+bash examples/train/run.sh \
+    examples/train/scenario/ode_init_self_forcing_wan_causal/step1_kd.yaml
+
+# Step 2 — Export ODE-init weights to diffusers format
+bash examples/train/scenario/ode_init_self_forcing_wan_causal/step2_export.sh
+
+# Step 3 — Self-forcing with the exported ODE-init
+bash examples/train/run.sh \
+    examples/train/scenario/ode_init_self_forcing_wan_causal/step3_self_forcing.yaml
+```
@@ -0,0 +1,20 @@
+{
+  "data": [
+    {
+      "caption": "A stylish woman walks down a Tokyo street filled with warm glowing neon and animated city signage. She wears a black leather jacket, a long red dress, and black boots, and carries a black purse. She walks confidently and casually. The street is damp and reflective, creating a mirror effect of the colorful lights.",
+      "video_path": null,
+      "num_inference_steps": 4,
+      "height": 448,
+      "width": 832,
+      "num_frames": 69
+    },
+    {
+      "caption": "A white and orange tabby cat is seen happily darting through a dense garden, as if chasing something. Its eyes are wide and happy as it jogs forward, scanning the branches, flowers, and leaves as it walks. The scene is captured from a ground-level angle, giving a low and intimate perspective.",
+      "video_path": null,
+      "num_inference_steps": 4,
+      "height": 448,
+      "width": 832,
+      "num_frames": 69
+    }
+  ]
+}
@@ -1,10 +1,7 @@
 # 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.backends.abstract import (AttentionBackend, AttentionMetadata, AttentionMetadataBuilder)
+from fastvideo.attention.layer import (DistributedAttention, DistributedAttention_VSA, LocalAttention)
 from fastvideo.attention.selector import get_attn_backend

 __all__ = [
@@ -54,17 +54,13 @@ class AttentionMetadata:
    # Current step of diffusion process
    current_timestep: int

-    def asdict_zerocopy(self,
-                        skip_fields: set[str] | None = None) -> dict[str, Any]:
+    def asdict_zerocopy(self, skip_fields: set[str] | None = None) -> dict[str, Any]:
        """Similar to dataclasses.asdict, but avoids deepcopying."""
        if skip_fields is None:
            skip_fields = set()
        # Note that if we add dataclasses as fields, they will need
        # similar handling.
-        return {
-            field.name: getattr(self, field.name)
-            for field in fields(self) if field.name not in skip_fields
-        }
+        return {field.name: getattr(self, field.name) for field in fields(self) if field.name not in skip_fields}


 T = TypeVar("T", bound=AttentionMetadata)
@@ -125,8 +121,7 @@ class AttentionImpl(ABC, Generic[T]):
    ) -> None:
        raise NotImplementedError

-    def preprocess_qkv(self, qkv: torch.Tensor,
-                       attn_metadata: T) -> torch.Tensor:
+    def preprocess_qkv(self, qkv: torch.Tensor, attn_metadata: T) -> torch.Tensor:
        """Preprocess QKV tensor before performing attention operation.

        Default implementation returns the tensor unchanged.
@@ -72,12 +72,11 @@ class FlashAttnMetadataBuilder(AttentionMetadataBuilder):
        pass

    def build(  # type: ignore
-        self,
-        current_timestep: int,
-        attn_mask: torch.Tensor,
+            self,
+            current_timestep: int,
+            attn_mask: torch.Tensor,
    ) -> FlashAttnMetadata:
-        return FlashAttnMetadata(current_timestep=current_timestep,
-                                 attn_mask=attn_mask)
+        return FlashAttnMetadata(current_timestep=current_timestep, attn_mask=attn_mask)


 class FlashAttentionImpl(AttentionImpl):
@@ -103,32 +102,24 @@ class FlashAttentionImpl(AttentionImpl):
        attn_metadata: FlashAttnMetadata,
    ):

-        def _key_padding_mask_from_attn_mask(attn_mask: torch.Tensor,
-                                             key_len: int) -> torch.Tensor:
+        def _key_padding_mask_from_attn_mask(attn_mask: torch.Tensor, key_len: int) -> torch.Tensor:
            # Normalize attn_mask to [B, key_len] where True means valid token.
            if attn_mask.dim() == 4:
                attn_mask = attn_mask[:, 0, 0, :]
            elif attn_mask.dim() == 3:
                attn_mask = attn_mask[:, 0, :]
            elif attn_mask.dim() != 2:
-                raise ValueError(
-                    f"Unsupported attn_mask shape for FLASH_ATTN: {attn_mask.shape}"
-                )
+                raise ValueError(f"Unsupported attn_mask shape for FLASH_ATTN: {attn_mask.shape}")

-            if attn_mask.dtype == torch.bool:
-                key_padding_mask = attn_mask
-            else:
-                # SDPA additive mask convention: valid=0, masked=-inf/large negative.
-                key_padding_mask = attn_mask >= 0
+            # SDPA additive mask convention: valid=0, masked=-inf/large negative.
+            key_padding_mask = attn_mask if attn_mask.dtype == torch.bool else attn_mask >= 0

            if key_padding_mask.shape[-1] != key_len:
-                raise ValueError(
-                    "Invalid key padding mask length for FLASH_ATTN: "
-                    f"expected {key_len}, got {key_padding_mask.shape[-1]}")
+                raise ValueError("Invalid key padding mask length for FLASH_ATTN: "
+                                 f"expected {key_len}, got {key_padding_mask.shape[-1]}")
            return key_padding_mask

-        if (attn_metadata is not None and hasattr(attn_metadata, "attn_mask")
-                and attn_metadata.attn_mask is not None):
+        if (attn_metadata is not None and hasattr(attn_metadata, "attn_mask") and attn_metadata.attn_mask is not None):
            from fastvideo.attention.utils.flash_attn_no_pad import (
                flash_attn_no_pad,
                flash_attn_varlen_qk_no_pad,
@@ -144,8 +135,7 @@ class FlashAttentionImpl(AttentionImpl):
                    dtype=torch.bool,
                    device=query.device,
                )
-                key_padding_mask = _key_padding_mask_from_attn_mask(
-                    attn_mask, key.shape[1]).to(device=key.device)
+                key_padding_mask = _key_padding_mask_from_attn_mask(attn_mask, key.shape[1]).to(device=key.device)
                return flash_attn_varlen_qk_no_pad(
                    query,
                    key,
@@ -159,13 +149,8 @@ class FlashAttentionImpl(AttentionImpl):

            qkv = torch.stack([query, key, value], dim=2)

-            attn_mask = F.pad(attn_mask, (qkv.shape[1] - attn_mask.shape[1], 0),
-                              value=True)
-            output = flash_attn_no_pad(qkv,
-                                       attn_mask,
-                                       causal=False,
-                                       dropout_p=0,
-                                       softmax_scale=None)
+            attn_mask = F.pad(attn_mask, (qkv.shape[1] - attn_mask.shape[1], 0), value=True)
+            output = flash_attn_no_pad(qkv, attn_mask, causal=False, dropout_p=0, softmax_scale=None)
        else:
            output = flash_attn_func(
                query,  # type: ignore[no-untyped-call]
@@ -3,9 +3,7 @@
 import torch
 from sageattn3 import sageattn3_blackwell

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

@@ -3,8 +3,7 @@
 import torch
 from dataclasses import dataclass
 from fastvideo.attention.backends.abstract import (  # FlashAttentionMetadata,
-    AttentionBackend, AttentionImpl, AttentionMetadata,
-    AttentionMetadataBuilder)
+    AttentionBackend, AttentionImpl, AttentionMetadata, AttentionMetadataBuilder)
 from fastvideo.logger import init_logger

 logger = init_logger(__name__)
@@ -46,12 +45,11 @@ class SDPAMetadataBuilder(AttentionMetadataBuilder):
        pass

    def build(  # type: ignore
-        self,
-        current_timestep: int,
-        attn_mask: torch.Tensor,
+            self,
+            current_timestep: int,
+            attn_mask: torch.Tensor,
    ) -> SDPAMetadata:
-        return SDPAMetadata(current_timestep=current_timestep,
-                            attn_mask=attn_mask)
+        return SDPAMetadata(current_timestep=current_timestep, attn_mask=attn_mask)


 class SDPAImpl(AttentionImpl):
@@ -82,9 +80,8 @@ class SDPAImpl(AttentionImpl):
        key = key.transpose(1, 2)
        value = value.transpose(1, 2)

-        attn_mask = attn_metadata.attn_mask if (
-            attn_metadata is not None
-            and hasattr(attn_metadata, "attn_mask")) else None
+        attn_mask = attn_metadata.attn_mask if (attn_metadata is not None
+                                                and hasattr(attn_metadata, "attn_mask")) else None
        attn_kwargs = {
            "attn_mask": attn_mask,
            "dropout_p": self.dropout,
@@ -93,7 +90,6 @@ class SDPAImpl(AttentionImpl):
        }
        if query.shape[1] != key.shape[1]:
            attn_kwargs["enable_gqa"] = True
-        output = torch.nn.functional.scaled_dot_product_attention(
-            query, key, value, **attn_kwargs)
+        output = torch.nn.functional.scaled_dot_product_attention(query, key, value, **attn_kwargs)
        output = output.transpose(1, 2)
        return output
@@ -55,13 +55,11 @@ def compress_kernel(

    x_offset = idx_bh * L * D
    xm_offset = idx_bh * ((L + BLOCK_L - 1) // BLOCK_L) * D
-    x = tl.load(X + x_offset + offs_l[:, None] * D + offs_d[None, :],
-                mask=offs_l[:, None] < L)
+    x = tl.load(X + x_offset + offs_l[:, None] * D + offs_d[None, :], mask=offs_l[:, None] < L)

    nx = min(BLOCK_L, L - idx_l * BLOCK_L)
    x_mean = tl.sum(x, axis=0, dtype=tl.float32) / nx
-    tl.store(XM + xm_offset + idx_l * D + offs_d,
-             x_mean.to(XM.dtype.element_ty))
+    tl.store(XM + xm_offset + idx_l * D + offs_d, x_mean.to(XM.dtype.element_ty))


 def mean_pool(x: torch.Tensor, BLK: int) -> torch.Tensor:
@@ -98,8 +96,7 @@ def get_block_map(
        lut: Top-k indices of shape (B, H, num_q_blocks, topk)
        topk: Number of key blocks selected
    """
-    arg_k = k - torch.mean(
-        k, dim=-2, keepdim=True)  # smooth-k technique from SageAttention
+    arg_k = k - torch.mean(k, dim=-2, keepdim=True)  # smooth-k technique from SageAttention
    pooled_qblocks = mean_pool(q, BLKQ)
    pooled_kblocks = mean_pool(arg_k, BLKK)
    pooled_score = pooled_qblocks @ pooled_kblocks.transpose(-1, -2)
@@ -300,11 +297,7 @@ class SLAAttentionImpl(AttentionImpl, nn.Module):
            topk_ratio = attn_metadata.topk_ratio  # type: ignore[union-attr]

        # Compute block-sparse attention pattern
-        sparse_map, lut, real_topk = get_block_map(q,
-                                                   k,
-                                                   topk_ratio=topk_ratio,
-                                                   BLKQ=self.BLKQ,
-                                                   BLKK=self.BLKK)
+        sparse_map, lut, real_topk = get_block_map(q, k, topk_ratio=topk_ratio, BLKQ=self.BLKQ, BLKK=self.BLKK)

        # Convert to compute dtype
        q = q.to(self.dtype)
@@ -312,8 +305,7 @@ class SLAAttentionImpl(AttentionImpl, nn.Module):
        v = v.to(self.dtype)

        # Sparse attention
-        o_s = _attention.apply(q, k, v, sparse_map, lut, real_topk, self.BLKQ,
-                               self.BLKK)
+        o_s = _attention.apply(q, k, v, sparse_map, lut, real_topk, self.BLKQ, self.BLKK)

        # Linear attention with feature maps
        q_linear = self.feature_map_q(q).contiguous().to(self.dtype)
@@ -412,14 +404,10 @@ class SageSLAAttentionImpl(AttentionImpl, nn.Module):
        nn.Module.__init__(self)

        if not SAGESLA_ENABLED:
-            raise ImportError(
-                "SageSLA requires spas_sage_attn. "
-                "Install with: pip install git+https://github.com/thu-ml/SpargeAttn.git"
-            )
+            raise ImportError("SageSLA requires spas_sage_attn. "
+                              "Install with: pip install git+https://github.com/thu-ml/SpargeAttn.git")

-        assert head_size in [
-            64, 128
-        ], f"SageSLA requires head_size in [64, 128], got {head_size}"
+        assert head_size in [64, 128], f"SageSLA requires head_size in [64, 128], got {head_size}"

        self.num_heads = num_heads
        self.head_size = head_size
@@ -509,11 +497,7 @@ class SageSLAAttentionImpl(AttentionImpl, nn.Module):
            BLKQ, BLKK = 128, 64

        # Compute block-sparse attention pattern
-        sparse_map, lut, real_topk = get_block_map(q,
-                                                   k,
-                                                   topk_ratio=topk_ratio,
-                                                   BLKQ=BLKQ,
-                                                   BLKK=BLKK)
+        sparse_map, lut, real_topk = get_block_map(q, k, topk_ratio=topk_ratio, BLKQ=BLKQ, BLKK=BLKK)

        # Convert to compute dtype
        q = q.to(self.dtype)
@@ -526,47 +510,33 @@ class SageSLAAttentionImpl(AttentionImpl, nn.Module):
        scale = 1.0 / (headdim**0.5)

        # Quantize Q, K to INT8
-        q_int8, q_scale, k_int8, k_scale = get_vanilla_qk_quant(
-            q, k, km, BLKQ, BLKK)
+        q_int8, q_scale, k_int8, k_scale = get_vanilla_qk_quant(q, k, km, BLKQ, BLKK)
        lut_triton, valid_block_num = block_map_lut_triton(sparse_map)

        # Quantize V to FP8
        b, h_kv, kv_len, head_dim = v.shape
        padded_len = (kv_len + 127) // 128 * 128
-        v_transposed_permutted = torch.empty((b, h_kv, head_dim, padded_len),
-                                             dtype=v.dtype,
-                                             device=v.device)
+        v_transposed_permutted = torch.empty((b, h_kv, head_dim, padded_len), dtype=v.dtype, device=v.device)
        fused.transpose_pad_permute_cuda(v, v_transposed_permutted, 1)
-        v_fp8 = torch.empty(v_transposed_permutted.shape,
-                            dtype=torch.float8_e4m3fn,
-                            device=v.device)
-        v_scale = torch.empty((b, h_kv, head_dim),
-                              dtype=torch.float32,
-                              device=v.device)
-        fused.scale_fuse_quant_cuda(v_transposed_permutted, v_fp8, v_scale,
-                                    kv_len, 2.25, 1)
+        v_fp8 = torch.empty(v_transposed_permutted.shape, dtype=torch.float8_e4m3fn, device=v.device)
+        v_scale = torch.empty((b, h_kv, head_dim), dtype=torch.float32, device=v.device)
+        fused.scale_fuse_quant_cuda(v_transposed_permutted, v_fp8, v_scale, kv_len, 2.25, 1)

        # Sparse attention with quantized kernels
        o_s = torch.empty_like(q)
        if arch == "sm90":
            qattn.qk_int8_sv_f8_accum_f32_block_sparse_attn_inst_buf_fuse_v_scale_sm90(
-                q_int8, k_int8, v_fp8, o_s, lut_triton, valid_block_num,
-                q_scale, k_scale, v_scale, 1, False, 1, scale)
+                q_int8, k_int8, v_fp8, o_s, lut_triton, valid_block_num, q_scale, k_scale, v_scale, 1, False, 1, scale)
        else:
-            pvthreshold = torch.full((q.shape[-3], ),
-                                     1e6,
-                                     dtype=torch.float32,
-                                     device=q.device)
+            pvthreshold = torch.full((q.shape[-3], ), 1e6, dtype=torch.float32, device=q.device)
            if SAGE2PP_ENABLED:
                qk_int8_sv_f8_accum_f16_block_sparse_attn_inst_buf_fuse_v_scale_with_pv_threshold(
-                    q_int8, k_int8, v_fp8, o_s, lut_triton, valid_block_num,
-                    pvthreshold, q_scale, k_scale, v_scale, 1, False, 1, scale,
-                    0)
+                    q_int8, k_int8, v_fp8, o_s, lut_triton, valid_block_num, pvthreshold, q_scale, k_scale, v_scale, 1,
+                    False, 1, scale, 0)
            else:
                qattn.qk_int8_sv_f8_accum_f32_block_sparse_attn_inst_buf_fuse_v_scale_with_pv_threshold(
-                    q_int8, k_int8, v_fp8, o_s, lut_triton, valid_block_num,
-                    pvthreshold, q_scale, k_scale, v_scale, 1, False, 1, scale,
-                    0)
+                    q_int8, k_int8, v_fp8, o_s, lut_triton, valid_block_num, pvthreshold, q_scale, k_scale, v_scale, 1,
+                    False, 1, scale, 0)
        # ========== END SPARGE ==========

        # Linear attention with feature maps
@@ -12,9 +12,7 @@ except ImportError:

 from typing import Any

-from fastvideo.attention.backends.abstract import (AttentionBackend,
-                                                   AttentionImpl,
-                                                   AttentionMetadata,
+from fastvideo.attention.backends.abstract import (AttentionBackend, AttentionImpl, AttentionMetadata,
                                                   AttentionMetadataBuilder)
 from fastvideo.distributed import get_sp_group
 from fastvideo.logger import init_logger
@@ -31,14 +29,12 @@ def get_tile_partition_indices(
 ) -> torch.LongTensor:
    T, H, W = dit_seq_shape
    ts, hs, ws = tile_size
-    indices = torch.arange(T * H * W, device=device,
-                           dtype=torch.long).reshape(T, H, W)
+    indices = torch.arange(T * H * W, device=device, dtype=torch.long).reshape(T, H, W)
    ls = []
    for t in range(math.ceil(T / ts)):
        for h in range(math.ceil(H / hs)):
            for w in range(math.ceil(W / ws)):
-                ls.append(indices[t * ts:min(t * ts + ts, T),
-                                  h * hs:min(h * hs + hs, H),
+                ls.append(indices[t * ts:min(t * ts + ts, T), h * hs:min(h * hs + hs, H),
                                  w * ws:min(w * ws + ws, W)].flatten())
    index = torch.cat(ls, dim=0)
    return index
@@ -50,8 +46,7 @@ def get_reverse_tile_partition_indices(
    tile_size: tuple[int, int, int],
    device: torch.device,
 ) -> torch.LongTensor:
-    return torch.argsort(
-        get_tile_partition_indices(dit_seq_shape, tile_size, device))
+    return torch.argsort(get_tile_partition_indices(dit_seq_shape, tile_size, device))


@functools.lru_cache(maxsize=10)
@@ -104,10 +99,8 @@ def get_non_pad_index(
    n_win = variable_block_sizes.shape[0]
    device = variable_block_sizes.device
    starts_pad = torch.arange(n_win, device=device) * max_block_size
-    index_pad = starts_pad[:, None] + torch.arange(max_block_size,
-                                                   device=device)[None, :]
-    index_mask = torch.arange(
-        max_block_size, device=device)[None, :] < variable_block_sizes[:, None]
+    index_pad = starts_pad[:, None] + torch.arange(max_block_size, device=device)[None, :]
+    index_mask = torch.arange(max_block_size, device=device)[None, :] < variable_block_sizes[:, None]
    return index_pad[index_mask]


@@ -167,23 +160,17 @@ class VideoSparseAttentionMetadataBuilder(AttentionMetadataBuilder):
        **kwargs: dict[str, Any],
    ) -> VideoSparseAttentionMetadata:
        patch_size = patch_size
-        dit_seq_shape = (raw_latent_shape[0] // patch_size[0],
-                         raw_latent_shape[1] // patch_size[1],
+        dit_seq_shape = (raw_latent_shape[0] // patch_size[0], raw_latent_shape[1] // patch_size[1],
                         raw_latent_shape[2] // patch_size[2])

-        num_tiles = (math.ceil(dit_seq_shape[0] / VSA_TILE_SIZE[0]),
-                     math.ceil(dit_seq_shape[1] / VSA_TILE_SIZE[1]),
+        num_tiles = (math.ceil(dit_seq_shape[0] / VSA_TILE_SIZE[0]), math.ceil(dit_seq_shape[1] / VSA_TILE_SIZE[1]),
                     math.ceil(dit_seq_shape[2] / VSA_TILE_SIZE[2]))
        total_seq_length = math.prod(dit_seq_shape)

-        tile_partition_indices = get_tile_partition_indices(
-            dit_seq_shape, VSA_TILE_SIZE, device)
-        reverse_tile_partition_indices = get_reverse_tile_partition_indices(
-            dit_seq_shape, VSA_TILE_SIZE, device)
-        variable_block_sizes = construct_variable_block_sizes(
-            dit_seq_shape, num_tiles, device)
-        non_pad_index = get_non_pad_index(variable_block_sizes,
-                                          math.prod(VSA_TILE_SIZE))
+        tile_partition_indices = get_tile_partition_indices(dit_seq_shape, VSA_TILE_SIZE, device)
+        reverse_tile_partition_indices = get_reverse_tile_partition_indices(dit_seq_shape, VSA_TILE_SIZE, device)
+        variable_block_sizes = construct_variable_block_sizes(dit_seq_shape, num_tiles, device)
+        non_pad_index = get_non_pad_index(variable_block_sizes, math.prod(VSA_TILE_SIZE))

        return VideoSparseAttentionMetadata(
            current_timestep=current_timestep,
@@ -213,23 +200,19 @@ class VideoSparseAttentionImpl(AttentionImpl):
        sp_group = get_sp_group()
        self.sp_size = sp_group.world_size

-    def tile(self, x: torch.Tensor, num_tiles: list[int],
-             tile_partition_indices: torch.LongTensor,
+    def tile(self, x: torch.Tensor, num_tiles: list[int], tile_partition_indices: torch.LongTensor,
             non_pad_index: torch.LongTensor) -> torch.Tensor:
        t_padded_size = num_tiles[0] * VSA_TILE_SIZE[0]
        h_padded_size = num_tiles[1] * VSA_TILE_SIZE[1]
        w_padded_size = num_tiles[2] * VSA_TILE_SIZE[2]

-        x_padded = torch.zeros(
-            (x.shape[0], t_padded_size * h_padded_size * w_padded_size,
-             x.shape[-2], x.shape[-1]),
-            device=x.device,
-            dtype=x.dtype)
+        x_padded = torch.zeros((x.shape[0], t_padded_size * h_padded_size * w_padded_size, x.shape[-2], x.shape[-1]),
+                               device=x.device,
+                               dtype=x.dtype)
        x_padded[:, non_pad_index] = x[:, tile_partition_indices]
        return x_padded

-    def untile(self, x: torch.Tensor,
-               reverse_tile_partition_indices: torch.LongTensor,
+    def untile(self, x: torch.Tensor, reverse_tile_partition_indices: torch.LongTensor,
               non_pad_index: torch.LongTensor) -> torch.Tensor:
        x = x[:, non_pad_index][:, reverse_tile_partition_indices]
        return x
@@ -239,8 +222,7 @@ class VideoSparseAttentionImpl(AttentionImpl):
        qkv: torch.Tensor,
        attn_metadata: VideoSparseAttentionMetadata,
    ) -> torch.Tensor:
-        return self.tile(qkv, attn_metadata.num_tiles,
-                         attn_metadata.tile_partition_indices,
+        return self.tile(qkv, attn_metadata.num_tiles, attn_metadata.tile_partition_indices,
                         attn_metadata.non_pad_index)

    def postprocess_output(
@@ -248,8 +230,7 @@ class VideoSparseAttentionImpl(AttentionImpl):
        output: torch.Tensor,
        attn_metadata: VideoSparseAttentionMetadata,
    ) -> torch.Tensor:
-        return self.untile(output, attn_metadata.reverse_tile_partition_indices,
-                           attn_metadata.non_pad_index)
+        return self.untile(output, attn_metadata.reverse_tile_partition_indices, attn_metadata.non_pad_index)

    def forward(  # type: ignore[override]
        self,
@@ -266,20 +247,17 @@ class VideoSparseAttentionImpl(AttentionImpl):

        VSA_sparsity = attn_metadata.VSA_sparsity

-        cur_topk = math.ceil(
-            (1 - VSA_sparsity) *
-            (attn_metadata.total_seq_length / math.prod(VSA_TILE_SIZE)))
+        cur_topk = math.ceil((1 - VSA_sparsity) * (attn_metadata.total_seq_length / math.prod(VSA_TILE_SIZE)))

        if video_sparse_attn is None:
            raise NotImplementedError("video_sparse_attn is not installed")
-        hidden_states = video_sparse_attn(
-            query,
-            key,
-            value,
-            attn_metadata.variable_block_sizes,
-            attn_metadata.variable_block_sizes,
-            cur_topk,
-            block_size=VSA_TILE_SIZE,
-            compress_attn_weight=gate_compress).transpose(1, 2)
+        hidden_states = video_sparse_attn(query,
+                                          key,
+                                          value,
+                                          attn_metadata.variable_block_sizes,
+                                          attn_metadata.variable_block_sizes,
+                                          cur_topk,
+                                          block_size=VSA_TILE_SIZE,
+                                          compress_attn_weight=gate_compress).transpose(1, 2)

        return hidden_states
@@ -6,11 +6,8 @@ from dataclasses import dataclass
 import torch
 from einops import rearrange

-from fastvideo_kernel import (moba_attn_varlen, process_moba_input,
-                              process_moba_output)
-from fastvideo.attention.backends.abstract import (AttentionBackend,
-                                                   AttentionImpl,
-                                                   AttentionMetadata,
+from fastvideo_kernel import (moba_attn_varlen, process_moba_input, process_moba_output)
+from fastvideo.attention.backends.abstract import (AttentionBackend, AttentionImpl, AttentionMetadata,
                                                   AttentionMetadataBuilder)
 from fastvideo.logger import init_logger

@@ -94,12 +91,10 @@ class VideoMobaAttentionMetadataBuilder(AttentionMetadataBuilder):
    ) -> VideoMobaAttentionMetadata:
        if device is None:
            device = torch.device("cpu")
-        assert raw_latent_shape[0] % patch_size[0] == 0 and raw_latent_shape[
-            1] % patch_size[1] == 0 and raw_latent_shape[2] % patch_size[
+        assert raw_latent_shape[0] % patch_size[0] == 0 and raw_latent_shape[1] % patch_size[
+            1] == 0 and raw_latent_shape[2] % patch_size[
                2] == 0, f"spatial patch_resolution {raw_latent_shape} should be divisible by patch_size {patch_size}"
-        patch_resolution = [
-            t // pt for t, pt in zip(raw_latent_shape, patch_size, strict=False)
-        ]
+        patch_resolution = [t // pt for t, pt in zip(raw_latent_shape, patch_size, strict=False)]

        return VideoMobaAttentionMetadata(
            current_timestep=current_timestep,
@@ -170,19 +165,11 @@ class VMOBAAttentionImpl(AttentionImpl):
            moba_chunk_size = attn_metadata.st_chunk_size
            moba_topk = attn_metadata.st_topk

-        query, chunk_size = process_moba_input(query,
-                                               attn_metadata.patch_resolution,
-                                               moba_chunk_size)
-        key, chunk_size = process_moba_input(key,
-                                             attn_metadata.patch_resolution,
-                                             moba_chunk_size)
-        value, chunk_size = process_moba_input(value,
-                                               attn_metadata.patch_resolution,
-                                               moba_chunk_size)
+        query, chunk_size = process_moba_input(query, attn_metadata.patch_resolution, moba_chunk_size)
+        key, chunk_size = process_moba_input(key, attn_metadata.patch_resolution, moba_chunk_size)
+        value, chunk_size = process_moba_input(value, attn_metadata.patch_resolution, moba_chunk_size)
        max_seqlen = query.shape[1]
-        indices_q = torch.arange(0,
-                                 query.shape[0] * query.shape[1],
-                                 device=query.device)
+        indices_q = torch.arange(0, query.shape[0] * query.shape[1], device=query.device)
        cu_seqlens = torch.arange(0,
                                  query.shape[0] * query.shape[1] + 1,
                                  query.shape[1],
@@ -205,10 +192,7 @@ class VMOBAAttentionImpl(AttentionImpl):
            simsum_threshold=attn_metadata.moba_threshold,
            threshold_type=attn_metadata.moba_threshold_type,
        )
-        hidden_states = self.pad_input(hidden_states, indices_q, batch_size,
-                                       sequence_length)
-        hidden_states = process_moba_output(hidden_states,
-                                            attn_metadata.patch_resolution,
-                                            moba_chunk_size)
+        hidden_states = self.pad_input(hidden_states, indices_q, batch_size, sequence_length)
+        hidden_states = process_moba_output(hidden_states, attn_metadata.patch_resolution, moba_chunk_size)

        return hidden_states
@@ -4,10 +4,9 @@ import torch
 import torch.nn as nn

 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.distributed.parallel_state import (get_sp_parallel_rank,
-                                                  get_sp_world_size)
+from fastvideo.distributed.communication_op import (sequence_model_parallel_all_gather,
+                                                    sequence_model_parallel_all_to_all_4D)
+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
@@ -38,10 +37,7 @@ class DistributedAttention(nn.Module):
            num_kv_heads = num_heads

        dtype = get_compute_dtype()
-        attn_backend = get_attn_backend(
-            head_size,
-            dtype,
-            supported_attention_backends=supported_attention_backends)
+        attn_backend = get_attn_backend(head_size, dtype, supported_attention_backends=supported_attention_backends)
        impl_cls = attn_backend.get_impl_cls()
        self.attn_impl = impl_cls(num_heads=num_heads,
                                  head_size=head_size,
@@ -89,8 +85,7 @@ class DistributedAttention(nn.Module):
                - replicated_o (Optional[torch.Tensor]): Output tensor for replicated tokens, if provided
        """
        # Check input shapes
-        assert q.dim() == 4 and k.dim() == 4 and v.dim(
-        ) == 4, "Expected 4D tensors"
+        assert q.dim() == 4 and k.dim() == 4 and v.dim() == 4, "Expected 4D tensors"
        batch_size, _, num_heads, _ = q.shape
        local_rank = get_sp_parallel_rank()
        world_size = get_sp_world_size()
@@ -99,13 +94,10 @@ class DistributedAttention(nn.Module):
        ctx_attn_metadata = forward_context.attn_metadata

        # Stack QKV
-        qkv = torch.cat([q, k, v],
-                        dim=0)  # [3*batch, seq_len, num_heads, head_dim]
+        qkv = torch.cat([q, k, v], dim=0)  # [3*batch, seq_len, num_heads, head_dim]

        # Redistribute heads across sequence dimension
-        qkv = sequence_model_parallel_all_to_all_4D(qkv,
-                                                    scatter_dim=2,
-                                                    gather_dim=1)
+        qkv = sequence_model_parallel_all_to_all_4D(qkv, scatter_dim=2, gather_dim=1)

        # After all-to-all, each rank has the full sequence but only a subset of heads.
        # Trim away SP padding for attention compute, then pad back before returning.
@@ -115,23 +107,17 @@ class DistributedAttention(nn.Module):

        if freqs_cis is not None:
            cos, sin = freqs_cis
-            qkv[:batch_size * 2] = _apply_rotary_emb(qkv[:batch_size * 2],
-                                                     cos,
-                                                     sin,
-                                                     is_neox_style=False)
+            qkv[:batch_size * 2] = _apply_rotary_emb(qkv[:batch_size * 2], cos, sin, is_neox_style=False)
        # Apply backend-specific preprocess_qkv
        qkv = self.attn_impl.preprocess_qkv(qkv, ctx_attn_metadata)

        # Concatenate with replicated QKV if provided
        if replicated_q is not None:
            assert replicated_k is not None and replicated_v is not None
-            replicated_qkv = torch.cat(
-                [replicated_q, replicated_k, replicated_v],
-                dim=0)  # [3, seq_len, num_heads, head_dim]
+            replicated_qkv = torch.cat([replicated_q, replicated_k, replicated_v],
+                                       dim=0)  # [3, seq_len, num_heads, head_dim]
            heads_per_rank = num_heads // world_size
-            replicated_qkv = replicated_qkv[:, :, local_rank *
-                                            heads_per_rank:(local_rank + 1) *
-                                            heads_per_rank]
+            replicated_qkv = replicated_qkv[:, :, local_rank * heads_per_rank:(local_rank + 1) * heads_per_rank]
            qkv = torch.cat([qkv, replicated_qkv], dim=1)

        q, k, v = qkv.chunk(3, dim=0)
@@ -145,16 +131,13 @@ class DistributedAttention(nn.Module):
            replicated_output = output[:, split_idx:]
            output = output[:, :split_idx]
            # TODO: make this asynchronous
-            replicated_output = sequence_model_parallel_all_gather(
-                replicated_output.contiguous(), dim=2)
+            replicated_output = sequence_model_parallel_all_gather(replicated_output.contiguous(), dim=2)
        # Apply backend-specific postprocess_output
        output = self.attn_impl.postprocess_output(output, ctx_attn_metadata)

        output = torch.nn.functional.pad(output, (0, 0, 0, 0, 0, pad_seq_len))

-        output = sequence_model_parallel_all_to_all_4D(output,
-                                                       scatter_dim=1,
-                                                       gather_dim=2)
+        output = sequence_model_parallel_all_to_all_4D(output, scatter_dim=1, gather_dim=2)

        return output, replicated_output

@@ -196,23 +179,19 @@ class DistributedAttention_VSA(DistributedAttention):
        # Check text tokens are not supported for VSA now
        assert replicated_q is None and replicated_k is None and replicated_v is None, "Replicated QKV is not supported for VSA now"
        # Check input shapes
-        assert q.dim() == 4 and k.dim() == 4 and v.dim(
-        ) == 4, "Expected 4D tensors"
+        assert q.dim() == 4 and k.dim() == 4 and v.dim() == 4, "Expected 4D tensors"

        forward_context: ForwardContext = get_forward_context()
        ctx_attn_metadata = forward_context.attn_metadata

        batch_size, seq_len, num_heads, head_dim = q.shape
        # Stack QKV
-        qkvg = torch.cat([q, k, v, gate_compress],
-                         dim=0)  # [4*batch, seq_len, num_heads, head_dim]
+        qkvg = torch.cat([q, k, v, gate_compress], dim=0)  # [4*batch, seq_len, num_heads, head_dim]

        # Redistribute heads across sequence dimension
        # Before: [4*batch, shard_seq_len, num_heads, head_dim]
        # After:  [4*batch, full_seq_len, shard_num_heads, head_dim]
-        qkvg = sequence_model_parallel_all_to_all_4D(qkvg,
-                                                     scatter_dim=2,
-                                                     gather_dim=1)
+        qkvg = sequence_model_parallel_all_to_all_4D(qkvg, scatter_dim=2, gather_dim=1)

        # After all-to-all, each rank has the full sequence but only a subset of heads
        pad_seq_len = qkvg.shape[1] - original_seq_len
@@ -220,16 +199,12 @@ class DistributedAttention_VSA(DistributedAttention):

        if freqs_cis is not None:
            cos, sin = freqs_cis
-            qkvg[:batch_size * 2] = _apply_rotary_emb(qkvg[:batch_size * 2],
-                                                      cos,
-                                                      sin,
-                                                      is_neox_style=False)
+            qkvg[:batch_size * 2] = _apply_rotary_emb(qkvg[:batch_size * 2], cos, sin, is_neox_style=False)

        qkvg = self.attn_impl.preprocess_qkv(qkvg, ctx_attn_metadata)

        q, k, v, gate_compress = qkvg.chunk(4, dim=0)
-        output = self.attn_impl.forward(
-            q, k, v, gate_compress, ctx_attn_metadata)  # type: ignore[call-arg]
+        output = self.attn_impl.forward(q, k, v, gate_compress, ctx_attn_metadata)  # type: ignore[call-arg]

        # Redistribute back if using sequence parallelism
        replicated_output = None
@@ -239,9 +214,7 @@ class DistributedAttention_VSA(DistributedAttention):

        output = torch.nn.functional.pad(output, (0, 0, 0, 0, 0, pad_seq_len))

-        output = sequence_model_parallel_all_to_all_4D(output,
-                                                       scatter_dim=1,
-                                                       gather_dim=2)
+        output = sequence_model_parallel_all_to_all_4D(output, scatter_dim=1, gather_dim=2)
        return output, replicated_output


@@ -267,10 +240,7 @@ class LocalAttention(nn.Module):
            num_kv_heads = num_heads

        dtype = get_compute_dtype()
-        attn_backend = get_attn_backend(
-            head_size,
-            dtype,
-            supported_attention_backends=supported_attention_backends)
+        attn_backend = get_attn_backend(head_size, dtype, supported_attention_backends=supported_attention_backends)
        impl_cls = attn_backend.get_impl_cls()
        self.attn_impl = impl_cls(num_heads=num_heads,
                                  head_size=head_size,
@@ -303,8 +273,7 @@ class LocalAttention(nn.Module):
            torch.Tensor: Output tensor after local attention
        """
        # Check input shapes
-        assert q.dim() == 4 and k.dim() == 4 and v.dim(
-        ) == 4, "Expected 4D tensors"
+        assert q.dim() == 4 and k.dim() == 4 and v.dim() == 4, "Expected 4D tensors"

        forward_context: ForwardContext = get_forward_context()
        ctx_attn_metadata = forward_context.attn_metadata
@@ -43,8 +43,7 @@ def get_env_variable_attn_backend() -> AttentionBackendEnum | None:
    * None otherwise
    '''
    backend_name = os.environ.get(STR_BACKEND_ENV_VAR)
-    return (None
-            if backend_name is None else backend_name_to_enum(backend_name))
+    return (None if backend_name is None else backend_name_to_enum(backend_name))


 # Global state allows a particular choice of backend
@@ -57,8 +56,7 @@ def get_env_variable_attn_backend() -> AttentionBackendEnum | None:
 forced_attn_backend: AttentionBackendEnum | None = None


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

@@ -87,8 +85,7 @@ def get_attn_backend(
    supported_attention_backends: tuple[AttentionBackendEnum, ...]
    | None = None,
 ) -> type[AttentionBackend]:
-    return _cached_get_attn_backend(head_size, dtype,
-                                    supported_attention_backends)
+    return _cached_get_attn_backend(head_size, dtype, supported_attention_backends)


@cache
@@ -106,8 +103,7 @@ def _cached_get_attn_backend(
    if not supported_attention_backends:
        raise ValueError("supported_attention_backends is empty")
    selected_backend = None
-    backend_by_global_setting: AttentionBackendEnum | None = (
-        get_global_forced_attn_backend())
+    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:
@@ -121,17 +117,14 @@ def _cached_get_attn_backend(

    if selected_backend not in supported_attention_backends:
        selected_backend = None
-    attention_cls = current_platform.get_attn_backend_cls(
-        selected_backend, head_size, dtype)
+    attention_cls = current_platform.get_attn_backend_cls(selected_backend, head_size, dtype)
    if not attention_cls:
-        raise ValueError(
-            f"Invalid attention backend for {current_platform.device_name}")
+        raise ValueError(f"Invalid attention backend for {current_platform.device_name}")
    return cast(type[AttentionBackend], resolve_obj_by_qualname(attention_cls))


@contextmanager
-def global_force_attn_backend_context_manager(
-        attn_backend: AttentionBackendEnum) -> Generator[None, None, None]:
+def global_force_attn_backend_context_manager(attn_backend: AttentionBackendEnum) -> Generator[None, None, None]:
    '''
    Globally force a FastVideo attention backend override within a
    context manager, reverting the global attention backend
@@ -5,16 +5,12 @@ if torch.cuda.is_available():
    from flash_attn.cute.interface import _flash_attn_bwd, _flash_attn_fwd
 else:
    # This error will be caught in flash_attn.py or flash_attn_no_pad.py
-    raise ImportError(
-        "flash_attn.cute is only available on CUDA devices; this error must be handled internally"
-    )
+    raise ImportError("flash_attn.cute is only available on CUDA devices; this error must be handled internally")


 def _check_dropout(dropout_p: float) -> None:
    if dropout_p != 0.0:
-        raise NotImplementedError(
-            f"flash_attn.cute does not support dropout (got dropout_p={dropout_p})"
-        )
+        raise NotImplementedError(f"flash_attn.cute does not support dropout (got dropout_p={dropout_p})")


@torch.library.custom_op(
@@ -61,8 +57,7 @@ def _flash_attn_cute_forward_fake(
    return out, lse


-def _flash_attn_cute_setup_context(ctx: torch.autograd.function.FunctionCtx,
-                                   inputs, output) -> None:
+def _flash_attn_cute_setup_context(ctx: torch.autograd.function.FunctionCtx, inputs, output) -> None:
    q, k, v, softmax_scale, causal, deterministic = inputs
    out, lse = output
    ctx.save_for_backward(q, k, v, out, lse)
@@ -160,8 +155,7 @@ def _flash_attn_cute_varlen_forward_fake(
    return out, lse


-def _flash_attn_cute_varlen_setup_context(
-        ctx: torch.autograd.function.FunctionCtx, inputs, output) -> None:
+def _flash_attn_cute_varlen_setup_context(ctx: torch.autograd.function.FunctionCtx, inputs, output) -> None:
    (
        q,
        k,
@@ -229,8 +223,7 @@ def flash_attn_func(
 ) -> torch.Tensor:
    """Only returns the output, not the lse."""
    _check_dropout(dropout_p)
-    out, _ = torch.ops.fastvideo._flash_attn_cute_forward(
-        q, k, v, softmax_scale, causal, deterministic)
+    out, _ = torch.ops.fastvideo._flash_attn_cute_forward(q, k, v, softmax_scale, causal, deterministic)
    return out


@@ -42,13 +42,9 @@ def flash_attn_no_pad(
    seqlen = qkv.shape[1]
    nheads = qkv.shape[-2]
    x = rearrange(qkv, "b s three h d -> b s (three h d)")
-    x_unpad, indices, cu_seqlens, max_s, used_seqlens_in_batch = unpad_input(
-        x, key_padding_mask)
+    x_unpad, indices, cu_seqlens, max_s, used_seqlens_in_batch = unpad_input(x, key_padding_mask)

-    x_unpad = rearrange(x_unpad,
-                        "nnz (three h d) -> nnz three h d",
-                        three=3,
-                        h=nheads)
+    x_unpad = rearrange(x_unpad, "nnz (three h d) -> nnz three h d", three=3, h=nheads)
    output_unpad = flash_attn_varlen_qkvpacked_func(
        x_unpad,
        cu_seqlens,
@@ -88,12 +84,10 @@ def flash_attn_no_pad_v3(
    batch_size, seqlen, _, nheads, head_dim = qkv.shape
    query, key, value = qkv.unbind(dim=2)

-    query_unpad, indices, cu_seqlens_q, max_seqlen_q, _ = unpad_input(
-        rearrange(query, "b s h d -> b s (h d)"), key_padding_mask)
-    key_unpad, _, cu_seqlens_k, _, _ = unpad_input(
-        rearrange(key, "b s h d -> b s (h d)"), key_padding_mask)
-    value_unpad, _, _, _, _ = unpad_input(
-        rearrange(value, "b s h d -> b s (h d)"), key_padding_mask)
+    query_unpad, indices, cu_seqlens_q, max_seqlen_q, _ = unpad_input(rearrange(query, "b s h d -> b s (h d)"),
+                                                                      key_padding_mask)
+    key_unpad, _, cu_seqlens_k, _, _ = unpad_input(rearrange(key, "b s h d -> b s (h d)"), key_padding_mask)
+    value_unpad, _, _, _, _ = unpad_input(rearrange(value, "b s h d -> b s (h d)"), key_padding_mask)

    query_unpad = rearrange(query_unpad, "nnz (h d) -> nnz h d", h=nheads)
    key_unpad = rearrange(key_unpad, "nnz (h d) -> nnz h d", h=nheads)
@@ -138,12 +132,10 @@ def flash_attn_varlen_qk_no_pad(
 ):
    batch_size, q_seqlen, nheads, _ = query.shape

-    query_unpad, q_indices, cu_seqlens_q, max_seqlen_q, _ = unpad_input(
-        rearrange(query, "b s h d -> b s (h d)"), query_padding_mask)
-    key_unpad, _, cu_seqlens_k, max_seqlen_k, _ = unpad_input(
-        rearrange(key, "b s h d -> b s (h d)"), key_padding_mask)
-    value_unpad, _, _, _, _ = unpad_input(
-        rearrange(value, "b s h d -> b s (h d)"), key_padding_mask)
+    query_unpad, q_indices, cu_seqlens_q, max_seqlen_q, _ = unpad_input(rearrange(query, "b s h d -> b s (h d)"),
+                                                                        query_padding_mask)
+    key_unpad, _, cu_seqlens_k, max_seqlen_k, _ = unpad_input(rearrange(key, "b s h d -> b s (h d)"), key_padding_mask)
+    value_unpad, _, _, _, _ = unpad_input(rearrange(value, "b s h d -> b s (h d)"), key_padding_mask)

    query_unpad = rearrange(query_unpad, "nnz (h d) -> nnz h d", h=nheads)
    key_unpad = rearrange(key_unpad, "nnz (h d) -> nnz h d", h=nheads)
@@ -23,8 +23,7 @@ class DatasetType(str, Enum):
            return cls(value.lower())
        except ValueError:
            raise ValueError(
-                f"Invalid dataset type: {value}. Must be one of: {', '.join([m.value for m in cls])}"
-            ) from None
+                f"Invalid dataset type: {value}. Must be one of: {', '.join([m.value for m in cls])}") from None

    @classmethod
    def choices(cls) -> list[str]:
@@ -46,8 +45,7 @@ class VideoLoaderType(str, Enum):
            return cls(value.lower())
        except ValueError:
            raise ValueError(
-                f"Invalid video loader: {value}. Must be one of: {', '.join([m.value for m in cls])}"
-            ) from None
+                f"Invalid video loader: {value}. Must be one of: {', '.join([m.value for m in cls])}") from None

    @classmethod
    def choices(cls) -> list[str]:
@@ -92,8 +90,7 @@ class PreprocessConfig:
    seed: int = 42

    @staticmethod
-    def add_cli_args(parser: FlexibleArgumentParser,
-                     prefix: str = "preprocess") -> FlexibleArgumentParser:
+    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 ""

@@ -103,22 +100,19 @@ class PreprocessConfig:
                                     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-type",
-            type=str,
-            choices=DatasetType.choices(),
-            default=PreprocessConfig.dataset_type.value,
-            help="Type of the dataset")
-        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.")
+        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-type",
+                                     type=str,
+                                     choices=DatasetType.choices(),
+                                     default=PreprocessConfig.dataset_type.value,
+                                     help="Type of the dataset")
+        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(
@@ -126,13 +120,11 @@ class PreprocessConfig:
            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.")
+            "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",
@@ -145,12 +137,11 @@ class PreprocessConfig:
                                     help="How often to save to parquet files")

        # Video processing parameters
-        preprocess_args.add_argument(
-            f"--{prefix_with_dot}video-loader-type",
-            type=str,
-            choices=VideoLoaderType.choices(),
-            default=PreprocessConfig.video_loader_type.value,
-            help="Type of the video loader")
+        preprocess_args.add_argument(f"--{prefix_with_dot}video-loader-type",
+                                     type=str,
+                                     choices=VideoLoaderType.choices(),
+                                     default=PreprocessConfig.video_loader_type.value,
+                                     help="Type of the video loader")
        preprocess_args.add_argument(f"--{prefix_with_dot}max-height",
                                     type=int,
                                     default=PreprocessConfig.max_height,
@@ -163,11 +154,10 @@ class PreprocessConfig:
                                     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}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,
@@ -180,11 +170,10 @@ class PreprocessConfig:
                                     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")
+        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",
@@ -203,20 +192,15 @@ class PreprocessConfig:
        return parser

    @classmethod
-    def from_kwargs(cls, kwargs: dict[str,
-                                      Any]) -> Optional["PreprocessConfig"]:
+    def from_kwargs(cls, kwargs: dict[str, Any]) -> Optional["PreprocessConfig"]:
        """Create PreprocessConfig from keyword arguments."""
        if 'dataset_type' in kwargs and isinstance(kwargs['dataset_type'], str):
-            kwargs['dataset_type'] = DatasetType.from_string(
-                kwargs['dataset_type'])
-        if 'video_loader_type' in kwargs and isinstance(
-                kwargs['video_loader_type'], str):
-            kwargs['video_loader_type'] = VideoLoaderType.from_string(
-                kwargs['video_loader_type'])
+            kwargs['dataset_type'] = DatasetType.from_string(kwargs['dataset_type'])
+        if 'video_loader_type' in kwargs and isinstance(kwargs['video_loader_type'], str):
+            kwargs['video_loader_type'] = VideoLoaderType.from_string(kwargs['video_loader_type'])

        preprocess_config = cls()
-        if not update_config_from_args(
-                preprocess_config, kwargs, prefix="preprocess", pop_args=True):
+        if not update_config_from_args(preprocess_config, kwargs, prefix="preprocess", pop_args=True):
            return None
        return preprocess_config

@@ -2,9 +2,7 @@ from fastvideo.configs.models.base import ModelConfig
 from fastvideo.configs.models.dits.base import DiTConfig
 from fastvideo.configs.models.encoders.base import EncoderConfig
 from fastvideo.configs.models.vaes.base import VAEConfig
-from fastvideo.configs.models.audio import (LTX2AudioDecoderConfig,
-                                            LTX2AudioEncoderConfig,
-                                            LTX2VocoderConfig)
+from fastvideo.configs.models.audio import (LTX2AudioDecoderConfig, LTX2AudioEncoderConfig, LTX2VocoderConfig)
 from fastvideo.configs.models.upsamplers.base import UpsamplerConfig

 __all__ = [
@@ -15,17 +15,14 @@ class LTX2AudioArchConfig(ArchConfig):

@dataclass
 class LTX2AudioEncoderConfig(ModelConfig):
-    arch_config: ArchConfig = field(default_factory=lambda: LTX2AudioArchConfig(
-        architectures=["LTX2AudioEncoder"]))
+    arch_config: ArchConfig = field(default_factory=lambda: LTX2AudioArchConfig(architectures=["LTX2AudioEncoder"]))


@dataclass
 class LTX2AudioDecoderConfig(ModelConfig):
-    arch_config: ArchConfig = field(default_factory=lambda: LTX2AudioArchConfig(
-        architectures=["LTX2AudioDecoder"]))
+    arch_config: ArchConfig = field(default_factory=lambda: LTX2AudioArchConfig(architectures=["LTX2AudioDecoder"]))


@dataclass
 class LTX2VocoderConfig(ModelConfig):
-    arch_config: ArchConfig = field(default_factory=lambda: LTX2AudioArchConfig(
-        architectures=["LTX2Vocoder"]))
+    arch_config: ArchConfig = field(default_factory=lambda: LTX2AudioArchConfig(architectures=["LTX2Vocoder"]))
@@ -13,8 +13,7 @@ logger = init_logger(__name__)
@dataclass
 class ArchConfig:
    stacked_params_mapping: list[tuple[str, str, str]] = field(
-        default_factory=list
-    )  # mapping from huggingface weight names to custom names
+        default_factory=list)  # mapping from huggingface weight names to custom names


@dataclass
@@ -29,8 +28,7 @@ class ModelConfig:
        # Only called if 'name' is not found in ModelConfig directly
        if hasattr(self.arch_config, name):
            return getattr(self.arch_config, name)
-        raise AttributeError(
-            f"'{type(self).__name__}' object has no attribute '{name}'")
+        raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")

    def __getstate__(self):
        # Return a dictionary of attributes to pickle
@@ -51,8 +49,7 @@ class ModelConfig:
            if key in valid_fields:
                setattr(arch_config, key, value)
            else:
-                raise AttributeError(
-                    f"{type(arch_config).__name__} has no field '{key}'")
+                raise AttributeError(f"{type(arch_config).__name__} has no field '{key}'")

        if hasattr(arch_config, "__post_init__"):
            arch_config.__post_init__()
@@ -66,8 +63,7 @@ class ModelConfig:
            if key in valid_fields:
                setattr(self, key, value)
            else:
-                logger.warning("%s does not contain field '%s'!",
-                               type(self).__name__, key)
+                logger.warning("%s does not contain field '%s'!", type(self).__name__, key)
                raise AttributeError(f"Invalid field: {key}")

        if hasattr(self, "__post_init__"):
@@ -7,9 +7,9 @@ from fastvideo.configs.models.dits.longcat import LongCatVideoConfig
 from fastvideo.configs.models.dits.ltx2 import LTX2VideoConfig
 from fastvideo.configs.models.dits.wanvideo import WanVideoConfig
 from fastvideo.configs.models.dits.hyworld import HYWorldConfig
+from fastvideo.configs.models.dits.kandinsky5 import Kandinsky5VideoConfig

 __all__ = [
-    "HunyuanVideoConfig", "HunyuanVideo15Config", "HunyuanGameCraftConfig",
-    "WanVideoConfig", "CosmosVideoConfig", "Cosmos25VideoConfig",
-    "LongCatVideoConfig", "LTX2VideoConfig", "HYWorldConfig"
+    "HunyuanVideoConfig", "HunyuanVideo15Config", "HunyuanGameCraftConfig", "WanVideoConfig", "CosmosVideoConfig",
+    "Cosmos25VideoConfig", "LongCatVideoConfig", "LTX2VideoConfig", "HYWorldConfig", "Kandinsky5VideoConfig"
 ]
@@ -14,11 +14,12 @@ 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, ...] = (
-        AttentionBackendEnum.SAGE_ATTN, AttentionBackendEnum.FLASH_ATTN,
-        AttentionBackendEnum.TORCH_SDPA, AttentionBackendEnum.VIDEO_SPARSE_ATTN,
-        AttentionBackendEnum.VMOBA_ATTN, AttentionBackendEnum.SAGE_ATTN_THREE,
-        AttentionBackendEnum.SLA_ATTN, AttentionBackendEnum.SAGE_SLA_ATTN)
+    _supported_attention_backends: tuple[AttentionBackendEnum,
+                                         ...] = (AttentionBackendEnum.SAGE_ATTN, AttentionBackendEnum.FLASH_ATTN,
+                                                 AttentionBackendEnum.TORCH_SDPA,
+                                                 AttentionBackendEnum.VIDEO_SPARSE_ATTN,
+                                                 AttentionBackendEnum.VMOBA_ATTN, AttentionBackendEnum.SAGE_ATTN_THREE,
+                                                 AttentionBackendEnum.SLA_ATTN, AttentionBackendEnum.SAGE_SLA_ATTN)

    hidden_size: int = 0
    num_attention_heads: int = 0
@@ -10,8 +10,7 @@ def is_transformer_blocks(n: str, m) -> bool:

@dataclass
 class CosmosArchConfig(DiTArchConfig):
-    _fsdp_shard_conditions: list = field(
-        default_factory=lambda: [is_transformer_blocks])
+    _fsdp_shard_conditions: list = field(default_factory=lambda: [is_transformer_blocks])

    param_names_mapping: dict = field(
        default_factory=lambda: {
@@ -19,58 +18,35 @@ class CosmosArchConfig(DiTArchConfig):
            r"^time_embed\.time_proj\.(.*)$": r"time_embed.time_proj.\1",
            r"^time_embed\.t_embedder\.(.*)$": r"time_embed.t_embedder.\1",
            r"^time_embed\.norm\.(.*)$": r"time_embed.norm.\1",
-            r"^transformer_blocks\.(\d+)\.attn1\.to_q\.(.*)$":
-            r"transformer_blocks.\1.attn1.to_q.\2",
-            r"^transformer_blocks\.(\d+)\.attn1\.to_k\.(.*)$":
-            r"transformer_blocks.\1.attn1.to_k.\2",
-            r"^transformer_blocks\.(\d+)\.attn1\.to_v\.(.*)$":
-            r"transformer_blocks.\1.attn1.to_v.\2",
-            r"^transformer_blocks\.(\d+)\.attn1\.to_out\.0\.(.*)$":
-            r"transformer_blocks.\1.attn1.to_out.\2",
-            r"^transformer_blocks\.(\d+)\.attn1\.norm_q\.(.*)$":
-            r"transformer_blocks.\1.attn1.norm_q.\2",
-            r"^transformer_blocks\.(\d+)\.attn1\.norm_k\.(.*)$":
-            r"transformer_blocks.\1.attn1.norm_k.\2",
-            r"^transformer_blocks\.(\d+)\.attn2\.to_q\.(.*)$":
-            r"transformer_blocks.\1.attn2.to_q.\2",
-            r"^transformer_blocks\.(\d+)\.attn2\.to_k\.(.*)$":
-            r"transformer_blocks.\1.attn2.to_k.\2",
-            r"^transformer_blocks\.(\d+)\.attn2\.to_v\.(.*)$":
-            r"transformer_blocks.\1.attn2.to_v.\2",
-            r"^transformer_blocks\.(\d+)\.attn2\.to_out\.0\.(.*)$":
-            r"transformer_blocks.\1.attn2.to_out.\2",
-            r"^transformer_blocks\.(\d+)\.attn2\.norm_q\.(.*)$":
-            r"transformer_blocks.\1.attn2.norm_q.\2",
-            r"^transformer_blocks\.(\d+)\.attn2\.norm_k\.(.*)$":
-            r"transformer_blocks.\1.attn2.norm_k.\2",
-            r"^transformer_blocks\.(\d+)\.ff\.net\.0\.proj\.(.*)$":
-            r"transformer_blocks.\1.ff.fc_in.\2",
-            r"^transformer_blocks\.(\d+)\.ff\.net\.2\.(.*)$":
-            r"transformer_blocks.\1.ff.fc_out.\2",
+            r"^transformer_blocks\.(\d+)\.attn1\.to_q\.(.*)$": r"transformer_blocks.\1.attn1.to_q.\2",
+            r"^transformer_blocks\.(\d+)\.attn1\.to_k\.(.*)$": r"transformer_blocks.\1.attn1.to_k.\2",
+            r"^transformer_blocks\.(\d+)\.attn1\.to_v\.(.*)$": r"transformer_blocks.\1.attn1.to_v.\2",
+            r"^transformer_blocks\.(\d+)\.attn1\.to_out\.0\.(.*)$": r"transformer_blocks.\1.attn1.to_out.\2",
+            r"^transformer_blocks\.(\d+)\.attn1\.norm_q\.(.*)$": r"transformer_blocks.\1.attn1.norm_q.\2",
+            r"^transformer_blocks\.(\d+)\.attn1\.norm_k\.(.*)$": r"transformer_blocks.\1.attn1.norm_k.\2",
+            r"^transformer_blocks\.(\d+)\.attn2\.to_q\.(.*)$": r"transformer_blocks.\1.attn2.to_q.\2",
+            r"^transformer_blocks\.(\d+)\.attn2\.to_k\.(.*)$": r"transformer_blocks.\1.attn2.to_k.\2",
+            r"^transformer_blocks\.(\d+)\.attn2\.to_v\.(.*)$": r"transformer_blocks.\1.attn2.to_v.\2",
+            r"^transformer_blocks\.(\d+)\.attn2\.to_out\.0\.(.*)$": r"transformer_blocks.\1.attn2.to_out.\2",
+            r"^transformer_blocks\.(\d+)\.attn2\.norm_q\.(.*)$": r"transformer_blocks.\1.attn2.norm_q.\2",
+            r"^transformer_blocks\.(\d+)\.attn2\.norm_k\.(.*)$": r"transformer_blocks.\1.attn2.norm_k.\2",
+            r"^transformer_blocks\.(\d+)\.ff\.net\.0\.proj\.(.*)$": r"transformer_blocks.\1.ff.fc_in.\2",
+            r"^transformer_blocks\.(\d+)\.ff\.net\.2\.(.*)$": r"transformer_blocks.\1.ff.fc_out.\2",
            r"^norm_out\.(.*)$": r"norm_out.\1",
            r"^proj_out\.(.*)$": r"proj_out.\1",
        })

    lora_param_names_mapping: dict = field(
        default_factory=lambda: {
-            r"^transformer_blocks\.(\d+)\.attn1\.to_q\.(.*)$":
-            r"transformer_blocks.\1.attn1.to_q.\2",
-            r"^transformer_blocks\.(\d+)\.attn1\.to_k\.(.*)$":
-            r"transformer_blocks.\1.attn1.to_k.\2",
-            r"^transformer_blocks\.(\d+)\.attn1\.to_v\.(.*)$":
-            r"transformer_blocks.\1.attn1.to_v.\2",
-            r"^transformer_blocks\.(\d+)\.attn1\.to_out\.(.*)$":
-            r"transformer_blocks.\1.attn1.to_out.\2",
-            r"^transformer_blocks\.(\d+)\.attn2\.to_q\.(.*)$":
-            r"transformer_blocks.\1.attn2.to_q.\2",
-            r"^transformer_blocks\.(\d+)\.attn2\.to_k\.(.*)$":
-            r"transformer_blocks.\1.attn2.to_k.\2",
-            r"^transformer_blocks\.(\d+)\.attn2\.to_v\.(.*)$":
-            r"transformer_blocks.\1.attn2.to_v.\2",
-            r"^transformer_blocks\.(\d+)\.attn2\.to_out\.(.*)$":
-            r"transformer_blocks.\1.attn2.to_out.\2",
-            r"^transformer_blocks\.(\d+)\.ff\.(.*)$":
-            r"transformer_blocks.\1.ff.\2",
+            r"^transformer_blocks\.(\d+)\.attn1\.to_q\.(.*)$": r"transformer_blocks.\1.attn1.to_q.\2",
+            r"^transformer_blocks\.(\d+)\.attn1\.to_k\.(.*)$": r"transformer_blocks.\1.attn1.to_k.\2",
+            r"^transformer_blocks\.(\d+)\.attn1\.to_v\.(.*)$": r"transformer_blocks.\1.attn1.to_v.\2",
+            r"^transformer_blocks\.(\d+)\.attn1\.to_out\.(.*)$": r"transformer_blocks.\1.attn1.to_out.\2",
+            r"^transformer_blocks\.(\d+)\.attn2\.to_q\.(.*)$": r"transformer_blocks.\1.attn2.to_q.\2",
+            r"^transformer_blocks\.(\d+)\.attn2\.to_k\.(.*)$": r"transformer_blocks.\1.attn2.to_k.\2",
+            r"^transformer_blocks\.(\d+)\.attn2\.to_v\.(.*)$": r"transformer_blocks.\1.attn2.to_v.\2",
+            r"^transformer_blocks\.(\d+)\.attn2\.to_out\.(.*)$": r"transformer_blocks.\1.attn2.to_out.\2",
+            r"^transformer_blocks\.(\d+)\.ff\.(.*)$": r"transformer_blocks.\1.ff.\2",
        })

    # Cosmos-specific config parameters based on transformer_cosmos.py
@@ -12,46 +12,33 @@ def is_transformer_blocks(n: str, m) -> bool:
 class Cosmos25ArchConfig(DiTArchConfig):
    """Configuration for Cosmos 2.5 architecture (MiniTrainDIT)."""

-    _fsdp_shard_conditions: list = field(
-        default_factory=lambda: [is_transformer_blocks])
+    _fsdp_shard_conditions: list = field(default_factory=lambda: [is_transformer_blocks])

    param_names_mapping: dict = field(
        default_factory=lambda: {
            # Remove "net." prefix and map official structure to FastVideo
            # Patch embedding: net.x_embedder.proj.1.weight -> patch_embed.proj.weight
-            r"^net\.x_embedder\.proj\.1\.(.*)$":
-            r"patch_embed.proj.\1",
+            r"^net\.x_embedder\.proj\.1\.(.*)$": r"patch_embed.proj.\1",

            # Time embedding: net.t_embedder.1.linear_1.weight -> time_embed.t_embedder.linear_1.weight
-            r"^net\.t_embedder\.1\.linear_1\.(.*)$":
-            r"time_embed.t_embedder.linear_1.\1",
-            r"^net\.t_embedder\.1\.linear_2\.(.*)$":
-            r"time_embed.t_embedder.linear_2.\1",
+            r"^net\.t_embedder\.1\.linear_1\.(.*)$": r"time_embed.t_embedder.linear_1.\1",
+            r"^net\.t_embedder\.1\.linear_2\.(.*)$": r"time_embed.t_embedder.linear_2.\1",
            # Time embedding norm: net.t_embedding_norm.weight -> time_embed.norm.weight
            # Note: This also handles _extra_state if present
-            r"^net\.t_embedding_norm\.(.*)$":
-            r"time_embed.norm.\1",
+            r"^net\.t_embedding_norm\.(.*)$": r"time_embed.norm.\1",

            # Cross-attention projection (optional): net.crossattn_proj.0.weight -> crossattn_proj.0.weight
-            r"^net\.crossattn_proj\.0\.weight$":
-            r"crossattn_proj.0.weight",
-            r"^net\.crossattn_proj\.0\.bias$":
-            r"crossattn_proj.0.bias",
+            r"^net\.crossattn_proj\.0\.weight$": r"crossattn_proj.0.weight",
+            r"^net\.crossattn_proj\.0\.bias$": r"crossattn_proj.0.bias",

            # Transformer blocks: net.blocks.N -> transformer_blocks.N
            # Self-attention (self_attn -> attn1)
-            r"^net\.blocks\.(\d+)\.self_attn\.q_proj\.(.*)$":
-            r"transformer_blocks.\1.attn1.to_q.\2",
-            r"^net\.blocks\.(\d+)\.self_attn\.k_proj\.(.*)$":
-            r"transformer_blocks.\1.attn1.to_k.\2",
-            r"^net\.blocks\.(\d+)\.self_attn\.v_proj\.(.*)$":
-            r"transformer_blocks.\1.attn1.to_v.\2",
-            r"^net\.blocks\.(\d+)\.self_attn\.output_proj\.(.*)$":
-            r"transformer_blocks.\1.attn1.to_out.\2",
-            r"^net\.blocks\.(\d+)\.self_attn\.q_norm\.weight$":
-            r"transformer_blocks.\1.attn1.norm_q.weight",
-            r"^net\.blocks\.(\d+)\.self_attn\.k_norm\.weight$":
-            r"transformer_blocks.\1.attn1.norm_k.weight",
+            r"^net\.blocks\.(\d+)\.self_attn\.q_proj\.(.*)$": r"transformer_blocks.\1.attn1.to_q.\2",
+            r"^net\.blocks\.(\d+)\.self_attn\.k_proj\.(.*)$": r"transformer_blocks.\1.attn1.to_k.\2",
+            r"^net\.blocks\.(\d+)\.self_attn\.v_proj\.(.*)$": r"transformer_blocks.\1.attn1.to_v.\2",
+            r"^net\.blocks\.(\d+)\.self_attn\.output_proj\.(.*)$": r"transformer_blocks.\1.attn1.to_out.\2",
+            r"^net\.blocks\.(\d+)\.self_attn\.q_norm\.weight$": r"transformer_blocks.\1.attn1.norm_q.weight",
+            r"^net\.blocks\.(\d+)\.self_attn\.k_norm\.weight$": r"transformer_blocks.\1.attn1.norm_k.weight",
            # RMSNorm _extra_state keys (internal PyTorch state, will be recomputed automatically)
            r"^net\.blocks\.(\d+)\.self_attn\.q_norm\._extra_state$":
            r"transformer_blocks.\1.attn1.norm_q._extra_state",
@@ -59,18 +46,12 @@ class Cosmos25ArchConfig(DiTArchConfig):
            r"transformer_blocks.\1.attn1.norm_k._extra_state",

            # Cross-attention (cross_attn -> attn2)
-            r"^net\.blocks\.(\d+)\.cross_attn\.q_proj\.(.*)$":
-            r"transformer_blocks.\1.attn2.to_q.\2",
-            r"^net\.blocks\.(\d+)\.cross_attn\.k_proj\.(.*)$":
-            r"transformer_blocks.\1.attn2.to_k.\2",
-            r"^net\.blocks\.(\d+)\.cross_attn\.v_proj\.(.*)$":
-            r"transformer_blocks.\1.attn2.to_v.\2",
-            r"^net\.blocks\.(\d+)\.cross_attn\.output_proj\.(.*)$":
-            r"transformer_blocks.\1.attn2.to_out.\2",
-            r"^net\.blocks\.(\d+)\.cross_attn\.q_norm\.weight$":
-            r"transformer_blocks.\1.attn2.norm_q.weight",
-            r"^net\.blocks\.(\d+)\.cross_attn\.k_norm\.weight$":
-            r"transformer_blocks.\1.attn2.norm_k.weight",
+            r"^net\.blocks\.(\d+)\.cross_attn\.q_proj\.(.*)$": r"transformer_blocks.\1.attn2.to_q.\2",
+            r"^net\.blocks\.(\d+)\.cross_attn\.k_proj\.(.*)$": r"transformer_blocks.\1.attn2.to_k.\2",
+            r"^net\.blocks\.(\d+)\.cross_attn\.v_proj\.(.*)$": r"transformer_blocks.\1.attn2.to_v.\2",
+            r"^net\.blocks\.(\d+)\.cross_attn\.output_proj\.(.*)$": r"transformer_blocks.\1.attn2.to_out.\2",
+            r"^net\.blocks\.(\d+)\.cross_attn\.q_norm\.weight$": r"transformer_blocks.\1.attn2.norm_q.weight",
+            r"^net\.blocks\.(\d+)\.cross_attn\.k_norm\.weight$": r"transformer_blocks.\1.attn2.norm_k.weight",
            # RMSNorm _extra_state keys for cross-attention
            r"^net\.blocks\.(\d+)\.cross_attn\.q_norm\._extra_state$":
            r"transformer_blocks.\1.attn2.norm_q._extra_state",
@@ -78,10 +59,8 @@ class Cosmos25ArchConfig(DiTArchConfig):
            r"transformer_blocks.\1.attn2.norm_k._extra_state",

            # MLP: net.blocks.N.mlp.layer1 -> transformer_blocks.N.mlp.fc_in
-            r"^net\.blocks\.(\d+)\.mlp\.layer1\.(.*)$":
-            r"transformer_blocks.\1.mlp.fc_in.\2",
-            r"^net\.blocks\.(\d+)\.mlp\.layer2\.(.*)$":
-            r"transformer_blocks.\1.mlp.fc_out.\2",
+            r"^net\.blocks\.(\d+)\.mlp\.layer1\.(.*)$": r"transformer_blocks.\1.mlp.fc_in.\2",
+            r"^net\.blocks\.(\d+)\.mlp\.layer2\.(.*)$": r"transformer_blocks.\1.mlp.fc_out.\2",

            # AdaLN-LoRA modulations: net.blocks.N.adaln_modulation_* -> transformer_blocks.N.adaln_modulation_*
            # These are now at the block level, not inside norm layers
@@ -93,27 +72,21 @@ class Cosmos25ArchConfig(DiTArchConfig):
            r"transformer_blocks.\1.adaln_modulation_cross_attn.1.\2",
            r"^net\.blocks\.(\d+)\.adaln_modulation_cross_attn\.2\.(.*)$":
            r"transformer_blocks.\1.adaln_modulation_cross_attn.2.\2",
-            r"^net\.blocks\.(\d+)\.adaln_modulation_mlp\.1\.(.*)$":
-            r"transformer_blocks.\1.adaln_modulation_mlp.1.\2",
-            r"^net\.blocks\.(\d+)\.adaln_modulation_mlp\.2\.(.*)$":
-            r"transformer_blocks.\1.adaln_modulation_mlp.2.\2",
+            r"^net\.blocks\.(\d+)\.adaln_modulation_mlp\.1\.(.*)$": r"transformer_blocks.\1.adaln_modulation_mlp.1.\2",
+            r"^net\.blocks\.(\d+)\.adaln_modulation_mlp\.2\.(.*)$": r"transformer_blocks.\1.adaln_modulation_mlp.2.\2",

            # Layer norms: net.blocks.N.layer_norm_* -> transformer_blocks.N.norm*.norm
            r"^net\.blocks\.(\d+)\.layer_norm_self_attn\._extra_state$":
            r"transformer_blocks.\1.norm1.norm._extra_state",
            r"^net\.blocks\.(\d+)\.layer_norm_cross_attn\._extra_state$":
            r"transformer_blocks.\1.norm2.norm._extra_state",
-            r"^net\.blocks\.(\d+)\.layer_norm_mlp\._extra_state$":
-            r"transformer_blocks.\1.norm3.norm._extra_state",
+            r"^net\.blocks\.(\d+)\.layer_norm_mlp\._extra_state$": r"transformer_blocks.\1.norm3.norm._extra_state",

            # Final layer: net.final_layer.linear -> final_layer.proj_out
-            r"^net\.final_layer\.linear\.(.*)$":
-            r"final_layer.proj_out.\1",
+            r"^net\.final_layer\.linear\.(.*)$": r"final_layer.proj_out.\1",
            # Final layer AdaLN-LoRA: net.final_layer.adaln_modulation -> final_layer.linear_*
-            r"^net\.final_layer\.adaln_modulation\.1\.(.*)$":
-            r"final_layer.linear_1.\1",
-            r"^net\.final_layer\.adaln_modulation\.2\.(.*)$":
-            r"final_layer.linear_2.\1",
+            r"^net\.final_layer\.adaln_modulation\.1\.(.*)$": r"final_layer.linear_1.\1",
+            r"^net\.final_layer\.adaln_modulation\.2\.(.*)$": r"final_layer.linear_2.\1",

            # Note: The following keys from official checkpoint are NOT mapped and can be safely ignored:
            # - net.pos_embedder.* (seq, dim_spatial_range, dim_temporal_range) - These are computed dynamically
@@ -123,24 +96,15 @@ class Cosmos25ArchConfig(DiTArchConfig):

    lora_param_names_mapping: dict = field(
        default_factory=lambda: {
-            r"^transformer_blocks\.(\d+)\.attn1\.to_q\.(.*)$":
-            r"transformer_blocks.\1.attn1.to_q.\2",
-            r"^transformer_blocks\.(\d+)\.attn1\.to_k\.(.*)$":
-            r"transformer_blocks.\1.attn1.to_k.\2",
-            r"^transformer_blocks\.(\d+)\.attn1\.to_v\.(.*)$":
-            r"transformer_blocks.\1.attn1.to_v.\2",
-            r"^transformer_blocks\.(\d+)\.attn1\.to_out\.(.*)$":
-            r"transformer_blocks.\1.attn1.to_out.\2",
-            r"^transformer_blocks\.(\d+)\.attn2\.to_q\.(.*)$":
-            r"transformer_blocks.\1.attn2.to_q.\2",
-            r"^transformer_blocks\.(\d+)\.attn2\.to_k\.(.*)$":
-            r"transformer_blocks.\1.attn2.to_k.\2",
-            r"^transformer_blocks\.(\d+)\.attn2\.to_v\.(.*)$":
-            r"transformer_blocks.\1.attn2.to_v.\2",
-            r"^transformer_blocks\.(\d+)\.attn2\.to_out\.(.*)$":
-            r"transformer_blocks.\1.attn2.to_out.\2",
-            r"^transformer_blocks\.(\d+)\.mlp\.(.*)$":
-            r"transformer_blocks.\1.mlp.\2",
+            r"^transformer_blocks\.(\d+)\.attn1\.to_q\.(.*)$": r"transformer_blocks.\1.attn1.to_q.\2",
+            r"^transformer_blocks\.(\d+)\.attn1\.to_k\.(.*)$": r"transformer_blocks.\1.attn1.to_k.\2",
+            r"^transformer_blocks\.(\d+)\.attn1\.to_v\.(.*)$": r"transformer_blocks.\1.attn1.to_v.\2",
+            r"^transformer_blocks\.(\d+)\.attn1\.to_out\.(.*)$": r"transformer_blocks.\1.attn1.to_out.\2",
+            r"^transformer_blocks\.(\d+)\.attn2\.to_q\.(.*)$": r"transformer_blocks.\1.attn2.to_q.\2",
+            r"^transformer_blocks\.(\d+)\.attn2\.to_k\.(.*)$": r"transformer_blocks.\1.attn2.to_k.\2",
+            r"^transformer_blocks\.(\d+)\.attn2\.to_v\.(.*)$": r"transformer_blocks.\1.attn2.to_v.\2",
+            r"^transformer_blocks\.(\d+)\.attn2\.to_out\.(.*)$": r"transformer_blocks.\1.attn2.to_out.\2",
+            r"^transformer_blocks\.(\d+)\.mlp\.(.*)$": r"transformer_blocks.\1.mlp.\2",
        })

    # Cosmos 2.5 specific config parameters
@@ -45,63 +45,45 @@ class HunyuanGameCraftArchConfig(DiTArchConfig):
    camera_net: bool = True

    _fsdp_shard_conditions: list = field(
-        default_factory=lambda:
-        [is_double_block, is_single_block, is_refiner_block, is_camera_net])
+        default_factory=lambda: [is_double_block, is_single_block, is_refiner_block, is_camera_net])

-    _compile_conditions: list = field(
-        default_factory=lambda: [is_double_block, is_single_block, is_txt_in])
+    _compile_conditions: list = field(default_factory=lambda: [is_double_block, is_single_block, is_txt_in])

    # Parameter names mapping from official checkpoint to FastVideo naming
    # GameCraft weights are already close to FastVideo format with minor adjustments
    param_names_mapping: dict = field(
        default_factory=lambda: {
            # MLP naming: fc1 -> fc_in, fc2 -> fc_out
-            r"^(.*)\.img_mlp\.fc1\.(.*)$":
-            r"\1.img_mlp.fc_in.\2",
-            r"^(.*)\.img_mlp\.fc2\.(.*)$":
-            r"\1.img_mlp.fc_out.\2",
-            r"^(.*)\.txt_mlp\.fc1\.(.*)$":
-            r"\1.txt_mlp.fc_in.\2",
-            r"^(.*)\.txt_mlp\.fc2\.(.*)$":
-            r"\1.txt_mlp.fc_out.\2",
+            r"^(.*)\.img_mlp\.fc1\.(.*)$": r"\1.img_mlp.fc_in.\2",
+            r"^(.*)\.img_mlp\.fc2\.(.*)$": r"\1.img_mlp.fc_out.\2",
+            r"^(.*)\.txt_mlp\.fc1\.(.*)$": r"\1.txt_mlp.fc_in.\2",
+            r"^(.*)\.txt_mlp\.fc2\.(.*)$": r"\1.txt_mlp.fc_out.\2",

            # Single block MLP naming
-            r"^single_blocks\.(\d+)\.mlp\.fc1\.(.*)$":
-            r"single_blocks.\1.mlp.fc_in.\2",
-            r"^single_blocks\.(\d+)\.mlp\.fc2\.(.*)$":
-            r"single_blocks.\1.mlp.fc_out.\2",
+            r"^single_blocks\.(\d+)\.mlp\.fc1\.(.*)$": r"single_blocks.\1.mlp.fc_in.\2",
+            r"^single_blocks\.(\d+)\.mlp\.fc2\.(.*)$": r"single_blocks.\1.mlp.fc_out.\2",

            # Token refiner naming
-            r"^txt_in\.individual_token_refiner\.blocks\.(\d+)\.(.*)$":
-            r"txt_in.refiner_blocks.\1.\2",
+            r"^txt_in\.individual_token_refiner\.blocks\.(\d+)\.(.*)$": r"txt_in.refiner_blocks.\1.\2",

            # Vector in naming
-            r"^vector_in\.in_layer\.(.*)$":
-            r"vector_in.fc_in.\1",
-            r"^vector_in\.out_layer\.(.*)$":
-            r"vector_in.fc_out.\1",
+            r"^vector_in\.in_layer\.(.*)$": r"vector_in.fc_in.\1",
+            r"^vector_in\.out_layer\.(.*)$": r"vector_in.fc_out.\1",

            # Time embedder naming
-            r"^time_in\.mlp\.0\.(.*)$":
-            r"time_in.mlp.fc_in.\1",
-            r"^time_in\.mlp\.2\.(.*)$":
-            r"time_in.mlp.fc_out.\1",
+            r"^time_in\.mlp\.0\.(.*)$": r"time_in.mlp.fc_in.\1",
+            r"^time_in\.mlp\.2\.(.*)$": r"time_in.mlp.fc_out.\1",

            # Guidance embedder naming (if present)
-            r"^guidance_in\.mlp\.0\.(.*)$":
-            r"guidance_in.mlp.fc_in.\1",
-            r"^guidance_in\.mlp\.2\.(.*)$":
-            r"guidance_in.mlp.fc_out.\1",
+            r"^guidance_in\.mlp\.0\.(.*)$": r"guidance_in.mlp.fc_in.\1",
+            r"^guidance_in\.mlp\.2\.(.*)$": r"guidance_in.mlp.fc_out.\1",

            # Final layer adaLN modulation
-            r"^final_layer\.adaLN_modulation\.1\.(.*)$":
-            r"final_layer.adaLN_modulation.linear.\1",
+            r"^final_layer\.adaLN_modulation\.1\.(.*)$": r"final_layer.adaLN_modulation.linear.\1",

            # Refiner block MLP naming
-            r"^txt_in\.refiner_blocks\.(\d+)\.mlp\.fc1\.(.*)$":
-            r"txt_in.refiner_blocks.\1.mlp.fc_in.\2",
-            r"^txt_in\.refiner_blocks\.(\d+)\.mlp\.fc2\.(.*)$":
-            r"txt_in.refiner_blocks.\1.mlp.fc_out.\2",
+            r"^txt_in\.refiner_blocks\.(\d+)\.mlp\.fc1\.(.*)$": r"txt_in.refiner_blocks.\1.mlp.fc_in.\2",
+            r"^txt_in\.refiner_blocks\.(\d+)\.mlp\.fc2\.(.*)$": r"txt_in.refiner_blocks.\1.mlp.fc_out.\2",

            # Camera net weights are already correctly named
        })
@@ -136,8 +118,7 @@ class HunyuanGameCraftArchConfig(DiTArchConfig):

    # Layers to exclude from LoRA
    exclude_lora_layers: list[str] = field(
-        default_factory=lambda:
-        ["img_in", "txt_in", "time_in", "vector_in", "camera_net"])
+        default_factory=lambda: ["img_in", "txt_in", "time_in", "vector_in", "camera_net"])

    def __post_init__(self):
        super().__post_init__()
@@ -157,7 +138,6 @@ class HunyuanGameCraftArchConfig(DiTArchConfig):
 class HunyuanGameCraftConfig(DiTConfig):
    """Full config for HunyuanGameCraft model."""

-    arch_config: DiTArchConfig = field(
-        default_factory=HunyuanGameCraftArchConfig)
+    arch_config: DiTArchConfig = field(default_factory=HunyuanGameCraftArchConfig)

    prefix: str = "HunyuanGameCraft"
@@ -24,12 +24,9 @@ def is_txt_in(n: str, m) -> bool:

@dataclass
 class HunyuanVideoArchConfig(DiTArchConfig):
-    _fsdp_shard_conditions: list = field(
-        default_factory=lambda:
-        [is_double_block, is_single_block, is_refiner_block])
+    _fsdp_shard_conditions: list = field(default_factory=lambda: [is_double_block, is_single_block, is_refiner_block])

-    _compile_conditions: list = field(
-        default_factory=lambda: [is_double_block, is_single_block, is_txt_in])
+    _compile_conditions: list = field(default_factory=lambda: [is_double_block, is_single_block, is_txt_in])

    param_names_mapping: dict = field(
        default_factory=lambda: {
@@ -90,18 +87,12 @@ class HunyuanVideoArchConfig(DiTArchConfig):
            r"double_blocks.\1.img_attn_q_norm.\2",
            r"^transformer_blocks\.(\d+)\.attn\.norm_k\.(.*)$":
            r"double_blocks.\1.img_attn_k_norm.\2",
-            r"^transformer_blocks\.(\d+)\.attn\.to_q\.(.*)$":
-            (r"double_blocks.\1.img_attn_qkv.\2", 0, 3),
-            r"^transformer_blocks\.(\d+)\.attn\.to_k\.(.*)$":
-            (r"double_blocks.\1.img_attn_qkv.\2", 1, 3),
-            r"^transformer_blocks\.(\d+)\.attn\.to_v\.(.*)$":
-            (r"double_blocks.\1.img_attn_qkv.\2", 2, 3),
-            r"^transformer_blocks\.(\d+)\.attn\.add_q_proj\.(.*)$":
-            (r"double_blocks.\1.txt_attn_qkv.\2", 0, 3),
-            r"^transformer_blocks\.(\d+)\.attn\.add_k_proj\.(.*)$":
-            (r"double_blocks.\1.txt_attn_qkv.\2", 1, 3),
-            r"^transformer_blocks\.(\d+)\.attn\.add_v_proj\.(.*)$":
-            (r"double_blocks.\1.txt_attn_qkv.\2", 2, 3),
+            r"^transformer_blocks\.(\d+)\.attn\.to_q\.(.*)$": (r"double_blocks.\1.img_attn_qkv.\2", 0, 3),
+            r"^transformer_blocks\.(\d+)\.attn\.to_k\.(.*)$": (r"double_blocks.\1.img_attn_qkv.\2", 1, 3),
+            r"^transformer_blocks\.(\d+)\.attn\.to_v\.(.*)$": (r"double_blocks.\1.img_attn_qkv.\2", 2, 3),
+            r"^transformer_blocks\.(\d+)\.attn\.add_q_proj\.(.*)$": (r"double_blocks.\1.txt_attn_qkv.\2", 0, 3),
+            r"^transformer_blocks\.(\d+)\.attn\.add_k_proj\.(.*)$": (r"double_blocks.\1.txt_attn_qkv.\2", 1, 3),
+            r"^transformer_blocks\.(\d+)\.attn\.add_v_proj\.(.*)$": (r"double_blocks.\1.txt_attn_qkv.\2", 2, 3),
            r"^transformer_blocks\.(\d+)\.attn\.to_out\.0\.(.*)$":
            r"double_blocks.\1.img_attn_proj.\2",
            # Corrected: merge attn.to_add_out into the main projection.
@@ -125,14 +116,10 @@ class HunyuanVideoArchConfig(DiTArchConfig):
            r"single_blocks.\1.q_norm.\2",
            r"^single_transformer_blocks\.(\d+)\.attn\.norm_k\.(.*)$":
            r"single_blocks.\1.k_norm.\2",
-            r"^single_transformer_blocks\.(\d+)\.attn\.to_q\.(.*)$":
-            (r"single_blocks.\1.linear1.\2", 0, 4),
-            r"^single_transformer_blocks\.(\d+)\.attn\.to_k\.(.*)$":
-            (r"single_blocks.\1.linear1.\2", 1, 4),
-            r"^single_transformer_blocks\.(\d+)\.attn\.to_v\.(.*)$":
-            (r"single_blocks.\1.linear1.\2", 2, 4),
-            r"^single_transformer_blocks\.(\d+)\.proj_mlp\.(.*)$":
-            (r"single_blocks.\1.linear1.\2", 3, 4),
+            r"^single_transformer_blocks\.(\d+)\.attn\.to_q\.(.*)$": (r"single_blocks.\1.linear1.\2", 0, 4),
+            r"^single_transformer_blocks\.(\d+)\.attn\.to_k\.(.*)$": (r"single_blocks.\1.linear1.\2", 1, 4),
+            r"^single_transformer_blocks\.(\d+)\.attn\.to_v\.(.*)$": (r"single_blocks.\1.linear1.\2", 2, 4),
+            r"^single_transformer_blocks\.(\d+)\.proj_mlp\.(.*)$": (r"single_blocks.\1.linear1.\2", 3, 4),
            # Corrected: map proj_out to modulation.linear rather than a separate proj_out branch.
            r"^single_transformer_blocks\.(\d+)\.proj_out\.(.*)$":
            r"single_blocks.\1.linear2.\2",
@@ -166,8 +153,7 @@ class HunyuanVideoArchConfig(DiTArchConfig):
    pooled_projection_dim: int = 768
    rope_theta: int = 256
    qk_norm: str = "rms_norm"
-    exclude_lora_layers: list[str] = field(
-        default_factory=lambda: ["img_in", "txt_in", "time_in", "vector_in"])
+    exclude_lora_layers: list[str] = field(default_factory=lambda: ["img_in", "txt_in", "time_in", "vector_in"])

    def __post_init__(self):
        super().__post_init__()
@@ -18,11 +18,9 @@ def is_txt_in(n: str, m) -> bool:

@dataclass
 class HunyuanVideo15ArchConfig(DiTArchConfig):
-    _fsdp_shard_conditions: list = field(
-        default_factory=lambda: [is_double_block, is_refiner_block])
+    _fsdp_shard_conditions: list = field(default_factory=lambda: [is_double_block, is_refiner_block])

-    _compile_conditions: list = field(
-        default_factory=lambda: [is_double_block, is_refiner_block, is_txt_in])
+    _compile_conditions: list = field(default_factory=lambda: [is_double_block, is_refiner_block, is_txt_in])

    param_names_mapping: dict = field(
        default_factory=lambda: {
@@ -85,18 +83,12 @@ class HunyuanVideo15ArchConfig(DiTArchConfig):
            r"double_blocks.\1.img_attn_q_norm.\2",
            r"^transformer_blocks\.(\d+)\.attn\.norm_k\.(.*)$":
            r"double_blocks.\1.img_attn_k_norm.\2",
-            r"^transformer_blocks\.(\d+)\.attn\.to_q\.(.*)$":
-            (r"double_blocks.\1.img_attn_qkv.\2", 0, 3),
-            r"^transformer_blocks\.(\d+)\.attn\.to_k\.(.*)$":
-            (r"double_blocks.\1.img_attn_qkv.\2", 1, 3),
-            r"^transformer_blocks\.(\d+)\.attn\.to_v\.(.*)$":
-            (r"double_blocks.\1.img_attn_qkv.\2", 2, 3),
-            r"^transformer_blocks\.(\d+)\.attn\.add_q_proj\.(.*)$":
-            (r"double_blocks.\1.txt_attn_qkv.\2", 0, 3),
-            r"^transformer_blocks\.(\d+)\.attn\.add_k_proj\.(.*)$":
-            (r"double_blocks.\1.txt_attn_qkv.\2", 1, 3),
-            r"^transformer_blocks\.(\d+)\.attn\.add_v_proj\.(.*)$":
-            (r"double_blocks.\1.txt_attn_qkv.\2", 2, 3),
+            r"^transformer_blocks\.(\d+)\.attn\.to_q\.(.*)$": (r"double_blocks.\1.img_attn_qkv.\2", 0, 3),
+            r"^transformer_blocks\.(\d+)\.attn\.to_k\.(.*)$": (r"double_blocks.\1.img_attn_qkv.\2", 1, 3),
+            r"^transformer_blocks\.(\d+)\.attn\.to_v\.(.*)$": (r"double_blocks.\1.img_attn_qkv.\2", 2, 3),
+            r"^transformer_blocks\.(\d+)\.attn\.add_q_proj\.(.*)$": (r"double_blocks.\1.txt_attn_qkv.\2", 0, 3),
+            r"^transformer_blocks\.(\d+)\.attn\.add_k_proj\.(.*)$": (r"double_blocks.\1.txt_attn_qkv.\2", 1, 3),
+            r"^transformer_blocks\.(\d+)\.attn\.add_v_proj\.(.*)$": (r"double_blocks.\1.txt_attn_qkv.\2", 2, 3),
            r"^transformer_blocks\.(\d+)\.attn\.to_out\.0\.(.*)$":
            r"double_blocks.\1.img_attn_proj.\2",
            # Corrected: merge attn.to_add_out into the main projection.
@@ -143,8 +135,7 @@ class HunyuanVideo15ArchConfig(DiTArchConfig):
    target_size: int = 640
    task_type: str = "i2v"
    use_meanflow: bool = False
-    exclude_lora_layers: list[str] = field(
-        default_factory=lambda: ["img_in", "txt_in", "time_in", "vector_in"])
+    exclude_lora_layers: list[str] = field(default_factory=lambda: ["img_in", "txt_in", "time_in", "vector_in"])

    def __post_init__(self):
        super().__post_init__()
@@ -18,27 +18,19 @@ def is_txt_in(n: str, m) -> bool:

@dataclass
 class HYWorldArchConfig(DiTArchConfig):
-    _fsdp_shard_conditions: list = field(
-        default_factory=lambda: [is_double_block, is_refiner_block])
+    _fsdp_shard_conditions: list = field(default_factory=lambda: [is_double_block, is_refiner_block])

-    _compile_conditions: list = field(
-        default_factory=lambda: [is_double_block, is_refiner_block, is_txt_in])
+    _compile_conditions: list = field(default_factory=lambda: [is_double_block, is_refiner_block, is_txt_in])

    param_names_mapping: dict = field(
        default_factory=lambda: {
            # 1. txt_in submodules (text embedder, refiner blocks):
-            r"^txt_in\.t_embedder\.mlp\.0\.(.*)$":
-            r"txt_in.t_embedder.mlp.fc_in.\1",
-            r"^txt_in\.t_embedder\.mlp\.2\.(.*)$":
-            r"txt_in.t_embedder.mlp.fc_out.\1",
-            r"^txt_in\.c_embedder\.linear_1\.(.*)$":
-            r"txt_in.c_embedder.fc_in.\1",
-            r"^txt_in\.c_embedder\.linear_2\.(.*)$":
-            r"txt_in.c_embedder.fc_out.\1",
-            r"^txt_in\.individual_token_refiner\.blocks\.(\d+)\.norm1\.(.*)$":
-            r"txt_in.refiner_blocks.\1.norm1.\2",
-            r"^txt_in\.individual_token_refiner\.blocks\.(\d+)\.norm2\.(.*)$":
-            r"txt_in.refiner_blocks.\1.norm2.\2",
+            r"^txt_in\.t_embedder\.mlp\.0\.(.*)$": r"txt_in.t_embedder.mlp.fc_in.\1",
+            r"^txt_in\.t_embedder\.mlp\.2\.(.*)$": r"txt_in.t_embedder.mlp.fc_out.\1",
+            r"^txt_in\.c_embedder\.linear_1\.(.*)$": r"txt_in.c_embedder.fc_in.\1",
+            r"^txt_in\.c_embedder\.linear_2\.(.*)$": r"txt_in.c_embedder.fc_out.\1",
+            r"^txt_in\.individual_token_refiner\.blocks\.(\d+)\.norm1\.(.*)$": r"txt_in.refiner_blocks.\1.norm1.\2",
+            r"^txt_in\.individual_token_refiner\.blocks\.(\d+)\.norm2\.(.*)$": r"txt_in.refiner_blocks.\1.norm2.\2",
            r"^txt_in\.individual_token_refiner\.blocks\.(\d+)\.self_attn_qkv\.(.*)$":
            r"txt_in.refiner_blocks.\1.self_attn_qkv.\2",
            r"^txt_in\.individual_token_refiner\.blocks\.(\d+)\.self_attn_proj\.(.*)$":
@@ -51,66 +43,42 @@ class HYWorldArchConfig(DiTArchConfig):
            r"txt_in.refiner_blocks.\1.adaLN_modulation.linear.\2",

            # 2. time_in mappings:
-            r"^time_in\.mlp\.0\.(.*)$":
-            r"time_in.timestep_embedder.mlp.fc_in.\1",
-            r"^time_in\.mlp\.2\.(.*)$":
-            r"time_in.timestep_embedder.mlp.fc_out.\1",
+            r"^time_in\.mlp\.0\.(.*)$": r"time_in.timestep_embedder.mlp.fc_in.\1",
+            r"^time_in\.mlp\.2\.(.*)$": r"time_in.timestep_embedder.mlp.fc_out.\1",

            # 3. action_in mappings:
-            r"^action_in\.mlp\.0\.(.*)$":
-            r"action_in.mlp.fc_in.\1",
-            r"^action_in\.mlp\.2\.(.*)$":
-            r"action_in.mlp.fc_out.\1",
+            r"^action_in\.mlp\.0\.(.*)$": r"action_in.mlp.fc_in.\1",
+            r"^action_in\.mlp\.2\.(.*)$": r"action_in.mlp.fc_out.\1",

            # 4. byt5_in -> txt_in_2 mappings:
-            r"^byt5_in\.layernorm\.(.*)$":
-            r"txt_in_2.norm.\1",
-            r"^byt5_in\.fc1\.(.*)$":
-            r"txt_in_2.linear_1.\1",
-            r"^byt5_in\.fc2\.(.*)$":
-            r"txt_in_2.linear_2.\1",
-            r"^byt5_in\.fc3\.(.*)$":
-            r"txt_in_2.linear_3.\1",
+            r"^byt5_in\.layernorm\.(.*)$": r"txt_in_2.norm.\1",
+            r"^byt5_in\.fc1\.(.*)$": r"txt_in_2.linear_1.\1",
+            r"^byt5_in\.fc2\.(.*)$": r"txt_in_2.linear_2.\1",
+            r"^byt5_in\.fc3\.(.*)$": r"txt_in_2.linear_3.\1",

            # 5. cond_type_embedding -> cond_type_embed:
-            r"^cond_type_embedding\.(.*)$":
-            r"cond_type_embed.\1",
+            r"^cond_type_embedding\.(.*)$": r"cond_type_embed.\1",

            # 6. vision_in -> image_embedder mappings:
-            r"^vision_in\.proj\.0\.(.*)$":
-            r"image_embedder.norm_in.\1",
-            r"^vision_in\.proj\.1\.(.*)$":
-            r"image_embedder.linear_1.\1",
-            r"^vision_in\.proj\.3\.(.*)$":
-            r"image_embedder.linear_2.\1",
-            r"^vision_in\.proj\.4\.(.*)$":
-            r"image_embedder.norm_out.\1",
+            r"^vision_in\.proj\.0\.(.*)$": r"image_embedder.norm_in.\1",
+            r"^vision_in\.proj\.1\.(.*)$": r"image_embedder.linear_1.\1",
+            r"^vision_in\.proj\.3\.(.*)$": r"image_embedder.linear_2.\1",
+            r"^vision_in\.proj\.4\.(.*)$": r"image_embedder.norm_out.\1",

            # 7. double_blocks mapping:
-            r"^double_blocks\.(\d+)\.img_attn_q\.(.*)$":
-            (r"double_blocks.\1.img_attn_qkv.\2", 0, 3),
-            r"^double_blocks\.(\d+)\.img_attn_k\.(.*)$":
-            (r"double_blocks.\1.img_attn_qkv.\2", 1, 3),
-            r"^double_blocks\.(\d+)\.img_attn_v\.(.*)$":
-            (r"double_blocks.\1.img_attn_qkv.\2", 2, 3),
-            r"^double_blocks\.(\d+)\.txt_attn_q\.(.*)$":
-            (r"double_blocks.\1.txt_attn_qkv.\2", 0, 3),
-            r"^double_blocks\.(\d+)\.txt_attn_k\.(.*)$":
-            (r"double_blocks.\1.txt_attn_qkv.\2", 1, 3),
-            r"^double_blocks\.(\d+)\.txt_attn_v\.(.*)$":
-            (r"double_blocks.\1.txt_attn_qkv.\2", 2, 3),
-            r"^double_blocks\.(\d+)\.img_mlp\.fc1\.(.*)$":
-            r"double_blocks.\1.img_mlp.fc_in.\2",
-            r"^double_blocks\.(\d+)\.img_mlp\.fc2\.(.*)$":
-            r"double_blocks.\1.img_mlp.fc_out.\2",
-            r"^double_blocks\.(\d+)\.txt_mlp\.fc1\.(.*)$":
-            r"double_blocks.\1.txt_mlp.fc_in.\2",
-            r"^double_blocks\.(\d+)\.txt_mlp\.fc2\.(.*)$":
-            r"double_blocks.\1.txt_mlp.fc_out.\2",
+            r"^double_blocks\.(\d+)\.img_attn_q\.(.*)$": (r"double_blocks.\1.img_attn_qkv.\2", 0, 3),
+            r"^double_blocks\.(\d+)\.img_attn_k\.(.*)$": (r"double_blocks.\1.img_attn_qkv.\2", 1, 3),
+            r"^double_blocks\.(\d+)\.img_attn_v\.(.*)$": (r"double_blocks.\1.img_attn_qkv.\2", 2, 3),
+            r"^double_blocks\.(\d+)\.txt_attn_q\.(.*)$": (r"double_blocks.\1.txt_attn_qkv.\2", 0, 3),
+            r"^double_blocks\.(\d+)\.txt_attn_k\.(.*)$": (r"double_blocks.\1.txt_attn_qkv.\2", 1, 3),
+            r"^double_blocks\.(\d+)\.txt_attn_v\.(.*)$": (r"double_blocks.\1.txt_attn_qkv.\2", 2, 3),
+            r"^double_blocks\.(\d+)\.img_mlp\.fc1\.(.*)$": r"double_blocks.\1.img_mlp.fc_in.\2",
+            r"^double_blocks\.(\d+)\.img_mlp\.fc2\.(.*)$": r"double_blocks.\1.img_mlp.fc_out.\2",
+            r"^double_blocks\.(\d+)\.txt_mlp\.fc1\.(.*)$": r"double_blocks.\1.txt_mlp.fc_in.\2",
+            r"^double_blocks\.(\d+)\.txt_mlp\.fc2\.(.*)$": r"double_blocks.\1.txt_mlp.fc_out.\2",

            # 8. Final layer mapping:
-            r"^final_layer\.adaLN_modulation\.1\.(.*)$":
-            r"final_layer.adaLN_modulation.linear.\1",
+            r"^final_layer\.adaLN_modulation\.1\.(.*)$": r"final_layer.adaLN_modulation.linear.\1",
        })

    # Reverse mapping for saving checkpoints: custom -> hf
@@ -150,8 +118,7 @@ class HYWorldArchConfig(DiTArchConfig):
    ideal_resolution: str = "480p"
    ideal_task: str = "i2v"
    task_type: str = "i2v"
-    exclude_lora_layers: list[str] = field(
-        default_factory=lambda: ["img_in", "txt_in", "time_in", "vector_in"])
+    exclude_lora_layers: list[str] = field(default_factory=lambda: ["img_in", "txt_in", "time_in", "vector_in"])

    def __post_init__(self):
        super().__post_init__()
@@ -0,0 +1,65 @@
+# SPDX-License-Identifier: Apache-2.0
+from dataclasses import dataclass, field
+
+from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig
+
+
+@dataclass
+class Kandinsky5ArchConfig(DiTArchConfig):
+    _fsdp_shard_conditions: list = field(default_factory=lambda: [
+        lambda n, m:
+        ("text_transformer_blocks" in n or "visual_transformer_blocks" in n) and n.split(".")[-1].isdigit()
+    ])
+
+    # Native FastVideo implementation uses the same parameter names as diffusers
+    # except FFN internals: Diffusers FFN uses `in_layer/out_layer`, while
+    # FastVideo uses MLP `fc_in/fc_out`.
+    param_names_mapping: dict = field(
+        default_factory=lambda: {
+            r"^(.*feed_forward)\.in_layer\.(weight|bias)$": r"\1.mlp.fc_in.\2",
+            r"^(.*feed_forward)\.out_layer\.(weight|bias)$": r"\1.mlp.fc_out.\2",
+        })
+
+    reverse_param_names_mapping: dict = field(default_factory=lambda: {})
+
+    # Diffusers Kandinsky5Transformer3DModel config fields.
+    in_visual_dim: int = 4
+    in_text_dim: int = 3584
+    in_text_dim2: int = 768
+    time_dim: int = 512
+    out_visual_dim: int = 4
+    patch_size: tuple[int, int, int] = (1, 2, 2)
+    model_dim: int = 2048
+    ff_dim: int = 5120
+    num_text_blocks: int = 2
+    num_visual_blocks: int = 32
+    axes_dims: tuple[int, int, int] = (16, 24, 24)
+    visual_cond: bool = False
+    attention_type: str = "regular"
+    attention_causal: bool | None = None
+    attention_local: bool | None = None
+    attention_glob: bool | None = None
+    attention_window: int | None = None
+    attention_P: float | None = None
+    attention_wT: int | None = None
+    attention_wW: int | None = None
+    attention_wH: int | None = None
+    attention_add_sta: bool | None = None
+    attention_method: str | None = None
+
+    def __post_init__(self):
+        super().__post_init__()
+        head_dim = sum(self.axes_dims)
+        if self.model_dim % head_dim != 0:
+            raise ValueError(f"model_dim ({self.model_dim}) must be divisible by head_dim ({head_dim})")
+        self.hidden_size = self.model_dim
+        self.num_attention_heads = self.model_dim // head_dim
+        self.in_channels = self.in_visual_dim
+        self.out_channels = self.out_visual_dim
+        self.num_channels_latents = self.in_visual_dim
+
+
+@dataclass
+class Kandinsky5VideoConfig(DiTConfig):
+    arch_config: DiTArchConfig = field(default_factory=Kandinsky5ArchConfig)
+    prefix: str = "Kandinsky5"
@@ -15,20 +15,14 @@ class LingBotWorldArchConfig(DiTArchConfig):
    param_names_mapping: dict = field(
        default_factory=lambda: {
            r"^patch_embedding\.(.*)$": r"patch_embedding.proj.\1",
-            r"^patch_embedding_wancamctrl\.(.*)$":
-            r"patch_embedding_wancamctrl.proj.\1",
+            r"^patch_embedding_wancamctrl\.(.*)$": r"patch_embedding_wancamctrl.proj.\1",
            r"^c2ws_hidden_states_layer1\.(.*)$": r"c2ws_mlp.fc_in.\1",
            r"^c2ws_hidden_states_layer2\.(.*)$": r"c2ws_mlp.fc_out.\1",
-            r"^text_embedding\.0\.(.*)$":
-            r"condition_embedder.text_embedder.fc_in.\1",
-            r"^text_embedding\.2\.(.*)$":
-            r"condition_embedder.text_embedder.fc_out.\1",
-            r"^time_embedding\.0\.(.*)$":
-            r"condition_embedder.time_embedder.mlp.fc_in.\1",
-            r"^time_embedding\.2\.(.*)$":
-            r"condition_embedder.time_embedder.mlp.fc_out.\1",
-            r"^time_projection\.1\.(.*)$":
-            r"condition_embedder.time_modulation.linear.\1",
+            r"^text_embedding\.0\.(.*)$": r"condition_embedder.text_embedder.fc_in.\1",
+            r"^text_embedding\.2\.(.*)$": r"condition_embedder.text_embedder.fc_out.\1",
+            r"^time_embedding\.0\.(.*)$": r"condition_embedder.time_embedder.mlp.fc_in.\1",
+            r"^time_embedding\.2\.(.*)$": r"condition_embedder.time_embedder.mlp.fc_out.\1",
+            r"^time_projection\.1\.(.*)$": r"condition_embedder.time_modulation.linear.\1",
            r"^blocks\.(\d+)\.modulation$": r"blocks.\1.scale_shift_table",
            r"^blocks\.(\d+)\.self_attn\.q\.(.*)$": r"blocks.\1.to_q.\2",
            r"^blocks\.(\d+)\.self_attn\.k\.(.*)$": r"blocks.\1.to_k.\2",
@@ -36,27 +30,19 @@ class LingBotWorldArchConfig(DiTArchConfig):
            r"^blocks\.(\d+)\.self_attn\.o\.(.*)$": r"blocks.\1.to_out.\2",
            r"^blocks\.(\d+)\.self_attn\.norm_q\.(.*)$": r"blocks.\1.norm_q.\2",
            r"^blocks\.(\d+)\.self_attn\.norm_k\.(.*)$": r"blocks.\1.norm_k.\2",
-            r"^blocks\.(\d+)\.norm3\.(.*)$":
-            r"blocks.\1.self_attn_residual_norm.norm.\2",
+            r"^blocks\.(\d+)\.norm3\.(.*)$": r"blocks.\1.self_attn_residual_norm.norm.\2",
            r"^blocks\.(\d+)\.cross_attn\.q\.(.*)$": r"blocks.\1.attn2.to_q.\2",
            r"^blocks\.(\d+)\.cross_attn\.k\.(.*)$": r"blocks.\1.attn2.to_k.\2",
            r"^blocks\.(\d+)\.cross_attn\.v\.(.*)$": r"blocks.\1.attn2.to_v.\2",
-            r"^blocks\.(\d+)\.cross_attn\.o\.(.*)$":
-            r"blocks.\1.attn2.to_out.\2",
-            r"^blocks\.(\d+)\.cross_attn\.norm_q\.(.*)$":
-            r"blocks.\1.attn2.norm_q.\2",
-            r"^blocks\.(\d+)\.cross_attn\.norm_k\.(.*)$":
-            r"blocks.\1.attn2.norm_k.\2",
+            r"^blocks\.(\d+)\.cross_attn\.o\.(.*)$": r"blocks.\1.attn2.to_out.\2",
+            r"^blocks\.(\d+)\.cross_attn\.norm_q\.(.*)$": r"blocks.\1.attn2.norm_q.\2",
+            r"^blocks\.(\d+)\.cross_attn\.norm_k\.(.*)$": r"blocks.\1.attn2.norm_k.\2",
            r"^blocks\.(\d+)\.ffn\.0\.(.*)$": r"blocks.\1.ffn.fc_in.\2",
            r"^blocks\.(\d+)\.ffn\.2\.(.*)$": r"blocks.\1.ffn.fc_out.\2",
-            r"^blocks\.(\d+)\.cam_injector_layer1\.(.*)$":
-            r"blocks.\1.cam_conditioner.cam_injector.fc_in.\2",
-            r"^blocks\.(\d+)\.cam_injector_layer2\.(.*)$":
-            r"blocks.\1.cam_conditioner.cam_injector.fc_out.\2",
-            r"^blocks\.(\d+)\.cam_scale_layer\.(.*)$":
-            r"blocks.\1.cam_conditioner.cam_scale_layer.\2",
-            r"^blocks\.(\d+)\.cam_shift_layer\.(.*)$":
-            r"blocks.\1.cam_conditioner.cam_shift_layer.\2",
+            r"^blocks\.(\d+)\.cam_injector_layer1\.(.*)$": r"blocks.\1.cam_conditioner.cam_injector.fc_in.\2",
+            r"^blocks\.(\d+)\.cam_injector_layer2\.(.*)$": r"blocks.\1.cam_conditioner.cam_injector.fc_out.\2",
+            r"^blocks\.(\d+)\.cam_scale_layer\.(.*)$": r"blocks.\1.cam_conditioner.cam_scale_layer.\2",
+            r"^blocks\.(\d+)\.cam_shift_layer\.(.*)$": r"blocks.\1.cam_conditioner.cam_shift_layer.\2",
            r"^head\.modulation$": r"scale_shift_table",
            r"^head\.head\.(.*)$": r"proj_out.\1",
        })
@@ -18,12 +18,10 @@ def is_longcat_blocks(n: str, m) -> bool:
 class LongCatVideoArchConfig(DiTArchConfig):
    """Architecture configuration for native LongCat Video DiT."""

-    _fsdp_shard_conditions: list = field(
-        default_factory=lambda: [is_longcat_blocks])
+    _fsdp_shard_conditions: list = field(default_factory=lambda: [is_longcat_blocks])

    # Enable torch.compile for transformer blocks (major speedup!)
-    _compile_conditions: list = field(
-        default_factory=lambda: [is_longcat_blocks])
+    _compile_conditions: list = field(default_factory=lambda: [is_longcat_blocks])

    # Parameter name mapping for weight conversion
    param_names_mapping: dict = field(
@@ -36,37 +34,27 @@ class LongCatVideoArchConfig(DiTArchConfig):
            r"^y_embedder\.y_proj\.2\.(.*)$": r"caption_embedder.linear_2.\1",

            # Transformer blocks - AdaLN modulation
-            r"^blocks\.(\d+)\.adaLN_modulation\.1\.(.*)$":
-            r"blocks.\1.adaln_linear_1.\2",
+            r"^blocks\.(\d+)\.adaLN_modulation\.1\.(.*)$": r"blocks.\1.adaln_linear_1.\2",

            # Transformer blocks - Normalization
            r"^blocks\.(\d+)\.mod_norm_attn\.(.*)$": r"blocks.\1.norm_attn.\2",
            r"^blocks\.(\d+)\.mod_norm_ffn\.(.*)$": r"blocks.\1.norm_ffn.\2",
-            r"^blocks\.(\d+)\.pre_crs_attn_norm\.(.*)$":
-            r"blocks.\1.norm_cross.\2",
+            r"^blocks\.(\d+)\.pre_crs_attn_norm\.(.*)$": r"blocks.\1.norm_cross.\2",

            # Self-attention: QKV fused -> separate (will need splitting in converter)
            # Original has attn.qkv.weight -> need to split into to_q, to_k, to_v
-            r"^blocks\.(\d+)\.attn\.qkv\.(.*)$":
-            r"blocks.\1.self_attn.qkv_fused.\2",  # Marker for splitting
-            r"^blocks\.(\d+)\.attn\.proj\.(.*)$":
-            r"blocks.\1.self_attn.to_out.\2",
-            r"^blocks\.(\d+)\.attn\.q_norm\.(.*)$":
-            r"blocks.\1.self_attn.q_norm.\2",
-            r"^blocks\.(\d+)\.attn\.k_norm\.(.*)$":
-            r"blocks.\1.self_attn.k_norm.\2",
+            r"^blocks\.(\d+)\.attn\.qkv\.(.*)$": r"blocks.\1.self_attn.qkv_fused.\2",  # Marker for splitting
+            r"^blocks\.(\d+)\.attn\.proj\.(.*)$": r"blocks.\1.self_attn.to_out.\2",
+            r"^blocks\.(\d+)\.attn\.q_norm\.(.*)$": r"blocks.\1.self_attn.q_norm.\2",
+            r"^blocks\.(\d+)\.attn\.k_norm\.(.*)$": r"blocks.\1.self_attn.k_norm.\2",

            # Cross-attention
-            r"^blocks\.(\d+)\.cross_attn\.q_linear\.(.*)$":
-            r"blocks.\1.cross_attn.to_q.\2",
+            r"^blocks\.(\d+)\.cross_attn\.q_linear\.(.*)$": r"blocks.\1.cross_attn.to_q.\2",
            r"^blocks\.(\d+)\.cross_attn\.kv_linear\.(.*)$":
            r"blocks.\1.cross_attn.kv_fused.\2",  # Marker for splitting
-            r"^blocks\.(\d+)\.cross_attn\.proj\.(.*)$":
-            r"blocks.\1.cross_attn.to_out.\2",
-            r"^blocks\.(\d+)\.cross_attn\.q_norm\.(.*)$":
-            r"blocks.\1.cross_attn.q_norm.\2",
-            r"^blocks\.(\d+)\.cross_attn\.k_norm\.(.*)$":
-            r"blocks.\1.cross_attn.k_norm.\2",
+            r"^blocks\.(\d+)\.cross_attn\.proj\.(.*)$": r"blocks.\1.cross_attn.to_out.\2",
+            r"^blocks\.(\d+)\.cross_attn\.q_norm\.(.*)$": r"blocks.\1.cross_attn.q_norm.\2",
+            r"^blocks\.(\d+)\.cross_attn\.k_norm\.(.*)$": r"blocks.\1.cross_attn.k_norm.\2",

            # FFN (SwiGLU)
            r"^blocks\.(\d+)\.ffn\.w1\.(.*)$": r"blocks.\1.ffn.w1.\2",  # gate
@@ -74,8 +62,7 @@ class LongCatVideoArchConfig(DiTArchConfig):
            r"^blocks\.(\d+)\.ffn\.w3\.(.*)$": r"blocks.\1.ffn.w3.\2",  # up

            # Final layer
-            r"^final_layer\.adaLN_modulation\.1\.(.*)$":
-            r"final_layer.adaln_linear.\1",
+            r"^final_layer\.adaLN_modulation\.1\.(.*)$": r"final_layer.adaln_linear.\1",
            r"^final_layer\.norm_final\.(.*)$": r"final_layer.norm.\1",
            r"^final_layer\.linear\.(.*)$": r"final_layer.proj.\1",
        })
@@ -97,8 +84,7 @@ class LongCatVideoArchConfig(DiTArchConfig):
    num_channels_latents: int = 16

    # Patch embedding
-    patch_size: tuple[int, int,
-                      int] = (1, 2, 2)  # [T, H, W] - no temporal compression
+    patch_size: tuple[int, int, int] = (1, 2, 2)  # [T, H, W] - no temporal compression

    # Text/caption embedding
    caption_channels: int = 4096  # UMT5 d_model
@@ -121,13 +107,12 @@ class LongCatVideoArchConfig(DiTArchConfig):

    # Block Sparse Attention (BSA)
    enable_bsa: bool = False
-    bsa_params: dict | None = field(
-        default_factory=lambda: {
-            "sparsity": 0.9375,
-            "cdf_threshold": None,
-            "chunk_3d_shape_q": [4, 4, 4],
-            "chunk_3d_shape_k": [4, 4, 4],
-        })
+    bsa_params: dict | None = field(default_factory=lambda: {
+        "sparsity": 0.9375,
+        "cdf_threshold": None,
+        "chunk_3d_shape_q": [4, 4, 4],
+        "chunk_3d_shape_k": [4, 4, 4],
+    })

    # LoRA exclusions
    exclude_lora_layers: list[str] = field(default_factory=lambda: [])
@@ -19,8 +19,7 @@ def is_ltx2_blocks(name: str, _module) -> bool:
 class LTX2VideoArchConfig(DiTArchConfig):
    """Architecture configuration for LTX-2 video transformer."""

-    _fsdp_shard_conditions: list = field(
-        default_factory=lambda: [is_ltx2_blocks])
+    _fsdp_shard_conditions: list = field(default_factory=lambda: [is_ltx2_blocks])
    _compile_conditions: list = field(default_factory=lambda: [is_ltx2_blocks])

    # Parameter name mapping for weight conversion (hf/comfy -> FastVideo)
@@ -47,8 +46,7 @@ class LTX2VideoArchConfig(DiTArchConfig):
    double_precision_rope: bool = True

    positional_embedding_theta: float = 10000.0
-    positional_embedding_max_pos: list[int] = field(
-        default_factory=lambda: [20, 2048, 2048])
+    positional_embedding_max_pos: list[int] = field(default_factory=lambda: [20, 2048, 2048])
    timestep_scale_multiplier: int = 1000
    use_middle_indices_grid: bool = True

@@ -64,14 +62,12 @@ class LTX2VideoArchConfig(DiTArchConfig):
    audio_in_channels: int = 128
    audio_out_channels: int = 128
    audio_cross_attention_dim: int = 2048
-    audio_positional_embedding_max_pos: list[int] = field(
-        default_factory=lambda: [20])
+    audio_positional_embedding_max_pos: list[int] = field(default_factory=lambda: [20])
    av_ca_timestep_scale_multiplier: int = 1

    def __post_init__(self):
        super().__post_init__()
-        patch_volume = self.patch_size[0] * self.patch_size[
-            1] * self.patch_size[2]
+        patch_volume = self.patch_size[0] * self.patch_size[1] * self.patch_size[2]
        if self.in_channels is None:
            self.in_channels = self.num_channels_latents * patch_volume
        if self.out_channels is None:
@@ -10,42 +10,26 @@ class MatrixGameWanVideoArchConfig(WanVideoArchConfig):
    # because MatrixGame checkpoints already have patch_embedding.proj format
    param_names_mapping: dict = field(
        default_factory=lambda: {
-            r"^patch_embedding\.(?!proj\.)(.*)$":
-            r"patch_embedding.proj.\1",
-            r"^condition_embedder\.text_embedder\.linear_1\.(.*)$":
-            r"condition_embedder.text_embedder.fc_in.\1",
-            r"^condition_embedder\.text_embedder\.linear_2\.(.*)$":
-            r"condition_embedder.text_embedder.fc_out.\1",
-            r"^condition_embedder\.time_embedder\.linear_1\.(.*)$":
-            r"condition_embedder.time_embedder.mlp.fc_in.\1",
-            r"^condition_embedder\.time_embedder\.linear_2\.(.*)$":
-            r"condition_embedder.time_embedder.mlp.fc_out.\1",
-            r"^condition_embedder\.time_proj\.(.*)$":
-            r"condition_embedder.time_modulation.linear.\1",
+            r"^patch_embedding\.(?!proj\.)(.*)$": r"patch_embedding.proj.\1",
+            r"^condition_embedder\.text_embedder\.linear_1\.(.*)$": r"condition_embedder.text_embedder.fc_in.\1",
+            r"^condition_embedder\.text_embedder\.linear_2\.(.*)$": r"condition_embedder.text_embedder.fc_out.\1",
+            r"^condition_embedder\.time_embedder\.linear_1\.(.*)$": r"condition_embedder.time_embedder.mlp.fc_in.\1",
+            r"^condition_embedder\.time_embedder\.linear_2\.(.*)$": r"condition_embedder.time_embedder.mlp.fc_out.\1",
+            r"^condition_embedder\.time_proj\.(.*)$": r"condition_embedder.time_modulation.linear.\1",
            r"^condition_embedder\.image_embedder\.ff\.net\.0\.proj\.(.*)$":
            r"condition_embedder.image_embedder.ff.fc_in.\1",
            r"^condition_embedder\.image_embedder\.ff\.net\.2\.(.*)$":
            r"condition_embedder.image_embedder.ff.fc_out.\1",
-            r"^blocks\.(\d+)\.attn1\.to_q\.(.*)$":
-            r"blocks.\1.to_q.\2",
-            r"^blocks\.(\d+)\.attn1\.to_k\.(.*)$":
-            r"blocks.\1.to_k.\2",
-            r"^blocks\.(\d+)\.attn1\.to_v\.(.*)$":
-            r"blocks.\1.to_v.\2",
-            r"^blocks\.(\d+)\.attn1\.to_out\.0\.(.*)$":
-            r"blocks.\1.to_out.\2",
-            r"^blocks\.(\d+)\.attn1\.norm_q\.(.*)$":
-            r"blocks.\1.norm_q.\2",
-            r"^blocks\.(\d+)\.attn1\.norm_k\.(.*)$":
-            r"blocks.\1.norm_k.\2",
-            r"^blocks\.(\d+)\.attn2\.to_out\.0\.(.*)$":
-            r"blocks.\1.attn2.to_out.\2",
-            r"^blocks\.(\d+)\.ffn\.net\.0\.proj\.(.*)$":
-            r"blocks.\1.ffn.fc_in.\2",
-            r"^blocks\.(\d+)\.ffn\.net\.2\.(.*)$":
-            r"blocks.\1.ffn.fc_out.\2",
-            r"^blocks\.(\d+)\.norm2\.(.*)$":
-            r"blocks.\1.self_attn_residual_norm.norm.\2",
+            r"^blocks\.(\d+)\.attn1\.to_q\.(.*)$": r"blocks.\1.to_q.\2",
+            r"^blocks\.(\d+)\.attn1\.to_k\.(.*)$": r"blocks.\1.to_k.\2",
+            r"^blocks\.(\d+)\.attn1\.to_v\.(.*)$": r"blocks.\1.to_v.\2",
+            r"^blocks\.(\d+)\.attn1\.to_out\.0\.(.*)$": r"blocks.\1.to_out.\2",
+            r"^blocks\.(\d+)\.attn1\.norm_q\.(.*)$": r"blocks.\1.norm_q.\2",
+            r"^blocks\.(\d+)\.attn1\.norm_k\.(.*)$": r"blocks.\1.norm_k.\2",
+            r"^blocks\.(\d+)\.attn2\.to_out\.0\.(.*)$": r"blocks.\1.attn2.to_out.\2",
+            r"^blocks\.(\d+)\.ffn\.net\.0\.proj\.(.*)$": r"blocks.\1.ffn.fc_in.\2",
+            r"^blocks\.(\d+)\.ffn\.net\.2\.(.*)$": r"blocks.\1.ffn.fc_out.\2",
+            r"^blocks\.(\d+)\.norm2\.(.*)$": r"blocks.\1.self_attn_residual_norm.norm.\2",
        })

    action_config: dict = field(
@@ -84,8 +68,6 @@ def _is_transformer_block(param_name: str, module: torch.nn.Module) -> bool:

@dataclass
 class MatrixGameWanVideoConfig(WanVideoConfig):
-    arch_config: MatrixGameWanVideoArchConfig = field(
-        default_factory=MatrixGameWanVideoArchConfig)
+    arch_config: MatrixGameWanVideoArchConfig = field(default_factory=MatrixGameWanVideoArchConfig)
    prefix: str = "Wan"
-    _compile_conditions: list = field(
-        default_factory=lambda: [_is_transformer_block])
+    _compile_conditions: list = field(default_factory=lambda: [_is_transformer_block])
@@ -16,8 +16,7 @@ class SD3Transformer2DArchConfig(DiTArchConfig):
    caption_projection_dim: int = 1536
    pooled_projection_dim: int = 2048
    pos_embed_max_size: int = 384
-    dual_attention_layers: list[int] = field(
-        default_factory=lambda: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12])
+    dual_attention_layers: list[int] = field(default_factory=lambda: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12])
    qk_norm: str = "rms_norm"
    in_channels: int = 16
    out_channels: int = 16
@@ -26,6 +25,5 @@ class SD3Transformer2DArchConfig(DiTArchConfig):

@dataclass
 class SD3DiTConfig(DiTConfig):
-    arch_config: DiTArchConfig = field(
-        default_factory=SD3Transformer2DArchConfig)
+    arch_config: DiTArchConfig = field(default_factory=SD3Transformer2DArchConfig)
    prefix: str = "sd3"
@@ -14,42 +14,26 @@ class WanVideoArchConfig(DiTArchConfig):

    param_names_mapping: dict = field(
        default_factory=lambda: {
-            r"^patch_embedding\.(.*)$":
-            r"patch_embedding.proj.\1",
-            r"^condition_embedder\.text_embedder\.linear_1\.(.*)$":
-            r"condition_embedder.text_embedder.fc_in.\1",
-            r"^condition_embedder\.text_embedder\.linear_2\.(.*)$":
-            r"condition_embedder.text_embedder.fc_out.\1",
-            r"^condition_embedder\.time_embedder\.linear_1\.(.*)$":
-            r"condition_embedder.time_embedder.mlp.fc_in.\1",
-            r"^condition_embedder\.time_embedder\.linear_2\.(.*)$":
-            r"condition_embedder.time_embedder.mlp.fc_out.\1",
-            r"^condition_embedder\.time_proj\.(.*)$":
-            r"condition_embedder.time_modulation.linear.\1",
+            r"^patch_embedding\.(.*)$": r"patch_embedding.proj.\1",
+            r"^condition_embedder\.text_embedder\.linear_1\.(.*)$": r"condition_embedder.text_embedder.fc_in.\1",
+            r"^condition_embedder\.text_embedder\.linear_2\.(.*)$": r"condition_embedder.text_embedder.fc_out.\1",
+            r"^condition_embedder\.time_embedder\.linear_1\.(.*)$": r"condition_embedder.time_embedder.mlp.fc_in.\1",
+            r"^condition_embedder\.time_embedder\.linear_2\.(.*)$": r"condition_embedder.time_embedder.mlp.fc_out.\1",
+            r"^condition_embedder\.time_proj\.(.*)$": r"condition_embedder.time_modulation.linear.\1",
            r"^condition_embedder\.image_embedder\.ff\.net\.0\.proj\.(.*)$":
            r"condition_embedder.image_embedder.ff.fc_in.\1",
            r"^condition_embedder\.image_embedder\.ff\.net\.2\.(.*)$":
            r"condition_embedder.image_embedder.ff.fc_out.\1",
-            r"^blocks\.(\d+)\.attn1\.to_q\.(.*)$":
-            r"blocks.\1.to_q.\2",
-            r"^blocks\.(\d+)\.attn1\.to_k\.(.*)$":
-            r"blocks.\1.to_k.\2",
-            r"^blocks\.(\d+)\.attn1\.to_v\.(.*)$":
-            r"blocks.\1.to_v.\2",
-            r"^blocks\.(\d+)\.attn1\.to_out\.0\.(.*)$":
-            r"blocks.\1.to_out.\2",
-            r"^blocks\.(\d+)\.attn1\.norm_q\.(.*)$":
-            r"blocks.\1.norm_q.\2",
-            r"^blocks\.(\d+)\.attn1\.norm_k\.(.*)$":
-            r"blocks.\1.norm_k.\2",
-            r"^blocks\.(\d+)\.attn2\.to_out\.0\.(.*)$":
-            r"blocks.\1.attn2.to_out.\2",
-            r"^blocks\.(\d+)\.ffn\.net\.0\.proj\.(.*)$":
-            r"blocks.\1.ffn.fc_in.\2",
-            r"^blocks\.(\d+)\.ffn\.net\.2\.(.*)$":
-            r"blocks.\1.ffn.fc_out.\2",
-            r"^blocks\.(\d+)\.norm2\.(.*)$":
-            r"blocks.\1.self_attn_residual_norm.norm.\2",
+            r"^blocks\.(\d+)\.attn1\.to_q\.(.*)$": r"blocks.\1.to_q.\2",
+            r"^blocks\.(\d+)\.attn1\.to_k\.(.*)$": r"blocks.\1.to_k.\2",
+            r"^blocks\.(\d+)\.attn1\.to_v\.(.*)$": r"blocks.\1.to_v.\2",
+            r"^blocks\.(\d+)\.attn1\.to_out\.0\.(.*)$": r"blocks.\1.to_out.\2",
+            r"^blocks\.(\d+)\.attn1\.norm_q\.(.*)$": r"blocks.\1.norm_q.\2",
+            r"^blocks\.(\d+)\.attn1\.norm_k\.(.*)$": r"blocks.\1.norm_k.\2",
+            r"^blocks\.(\d+)\.attn2\.to_out\.0\.(.*)$": r"blocks.\1.attn2.to_out.\2",
+            r"^blocks\.(\d+)\.ffn\.net\.0\.proj\.(.*)$": r"blocks.\1.ffn.fc_in.\2",
+            r"^blocks\.(\d+)\.ffn\.net\.2\.(.*)$": r"blocks.\1.ffn.fc_out.\2",
+            r"^blocks\.(\d+)\.norm2\.(.*)$": r"blocks.\1.self_attn_residual_norm.norm.\2",
        })

    # Reverse mapping for saving checkpoints: custom -> hf
@@ -62,13 +46,11 @@ class WanVideoArchConfig(DiTArchConfig):
            r"^blocks\.(\d+)\.self_attn\.q\.(.*)$": r"blocks.\1.attn1.to_q.\2",
            r"^blocks\.(\d+)\.self_attn\.k\.(.*)$": r"blocks.\1.attn1.to_k.\2",
            r"^blocks\.(\d+)\.self_attn\.v\.(.*)$": r"blocks.\1.attn1.to_v.\2",
-            r"^blocks\.(\d+)\.self_attn\.o\.(.*)$":
-            r"blocks.\1.attn1.to_out.0.\2",
+            r"^blocks\.(\d+)\.self_attn\.o\.(.*)$": r"blocks.\1.attn1.to_out.0.\2",
            r"^blocks\.(\d+)\.cross_attn\.q\.(.*)$": r"blocks.\1.attn2.to_q.\2",
            r"^blocks\.(\d+)\.cross_attn\.k\.(.*)$": r"blocks.\1.attn2.to_k.\2",
            r"^blocks\.(\d+)\.cross_attn\.v\.(.*)$": r"blocks.\1.attn2.to_v.\2",
-            r"^blocks\.(\d+)\.cross_attn\.o\.(.*)$":
-            r"blocks.\1.attn2.to_out.0.\2",
+            r"^blocks\.(\d+)\.cross_attn\.o\.(.*)$": r"blocks.\1.attn2.to_out.0.\2",
            r"^blocks\.(\d+)\.ffn\.0\.(.*)$": r"blocks.\1.ffn.fc_in.\2",
            r"^blocks\.(\d+)\.ffn\.2\.(.*)$": r"blocks.\1.ffn.fc_out.\2",
        })
@@ -1,9 +1,6 @@
-from fastvideo.configs.models.encoders.base import (BaseEncoderOutput,
-                                                    EncoderConfig,
-                                                    ImageEncoderConfig,
+from fastvideo.configs.models.encoders.base import (BaseEncoderOutput, EncoderConfig, ImageEncoderConfig,
                                                    TextEncoderConfig)
-from fastvideo.configs.models.encoders.clip import (
-    CLIPTextConfig, CLIPVisionConfig, WAN2_1ControlCLIPVisionConfig)
+from fastvideo.configs.models.encoders.clip import (CLIPTextConfig, CLIPVisionConfig, WAN2_1ControlCLIPVisionConfig)
 from fastvideo.configs.models.encoders.llama import LlamaConfig
 from fastvideo.configs.models.encoders.t5 import T5Config, T5LargeConfig
 from fastvideo.configs.models.encoders.qwen2_5 import Qwen2_5_VLConfig
@@ -12,9 +9,7 @@ from fastvideo.configs.models.encoders.reason1 import Reason1ArchConfig, Reason1
 from fastvideo.configs.models.encoders.gemma import LTX2GemmaConfig

 __all__ = [
-    "EncoderConfig", "TextEncoderConfig", "ImageEncoderConfig",
-    "BaseEncoderOutput", "CLIPTextConfig", "CLIPVisionConfig",
-    "WAN2_1ControlCLIPVisionConfig", "LlamaConfig", "T5Config", "T5LargeConfig",
-    "Qwen2_5_VLConfig", "Reason1ArchConfig", "Reason1Config", "LTX2GemmaConfig",
-    "SiglipVisionConfig"
+    "EncoderConfig", "TextEncoderConfig", "ImageEncoderConfig", "BaseEncoderOutput", "CLIPTextConfig",
+    "CLIPVisionConfig", "WAN2_1ControlCLIPVisionConfig", "LlamaConfig", "T5Config", "T5LargeConfig", "Qwen2_5_VLConfig",
+    "Reason1ArchConfig", "Reason1Config", "LTX2GemmaConfig", "SiglipVisionConfig"
 ]
@@ -12,8 +12,8 @@ from fastvideo.platforms import AttentionBackendEnum
@dataclass
 class EncoderArchConfig(ArchConfig):
    architectures: list[str] = field(default_factory=lambda: [])
-    _supported_attention_backends: tuple[AttentionBackendEnum, ...] = (
-        AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TORCH_SDPA)
+    _supported_attention_backends: tuple[AttentionBackendEnum,
+                                         ...] = (AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TORCH_SDPA)
    output_hidden_states: bool = False
    use_return_dict: bool = True

@@ -33,8 +33,7 @@ class TextEncoderArchConfig(EncoderArchConfig):
    scalable_attention: bool = True
    tie_word_embeddings: bool = False
    stacked_params_mapping: list[tuple[str, str, str]] = field(
-        default_factory=list
-    )  # mapping from huggingface weight names to custom names
+        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: [])

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

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


@@ -34,15 +32,13 @@ 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])
+    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
@@ -61,19 +57,17 @@ 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"),
-                                       ])
+    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
 class CLIPTextConfig(TextEncoderConfig):
-    arch_config: TextEncoderArchConfig = field(
-        default_factory=CLIPTextArchConfig)
+    arch_config: TextEncoderArchConfig = field(default_factory=CLIPTextArchConfig)

    num_hidden_layers_override: int | None = None
    require_post_norm: bool | None = None
@@ -84,8 +78,7 @@ class CLIPTextConfig(TextEncoderConfig):

@dataclass
 class CLIPVisionConfig(ImageEncoderConfig):
-    arch_config: ImageEncoderArchConfig = field(
-        default_factory=CLIPVisionArchConfig)
+    arch_config: ImageEncoderArchConfig = field(default_factory=CLIPVisionArchConfig)

    num_hidden_layers_override: int | None = 31
    require_post_norm: bool | None = None
@@ -12,8 +12,7 @@ def _is_feature_extractor_linear(n: str, m) -> bool:


 def _is_embeddings(n: str, m) -> bool:
-    return n.endswith("embeddings_connector") or n.endswith(
-        "audio_embeddings_connector")
+    return n.endswith("embeddings_connector") or n.endswith("audio_embeddings_connector")


 def _is_gemma_model(n: str, m) -> bool:
@@ -22,8 +21,7 @@ def _is_gemma_model(n: str, m) -> bool:

@dataclass
 class LTX2GemmaArchConfig(TextEncoderArchConfig):
-    architectures: list[str] = field(
-        default_factory=lambda: ["LTX2GemmaTextEncoderModel"])
+    architectures: list[str] = field(default_factory=lambda: ["LTX2GemmaTextEncoderModel"])
    hidden_size: int = 3840
    num_hidden_layers: int = 48
    num_attention_heads: int = 30
@@ -42,15 +40,13 @@ class LTX2GemmaArchConfig(TextEncoderArchConfig):
    connector_attention_head_dim: int = 128
    connector_num_layers: int = 2
    connector_positional_embedding_theta: float = 10000.0
-    connector_positional_embedding_max_pos: list[int] = field(
-        default_factory=lambda: [4096])
+    connector_positional_embedding_max_pos: list[int] = field(default_factory=lambda: [4096])
    connector_rope_type: str = "split"
    connector_double_precision_rope: bool = False
    connector_num_learnable_registers: int | None = 128

    _fsdp_shard_conditions: list = field(
-        default_factory=lambda:
-        [_is_feature_extractor_linear, _is_embeddings, _is_gemma_model])
+        default_factory=lambda: [_is_feature_extractor_linear, _is_embeddings, _is_gemma_model])

    def __post_init__(self) -> None:
        super().__post_init__()
@@ -59,7 +55,6 @@ class LTX2GemmaArchConfig(TextEncoderArchConfig):

@dataclass
 class LTX2GemmaConfig(TextEncoderConfig):
-    arch_config: TextEncoderArchConfig = field(
-        default_factory=LTX2GemmaArchConfig)
+    arch_config: TextEncoderArchConfig = field(default_factory=LTX2GemmaArchConfig)

    prefix: str = "ltx2_gemma"
@@ -1,8 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass, field

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


 def _is_transformer_layer(n: str, m) -> bool:
@@ -43,18 +42,16 @@ class LlamaArchConfig(TextEncoderArchConfig):
    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
-        ])
+    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])
+        default_factory=lambda: [_is_transformer_layer, _is_embeddings, _is_final_norm])


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

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


 def _is_transformer_layer(n: str, m) -> bool:
@@ -54,8 +53,7 @@ class Qwen2_5_VLArchConfig(TextEncoderArchConfig):
                                           (".gate_up_proj", ".up_proj", 1),
                                       ])
    _fsdp_shard_conditions: list = field(
-        default_factory=lambda:
-        [_is_transformer_layer, _is_embeddings, _is_final_norm])
+        default_factory=lambda: [_is_transformer_layer, _is_embeddings, _is_final_norm])

    def __post_init__(self):
        super().__post_init__()
@@ -65,8 +63,8 @@ class Qwen2_5_VLArchConfig(TextEncoderArchConfig):
            self.num_key_value_heads = self.num_attention_heads
        if self.layer_types is None:
            self.layer_types = [
-                "sliding_attention" if self.sliding_window is not None
-                and i >= self.max_window_layers else "full_attention"
+                "sliding_attention"
+                if self.sliding_window is not None and i >= self.max_window_layers else "full_attention"
                for i in range(self.num_hidden_layers)
            ]
        if self.rope_scaling is not None and "type" in self.rope_scaling:
@@ -86,7 +84,6 @@ class Qwen2_5_VLArchConfig(TextEncoderArchConfig):

@dataclass
 class Qwen2_5_VLConfig(TextEncoderConfig):
-    arch_config: TextEncoderArchConfig = field(
-        default_factory=Qwen2_5_VLArchConfig)
+    arch_config: TextEncoderArchConfig = field(default_factory=Qwen2_5_VLArchConfig)
    prefix: str = "qwen2_5_vl"
    is_chat_model: bool = True
@@ -11,8 +11,7 @@ from fastvideo.configs.models.encoders.base import TextEncoderArchConfig, TextEn
 class Reason1ArchConfig(TextEncoderArchConfig):
    """Architecture settings (defaults match Qwen2.5-VL-7B-Instruct)."""

-    architectures: list[str] = field(
-        default_factory=lambda: ["Qwen2_5_VLForConditionalGeneration"])
+    architectures: list[str] = field(default_factory=lambda: ["Qwen2_5_VLForConditionalGeneration"])
    model_type: str = "qwen2_5_vl"

    vocab_size: int = 152064
@@ -3,8 +3,7 @@

 from dataclasses import dataclass, field

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


@dataclass
@@ -15,8 +14,7 @@ class SiglipVisionArchConfig(ImageEncoderArchConfig):
    """

    # From config.json
-    architectures: list[str] = field(
-        default_factory=lambda: ["SiglipVisionModel"])
+    architectures: list[str] = field(default_factory=lambda: ["SiglipVisionModel"])
    attention_dropout: float = 0.0
    dtype: str | None = None
    hidden_act: str = "gelu_pytorch_tanh"
@@ -42,8 +40,7 @@ class SiglipVisionArchConfig(ImageEncoderArchConfig):
 class SiglipVisionConfig(ImageEncoderConfig):
    """Configuration for SigLIP vision encoder."""

-    arch_config: ImageEncoderArchConfig = field(
-        default_factory=SiglipVisionArchConfig)
+    arch_config: ImageEncoderArchConfig = field(default_factory=SiglipVisionArchConfig)

    # FastVideo specific
    num_hidden_layers_override: int | None = None
@@ -1,8 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass, field

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


 def _is_transformer_layer(n: str, m) -> bool:
@@ -42,16 +41,14 @@ class T5ArchConfig(TextEncoderArchConfig):
    text_len: int = 512
    dtype: str | None = None
    gradient_checkpointing: bool = False
-    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"),
-                                       ])
+    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])
+        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):
@@ -96,7 +93,6 @@ class T5Config(TextEncoderConfig):
@dataclass
 class T5LargeConfig(TextEncoderConfig):
    """T5 Large configuration for your specific model."""
-    arch_config: TextEncoderArchConfig = field(
-        default_factory=T5LargeArchConfig)
+    arch_config: TextEncoderArchConfig = field(default_factory=T5LargeArchConfig)

    prefix: str = "t5"
@@ -1,6 +1,4 @@
 from fastvideo.configs.models.upsamplers.hunyuan15 import SRTo720pUpsamplerConfig, SRTo1080pUpsamplerConfig
 from fastvideo.configs.models.upsamplers.base import UpsamplerConfig

-__all__ = [
-    "SRTo720pUpsamplerConfig", "SRTo1080pUpsamplerConfig", "UpsamplerConfig"
-]
+__all__ = ["SRTo720pUpsamplerConfig", "SRTo1080pUpsamplerConfig", "UpsamplerConfig"]
@@ -11,8 +11,7 @@ from fastvideo.configs.models.vaes.base import VAEArchConfig, VAEConfig
 class AutoencoderKLArchConfig(VAEArchConfig):
    _name_or_path: str = ""
    act_fn: str = "silu"
-    block_out_channels: tuple[int, ...] | list[int] = field(
-        default_factory=list)
+    block_out_channels: tuple[int, ...] | list[int] = field(default_factory=list)
    down_block_types: tuple[str, ...] | list[str] = field(default_factory=list)
    up_block_types: tuple[str, ...] | list[str] = field(default_factory=list)
    force_upcast: bool = True
@@ -138,16 +138,13 @@ class Cosmos25VAEArchConfig(VAEArchConfig):

        m = re.match(r"^(encoder|decoder)\.middle\.0\.(.*)$", key)
        if m:
-            return map_residual_subkey(f"{m.group(1)}.mid_block.resnets.0",
-                                       m.group(2))
+            return map_residual_subkey(f"{m.group(1)}.mid_block.resnets.0", m.group(2))
        m = re.match(r"^(encoder|decoder)\.middle\.1\.(.*)$", key)
        if m:
-            return map_attn_subkey(f"{m.group(1)}.mid_block.attentions.0",
-                                   m.group(2))
+            return map_attn_subkey(f"{m.group(1)}.mid_block.attentions.0", m.group(2))
        m = re.match(r"^(encoder|decoder)\.middle\.2\.(.*)$", key)
        if m:
-            return map_residual_subkey(f"{m.group(1)}.mid_block.resnets.1",
-                                       m.group(2))
+            return map_residual_subkey(f"{m.group(1)}.mid_block.resnets.1", m.group(2))

        m = re.match(r"^encoder\.downsamples\.(\d+)\.(.*)$", key)
        if m:
@@ -198,10 +195,8 @@ class Cosmos25VAEArchConfig(VAEArchConfig):
    spatial_compression_ratio: int = 8

    def __post_init__(self):
-        self.scaling_factor: torch.Tensor = 1.0 / torch.tensor(
-            self.latents_std).view(1, self.z_dim, 1, 1, 1)
-        self.shift_factor: torch.Tensor = torch.tensor(self.latents_mean).view(
-            1, self.z_dim, 1, 1, 1)
+        self.scaling_factor: torch.Tensor = 1.0 / torch.tensor(self.latents_std).view(1, self.z_dim, 1, 1, 1)
+        self.shift_factor: torch.Tensor = torch.tensor(self.latents_mean).view(1, self.z_dim, 1, 1, 1)
        self.temporal_compression_ratio = self.scale_factor_temporal
        self.spatial_compression_ratio = self.scale_factor_spatial

@@ -210,8 +205,7 @@ class Cosmos25VAEArchConfig(VAEArchConfig):
 class Cosmos25VAEConfig(VAEConfig):
    """Cosmos2.5 VAE config."""

-    arch_config: Cosmos25VAEArchConfig = field(
-        default_factory=Cosmos25VAEArchConfig)
+    arch_config: Cosmos25VAEArchConfig = field(default_factory=Cosmos25VAEArchConfig)

    use_feature_cache: bool = True
    use_tiling: bool = False
@@ -219,5 +213,4 @@ class Cosmos25VAEConfig(VAEConfig):
    use_parallel_tiling: bool = False

    def __post_init__(self):
-        self.blend_num_frames = (self.tile_sample_min_num_frames -
-                                 self.tile_sample_stride_num_frames) * 2
+        self.blend_num_frames = (self.tile_sample_min_num_frames - self.tile_sample_stride_num_frames) * 2
@@ -64,18 +64,15 @@ class CosmosVAEArchConfig(VAEArchConfig):
    spatial_compression_ratio = 8

    def __post_init__(self):
-        self.scaling_factor: torch.Tensor = 1.0 / torch.tensor(
-            self.latents_std).view(1, self.z_dim, 1, 1, 1)
-        self.shift_factor: torch.Tensor = torch.tensor(self.latents_mean).view(
-            1, self.z_dim, 1, 1, 1)
+        self.scaling_factor: torch.Tensor = 1.0 / torch.tensor(self.latents_std).view(1, self.z_dim, 1, 1, 1)
+        self.shift_factor: torch.Tensor = torch.tensor(self.latents_mean).view(1, self.z_dim, 1, 1, 1)
        self.temporal_compression_ratio = self.scale_factor_temporal
        self.spatial_compression_ratio = self.scale_factor_spatial


@dataclass
 class CosmosVAEConfig(VAEConfig):
-    arch_config: CosmosVAEArchConfig = field(
-        default_factory=CosmosVAEArchConfig)
+    arch_config: CosmosVAEArchConfig = field(default_factory=CosmosVAEArchConfig)
    use_feature_cache: bool = True

    use_tiling: bool = False
@@ -83,5 +80,4 @@ class CosmosVAEConfig(VAEConfig):
    use_parallel_tiling: bool = False

    def __post_init__(self):
-        self.blend_num_frames = (self.tile_sample_min_num_frames -
-                                 self.tile_sample_stride_num_frames) * 2
+        self.blend_num_frames = (self.tile_sample_min_num_frames - self.tile_sample_stride_num_frames) * 2
@@ -18,8 +18,7 @@ class Hunyuan15VAEArchConfig(VAEArchConfig):
    scaling_factor: float = 1.03682

    def __post_init__(self):
-        self.spatial_compression_ratio: int = 2**(len(self.block_out_channels) -
-                                                  1)
+        self.spatial_compression_ratio: int = 2**(len(self.block_out_channels) - 1)


@dataclass
@@ -31,8 +31,7 @@ class HunyuanVAEArchConfig(VAEArchConfig):
    mid_block_add_attention: bool = True

    def __post_init__(self):
-        self.spatial_compression_ratio: int = 2**(len(self.block_out_channels) -
-                                                  1)
+        self.spatial_compression_ratio: int = 2**(len(self.block_out_channels) - 1)


@dataclass
@@ -61,10 +61,8 @@ class WanVAEArchConfig(VAEArchConfig):
    clip_output: bool = True

    def __post_init__(self):
-        self.scaling_factor: torch.tensor = 1.0 / torch.tensor(
-            self.latents_std).view(1, self.z_dim, 1, 1, 1)
-        self.shift_factor: torch.tensor = torch.tensor(self.latents_mean).view(
-            1, self.z_dim, 1, 1, 1)
+        self.scaling_factor: torch.tensor = 1.0 / torch.tensor(self.latents_std).view(1, self.z_dim, 1, 1, 1)
+        self.shift_factor: torch.tensor = torch.tensor(self.latents_mean).view(1, self.z_dim, 1, 1, 1)
        self.temporal_compression_ratio = self.scale_factor_temporal
        self.spatial_compression_ratio = self.scale_factor_spatial

@@ -79,5 +77,4 @@ class WanVAEConfig(VAEConfig):
    use_parallel_tiling: bool = False

    def __post_init__(self):
-        self.blend_num_frames = (self.tile_sample_min_num_frames -
-                                 self.tile_sample_stride_num_frames) * 2
+        self.blend_num_frames = (self.tile_sample_min_num_frames - self.tile_sample_stride_num_frames) * 2
@@ -7,15 +7,12 @@ from fastvideo.configs.pipelines.hunyuangamecraft import HunyuanGameCraftPipelin
 from fastvideo.configs.pipelines.hyworld import HYWorldConfig
 from fastvideo.configs.pipelines.ltx2 import LTX2T2VConfig
 from fastvideo.registry import get_pipeline_config_cls_from_name
-from fastvideo.configs.pipelines.wan import (SelfForcingWanT2V480PConfig,
-                                             WanI2V480PConfig, WanI2V720PConfig,
+from fastvideo.configs.pipelines.wan import (SelfForcingWanT2V480PConfig, WanI2V480PConfig, WanI2V720PConfig,
                                             WanT2V480PConfig, WanT2V720PConfig)

 __all__ = [
-    "HunyuanConfig", "FastHunyuanConfig", "HunyuanGameCraftPipelineConfig",
-    "PipelineConfig", "Hunyuan15T2V480PConfig", "Hunyuan15T2V720PConfig",
-    "WanT2V480PConfig", "WanI2V480PConfig", "WanT2V720PConfig",
-    "WanI2V720PConfig", "SelfForcingWanT2V480PConfig", "CosmosConfig",
-    "Cosmos25Config", "LTX2T2VConfig", "HYWorldConfig",
+    "HunyuanConfig", "FastHunyuanConfig", "HunyuanGameCraftPipelineConfig", "PipelineConfig", "Hunyuan15T2V480PConfig",
+    "Hunyuan15T2V720PConfig", "WanT2V480PConfig", "WanI2V480PConfig", "WanT2V720PConfig", "WanI2V720PConfig",
+    "SelfForcingWanT2V480PConfig", "CosmosConfig", "Cosmos25Config", "LTX2T2VConfig", "HYWorldConfig",
    "get_pipeline_config_cls_from_name"
 ]
@@ -6,8 +6,7 @@ from typing import Any, cast

 import torch

-from fastvideo.configs.models import (DiTConfig, EncoderConfig, ModelConfig,
-                                      VAEConfig, UpsamplerConfig)
+from fastvideo.configs.models import (DiTConfig, EncoderConfig, ModelConfig, VAEConfig, UpsamplerConfig)
 from fastvideo.configs.models.encoders import BaseEncoderOutput
 from fastvideo.configs.utils import update_config_from_args
 from fastvideo.logger import init_logger
@@ -56,15 +55,11 @@ class PipelineConfig:

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

    # DMD parameters
    dmd_denoising_steps: list[int] | None = field(default=None)
@@ -77,8 +72,7 @@ class PipelineConfig:
    # enable_torch_compile: bool = False

    @staticmethod
-    def add_cli_args(parser: FlexibleArgumentParser,
-                     prefix: str = "") -> FlexibleArgumentParser:
+    def add_cli_args(parser: FlexibleArgumentParser, prefix: str = "") -> FlexibleArgumentParser:
        prefix_with_dot = f"{prefix}." if (prefix.strip() != "") else ""

        # model_path will be conflicting with the model_path in FastVideoArgs,
@@ -172,8 +166,7 @@ class PipelineConfig:
            f"--{prefix_with_dot}dmd-denoising-steps",
            type=parse_int_list,
            default=PipelineConfig.dmd_denoising_steps,
-            help=
-            "Comma-separated list of denoising steps (e.g., '1000,757,522')",
+            help="Comma-separated list of denoising steps (e.g., '1000,757,522')",
        )

        # Add VAE configuration arguments
@@ -186,19 +179,11 @@ class PipelineConfig:

        return parser

-    def update_config_from_dict(self,
-                                args: dict[str, Any],
-                                prefix: str = "") -> None:
+    def update_config_from_dict(self, 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)
-        update_config_from_args(self.vae_config,
-                                args,
-                                f"{prefix_with_dot}vae_config",
-                                pop_args=True)
-        update_config_from_args(self.dit_config,
-                                args,
-                                f"{prefix_with_dot}dit_config",
-                                pop_args=True)
+        update_config_from_args(self.vae_config, args, f"{prefix_with_dot}vae_config", pop_args=True)
+        update_config_from_args(self.dit_config, args, f"{prefix_with_dot}dit_config", pop_args=True)

    @classmethod
    def from_pretrained(cls, model_path: str) -> "PipelineConfig":
@@ -211,9 +196,7 @@ class PipelineConfig:
        return cast(PipelineConfig, pipeline_config_cls(model_path=model_path))

    @classmethod
-    def from_kwargs(cls,
-                    kwargs: dict[str, Any],
-                    config_cli_prefix: str = "") -> "PipelineConfig":
+    def from_kwargs(cls, kwargs: dict[str, Any], config_cli_prefix: str = "") -> "PipelineConfig":
        """
        Load PipelineConfig from kwargs Dictionary.
        kwargs: dictionary of kwargs
@@ -221,13 +204,10 @@ class PipelineConfig:
        """
        from fastvideo.registry import get_pipeline_config_cls_from_name

-        prefix_with_dot = f"{config_cli_prefix}." if (config_cli_prefix.strip()
-                                                      != "") else ""
-        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')
+        prefix_with_dot = f"{config_cli_prefix}." if (config_cli_prefix.strip() != "") else ""
+        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")

@@ -236,9 +216,7 @@ class PipelineConfig:

        # 2. Instantiate PipelineConfig
        if pipeline_config_cls is None:
-            logger.warning(
-                "Couldn't find pipeline config for %s. Using the default pipeline config.",
-                model_path)
+            logger.warning("Couldn't find pipeline config for %s. Using the default pipeline config.", model_path)
            pipeline_config = cls()
        else:
            pipeline_config = pipeline_config_cls()
@@ -246,8 +224,7 @@ class PipelineConfig:
        # 3. Load PipelineConfig from a json file or a PipelineConfig object if provided
        if isinstance(pipeline_config_or_path, str):
            pipeline_config.load_from_json(pipeline_config_or_path)
-            kwargs[prefix_with_dot +
-                   'pipeline_config_path'] = pipeline_config_or_path
+            kwargs[prefix_with_dot + 'pipeline_config_path'] = pipeline_config_or_path
        elif isinstance(pipeline_config_or_path, PipelineConfig):
            pipeline_config = pipeline_config_or_path
        elif isinstance(pipeline_config_or_path, dict):
@@ -261,9 +238,7 @@ class PipelineConfig:

    def check_pipeline_config(self) -> None:
        if self.vae_sp and not self.vae_tiling:
-            raise ValueError(
-                "Currently enabling vae_sp requires enabling vae_tiling, please set --vae-tiling to True."
-            )
+            raise ValueError("Currently enabling vae_sp requires enabling vae_tiling, please set --vae-tiling to True.")

        if len(self.text_encoder_configs) != len(self.text_encoder_precisions):
            raise ValueError(
@@ -289,8 +264,7 @@ class PipelineConfig:
                # Model Arch Config should be hidden away from the users
                model_dict.pop("arch_config")
                output_dict[key] = model_dict
-            elif isinstance(value, tuple) and all(
-                    isinstance(v, ModelConfig) for v in value):
+            elif isinstance(value, tuple) and all(isinstance(v, ModelConfig) for v in value):
                model_dicts = []
                for v in value:
                    model_dict = asdict(v)
@@ -313,8 +287,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,
-                                                                Any]) -> None:
+    def update_pipeline_config(self, source_pipeline_dict: dict[str, Any]) -> None:
        for f in fields(self):
            key = f.name
            if key in source_pipeline_dict:
@@ -324,14 +297,10 @@ class PipelineConfig:
                # If it's a nested ModelConfig, update it recursively
                if isinstance(current_value, ModelConfig):
                    current_value.update_model_config(new_value)
-                elif isinstance(current_value, tuple) and all(
-                        isinstance(v, ModelConfig) for v in current_value):
+                elif isinstance(current_value, tuple) and all(isinstance(v, ModelConfig) for v in current_value):
                    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,
-                                                            strict=True):
+                        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, strict=True):
                        target_config.update_model_config(source_config)
                else:
                    setattr(self, key, new_value)
@@ -43,16 +43,13 @@ class CosmosConfig(PipelineConfig):

    vae_config: VAEConfig = field(default_factory=CosmosVAEConfig)

-    text_encoder_configs: tuple[EncoderConfig, ...] = field(
-        default_factory=lambda: (T5LargeConfig(), ))
+    text_encoder_configs: tuple[EncoderConfig, ...] = field(default_factory=lambda: (T5LargeConfig(), ))
    postprocess_text_funcs: tuple[Callable[[BaseEncoderOutput], torch.Tensor],
-                                  ...] = field(default_factory=lambda:
-                                               (t5_large_postprocess_text, ))
+                                  ...] = field(default_factory=lambda: (t5_large_postprocess_text, ))

    dit_precision: str = "bf16"
    vae_precision: str = "fp16"
-    text_encoder_precisions: tuple[str, ...] = field(
-        default_factory=lambda: ("bf16", ))
+    text_encoder_precisions: tuple[str, ...] = field(default_factory=lambda: ("bf16", ))

    conditioning_strategy: str = "frame_replace"
    min_num_conditional_frames: int = 1
@@ -26,8 +26,7 @@ def reason1_postprocess_text(outputs: BaseEncoderOutput) -> torch.Tensor:
    normed = []
    for h in hs:
        h = h.float()
-        h = (h - h.mean(dim=-1, keepdim=True)) / (h.std(dim=-1, keepdim=True) +
-                                                  1e-8)
+        h = (h - h.mean(dim=-1, keepdim=True)) / (h.std(dim=-1, keepdim=True) + 1e-8)
        normed.append(h)
    return torch.cat(normed, dim=-1).to(hidden_states[0].dtype)

@@ -36,43 +35,39 @@ def reason1_postprocess_text(outputs: BaseEncoderOutput) -> torch.Tensor:
 class Cosmos25Config(PipelineConfig):
    """Configuration for Cosmos 2.5 (Predict2.5) video generation pipeline."""

-    dit_config: DiTConfig = field(default_factory=lambda: Cosmos25VideoConfig(
-        arch_config=Cosmos25ArchConfig(
-            num_attention_heads=16,
-            attention_head_dim=128,
-            in_channels=16,
-            out_channels=16,
-            num_layers=28,
-            patch_size=[1, 2, 2],
-            max_size=[128, 240, 240],
-            rope_scale=[1.0, 3.0, 3.0],
-            text_embed_dim=1024,
-            mlp_ratio=4.0,
-            adaln_lora_dim=256,
-            use_adaln_lora=True,
-            concat_padding_mask=True,
-            extra_pos_embed_type=None,
-            use_crossattn_projection=True,
-            rope_enable_fps_modulation=False,
-            qk_norm="rms_norm",
-        )))
+    dit_config: DiTConfig = field(default_factory=lambda: Cosmos25VideoConfig(arch_config=Cosmos25ArchConfig(
+        num_attention_heads=16,
+        attention_head_dim=128,
+        in_channels=16,
+        out_channels=16,
+        num_layers=28,
+        patch_size=[1, 2, 2],
+        max_size=[128, 240, 240],
+        rope_scale=[1.0, 3.0, 3.0],
+        text_embed_dim=1024,
+        mlp_ratio=4.0,
+        adaln_lora_dim=256,
+        use_adaln_lora=True,
+        concat_padding_mask=True,
+        extra_pos_embed_type=None,
+        use_crossattn_projection=True,
+        rope_enable_fps_modulation=False,
+        qk_norm="rms_norm",
+    )))

    vae_config: VAEConfig = field(default_factory=Cosmos25VAEConfig)

-    text_encoder_configs: tuple[EncoderConfig, ...] = field(
-        default_factory=lambda: (Reason1Config(arch_config=Reason1ArchConfig(
-            embedding_concat_strategy="full_concat")), ))
+    text_encoder_configs: tuple[EncoderConfig, ...] = field(default_factory=lambda: (Reason1Config(
+        arch_config=Reason1ArchConfig(embedding_concat_strategy="full_concat")), ))

-    preprocess_text_funcs: tuple[Callable[[str], str], ...] = field(
-        default_factory=lambda: (_identity_preprocess_text, ))
+    preprocess_text_funcs: tuple[Callable[[str], str],
+                                 ...] = field(default_factory=lambda: (_identity_preprocess_text, ))
    postprocess_text_funcs: tuple[Callable[[BaseEncoderOutput], torch.Tensor],
-                                  ...] = field(default_factory=lambda:
-                                               (reason1_postprocess_text, ))
+                                  ...] = field(default_factory=lambda: (reason1_postprocess_text, ))

    dit_precision: str = "bf16"
    vae_precision: str = "bf16"
-    text_encoder_precisions: tuple[str, ...] = field(
-        default_factory=lambda: ("bf16", ))
+    text_encoder_precisions: tuple[str, ...] = field(default_factory=lambda: ("bf16", ))

    embedded_cfg_scale: float = 0.0
    flow_shift: float = 5.0
@@ -7,8 +7,7 @@ import torch

 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.encoders import (BaseEncoderOutput, CLIPTextConfig, LlamaConfig)
 from fastvideo.configs.models.vaes import HunyuanVAEConfig
 from fastvideo.configs.pipelines.base import PipelineConfig

@@ -41,8 +40,7 @@ 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
-    last_hidden_state: torch.tensor = hidden_states[-(hidden_state_skip_layer +
-                                                      1)]
+    last_hidden_state: torch.tensor = hidden_states[-(hidden_state_skip_layer + 1)]
    crop_start = prompt_template_video.get("crop_start", -1)
    last_hidden_state = last_hidden_state[:, crop_start:]
    return last_hidden_state
@@ -71,20 +69,16 @@ class HunyuanConfig(PipelineConfig):
    flow_shift: int = 7

    # Text encoding stage
-    text_encoder_configs: tuple[EncoderConfig, ...] = field(
-        default_factory=lambda: (LlamaConfig(), CLIPTextConfig()))
-    preprocess_text_funcs: tuple[Callable[[str], str], ...] = field(
-        default_factory=lambda: (llama_preprocess_text, clip_preprocess_text))
-    postprocess_text_funcs: tuple[
-        Callable[[BaseEncoderOutput], torch.tensor],
-        ...] = field(default_factory=lambda:
-                     (llama_postprocess_text, clip_postprocess_text))
+    text_encoder_configs: tuple[EncoderConfig, ...] = field(default_factory=lambda: (LlamaConfig(), CLIPTextConfig()))
+    preprocess_text_funcs: tuple[Callable[[str], str],
+                                 ...] = field(default_factory=lambda: (llama_preprocess_text, clip_preprocess_text))
+    postprocess_text_funcs: tuple[Callable[[BaseEncoderOutput], torch.tensor],
+                                  ...] = field(default_factory=lambda: (llama_postprocess_text, clip_postprocess_text))

    # Precision for each component
    dit_precision: str = "bf16"
    vae_precision: str = "fp16"
-    text_encoder_precisions: tuple[str, ...] = field(
-        default_factory=lambda: ("fp16", "fp16"))
+    text_encoder_precisions: tuple[str, ...] = field(default_factory=lambda: ("fp16", "fp16"))

    def __post_init__(self):
        self.vae_config.load_encoder = False
--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
Jinzhe Pan	c2a12fd087	Revert "[feat] Job Runner UI (#1172 )" This reverts commit `f72618dafb`.	2026-03-29 16:42:04 +08:00
Darren	f72618dafb	[feat] Job Runner UI (#1172 )	2026-03-29 16:17:56 +08:00
alexzms	b3edfacdd8	[bugfix]: fix I2V preprocessing crash for models without CLIP (Wan2.2 I2V) (#1184 )	2026-03-28 10:28:50 +08:00
alexzms	30129a3350	[misc]: reorganize training configs and add documentation (#1177 )	2026-03-26 16:38:31 -07:00
jaisurya27	4d49f7b0aa	Kandinsky5 lite dit clean (#1088 )	2026-03-26 08:07:11 +08:00
alexzms	71bfc13d75	[feat]: add HunyuanVideo model plugin for fastvideo/train framework (#1175 )	2026-03-24 16:49:51 -07:00
Kaiqin Kong	74db6e18d1	[misc] update action loading in validation and preprocess (#1143 )	2026-03-24 15:10:03 -07:00
Kaiqin Kong	7d263c6a36	[bugfix] self-forcing train/validation step mismatch (#1173 )	2026-03-20 00:52:45 -07:00
Zhang Peiyuan	454c32d1d1	Update README.md	2026-03-17 14:26:07 -07:00
Jinzhe Pan	d1240b9238	[CI] add contributor interaction automation (#1170 )	2026-03-17 12:05:40 +08:00
Hao Zhang	4105094fa5	[docs] Update README with realtime demo announcement (#1169 )	2026-03-13 15:52:02 -07:00
alexzms	f036469d3d	[feat]: Knowledge Distillation training method for ODE-init (KDMethod + KDCausalMethod) (#1166 )	2026-03-11 20:43:24 -07:00
alexzms	14261bc98c	[feat] pre-commit support 120 col num (#1167 )	2026-03-11 20:19:30 -07:00