[misc]: drop .agents/exploration journals added by #1453

AGENTS.md bans experiment journals and branch-state snapshots under
.agents/.

.agents/skills/reseed-performance-baseline/SKILL.md

@@ -69,7 +69,7 @@ approval, then upload reviewed accepted baseline records.
 | `model_id` | Yes | Benchmark id, e.g. `wan-t2v-1.3b-2gpu`. This maps to the HF subdirectory after `sanitize(model_id)`. |
 | `gpu_type` | Yes | Exact GPU device string from the performance record, e.g. the L40S device name emitted by CI. Baselines are GPU-specific. |
 | `source_results` | Yes | One or more local paths or Buildkite artifact URLs for accepted shifted performance JSONs. Prefer normalized `normalized_perf_*.json` artifacts emitted by `compare_baseline.py`. Accept `source_result` as an alias only for a single JSON. |
-| `max_intra_batch_regression` | No | Maximum allowed regression of any source JSON against the source batch median. Default: `0.05` (5%). |
+| `max_intra_batch_regression` | No | Maximum allowed regression of any source JSON against the source batch median. Default: `PERF_MAX_REGRESSION` if set, otherwise `0.05` (5%). |
 | `intent_rationale` | Yes | One-line explanation for why the baseline shift is legitimate. This is written into provenance and should be reused in the PR. |
 
 Hardcoded defaults:
@@ -148,7 +148,8 @@ For each metric with at least two non-null source values:
 4. Stop if any source record regresses against the batch median by more than
    `max_intra_batch_regression`.
 
-Default `max_intra_batch_regression` to `0.05`. Print a table with per-source values, batch median, and
+Default `max_intra_batch_regression` to `PERF_MAX_REGRESSION` when set,
+otherwise `0.05`. Print a table with per-source values, batch median, and
 worst intra-batch regression.
 
 This check prevents uploading a mixed batch where one JSON is materially
@@ -182,7 +183,7 @@ present, that run is not a valid source for baseline reseeding.
 
 ### 2. Sync and back up existing HF records under /tmp
 
-Use `fastvideo/performance/hf_store.py` helpers directly. Do **not** use
+Use `fastvideo/tests/performance/hf_store.py` helpers directly. Do **not** use
 `compare_baseline.py` as a sync shortcut; on full main runs it can persist
 records, while this step must only fetch and back up existing history.
 
@@ -191,7 +192,7 @@ The sync command pattern is:
 ```bash
 export PERFORMANCE_TRACKING_ROOT="${PERFORMANCE_TRACKING_ROOT:-/tmp/perf-tracking}"
 export HF_REPO_ID="${HF_REPO_ID:-FastVideo/performance-tracking}"
-python -c 'from fastvideo.performance.hf_store import sync_from_hf; import os; sync_from_hf(os.environ["PERFORMANCE_TRACKING_ROOT"], strict=True)'
+PYTHONPATH=fastvideo/tests/performance python -c 'from hf_store import sync_from_hf; import os; sync_from_hf(os.environ["PERFORMANCE_TRACKING_ROOT"], strict=True)'
 ```
 
 Then back up only the sanitized model directory under `/tmp`:
@@ -199,8 +200,8 @@ Then back up only the sanitized model directory under `/tmp`:
 ```bash
 SHORT_COMMIT=$(git rev-parse --short=12 HEAD)
 TIMESTAMP=$(date -u +%Y%m%d_%H%M%S)
-MODEL_SAFE=$(python - <<'PY'
-from fastvideo.performance.hf_store import sanitize
+MODEL_SAFE=$(PYTHONPATH=fastvideo/tests/performance python - <<'PY'
+from hf_store import sanitize
 print(sanitize("<model_id>"))
 PY
 )
@@ -234,7 +235,7 @@ first baseline seed. Continue, but report that baseline history was empty.
 Load the last 5 successful records for the target:
 
 ```python
-from fastvideo.performance.hf_store import load_records_for_model
+from hf_store import load_records_for_model
 
 records = load_records_for_model(
     "/tmp/perf-tracking",
@@ -371,7 +372,7 @@ prepared records plus backup on disk.
 Use the shared storage helper so the path and repo type match CI:
 
 ```python
-from fastvideo.performance.hf_store import upload_record
+from hf_store import upload_record
 
 upload_record("<local_record_path>", record, strict=True)
 ```
@@ -459,7 +460,7 @@ directories created for this reseed. Never remove unrelated `/tmp` contents.
   intentional baseline replacement.
 - `fastvideo/tests/performance/compare_baseline.py` — normalization, rolling
   median comparison, and persistence rules.
-- `fastvideo/performance/hf_store.py` — HF sync, record loading,
+- `fastvideo/tests/performance/hf_store.py` — HF sync, record loading,
   `sanitize()`, and `upload_record()`.
 - `fastvideo/tests/performance/test_inference_performance.py` — source result
   JSON schema.

.buildkite/performance-benchmarks/tests/wan-t2v-1.3b.json

@@ -1,9 +1,9 @@
 {
   "benchmark_id": "wan-t2v-1.3b-2gpu",
   "config_schema_version": 2,
-  "workload_id": "wan-t2v",
-  "variant_id": "1.3b-sp2",
-  "benchmark_version": 2,
+  "workload_id": "wan-t2v-1.3b",
+  "variant_id": "canonical",
+  "benchmark_version": 1,
   "description": "Wan2.1 T2V 1.3B inference performance",
   "model": {
     "model_path": "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",

.github/workflows/ci-precommit.yml

@@ -1,11 +1,7 @@
 name: pre-commit
 
 on:
-  # pull_request_target instead of pull_request: the workflow definition and
-  # the hook config are always taken from the BASE branch, so fork /
-  # first-time-contributor PRs run immediately without a maintainer clicking
-  # "Approve and run". The PR head is checked out as data only.
-  pull_request_target:
+  pull_request:
     branches: [main]
   workflow_call:
     inputs:
@@ -19,25 +15,12 @@ permissions:
 
 jobs:
   pre-commit:
-    if: github.event.pull_request.draft != true
+    if: github.event_name == 'workflow_call' || github.event.pull_request.draft != true
     runs-on: ubuntu-latest
     steps:
     - uses: actions/checkout@v4
       with:
         ref: ${{ inputs.ref || '' }}
-    # For PR events, lint the PR head — but keep the hook definitions from
-    # the base branch so an untrusted PR cannot alter what gets executed.
-    - name: Save trusted hook config
-      if: github.event_name == 'pull_request_target'
-      run: cp .pre-commit-config.yaml "$RUNNER_TEMP/trusted-pre-commit-config.yaml"
-    - uses: actions/checkout@v4
-      if: github.event_name == 'pull_request_target'
-      with:
-        ref: ${{ github.event.pull_request.head.sha }}
-        persist-credentials: false
-    - name: Restore trusted hook config
-      if: github.event_name == 'pull_request_target'
-      run: cp "$RUNNER_TEMP/trusted-pre-commit-config.yaml" .pre-commit-config.yaml
     - uses: actions/setup-python@v5
       with:
         python-version: "3.12"

apps/performance_dashboard/frontend/src/App.tsx

@@ -1,7 +1,6 @@
 import { useEffect, useMemo, useState } from "react";
 
-import { fetchSummary, fetchTrends, refreshData } from "./api";
-import type { CohortValue, RunSource, SummaryResponse, TrendGroup, TrendPoint } from "./api";
+import { fetchSummary, fetchTrends, refreshData, RunSource, SummaryResponse, TrendGroup, TrendPoint } from "./api";
 
 const METRIC_KEYS = ["latency", "throughput", "memory", "text_encoder_time_s", "dit_time_s", "vae_decode_time_s"];
 const RUN_SOURCES: Array<{ value: "" | RunSource; label: string }> = [
@@ -110,61 +109,6 @@ function metricLabel(metricKey: string) {
   return METRIC_DEFINITIONS[metricKey]?.label ?? metricKey;
 }
 
-type CohortFields = {
-  model_id: string;
-  gpu_type: string;
-  workload_id: CohortValue;
-  variant_id: CohortValue;
-  benchmark_version: CohortValue;
-  recipe_fingerprint: CohortValue;
-  hardware_profile_id: CohortValue;
-  software_profile_id: CohortValue;
-};
-
-function cohortValue(value: CohortValue) {
-  if (value === null || value === undefined || value === "") {
-    return "legacy";
-  }
-  return String(value);
-}
-
-function shortCohortValue(value: CohortValue) {
-  const text = cohortValue(value);
-  if (text === "legacy" || text.length <= 14) {
-    return text;
-  }
-  return text.slice(0, 12);
-}
-
-function cohortKey(cohort: CohortFields) {
-  return [
-    cohort.model_id,
-    cohort.gpu_type,
-    cohortValue(cohort.workload_id),
-    cohortValue(cohort.variant_id),
-    cohortValue(cohort.benchmark_version),
-    cohortValue(cohort.recipe_fingerprint),
-    cohortValue(cohort.hardware_profile_id),
-    cohortValue(cohort.software_profile_id)
-  ].join("|");
-}
-
-function cohortTitle(cohort: CohortFields) {
-  const workload = cohortValue(cohort.workload_id);
-  const variant = cohortValue(cohort.variant_id);
-  const version = cohortValue(cohort.benchmark_version);
-  const versionLabel = version === "legacy" ? version : `v${version}`;
-  return `${workload} / ${variant} / ${versionLabel}`;
-}
-
-function cohortDetail(cohort: CohortFields) {
-  return [
-    `recipe ${shortCohortValue(cohort.recipe_fingerprint)}`,
-    shortCohortValue(cohort.hardware_profile_id),
-    shortCohortValue(cohort.software_profile_id)
-  ].join(" | ");
-}
-
 function formatMetricValue(metricKey: string, value: number | null | undefined, tooltip = false) {
   const definition = METRIC_DEFINITIONS[metricKey];
   if (!definition) {
@@ -227,9 +171,7 @@ function TrendChart({ group, metricKey }: { group: TrendGroup; metricKey: string
         top: `${(activePoint.y / height) * 100}%`
       }
     : undefined;
-  const ariaLabel = `${metricLabel(metricKey)} trend for ${group.model_id} on ${group.gpu_type}, ${cohortTitle(
-    group
-  )}`;
+  const ariaLabel = `${metricLabel(metricKey)} trend for ${group.model_id} on ${group.gpu_type}`;
 
   return (
     <div className="chart-shell">
@@ -477,7 +419,7 @@ export default function App() {
       <section className="panel">
         <div className="panel-header">
           <h2>Latest Status</h2>
-          <span>{latestRows.length} comparison cohorts</span>
+          <span>{latestRows.length} model/GPU groups</span>
         </div>
         {latestRows.length === 0 ? (
           <div className="empty">No records match the selected filters.</div>
@@ -490,7 +432,6 @@ export default function App() {
                   <th>Recomputed</th>
                   <th>Model</th>
                   <th>GPU</th>
-                  <th>Cohort</th>
                   <th>Commit</th>
                   <th>Source</th>
                   <th>Baseline</th>
@@ -505,7 +446,7 @@ export default function App() {
               </thead>
               <tbody>
                 {latestRows.map((row) => (
-                  <tr key={cohortKey(row)}>
+                  <tr key={`${row.model_id}-${row.gpu_type}`}>
                     <td>
                       <span className={`badge ${row.status}`}>{row.status}</span>
                     </td>
@@ -516,12 +457,6 @@ export default function App() {
                     </td>
                     <td>{row.model_id}</td>
                     <td>{row.gpu_type}</td>
-                    <td>
-                      <div className="cohort-cell">
-                        <strong>{cohortTitle(row)}</strong>
-                        <span>{cohortDetail(row)}</span>
-                      </div>
-                    </td>
                     <td>{shortSha(row.commit_sha)}</td>
                     <td>
                       <span className={`source-badge source-${row.run_source}`}>{runSourceLabel(row.run_source)}</span>
@@ -560,13 +495,11 @@ export default function App() {
           ) : (
             trends.map((group) =>
               METRIC_KEYS.map((metricKey) => (
-                <article className="trend-card" key={`${cohortKey(group)}-${metricKey}`}>
+                <article className="trend-card" key={`${group.model_id}-${group.gpu_type}-${metricKey}`}>
                   <div>
                     <h3>{metricLabel(metricKey)}</h3>
                     <p>
                       {group.model_id} | {group.gpu_type}
-                      <span>{cohortTitle(group)}</span>
-                      <span>{cohortDetail(group)}</span>
                     </p>
                   </div>
                   <TrendChart group={group} metricKey={metricKey} />

apps/performance_dashboard/frontend/src/api.ts

@@ -13,17 +13,6 @@ export type MetricValue = {
   precision: number;
 };
 
-export type CohortValue = string | number | null;
-
-export type ComparisonCohort = {
-  workload_id: CohortValue;
-  variant_id: CohortValue;
-  benchmark_version: CohortValue;
-  recipe_fingerprint: CohortValue;
-  hardware_profile_id: CohortValue;
-  software_profile_id: CohortValue;
-};
-
 export type SummaryRow = {
   model_id: string;
   gpu_type: string;
@@ -45,7 +34,7 @@ export type SummaryRow = {
   build_id: string;
   job_id: string;
   metrics: Record<string, MetricValue>;
-} & ComparisonCohort;
+};
 
 export type RunSource = "pr" | "local" | "scheduled_main" | "unknown";
 
@@ -79,13 +68,13 @@ export type TrendPoint = {
   build_id: string;
   job_id: string;
   metrics: Record<string, number | null>;
-} & ComparisonCohort;
+};
 
 export type TrendGroup = {
   model_id: string;
   gpu_type: string;
   points: TrendPoint[];
-} & ComparisonCohort;
+};
 
 export type TrendsResponse = {
   groups: TrendGroup[];

apps/performance_dashboard/frontend/src/styles.css

@@ -149,11 +149,6 @@ h3 {
   font-size: 0.82rem;
 }
 
-.trend-card p {
-  display: grid;
-  gap: 2px;
-}
-
 .stat strong {
   display: block;
   margin-top: 8px;
@@ -191,7 +186,7 @@ h3 {
 
 table {
   width: 100%;
-  min-width: 1260px;
+  min-width: 1120px;
   border-collapse: collapse;
 }
 
@@ -214,25 +209,6 @@ td {
   font-size: 0.9rem;
 }
 
-.cohort-cell {
-  display: grid;
-  gap: 2px;
-}
-
-.cohort-cell strong,
-.trend-card p span {
-  color: #1b2836;
-  font-size: 0.78rem;
-  font-weight: 700;
-}
-
-.cohort-cell span,
-.trend-card p span + span {
-  color: #607080;
-  font-family: ui-monospace, SFMono-Regular, Menlo, Monaco, Consolas, "Liberation Mono", monospace;
-  font-size: 0.72rem;
-}
-
 .badge {
   display: inline-flex;
   align-items: center;

docker/Dockerfile

@@ -68,7 +68,7 @@ ARG FLASH_ATTN_WHEEL_RELEASE_ARM64=https://github.com/mjun0812/flash-attention-p
 # cutlass-4.4 `cute.core.ThrMma` API, which crashes on the cutlass-dsl 4.5 that
 # flashinfer/quack pull in. After the wheel install we overlay this cutlass-4.5-safe
 # upstream cute (flash-attn-4) so the image runs FA4 instead of the FA2 fallback.
-ARG FA4_CUTE_REF=82d6441eec5d4dfec120153db2c0145ae855a083
+ARG FA4_CUTE_REF=940cd9680f3315f2f06b43ab5bea2c2cf2d96806
 
 # Provided automatically by BuildKit/buildx (e.g. "amd64" / "arm64") and used to
 # select the prebuilt flash-attn wheel. Empty under a plain `docker build` without
@@ -161,7 +161,7 @@ RUN --mount=type=cache,target=/opt/uv/cache \
         uv pip install flash-attn==${FLASH_ATTN_VERSION} --no-build-isolation; \
     fi
 
-# Overlay the CuTe-DSL-4.6-compatible upstream FA4 cute (FA4_CUTE_REF) over the
+# Overlay the cutlass-4.5-safe upstream FA4 cute (FA4_CUTE_REF) over the
 # wheel/source one so the image runs FA4, not the FA2 fallback. This pulls the FA4
 # runtime stack (cutlass-dsl, quack-kernels, apache-tvm-ffi, torch-c-dlpack-ext) --
 # the same deps the [dreamverse] extra already installs in CI; the installed torch

docs/contributing/performance_benchmarks.md

@@ -112,11 +112,10 @@ per-metric policy with direction, percent threshold, absolute threshold, and a
 while it runs so pipeline stage execution times are available in
 `generate_video(...).logging_info`. Stage logs use pipeline-unique keys such as
 `prompt_encoding_stage` so duplicate stage classes do not collide. For
-`PipelineStage` entries, shared component stage bases emit a stable
-`component_metric`: text encoding stages map to `text_encoder_time_s`,
-denoising stages and subclasses map to `dit_time_s`, and decoding stages map to
-`vae_decode_time_s`. The extractor falls back to known `stage_class` names for
-older logs that do not include `component_metric` or that used the class name as
+`PipelineStage` entries, the extractor maps the `stage_class` field:
+`TextEncodingStage` maps to `text_encoder_time_s`, `DenoisingStage` and
+`DmdDenoisingStage` map to `dit_time_s`, and `DecodingStage` maps to
+`vae_decode_time_s`, with a fallback for older logs that used the class name as
 the stage key. Generator-side timings such as `PostDecodeFrameProcessStage`,
 `VideoSaveStage`, and `AudioMuxStage` are intentionally ignored. If a pipeline
 does not report one of the mapped stages, that component metric is stored as
@@ -201,9 +200,9 @@ configs and remain loadable. New or migrated configs should use
 {
   "benchmark_id": "wan-t2v-1.3b-2gpu",
   "config_schema_version": 2,
-  "workload_id": "wan-t2v",
-  "variant_id": "1.3b-sp2",
-  "benchmark_version": 2
+  "workload_id": "wan-t2v-1.3b",
+  "variant_id": "canonical",
+  "benchmark_version": 1
 }
 ```
 
@@ -214,25 +213,21 @@ metadata that make the measured workload explicit:
 
 | Field | Purpose |
 |---|---|
-| `workload_id` | Stable benchmark family, such as `wan-t2v`. |
-| `variant_id` | Intentional recipe family, including model size and parallelism config, such as `1.3b-sp2`. |
+| `workload_id` | Stable benchmark family, such as `wan-t2v-1.3b`. |
+| `variant_id` | Intentional recipe family, such as `canonical`. |
 | `benchmark_version` | Version of the measurement protocol and comparison policy. |
 
 If a config declares `config_schema_version: 2`, loading fails clearly when any
 required v2 identity field is missing. If v2 identity or metadata fields are
 added without `config_schema_version: 2`, loading also fails so partial
 migrations do not silently run as v1 configs. Optional v2 metadata fields
-reserved for follow-up work, such as `metric_threshold_policy` and
-`quality_metadata`, must be JSON objects when present. (`recipe` is emitted
-by the harness and is not config-declarable.)
+reserved for follow-up work, such as `recipe`, `metric_threshold_policy`, and
+`quality_metadata`, must be JSON objects when present.
 
 Recipe fingerprinting, hardware/software profile IDs, exact-identity
-comparison, and dashboard cohort grouping land with this change: v2 records
-compare only within their identity cohort, and a record that opens a NEW
-cohort is marked `baseline_status: "initialized_new_cohort"` (regression
-gating starts once that cohort accumulates history). Legacy v1 configs keep
-the `(model_id, gpu_type)` rolling-baseline comparison. Metric-specific
-threshold policies and promoted baselines remain separate follow-ups.
+comparison, metric-specific threshold policy behavior, promoted baselines, and
+dashboard regrouping are separate follow-up changes. Until those land, rolling
+baseline comparison remains keyed by `(model_id, gpu_type)`.
 
 ### Raw record (`results/perf_*.json`)
 
@@ -242,9 +237,9 @@ Written by `test_inference_performance.py`. One file per benchmark run.
 {
   "benchmark_id": "wan-t2v-1.3b-2gpu",
   "config_schema_version": 2,
-  "workload_id": "wan-t2v",
-  "variant_id": "1.3b-sp2",
-  "benchmark_version": 2,
+  "workload_id": "wan-t2v-1.3b",
+  "variant_id": "canonical",
+  "benchmark_version": 1,
   "model_short_name": "Wan2.1-T2V-1.3B-Diffusers",
   "device": "NVIDIA L40S",
   "num_gpus": 2,
@@ -274,43 +269,7 @@ Written by `test_inference_performance.py`. One file per benchmark run.
   "timestamp": "2026-05-08T22:00:00+00:00",
   "text_encoder_time_s": 2.141,
   "dit_time_s": 8.437,
-  "vae_decode_time_s": 3.208,
-  "recipe": {
-    "recipe_schema_version": 1,
-    "benchmark": {
-      "benchmark_id": "wan-t2v-1.3b-2gpu",
-      "workload_id": "wan-t2v",
-      "variant_id": "1.3b-sp2",
-      "benchmark_version": 2
-    },
-    "model": { "model_path": "Wan-AI/Wan2.1-T2V-1.3B-Diffusers" },
-    "init_kwargs": { "num_gpus": 2, "sp_size": 2, "tp_size": 1 },
-    "generation_kwargs": { "height": 480, "width": 832, "num_frames": 45 },
-    "inputs": { "prompt_count": 1, "prompt_sha256": ["<measured-prompt-sha256>"] },
-    "attention": { "requested_backend": "FLASH_ATTN", "resolved_backend": "FLASH_ATTN" }
-  },
-  "recipe_fingerprint": "<sha256>",
-  "hardware_profile": {
-    "device_type": "cuda",
-    "gpu_count": 2,
-    "gpus": [{ "name": "NVIDIA L40S", "memory_gb": 48, "compute_capability": "8.9" }],
-    "interconnect": "none_or_partial"
-  },
-  "hardware_profile_id": "hw-<sha256-prefix>",
-  "software_profile": {
-    "python": "3.12",
-    "pytorch": "2.12",
-    "cuda": "13.0",
-    "packages": {
-      "fastvideo_kernel": "0.3.2",
-      "flashinfer": "0.2.11",
-      "nvidia_cutlass_dsl": "4.5.0",
-      "triton": "3.4.1"
-    }
-  },
-  "software_profile_id": "sw-<sha256-prefix>",
-  "environment_metadata": { "env": { "IMAGE_VERSION": "py3.12-cuda13.0.0" } },
-  "environment_fingerprint": "env-<sha256-prefix>"
+  "vae_decode_time_s": 3.208
 }
 ```
 
@@ -322,9 +281,6 @@ result, used as the rolling-baseline source of truth.
 ```jsonc
 {
   "model_id": "wan-t2v-1.3b-2gpu",
-  "workload_id": "wan-t2v",
-  "variant_id": "1.3b-sp2",
-  "benchmark_version": 2,
   "timestamp": "2026-05-08T22:00:00+00:00",
   "commit_sha": "<full sha>",
   "gpu_type": "NVIDIA L40S",
@@ -341,10 +297,6 @@ result, used as the rolling-baseline source of truth.
       "gated": true
     }
   },
-  "recipe_fingerprint": "<sha256>",
-  "hardware_profile_id": "hw-<sha256-prefix>",
-  "software_profile_id": "sw-<sha256-prefix>",
-  "environment_fingerprint": "env-<sha256-prefix>",
   "success": true
 }
 ```
@@ -357,18 +309,6 @@ dashboard lists skipped plots for metric series that have no non-null values.
 Records missing both `run_source` and `baseline_eligible` are treated as legacy
 successful main/full-suite uploads and remain eligible for rolling baselines.
 
-New records compare only against the same `model_id`, `gpu_type`,
-`workload_id`, `variant_id`, `benchmark_version`, `recipe_fingerprint`,
-`hardware_profile_id`, and `software_profile_id` cohort.
-`environment_metadata` and `environment_fingerprint` are audit data and are not
-part of the comparison key.
-The recipe prompt digests describe the prompts actually measured by the
-benchmark run; extra configured prompts are ignored unless the benchmark runner
-executes them.
-Software profile package cohorts keep exact versions for relevant
-attention/kernel packages, including FastVideo kernels, FlashAttention,
-FlashInfer, Cutlass DSL, SageAttention, Triton, and xFormers when installed.
-
 ## Environment variable reference
 
 | Variable | Default | Used by | Purpose |
@@ -423,7 +363,7 @@ When the rolling-baseline phase runs, it emits:
      "benchmark_id": "<unique-id>",
      "config_schema_version": 2,
      "workload_id": "<stable-workload-id>",
-     "variant_id": "<variant, e.g. 1.3b-sp2>",
+     "variant_id": "canonical",
      "benchmark_version": 1,
      "model": { "model_path": "...", "model_short_name": "..." },
      "init_kwargs": { "num_gpus": 1, ... },
@@ -450,9 +390,7 @@ When the rolling-baseline phase runs, it emits:
 
    Legacy v1 configs without `config_schema_version` still load, but should not
    gain v2 identity or metadata fields until they are migrated to
-   `config_schema_version: 2`. For v2 configs, `workload_id`, `variant_id`,
-   and `benchmark_version` are part of the comparison key; benchmark runs
-   fail if any of these identity fields are missing.
+   `config_schema_version: 2`.
 
 2. The pytest test auto-discovers all configs — no test code needed. CI
    picks it up on the next `/test performance` run.
@@ -494,5 +432,5 @@ pipelines that did not report a mapped component stage.
 
 **Component timing is `null`** — the generated result did not include a mapped
 stage in `logging_info.stages`. Check that the pipeline emits stage logging
-and that the stage emits `component_metric` or is covered by the legacy
-`STAGE_METRIC_MAP` fallback in `test_inference_performance.py`.
+and that the stage name is listed in `STAGE_METRIC_MAP` in
+`test_inference_performance.py`.

docs/design/inference_schema_parity_inventory.yaml

@@ -41,6 +41,11 @@ surfaces:
       disable_autocast: generator.engine.disable_autocast
       enable_stage_verification: generator.engine.enable_stage_verification
       prompt_txt: request.inputs.prompt_path
+      batching_mode: generator.engine.batching.mode
+      batching_max_size: generator.engine.batching.max_size
+      batching_delay_ms: generator.engine.batching.delay_ms
+      batching_config: generator.engine.batching.config_path
+      enable_batching_metrics: generator.engine.batching.enable_metrics
       override_text_encoder_safetensors: generator.pipeline.components.text_encoder_weights
       override_text_encoder_quant: generator.engine.quantization.text_encoder_quant
       transformer_quant: generator.engine.quantization.transformer_quant
@@ -191,9 +196,6 @@ surfaces:
         sources:
           - fastvideo.configs.pipelines.gen3c.Gen3CConfig
           - fastvideo.configs.pipelines.gen3c.Gen3CInferenceConfig
-      color_correction_strength:
-        sources:
-          - fastvideo.configs.pipelines.dreamx_world.DreamXWorld5BARPipelineConfig
       default_camera_rotation:
         sources:
           - fastvideo.configs.pipelines.gen3c.Gen3CConfig

docs/inference/support_matrix.md

@@ -58,8 +58,6 @@ pipeline initialization and sampling.
 | FastWan2.1 T2V 1.3B | `FastVideo/FastWan2.1-T2V-1.3B-Diffusers` | 480P | ⭕ | ⭕ | ⭕ | ✅ | ⭕ |
 | FastWan2.2 TI2V 5B Full Attn* | `FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers` | 720P | ⭕ | ⭕ | ⭕ | ✅ | ⭕ |
 | Wan2.2 TI2V 5B | `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | 720P | ⭕ | ⭕ | ✅ | ⭕ | ⭕ |
-| DreamX-World 5B Cam | `FastVideo/DreamX-World-5B-Cam-Diffusers` | 480P | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ |
-| DreamX-World 5B AR | `FastVideo/DreamX-World-5B-Diffusers` | 704px1280p | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ |
 | Lucy Edit Dev 5B*** | `decart-ai/Lucy-Edit-Dev` | 480P | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ |
 | Wan2.2 T2V A14B | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | 480P<br>720P | ❌ | ❌ | ✅ | ⭕ | ⭕ |
 | Wan2.2 I2V A14B | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | 480P<br>720P | ❌ | ❌ | ✅ | ⭕ | ⭕ |

docs/training/trackers.md

@@ -1,91 +0,0 @@
-# Training Trackers
-
-FastVideo can send training metrics and validation media to Weights & Biases
-or SwanLab. Tracking runs only on global rank 0, and local tracker files are
-stored under `<output_dir>/tracker`.
-
-## Supported Trackers
-
-| Value | Backend | Installation |
-|-------|---------|--------------|
-| `wandb` | Weights & Biases | Included with FastVideo |
-| `swanlab` | SwanLab | Install the optional `swanlab` dependency |
-| `none` | Disable external tracking | No additional package |
-
-You can enable more than one backend, for example `trackers: [wandb, swanlab]`.
-Metrics and validation media are converted to the artifact type required by
-each backend.
-
-## Install SwanLab
-
-For a published FastVideo installation, install the SwanLab extra:
-
-```bash
-uv pip install "fastvideo[swanlab]"
-```
-
-For an editable source checkout, include the same extra during installation:
-
-```bash
-uv pip install -e ".[swanlab]"
-```
-
-If FastVideo is already installed, you can install the compatible SDK directly:
-
-```bash
-uv pip install "swanlab>=0.6.7"
-```
-
-Authenticate once before starting a training run:
-
-```bash
-swanlab login
-```
-
-See the [SwanLab login documentation](https://docs.swanlab.cn/en/api/cli-swanlab-login.html)
-for non-interactive and self-hosted setups.
-
-## Configure Tracking
-
-Select SwanLab in the YAML config used by the modular training framework:
-
-```yaml
-training:
-  checkpoint:
-    output_dir: outputs/my_run
-  tracker:
-    trackers: [swanlab]
-    project_name: my_project
-    run_name: my_run
-```
-
-To log to both supported services:
-
-```yaml
-training:
-  tracker:
-    trackers: [wandb, swanlab]
-    project_name: my_project
-    run_name: my_run
-```
-
-An empty or omitted `trackers` list selects W&B when `project_name` is set.
-Use an explicit `none` entry to disable external tracking:
-
-```yaml
-training:
-  tracker:
-    trackers: [none]
-```
-
-## Validation Videos
-
-SwanLab currently accepts GIF video artifacts. FastVideo converts validation
-MP4 files and in-memory video arrays to GIF automatically before logging them.
-For video files, FastVideo uses the sampling frame rate supplied by the caller,
-or the source file's frame rate when no value is supplied. In-memory arrays use
-the frame rate supplied by the caller. Both forms fall back to 16 FPS when no
-frame rate is available.
-
-For details about configuring validation callbacks, see
-[Training Infrastructure](train_infra.md#callbacks-pluggable-hooks).

docs/training/train_infra.md

@@ -161,9 +161,6 @@ training:
     decay_interval_steps: 0
 ```
 
-See [Training Trackers](trackers.md) to configure Weights & Biases or SwanLab,
-including SwanLab installation and authentication.
-
 ### `callbacks` — Pluggable hooks
 
 Callbacks run at specific points in the training loop (before/after optimizer

examples/inference/basic/basic_dreamx_world.py

@@ -1,64 +0,0 @@
-import os
-
-from fastvideo import VideoGenerator
-
-
-OUTPUT_PATH = os.getenv("DREAMX_WORLD_OUTPUT_PATH", "video_samples_dreamx_world")
-
-
-def _env_int(name: str, default: int) -> int:
-    return int(os.getenv(name, str(default)))
-
-
-def _env_float(name: str, default: float) -> float:
-    return float(os.getenv(name, str(default)))
-
-
-def main():
-    model_name = os.getenv("DREAMX_WORLD_MODEL_DIR", "FastVideo/DreamX-World-5B-Cam-Diffusers")
-    generator = VideoGenerator.from_pretrained(
-        model_name,
-        num_gpus=1,
-        use_fsdp_inference=False,
-        dit_cpu_offload=False,
-        vae_cpu_offload=True,
-        text_encoder_cpu_offload=True,
-        pin_cpu_memory=False,
-        override_pipeline_cls_name="DreamXWorldPipeline",
-    )
-
-    prompt = os.getenv(
-        "DREAMX_WORLD_PROMPT",
-        "A cinematic first-person drive through a futuristic coastal city at "
-        "sunrise, reflective glass towers, clean streets, soft volumetric light.",
-    )
-    image_path = os.getenv(
-        "DREAMX_WORLD_IMAGE_PATH",
-        "https://huggingface.co/datasets/YiYiXu/testing-images/resolve/main/wan_i2v_input.JPG",
-    )
-
-    kwargs = {
-        "output_path": OUTPUT_PATH,
-        "save_video": os.getenv("DREAMX_WORLD_SAVE_VIDEO", "1") != "0",
-        "height": _env_int("DREAMX_WORLD_HEIGHT", 480),
-        "width": _env_int("DREAMX_WORLD_WIDTH", 832),
-        "num_frames": _env_int("DREAMX_WORLD_NUM_FRAMES", 161),
-        "num_inference_steps": _env_int("DREAMX_WORLD_STEPS", 30),
-        "guidance_scale": _env_float("DREAMX_WORLD_GUIDANCE", 5.0),
-        "action_list": os.getenv("DREAMX_WORLD_ACTIONS", "w,d,w").split(","),
-        "action_speed_list": [
-            float(value)
-            for value in os.getenv("DREAMX_WORLD_ACTION_SPEEDS", "4.0,2.0,4.0").split(",")
-        ],
-    }
-    if image_path:
-        kwargs["image_path"] = image_path
-
-    try:
-        generator.generate_video(prompt, **kwargs)
-    finally:
-        generator.shutdown()
-
-
-if __name__ == "__main__":
-    main()

examples/train/configs/consistency_model/wan/causal_cd_t2v.yaml

@@ -1,78 +0,0 @@
-# Causal Consistency Distillation: Wan 2.1 T2V 1.3B Causal
-#
-# ODE-data-free distillation. A frozen teacher takes a single CFG Euler step;
-# the student matches an EMA copy of itself at the next timestep, all under
-# clean-history teacher forcing.
-#
-# All three roles initialize from the SAME checkpoint (the teacher-forcing
-# AR-diffusion model). Point init_from at that checkpoint for a real run.
-
-models:
-  student:
-    _target_: fastvideo.train.models.wan.WanCausalModel
-    init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers
-    trainable: true
-  teacher:
-    _target_: fastvideo.train.models.wan.WanCausalModel
-    init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers
-    trainable: false
-  ema:
-    _target_: fastvideo.train.models.wan.WanCausalModel
-    init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers
-    trainable: false
-
-method:
-  _target_: fastvideo.train.methods.consistency_model.causal_cd.CausalConsistencyDistillationMethod
-  discrete_cd_N: 48
-  guidance_scale: 3.0
-  ema_decay: 0.99
-  ema_start_step: 200
-
-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: 1
-    training_cfg_rate: 0.0
-    seed: 1000
-    num_latent_t: 18
-    num_height: 448
-    num_width: 832
-    num_frames: 69
-
-  optimizer:
-    learning_rate: 2.0e-6
-    betas: [0.9, 0.999]
-    weight_decay: 0.01
-    lr_scheduler: constant
-    lr_warmup_steps: 0
-
-  loop:
-    max_train_steps: 3000
-    gradient_accumulation_steps: 1
-
-  checkpoint:
-    output_dir: outputs/wan2.1_causal_cd
-    training_state_checkpointing_steps: 1000
-    checkpoints_total_limit: 3
-
-  tracker:
-    project_name: distillation_wan_r
-    run_name: wan2.1_causal_cd_shift5
-
-  model:
-    enable_gradient_checkpointing_type: full
-
-callbacks:
-  grad_clip:
-    max_grad_norm: 1.0
-
-pipeline:
-  flow_shift: 5

examples/train/configs/fine_tuning/wan/dfsft_causal_t2v_framewise.yaml

@@ -1,73 +0,0 @@
-# DFSFT (Diffusion-Forcing SFT), frame-wise: Wan 2.1 T2V 1.3B Causal
-#
-# - Student: trainable causal Wan model with a block size of 1 frame
-# - Training: each frame gets its own independent noise level (frame-wise
-#   diffusion forcing), versus the chunk-wise variant that shares one noise
-#   level across num_frames_per_block frames.
-
-models:
-  student:
-    _target_: fastvideo.train.models.wan.WanCausalModel
-    init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers
-    trainable: true
-    num_frames_per_block: 1
-
-method:
-  _target_: fastvideo.train.methods.fine_tuning.dfsft.DiffusionForcingSFTMethod
-  chunk_size: 1
-
-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: 1
-    training_cfg_rate: 0.0
-    seed: 1000
-    num_latent_t: 18
-    num_height: 448
-    num_width: 832
-    num_frames: 69
-
-  optimizer:
-    learning_rate: 2.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/wan2.1_causal_dfsft_framewise
-    training_state_checkpointing_steps: 1000
-    checkpoints_total_limit: 3
-
-  tracker:
-    project_name: distillation_wan_r
-    run_name: wan2.1_causal_dfsft_framewise_shift5_gauss_weight
-
-  model:
-    enable_gradient_checkpointing_type: full
-
-callbacks:
-  grad_clip:
-    max_grad_norm: 1.0
-  validation:
-    pipeline_target: fastvideo.pipelines.basic.wan.wan_causal_pipeline.WanCausalPipeline
-    dataset_file: examples/training/finetune/Wan2.1-VSA/Wan-Syn-Data/validation_4.json
-    every_steps: 50
-    sampling_steps: [40]
-    guidance_scale: 6.0
-    num_frames: 69
-
-pipeline:
-  flow_shift: 5

examples/train/configs/fine_tuning/wan/tfsft_causal_t2v.yaml

@@ -1,72 +0,0 @@
-# TFSFT (Teacher-Forcing SFT): Wan 2.1 T2V 1.3B Causal
-#
-# - Student: trainable causal Wan model
-# - Training: inhomogeneous timesteps per chunk, but the causal transformer
-#   denoises the current block while attending to *clean* history (clean_x),
-#   not its own noisy rollout.
-
-models:
-  student:
-    _target_: fastvideo.train.models.wan.WanCausalModel
-    init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers
-    trainable: true
-
-method:
-  _target_: fastvideo.train.methods.fine_tuning.tfsft.TeacherForcingSFTMethod
-  chunk_size: 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_77x448x832_600k
-    dataloader_num_workers: 4
-    train_batch_size: 1
-    training_cfg_rate: 0.0
-    seed: 1000
-    num_latent_t: 18
-    num_height: 448
-    num_width: 832
-    num_frames: 69
-
-  optimizer:
-    learning_rate: 2.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/wan2.1_causal_tfsft
-    training_state_checkpointing_steps: 1000
-    checkpoints_total_limit: 3
-
-  tracker:
-    project_name: distillation_wan_r
-    run_name: wan2.1_causal_tfsft_shift5_gauss_weight
-
-  model:
-    enable_gradient_checkpointing_type: full
-
-callbacks:
-  grad_clip:
-    max_grad_norm: 1.0
-  validation:
-    pipeline_target: fastvideo.pipelines.basic.wan.wan_causal_pipeline.WanCausalPipeline
-    dataset_file: examples/training/finetune/Wan2.1-VSA/Wan-Syn-Data/validation_4.json
-    every_steps: 50
-    sampling_steps: [40]
-    guidance_scale: 6.0
-    num_frames: 69
-
-pipeline:
-  flow_shift: 5

fastvideo/api/compat.py

@@ -159,6 +159,16 @@ def legacy_from_pretrained_to_config(
             preset_refine["guidance_scale"] = value
         elif key in {"enable_stage_verification", "use_fsdp_inference", "disable_autocast"}:
             engine[key] = value
+        elif key == "batching_mode":
+            engine.setdefault("batching", {})["mode"] = value
+        elif key == "batching_max_size":
+            engine.setdefault("batching", {})["max_size"] = value
+        elif key == "batching_delay_ms":
+            engine.setdefault("batching", {})["delay_ms"] = value
+        elif key == "batching_config":
+            engine.setdefault("batching", {})["config_path"] = value
+        elif key == "enable_batching_metrics":
+            engine.setdefault("batching", {})["enable_metrics"] = value
         elif key == "override_text_encoder_quant":
             quantization["text_encoder_quant"] = value
         elif key == "workload_type":
@@ -244,6 +254,11 @@ def generator_config_to_fastvideo_args(config: GeneratorConfig | Mapping[str, An
         "enable_stage_verification": engine.enable_stage_verification,
         "use_fsdp_inference": engine.use_fsdp_inference,
         "disable_autocast": engine.disable_autocast,
+        "batching_mode": engine.batching.mode,
+        "batching_max_size": engine.batching.max_size,
+        "batching_delay_ms": engine.batching.delay_ms,
+        "batching_config": engine.batching.config_path,
+        "enable_batching_metrics": engine.batching.enable_metrics,
     }
     if normalized.pipeline.workload_type is not None:
         kwargs["workload_type"] = normalized.pipeline.workload_type

fastvideo/api/schema.py

@@ -71,6 +71,15 @@ class QuantizationConfig:
     transformer_quant: str | None = None
 
 
+@dataclass
+class BatchingConfig:
+    mode: Literal["disabled", "dynamic"] = "disabled"
+    max_size: int = 1
+    delay_ms: float = 0.0
+    config_path: str | None = None
+    enable_metrics: bool = False
+
+
 @dataclass
 class EngineConfig:
     num_gpus: int = 1
@@ -82,6 +91,7 @@ class EngineConfig:
     use_fsdp_inference: bool = False
     disable_autocast: bool = False
     quantization: QuantizationConfig | None = None
+    batching: BatchingConfig = field(default_factory=BatchingConfig)
 
 
 @dataclass
@@ -280,6 +290,7 @@ class ServeConfig:
 
 
 __all__ = [
+    "BatchingConfig",
     "CompileConfig",
     "ComponentConfig",
     "ContinuationState",

fastvideo/batching/__init__.py

@@ -0,0 +1,26 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Dynamic generation batching helpers."""
+
+from fastvideo.batching.admission import (
+    AdmissionLimit,
+    BatchAdmissionController,
+    BatchingRule,
+    load_batching_config,
+)
+from fastvideo.batching.signature import (
+    BatchCompatibility,
+    can_dynamic_batch,
+    dynamic_batch_signature,
+    resolution_key,
+)
+
+__all__ = [
+    "AdmissionLimit",
+    "BatchAdmissionController",
+    "BatchCompatibility",
+    "BatchingRule",
+    "can_dynamic_batch",
+    "dynamic_batch_signature",
+    "load_batching_config",
+    "resolution_key",
+]

fastvideo/batching/admission.py

@@ -0,0 +1,297 @@
+# SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
+
+import json
+import os
+from dataclasses import dataclass
+from difflib import get_close_matches
+from typing import Any
+
+from fastvideo.batching.signature import resolution_key
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.logger import init_logger
+
+logger = init_logger(__name__)
+
+_BYTES_PER_GB = 1024**3
+_BATCHING_RULE_KEYS = frozenset({
+    "model",
+    "model_contains",
+    "resolution",
+    "device_memory_gb_min",
+    "device_memory_gb_max",
+    "offload",
+    "max_batch_size",
+    "max_cost",
+    "calibration",
+})
+
+
+@dataclass(frozen=True)
+class AdmissionLimit:
+    max_batch_size: int
+    max_cost: float | None = None
+    cap_reason: str | None = None
+
+    def reject_reason(self, *, batch_size: int, batch_cost: float) -> str | None:
+        if batch_size > self.max_batch_size:
+            return self.cap_reason or f"config_cap:{self.max_batch_size}"
+        if self.max_cost is not None and batch_cost > self.max_cost:
+            return f"cost_budget:{batch_cost:.0f}>{self.max_cost:.0f}"
+        return None
+
+    def stop_reason_for_next_cost(self, next_batch_cost: float) -> str | None:
+        if self.max_cost is not None and next_batch_cost > self.max_cost:
+            return f"cost_budget_next:{next_batch_cost:.0f}>{self.max_cost:.0f}"
+        return None
+
+
+@dataclass(frozen=True)
+class BatchingRule:
+    model: str | None = None
+    model_contains: str | None = None
+    resolution: str | None = None
+    device_memory_gb_min: float | None = None
+    device_memory_gb_max: float | None = None
+    offload: bool | None = None
+    max_batch_size: int = 1
+    max_cost: float | None = None
+    source: str = "user"
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any], *, source: str) -> BatchingRule:
+        if not isinstance(data, dict):
+            raise ValueError(f"batching config rule from {source} must be an object, got {type(data).__name__}")
+        _validate_rule_keys(data, source=source)
+        if "max_batch_size" not in data:
+            raise ValueError("batching config rule requires max_batch_size")
+
+        rule = cls(
+            model=_optional_str(data.get("model")),
+            model_contains=_optional_str(data.get("model_contains")),
+            resolution=_optional_str(data.get("resolution")),
+            device_memory_gb_min=_optional_float(data.get("device_memory_gb_min")),
+            device_memory_gb_max=_optional_float(data.get("device_memory_gb_max")),
+            offload=_optional_bool(data.get("offload")),
+            max_batch_size=int(data["max_batch_size"]),
+            max_cost=_optional_float(data.get("max_cost")),
+            source=source,
+        )
+        rule.validate()
+        return rule
+
+    def validate(self) -> None:
+        if self.model is not None and self.model_contains is not None:
+            raise ValueError("batching config rule cannot set both model and model_contains")
+        if self.model is None and self.model_contains is None:
+            raise ValueError("batching config rule requires model or model_contains")
+        if self.max_batch_size < 1:
+            raise ValueError("batching config rule max_batch_size must be >= 1")
+        if self.max_cost is not None and self.max_cost <= 0.0:
+            raise ValueError("batching config rule max_cost must be > 0")
+        if (self.device_memory_gb_min is not None and self.device_memory_gb_max is not None
+                and self.device_memory_gb_min > self.device_memory_gb_max):
+            raise ValueError("batching config rule device_memory_gb_min must be <= device_memory_gb_max")
+
+    def matches(
+        self,
+        *,
+        model_path: str,
+        resolution: str | None,
+        device_memory_gb: float | None,
+        offload: bool,
+    ) -> bool:
+        if self.model is not None and self.model != model_path:
+            return False
+        if self.model_contains is not None and self.model_contains not in model_path:
+            return False
+        if self.resolution not in (None, "*") and self.resolution != resolution:
+            return False
+        if self.offload is not None and self.offload != offload:
+            return False
+        if device_memory_gb is None:
+            return True
+        if self.device_memory_gb_min is not None and device_memory_gb < self.device_memory_gb_min:
+            return False
+        return not (self.device_memory_gb_max is not None and device_memory_gb > self.device_memory_gb_max)
+
+
+class BatchAdmissionController:
+
+    def __init__(self, fastvideo_args: FastVideoArgs, *, gpu_id: int = 0):
+        self._mode = fastvideo_args.batching_mode
+        self._user_max_batch_size = max(1, int(fastvideo_args.batching_max_size))
+        self._model_path = fastvideo_args.model_path
+        self._offload = bool(fastvideo_args.dit_cpu_offload or fastvideo_args.dit_layerwise_offload)
+        self._device_memory_gb = self._get_device_memory_gb(gpu_id)
+        self._rules = load_batching_config(fastvideo_args.batching_config)
+        self._pipeline_config = fastvideo_args.pipeline_config
+
+        if self.enabled:
+            logger.info(
+                "Batch admission enabled: user_max=%d, device_memory=%.1fGiB, rules=%d",
+                self._user_max_batch_size,
+                self._device_memory_gb or 0.0,
+                len(self._rules),
+            )
+
+    @property
+    def enabled(self) -> bool:
+        return self._mode == "dynamic" and self._user_max_batch_size > 1
+
+    def reject_reason_for_candidate(self, current_requests: list[Any], candidate_request: Any) -> str | None:
+        if not self.enabled:
+            return None
+        proposed = current_requests + [candidate_request]
+        limit = self.limit_for(proposed[0])
+        return limit.reject_reason(
+            batch_size=len(proposed),
+            batch_cost=self.estimate_batch_cost(proposed),
+        )
+
+    def batch_is_full(self, requests: list[Any]) -> bool:
+        if not self.enabled or not requests:
+            return len(requests) >= self._user_max_batch_size
+
+        limit = self.limit_for(requests[0])
+        if len(requests) >= limit.max_batch_size:
+            return True
+
+        next_cost = self.estimate_batch_cost(requests + [requests[0]])
+        return limit.max_cost is not None and next_cost > limit.max_cost
+
+    def limit_reason_for_batch(self, requests: list[Any]) -> str | None:
+        if not self.enabled or not requests:
+            return None
+
+        limit = self.limit_for(requests[0])
+        if len(requests) >= limit.max_batch_size:
+            return limit.cap_reason or f"config_cap:{limit.max_batch_size}"
+
+        next_cost = self.estimate_batch_cost(requests + [requests[0]])
+        return limit.stop_reason_for_next_cost(next_cost)
+
+    def max_admissible_batch_size(self, request: Any) -> int:
+        return self.limit_for(request).max_batch_size
+
+    def limit_for(self, request: Any) -> AdmissionLimit:
+        rules = self._matching_rules(request)
+        if not rules:
+            return AdmissionLimit(max_batch_size=self._user_max_batch_size)
+
+        config_cap = min(rule.max_batch_size for rule in rules)
+        max_batch_size = min(self._user_max_batch_size, config_cap)
+        cap_reason = f"config_cap:{max_batch_size}" if max_batch_size < self._user_max_batch_size else None
+        costs = [rule.max_cost for rule in rules if rule.max_cost is not None]
+        return AdmissionLimit(
+            max_batch_size=max(1, max_batch_size),
+            max_cost=min(costs) if costs else None,
+            cap_reason=cap_reason,
+        )
+
+    def estimate_batch_cost(self, requests: list[Any]) -> float:
+        return sum(float(self._pipeline_config.estimate_request_cost(request)) for request in requests)
+
+    def _matching_rules(self, request: Any) -> list[BatchingRule]:
+        return [
+            rule for rule in self._rules if rule.matches(
+                model_path=self._model_path,
+                resolution=resolution_key(request),
+                device_memory_gb=self._device_memory_gb,
+                offload=self._offload,
+            )
+        ]
+
+    @staticmethod
+    def _get_device_memory_gb(gpu_id: int) -> float | None:
+        try:
+            from fastvideo.platforms import current_platform
+
+            return current_platform.get_device_total_memory(gpu_id) / _BYTES_PER_GB
+        except Exception:
+            return None
+
+
+def load_batching_config(path: str | None) -> list[BatchingRule]:
+    if path is None:
+        return []
+
+    with open(path, encoding="utf-8") as f:
+        payload = json.load(f)
+
+    source = os.path.abspath(path)
+    entries = _config_entries(payload)
+    rules = [BatchingRule.from_dict(entry, source=source) for entry in entries]
+    if not rules:
+        raise ValueError(f"batching config {source} does not contain any rules")
+    return rules
+
+
+def _config_entries(payload: Any) -> list[dict[str, Any]]:
+    if isinstance(payload, dict) and payload.get("schema_version") not in (None, 1):
+        raise ValueError("batching config schema_version must be 1")
+    if isinstance(payload, dict) and isinstance(payload.get("rules"), list):
+        return payload["rules"]
+    if isinstance(payload, list):
+        return payload
+    if isinstance(payload, dict):
+        entries: list[dict[str, Any]] = []
+        for key, value in payload.items():
+            if key == "schema_version" or not isinstance(value, dict):
+                continue
+            model, _sep, resolution = key.partition("|")
+            entry = dict(value)
+            if model:
+                entry.setdefault("model", model)
+            if resolution:
+                entry.setdefault("resolution", resolution)
+            entries.append(entry)
+        return entries
+    raise ValueError("batching config must be a {'schema_version': 1, 'rules': [...]} object, "
+                     "a list of rules, or a mapping keyed by model|resolution")
+
+
+def _validate_rule_keys(data: dict[str, Any], *, source: str) -> None:
+    unknown = sorted(set(data) - _BATCHING_RULE_KEYS)
+    if not unknown:
+        return
+
+    hints = []
+    for key in unknown:
+        matches = get_close_matches(key, _BATCHING_RULE_KEYS, n=1)
+        if matches:
+            hints.append(f"{key!r} (did you mean {matches[0]!r}?)")
+        else:
+            hints.append(repr(key))
+    raise ValueError(f"batching config rule from {source} contains unknown key(s): {', '.join(hints)}")
+
+
+def _optional_str(value: Any) -> str | None:
+    if value is None:
+        return None
+    return str(value)
+
+
+def _optional_float(value: Any) -> float | None:
+    if value is None:
+        return None
+    return float(value)
+
+
+def _optional_bool(value: Any) -> bool | None:
+    if value is None:
+        return None
+    if isinstance(value, bool):
+        return value
+    if isinstance(value, int | float):
+        if value == 1.0:
+            return True
+        if value == 0.0:
+            return False
+    if isinstance(value, str):
+        lowered = value.strip().lower()
+        if lowered in ("1", "true", "yes", "y", "on"):
+            return True
+        if lowered in ("0", "false", "no", "n", "off"):
+            return False
+    raise ValueError(f"cannot parse boolean batching config value: {value!r}")

fastvideo/batching/signature.py

@@ -0,0 +1,178 @@
+# SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
+
+import dataclasses
+from dataclasses import dataclass
+from enum import Enum
+from typing import Any
+
+from fastvideo.api.sampling_param import SamplingParam
+
+_SIGNATURE_EXCLUDED_FIELDS = frozenset({
+    "prompt",
+    "prompt_path",
+    "output_path",
+    "output_video_name",
+    "seed",
+    "save_video",
+    "return_frames",
+})
+
+_UNSUPPORTED_DYNAMIC_BATCH_FIELDS = frozenset({
+    "image_path",
+    "pil_image",
+    "video_path",
+    "mouse_cond",
+    "keyboard_cond",
+    "grid_sizes",
+    "pose",
+    "camera_states",
+    "camera_trajectory",
+    "action_list",
+    "action_speed_list",
+    "gt_latents",
+    "conditioning_mask",
+    "c2ws_plucker_emb",
+    "refine_from",
+    "stage1_video",
+    "trajectory_type",
+    "movement_distance",
+    "camera_rotation",
+    "ltx2_images",
+    "ltx2_conditioning_latent_stage1",
+    "ltx2_conditioning_latent_stage2",
+    "ltx2_video_conditions",
+    "init_audio",
+    "inpaint_audio",
+    "inpaint_mask",
+    "continuation_state",
+})
+
+_UNSUPPORTED_EXTRA_KEYS = frozenset({
+    "ltx2_audio_latents",
+    "ltx2_audio_clean_latent",
+    "ltx2_audio_denoise_mask",
+    "audio_num_frames",
+    "video_position_offset_sec",
+})
+
+
+@dataclass(frozen=True)
+class BatchCompatibility:
+    can_batch: bool
+    reason: str | None = None
+
+
+def resolution_key(request: Any) -> str:
+    height = _first_scalar(getattr(request, "height", None))
+    width = _first_scalar(getattr(request, "width", None))
+    num_frames = _first_scalar(getattr(request, "num_frames", None))
+    return f"{height}x{width}x{num_frames}"
+
+
+def dynamic_batch_signature(
+    request: SamplingParam,
+    *,
+    extra: dict[str, Any] | None = None,
+) -> tuple[tuple[str, Any], ...]:
+    """Build a hashable compatibility signature for a generation request."""
+    signature_items: list[tuple[str, Any]] = []
+    for field in dataclasses.fields(request):
+        if field.name in _SIGNATURE_EXCLUDED_FIELDS:
+            continue
+        signature_items.append((field.name, _freeze_signature_value(getattr(request, field.name, None))))
+    if extra:
+        signature_items.append(("extra", _freeze_signature_value(extra)))
+    return tuple(signature_items)
+
+
+def can_dynamic_batch(
+    base: SamplingParam,
+    candidate: SamplingParam,
+    *,
+    base_extra: dict[str, Any] | None = None,
+    candidate_extra: dict[str, Any] | None = None,
+) -> BatchCompatibility:
+    """Return whether two FastVideo generation requests can be merged."""
+    base_ready = _request_is_batchable(base, extra=base_extra)
+    if not base_ready.can_batch:
+        return base_ready
+    candidate_ready = _request_is_batchable(candidate, extra=candidate_extra)
+    if not candidate_ready.can_batch:
+        return candidate_ready
+
+    base_sig = dynamic_batch_signature(base, extra=base_extra)
+    candidate_sig = dynamic_batch_signature(candidate, extra=candidate_extra)
+    if base_sig == candidate_sig:
+        return BatchCompatibility(can_batch=True)
+
+    mismatch = _first_mismatch(base_sig, candidate_sig)
+    return BatchCompatibility(can_batch=False, reason=mismatch or "signature_mismatch")
+
+
+def _request_is_batchable(
+    request: SamplingParam,
+    *,
+    extra: dict[str, Any] | None = None,
+) -> BatchCompatibility:
+    if not isinstance(request.prompt, str):
+        return BatchCompatibility(can_batch=False, reason="prompt_type")
+    if request.prompt_path is not None:
+        return BatchCompatibility(can_batch=False, reason="prompt_path")
+    if request.num_videos_per_prompt != 1:
+        return BatchCompatibility(can_batch=False, reason="num_videos_per_prompt")
+    if request.return_continuation_state:
+        return BatchCompatibility(can_batch=False, reason="return_continuation_state")
+
+    for name in _UNSUPPORTED_DYNAMIC_BATCH_FIELDS:
+        value = getattr(request, name, None)
+        if _is_present(value):
+            return BatchCompatibility(can_batch=False, reason=name)
+
+    if extra:
+        unsupported = sorted(set(extra) & _UNSUPPORTED_EXTRA_KEYS)
+        if unsupported:
+            return BatchCompatibility(can_batch=False, reason=f"extra.{unsupported[0]}")
+
+    return BatchCompatibility(can_batch=True)
+
+
+def _freeze_signature_value(value: Any) -> Any:
+    if isinstance(value, str | int | float | bool | type(None)):
+        return value
+    if isinstance(value, Enum):
+        return value.value
+    if isinstance(value, dict):
+        return tuple(
+            (str(key), _freeze_signature_value(item)) for key, item in sorted(value.items(), key=lambda kv: str(kv[0])))
+    if isinstance(value, list | tuple):
+        return tuple(_freeze_signature_value(item) for item in value)
+    return repr(value)
+
+
+def _is_present(value: Any) -> bool:
+    if value is None:
+        return False
+    if value is False:
+        return False
+    return not (isinstance(value, list | tuple | dict | set) and not value)
+
+
+def _first_scalar(value: Any) -> Any:
+    if isinstance(value, list | tuple):
+        return value[0] if value else None
+    return value
+
+
+def _first_mismatch(
+    base_sig: tuple[tuple[str, Any], ...],
+    candidate_sig: tuple[tuple[str, Any], ...],
+) -> str | None:
+    if len(base_sig) != len(candidate_sig):
+        return "sampling_params"
+    for (name, base_value), (candidate_name, candidate_value) in zip(base_sig, candidate_sig, strict=True):
+        if name != candidate_name:
+            return "sampling_params"
+        if base_value != candidate_value:
+            return f"sampling_params.{name}"
+    return None

fastvideo/configs/models/dits/__init__.py

@@ -1,6 +1,5 @@
 from fastvideo.configs.models.dits.cosmos import CosmosVideoConfig
 from fastvideo.configs.models.dits.cosmos2_5 import Cosmos25VideoConfig
-from fastvideo.configs.models.dits.dreamx_world import DreamXWorldARConfig, DreamXWorldConfig
 from fastvideo.configs.models.dits.flux_2 import Flux2Config
 from fastvideo.configs.models.dits.hunyuangamecraft import HunyuanGameCraftConfig
 from fastvideo.configs.models.dits.hunyuanvideo import HunyuanVideoConfig
@@ -14,7 +13,7 @@ from fastvideo.configs.models.dits.hyworld import HYWorldConfig
 from fastvideo.configs.models.dits.kandinsky5 import Kandinsky5VideoConfig
 
 __all__ = [
-    "HunyuanVideoConfig", "HunyuanVideo15Config", "HunyuanGameCraftConfig", "WanVideoConfig", "DreamXWorldConfig",
-    "DreamXWorldARConfig", "CosmosVideoConfig", "Cosmos25VideoConfig", "LongCatVideoConfig", "LTX2VideoConfig",
-    "HYWorldConfig", "Kandinsky5VideoConfig", "MagiHumanVideoConfig", "StableAudioConfig", "Flux2Config"
+    "HunyuanVideoConfig", "HunyuanVideo15Config", "HunyuanGameCraftConfig", "WanVideoConfig", "CosmosVideoConfig",
+    "Cosmos25VideoConfig", "LongCatVideoConfig", "LTX2VideoConfig", "HYWorldConfig", "Kandinsky5VideoConfig",
+    "MagiHumanVideoConfig", "StableAudioConfig", "Flux2Config"
 ]

fastvideo/configs/models/dits/dreamx_world.py

@@ -1,71 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-from dataclasses import dataclass, field
-
-from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig
-from fastvideo.configs.models.dits.wanvideo import WanVideoArchConfig
-
-
-@dataclass
-class DreamXWorldArchConfig(WanVideoArchConfig):
-    """DreamX-World DiT config with camera PRoPE control fields."""
-
-    add_control_adapter: bool = True
-    cam_method: str | None = "prope"
-    attn_compress: int = 1
-    cam_self_attn_layers: tuple[int, ...] | None = None
-
-
-@dataclass
-class DreamXWorldConfig(DiTConfig):
-    arch_config: DiTArchConfig = field(default_factory=DreamXWorldArchConfig)
-
-    prefix: str = "Wan"
-
-
-@dataclass
-class DreamXWorldARArchConfig(DreamXWorldArchConfig):
-    """DreamX-World-5B autoregressive causal DiT config."""
-
-    model_type: str = "ti2v"
-    patch_size: tuple[int, int, int] = (1, 2, 2)
-    text_len: int = 512
-    text_dim: int = 4096
-    freq_dim: int = 256
-    attn_compress: int = 4
-    cam_self_attn_layers: tuple[int, ...] | None = tuple(range(30))
-    local_attn_size: int = 12
-    sink_size: int = 3
-    num_frames_per_block: int = 3
-    rope_cache_policy: str = "block_relativistic"
-    # The official AR checkpoint (AMAP-ML/DreamX-World ``model.safetensors``)
-    # already uses FastVideo's native key names and the converter copies the
-    # tensors verbatim, so every rule is an identity. The rules enumerate the
-    # full state-dict surface of ``DreamXWorldARTransformer3DModel`` (norm1 /
-    # norm2 / head.norm are affine-free and have no parameters).
-    param_names_mapping: dict = field(
-        default_factory=lambda: {
-            r"^patch_embedding\.(.*)$": r"patch_embedding.\1",
-            r"^text_embedding\.([02])\.(.*)$": r"text_embedding.\1.\2",
-            r"^time_embedding\.([02])\.(.*)$": r"time_embedding.\1.\2",
-            r"^time_projection\.1\.(.*)$": r"time_projection.1.\1",
-            r"^blocks\.(\d+)\.self_attn\.(q|k|v|o)\.(.*)$": r"blocks.\1.self_attn.\2.\3",
-            r"^blocks\.(\d+)\.self_attn\.norm_(q|k)\.weight$": r"blocks.\1.self_attn.norm_\2.weight",
-            r"^blocks\.(\d+)\.cross_attn\.(q|k|v|o)\.(.*)$": r"blocks.\1.cross_attn.\2.\3",
-            r"^blocks\.(\d+)\.cross_attn\.norm_(q|k)\.weight$": r"blocks.\1.cross_attn.norm_\2.weight",
-            r"^blocks\.(\d+)\.cam_self_attn\.(q_proj|k_proj|v_proj|out_proj)\.(.*)$": r"blocks.\1.cam_self_attn.\2.\3",
-            r"^blocks\.(\d+)\.cam_self_attn\.norm_(q|k)\.weight$": r"blocks.\1.cam_self_attn.norm_\2.weight",
-            r"^blocks\.(\d+)\.norm3\.(.*)$": r"blocks.\1.norm3.\2",
-            r"^blocks\.(\d+)\.ffn\.([02])\.(.*)$": r"blocks.\1.ffn.\2.\3",
-            r"^blocks\.(\d+)\.modulation$": r"blocks.\1.modulation",
-            r"^head\.head\.(.*)$": r"head.head.\1",
-            r"^head\.modulation$": r"head.modulation",
-        })
-    reverse_param_names_mapping: dict = field(default_factory=dict)
-    lora_param_names_mapping: dict = field(default_factory=dict)
-
-
-@dataclass
-class DreamXWorldARConfig(DreamXWorldConfig):
-    arch_config: DiTArchConfig = field(default_factory=DreamXWorldARArchConfig)
-
-    prefix: str = "Wan"

fastvideo/configs/pipelines/__init__.py

@@ -1,7 +1,6 @@
 from fastvideo.configs.pipelines.base import PipelineConfig
 from fastvideo.configs.pipelines.cosmos import CosmosConfig
 from fastvideo.configs.pipelines.cosmos2_5 import Cosmos25Config
-from fastvideo.configs.pipelines.dreamx_world import DreamXWorld5BARPipelineConfig, DreamXWorld5BCamPipelineConfig
 from fastvideo.configs.pipelines.hunyuan import FastHunyuanConfig, HunyuanConfig
 from fastvideo.configs.pipelines.hunyuan15 import Hunyuan15T2V480PConfig, Hunyuan15T2V720PConfig
 from fastvideo.configs.pipelines.hunyuangamecraft import HunyuanGameCraftPipelineConfig
@@ -17,6 +16,5 @@ __all__ = [
     "HunyuanConfig", "FastHunyuanConfig", "HunyuanGameCraftPipelineConfig", "PipelineConfig", "Hunyuan15T2V480PConfig",
     "Hunyuan15T2V720PConfig", "WanT2V480PConfig", "WanI2V480PConfig", "WanT2V720PConfig", "WanI2V720PConfig",
     "SelfForcingWanT2V480PConfig", "LucyEditDevConfig", "CosmosConfig", "Cosmos25Config", "LTX2T2VConfig",
-    "DreamXWorld5BCamPipelineConfig", "DreamXWorld5BARPipelineConfig", "HYWorldConfig", "MatrixGame2I2V480PConfig",
-    "MatrixGame3I2V720PConfig", "get_pipeline_config_cls_from_name"
+    "HYWorldConfig", "MatrixGame2I2V480PConfig", "MatrixGame3I2V720PConfig", "get_pipeline_config_cls_from_name"
 ]

fastvideo/configs/pipelines/base.py

@@ -276,6 +276,27 @@ class PipelineConfig:
                 f"Length of text postprocess functions ({len(self.postprocess_text_funcs)}) must be equal to length of text preprocessing functions ({len(self.preprocess_text_funcs)})"
             )
 
+    def estimate_request_cost(self, request: Any) -> float:
+        """Estimate relative memory/compute cost for batching admission.
+
+        The default is intentionally simple and model-agnostic: pixel count
+        times frame count. Pipeline subclasses can override this when they have
+        calibrated costs.
+        """
+        height = getattr(request, "height", None)
+        width = getattr(request, "width", None)
+        num_frames = getattr(request, "num_frames", None)
+        if isinstance(height, list):
+            height = height[0] if height else None
+        if isinstance(width, list):
+            width = width[0] if width else None
+        if isinstance(num_frames, list):
+            num_frames = num_frames[0] if num_frames else None
+        height = int(height or 1)
+        width = int(width or 1)
+        num_frames = int(num_frames or 1)
+        return float(max(1, height) * max(1, width) * max(1, num_frames))
+
     def dump_to_json(self, file_path: str):
         output_dict = shallow_asdict(self)
         del_keys = []

fastvideo/configs/pipelines/dreamx_world.py

@@ -1,127 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""DreamX-World-5B-Cam FastVideo model configuration helpers."""
-
-from collections.abc import Callable
-from dataclasses import dataclass, field
-
-import torch
-
-from fastvideo.configs.models import DiTConfig, EncoderConfig, VAEConfig
-from fastvideo.configs.models.dits.dreamx_world import (DreamXWorldARArchConfig, DreamXWorldARConfig,
-                                                        DreamXWorldArchConfig, DreamXWorldConfig)
-from fastvideo.configs.models.encoders import BaseEncoderOutput, T5Config
-from fastvideo.configs.models.encoders.t5 import T5ArchConfig
-from fastvideo.configs.models.vaes import WanVAEConfig
-from fastvideo.configs.pipelines.base import PipelineConfig
-from fastvideo.configs.pipelines.wan import LucyEditDevConfig, t5_postprocess_text
-
-
-def make_dreamx_world_5b_cam_dit_config() -> DreamXWorldConfig:
-    """Return the DreamX-World DiT config matching DreamX-World-5B-Cam."""
-    return DreamXWorldConfig(arch_config=DreamXWorldArchConfig(
-        num_attention_heads=24,
-        attention_head_dim=128,
-        in_channels=48,
-        out_channels=48,
-        ffn_dim=14336,
-        num_layers=30,
-        cross_attn_norm=True,
-        qk_norm="rms_norm_across_heads",
-        add_control_adapter=True,
-        cam_method="prope",
-        attn_compress=1,
-        cam_self_attn_layers=None,
-    ))
-
-
-def make_dreamx_world_5b_ar_dit_config() -> DreamXWorldARConfig:
-    """Return the DreamX-World-5B autoregressive causal DiT config."""
-    return DreamXWorldARConfig(arch_config=DreamXWorldARArchConfig(
-        model_type="ti2v",
-        num_attention_heads=24,
-        attention_head_dim=128,
-        in_channels=48,
-        out_channels=48,
-        ffn_dim=14336,
-        num_layers=30,
-        cross_attn_norm=True,
-        qk_norm=True,
-        add_control_adapter=True,
-        cam_method="prope",
-        attn_compress=4,
-        cam_self_attn_layers=tuple(range(30)),
-        local_attn_size=12,
-        sink_size=3,
-        num_frames_per_block=3,
-    ))
-
-
-def make_dreamx_world_5b_cam_vae_config() -> WanVAEConfig:
-    """Return the Wan2.2 48-channel VAE config used by DreamX-World-5B-Cam."""
-    return LucyEditDevConfig().vae_config
-
-
-def make_dreamx_world_5b_cam_text_encoder_config() -> T5Config:
-    """Return the UMT5-XXL text encoder config used by DreamX-World-5B-Cam."""
-    return T5Config(
-        arch_config=T5ArchConfig(
-            vocab_size=256384,
-            d_model=4096,
-            d_kv=64,
-            d_ff=10240,
-            num_layers=24,
-            num_decoder_layers=None,
-            num_heads=64,
-            relative_attention_num_buckets=32,
-            dropout_rate=0.0,
-            text_len=512,
-            feed_forward_proj="gelu",
-            is_encoder_decoder=False,
-        ),
-        prefix="umt5",
-    )
-
-
-@dataclass
-class DreamXWorld5BCamPipelineConfig(PipelineConfig):
-    """Pipeline config for the first-scope DreamX-World-5B-Cam mode."""
-
-    dit_config: DiTConfig = field(default_factory=make_dreamx_world_5b_cam_dit_config)
-    vae_config: VAEConfig = field(default_factory=make_dreamx_world_5b_cam_vae_config)
-    text_encoder_configs: tuple[EncoderConfig,
-                                ...] = field(default_factory=lambda: (make_dreamx_world_5b_cam_text_encoder_config(), ))
-    postprocess_text_funcs: tuple[Callable[[BaseEncoderOutput], torch.Tensor],
-                                  ...] = field(default_factory=lambda: (t5_postprocess_text, ))
-    text_encoder_precisions: tuple[str, ...] = field(default_factory=lambda: ("bf16", ))
-    flow_shift: float | None = 3.0
-    ti2v_task: bool = True
-    expand_timesteps: bool = True
-    vae_tiling: bool = False
-    vae_sp: bool = False
-    vae_precision: str = "fp32"
-    vae_decode_precision: str | None = "bf16"
-    dit_precision: str = "bf16"
-
-    def __post_init__(self) -> None:
-        self.vae_config.load_encoder = True
-        self.vae_config.load_decoder = True
-        self.dit_config.expand_timesteps = self.expand_timesteps
-
-
-@dataclass
-class DreamXWorld5BARPipelineConfig(DreamXWorld5BCamPipelineConfig):
-    """Pipeline config for DreamX-World-5B autoregressive forcing."""
-
-    dit_config: DiTConfig = field(default_factory=make_dreamx_world_5b_ar_dit_config)
-    flow_shift: float | None = 5.0
-    ti2v_task: bool = True
-    is_causal: bool = True
-    dmd_denoising_steps: tuple[int, ...] = (1000, 750, 500, 250)
-    warp_denoising_step: bool = True
-    context_noise: float = 0.1
-    num_frames_per_block: int = 3
-    color_correction_strength: float = 1.0
-
-    def __post_init__(self) -> None:
-        super().__post_init__()
-        self.dit_config.expand_timesteps = True

fastvideo/distributed/parallel_state.py

@@ -23,7 +23,6 @@ If you only need to use the distributed environment without model parallelism,
  you can skip the model parallel initialization and destruction steps.
 """
 import contextlib
-import gc
 import os
 import pickle
 import weakref
@@ -1007,13 +1006,6 @@ def cleanup_dist_env_and_memory(shutdown_ray: bool = False):
     if shutdown_ray:
         import ray  # Lazy import Ray
         ray.shutdown()
-    # Actually free GPU memory, as the name promises. FSDP-wrapped modules sit
-    # in reference cycles (module <-> FSDP state <-> hooks), so dropping the
-    # last user reference does not free their parameters until a gc pass runs.
-    # Without this, back-to-back model-load tests in one pytest session OOM.
-    gc.collect()
-    if torch.cuda.is_available():
-        torch.cuda.empty_cache()
 
 
 def get_same_node_ranks(pg: ProcessGroup | StatelessProcessGroup, source_rank: int = 0) -> list[int]:

fastvideo/entrypoints/openai/api_server.py

@@ -10,6 +10,7 @@ from fastapi.middleware.cors import CORSMiddleware
 
 from fastvideo.api.presets import validate_preset_selection
 from fastvideo.api.schema import GenerationRequest
+from fastvideo.entrypoints.openai.batching import VideoBatchScheduler
 from fastvideo.entrypoints.openai.state import (
     DEFAULT_OUTPUT_DIR,
     clear_state,
@@ -59,11 +60,29 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]:
     generator = VideoGenerator.from_fastvideo_args(args)
     logger.info("Model loaded successfully.")
 
-    set_state(generator, args, output_dir, default_request=default_request)
+    video_batch_scheduler: VideoBatchScheduler | None = None
+    if args.batching_mode == "dynamic" and args.batching_max_size > 1:
+        video_batch_scheduler = VideoBatchScheduler(generator, args)
+        await video_batch_scheduler.start()
+        logger.info(
+            "Started dynamic video batch scheduler: max_size=%d delay_ms=%.2f",
+            args.batching_max_size,
+            args.batching_delay_ms,
+        )
+
+    set_state(
+        generator,
+        args,
+        output_dir,
+        default_request=default_request,
+        video_batch_scheduler=video_batch_scheduler,
+    )
 
     yield  # server is running
 
     logger.info("Shutting down — releasing model resources ...")
+    if video_batch_scheduler is not None:
+        await video_batch_scheduler.stop()
     generator.shutdown()
     clear_state()
     logger.info("Shutdown complete.")

fastvideo/entrypoints/openai/batching.py

@@ -0,0 +1,183 @@
+# SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
+
+import asyncio
+import time
+from collections import deque
+from dataclasses import dataclass
+from typing import Any
+
+from fastvideo.api.sampling_param import SamplingParam
+from fastvideo.batching.signature import can_dynamic_batch
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+@dataclass
+class _VideoBatchJob:
+    request_id: str
+    kwargs: dict[str, Any]
+    future: asyncio.Future
+    enqueue_time: float
+
+
+class VideoBatchScheduler:
+    """Async FIFO scheduler for OpenAI-compatible video generation."""
+
+    def __init__(self, generator: Any, fastvideo_args: FastVideoArgs) -> None:
+        self._generator = generator
+        self._fastvideo_args = fastvideo_args
+        self._queue: asyncio.Queue[_VideoBatchJob | None] = asyncio.Queue()
+        self._pending: deque[_VideoBatchJob] = deque()
+        self._task: asyncio.Task | None = None
+        self._stopped = False
+
+    @property
+    def enabled(self) -> bool:
+        return self._fastvideo_args.batching_mode == "dynamic" and self._fastvideo_args.batching_max_size > 1
+
+    async def start(self) -> None:
+        if self._task is not None:
+            return
+        self._task = asyncio.create_task(self._run(), name="fastvideo-video-batch-scheduler")
+
+    async def stop(self) -> None:
+        self._stopped = True
+        await self._queue.put(None)
+        if self._task is not None:
+            await self._task
+            self._task = None
+
+    async def submit(self, request_id: str, kwargs: dict[str, Any]) -> Any:
+        if self._stopped:
+            raise RuntimeError("Video batch scheduler is stopped; cannot submit new requests")
+        loop = asyncio.get_running_loop()
+        future = loop.create_future()
+        await self._queue.put(
+            _VideoBatchJob(
+                request_id=request_id,
+                kwargs=dict(kwargs),
+                future=future,
+                enqueue_time=time.perf_counter(),
+            ))
+        return await future
+
+    async def _run(self) -> None:
+        while not self._stopped:
+            job = await self._get_next_job()
+            if job is None:
+                break
+            batch = await self._collect_batch(job)
+            await self._dispatch(batch)
+
+        while self._pending:
+            pending = self._pending.popleft()
+            if not pending.future.done():
+                pending.future.set_exception(RuntimeError("Video batch scheduler stopped before dispatch"))
+
+    async def _get_next_job(self) -> _VideoBatchJob | None:
+        if self._pending:
+            return self._pending.popleft()
+        return await self._queue.get()
+
+    async def _collect_batch(self, first: _VideoBatchJob) -> list[_VideoBatchJob]:
+        batch = [first]
+        max_size = self._fastvideo_args.batching_max_size
+        delay_s = max(0.0, self._fastvideo_args.batching_delay_ms / 1000.0)
+        deadline = first.enqueue_time + delay_s
+
+        while len(batch) < max_size:
+            if delay_s > 0:
+                timeout = deadline - time.perf_counter()
+                if timeout <= 0:
+                    break
+                try:
+                    candidate = await asyncio.wait_for(self._get_next_job(), timeout=timeout)
+                except TimeoutError:
+                    break
+            else:
+                # delay=0 means "don't wait": greedily drain whatever is
+                # already queued so max_size still coalesces.
+                if self._pending:
+                    candidate = self._pending.popleft()
+                else:
+                    try:
+                        candidate = self._queue.get_nowait()
+                    except asyncio.QueueEmpty:
+                        break
+            if candidate is None:
+                await self._queue.put(None)
+                break
+            if self._jobs_are_compatible(batch[0], candidate):
+                batch.append(candidate)
+                continue
+            self._pending.appendleft(candidate)
+            break
+        return batch
+
+    async def _dispatch(self, batch: list[_VideoBatchJob]) -> None:
+        loop = asyncio.get_running_loop()
+        request_ids = [job.request_id for job in batch]
+        queue_wait_ms = (time.perf_counter() - min(job.enqueue_time for job in batch)) * 1000.0
+        if self._fastvideo_args.enable_batching_metrics:
+            logger.info(
+                "Dispatching video batch: request_ids=%s size=%d queue_wait_ms=%.2f",
+                request_ids,
+                len(batch),
+                queue_wait_ms,
+            )
+
+        try:
+            results = await loop.run_in_executor(
+                None,
+                lambda: self._generator.generate_video_batch([job.kwargs for job in batch]),
+            )
+        except Exception as exc:
+            for job in batch:
+                if not job.future.done():
+                    job.future.set_exception(exc)
+            return
+
+        if len(results) != len(batch):
+            error = RuntimeError(f"Video batch returned {len(results)} results for {len(batch)} requests")
+            for job in batch:
+                if not job.future.done():
+                    job.future.set_exception(error)
+            return
+
+        for job, result in zip(batch, results, strict=True):
+            if not job.future.done():
+                job.future.set_result(result)
+
+    def _jobs_are_compatible(self, base: _VideoBatchJob, candidate: _VideoBatchJob) -> bool:
+        try:
+            base_sampling, base_extra = self._sampling_param_from_kwargs(base.kwargs)
+            candidate_sampling, candidate_extra = self._sampling_param_from_kwargs(candidate.kwargs)
+        except Exception:
+            return False
+        return can_dynamic_batch(
+            base_sampling,
+            candidate_sampling,
+            base_extra=base_extra,
+            candidate_extra=candidate_extra,
+        ).can_batch
+
+    def _sampling_param_from_kwargs(self, kwargs: dict[str, Any]) -> tuple[SamplingParam, dict[str, Any]]:
+        sampling_param = SamplingParam.from_pretrained(self._fastvideo_args.model_path)
+        updates = dict(kwargs)
+        prompt = updates.pop("prompt", None)
+        extra: dict[str, Any] = {}
+        for key in (
+                "ltx2_audio_latents",
+                "ltx2_audio_clean_latent",
+                "ltx2_audio_denoise_mask",
+                "audio_num_frames",
+                "video_position_offset_sec",
+        ):
+            if key in updates:
+                extra[key] = updates.pop(key)
+        sampling_param.update(updates)
+        sampling_param.prompt = prompt
+        return sampling_param, extra

fastvideo/entrypoints/openai/state.py

@@ -11,6 +11,7 @@ from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
     from fastvideo.api.schema import GenerationRequest
+    from fastvideo.entrypoints.openai.batching import VideoBatchScheduler
     from fastvideo.entrypoints.video_generator import VideoGenerator
     from fastvideo.fastvideo_args import FastVideoArgs
 
@@ -20,6 +21,7 @@ _generator: VideoGenerator | None = None
 _fastvideo_args: FastVideoArgs | None = None
 _output_dir: str = DEFAULT_OUTPUT_DIR
 _default_request: GenerationRequest | None = None
+_video_batch_scheduler: VideoBatchScheduler | None = None
 
 
 def get_generator() -> VideoGenerator:
@@ -44,23 +46,31 @@ def get_default_request() -> GenerationRequest | None:
     return _default_request
 
 
+def get_video_batch_scheduler() -> VideoBatchScheduler | None:
+    """Return the video batch scheduler when dynamic batching is enabled."""
+    return _video_batch_scheduler
+
+
 def set_state(
     generator: VideoGenerator,
     fastvideo_args: FastVideoArgs,
     output_dir: str,
     default_request: GenerationRequest | None = None,
+    video_batch_scheduler: VideoBatchScheduler | None = None,
 ) -> None:
     """Set all server state at once (called from lifespan)."""
-    global _generator, _fastvideo_args, _output_dir, _default_request
+    global _generator, _fastvideo_args, _output_dir, _default_request, _video_batch_scheduler
     _generator = generator
     _fastvideo_args = fastvideo_args
     _output_dir = output_dir
     _default_request = default_request
+    _video_batch_scheduler = video_batch_scheduler
 
 
 def clear_state() -> None:
     """Clear server state on shutdown."""
-    global _generator, _fastvideo_args, _default_request
+    global _generator, _fastvideo_args, _default_request, _video_batch_scheduler
     _generator = None
     _fastvideo_args = None
     _default_request = None
+    _video_batch_scheduler = None

fastvideo/entrypoints/openai/video_api.py

@@ -26,6 +26,7 @@ from fastvideo.entrypoints.openai.state import (
     get_generator,
     get_output_dir,
     get_server_args,
+    get_video_batch_scheduler,
 )
 from fastvideo.entrypoints.openai.protocol import (
     VideoGenerationsRequest,
@@ -151,15 +152,19 @@ async def _run_generation(request_id: str, kwargs: dict[str, Any]) -> None:
     is synchronous) and update the store on completion or failure.
     """
     generator = get_generator()
+    scheduler = get_video_batch_scheduler()
     loop = asyncio.get_running_loop()
 
     try:
         start = time.perf_counter()
 
-        result = await loop.run_in_executor(
-            None,
-            lambda: generator.generate_video(**kwargs),
-        )
+        if scheduler is not None and scheduler.enabled:
+            result = await scheduler.submit(request_id, kwargs)
+        else:
+            result = await loop.run_in_executor(
+                None,
+                lambda: generator.generate_video(**kwargs),
+            )
 
         elapsed = time.perf_counter() - start
         update: dict[str, Any] = {

fastvideo/entrypoints/video_generator.py

@@ -18,6 +18,7 @@ import warnings
 from collections.abc import Mapping
 from contextlib import suppress
 from copy import deepcopy
+from dataclasses import dataclass
 from typing import Any
 
 import imageio
@@ -26,6 +27,8 @@ import torch
 import torchvision
 from einops import rearrange
 
+from fastvideo.batching.admission import BatchAdmissionController
+from fastvideo.batching.signature import can_dynamic_batch
 from fastvideo.api.compat import (
     expand_request_prompt_batch,
     generator_config_to_fastvideo_args,
@@ -94,6 +97,11 @@ _FROM_PRETRAINED_CONVENIENCE_KWARGS = frozenset({
     "pin_cpu_memory",
     "enable_torch_compile",
     "torch_compile_kwargs",
+    "batching_mode",
+    "batching_max_size",
+    "batching_delay_ms",
+    "batching_config",
+    "enable_batching_metrics",
     "output_type",
     "nvfp4_fa4",
 })
@@ -112,6 +120,17 @@ def _infer_latent_batch_size(batch: ForwardBatch) -> int:
     return latent_batch_size
 
 
+@dataclass
+class _GenerationWorkItem:
+    prompt: str
+    sampling_param: SamplingParam
+    fastvideo_args: FastVideoArgs
+    batch: ForwardBatch
+    output_path: str
+    target_height: int
+    target_width: int
+
+
 class VideoGenerator:
     """
     A unified class for generating videos using diffusion models.
@@ -439,6 +458,70 @@ class VideoGenerator:
             if log_queue:
                 self.executor.clear_log_queue()
 
+    def generate_video_batch(self, request_kwargs: list[dict[str, Any]]) -> list[dict[str, Any]]:
+        """Generate multiple legacy video requests, batching compatible items."""
+        work_items: list[_GenerationWorkItem] = []
+        reserved_output_paths: set[str] = set()
+        fastvideo_args_by_pipeline_override: dict[tuple[tuple[str, str], ...], FastVideoArgs] = {
+            (): self.fastvideo_args
+        }
+        for raw_kwargs in request_kwargs:
+            kwargs = dict(raw_kwargs)
+            prompt = kwargs.pop("prompt", None)
+            if prompt is None:
+                raise ValueError("Each batched generation request must include prompt")
+            if not isinstance(prompt, str):
+                raise TypeError(f"`prompt` must be a string, but got {type(prompt)}")
+
+            sampling_param = kwargs.pop("sampling_param", None)
+            if sampling_param is None:
+                sampling_param = SamplingParam.from_pretrained(self.fastvideo_args.model_path)
+            else:
+                sampling_param = deepcopy(sampling_param)
+
+            extra_overrides: dict[str, Any] = {}
+            for _ek in _BATCH_EXTRA_PASSTHROUGH_KEYS:
+                if _ek in kwargs:
+                    extra_overrides[_ek] = kwargs.pop(_ek)
+
+            request = legacy_generate_call_to_request(
+                prompt,
+                sampling_param,
+                legacy_kwargs=kwargs,
+            )
+            if not isinstance(request.prompt, str):
+                raise TypeError(f"`prompt` must be a string, but got {type(request.prompt)}")
+
+            fastvideo_args = self.fastvideo_args
+            pipeline_overrides = request_to_pipeline_overrides(request)
+            if pipeline_overrides:
+                override_key = tuple((key, repr(value)) for key, value in sorted(pipeline_overrides.items()))
+                fastvideo_args = fastvideo_args_by_pipeline_override.get(override_key)
+                if fastvideo_args is None:
+                    fastvideo_args = deepcopy(self.fastvideo_args)
+                    for key, value in pipeline_overrides.items():
+                        if not hasattr(fastvideo_args.pipeline_config, key):
+                            raise ValueError(f"Request field {key!r} is not supported by pipeline config overrides")
+                        setattr(fastvideo_args.pipeline_config, key, deepcopy(value))
+                    fastvideo_args_by_pipeline_override[override_key] = fastvideo_args
+
+            resolved_sampling_param = request_to_sampling_param(
+                request,
+                model_path=self.fastvideo_args.model_path,
+            )
+            output_path = self._prepare_output_path(resolved_sampling_param.output_path, request.prompt,
+                                                    reserved_output_paths)
+            work_items.append(
+                self._prepare_generation_work_item(
+                    prompt=request.prompt,
+                    sampling_param=resolved_sampling_param,
+                    fastvideo_args=fastvideo_args,
+                    output_path=output_path,
+                    _extra_overrides=extra_overrides,
+                ))
+
+        return self._generate_prepared_work_items(work_items)
+
     def _generate_request_impl(
         self,
         request: GenerationRequest,
@@ -535,6 +618,29 @@ class VideoGenerator:
 
             logger.info("Found %d prompts in %s", len(prompts), prompt_txt_path)
 
+            if self._dynamic_batching_enabled(fastvideo_args):
+                work_items: list[_GenerationWorkItem] = []
+                reserved_output_paths: set[str] = set()
+                for batch_prompt in prompts:
+                    item_kwargs = dict(kwargs)
+                    item_kwargs["output_path"] = self._prepare_output_path(sampling_param.output_path, batch_prompt,
+                                                                           reserved_output_paths)
+                    work_items.append(
+                        self._prepare_generation_work_item(
+                            prompt=batch_prompt,
+                            sampling_param=sampling_param,
+                            fastvideo_args=fastvideo_args,
+                            **item_kwargs,
+                        ))
+
+                results = self._generate_prepared_work_items(work_items, tolerate_failures=True)
+                for i, (result, batch_prompt) in enumerate(zip(results, prompts, strict=True)):
+                    result["prompt_index"] = i
+                    result["prompt"] = batch_prompt
+                logger.info("Completed batch processing. Generated %d videos successfully.",
+                            sum(1 for result in results if "error" not in result))
+                return results
+
             results = []
             for i, batch_prompt in enumerate(prompts):
                 logger.info("Processing prompt %d/%d: %s...", i + 1, len(prompts), batch_prompt[:100])
@@ -588,6 +694,7 @@ class VideoGenerator:
         self,
         output_path: str,
         prompt: str,
+        reserved_paths: set[str] | None = None,
     ) -> str:
         """Build a unique, sanitized output file path.
 
@@ -602,6 +709,9 @@ class VideoGenerator:
         - Invalid filename characters are removed; if the name changes, a
           warning is logged.
         - If the target path already exists, a numeric suffix is appended.
+        - ``reserved_paths`` lets batch callers resolve every path before any
+          file is written: paths in the set are treated as taken, and the
+          chosen path is added to the set.
         """
         target_ext = ".png" if self._is_image_workload() else ".mp4"
 
@@ -646,15 +756,416 @@ class VideoGenerator:
         if output_dir:
             os.makedirs(output_dir, exist_ok=True)
 
+        def _is_taken(path: str) -> bool:
+            return os.path.exists(path) or (reserved_paths is not None and path in reserved_paths)
+
         new_output_path = os.path.join(output_dir, out_name)
         counter = 1
-        while os.path.exists(new_output_path):
+        while _is_taken(new_output_path):
             name_part, ext_part = os.path.splitext(out_name)
             new_name = f"{name_part}_{counter}{ext_part}"
             new_output_path = os.path.join(output_dir, new_name)
             counter += 1
+        if reserved_paths is not None:
+            reserved_paths.add(new_output_path)
         return new_output_path
 
+    def _dynamic_batching_enabled(self, fastvideo_args: FastVideoArgs) -> bool:
+        batching_mode = getattr(fastvideo_args, "batching_mode", "disabled")
+        batching_max_size = getattr(fastvideo_args, "batching_max_size", 1)
+        return batching_mode == "dynamic" and batching_max_size > 1
+
+    def _prepare_generation_work_item(
+        self,
+        prompt: str | list[str],
+        sampling_param: SamplingParam,
+        fastvideo_args: FastVideoArgs,
+        **kwargs,
+    ) -> _GenerationWorkItem:
+        if isinstance(prompt, str):
+            prompt_for_output = prompt.strip()
+            prompt_value: str | list[str] = prompt_for_output
+        elif isinstance(prompt, list) and all(isinstance(item, str) for item in prompt):
+            prompt_value = [item.strip() for item in prompt]
+            prompt_for_output = prompt_value[0] if prompt_value else ""
+        else:
+            raise TypeError(f"`prompt` must be a string or list of strings, but got {type(prompt)}")
+
+        sampling_param = deepcopy(sampling_param)
+        output_path = kwargs["output_path"]
+        sampling_param.prompt = prompt_value
+        if sampling_param.negative_prompt is not None:
+            sampling_param.negative_prompt = sampling_param.negative_prompt.strip()
+
+        if sampling_param.height <= 0 or sampling_param.width <= 0 or sampling_param.num_frames <= 0:
+            raise ValueError(f"Height, width, and num_frames must be positive integers, got "
+                             f"height={sampling_param.height}, width={sampling_param.width}, "
+                             f"num_frames={sampling_param.num_frames}")
+
+        target_height = align_to(sampling_param.height, 16)
+        target_width = align_to(sampling_param.width, 16)
+        latents_size = [(sampling_param.num_frames - 1) // 4 + 1, sampling_param.height // 8, sampling_param.width // 8]
+        n_tokens = latents_size[0] * latents_size[1] * latents_size[2]
+
+        debug_str = f"""
+                      height: {target_height}
+                       width: {target_width}
+                video_length: {sampling_param.num_frames}
+                      prompt: {sampling_param.prompt}
+                      image_path: {sampling_param.image_path}
+                  neg_prompt: {sampling_param.negative_prompt}
+                        seed: {sampling_param.seed}
+                 infer_steps: {sampling_param.num_inference_steps}
+       num_videos_per_prompt: {sampling_param.num_videos_per_prompt}
+              guidance_scale: {sampling_param.guidance_scale}
+                    n_tokens: {n_tokens}
+                  flow_shift: {fastvideo_args.pipeline_config.flow_shift}
+     embedded_guidance_scale: {fastvideo_args.pipeline_config.embedded_cfg_scale}
+                  save_video: {sampling_param.save_video}
+                  output_path: {output_path}
+        """ # type: ignore[attr-defined]
+        logger.info(debug_str)
+
+        batch = ForwardBatch(
+            **shallow_asdict(sampling_param),
+            eta=0.0,
+            n_tokens=n_tokens,
+            VSA_sparsity=fastvideo_args.VSA_sparsity,
+        )
+
+        extra_overrides = kwargs.get("_extra_overrides", {})
+        for _ek, _ev in extra_overrides.items():
+            batch.extra[_ek] = _ev
+
+        return _GenerationWorkItem(
+            prompt=prompt_for_output,
+            sampling_param=sampling_param,
+            fastvideo_args=fastvideo_args,
+            batch=batch,
+            output_path=output_path,
+            target_height=target_height,
+            target_width=target_width,
+        )
+
+    def _run_forward_batch(
+        self,
+        batch: ForwardBatch,
+        fastvideo_args: FastVideoArgs,
+    ) -> tuple[ForwardBatch, float, float]:
+        start_time = time.perf_counter()
+        result_container = {"output_batch": ForwardBatch(data_type=batch.data_type)}
+        thread_error: dict[str, BaseException | None] = {"error": None}
+        thread_error_traceback: dict[str, str] = {"traceback": ""}
+
+        def execute_forward_thread():
+            import traceback
+            try:
+                result_container["output_batch"] = self.executor.execute_forward(batch, fastvideo_args)
+            except BaseException as error:  # noqa: BLE001
+                thread_error["error"] = error
+                thread_error_traceback["traceback"] = traceback.format_exc()
+
+        thread = threading.Thread(target=execute_forward_thread)
+        thread.start()
+        thread.join()
+
+        if thread_error["error"] is not None:
+            raise RuntimeError("Forward execution thread failed.\n"
+                               f"{thread_error_traceback['traceback']}") from thread_error["error"]
+
+        output_batch = result_container["output_batch"]
+        if output_batch.output is None:
+            raise RuntimeError("Forward execution returned no output tensor. "
+                               "This usually means the executor/pipeline failed earlier.")
+
+        gen_time = time.perf_counter() - start_time
+        logger.info("Generated successfully in %.2f seconds", gen_time)
+        return output_batch, gen_time, start_time
+
+    def _samples_from_output(
+        self,
+        work_item: _GenerationWorkItem,
+        output_batch: ForwardBatch,
+    ) -> torch.Tensor:
+        output = output_batch.output
+        if output is None:
+            raise RuntimeError("Forward execution returned no output tensor.")
+        fastvideo_args = work_item.fastvideo_args
+        sampling_param = work_item.sampling_param
+        latent_batch_size = _infer_latent_batch_size(work_item.batch)
+        skip_pixel_prealloc = fastvideo_args.output_type == "latent"
+        expected_shape = (
+            latent_batch_size,
+            3,
+            sampling_param.num_frames,
+            sampling_param.height,
+            sampling_param.width,
+        )
+        if skip_pixel_prealloc:
+            return output.cpu()
+        samples = torch.empty(expected_shape, device="cpu", pin_memory=fastvideo_args.pin_cpu_memory)
+        if output.shape == samples.shape:
+            samples.copy_(output)
+            return samples
+        logger.warning("Output shape %s does not match expected shape %s; use slow path", output.shape, samples.shape)
+        return output.cpu()
+
+    def _postprocess_generation_output(
+        self,
+        work_item: _GenerationWorkItem,
+        output_batch: ForwardBatch,
+        gen_time: float,
+        start_time: float,
+    ) -> dict[str, Any]:
+        batch = work_item.batch
+        fastvideo_args = work_item.fastvideo_args
+        output_path = work_item.output_path
+        samples = self._samples_from_output(work_item, output_batch)
+        logging_info = output_batch.logging_info
+
+        is_latent_output = fastvideo_args.output_type == "latent"
+        audio_only = bool(output_batch.extra.get("audio_only"))
+
+        postprocess_start = time.perf_counter()
+        frames: list[np.ndarray] | None
+        if is_latent_output or audio_only:
+            frames = None if is_latent_output else []
+        else:
+            videos = rearrange(samples, "b c t h w -> t b c h w")
+            frames = []
+            for x in videos:
+                x = torchvision.utils.make_grid(x, nrow=6)
+                x = x.permute(1, 2, 0).squeeze(-1)
+                x = (x * 255).to(torch.uint8)
+                frames.append(x.contiguous().cpu().numpy())
+        postprocess_time = time.perf_counter() - postprocess_start
+        logger.info("PostDecodeFrameProcessStage completed in %.3f s", postprocess_time)
+        if logging_info is not None:
+            logging_info.add_stage_execution_time("PostDecodeFrameProcessStage", postprocess_time)
+
+        save_to_disk = batch.save_video and not is_latent_output
+        save_video_time = 0.0
+        audio_mux_time = 0.0
+        if save_to_disk:
+            if audio_only:
+                output_path = self._rewrite_extension(output_path, ".wav")
+                save_start = time.perf_counter()
+                self._write_pcm_wav(
+                    output_path,
+                    output_batch.extra["audio"],
+                    int(output_batch.extra["audio_sample_rate"]),
+                )
+                save_video_time = time.perf_counter() - save_start
+                logger.info("Saved audio to %s", output_path)
+            elif self._is_image_workload():
+                assert frames is not None
+                save_start = time.perf_counter()
+                imageio.imwrite(output_path, frames[0])
+                save_video_time = time.perf_counter() - save_start
+                logger.info("Saved image to %s", output_path)
+            else:
+                assert frames is not None
+                audio = output_batch.extra.get("audio")
+                audio_sample_rate = output_batch.extra.get("audio_sample_rate")
+                if audio is not None and audio_sample_rate is not None:
+                    save_start = time.perf_counter()
+                    save_ok = self._save_video_with_audio_ffmpeg_pipe(
+                        output_path=output_path,
+                        frames=frames,
+                        fps=batch.fps,
+                        audio=audio,
+                        sample_rate=int(audio_sample_rate),
+                    )
+                    if not save_ok:
+                        logger.warning("ffmpeg pipe save failed; trying PyAV single-pass save.")
+                        save_ok = self._save_video_with_audio_single_pass(
+                            output_path=output_path,
+                            frames=frames,
+                            fps=batch.fps,
+                            audio=audio,
+                            sample_rate=int(audio_sample_rate),
+                        )
+                    save_video_time = time.perf_counter() - save_start
+                    if save_ok:
+                        audio_mux_time = 0.0
+                    else:
+                        logger.warning("Single-pass save failed; falling back to two-step save/mux.")
+                        save_start = time.perf_counter()
+                        imageio.mimsave(output_path, frames, fps=batch.fps, format="mp4")
+                        save_video_time = time.perf_counter() - save_start
+                        mux_start = time.perf_counter()
+                        mux_ok = self._mux_audio(output_path, audio, int(audio_sample_rate))
+                        audio_mux_time = time.perf_counter() - mux_start
+                        if not mux_ok:
+                            logger.warning("Audio mux failed; saved video without audio.")
+                else:
+                    save_start = time.perf_counter()
+                    imageio.mimsave(output_path, frames, fps=batch.fps, format="mp4")
+                    save_video_time = time.perf_counter() - save_start
+                    audio_mux_time = 0.0
+                logger.info("Saved video to %s", output_path)
+
+            logger.info("VideoSaveStage completed in %.3f s", save_video_time)
+            if logging_info is not None:
+                logging_info.add_stage_execution_time("VideoSaveStage", save_video_time)
+            logger.info("AudioMuxStage completed in %.3f s", audio_mux_time)
+            if logging_info is not None:
+                logging_info.add_stage_execution_time("AudioMuxStage", audio_mux_time)
+
+        e2e_time = time.perf_counter() - start_time
+        logger.info("End-to-end latency: %.2f seconds", e2e_time)
+
+        return {
+            "prompts": work_item.prompt,
+            "samples": samples if batch.return_frames else None,
+            "frames": frames if batch.return_frames else None,
+            "audio": output_batch.extra.get("audio"),
+            "audio_sample_rate": output_batch.extra.get("audio_sample_rate"),
+            "ltx2_audio_latents": output_batch.extra.get("ltx2_audio_latents"),
+            "size": (work_item.target_height, work_item.target_width, batch.num_frames),
+            "generation_time": gen_time,
+            "e2e_latency": e2e_time,
+            "logging_info": logging_info,
+            "trajectory": output_batch.trajectory_latents,
+            "trajectory_timesteps": output_batch.trajectory_timesteps,
+            "trajectory_decoded": output_batch.trajectory_decoded,
+            "video_path": output_path if save_to_disk else None,
+            "peak_memory_mb": output_batch.extra.get("peak_memory_mb"),
+        }
+
+    def _split_output_batch(
+        self,
+        output_batch: ForwardBatch,
+        *,
+        index: int,
+        batch_size: int,
+    ) -> ForwardBatch:
+        extra = {}
+        for key, value in (output_batch.extra or {}).items():
+            if torch.is_tensor(value) and value.ndim > 0 and value.shape[0] == batch_size:
+                extra[key] = value[index:index + 1]
+            elif isinstance(value, list) and len(value) == batch_size:
+                extra[key] = value[index]
+            else:
+                extra[key] = value
+
+        result = ForwardBatch(
+            data_type=output_batch.data_type,
+            output=(output_batch.output[index:index + 1] if output_batch.output is not None else None),
+            logging_info=output_batch.logging_info,
+            extra=extra,
+        )
+        if output_batch.trajectory_latents is not None:
+            result.trajectory_latents = output_batch.trajectory_latents[index:index + 1]
+        result.trajectory_timesteps = output_batch.trajectory_timesteps
+        if output_batch.trajectory_decoded is not None:
+            result.trajectory_decoded = [
+                decoded[index:index + 1] if torch.is_tensor(decoded) and decoded.shape[0] == batch_size else decoded
+                for decoded in output_batch.trajectory_decoded
+            ]
+        return result
+
+    def _merge_work_items(self, work_items: list[_GenerationWorkItem]) -> _GenerationWorkItem:
+        first = work_items[0]
+        sampling_param = deepcopy(first.sampling_param)
+        prompts = [item.prompt for item in work_items]
+        sampling_param.prompt = prompts
+        sampling_param.seed = work_items[0].sampling_param.seed
+
+        merged = self._prepare_generation_work_item(
+            prompts,
+            sampling_param,
+            first.fastvideo_args,
+            output_path=first.output_path,
+            _extra_overrides=first.batch.extra,
+        )
+        merged.batch.seeds = [int(item.sampling_param.seed) for item in work_items]
+        merged.batch.extra["dynamic_batch_size"] = len(work_items)
+        merged.batch.extra["dynamic_batch_output_paths"] = [item.output_path for item in work_items]
+        return merged
+
+    def _can_merge_work_items(
+        self,
+        base: _GenerationWorkItem,
+        candidate: _GenerationWorkItem,
+        admission: BatchAdmissionController,
+        current_group: list[_GenerationWorkItem],
+    ) -> bool:
+        if candidate.fastvideo_args is not base.fastvideo_args:
+            return False
+        compatibility = can_dynamic_batch(
+            base.sampling_param,
+            candidate.sampling_param,
+            base_extra=base.batch.extra,
+            candidate_extra=candidate.batch.extra,
+        )
+        if not compatibility.can_batch:
+            return False
+        current_requests = [item.sampling_param for item in current_group]
+        return admission.reject_reason_for_candidate(current_requests, candidate.sampling_param) is None
+
+    def _run_work_item_group(self, group: list[_GenerationWorkItem]) -> list[dict[str, Any]]:
+        if len(group) == 1:
+            return [self._execute_single_work_item(group[0])]
+        merged = self._merge_work_items(group)
+        output_batch, gen_time, start_time = self._run_forward_batch(merged.batch, merged.fastvideo_args)
+        return [
+            self._postprocess_generation_output(
+                item,
+                self._split_output_batch(output_batch, index=item_index, batch_size=len(group)),
+                gen_time,
+                start_time,
+            ) for item_index, item in enumerate(group)
+        ]
+
+    def _generate_prepared_work_items(
+        self,
+        work_items: list[_GenerationWorkItem],
+        tolerate_failures: bool = False,
+    ) -> list[dict[str, Any]]:
+        """Execute prepared work items, batching compatible neighbors.
+
+        With ``tolerate_failures`` (prompt-file semantics), a failed group
+        yields one ``{"error": ..., "prompt": ...}`` entry per work item so
+        completed results survive and stay aligned with the inputs; otherwise
+        the exception propagates.
+        """
+        if not work_items:
+            return []
+
+        def run_group(group: list[_GenerationWorkItem]) -> list[dict[str, Any]]:
+            if not tolerate_failures:
+                return self._run_work_item_group(group)
+            try:
+                return self._run_work_item_group(group)
+            except Exception as e:
+                logger.error("Failed to generate videos for batched prompts %s: %s",
+                             [item.prompt[:100] for item in group], e)
+                return [{"error": str(e), "prompt": item.prompt} for item in group]
+
+        fastvideo_args = work_items[0].fastvideo_args
+        if not self._dynamic_batching_enabled(fastvideo_args):
+            return [result for item in work_items for result in run_group([item])]
+
+        admission = BatchAdmissionController(fastvideo_args)
+        results: list[dict[str, Any]] = []
+        index = 0
+        while index < len(work_items):
+            group = [work_items[index]]
+            index += 1
+            while index < len(work_items) and len(group) < fastvideo_args.batching_max_size:
+                candidate = work_items[index]
+                if not self._can_merge_work_items(group[0], candidate, admission, group):
+                    break
+                group.append(candidate)
+                index += 1
+            results.extend(run_group(group))
+        return results
+
+    def _execute_single_work_item(self, work_item: _GenerationWorkItem) -> dict[str, Any]:
+        output_batch, gen_time, start_time = self._run_forward_batch(work_item.batch, work_item.fastvideo_args)
+        return self._postprocess_generation_output(work_item, output_batch, gen_time, start_time)
+
     def _generate_single_video(
         self,
         prompt: str,

fastvideo/fastvideo_args.py

@@ -169,6 +169,14 @@ class FastVideoArgs:
     # Prompt text file for batch processing
     prompt_txt: str | None = None
 
+    # Dynamic multimodal generation batching. Defaults preserve the historical
+    # one-request-at-a-time execution path.
+    batching_mode: str = "disabled"
+    batching_max_size: int = 1
+    batching_delay_ms: float = 0.0
+    batching_config: str | None = None
+    enable_batching_metrics: bool = False
+
     # LTX-2 VAE tiling overrides
     ltx2_vae_tiling: bool | None = None
     ltx2_vae_spatial_tile_size_in_pixels: int | None = None
@@ -446,6 +454,37 @@ class FastVideoArgs:
             default=FastVideoArgs.prompt_txt,
             help="Path to a text file containing prompts (one per line) for batch processing",
         )
+        parser.add_argument(
+            "--batching-mode",
+            type=str,
+            choices=["disabled", "dynamic"],
+            default=FastVideoArgs.batching_mode,
+            help="Request batching mode for inference serving.",
+        )
+        parser.add_argument(
+            "--batching-max-size",
+            type=int,
+            default=FastVideoArgs.batching_max_size,
+            help="Maximum number of compatible generation requests to execute as one batch.",
+        )
+        parser.add_argument(
+            "--batching-delay-ms",
+            type=float,
+            default=FastVideoArgs.batching_delay_ms,
+            help="Maximum queue delay in milliseconds before dispatching a dynamic batch.",
+        )
+        parser.add_argument(
+            "--batching-config",
+            type=str,
+            default=FastVideoArgs.batching_config,
+            help="Optional JSON batching admission rule file.",
+        )
+        parser.add_argument(
+            "--enable-batching-metrics",
+            action=StoreBoolean,
+            default=FastVideoArgs.enable_batching_metrics,
+            help="Log dynamic batching utilization and rejection metrics.",
+        )
 
         # LTX-2 VAE tiling overrides
         parser.add_argument(
@@ -753,6 +792,13 @@ class FastVideoArgs:
                           WorkloadType), f"Workload type must be a WorkloadType enum, got {type(self.workload_type)}"
         assert self.workload_type in WorkloadType.choices(), f"Invalid workload type: {self.workload_type}"
 
+        if self.batching_mode not in {"disabled", "dynamic"}:
+            raise ValueError(f"batching_mode must be 'disabled' or 'dynamic', got {self.batching_mode!r}")
+        if self.batching_max_size < 1:
+            raise ValueError("batching_max_size must be >= 1")
+        if self.batching_delay_ms < 0:
+            raise ValueError("batching_delay_ms must be >= 0")
+
         if self.mode in [ExecutionMode.DISTILLATION, ExecutionMode.FINETUNING] and self.inference_mode:
             logger.warning("Mode is 'training' but inference_mode is True. Setting inference_mode to False.")
             self.inference_mode = False

fastvideo/models/dits/causal_wanvideo.py

@@ -437,7 +437,6 @@ class CausalWanTransformer3DModel(BaseDiT):
 
         # Causal-specific
         self.block_mask = None
-        self.teacher_forcing_block_mask = None
         self.num_frame_per_block = config.arch_config.num_frames_per_block
         assert self.num_frame_per_block <= 3
         self.independent_first_frame = False
@@ -501,70 +500,6 @@ class CausalWanTransformer3DModel(BaseDiT):
 
         return block_mask
 
-    @staticmethod
-    def _prepare_teacher_forcing_mask(
-        device: torch.device | str, num_frames: int = 21,
-        frame_seqlen: int = 1560, num_frame_per_block=1, local_attn_size=-1
-    ) -> BlockMask:
-        """Attention mask for the teacher-forcing ``[clean | noisy]`` sequence.
-
-        A noisy token attends to its own block plus the clean context of all
-        strictly previous blocks; clean tokens are block-wise causal.
-        """
-        if local_attn_size != -1:
-            raise NotImplementedError(
-                f"Teacher forcing ignores local_attn_size={local_attn_size}: "
-                "unlike the block-wise causal mask, this mask always attends "
-                "to the full clean context. Windowed teacher forcing is not "
-                "implemented; use local_attn_size=-1 for teacher-forcing "
-                "training.")
-        total_length = num_frames * frame_seqlen * 2
-        padded_length = math.ceil(total_length / 128) * 128 - total_length
-
-        clean_ends = num_frames * frame_seqlen
-        context_ends = torch.zeros(total_length + padded_length, device=device, dtype=torch.long)
-        noise_context_starts = torch.zeros(total_length + padded_length, device=device, dtype=torch.long)
-        noise_context_ends = torch.zeros(total_length + padded_length, device=device, dtype=torch.long)
-        noise_noise_starts = torch.zeros(total_length + padded_length, device=device, dtype=torch.long)
-        noise_noise_ends = torch.zeros(total_length + padded_length, device=device, dtype=torch.long)
-
-        attention_block_size = frame_seqlen * num_frame_per_block
-        frame_indices = torch.arange(
-            start=0, end=num_frames * frame_seqlen,
-            step=attention_block_size, device=device, dtype=torch.long
-        )
-        for start in frame_indices:
-            context_ends[start:start + attention_block_size] = start + attention_block_size
-
-        noisy_image_start_list = torch.arange(
-            num_frames * frame_seqlen, total_length,
-            step=attention_block_size, device=device, dtype=torch.long
-        )
-        noisy_image_end_list = noisy_image_start_list + attention_block_size
-        for block_index, (start, end) in enumerate(zip(noisy_image_start_list, noisy_image_end_list)):
-            noise_noise_starts[start:end] = start
-            noise_noise_ends[start:end] = end
-            noise_context_ends[start:end] = block_index * attention_block_size
-
-        def attention_mask(b, h, q_idx, kv_idx):
-            clean_mask = (q_idx < clean_ends) & (kv_idx < context_ends[q_idx])
-            c1 = (kv_idx < noise_noise_ends[q_idx]) & (kv_idx >= noise_noise_starts[q_idx])
-            c2 = (kv_idx < noise_context_ends[q_idx]) & (kv_idx >= noise_context_starts[q_idx])
-            noise_mask = (q_idx >= clean_ends) & (c1 | c2)
-            eye_mask = q_idx == kv_idx
-            return eye_mask | clean_mask | noise_mask
-
-        block_mask = create_block_mask(
-            attention_mask, B=None, H=None,
-            Q_LEN=total_length + padded_length, KV_LEN=total_length + padded_length,
-            _compile=False, device=device)
-
-        if not dist.is_initialized() or dist.get_rank() == 0:
-            print(f" cache a teacher-forcing mask with block size of {num_frame_per_block} frames")
-            print(block_mask)
-
-        return block_mask
-
     def _forward_inference(
                 self,
                 hidden_states: torch.Tensor,
@@ -693,12 +628,9 @@ class CausalWanTransformer3DModel(BaseDiT):
                 encoder_hidden_states_image: torch.Tensor | list[torch.Tensor]
                 | None = None,
                 start_frame: int = 0,
-                clean_x: torch.Tensor | None = None,
-                aug_t: torch.Tensor | None = None,
                 **kwargs) -> torch.Tensor:
 
         orig_dtype = hidden_states.dtype
-        teacher_forcing = clean_x is not None
         if not isinstance(encoder_hidden_states, torch.Tensor):
             encoder_hidden_states = encoder_hidden_states[0]
         if isinstance(encoder_hidden_states_image,
@@ -731,26 +663,15 @@ class CausalWanTransformer3DModel(BaseDiT):
         freqs_cis = (freqs_cos,
                      freqs_sin) if freqs_cos is not None else None
 
-        if teacher_forcing:
-            if self.teacher_forcing_block_mask is None:
-                self.teacher_forcing_block_mask = self._prepare_teacher_forcing_mask(
-                    device=hidden_states.device,
-                    num_frames=num_frames,
-                    frame_seqlen=post_patch_height * post_patch_width,
-                    num_frame_per_block=self.num_frame_per_block,
-                    local_attn_size=self.local_attn_size,
-                )
-            block_mask = self.teacher_forcing_block_mask
-        else:
-            if self.block_mask is None:
-                self.block_mask = self._prepare_blockwise_causal_attn_mask(
-                    device=hidden_states.device,
-                    num_frames=num_frames,
-                    frame_seqlen=post_patch_height * post_patch_width,
-                    num_frame_per_block=self.num_frame_per_block,
-                    local_attn_size=self.local_attn_size
-                )
-            block_mask = self.block_mask
+        # Construct blockwise causal attn mask
+        if self.block_mask is None:
+            self.block_mask = self._prepare_blockwise_causal_attn_mask(
+                device=hidden_states.device,
+                num_frames=num_frames,
+                frame_seqlen=post_patch_height * post_patch_width,
+                num_frame_per_block=self.num_frame_per_block,
+                local_attn_size=self.local_attn_size
+            )
 
         hidden_states = self.patch_embedding(hidden_states)
         grid_sizes = torch.stack(
@@ -758,7 +679,6 @@ class CausalWanTransformer3DModel(BaseDiT):
         hidden_states = hidden_states.flatten(2).transpose(1, 2)
 
         encoder_hidden_states = torch.cat([encoder_hidden_states, encoder_hidden_states.new_zeros(1, self.text_len - encoder_hidden_states.size(1), encoder_hidden_states.size(2))], dim=1)
-        encoder_hidden_states_text = encoder_hidden_states
 
         temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = self.condition_embedder(
                         timestep.flatten(), encoder_hidden_states, encoder_hidden_states_image)
@@ -774,35 +694,18 @@ class CausalWanTransformer3DModel(BaseDiT):
 
         assert encoder_hidden_states.dtype == orig_dtype
 
-        if teacher_forcing:
-            # Tile RoPE/modulation so clean frame i and noisy frame i share a position.
-            clean_tokens = self.patch_embedding(clean_x).flatten(2).transpose(1, 2)
-            hidden_states = torch.cat([clean_tokens, hidden_states], dim=1)
-            if aug_t is None:
-                aug_t = torch.zeros_like(timestep)
-            _, timestep_proj_clean, _, _ = self.condition_embedder(
-                aug_t.flatten(), encoder_hidden_states_text, None)
-            timestep_proj_clean = timestep_proj_clean.unflatten(
-                1, (6, self.hidden_size)).unflatten(dim=0, sizes=timestep.shape)
-            timestep_proj = torch.cat([timestep_proj_clean, timestep_proj], dim=1)
-            freqs_cis = (torch.cat([freqs_cos, freqs_cos], dim=0),
-                         torch.cat([freqs_sin, freqs_sin], dim=0))
-
         # 4. Transformer blocks
         if torch.is_grad_enabled() and self.gradient_checkpointing:
             for block in self.blocks:
                 hidden_states = self._gradient_checkpointing_func(
                     block, hidden_states, encoder_hidden_states,
                     timestep_proj, freqs_cis,
-                    block_mask=block_mask)
+                    block_mask=self.block_mask)
         else:
             for block in self.blocks:
                 hidden_states = block(hidden_states, encoder_hidden_states,
                                         timestep_proj, freqs_cis,
-                                        block_mask=block_mask)
-
-        if teacher_forcing:
-            hidden_states = hidden_states[:, hidden_states.shape[1] // 2:]
+                                        block_mask=self.block_mask)
 
         # 5. Output norm, projection & unpatchify
         temb = temb.unflatten(dim=0, sizes=timestep.shape).unsqueeze(2)

fastvideo/models/dits/dreamx_world.py

@@ -1,511 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-
-import math
-from typing import Any
-
-import torch
-import torch.nn as nn
-
-from fastvideo.configs.models.dits.dreamx_world import DreamXWorldConfig
-from fastvideo.distributed.communication_op import (
-    sequence_model_parallel_all_gather_with_unpad,
-    sequence_model_parallel_shard)
-from fastvideo.distributed.parallel_state import get_sp_world_size
-from fastvideo.layers.layernorm import RMSNorm
-from fastvideo.layers.linear import ReplicatedLinear
-from fastvideo.layers.rotary_embedding import get_rotary_pos_embed
-from fastvideo.models.dits.wanvideo import (LayerNormScaleShift,
-                                            PatchEmbed,
-                                            WanTimeTextImageEmbedding,
-                                            WanTransformer3DModel,
-                                            WanTransformerBlock)
-from fastvideo.platforms import AttentionBackendEnum, current_platform
-from fastvideo.attention import LocalAttention
-from fastvideo.distributed.parallel_state import get_sp_world_size
-from fastvideo.layers.quantization import QuantizationConfig
-from fastvideo.models.dits.base import BaseDiT
-
-
-def _dreamx_invert_se3(transforms: torch.Tensor) -> torch.Tensor:
-    assert transforms.shape[-2:] == (4, 4)
-    rot_inv = transforms[..., :3, :3].transpose(-1, -2)
-    out = torch.zeros_like(transforms)
-    out[..., :3, :3] = rot_inv
-    out[..., :3, 3] = -torch.einsum("...ij,...j->...i", rot_inv,
-                                    transforms[..., :3, 3])
-    out[..., 3, 3] = 1.0
-    return out.to(dtype=transforms.dtype)
-
-
-def _dreamx_lift_k(intrinsics: torch.Tensor) -> torch.Tensor:
-    assert intrinsics.shape[-2:] == (3, 3)
-    out = torch.zeros(intrinsics.shape[:-2] + (4, 4),
-                      device=intrinsics.device,
-                      dtype=intrinsics.dtype)
-    out[..., :3, :3] = intrinsics
-    out[..., 3, 3] = 1.0
-    return out
-
-
-def _dreamx_invert_k(intrinsics: torch.Tensor) -> torch.Tensor:
-    assert intrinsics.shape[-2:] == (3, 3)
-    out = torch.zeros_like(intrinsics)
-    out[..., 0, 0] = 1.0 / intrinsics[..., 0, 0]
-    out[..., 1, 1] = 1.0 / intrinsics[..., 1, 1]
-    out[..., 0, 2] = -intrinsics[..., 0, 2] / intrinsics[..., 0, 0]
-    out[..., 1, 2] = -intrinsics[..., 1, 2] / intrinsics[..., 1, 1]
-    out[..., 2, 2] = 1.0
-    return out.to(dtype=intrinsics.dtype)
-
-
-def _dreamx_apply_tiled_projmat(feats: torch.Tensor,
-                                matrix: torch.Tensor) -> torch.Tensor:
-    batch, num_heads, seq_len, feat_dim = feats.shape
-    proj_dim = matrix.shape[-1]
-    assert feat_dim % proj_dim == 0
-
-    if matrix.shape[1] == seq_len:
-        feats = feats.view(batch, num_heads, seq_len, feat_dim // proj_dim,
-                           proj_dim)
-        out = torch.einsum("btij,bntpj->bntpi", matrix, feats)
-        return out.reshape(batch, num_heads, seq_len, feat_dim)
-
-    cameras = matrix.shape[1]
-    assert seq_len > cameras and seq_len % cameras == 0
-    feats = feats.reshape(batch, num_heads, cameras, -1,
-                          feat_dim // proj_dim, proj_dim)
-    out = torch.einsum("bcij,bncpkj->bncpki", matrix, feats)
-    return out.reshape(batch, num_heads, seq_len, feat_dim)
-
-
-def _dreamx_prope_qkv(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
-                      viewmats: torch.Tensor, intrinsics: torch.Tensor):
-    batch, num_heads, seq_len, head_dim = q.shape
-    cameras = viewmats.shape[1]
-    assert q.shape == k.shape == v.shape
-    assert viewmats.shape == (batch, cameras, 4, 4)
-    assert intrinsics.shape == (batch, cameras, 3, 3)
-    assert head_dim % 4 == 0
-
-    intrinsics_norm = torch.zeros_like(intrinsics)
-    intrinsics_norm[..., 0, 0] = intrinsics[..., 0, 0]
-    intrinsics_norm[..., 1, 1] = intrinsics[..., 1, 1]
-    intrinsics_norm[..., 2, 2] = 1.0
-
-    proj = torch.einsum("...ij,...jk->...ik",
-                        _dreamx_lift_k(intrinsics_norm), viewmats)
-    proj_t = proj.transpose(-1, -2).to(dtype=viewmats.dtype)
-    proj_inv = torch.einsum(
-        "...ij,...jk->...ik",
-        _dreamx_invert_se3(viewmats),
-        _dreamx_lift_k(_dreamx_invert_k(intrinsics_norm)),
-    ).to(dtype=viewmats.dtype)
-
-    q = _dreamx_apply_tiled_projmat(q, proj_t)
-    k = _dreamx_apply_tiled_projmat(k, proj_inv)
-    v = _dreamx_apply_tiled_projmat(v, proj_inv)
-    return q, k, v, proj
-
-
-class DreamXPropeSelfAttention(nn.Module):
-    """DreamX-World parallel PRoPE camera self-attention branch."""
-
-    def __init__(self,
-                 dim: int,
-                 attn_dim: int,
-                 num_heads: int,
-                 qk_norm: str | bool = True,
-                 eps: float = 1e-6,
-                 quant_config: QuantizationConfig | None = None,
-                 prefix: str = ""):
-        super().__init__()
-        assert attn_dim % num_heads == 0
-        self.attn_dim = attn_dim
-        self.num_heads = num_heads
-        self.head_dim = attn_dim // num_heads
-        self.qk_norm = qk_norm
-
-        self.q_proj = ReplicatedLinear(dim,
-                                       attn_dim,
-                                       quant_config=quant_config,
-                                       prefix=f"{prefix}.q_proj")
-        self.k_proj = ReplicatedLinear(dim,
-                                       attn_dim,
-                                       quant_config=quant_config,
-                                       prefix=f"{prefix}.k_proj")
-        self.v_proj = ReplicatedLinear(dim,
-                                       attn_dim,
-                                       quant_config=quant_config,
-                                       prefix=f"{prefix}.v_proj")
-        self.out_proj = ReplicatedLinear(attn_dim,
-                                         dim,
-                                         quant_config=quant_config,
-                                         prefix=f"{prefix}.out_proj")
-
-        if qk_norm == "rms_norm":
-            self.norm_q = RMSNorm(self.head_dim, eps=eps)
-            self.norm_k = RMSNorm(self.head_dim, eps=eps)
-        elif qk_norm in (True, "rms_norm_across_heads"):
-            self.norm_q = RMSNorm(attn_dim, eps=eps)
-            self.norm_k = RMSNorm(attn_dim, eps=eps)
-        elif qk_norm is False:
-            self.norm_q = nn.Identity()
-            self.norm_k = nn.Identity()
-        else:
-            raise ValueError(f"Unsupported qk_norm for DreamX PRoPE: {qk_norm}")
-
-        nn.init.zeros_(self.out_proj.weight)
-        if self.out_proj.bias is not None:
-            nn.init.zeros_(self.out_proj.bias)
-
-        self.attn = LocalAttention(
-            num_heads=num_heads,
-            head_size=self.head_dim,
-            dropout_rate=0,
-            softmax_scale=None,
-            causal=False,
-            supported_attention_backends=(AttentionBackendEnum.FLASH_ATTN,
-                                          AttentionBackendEnum.TORCH_SDPA))
-
-    def forward(self, hidden_states: torch.Tensor,
-                y_camera: dict[str, torch.Tensor]) -> torch.Tensor:
-        if get_sp_world_size() > 1:
-            # The transformer shards the sequence before the block loop and
-            # this branch uses LocalAttention (no all-to-all): under
-            # sequence parallelism each rank would attend only within its
-            # own shard — silently wrong output. Fail loudly until this
-            # path is ported to DistributedAttention and validated.
-            raise NotImplementedError(
-                "DreamXPropeSelfAttention does not support sequence "
-                "parallelism yet (LocalAttention on a sharded sequence "
-                "corrupts output). Run with sp_size=1.")
-        batch_size, seq_len, _ = hidden_states.shape
-
-        query, _ = self.q_proj(hidden_states)
-        key, _ = self.k_proj(hidden_states)
-        value, _ = self.v_proj(hidden_states)
-
-        if self.qk_norm == "rms_norm":
-            query = query.view(batch_size, seq_len, self.num_heads,
-                               self.head_dim)
-            key = key.view(batch_size, seq_len, self.num_heads, self.head_dim)
-            query = self.norm_q(query)
-            key = self.norm_k(key)
-        else:
-            query = self.norm_q(query).view(batch_size, seq_len,
-                                            self.num_heads, self.head_dim)
-            key = self.norm_k(key).view(batch_size, seq_len, self.num_heads,
-                                        self.head_dim)
-
-        value = value.view(batch_size, seq_len, self.num_heads, self.head_dim)
-        query = query.transpose(1, 2)
-        key = key.transpose(1, 2)
-        value = value.transpose(1, 2)
-
-        query, key, value, output_projection = _dreamx_prope_qkv(
-            query,
-            key,
-            value,
-            viewmats=y_camera["viewmats"],
-            intrinsics=y_camera["K"],
-        )
-
-        out = self.attn(query.transpose(1, 2), key.transpose(1, 2),
-                        value.transpose(1, 2))
-        out = _dreamx_apply_tiled_projmat(out.transpose(1, 2),
-                                          output_projection).transpose(1, 2)
-        out = out.flatten(2)
-        out, _ = self.out_proj(out)
-        return out
-
-
-class DreamXWorldTransformerBlock(WanTransformerBlock):
-
-    def __init__(self,
-                 dim: int,
-                 ffn_dim: int,
-                 num_heads: int,
-                 qk_norm: str = "rms_norm_across_heads",
-                 cross_attn_norm: bool = False,
-                 eps: float = 1e-6,
-                 added_kv_proj_dim: int | None = None,
-                 supported_attention_backends: tuple[AttentionBackendEnum, ...]
-                 | None = None,
-                 quant_config: QuantizationConfig | None = None,
-                 prefix: str = "",
-                 add_control_adapter: bool = True,
-                 cam_method: str | None = "prope",
-                 attn_compress: int = 1,
-                 cam_self_attn_layers: tuple[int, ...] | None = None,
-                 layer_idx: int | None = None):
-        super().__init__(dim, ffn_dim, num_heads, qk_norm, cross_attn_norm,
-                         eps, added_kv_proj_dim,
-                         supported_attention_backends, quant_config, prefix)
-        self.cam_self_attn = None
-        add_cam_attn = add_control_adapter and cam_method == "prope"
-        if add_cam_attn and cam_self_attn_layers is not None:
-            add_cam_attn = layer_idx in cam_self_attn_layers
-        if add_cam_attn:
-            if num_heads % attn_compress != 0 or dim % attn_compress != 0:
-                raise ValueError("DreamX attn_compress must divide dim and num_heads")
-            self.cam_self_attn = DreamXPropeSelfAttention(
-                dim,
-                dim // attn_compress,
-                num_heads // attn_compress,
-                qk_norm=qk_norm,
-                eps=eps,
-                quant_config=quant_config,
-                prefix=f"{prefix}.cam_self_attn")
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        encoder_hidden_states: torch.Tensor,
-        temb: torch.Tensor,
-        freqs_cis: tuple[torch.Tensor, torch.Tensor],
-        original_seq_len: int,
-        y_camera: dict[str, torch.Tensor] | None = None,
-    ) -> torch.Tensor:
-        if hidden_states.dim() == 4:
-            hidden_states = hidden_states.squeeze(1)
-        orig_dtype = hidden_states.dtype
-
-        if temb.dim() == 4:
-            shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = (
-                self.scale_shift_table.unsqueeze(0) + temb.float()).chunk(
-                    6, dim=2)
-            shift_msa = shift_msa.squeeze(2)
-            scale_msa = scale_msa.squeeze(2)
-            gate_msa = gate_msa.squeeze(2)
-            c_shift_msa = c_shift_msa.squeeze(2)
-            c_scale_msa = c_scale_msa.squeeze(2)
-            c_gate_msa = c_gate_msa.squeeze(2)
-        else:
-            e = self.scale_shift_table + temb.float()
-            shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = e.chunk(
-                6, dim=1)
-        assert shift_msa.dtype == torch.float32
-
-        norm_hidden_states = (self.norm1(hidden_states.float()) *
-                              (1 + scale_msa) + shift_msa).to(orig_dtype)
-        query, _ = self.to_q(norm_hidden_states)
-        key, _ = self.to_k(norm_hidden_states)
-        value, _ = self.to_v(norm_hidden_states)
-
-        if self.norm_q is not None:
-            query = self.norm_q(query)
-        if self.norm_k is not None:
-            key = self.norm_k(key)
-
-        query = query.squeeze(1).unflatten(2, (self.num_attention_heads, -1))
-        key = key.squeeze(1).unflatten(2, (self.num_attention_heads, -1))
-        value = value.squeeze(1).unflatten(2, (self.num_attention_heads, -1))
-
-        attn_output, _ = self.attn1(
-            query,
-            key,
-            value,
-            original_seq_len,
-            freqs_cis=freqs_cis,
-        )
-        attn_output = attn_output.flatten(2)
-        attn_output, _ = self.to_out(attn_output)
-        attn_output = attn_output.squeeze(1)
-        if self.cam_self_attn is not None and y_camera is not None:
-            attn_output = attn_output + self.cam_self_attn(
-                norm_hidden_states, y_camera)
-
-        null_shift = null_scale = torch.tensor([0], device=hidden_states.device)
-        norm_hidden_states, hidden_states = self.self_attn_residual_norm(
-            hidden_states, attn_output, gate_msa, null_shift, null_scale)
-        norm_hidden_states, hidden_states = norm_hidden_states.to(
-            orig_dtype), hidden_states.to(orig_dtype)
-
-        attn_output = self.attn2(norm_hidden_states,
-                                 context=encoder_hidden_states,
-                                 context_lens=None)
-        norm_hidden_states, hidden_states = self.cross_attn_residual_norm(
-            hidden_states, attn_output, 1, c_shift_msa, c_scale_msa)
-        norm_hidden_states, hidden_states = norm_hidden_states.to(
-            orig_dtype), hidden_states.to(orig_dtype)
-
-        ff_output = self.ffn(norm_hidden_states)
-        hidden_states = self.mlp_residual(hidden_states, ff_output, c_gate_msa)
-        hidden_states = hidden_states.to(orig_dtype)
-
-        return hidden_states
-
-
-class DreamXWorldTransformer3DModel(WanTransformer3DModel):
-    _fsdp_shard_conditions = DreamXWorldConfig()._fsdp_shard_conditions
-    _compile_conditions = DreamXWorldConfig()._compile_conditions
-    _supported_attention_backends = DreamXWorldConfig(
-    )._supported_attention_backends
-    param_names_mapping = DreamXWorldConfig().param_names_mapping
-    reverse_param_names_mapping = DreamXWorldConfig().reverse_param_names_mapping
-    lora_param_names_mapping = DreamXWorldConfig().lora_param_names_mapping
-
-    def __init__(self, config: DreamXWorldConfig, hf_config: dict[str,
-                                                                  Any]) -> None:
-        BaseDiT.__init__(self, config=config, hf_config=hf_config)
-        self.quant_config = config.quant_config
-
-        inner_dim = config.num_attention_heads * config.attention_head_dim
-        self.hidden_size = config.hidden_size
-        self.num_attention_heads = config.num_attention_heads
-        self.in_channels = config.in_channels
-        self.out_channels = config.out_channels
-        self.num_channels_latents = config.num_channels_latents
-        self.patch_size = config.patch_size
-        self.text_len = config.text_len
-
-        assert config.num_attention_heads % get_sp_world_size() == 0, f"The number of attention heads ({config.num_attention_heads}) must be divisible by the sequence parallel size ({get_sp_world_size()})"
-
-        self.patch_embedding = PatchEmbed(in_chans=config.in_channels,
-                                          embed_dim=inner_dim,
-                                          patch_size=config.patch_size,
-                                          flatten=False)
-        self.condition_embedder = WanTimeTextImageEmbedding(
-            dim=inner_dim,
-            time_freq_dim=config.freq_dim,
-            text_embed_dim=config.text_dim,
-            image_embed_dim=config.image_dim,
-        )
-        self.blocks = nn.ModuleList([
-            DreamXWorldTransformerBlock(
-                inner_dim,
-                config.ffn_dim,
-                config.num_attention_heads,
-                config.qk_norm,
-                config.cross_attn_norm,
-                config.eps,
-                config.added_kv_proj_dim,
-                self._supported_attention_backends,
-                quant_config=config.quant_config,
-                prefix=f"{config.prefix}.blocks.{i}",
-                add_control_adapter=config.add_control_adapter,
-                cam_method=config.cam_method,
-                attn_compress=config.attn_compress,
-                cam_self_attn_layers=config.cam_self_attn_layers,
-                layer_idx=i)
-            for i in range(config.num_layers)
-        ])
-        self.norm_out = LayerNormScaleShift(inner_dim,
-                                            norm_type="layer",
-                                            eps=config.eps,
-                                            elementwise_affine=False,
-                                            dtype=torch.float32,
-                                            compute_dtype=torch.float32)
-        self.proj_out = nn.Linear(
-            inner_dim, config.out_channels * math.prod(config.patch_size))
-        self.scale_shift_table = nn.Parameter(
-            torch.randn(1, 2, inner_dim) / inner_dim**0.5)
-        self.gradient_checkpointing = False
-        self.__post_init__()
-
-    def forward(self,
-                hidden_states: torch.Tensor,
-                encoder_hidden_states: torch.Tensor | list[torch.Tensor],
-                timestep: torch.LongTensor,
-                encoder_hidden_states_image: torch.Tensor | list[torch.Tensor]
-                | None = None,
-                guidance=None,
-                y_camera: dict[str, torch.Tensor] | None = None,
-                **kwargs) -> torch.Tensor:
-        orig_dtype = hidden_states.dtype
-        if encoder_hidden_states is not None and not isinstance(
-                encoder_hidden_states, torch.Tensor):
-            encoder_hidden_states = encoder_hidden_states[0]
-        if isinstance(encoder_hidden_states_image,
-                      list) and len(encoder_hidden_states_image) > 0:
-            encoder_hidden_states_image = encoder_hidden_states_image[0]
-        else:
-            encoder_hidden_states_image = None
-
-        batch_size, _, num_frames, height, width = hidden_states.shape
-        p_t, p_h, p_w = self.patch_size
-        post_patch_num_frames = num_frames // p_t
-        post_patch_height = height // p_h
-        post_patch_width = width // p_w
-
-        d = self.hidden_size // self.num_attention_heads
-        rope_dim_list = [d - 4 * (d // 6), 2 * (d // 6), 2 * (d // 6)]
-        freqs_cos, freqs_sin = get_rotary_pos_embed(
-            (post_patch_num_frames, post_patch_height, post_patch_width),
-            self.hidden_size,
-            self.num_attention_heads,
-            rope_dim_list,
-            dtype=torch.float32 if current_platform.is_mps() else torch.float64,
-            rope_theta=10000)
-        freqs_cis = (freqs_cos.to(hidden_states.device).float(),
-                     freqs_sin.to(hidden_states.device).float())
-
-        hidden_states = self.patch_embedding(hidden_states)
-        hidden_states = hidden_states.flatten(2).transpose(1, 2)
-        hidden_states, original_seq_len = sequence_model_parallel_shard(
-            hidden_states, dim=1)
-
-        if timestep.dim() == 2:
-            ts_seq_len = timestep.shape[1]
-            timestep = timestep.flatten()
-        else:
-            ts_seq_len = None
-
-        temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = self.condition_embedder(
-            timestep,
-            encoder_hidden_states,
-            encoder_hidden_states_image,
-            timestep_seq_len=ts_seq_len)
-        if ts_seq_len is not None:
-            timestep_proj = timestep_proj.unflatten(2, (6, -1))
-        else:
-            timestep_proj = timestep_proj.unflatten(1, (6, -1))
-
-        if encoder_hidden_states_image is not None:
-            if encoder_hidden_states is not None:
-                encoder_hidden_states = torch.concat(
-                    [encoder_hidden_states_image, encoder_hidden_states], dim=1)
-            else:
-                encoder_hidden_states = encoder_hidden_states_image
-
-        if current_platform.is_mps() or current_platform.is_npu():
-            encoder_hidden_states = encoder_hidden_states.to(orig_dtype)
-
-        assert encoder_hidden_states.dtype == orig_dtype
-
-        if torch.is_grad_enabled() and self.gradient_checkpointing:
-            for block in self.blocks:
-                hidden_states = self._gradient_checkpointing_func(
-                    block, hidden_states, encoder_hidden_states,
-                    timestep_proj, freqs_cis, original_seq_len, y_camera)
-        else:
-            for block in self.blocks:
-                hidden_states = block(hidden_states, encoder_hidden_states,
-                                      timestep_proj, freqs_cis, original_seq_len,
-                                      y_camera=y_camera)
-
-        if temb.dim() == 3:
-            shift, scale = (self.scale_shift_table.unsqueeze(0) +
-                            temb.unsqueeze(2)).chunk(2, dim=2)
-            shift = shift.squeeze(2)
-            scale = scale.squeeze(2)
-        else:
-            shift, scale = (self.scale_shift_table +
-                            temb.unsqueeze(1)).chunk(2, dim=1)
-
-        hidden_states = self.norm_out(hidden_states, shift, scale)
-        hidden_states = sequence_model_parallel_all_gather_with_unpad(
-            hidden_states, original_seq_len, dim=1)
-        hidden_states = self.proj_out(hidden_states)
-
-        hidden_states = hidden_states.reshape(batch_size, post_patch_num_frames,
-                                              post_patch_height,
-                                              post_patch_width, p_t, p_h, p_w,
-                                              -1)
-        hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6)
-        output = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3)
-
-        return output
-
-
-EntryClass = DreamXWorldTransformer3DModel

fastvideo/models/dits/dreamx_world_ar.py

@@ -1,920 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""DreamX-World autoregressive causal DiT.
-
-Adapted from DreamX-World's Apache-2.0
-``wan/modules/causal_camera_model_2_2_prope_infinity.py``. The implementation is
-kept native to FastVideo: no production import from DreamX, Diffusers, or
-Transformers is required.
-"""
-
-from typing import Any
-
-import math
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-from fastvideo.configs.models.dits.dreamx_world import DreamXWorldARConfig
-from fastvideo.layers.linear import ReplicatedLinear
-from fastvideo.models.dits.base import BaseDiT
-from fastvideo.models.dits.dreamx_world import (_dreamx_apply_tiled_projmat,
-                                               _dreamx_prope_qkv)
-
-
-def attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
-    # Deliberately raw SDPA rather than fastvideo.attention.LocalAttention:
-    # (1) LocalAttention dispatches through the attention-backend registry, so
-    #     FLASH_ATTN could be selected and its kernel is not bit-identical to
-    #     torch SDPA — the AR KV-cache rollout must stay numerically frozen;
-    # (2) LocalAttention requires an active ForwardContext, which direct
-    #     transformer invocations (parity tests) do not set;
-    # (3) the sibling causal model keeps raw SDPA in the same KV-cache window
-    #     path (matrixgame2/causal_model.py).
-    # Sequence-parallel gap: this model never shards the sequence; run with
-    # sp_size=1 (see fastvideo/layers/AGENTS.md on documenting raw SDPA).
-    q_bhld = q.transpose(1, 2)
-    k_bhld = k.transpose(1, 2)
-    v_bhld = v.transpose(1, 2)
-    out = F.scaled_dot_product_attention(q_bhld, k_bhld, v_bhld, dropout_p=0.0)
-    return out.transpose(1, 2)
-
-
-def prope_qkv(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
-              viewmats: torch.Tensor, Ks: torch.Tensor):
-    q, k, v, output_projection = _dreamx_prope_qkv(q, k, v, viewmats, Ks)
-
-    def apply_fn_o(x: torch.Tensor) -> torch.Tensor:
-        return _dreamx_apply_tiled_projmat(x, output_projection)
-
-    return q, k, v, apply_fn_o
-
-
-def sinusoidal_embedding_1d(dim, position):
-    assert dim % 2 == 0
-    half = dim // 2
-    position = position.type(torch.float64)
-    sinusoid = torch.outer(
-        position, torch.pow(10000, -torch.arange(half).to(position).div(half)))
-    return torch.cat([torch.cos(sinusoid), torch.sin(sinusoid)], dim=1)
-
-
-def rope_params(max_seq_len, dim, theta=10000):
-    assert dim % 2 == 0
-    freqs = torch.outer(
-        torch.arange(max_seq_len),
-        1.0 / torch.pow(theta,
-                        torch.arange(0, dim, 2).to(torch.float64).div(dim)))
-    return torch.polar(torch.ones_like(freqs), freqs)
-
-
-class WanRMSNorm(nn.Module):
-    """Kept private instead of fastvideo.layers.layernorm.RMSNorm.
-
-    The official DreamX-World ``model_2_2.py`` computes the RMS statistics in
-    the *input* dtype — the upstream code has the fp32 upcast explicitly
-    commented out (``# return self._norm(x.float())...``). FastVideo's RMSNorm
-    always normalizes in fp32, which is not bit-identical under bf16, so the
-    verbatim implementation stays.
-    """
-
-    def __init__(self, dim, eps=1e-5):
-        super().__init__()
-        self.dim = dim
-        self.eps = eps
-        self.weight = nn.Parameter(torch.ones(dim))
-
-    def forward(self, x):
-        return self._norm(x).type_as(x) * self.weight
-
-    def _norm(self, x):
-        return x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps)
-
-
-class WanLayerNorm(nn.LayerNorm):
-    """Kept private instead of fastvideo.layers.layernorm.FP32LayerNorm.
-
-    The official DreamX-World ``model_2_2.py`` normalizes in the *input* dtype
-    (no ``x.float()`` upcast, unlike Wan2.1). FP32LayerNorm casts input and
-    affine params to fp32, which is not bit-identical under bf16, so the
-    verbatim implementation stays.
-    """
-
-    def __init__(self, dim, eps=1e-6, elementwise_affine=False):
-        super().__init__(dim, elementwise_affine=elementwise_affine, eps=eps)
-
-    def forward(self, x):
-        return super().forward(x).type_as(x)
-
-
-class WanCrossAttention(nn.Module):
-    def __init__(self, dim, num_heads, window_size=(-1, -1), qk_norm=True, eps=1e-6):
-        super().__init__()
-        self.num_heads = num_heads
-        self.head_dim = dim // num_heads
-        self.q = ReplicatedLinear(dim, dim)
-        self.k = ReplicatedLinear(dim, dim)
-        self.v = ReplicatedLinear(dim, dim)
-        self.o = ReplicatedLinear(dim, dim)
-        self.norm_q = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
-        self.norm_k = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
-
-    def _kv(self, context, b, n, d):
-        k, _ = self.k(context)
-        k = self.norm_k(k).view(b, -1, n, d)
-        v, _ = self.v(context)
-        v = v.view(b, -1, n, d)
-        return k, v
-
-    def forward(self, x, context, context_lens, crossattn_cache=None):
-        b, n, d = x.size(0), self.num_heads, self.head_dim
-        q, _ = self.q(x)
-        q = self.norm_q(q).view(b, -1, n, d)
-
-        if crossattn_cache is not None:
-            if not crossattn_cache["is_init"]:
-                crossattn_cache["is_init"] = True
-                k, v = self._kv(context, b, n, d)
-                crossattn_cache["k"] = k
-                crossattn_cache["v"] = v
-            else:
-                k = crossattn_cache["k"]
-                v = crossattn_cache["v"]
-        else:
-            k, v = self._kv(context, b, n, d)
-
-        x = attention(q, k, v)
-        x = x.flatten(2)
-        out, _ = self.o(x)
-        return out
-
-
-def block_relativistic_rope(x, grid_sizes, freqs, start_frame=0, relative_frame_indices=None):
-    """
-    Apply Block-Relativistic RoPE to input tensor.
-    Adapted from Infinity-RoPE (https://arxiv.org/abs/2511.20649).
-
-    Args:
-        x: Input tensor [B, L, num_heads, head_dim]
-        grid_sizes: Tensor [B, 3] containing (F, H, W)
-        freqs: RoPE frequencies
-        start_frame: Starting frame index for sequential RoPE
-        relative_frame_indices: Optional tensor [F] specifying explicit frame indices
-                               for Block-Relativistic RoPE. Overrides start_frame if provided.
-    """
-    n, c = x.size(2), x.size(3) // 2
-    freqs = freqs.split([c - 2 * (c // 3), c // 3, c // 3], dim=1)
-
-    output = []
-    for i, (f, h, w) in enumerate(grid_sizes.tolist()):
-        seq_len = f * h * w
-        x_i = torch.view_as_complex(x[i, :seq_len].to(torch.float64).reshape(
-            seq_len, n, -1, 2))
-
-        if relative_frame_indices is not None:
-            frame_indices = relative_frame_indices.long()
-            freqs_temporal = freqs[0][frame_indices].view(f, 1, 1, -1).expand(f, h, w, -1)
-        else:
-            freqs_temporal = freqs[0][start_frame:start_frame + f].view(f, 1, 1, -1).expand(f, h, w, -1)
-
-        freqs_i = torch.cat([
-            freqs_temporal,
-            freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1),
-            freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1)
-        ], dim=-1).reshape(seq_len, 1, -1)
-
-        x_i = torch.view_as_real(x_i * freqs_i).flatten(2)
-        x_i = torch.cat([x_i, x[i, seq_len:]])
-        output.append(x_i)
-
-    return torch.stack(output).type_as(x)
-
-
-class CausalWanSelfAttention(nn.Module):
-    """Self-attention with KV cache and Block-Relativistic RoPE for causal inference."""
-
-    def __init__(self, dim, num_heads, local_attn_size=6, sink_size=1,
-                 qk_norm=True, eps=1e-6):
-        assert dim % num_heads == 0
-        super().__init__()
-        self.dim = dim
-        self.num_heads = num_heads
-        self.head_dim = dim // num_heads
-        self.local_attn_size = local_attn_size
-        self.sink_size = sink_size
-        self.qk_norm = qk_norm
-        self.eps = eps
-        self.max_attention_size = 39600 if local_attn_size == -1 else local_attn_size * 880
-
-        self.q = ReplicatedLinear(dim, dim)
-        self.k = ReplicatedLinear(dim, dim)
-        self.v = ReplicatedLinear(dim, dim)
-        self.o = ReplicatedLinear(dim, dim)
-        self.norm_q = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
-        self.norm_k = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
-
-    def forward(self, x, seq_lens, grid_sizes, freqs, kv_cache,
-                current_start=0, cache_start=None, sink_recache_after_switch=False):
-        """
-        Args:
-            x: Shape [B, L, C]
-            seq_lens: Shape [B]
-            grid_sizes: Shape [B, 3] containing (F, H, W)
-            freqs: RoPE frequencies [1024, head_dim / 2]
-            kv_cache: Dict with 'k', 'v', 'global_end_index', 'local_end_index'
-            current_start: Current position in the global token sequence
-        """
-        b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim
-        if cache_start is None:
-            cache_start = current_start
-
-        q, _ = self.q(x)
-        q = self.norm_q(q).view(b, s, n, d)
-        k, _ = self.k(x)
-        k = self.norm_k(k).view(b, s, n, d)
-        v, _ = self.v(x)
-        v = v.view(b, s, n, d)
-
-        frame_seqlen = math.prod(grid_sizes[0][1:]).item()
-        num_new_frames = grid_sizes[0][0].item()
-        current_end = current_start + q.shape[1]
-        sink_tokens = self.sink_size * frame_seqlen
-        kv_cache_size = kv_cache["k"].shape[1]
-        num_new_tokens = q.shape[1]
-
-        cache_update_info = None
-        is_recompute = current_end <= kv_cache["global_end_index"].item() and current_start > 0
-
-        if self.local_attn_size != -1 and (current_end > kv_cache["global_end_index"].item()) and (
-                num_new_tokens + kv_cache["local_end_index"].item() > kv_cache_size):
-            # === ROLLING MODE: cache full, evict oldest non-sink tokens ===
-            num_evicted_tokens = num_new_tokens + kv_cache["local_end_index"].item() - kv_cache_size
-            num_rolled_tokens = kv_cache["local_end_index"].item() - num_evicted_tokens - sink_tokens
-            local_end_index = kv_cache["local_end_index"].item() + current_end - \
-                kv_cache["global_end_index"].item() - num_evicted_tokens
-            local_start_index = local_end_index - num_new_tokens
-
-            temp_k = kv_cache["k"].detach().clone()
-            temp_v = kv_cache["v"].detach().clone()
-            temp_k[:, sink_tokens:sink_tokens + num_rolled_tokens] = \
-                temp_k[:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone()
-            temp_v[:, sink_tokens:sink_tokens + num_rolled_tokens] = \
-                temp_v[:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone()
-
-            write_start_index = max(local_start_index, sink_tokens) if is_recompute else local_start_index
-            roped_offset = max(0, write_start_index - local_start_index)
-            write_len = max(0, local_end_index - write_start_index)
-            if write_len > 0:
-                temp_k[:, write_start_index:local_end_index] = k[:, roped_offset:roped_offset + write_len]
-                temp_v[:, write_start_index:local_end_index] = v[:, roped_offset:roped_offset + write_len]
-
-            # Block-Relativistic RoPE: query uses window-relative indices
-            query_relative_indices = torch.arange(
-                self.local_attn_size - num_new_frames, self.local_attn_size, device=q.device)
-            roped_query = block_relativistic_rope(
-                q, grid_sizes, freqs, relative_frame_indices=query_relative_indices).type_as(v)
-
-            # Block-Relativistic RoPE: cached K uses position-in-window indices
-            num_cache_frames = local_end_index // frame_seqlen
-            cache_relative_indices = torch.arange(0, num_cache_frames, device=k.device)
-            cache_grid_sizes = grid_sizes.clone()
-            cache_grid_sizes[0, 0] = num_cache_frames
-            roped_temp_k = block_relativistic_rope(
-                temp_k[:, :local_end_index].view(b, num_cache_frames, frame_seqlen, n, d).flatten(1, 2),
-                cache_grid_sizes, freqs, relative_frame_indices=cache_relative_indices).type_as(v)
-
-            cache_update_info = {
-                "action": "roll_and_insert",
-                "sink_tokens": sink_tokens,
-                "num_rolled_tokens": num_rolled_tokens,
-                "num_evicted_tokens": num_evicted_tokens,
-                "local_start_index": local_start_index,
-                "local_end_index": local_end_index,
-                "write_start_index": write_start_index,
-                "write_end_index": local_end_index,
-                "new_k": k[:, roped_offset:roped_offset + write_len],
-                "new_v": v[:, roped_offset:roped_offset + write_len],
-                "current_end": current_end,
-                "is_recompute": is_recompute
-            }
-        else:
-            # === DIRECT INSERT MODE: cache not yet full ===
-            local_end_index = kv_cache["local_end_index"].item() + current_end - kv_cache["global_end_index"].item()
-            local_start_index = local_end_index - num_new_tokens
-
-            temp_k = kv_cache["k"].detach().clone()
-            temp_v = kv_cache["v"].detach().clone()
-
-            write_start_index = max(local_start_index, sink_tokens) if is_recompute else local_start_index
-            if sink_recache_after_switch:
-                write_start_index = local_start_index
-            roped_offset = max(0, write_start_index - local_start_index)
-            write_len = max(0, local_end_index - write_start_index)
-            if write_len > 0:
-                temp_k[:, write_start_index:local_end_index] = k[:, roped_offset:roped_offset + write_len]
-                temp_v[:, write_start_index:local_end_index] = v[:, roped_offset:roped_offset + write_len]
-
-            # RoPE with relative indices (growing sequentially before cache fills)
-            current_frame_in_window = local_start_index // frame_seqlen
-            query_relative_indices = torch.arange(
-                current_frame_in_window, current_frame_in_window + num_new_frames, device=q.device)
-            roped_query = block_relativistic_rope(
-                q, grid_sizes, freqs, relative_frame_indices=query_relative_indices).type_as(v)
-
-            num_cache_frames = local_end_index // frame_seqlen
-            cache_relative_indices = torch.arange(0, num_cache_frames, device=k.device)
-            cache_grid_sizes = grid_sizes.clone()
-            cache_grid_sizes[0, 0] = num_cache_frames
-            roped_temp_k = block_relativistic_rope(
-                temp_k[:, :local_end_index].view(b, num_cache_frames, frame_seqlen, n, d).flatten(1, 2),
-                cache_grid_sizes, freqs, relative_frame_indices=cache_relative_indices).type_as(v)
-
-            cache_update_info = {
-                "action": "direct_insert",
-                "local_start_index": local_start_index,
-                "local_end_index": local_end_index,
-                "write_start_index": write_start_index,
-                "write_end_index": local_end_index,
-                "new_k": k[:, roped_offset:roped_offset + write_len],
-                "new_v": v[:, roped_offset:roped_offset + write_len],
-                "current_end": current_end,
-                "is_recompute": is_recompute
-            }
-
-        # Attention: sink tokens + local window
-        if sink_tokens > 0:
-            local_budget = self.max_attention_size - sink_tokens
-            k_sink = roped_temp_k[:, :sink_tokens]
-            v_sink = temp_v[:, :sink_tokens]
-            if local_budget > 0:
-                local_start_for_window = max(sink_tokens, local_end_index - local_budget)
-                k_local = roped_temp_k[:, local_start_for_window:local_end_index]
-                v_local = temp_v[:, local_start_for_window:local_end_index]
-                k_cat = torch.cat([k_sink, k_local], dim=1)
-                v_cat = torch.cat([v_sink, v_local], dim=1)
-            else:
-                k_cat = k_sink
-                v_cat = v_sink
-            x = attention(roped_query, k_cat, v_cat)
-        else:
-            window_start = max(0, local_end_index - self.max_attention_size)
-            x = attention(
-                roped_query,
-                roped_temp_k[:, window_start:local_end_index],
-                temp_v[:, window_start:local_end_index])
-
-        x = x.flatten(2)
-        x, _ = self.o(x)
-        return x, (current_end, local_end_index, cache_update_info)
-
-
-class CausalPropeSelfAttention(nn.Module):
-    """PRoPE self-attention with optional KV cache for camera-controlled inference."""
-
-    def __init__(self, dim, attn_dim, num_heads, window_size=(-1, -1),
-                 local_attn_size=-1, sink_size=0, qk_norm=True, eps=1e-6):
-        assert dim % num_heads == 0
-        assert attn_dim % num_heads == 0
-        super().__init__()
-        self.dim = dim
-        self.attn_dim = attn_dim
-        self.num_heads = num_heads
-        self.head_dim = attn_dim // num_heads
-        self.local_attn_size = local_attn_size
-        self.sink_size = sink_size
-        self.qk_norm = qk_norm
-        self.eps = eps
-        self.window_size = window_size
-        self.max_attention_size = 39600 if local_attn_size == -1 else local_attn_size * 880
-
-        self.q_proj = ReplicatedLinear(dim, attn_dim)
-        self.k_proj = ReplicatedLinear(dim, attn_dim)
-        self.v_proj = ReplicatedLinear(dim, attn_dim)
-        self.out_proj = ReplicatedLinear(attn_dim, dim)
-
-        self.norm_q = WanRMSNorm(attn_dim, eps=eps) if qk_norm else nn.Identity()
-        self.norm_k = WanRMSNorm(attn_dim, eps=eps) if qk_norm else nn.Identity()
-
-        nn.init.zeros_(self.out_proj.weight)
-        nn.init.zeros_(self.out_proj.bias)
-
-    def forward(self, x, cam_viewmats, cam_K, seq_lens, grid_sizes, freqs,
-                kv_cache=None, current_start=0, cache_start=None,
-                sink_recache_after_switch=False, cache_update_policy="commit_detached"):
-        """
-        Args:
-            x: Shape [B, L, C]
-            cam_viewmats: Camera view matrices
-            cam_K: Camera intrinsics
-            kv_cache: Optional KV cache dict. When None, runs full attention over current chunk.
-        """
-        b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim
-        if cache_start is None:
-            cache_start = current_start
-
-        q, _ = self.q_proj(x)
-        q = self.norm_q(q).view(b, s, n, d)
-        k, _ = self.k_proj(x)
-        k = self.norm_k(k).view(b, s, n, d)
-        v, _ = self.v_proj(x)
-        v = v.view(b, s, n, d)
-
-        # Apply PRoPE (Positional Rotary Position Embedding from camera parameters)
-        q_t, k_t, v_t, apply_fn_o = prope_qkv(
-            q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2),
-            viewmats=cam_viewmats, Ks=cam_K)
-        proped_q = q_t.transpose(1, 2)
-        proped_k = k_t.transpose(1, 2)
-        proped_v = v_t.transpose(1, 2)
-
-        if kv_cache is None:
-            # No cache: full attention over current chunk
-            x_out = attention(proped_q, proped_k, proped_v)
-        else:
-            # KV cache mode with rolling cache support
-            frame_seqlen = math.prod(grid_sizes[0][1:]).item()
-            num_new_tokens = s
-            current_end = current_start + num_new_tokens
-            sink_tokens = self.sink_size * frame_seqlen
-            kv_cache_size = kv_cache["k"].shape[1]
-            is_recompute = (current_end <= kv_cache["global_end_index"].item()) and (current_start > 0)
-
-            if self.local_attn_size != -1 and (current_end > kv_cache["global_end_index"].item()) and (
-                    num_new_tokens + kv_cache["local_end_index"].item() > kv_cache_size):
-                # === ROLLING MODE ===
-                num_evicted_tokens = num_new_tokens + kv_cache["local_end_index"].item() - kv_cache_size
-                num_rolled_tokens = kv_cache["local_end_index"].item() - num_evicted_tokens - sink_tokens
-                local_end_index = kv_cache["local_end_index"].item() + current_end - \
-                    kv_cache["global_end_index"].item() - num_evicted_tokens
-                local_start_index = local_end_index - num_new_tokens
-
-                if cache_update_policy != "none":
-                    with torch.no_grad():
-                        kv_cache["k"][:, sink_tokens:sink_tokens + num_rolled_tokens] = \
-                            kv_cache["k"][:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone()
-                        kv_cache["v"][:, sink_tokens:sink_tokens + num_rolled_tokens] = \
-                            kv_cache["v"][:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone()
-
-                    write_start_index = max(local_start_index, sink_tokens) if is_recompute else local_start_index
-                    roped_offset = max(0, write_start_index - local_start_index)
-                    write_len = max(0, local_end_index - write_start_index)
-                    if write_len > 0:
-                        with torch.no_grad():
-                            kv_cache["k"][:, write_start_index:local_end_index] = proped_k[:, roped_offset:roped_offset + write_len].detach()
-                            kv_cache["v"][:, write_start_index:local_end_index] = proped_v[:, roped_offset:roped_offset + write_len].detach()
-            else:
-                # === DIRECT INSERT MODE ===
-                local_end_index = kv_cache["local_end_index"].item() + current_end - kv_cache["global_end_index"].item()
-                local_start_index = local_end_index - num_new_tokens
-
-                if cache_update_policy != "none":
-                    write_start_index = max(local_start_index, sink_tokens) if is_recompute else local_start_index
-                    if sink_recache_after_switch:
-                        write_start_index = local_start_index
-                    roped_offset = max(0, write_start_index - local_start_index)
-                    write_len = max(0, local_end_index - write_start_index)
-                    if write_len > 0:
-                        with torch.no_grad():
-                            kv_cache["k"][:, write_start_index:local_end_index] = proped_k[:, roped_offset:roped_offset + write_len].detach()
-                            kv_cache["v"][:, write_start_index:local_end_index] = proped_v[:, roped_offset:roped_offset + write_len].detach()
-
-            # Attention: sink tokens + local window
-            if sink_tokens > 0:
-                local_budget = self.max_attention_size - sink_tokens
-                k_sink = kv_cache["k"][:, :sink_tokens].detach()
-                v_sink = kv_cache["v"][:, :sink_tokens].detach()
-                if local_budget > 0:
-                    local_start_for_window = max(sink_tokens, local_end_index - local_budget)
-                    k_local = kv_cache["k"][:, local_start_for_window:local_end_index].detach()
-                    v_local = kv_cache["v"][:, local_start_for_window:local_end_index].detach()
-                    k_cat = torch.cat([k_sink, k_local], dim=1)
-                    v_cat = torch.cat([v_sink, v_local], dim=1)
-                else:
-                    k_cat = k_sink
-                    v_cat = v_sink
-                x_out = attention(proped_q, k_cat, v_cat)
-            else:
-                window_start = max(0, local_end_index - self.max_attention_size)
-                x_out = attention(
-                    proped_q,
-                    kv_cache["k"][:, window_start:local_end_index].detach(),
-                    kv_cache["v"][:, window_start:local_end_index].detach())
-
-            if not is_recompute and cache_update_policy != "none":
-                kv_cache["global_end_index"].fill_(current_end)
-                kv_cache["local_end_index"].fill_(local_end_index)
-
-        # Apply inverse PRoPE
-        x = apply_fn_o(x_out.transpose(1, 2)).transpose(1, 2)
-        x = x.flatten(2)
-        x, _ = self.out_proj(x)
-        return x
-
-
-class CausalWanAttentionBlock(nn.Module):
-
-    def __init__(self, dim, ffn_dim, num_heads, local_attn_size=-1, sink_size=0,
-                 qk_norm=True, cross_attn_norm=False, eps=1e-6, **kwargs):
-        super().__init__()
-        self.dim = dim
-        self.ffn_dim = ffn_dim
-        self.num_heads = num_heads
-        self.local_attn_size = local_attn_size
-        self.qk_norm = qk_norm
-        self.cross_attn_norm = cross_attn_norm
-        self.eps = eps
-
-        self.add_control_adapter = kwargs.get('add_control_adapter', False)
-        self.cam_method = kwargs.get('cam_method')
-        self.attn_compress = kwargs.get('attn_compress', 1)
-        self.layer_idx = kwargs.get('layer_idx')
-        cam_self_attn_layers = kwargs.get('cam_self_attn_layers')
-
-        # layers
-        self.norm1 = WanLayerNorm(dim, eps)
-        self.self_attn = CausalWanSelfAttention(
-            dim, num_heads, local_attn_size, sink_size, qk_norm, eps)
-        self.norm3 = WanLayerNorm(
-            dim, eps, elementwise_affine=True) if cross_attn_norm else nn.Identity()
-        self.cross_attn = WanCrossAttention(dim, num_heads, (-1, -1), qk_norm, eps)
-        self.norm2 = WanLayerNorm(dim, eps)
-        # nn.Linear (not ReplicatedLinear) on purpose: the official checkpoint
-        # stores these as positional Sequential keys (ffn.0 / ffn.2) that the
-        # copy-only converter and the strict-load tests require verbatim, and
-        # ReplicatedLinear's (out, bias) tuple return cannot compose inside
-        # nn.Sequential without renaming the state-dict surface.
-        self.ffn = nn.Sequential(
-            nn.Linear(dim, ffn_dim), nn.GELU(approximate='tanh'),
-            nn.Linear(ffn_dim, dim))
-
-        # PRoPE self-attention branch for camera control
-        add_cam_attn = self.add_control_adapter and self.cam_method == 'prope'
-        if add_cam_attn and cam_self_attn_layers is not None:
-            add_cam_attn = self.layer_idx in cam_self_attn_layers
-        if add_cam_attn:
-            self.cam_self_attn = CausalPropeSelfAttention(
-                dim, dim // self.attn_compress, num_heads,
-                local_attn_size=local_attn_size, sink_size=sink_size,
-                qk_norm=qk_norm, eps=eps)
-
-        # modulation
-        self.modulation = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5)
-
-    def forward(self, x, e, seq_lens, grid_sizes, freqs, context, context_lens,
-                kv_cache, crossattn_cache=None, current_start=0, cache_start=None,
-                cam_viewmats=None, cam_K=None, sink_recache_after_switch=False,
-                cache_update_policy="commit_detached"):
-        num_frames, frame_seqlen = e.shape[1], x.shape[1] // e.shape[1]
-        e = (self.modulation.unsqueeze(1) + e).chunk(6, dim=2)
-
-        # self-attention
-        attn_input = (self.norm1(x).unflatten(
-            dim=1, sizes=(num_frames, frame_seqlen)) * (1 + e[1]) + e[0]).flatten(1, 2)
-        y, cache_update_info = self.self_attn(
-            attn_input, seq_lens, grid_sizes, freqs, kv_cache,
-            current_start, cache_start, sink_recache_after_switch)
-
-        # PRoPE camera attention (parallel branch)
-        if hasattr(self, 'cam_self_attn') and cam_viewmats is not None and cam_K is not None:
-            prope_kv_cache = None
-            if kv_cache is not None and "prope_k" in kv_cache:
-                prope_kv_cache = {
-                    "k": kv_cache["prope_k"],
-                    "v": kv_cache["prope_v"],
-                    "global_end_index": kv_cache["prope_global_end_index"],
-                    "local_end_index": kv_cache["prope_local_end_index"],
-                }
-            y = y + self.cam_self_attn(
-                attn_input, cam_viewmats, cam_K, seq_lens, grid_sizes, freqs,
-                kv_cache=prope_kv_cache, current_start=current_start,
-                cache_start=cache_start, cache_update_policy=cache_update_policy)
-
-        x = x + (y.unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * e[2]).flatten(1, 2)
-
-        # cross-attention & FFN
-        x = x + self.cross_attn(self.norm3(x), context, context_lens,
-                                crossattn_cache=crossattn_cache)
-        y = self.ffn(
-            (self.norm2(x).unflatten(dim=1, sizes=(num_frames, frame_seqlen))
-             * (1 + e[4]) + e[3]).flatten(1, 2))
-        x = x + (y.unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * e[5]).flatten(1, 2)
-
-        return x, cache_update_info
-
-
-class CausalHead(nn.Module):
-
-    def __init__(self, dim, out_dim, patch_size, eps=1e-6):
-        super().__init__()
-        self.dim = dim
-        self.out_dim = out_dim
-        self.patch_size = patch_size
-        self.eps = eps
-
-        out_dim = math.prod(patch_size) * out_dim
-        self.norm = WanLayerNorm(dim, eps)
-        self.head = ReplicatedLinear(dim, out_dim)
-        self.modulation = nn.Parameter(torch.randn(1, 2, dim) / dim**0.5)
-
-    def forward(self, x, e):
-        num_frames, frame_seqlen = e.shape[1], x.shape[1] // e.shape[1]
-        e = (self.modulation.unsqueeze(1) + e).chunk(2, dim=2)
-        x, _ = self.head(
-            self.norm(x).unflatten(dim=1, sizes=(num_frames, frame_seqlen))
-            * (1 + e[1]) + e[0])
-        return x
-
-
-class DreamXWorldARTransformer3DModel(BaseDiT):
-    """DreamX-World-5B autoregressive causal transformer."""
-
-    _fsdp_shard_conditions = DreamXWorldARConfig()._fsdp_shard_conditions
-    _compile_conditions = DreamXWorldARConfig()._compile_conditions
-    _supported_attention_backends = DreamXWorldARConfig()._supported_attention_backends
-    param_names_mapping = DreamXWorldARConfig().param_names_mapping
-    reverse_param_names_mapping = DreamXWorldARConfig().reverse_param_names_mapping
-    lora_param_names_mapping = DreamXWorldARConfig().lora_param_names_mapping
-    _no_split_modules = ["CausalWanAttentionBlock"]
-
-    def __init__(self, config: DreamXWorldARConfig, hf_config: dict[str, Any]) -> None:
-        super().__init__(config=config, hf_config=hf_config)
-
-        model_type = config.model_type
-        patch_size = config.patch_size
-        text_len = config.text_len
-        in_dim = config.in_channels
-        dim = config.hidden_size
-        ffn_dim = config.ffn_dim
-        freq_dim = config.freq_dim
-        text_dim = config.text_dim
-        out_dim = config.out_channels
-        num_heads = config.num_attention_heads
-        num_layers = config.num_layers
-        local_attn_size = config.local_attn_size
-        sink_size = config.sink_size
-        qk_norm = bool(config.qk_norm)
-        cross_attn_norm = config.cross_attn_norm
-        eps = config.eps
-        add_control_adapter = config.add_control_adapter
-        cam_method = config.cam_method
-        attn_compress = config.attn_compress
-        cam_self_attn_layers = config.cam_self_attn_layers
-
-        assert model_type in ['t2v', 'i2v', 'ti2v']
-        self.model_type = model_type
-        self.patch_size = patch_size
-        self.text_len = text_len
-        self.in_dim = in_dim
-        self.dim = dim
-        self.ffn_dim = ffn_dim
-        self.freq_dim = freq_dim
-        self.text_dim = text_dim
-        self.out_dim = out_dim
-        self.num_heads = num_heads
-        self.num_layers = num_layers
-        self.local_attn_size = local_attn_size
-        self.qk_norm = qk_norm
-        self.cross_attn_norm = cross_attn_norm
-        self.eps = eps
-
-        # embeddings — nn.Linear inside nn.Sequential on purpose: the official
-        # checkpoint keys are positional (text_embedding.0/.2, time_embedding.0/.2,
-        # time_projection.1) and must load verbatim (see ffn comment above).
-        self.patch_embedding = nn.Conv3d(
-            in_dim, dim, kernel_size=patch_size, stride=patch_size)
-        self.text_embedding = nn.Sequential(
-            nn.Linear(text_dim, dim), nn.GELU(approximate='tanh'),
-            nn.Linear(dim, dim))
-        self.time_embedding = nn.Sequential(
-            nn.Linear(freq_dim, dim), nn.SiLU(), nn.Linear(dim, dim))
-        self.time_projection = nn.Sequential(
-            nn.SiLU(), nn.Linear(dim, dim * 6))
-
-        # transformer blocks
-        self.blocks = nn.ModuleList([
-            CausalWanAttentionBlock(
-                dim, ffn_dim, num_heads, local_attn_size, sink_size,
-                qk_norm, cross_attn_norm, eps,
-                add_control_adapter=add_control_adapter,
-                cam_method=cam_method,
-                attn_compress=attn_compress,
-                layer_idx=layer_idx,
-                cam_self_attn_layers=cam_self_attn_layers)
-            for layer_idx in range(num_layers)
-        ])
-        for layer_idx, block in enumerate(self.blocks):
-            block.self_attn.layer_idx = layer_idx
-            block.self_attn.num_layers = self.num_layers
-
-        # head
-        self.head = CausalHead(dim, out_dim, patch_size, eps)
-
-        # RoPE frequencies
-        assert (dim % num_heads) == 0 and (dim // num_heads) % 2 == 0
-        d = dim // num_heads
-        self.freqs = torch.cat([
-            rope_params(1024, d - 4 * (d // 6)),
-            rope_params(1024, 2 * (d // 6)),
-            rope_params(1024, 2 * (d // 6))
-        ], dim=1)
-
-        self.num_attention_heads = num_heads
-        self.attention_head_dim = dim // num_heads
-        self.hidden_size = dim
-        self.in_channels = in_dim
-        self.out_channels = out_dim
-        self.num_channels_latents = out_dim
-        self.init_weights()
-        self.num_frame_per_block = config.arch_config.num_frames_per_block
-        self.__post_init__()
-
-    def forward(self, x=None, t=None, context=None, seq_len=None, y=None, y_camera=None,
-                kv_cache=None, crossattn_cache=None, current_start=0,
-                cache_start=0, cache_update_policy="commit_detached",
-                hidden_states=None, encoder_hidden_states=None, timestep=None, **kwargs):
-        """
-        Causal inference with KV caching.
-        See Algorithm 2 of CausVid (https://arxiv.org/abs/2412.07772).
-
-        Args:
-            x: List of input video tensors [C_in, F, H, W]
-            t: Timestep tensor [B, L]
-            context: List of text embeddings [L, C]
-            seq_len: Maximum sequence length for positional encoding
-            y: Optional conditional video inputs (I2V mode)
-            y_camera: Camera parameters dict {'viewmats': ..., 'K': ...}
-            kv_cache: List of KV cache dicts per transformer block
-            crossattn_cache: List of cross-attention cache dicts
-            current_start: Current position in global token sequence
-            cache_start: Cache start position
-            cache_update_policy: Cache update strategy ('commit_detached' or 'none')
-
-        Returns:
-            Stacked output tensors [B, C_out, F, H/8, W/8]
-        """
-        if x is None and hidden_states is not None:
-            x = [sample for sample in hidden_states]
-        if t is None and timestep is not None:
-            t = timestep
-        if context is None and encoder_hidden_states is not None:
-            if isinstance(encoder_hidden_states, torch.Tensor):
-                context = [sample for sample in encoder_hidden_states]
-            else:
-                context = encoder_hidden_states
-        if seq_len is None:
-            if torch.is_tensor(t):
-                seq_len = int(t.shape[1]) if t.dim() > 1 else int(t.numel())
-            elif x is not None:
-                sample = x[0]
-                seq_len = (sample.shape[1] // self.patch_size[0]) * (sample.shape[2] // self.patch_size[1]) * (sample.shape[3] // self.patch_size[2])
-        if x is None or t is None or context is None or seq_len is None:
-            raise ValueError("DreamXWorldARTransformer3DModel requires x/t/context/seq_len or FastVideo aliases")
-
-        device = self.patch_embedding.weight.device
-        if self.freqs.is_meta or self.freqs.device != device:
-            d = self.dim // self.num_heads
-            self.freqs = torch.cat([
-                rope_params(1024, d - 4 * (d // 6)),
-                rope_params(1024, 2 * (d // 6)),
-                rope_params(1024, 2 * (d // 6)),
-            ], dim=1).to(device)
-
-        if y is not None:
-            x = [torch.cat([u, v], dim=0) for u, v in zip(x, y, strict=True)]
-
-        # patch embedding
-        x = [self.patch_embedding(u.unsqueeze(0)) for u in x]
-        grid_sizes = torch.stack(
-            [torch.tensor(u.shape[2:], dtype=torch.long) for u in x])
-        x = [u.flatten(2).transpose(1, 2) for u in x]
-        seq_lens = torch.tensor([u.size(1) for u in x], dtype=torch.long)
-        assert seq_lens.max() <= seq_len
-        x = torch.cat(x)
-
-        # time embedding
-        e = self.time_embedding(
-            sinusoidal_embedding_1d(self.freq_dim, t.flatten()).type_as(x))
-        e0 = self.time_projection(e).unflatten(
-            1, (6, self.dim)).unflatten(dim=0, sizes=t.shape)
-
-        # text embedding
-        context_lens = None
-        context = self.text_embedding(
-            torch.stack([
-                torch.cat(
-                    [u, u.new_zeros(self.text_len - u.size(0), u.size(1))])
-                for u in context
-            ]))
-
-        # camera parameters
-        if y_camera is not None and isinstance(y_camera, dict):
-            cam_viewmats = y_camera['viewmats']
-            cam_K = y_camera['K']
-        else:
-            cam_viewmats = None
-            cam_K = None
-
-        block_kwargs = dict(
-            e=e0, seq_lens=seq_lens, grid_sizes=grid_sizes, freqs=self.freqs,
-            context=context, context_lens=context_lens,
-            cam_viewmats=cam_viewmats, cam_K=cam_K,
-            cache_update_policy=cache_update_policy,
-        )
-
-        cache_update_infos = []
-        for block_index, block in enumerate(self.blocks):
-            block_kwargs.update({
-                "kv_cache": kv_cache[block_index] if kv_cache is not None else None,
-                "crossattn_cache": crossattn_cache[block_index] if crossattn_cache is not None else None,
-                "current_start": current_start,
-                "cache_start": cache_start,
-            })
-            x, block_cache_update_info = block(x, **block_kwargs)
-            if kv_cache is not None:
-                cache_update_infos.append((block_index, block_cache_update_info))
-
-        # Apply deferred cache updates
-        if kv_cache is not None and cache_update_infos and cache_update_policy != "none":
-            self._apply_cache_updates(kv_cache, cache_update_infos)
-
-        # head & unpatchify
-        x = self.head(x, e.unflatten(dim=0, sizes=t.shape).unsqueeze(2))
-        x = self.unpatchify(x, grid_sizes)
-        return torch.stack(x)
-
-    def _apply_cache_updates(self, kv_cache, cache_update_infos):
-        """Apply deferred cache updates collected from all transformer blocks.
-
-        For Block-Relativistic RoPE, this stores un-roped K values in the cache.
-        RoPE is applied dynamically during attention based on each token's current
-        relative position in the sliding window.
-        """
-        with torch.no_grad():
-            for block_index, (current_end, local_end_index, update_info) in cache_update_infos:
-                if update_info is not None:
-                    cache = kv_cache[block_index]
-
-                    if update_info["action"] == "roll_and_insert":
-                        sink_tokens = update_info["sink_tokens"]
-                        num_rolled_tokens = update_info["num_rolled_tokens"]
-                        num_evicted_tokens = update_info["num_evicted_tokens"]
-                        write_start_index = update_info.get("write_start_index", update_info["local_start_index"])
-                        write_end_index = update_info.get("write_end_index", update_info["local_end_index"])
-                        new_k = update_info["new_k"].detach()
-                        new_v = update_info["new_v"].detach()
-
-                        cache["k"][:, sink_tokens:sink_tokens + num_rolled_tokens] = \
-                            cache["k"][:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone()
-                        cache["v"][:, sink_tokens:sink_tokens + num_rolled_tokens] = \
-                            cache["v"][:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone()
-
-                        if write_end_index > write_start_index and new_k.shape[1] == (write_end_index - write_start_index):
-                            cache["k"][:, write_start_index:write_end_index] = new_k
-                            cache["v"][:, write_start_index:write_end_index] = new_v
-
-                    elif update_info["action"] == "direct_insert":
-                        write_start_index = update_info.get("write_start_index", update_info["local_start_index"])
-                        write_end_index = update_info.get("write_end_index", update_info["local_end_index"])
-                        new_k = update_info["new_k"].detach()
-                        new_v = update_info["new_v"].detach()
-
-                        if write_end_index > write_start_index and new_k.shape[1] == (write_end_index - write_start_index):
-                            cache["k"][:, write_start_index:write_end_index] = new_k
-                            cache["v"][:, write_start_index:write_end_index] = new_v
-
-                is_recompute = False if update_info is None else update_info.get("is_recompute", False)
-                if not is_recompute:
-                    kv_cache[block_index]["global_end_index"].fill_(current_end)
-                    kv_cache[block_index]["local_end_index"].fill_(local_end_index)
-
-    def unpatchify(self, x, grid_sizes):
-        """Reconstruct video tensors from patch embeddings."""
-        c = self.out_dim
-        out = []
-        for u, v in zip(x, grid_sizes.tolist(), strict=True):
-            u = u[:math.prod(v)].view(*v, *self.patch_size, c)
-            u = torch.einsum('fhwpqrc->cfphqwr', u)
-            u = u.reshape(c, *[i * j for i, j in zip(v, self.patch_size, strict=True)])
-            out.append(u)
-        return out
-
-    def init_weights(self):
-        """Initialize model parameters using Xavier initialization."""
-        for m in self.modules():
-            if isinstance(m, (nn.Linear, ReplicatedLinear)):
-                nn.init.xavier_uniform_(m.weight)
-                if m.bias is not None:
-                    nn.init.zeros_(m.bias)
-
-        nn.init.xavier_uniform_(self.patch_embedding.weight.flatten(1))
-        for m in self.text_embedding.modules():
-            if isinstance(m, nn.Linear):
-                nn.init.normal_(m.weight, std=.02)
-        for m in self.time_embedding.modules():
-            if isinstance(m, nn.Linear):
-                nn.init.normal_(m.weight, std=.02)
-
-        nn.init.zeros_(self.head.head.weight)
-
-
-EntryClass = DreamXWorldARTransformer3DModel

fastvideo/models/dits/ltx2.py

@@ -72,12 +72,8 @@ def _rms_norm_dispatch(
     """Use QuACK RMSNorm only for LTX-2 refine stage, else Torch RMSNorm."""
     if _is_ltx2_refine_stage():
         return _quack_rmsnorm(x, weight=weight, eps=eps)
-    # torch 2.12 added rms_norm to autocast's float32 cast policy, so under
-    # autocast(bfloat16) this norm now returns float32 where torch <= 2.11
-    # returned bfloat16. Restore the input dtype so autocast-blind consumers
-    # downstream don't receive upcast activations.
     return torch.nn.functional.rms_norm(
-        x, (x.shape[-1], ), weight=weight, eps=eps).to(x.dtype)
+        x, (x.shape[-1], ), weight=weight, eps=eps)
 
 
 class StageAwareRMSNorm(nn.RMSNorm):

fastvideo/models/registry.py

@@ -11,7 +11,7 @@ import tempfile
 from abc import ABC, abstractmethod
 from collections.abc import Callable, Set
 from dataclasses import dataclass, field
-from functools import cache, lru_cache
+from functools import lru_cache
 from typing import NoReturn, TypeVar, cast
 
 import cloudpickle
@@ -32,8 +32,6 @@ _TEXT_TO_VIDEO_DIT_MODELS = {
     "HYWorldTransformer3DModel":
     ("dits", "hyworld", "HYWorldTransformer3DModel"),
     "WanTransformer3DModel": ("dits", "wanvideo", "WanTransformer3DModel"),
-    "DreamXWorldTransformer3DModel": ("dits", "dreamx_world", "DreamXWorldTransformer3DModel"),
-    "DreamXWorldARTransformer3DModel": ("dits", "dreamx_world_ar", "DreamXWorldARTransformer3DModel"),
     "CausalWanTransformer3DModel": ("dits", "causal_wanvideo", "CausalWanTransformer3DModel"),
     "CosmosTransformer3DModel": ("dits", "cosmos", "CosmosTransformer3DModel"),
     "Cosmos25Transformer3DModel": ("dits", "cosmos2_5", "Cosmos25Transformer3DModel"),
@@ -50,8 +48,6 @@ _TEXT_TO_VIDEO_DIT_MODELS = {
 _IMAGE_TO_VIDEO_DIT_MODELS = {
     # "HunyuanVideoTransformer3DModel": ("dits", "hunyuanvideo", "HunyuanVideoDiT"),
     "WanTransformer3DModel": ("dits", "wanvideo", "WanTransformer3DModel"),
-    "DreamXWorldTransformer3DModel": ("dits", "dreamx_world", "DreamXWorldTransformer3DModel"),
-    "DreamXWorldARTransformer3DModel": ("dits", "dreamx_world_ar", "DreamXWorldARTransformer3DModel"),
     "CausalWanTransformer3DModel": ("dits", "causal_wanvideo", "CausalWanTransformer3DModel"),
     "MatrixGame2WanModel": ("dits", "matrixgame2", "MatrixGame2WanModel"),
     "CausalMatrixGame2WanModel": ("dits", "matrixgame2", "CausalMatrixGame2WanModel"),
@@ -145,7 +141,7 @@ _LEGACY_FAST_VIDEO_MODELS = {
 MODELS_PATH = os.path.dirname(__file__)
 
 
-@cache
+@lru_cache(maxsize=None)
 def _discover_and_register_models() -> dict[str, tuple[str, str, str]]:
     discovered_models: dict[str, tuple[str, str, str]] = {}
     for root, dirs, files in os.walk(MODELS_PATH):
@@ -160,7 +156,7 @@ def _discover_and_register_models() -> dict[str, tuple[str, str, str]]:
 
             filepath = os.path.join(root, filename)
             try:
-                with open(filepath, encoding="utf-8") as f:
+                with open(filepath, "r", encoding="utf-8") as f:
                     source = f.read()
                 tree = ast.parse(source, filename=filename)
 

fastvideo/performance/hf_store.py

@@ -268,31 +268,16 @@ def load_records_for_model(
     model_id: str,
     gpu_type: str | None = None,
     *,
-    workload_id: str | None = None,
-    variant_id: str | None = None,
-    benchmark_version: str | None = None,
-    recipe_fingerprint: str | None = None,
-    hardware_profile_id: str | None = None,
-    software_profile_id: str | None = None,
     last_n: int | None = None,
     successful_only: bool = True,
     baseline_eligible_only: bool = False,
 ) -> list[dict[str, Any]]:
-    """Return records for a specific *model_id*, optionally filtered by cohort.
+    """Return records for a specific *model_id*, optionally filtered by GPU.
 
     Args:
         local_dir: Root directory previously populated by :func:`sync_from_hf`.
         model_id: Matches the ``model_id`` field inside each JSON record.
         gpu_type: When set, only records whose ``gpu_type`` matches are returned.
-        workload_id: When set, only records from the same workload are returned.
-        variant_id: When set, only records from the same workload variant are returned.
-        benchmark_version: When set, only records from the same benchmark version are returned.
-        recipe_fingerprint: When set, only records from the same benchmark
-            recipe are returned.
-        hardware_profile_id: When set, only records from the same hardware
-            cohort are returned.
-        software_profile_id: When set, only records from the same software
-            cohort are returned.
         last_n: When set, return only the most recent *n* records (after all
             other filters). Useful for sliding-window baseline calculations.
         successful_only: Passed through to :func:`load_records`.
@@ -314,18 +299,6 @@ def load_records_for_model(
     if gpu_type is not None:
         records = [r for r in records if r.get("gpu_type") == gpu_type]
 
-    identity_filters = {
-        "workload_id": workload_id,
-        "variant_id": variant_id,
-        "benchmark_version": benchmark_version,
-        "recipe_fingerprint": recipe_fingerprint,
-        "hardware_profile_id": hardware_profile_id,
-        "software_profile_id": software_profile_id,
-    }
-    for key, expected in identity_filters.items():
-        if expected is not None:
-            records = [r for r in records if str(r.get(key)) == str(expected)]
-
     if last_n is not None:
         records = records[-last_n:]
 
@@ -337,12 +310,8 @@ def load_records_for_model(
 # ---------------------------------------------------------------------------
 
 _NUMERIC_COLS = (
-    "latency",
-    "throughput",
-    "memory",
-    "text_encoder_time_s",
-    "dit_time_s",
-    "vae_decode_time_s",
+    "latency", "throughput", "memory",
+    "text_encoder_time_s", "dit_time_s", "vae_decode_time_s",
 )
 
 

fastvideo/performance/metric_policy.py

@@ -36,7 +36,6 @@ DEFAULT_METRIC_POLICIES: tuple[MetricPolicy, ...] = (
     MetricPolicy("vae_decode_time_s", "VAE Decode", 3, True, 0.05, 0.25),
 )
 
-
 def _optional_float(value: Any) -> float | None:
     if value is None or isinstance(value, bool):
         return None
@@ -58,7 +57,9 @@ def _optional_bool(value: Any) -> bool | None:
     return None
 
 
-def resolve_metric_policies(threshold_overrides: Mapping[str, Any] | None, ) -> tuple[MetricPolicy, ...]:
+def resolve_metric_policies(
+    threshold_overrides: Mapping[str, Any] | None,
+) -> tuple[MetricPolicy, ...]:
     """Return default metric policies with optional per-metric overrides."""
 
     if not isinstance(threshold_overrides, Mapping):
@@ -79,15 +80,25 @@ def resolve_metric_policies(threshold_overrides: Mapping[str, Any] | None, ) ->
                 label=base_policy.label,
                 precision=base_policy.precision,
                 lower_is_better=base_policy.lower_is_better,
-                threshold_percent=(base_policy.threshold_percent if threshold_percent is None else threshold_percent),
-                threshold_absolute=(base_policy.threshold_absolute
-                                    if threshold_absolute is None else threshold_absolute),
+                threshold_percent=(
+                    base_policy.threshold_percent
+                    if threshold_percent is None
+                    else threshold_percent
+                ),
+                threshold_absolute=(
+                    base_policy.threshold_absolute
+                    if threshold_absolute is None
+                    else threshold_absolute
+                ),
                 gated=base_policy.gated if gated is None else gated,
-            ))
+            )
+        )
     return tuple(policies)
 
 
-def serialize_metric_thresholds(policies: tuple[MetricPolicy, ...], ) -> dict[str, dict[str, float | bool]]:
+def serialize_metric_thresholds(
+    policies: tuple[MetricPolicy, ...],
+) -> dict[str, dict[str, float | bool]]:
     return {
         policy.key: {
             "threshold_percent": policy.threshold_percent,
@@ -107,7 +118,10 @@ def regression_delta(
         return None
     absolute_delta = current - baseline if policy.lower_is_better else baseline - current
     percent_delta = absolute_delta / baseline
-    threshold_exceeded = (percent_delta > policy.threshold_percent and absolute_delta > policy.threshold_absolute)
+    threshold_exceeded = (
+        percent_delta > policy.threshold_percent
+        and absolute_delta > policy.threshold_absolute
+    )
     return MetricDelta(
         absolute=absolute_delta,
         percent=percent_delta,

fastvideo/performance_dashboard/service.py

@@ -17,15 +17,6 @@ from fastvideo.performance.hf_store import is_baseline_eligible_record, safe_flo
 from fastvideo.performance.metric_policy import regression_delta, resolve_metric_policies
 
 Record = dict[str, Any]
-CohortKey = tuple[str, str, str, str, str, str, str, str]
-COMPARISON_COHORT_KEYS = (
-    "workload_id",
-    "variant_id",
-    "benchmark_version",
-    "recipe_fingerprint",
-    "hardware_profile_id",
-    "software_profile_id",
-)
 
 
 def parse_timestamp(value: Any) -> datetime | None:
@@ -86,44 +77,15 @@ def record_metadata(record: Record) -> Record:
     }
 
 
-def record_comparison_metadata(record: Record) -> Record:
-    return {key: record.get(key) or "" for key in COMPARISON_COHORT_KEYS}
-
-
-def comparison_cohort_key(record: Record) -> CohortKey:
-    return (
-        str(record.get("model_id") or "unknown"),
-        str(record.get("gpu_type") or "unknown"),
-        str(record.get("workload_id") or ""),
-        str(record.get("variant_id") or ""),
-        str(record.get("benchmark_version") or ""),
-        str(record.get("recipe_fingerprint") or ""),
-        str(record.get("hardware_profile_id") or ""),
-        str(record.get("software_profile_id") or ""),
-    )
-
-
-def group_by_comparison_cohort(records: list[Record]) -> dict[CohortKey, list[Record]]:
-    groups: dict[CohortKey, list[Record]] = defaultdict(list)
+def group_by_model_gpu(records: list[Record]) -> dict[tuple[str, str], list[Record]]:
+    groups: dict[tuple[str, str], list[Record]] = defaultdict(list)
     for record in records:
-        groups[comparison_cohort_key(record)].append(record)
+        model_id = str(record.get("model_id") or "unknown")
+        gpu_type = str(record.get("gpu_type") or "unknown")
+        groups[(model_id, gpu_type)].append(record)
     return {key: sorted(value, key=record_sort_key) for key, value in groups.items()}
 
 
-def comparison_sort_key(record: Record) -> CohortKey:
-    return comparison_cohort_key(record)
-
-
-def latest_row_sort_key(row: Record) -> tuple[Any, ...]:
-    return (row["status"] != "fail", *comparison_sort_key(row))
-
-
-def group_identity(key: CohortKey) -> tuple[str, str]:
-    model_id = key[0]
-    gpu_type = key[1]
-    return model_id, gpu_type
-
-
 def baseline_value(records: list[Record], metric_key: str) -> float | None:
     values = [safe_float(record.get(metric_key)) for record in records]
     values = [value for value in values if value is not None]
@@ -137,8 +99,7 @@ def build_latest_summary(records: list[Record],
                          baseline_window: int = 5,
                          run_source: str | None = None) -> list[Record]:
     rows: list[Record] = []
-    for key, group in group_by_comparison_cohort(records).items():
-        model_id, gpu_type = group_identity(key)
+    for (model_id, gpu_type), group in group_by_model_gpu(records).items():
         latest_candidates = group
         if run_source:
             latest_candidates = [record for record in group if record_run_source(record) == run_source]
@@ -195,7 +156,6 @@ def build_latest_summary(records: list[Record],
             "timestamp": latest.get("timestamp"),
             "commit_sha": latest.get("commit_sha"),
             **record_metadata(latest),
-            **record_comparison_metadata(latest),
             "success": success,
             "baseline_n": len(baseline_records),
             "worst_regression_pct": worst_regression,
@@ -206,13 +166,12 @@ def build_latest_summary(records: list[Record],
             "metrics": metrics,
         })
 
-    return sorted(rows, key=latest_row_sort_key)
+    return sorted(rows, key=lambda row: (row["status"] != "fail", row["model_id"], row["gpu_type"]))
 
 
 def build_trends(records: list[Record]) -> list[Record]:
     trends: list[Record] = []
-    for key, group in group_by_comparison_cohort(records).items():
-        model_id, gpu_type = group_identity(key)
+    for (model_id, gpu_type), group in group_by_model_gpu(records).items():
         points = []
         for record in group:
             metric_policies = resolve_metric_policies(record.get("regression_thresholds"))
@@ -220,7 +179,6 @@ def build_trends(records: list[Record]) -> list[Record]:
                 "timestamp": record.get("timestamp"),
                 "commit_sha": record.get("commit_sha"),
                 **record_metadata(record),
-                **record_comparison_metadata(record),
                 "success": bool(record.get("success", True)),
                 "metrics": {
                     policy.key: safe_float(record.get(policy.key))
@@ -231,7 +189,6 @@ def build_trends(records: list[Record]) -> list[Record]:
         trends.append({
             "model_id": model_id,
             "gpu_type": gpu_type,
-            **record_comparison_metadata(group[-1]),
             "points": points,
         })
-    return sorted(trends, key=comparison_sort_key)
+    return sorted(trends, key=lambda trend: (trend["model_id"], trend["gpu_type"]))

fastvideo/pipelines/basic/dreamx_world/__init__.py

@@ -1,28 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-from fastvideo.configs.pipelines.dreamx_world import (
-    DreamXWorld5BARPipelineConfig,
-    DreamXWorld5BCamPipelineConfig,
-    make_dreamx_world_5b_ar_dit_config,
-    make_dreamx_world_5b_cam_dit_config,
-    make_dreamx_world_5b_cam_text_encoder_config,
-    make_dreamx_world_5b_cam_vae_config,
-)
-from fastvideo.pipelines.basic.dreamx_world.dreamx_world_ar_pipeline import DreamXWorldARPipeline
-from fastvideo.pipelines.basic.dreamx_world.dreamx_world_pipeline import DreamXWorldPipeline
-from fastvideo.pipelines.basic.dreamx_world.stages import (
-    DREAMX_Y_CAMERA_KEY,
-    DreamXWorldCameraConditioningStage,
-)
-
-__all__ = [
-    "DREAMX_Y_CAMERA_KEY",
-    "DreamXWorld5BARPipelineConfig",
-    "DreamXWorld5BCamPipelineConfig",
-    "DreamXWorldCameraConditioningStage",
-    "DreamXWorldARPipeline",
-    "DreamXWorldPipeline",
-    "make_dreamx_world_5b_ar_dit_config",
-    "make_dreamx_world_5b_cam_dit_config",
-    "make_dreamx_world_5b_cam_text_encoder_config",
-    "make_dreamx_world_5b_cam_vae_config",
-]

fastvideo/pipelines/basic/dreamx_world/ar_denoising.py

@@ -1,219 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""DreamX-World autoregressive causal denoising stage."""
-
-from __future__ import annotations
-
-from typing import Any
-
-import torch
-from tqdm.auto import tqdm
-
-from fastvideo.fastvideo_args import FastVideoArgs
-from fastvideo.pipelines.basic.dreamx_world.stages import DREAMX_Y_CAMERA_KEY
-from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
-from fastvideo.pipelines.stages.denoising import DenoisingStage
-
-
-class DreamXWorldARCausalDenoisingStage(DenoisingStage):
-    """Official DreamX AR-forcing denoising loop with KV cache."""
-
-    _AR_NOISE_SEED_OFFSET = 1_000_003
-
-    def __init__(self, transformer, scheduler, pipeline=None, vae=None) -> None:
-        super().__init__(transformer=transformer, scheduler=scheduler, pipeline=pipeline, vae=vae)
-        self.num_transformer_blocks = len(self.transformer.blocks)
-        self.num_frame_per_block = int(getattr(self.transformer, "num_frame_per_block", 3))
-        self.local_attn_size = int(getattr(self.transformer, "local_attn_size", 12))
-
-    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
-        assert batch.latents is not None, "latents must be prepared before DreamX AR denoising"
-        assert batch.prompt_embeds, "prompt embeds must be prepared before DreamX AR denoising"
-        latents = batch.latents
-        device = latents.device
-        target_dtype = torch.bfloat16
-        autocast_enabled = device.type == "cuda" and not fastvideo_args.disable_autocast
-
-        frame_seq_length = (latents.shape[-2] // self.transformer.patch_size[1]) * (latents.shape[-1] //
-                                                                                    self.transformer.patch_size[2])
-        timesteps = torch.tensor(
-            tuple(getattr(fastvideo_args.pipeline_config, "dmd_denoising_steps", (1000, 750, 500, 250))),
-            dtype=torch.long,
-        ).cpu()
-        if getattr(fastvideo_args.pipeline_config, "warp_denoising_step", True):
-            self.scheduler.set_timesteps(1000)
-            scheduler_timesteps = torch.cat((self.scheduler.timesteps.cpu(), torch.tensor([0], dtype=torch.float32)))
-            timesteps = scheduler_timesteps[1000 - timesteps]
-        timesteps = timesteps.to(device)
-
-        if latents.shape[2] % self.num_frame_per_block != 0:
-            raise ValueError("DreamX AR latent frames must be divisible by num_frame_per_block")
-
-        y_camera = batch.extra.get(DREAMX_Y_CAMERA_KEY, batch.extra.get("y_camera"))
-        if isinstance(y_camera, dict):
-            y_camera = {
-                k: v.to(device=device, dtype=target_dtype) if torch.is_tensor(v) else v
-                for k, v in y_camera.items()
-            }
-
-        if batch.image_latent is not None and batch.image_latent.shape[1] == latents.shape[1]:
-            latents[:, :, :batch.image_latent.shape[2]] = batch.image_latent.to(device=device, dtype=latents.dtype)
-
-        kv_cache = self._initialize_kv_cache(latents.shape[0], target_dtype, device, frame_seq_length)
-        crossattn_cache = self._initialize_crossattn_cache(latents.shape[0], target_dtype, device)
-        prompt = batch.prompt_embeds[0]
-        if torch.is_tensor(prompt):
-            prompt = prompt.to(device=device, dtype=target_dtype)
-            context = [sample for sample in prompt]
-        else:
-            context = prompt
-
-        num_blocks = latents.shape[2] // self.num_frame_per_block
-        start = 0
-        first_frame_mask = torch.ones_like(latents)
-        first_frame_mask[:, :, 0] = 0
-        base_generator = batch.generator[0] if isinstance(batch.generator, list) else batch.generator
-        noise_generator = self._make_noise_generator(base_generator, device)
-
-        with tqdm(total=num_blocks * len(timesteps), desc="DreamX AR denoising", leave=False) as progress:
-            for _ in range(num_blocks):
-                current_num_frames = self.num_frame_per_block
-                block_latents = latents[:, :, start:start + current_num_frames]
-                noisy_input = block_latents.clone()
-                mask_block = first_frame_mask[:, :, start:start + current_num_frames]
-                camera_block = self._slice_camera(y_camera, start, current_num_frames)
-
-                for idx, current_timestep in enumerate(timesteps):
-                    timestep = torch.full(
-                        (latents.shape[0], current_num_frames * frame_seq_length),
-                        int(current_timestep.item()),
-                        device=device,
-                        dtype=torch.long,
-                    )
-                    if start == 0:
-                        timestep[:, :frame_seq_length] = 0
-                    with torch.autocast(device_type="cuda", dtype=target_dtype, enabled=autocast_enabled):
-                        denoised = self.transformer(
-                            hidden_states=block_latents.to(target_dtype),
-                            encoder_hidden_states=torch.stack(context).to(target_dtype),
-                            timestep=timestep,
-                            y_camera=camera_block,
-                            kv_cache=kv_cache,
-                            crossattn_cache=crossattn_cache,
-                            current_start=start * frame_seq_length,
-                        )
-                    denoised = denoised.to(latents.dtype)
-                    if idx < len(timesteps) - 1:
-                        next_timestep = torch.full((latents.shape[0], current_num_frames),
-                                                   int(timesteps[idx + 1].item()),
-                                                   device=device,
-                                                   dtype=torch.long)
-                        noise_kwargs = {"device": device, "dtype": denoised.dtype}
-                        if noise_generator is not None:
-                            noise_kwargs["generator"] = noise_generator
-                        noise = torch.randn(denoised.permute(0, 2, 1, 3, 4).shape, **noise_kwargs)
-                        block_btchw = self.scheduler.add_noise(
-                            denoised.permute(0, 2, 1, 3, 4).flatten(0, 1),
-                            noise.flatten(0, 1),
-                            next_timestep.flatten(),
-                        ).unflatten(0, (latents.shape[0], current_num_frames))
-                        block_latents = block_btchw.permute(0, 2, 1, 3, 4)
-                        block_latents = block_latents * mask_block + noisy_input * (1 - mask_block)
-                    else:
-                        block_latents = denoised * mask_block + noisy_input * (1 - mask_block)
-                    progress.update()
-
-                latents[:, :, start:start + current_num_frames] = block_latents
-                self._update_context_cache(block_latents, context, camera_block, kv_cache, crossattn_cache, start,
-                                           frame_seq_length, target_dtype, autocast_enabled,
-                                           float(getattr(fastvideo_args.pipeline_config, "context_noise", 0.1)))
-                start += current_num_frames
-
-        batch.latents = latents
-        return batch
-
-    def _make_noise_generator(self, generator: torch.Generator | None, device: torch.device) -> torch.Generator | None:
-        if generator is None:
-            return None
-        if getattr(generator, "device", None) == device:
-            return generator
-        seed = int(generator.initial_seed()) + self._AR_NOISE_SEED_OFFSET
-        return torch.Generator(device=device).manual_seed(seed)
-
-    @staticmethod
-    def _context_noise_timestep(context_noise: float) -> int:
-        if 0.0 < context_noise <= 1.0:
-            return int(context_noise * 1000)
-        return int(context_noise)
-
-    def _slice_camera(self, y_camera: Any, start: int, num_frames: int):
-        if not isinstance(y_camera, dict):
-            return y_camera
-        return {
-            "viewmats": y_camera["viewmats"][:, start:start + num_frames],
-            "K": y_camera["K"][:, start:start + num_frames],
-        }
-
-    def _initialize_kv_cache(self, batch_size: int, dtype: torch.dtype, device: torch.device,
-                             frame_seq_length: int) -> list[dict[str, Any]]:
-        size = self.local_attn_size * frame_seq_length if self.local_attn_size != -1 else 18480
-        heads = self.transformer.num_attention_heads
-        head_dim = self.transformer.attention_head_dim
-        cam_self_attn = next(
-            (getattr(block, "cam_self_attn", None)
-             for block in self.transformer.blocks if getattr(block, "cam_self_attn", None) is not None),
-            None,
-        )
-
-        caches = []
-        for _ in range(self.num_transformer_blocks):
-            cache = {
-                "k": torch.zeros(batch_size, size, heads, head_dim, dtype=dtype, device=device),
-                "v": torch.zeros(batch_size, size, heads, head_dim, dtype=dtype, device=device),
-                "global_end_index": torch.tensor([0], dtype=torch.long, device=device),
-                "local_end_index": torch.tensor([0], dtype=torch.long, device=device),
-            }
-            if cam_self_attn is not None:
-                cam_heads = int(cam_self_attn.num_heads)
-                cam_head_dim = int(cam_self_attn.head_dim)
-                cache.update({
-                    "prope_k":
-                    torch.zeros(batch_size, size, cam_heads, cam_head_dim, dtype=dtype, device=device),
-                    "prope_v":
-                    torch.zeros(batch_size, size, cam_heads, cam_head_dim, dtype=dtype, device=device),
-                    "prope_global_end_index":
-                    torch.tensor([0], dtype=torch.long, device=device),
-                    "prope_local_end_index":
-                    torch.tensor([0], dtype=torch.long, device=device),
-                })
-            caches.append(cache)
-        return caches
-
-    def _initialize_crossattn_cache(self, batch_size: int, dtype: torch.dtype, device: torch.device):
-        heads = self.transformer.num_attention_heads
-        head_dim = self.transformer.attention_head_dim
-        return [{
-            "k": torch.zeros(batch_size, 512, heads, head_dim, dtype=dtype, device=device),
-            "v": torch.zeros(batch_size, 512, heads, head_dim, dtype=dtype, device=device),
-            "is_init": False,
-        } for _ in range(self.num_transformer_blocks)]
-
-    def _update_context_cache(self, block_latents: torch.Tensor, context: Any, camera_block: Any,
-                              kv_cache: list[dict[str, Any]], crossattn_cache: list[dict[str, Any]], start: int,
-                              frame_seq_length: int, target_dtype: torch.dtype, autocast_enabled: bool,
-                              context_noise: float) -> None:
-        timestep = torch.full(
-            (block_latents.shape[0], block_latents.shape[2] * frame_seq_length),
-            self._context_noise_timestep(context_noise),
-            device=block_latents.device,
-            dtype=torch.long,
-        )
-        with torch.autocast(device_type="cuda", dtype=target_dtype, enabled=autocast_enabled):
-            self.transformer(
-                hidden_states=block_latents.to(target_dtype),
-                encoder_hidden_states=torch.stack(context).to(target_dtype),
-                timestep=timestep,
-                y_camera=camera_block,
-                kv_cache=kv_cache,
-                crossattn_cache=crossattn_cache,
-                current_start=start * frame_seq_length,
-            )

fastvideo/pipelines/basic/dreamx_world/camera_conditioning.py

@@ -1,228 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-from __future__ import annotations
-
-import math
-from dataclasses import dataclass
-
-import numpy as np
-import torch
-from scipy.interpolate import interp1d
-from scipy.spatial.transform import Rotation, Slerp
-
-_ACTION_TO_MOTION = {
-    "w": "forward",
-    "a": "left",
-    "d": "right",
-    "s": "backward",
-    "j": "left_rot",
-    "l": "right_rot",
-    "i": "up_rot",
-    "k": "down_rot",
-}
-_TRANSLATION_BASE_UNIT = 1.0
-_ROTATION_BASE_UNIT = 10.0
-_INTRINSIC_ROW = [0.8, 0.5, 0.5, 0.5]
-
-
-@dataclass
-class DreamXCamera:
-    fx: float
-    fy: float
-    cx: float
-    cy: float
-    w2c_mat: np.ndarray
-
-    @property
-    def c2w_mat(self) -> np.ndarray:
-        return np.linalg.inv(self.w2c_mat)
-
-    @classmethod
-    def from_pose_row(cls, row: list[float]) -> DreamXCamera:
-        w2c_mat = np.eye(4, dtype=np.float64)
-        w2c_mat[:3, :] = np.asarray(row[7:], dtype=np.float64).reshape(3, 4)
-        return cls(
-            fx=float(row[1]),
-            fy=float(row[2]),
-            cx=float(row[3]),
-            cy=float(row[4]),
-            w2c_mat=w2c_mat,
-        )
-
-
-def _translation_step(motion_type: str, current_pose: dict[str, np.ndarray], value: float, duration: int) -> np.ndarray:
-    if motion_type in ("forward", "backward"):
-        yaw = np.radians(current_pose["rotation"][1])
-        pitch = np.radians(current_pose["rotation"][0])
-        forward = np.array([-math.sin(yaw) * math.cos(pitch), math.sin(pitch), math.cos(yaw) * math.cos(pitch)])
-        direction = 1 if motion_type == "forward" else -1
-        return forward * value * direction / duration
-    if motion_type in ("left", "right"):
-        yaw = np.radians(current_pose["rotation"][1])
-        right = np.array([math.cos(yaw), 0.0, math.sin(yaw)])
-        direction = -1 if motion_type == "left" else 1
-        return right * value * direction / duration
-    return np.zeros(3)
-
-
-def _rotation_step(motion_type: str, value: float, duration: int) -> np.ndarray:
-    if not motion_type.endswith("rot"):
-        return np.zeros(3)
-    axis = motion_type.split("_")[0]
-    rotation = np.zeros(3)
-    if axis == "left":
-        rotation[1] = value
-    elif axis == "right":
-        rotation[1] = -value
-    elif axis == "up":
-        rotation[0] = -value
-    elif axis == "down":
-        rotation[0] = value
-    return rotation / duration
-
-
-def _euler_to_quaternion(angles: np.ndarray) -> list[float]:
-    pitch, yaw, roll = np.radians(angles)
-    cy = math.cos(yaw * 0.5)
-    sy = math.sin(yaw * 0.5)
-    cp = math.cos(pitch * 0.5)
-    sp = math.sin(pitch * 0.5)
-    cr = math.cos(roll * 0.5)
-    sr = math.sin(roll * 0.5)
-    return [
-        cy * cp * cr + sy * sp * sr,
-        cy * sp * cr + sy * cp * sr,
-        sy * cp * cr - cy * sp * sr,
-        cy * cp * sr - sy * sp * cr,
-    ]
-
-
-def _quaternion_to_rotation_matrix(quaternion: list[float]) -> np.ndarray:
-    qw, qx, qy, qz = quaternion
-    return np.array([
-        [1 - 2 * (qy**2 + qz**2), 2 * (qx * qy - qw * qz), 2 * (qx * qz + qw * qy)],
-        [2 * (qx * qy + qw * qz), 1 - 2 * (qx**2 + qz**2), 2 * (qy * qz - qw * qx)],
-        [2 * (qx * qz - qw * qy), 2 * (qy * qz + qw * qx), 1 - 2 * (qx**2 + qy**2)],
-    ])
-
-
-def _pose_rows_from_actions(action_seq: list[str], action_speed_list: list[float], duration: int) -> list[list[float]]:
-    if len(action_seq) != len(action_speed_list):
-        raise ValueError("action_seq and action_speed_list must have the same length")
-
-    positions: list[np.ndarray] = []
-    rotations: list[np.ndarray] = []
-    current_pose = {
-        "position": np.array([0.0, 0.0, 0.0]),
-        "rotation": np.array([0.0, 0.0, 0.0]),
-    }
-
-    for action_id, speed in zip(action_seq, action_speed_list, strict=True):
-        motion_types = [_ACTION_TO_MOTION[key] for key in list(action_id)]
-        translation_step = np.zeros(3)
-        rotation_step = np.zeros(3)
-        for motion_type in motion_types:
-            translation_step += _translation_step(motion_type, current_pose,
-                                                  float(speed) * _TRANSLATION_BASE_UNIT, duration)
-            rotation_step += _rotation_step(motion_type, float(speed) * _ROTATION_BASE_UNIT, duration)
-
-        segment_positions = []
-        segment_rotations = []
-        for index in range(1, duration + 1):
-            segment_positions.append(current_pose["position"] + translation_step * index)
-            segment_rotations.append(current_pose["rotation"] + rotation_step * index)
-        current_pose["position"] = segment_positions[-1].copy()
-        current_pose["rotation"] = segment_rotations[-1].copy()
-        positions.extend(segment_positions)
-        rotations.extend(segment_rotations)
-
-    rows: list[list[float]] = [[0.0] + _INTRINSIC_ROW + [0.0, 0.0] +
-                               [1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0]]
-    for index, (position, rotation) in enumerate(zip(positions, rotations, strict=False)):
-        rotation_matrix = _quaternion_to_rotation_matrix(_euler_to_quaternion(rotation))
-        translation = -rotation_matrix @ position
-        extrinsic = np.hstack([rotation_matrix, translation.reshape(3, 1)])
-        rows.append([float(index)] + _INTRINSIC_ROW + [0.0, 0.0] + extrinsic.flatten().tolist())
-    return rows
-
-
-def _interpolate_camera_poses(
-    cameras: list[DreamXCamera],
-    src_indices: np.ndarray,
-    tgt_indices: np.ndarray,
-) -> list[DreamXCamera]:
-    if len(cameras) <= 1:
-        return [cameras[0]] * len(tgt_indices) if cameras else []
-    src_rot_mat = np.array([camera.w2c_mat[:3, :3] for camera in cameras])
-    src_trans_vec = np.array([camera.w2c_mat[:3, 3] for camera in cameras])
-
-    dets = np.linalg.det(src_rot_mat)
-    flip_handedness = dets.size > 0 and np.median(dets) < 0.0
-    if flip_handedness:
-        flip_mat = np.diag([1.0, 1.0, -1.0]).astype(src_rot_mat.dtype)
-        src_rot_mat = src_rot_mat @ flip_mat
-
-    trans = interp1d(src_indices, src_trans_vec, axis=0, kind="linear", bounds_error=False,
-                     fill_value="extrapolate")(tgt_indices)
-    quats = Rotation.from_matrix(src_rot_mat).as_quat().copy()
-    for index in range(1, len(quats)):
-        if np.dot(quats[index], quats[index - 1]) < 0:
-            quats[index] = -quats[index]
-    rot = Slerp(src_indices, Rotation.from_quat(quats))(tgt_indices).as_matrix()
-    if flip_handedness:
-        rot = rot @ flip_mat
-
-    ref = cameras[0]
-    result = []
-    for index in range(len(tgt_indices)):
-        w2c_mat = np.eye(4, dtype=np.float64)
-        w2c_mat[:3, :] = np.hstack([rot[index], trans[index].reshape(3, 1)])
-        result.append(DreamXCamera(ref.fx, ref.fy, ref.cx, ref.cy, w2c_mat))
-    return result
-
-
-def _relative_c2w_poses(cameras: list[DreamXCamera]) -> np.ndarray:
-    abs_w2cs = [camera.w2c_mat for camera in cameras]
-    abs_c2ws = [camera.c2w_mat for camera in cameras]
-    target_cam_c2w = np.eye(4, dtype=np.float64)
-    abs2rel = target_cam_c2w @ abs_w2cs[0]
-    poses = [target_cam_c2w] + [abs2rel @ c2w for c2w in abs_c2ws[1:]]
-    return np.asarray(poses, dtype=np.float32)
-
-
-def _invert_se3(transforms: torch.Tensor) -> torch.Tensor:
-    rotation_inv = transforms[..., :3, :3].transpose(-1, -2)
-    output = torch.zeros_like(transforms)
-    output[..., :3, :3] = rotation_inv
-    output[..., :3, 3] = -torch.einsum("...ij,...j->...i", rotation_inv, transforms[..., :3, 3])
-    output[..., 3, 3] = 1.0
-    return output
-
-
-def build_dreamx_camera_condition(
-    action_seq: list[str],
-    action_speed_list: list[float],
-    *,
-    num_frames: int,
-    height: int,
-    width: int,
-    dtype: torch.dtype = torch.float32,
-    device: torch.device | str = "cpu",
-) -> dict[str, torch.Tensor]:
-    del height, width  # DreamX-World-5B-Cam uses fixed normalized intrinsics.
-    duration = math.ceil(num_frames / len(action_seq))
-    rows = _pose_rows_from_actions(action_seq, action_speed_list, duration)[:num_frames]
-    cameras = [DreamXCamera.from_pose_row(row) for row in rows]
-
-    latent_frame_count = 1 + (len(cameras) - 1) // 4
-    src_indices = np.arange(len(cameras), dtype=np.float64)
-    tgt_indices = np.linspace(0, len(cameras) - 1, latent_frame_count)
-    cameras = _interpolate_camera_poses(cameras, src_indices, tgt_indices)
-
-    c2ws = torch.as_tensor(_relative_c2w_poses(cameras), dtype=dtype, device=device)
-    viewmats = _invert_se3(c2ws)
-
-    intrinsics = torch.zeros((latent_frame_count, 3, 3), dtype=dtype, device=device)
-    intrinsics[:, 0, 0] = 969.6969696969696 / (960.0 * 2)
-    intrinsics[:, 1, 1] = 969.6969696969696 / (540.0 * 2)
-    intrinsics[:, 2, 2] = 1.0
-    return {"viewmats": viewmats, "K": intrinsics}

fastvideo/pipelines/basic/dreamx_world/config.py

@@ -1,20 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Compatibility exports for DreamX-World pipeline configs."""
-
-from fastvideo.configs.pipelines.dreamx_world import (
-    DreamXWorld5BARPipelineConfig,
-    DreamXWorld5BCamPipelineConfig,
-    make_dreamx_world_5b_ar_dit_config,
-    make_dreamx_world_5b_cam_dit_config,
-    make_dreamx_world_5b_cam_text_encoder_config,
-    make_dreamx_world_5b_cam_vae_config,
-)
-
-__all__ = [
-    "DreamXWorld5BARPipelineConfig",
-    "DreamXWorld5BCamPipelineConfig",
-    "make_dreamx_world_5b_ar_dit_config",
-    "make_dreamx_world_5b_cam_dit_config",
-    "make_dreamx_world_5b_cam_text_encoder_config",
-    "make_dreamx_world_5b_cam_vae_config",
-]

fastvideo/pipelines/basic/dreamx_world/dreamx_world_ar_pipeline.py

@@ -1,67 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""DreamX-World-5B autoregressive pipeline entrypoint."""
-
-from fastvideo.api.sampling_param import SamplingParam
-from fastvideo.configs.pipelines.dreamx_world import DreamXWorld5BARPipelineConfig
-from fastvideo.fastvideo_args import FastVideoArgs
-from fastvideo.logger import init_logger
-from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler
-from fastvideo.pipelines import ComposedPipelineBase, LoRAPipeline
-from fastvideo.pipelines.basic.dreamx_world.ar_denoising import DreamXWorldARCausalDenoisingStage
-from fastvideo.pipelines.basic.dreamx_world.stages import (
-    DreamXWorldCameraConditioningStage,
-    DreamXWorldImageVAEEncodingStage,
-)
-from fastvideo.pipelines.stages import (
-    ConditioningStage,
-    DecodingStage,
-    InputValidationStage,
-    LatentPreparationStage,
-    TextEncodingStage,
-    TimestepPreparationStage,
-)
-
-logger = init_logger(__name__)
-
-
-class DreamXWorldARPipeline(LoRAPipeline, ComposedPipelineBase):
-    """DreamX-World-5B autoregressive causal camera pipeline."""
-
-    _required_config_modules = ["text_encoder", "tokenizer", "vae", "transformer", "scheduler"]
-    pipeline_config_cls = DreamXWorld5BARPipelineConfig
-    sampling_params_cls = SamplingParam
-
-    def initialize_pipeline(self, fastvideo_args: FastVideoArgs):
-        self.modules["scheduler"] = FlowMatchEulerDiscreteScheduler(shift=fastvideo_args.pipeline_config.flow_shift)
-        self.modules["scheduler"].set_timesteps(1000)
-
-    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None:
-        self.add_stage(stage_name="input_validation_stage", stage=InputValidationStage())
-        self.add_stage(stage_name="prompt_encoding_stage",
-                       stage=TextEncodingStage(
-                           text_encoders=[self.get_module("text_encoder")],
-                           tokenizers=[self.get_module("tokenizer")],
-                       ))
-        self.add_stage(stage_name="conditioning_stage", stage=ConditioningStage())
-        self.add_stage(stage_name="timestep_preparation_stage",
-                       stage=TimestepPreparationStage(scheduler=self.get_module("scheduler")))
-        self.add_stage(stage_name="latent_preparation_stage",
-                       stage=LatentPreparationStage(
-                           scheduler=self.get_module("scheduler"),
-                           transformer=self.get_module("transformer", None),
-                       ))
-        self.add_stage(stage_name="image_latent_preparation_stage",
-                       stage=DreamXWorldImageVAEEncodingStage(vae=self.get_module("vae")))
-        self.add_stage(stage_name="dreamx_camera_conditioning_stage", stage=DreamXWorldCameraConditioningStage())
-        self.add_stage(stage_name="denoising_stage",
-                       stage=DreamXWorldARCausalDenoisingStage(
-                           transformer=self.get_module("transformer"),
-                           scheduler=self.get_module("scheduler"),
-                           vae=self.get_module("vae"),
-                           pipeline=self,
-                       ))
-        self.add_stage(stage_name="decoding_stage", stage=DecodingStage(vae=self.get_module("vae"), pipeline=self))
-        logger.info("DreamXWorldARPipeline initialized with autoregressive causal denoising")
-
-
-EntryClass = DreamXWorldARPipeline

fastvideo/pipelines/basic/dreamx_world/dreamx_world_pipeline.py

@@ -1,78 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""DreamX-World video pipeline entrypoint."""
-
-from fastvideo.api.sampling_param import SamplingParam
-from fastvideo.fastvideo_args import FastVideoArgs
-from fastvideo.logger import init_logger
-from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler
-from fastvideo.pipelines import ComposedPipelineBase, LoRAPipeline
-from fastvideo.configs.pipelines.dreamx_world import DreamXWorld5BCamPipelineConfig
-from fastvideo.pipelines.basic.dreamx_world.stages import DreamXWorldCameraConditioningStage
-from fastvideo.pipelines.stages import (
-    ConditioningStage,
-    DecodingStage,
-    DenoisingStage,
-    InputValidationStage,
-    LatentPreparationStage,
-    TextEncodingStage,
-    TimestepPreparationStage,
-)
-
-logger = init_logger(__name__)
-
-
-class DreamXWorldPipeline(LoRAPipeline, ComposedPipelineBase):
-    """DreamX-World-5B-Cam pipeline with native FastVideo camera conditioning."""
-
-    _required_config_modules = ["text_encoder", "tokenizer", "vae", "transformer", "scheduler"]
-    pipeline_config_cls = DreamXWorld5BCamPipelineConfig
-    sampling_params_cls = SamplingParam
-
-    def initialize_pipeline(self, fastvideo_args: FastVideoArgs):
-        self.modules["scheduler"] = FlowMatchEulerDiscreteScheduler(shift=fastvideo_args.pipeline_config.flow_shift)
-
-    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None:
-        self.add_stage(stage_name="input_validation_stage", stage=InputValidationStage())
-
-        self.add_stage(
-            stage_name="prompt_encoding_stage",
-            stage=TextEncodingStage(
-                text_encoders=[self.get_module("text_encoder")],
-                tokenizers=[self.get_module("tokenizer")],
-            ),
-        )
-
-        self.add_stage(stage_name="conditioning_stage", stage=ConditioningStage())
-
-        self.add_stage(
-            stage_name="timestep_preparation_stage",
-            stage=TimestepPreparationStage(scheduler=self.get_module("scheduler")),
-        )
-
-        self.add_stage(
-            stage_name="latent_preparation_stage",
-            stage=LatentPreparationStage(
-                scheduler=self.get_module("scheduler"),
-                transformer=self.get_module("transformer", None),
-            ),
-        )
-
-        self.add_stage(stage_name="dreamx_camera_conditioning_stage", stage=DreamXWorldCameraConditioningStage())
-
-        self.add_stage(
-            stage_name="denoising_stage",
-            stage=DenoisingStage(
-                transformer=self.get_module("transformer"),
-                transformer_2=self.get_module("transformer_2", None),
-                scheduler=self.get_module("scheduler"),
-                vae=self.get_module("vae"),
-                pipeline=self,
-            ),
-        )
-
-        self.add_stage(stage_name="decoding_stage", stage=DecodingStage(vae=self.get_module("vae"), pipeline=self))
-
-        logger.info("DreamXWorldPipeline initialized with native camera conditioning")
-
-
-EntryClass = DreamXWorldPipeline

fastvideo/pipelines/basic/dreamx_world/presets.py

@@ -1,58 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""DreamX-World model family pipeline presets."""
-
-from fastvideo.api.presets import InferencePreset, PresetStageSpec
-
-_NEGATIVE_PROMPT_CN = ("色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，"
-                       "静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，"
-                       "多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，"
-                       "形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，"
-                       "背景人很多，倒着走")
-
-_DENOISE_STAGE = PresetStageSpec(
-    name="denoise",
-    kind="denoising",
-    description="DreamX-World camera-conditioned denoising pass",
-    allowed_overrides=frozenset({
-        "num_inference_steps",
-        "guidance_scale",
-    }),
-)
-
-DREAMX_WORLD_5B_CAM = InferencePreset(
-    name="dreamx_world_5b_cam",
-    version=1,
-    model_family="dreamx_world",
-    description="DreamX-World 5B camera-control video generation",
-    workload_type="i2v",
-    stage_schemas=(_DENOISE_STAGE, ),
-    defaults={
-        "height": 480,
-        "width": 832,
-        "num_frames": 161,
-        "fps": 16,
-        "guidance_scale": 5.0,
-        "num_inference_steps": 30,
-        "negative_prompt": _NEGATIVE_PROMPT_CN,
-    },
-)
-
-DREAMX_WORLD_5B_AR = InferencePreset(
-    name="dreamx_world_5b_ar",
-    version=1,
-    model_family="dreamx_world",
-    description="DreamX-World 5B autoregressive camera-control generation",
-    workload_type="i2v",
-    stage_schemas=(_DENOISE_STAGE, ),
-    defaults={
-        "height": 704,
-        "width": 1280,
-        "num_frames": 1005,
-        "fps": 16,
-        "guidance_scale": 1.0,
-        "num_inference_steps": 4,
-        "negative_prompt": _NEGATIVE_PROMPT_CN,
-    },
-)
-
-ALL_PRESETS = (DREAMX_WORLD_5B_CAM, DREAMX_WORLD_5B_AR)

fastvideo/pipelines/basic/dreamx_world/stages.py

@@ -1,156 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""DreamX-World pipeline stages."""
-
-from __future__ import annotations
-
-import torch
-
-from fastvideo.fastvideo_args import FastVideoArgs
-from fastvideo.logger import init_logger
-from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
-from fastvideo.pipelines.stages.base import PipelineStage
-from fastvideo.pipelines.stages.validators import VerificationResult
-
-from fastvideo.pipelines.basic.dreamx_world.camera_conditioning import (
-    build_dreamx_camera_condition, )
-
-DREAMX_Y_CAMERA_KEY = "dreamx_y_camera"
-
-logger = init_logger(__name__)
-
-
-class DreamXWorldCameraConditioningStage(PipelineStage):
-    """Build PRoPE camera conditioning for DreamX-World-5B-Cam."""
-
-    def forward(
-        self,
-        batch: ForwardBatch,
-        fastvideo_args: FastVideoArgs,
-    ) -> ForwardBatch:
-        del fastvideo_args
-        if DREAMX_Y_CAMERA_KEY in batch.extra:
-            return batch
-
-        action_seq = batch.extra.get("dreamx_action_seq", batch.action_list)
-        action_speed_list = batch.extra.get("dreamx_action_speed_list", batch.action_speed_list)
-        if action_seq is None:
-            action_seq = ["w"]
-        if action_speed_list is None:
-            action_speed_list = [4]
-
-        if isinstance(action_seq, str):
-            action_seq = [action_seq]
-        if isinstance(action_speed_list, int | float):
-            action_speed_list = [action_speed_list]
-        if len(action_speed_list) == 1 and len(action_seq) > 1:
-            action_speed_list = list(action_speed_list) * len(action_seq)
-        action_speed_list = [float(speed) for speed in action_speed_list]
-
-        height = int(batch.height) if batch.height is not None else 704
-        width = int(batch.width) if batch.width is not None else 1280
-        num_frames = int(batch.num_frames)
-        dtype = batch.latents.dtype if torch.is_tensor(batch.latents) else torch.float32
-        device = batch.latents.device if torch.is_tensor(batch.latents) else "cpu"
-
-        y_camera = build_dreamx_camera_condition(
-            list(action_seq),
-            action_speed_list,
-            num_frames=num_frames,
-            height=height,
-            width=width,
-            dtype=dtype,
-            device=device,
-        )
-        batch.extra[DREAMX_Y_CAMERA_KEY] = {key: value.unsqueeze(0) for key, value in y_camera.items()}
-        return batch
-
-    def verify_output(
-        self,
-        batch: ForwardBatch,
-        fastvideo_args: FastVideoArgs,
-    ) -> VerificationResult:
-        del fastvideo_args
-        result = VerificationResult()
-        y_camera = batch.extra.get(DREAMX_Y_CAMERA_KEY)
-        result.add_check("dreamx_y_camera", y_camera, lambda value: isinstance(value, dict))
-        if isinstance(y_camera, dict):
-            result.add_check("dreamx_y_camera.viewmats", y_camera.get("viewmats"), torch.is_tensor)
-            result.add_check("dreamx_y_camera.K", y_camera.get("K"), torch.is_tensor)
-        return result
-
-
-class DreamXWorldImageVAEEncodingStage(PipelineStage):
-    """Encode the conditioning image into the first-frame latent.
-
-    Official AR-forcing flow (AMAP-ML/DreamX-World inference_ar_forcing.py):
-    the input image is resized, normalized to [-1, 1], VAE-encoded
-    deterministically, and written into frame 0 of the noise — the causal
-    denoiser then treats frame 0 as clean context. This stage produces
-    ``batch.image_latent`` ([B, C, 1, H_lat, W_lat]); the injection into
-    the latents happens in DreamXWorldARCausalDenoisingStage.
-    """
-
-    def __init__(self, vae) -> None:
-        super().__init__()
-        self.vae = vae
-
-    def forward(
-        self,
-        batch: ForwardBatch,
-        fastvideo_args: FastVideoArgs,
-    ) -> ForwardBatch:
-        if batch.pil_image is None:
-            # No conditioning image: the causal denoiser falls back to
-            # running from pure noise (frame 0 uninitialized). Warn loudly —
-            # this pipeline is registered I2V and the official flow always
-            # forces from a frame.
-            logger.warning("DreamXWorldARPipeline called without an input image; "
-                           "first-frame context will be noise (T2V-style). Pass an "
-                           "image for the official AR-forcing behavior.")
-            return batch
-
-        from fastvideo.platforms import get_local_torch_device
-        from fastvideo.utils import PRECISION_TO_TYPE
-
-        device = get_local_torch_device()
-        image = batch.pil_image
-        if not isinstance(image, torch.Tensor):
-            import numpy as np
-            import PIL.Image
-            assert isinstance(image, PIL.Image.Image)
-            width = batch.width if isinstance(batch.width, int) else batch.width[0]
-            height = batch.height if isinstance(batch.height, int) else batch.height[0]
-            image = image.convert("RGB").resize((width, height), PIL.Image.Resampling.LANCZOS)
-            arr = torch.from_numpy(np.asarray(image)).float().permute(2, 0, 1) / 255.0
-            image = (arr - 0.5) / 0.5  # official Normalize([0.5], [0.5])
-            image = image.unsqueeze(0)  # [1, C, H, W]
-        if image.dim() == 4:
-            image = image.unsqueeze(2)  # [B, C, 1, H, W]
-        elif image.dim() == 5:
-            image = image[:, :, :1]
-        image = image.to(device=device, dtype=torch.float32)
-
-        vae_dtype = PRECISION_TO_TYPE[fastvideo_args.pipeline_config.vae_precision]
-        vae_autocast_enabled = (vae_dtype != torch.float32) and not fastvideo_args.disable_autocast
-        self.vae = self.vae.to(device)
-        with torch.autocast(device_type="cuda", dtype=vae_dtype, enabled=vae_autocast_enabled):
-            if not vae_autocast_enabled:
-                image = image.to(vae_dtype)
-            encoder_output = self.vae.encode(image)
-
-        # Official encode_to_latent is deterministic ((mean - mu) / sigma per
-        # channel); the posterior mean + shift/scale is the FastVideo
-        # equivalent of that normalization.
-        latent = encoder_output.mean
-        if getattr(self.vae, "shift_factor", None) is not None:
-            shift = self.vae.shift_factor
-            latent = latent - (shift.to(latent.device, latent.dtype) if isinstance(shift, torch.Tensor) else shift)
-        scale = self.vae.scaling_factor
-        latent = latent * (scale.to(latent.device, latent.dtype) if isinstance(scale, torch.Tensor) else scale)
-
-        batch.image_latent = latent
-        return batch
-
-    def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult:
-        result = VerificationResult()
-        return result

fastvideo/pipelines/stages/base.py

@@ -34,7 +34,6 @@ class PipelineStage(ABC):
     composed with other stages to create a complete pipeline. Each stage is responsible
     for a specific part of the process, such as prompt encoding, latent preparation, etc.
     """
-    performance_component_metric: str | None = None
 
     def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult:
         """
@@ -156,9 +155,6 @@ class PipelineStage(ABC):
                 logger.info("[%s] Execution completed in %s ms", stage_name, execution_time * 1000)
                 batch.logging_info.add_stage_execution_time(stage_key, execution_time)
                 batch.logging_info.add_stage_metric(stage_key, "stage_class", stage_class_name)
-                component_metric = self.performance_component_metric
-                if component_metric is not None:
-                    batch.logging_info.add_stage_metric(stage_key, "component_metric", component_metric)
             except Exception as e:
                 torch.cuda.synchronize()
                 execution_time = time.perf_counter() - start_time

fastvideo/pipelines/stages/decoding.py

@@ -28,7 +28,6 @@ class DecodingStage(PipelineStage):
     This stage handles the decoding of latent representations into the final
     output format (e.g., pixel values).
     """
-    performance_component_metric = "vae_decode_time_s"
 
     def __init__(self, vae, pipeline=None) -> None:
         self.vae: ParallelTiledVAE = vae

fastvideo/pipelines/stages/denoising.py

@@ -51,7 +51,6 @@ class DenoisingStage(PipelineStage):
     This stage handles the iterative denoising process that transforms
     the initial noise into the final output.
     """
-    performance_component_metric = "dit_time_s"
 
     def __init__(self, transformer, scheduler, pipeline=None, transformer_2=None, vae=None) -> None:
         super().__init__()
@@ -191,19 +190,6 @@ class DenoisingStage(PipelineStage):
             },
         )
 
-        dreamx_y_camera = batch.extra.get("dreamx_y_camera", batch.extra.get("y_camera"))
-        if isinstance(dreamx_y_camera, dict):
-            dreamx_y_camera = {
-                key: value.to(device=local_device, dtype=target_dtype) if torch.is_tensor(value) else value
-                for key, value in dreamx_y_camera.items()
-            }
-        dreamx_camera_kwargs = self.prepare_extra_func_kwargs(
-            self.transformer.forward,
-            {
-                "y_camera": dreamx_y_camera,
-            },
-        )
-
         for key in ("flux2_txt_ids", "flux2_img_ids"):
             value = batch.extra.get(key)
             if torch.is_tensor(value):
@@ -248,18 +234,13 @@ class DenoisingStage(PipelineStage):
 
         boundary_timestep = boundary_ratio * self.scheduler.num_train_timesteps if boundary_ratio is not None else None
         latent_model_input = latents.to(target_dtype)
-        assert latent_model_input.shape[0] == 1, "only support batch size 1"
 
         if fastvideo_args.pipeline_config.ti2v_task and batch.pil_image is not None:
             # TI2V directly replaces the first frame of the latent with
             # the image latent instead of appending along the channel dim
             assert batch.image_latent is None, "TI2V task should not have image latents"
             assert self.vae is not None, "VAE is not provided for TI2V task"
-            vae_device = next(self.vae.parameters()).device
-            self.vae = self.vae.to(local_device)
             z = self.vae.encode(batch.pil_image).mean.float()
-            if getattr(fastvideo_args, "vae_cpu_offload", False):
-                self.vae = self.vae.to(vae_device)
             if (hasattr(self.vae, "shift_factor") and self.vae.shift_factor is not None):
                 if isinstance(self.vae.shift_factor, torch.Tensor):
                     z -= self.vae.shift_factor.to(z.device, z.dtype)
@@ -512,7 +493,6 @@ class DenoisingStage(PipelineStage):
                             **pos_cond_kwargs,
                             **action_kwargs,
                             **camera_kwargs,
-                            **dreamx_camera_kwargs,
                             **timesteps_r_kwarg,
                             **flux2_id_kwargs,
                         )
@@ -555,7 +535,6 @@ class DenoisingStage(PipelineStage):
                                     **neg_cond_kwargs,
                                     **action_kwargs,
                                     **camera_kwargs,
-                                    **dreamx_camera_kwargs,
                                     **timesteps_r_kwarg,
                                     **flux2_id_kwargs,
                                 )
@@ -1207,7 +1186,6 @@ class Cosmos25V2WDenoisingStage(Cosmos25DenoisingStage):
 
 class Cosmos25AutoDenoisingStage(PipelineStage):
     """Route Cosmos 2.5 denoising to T2W vs V2W/I2W."""
-    performance_component_metric = "dit_time_s"
 
     def __init__(self, transformer, scheduler) -> None:
         super().__init__()

fastvideo/pipelines/stages/input_validation.py

@@ -35,7 +35,11 @@ class InputValidationStage(PipelineStage):
         num_videos_per_prompt = batch.num_videos_per_prompt
 
         assert seed is not None
-        seeds = [seed + i for i in range(num_videos_per_prompt)]
+        if batch.seeds is not None:
+            seeds = batch.seeds
+        else:
+            prompt_count = len(batch.prompt) if isinstance(batch.prompt, list) else 1
+            seeds = [seed + i for i in range(prompt_count * num_videos_per_prompt)]
         batch.seeds = seeds
 
         # Peiyuan: using GPU seed will cause A100 and H100 to generate different results...

fastvideo/pipelines/stages/text_encoding.py

@@ -24,7 +24,6 @@ class TextEncodingStage(PipelineStage):
     This stage handles the encoding of text prompts into the embedding space
     expected by the diffusion model.
     """
-    performance_component_metric = "text_encoder_time_s"
 
     def __init__(self, text_encoders, tokenizers) -> None:
         """
@@ -66,12 +65,20 @@ class TextEncodingStage(PipelineStage):
         assert batch.prompt is not None
         prompt_text: str | list[str] = batch.prompt
         all_indices: list[int] = list(range(len(self.text_encoders)))
-        prompt_embeds_list, prompt_masks_list = self.encode_text(
-            prompt_text,
-            fastvideo_args,
-            encoder_index=all_indices,
-            return_attention_mask=True,
-        )
+        if isinstance(prompt_text, list):
+            prompt_embeds_list, prompt_masks_list = self._encode_prompt_list_individually(
+                prompt_text,
+                fastvideo_args,
+                encoder_index=all_indices,
+                return_attention_mask=True,
+            )
+        else:
+            prompt_embeds_list, prompt_masks_list = self.encode_text(
+                prompt_text,
+                fastvideo_args,
+                encoder_index=all_indices,
+                return_attention_mask=True,
+            )
         if self._last_audio_embeds is not None:
             batch.extra["ltx2_audio_prompt_embeds"] = self._last_audio_embeds
 
@@ -83,13 +90,24 @@ class TextEncodingStage(PipelineStage):
 
         # Encode negative prompt if CFG is enabled
         if batch.do_classifier_free_guidance:
-            assert isinstance(batch.negative_prompt, str)
-            neg_embeds_list, neg_masks_list = self.encode_text(
-                batch.negative_prompt,
-                fastvideo_args,
-                encoder_index=all_indices,
-                return_attention_mask=True,
-            )
+            assert isinstance(batch.negative_prompt, str | list)
+            negative_prompt: str | list[str] = batch.negative_prompt
+            if isinstance(batch.prompt, list) and isinstance(negative_prompt, str):
+                negative_prompt = [negative_prompt] * len(batch.prompt)
+            if isinstance(negative_prompt, list):
+                neg_embeds_list, neg_masks_list = self._encode_prompt_list_individually(
+                    negative_prompt,
+                    fastvideo_args,
+                    encoder_index=all_indices,
+                    return_attention_mask=True,
+                )
+            else:
+                neg_embeds_list, neg_masks_list = self.encode_text(
+                    negative_prompt,
+                    fastvideo_args,
+                    encoder_index=all_indices,
+                    return_attention_mask=True,
+                )
             if self._last_audio_embeds is not None:
                 batch.extra["ltx2_audio_negative_embeds"] = self._last_audio_embeds
 
@@ -102,6 +120,81 @@ class TextEncodingStage(PipelineStage):
 
         return batch
 
+    def _encode_prompt_list_individually(
+        self,
+        texts: list[str],
+        fastvideo_args: FastVideoArgs,
+        *,
+        encoder_index: list[int],
+        return_attention_mask: bool,
+    ) -> tuple[list[torch.Tensor], list[torch.Tensor]]:
+        per_prompt_embeds: list[list[torch.Tensor]] = []
+        per_prompt_masks: list[list[torch.Tensor]] = []
+        per_prompt_audio_embeds: list[list[torch.Tensor] | None] = []
+
+        for text in texts:
+            embeds, masks = self.encode_text(
+                text,
+                fastvideo_args,
+                encoder_index=encoder_index,
+                return_attention_mask=return_attention_mask,
+            )
+            per_prompt_embeds.append(embeds)
+            per_prompt_masks.append(masks)
+            per_prompt_audio_embeds.append(self._last_audio_embeds)
+
+        merged_embeds = [
+            self._cat_tensors([prompt_embeds[encoder_pos] for prompt_embeds in per_prompt_embeds])
+            for encoder_pos in range(len(per_prompt_embeds[0]))
+        ]
+        merged_masks = [
+            self._cat_attention_masks([prompt_masks[encoder_pos] for prompt_masks in per_prompt_masks])
+            for encoder_pos in range(len(per_prompt_masks[0]))
+        ]
+        if per_prompt_audio_embeds and all(audio_embeds is not None for audio_embeds in per_prompt_audio_embeds):
+            audio_embed_lists = [audio_embeds for audio_embeds in per_prompt_audio_embeds if audio_embeds is not None]
+            self._last_audio_embeds = [
+                self._cat_tensors([audio_embeds[encoder_pos] for audio_embeds in audio_embed_lists])
+                for encoder_pos in range(len(audio_embed_lists[0]))
+            ]
+        else:
+            self._last_audio_embeds = None
+        return merged_embeds, merged_masks
+
+    @staticmethod
+    def _cat_tensors(tensors: list[torch.Tensor]) -> torch.Tensor:
+        base_shape = tensors[0].shape[1:]
+        if all(tensor.shape[1:] == base_shape for tensor in tensors):
+            return torch.cat(tensors, dim=0)
+        if all(tensor.ndim == 3 for tensor in tensors):
+            base_trailing_shape = tensors[0].shape[2:]
+            if all(tensor.shape[2:] == base_trailing_shape for tensor in tensors):
+                max_length = max(tensor.shape[1] for tensor in tensors)
+                padded_tensors = []
+                for tensor in tensors:
+                    pad_width = max_length - tensor.shape[1]
+                    if pad_width > 0:
+                        tensor = torch.nn.functional.pad(tensor, (0, 0, 0, pad_width), value=0.0)
+                    padded_tensors.append(tensor)
+                return torch.cat(padded_tensors, dim=0)
+        raise ValueError(f"Cannot concatenate tensors with shapes: {[list(tensor.shape) for tensor in tensors]}")
+
+    @staticmethod
+    def _cat_attention_masks(masks: list[torch.Tensor]) -> torch.Tensor:
+        base_shape = masks[0].shape[1:]
+        if all(mask.shape[1:] == base_shape for mask in masks):
+            return torch.cat(masks, dim=0)
+        if all(mask.ndim == 2 for mask in masks):
+            max_length = max(mask.shape[1] for mask in masks)
+            padded_masks = []
+            for mask in masks:
+                pad_width = max_length - mask.shape[1]
+                if pad_width > 0:
+                    mask = torch.nn.functional.pad(mask, (0, pad_width), value=0)
+                padded_masks.append(mask)
+            return torch.cat(padded_masks, dim=0)
+        raise ValueError(f"Cannot concatenate attention masks with shapes: {[list(mask.shape) for mask in masks]}")
+
     def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult:
         """Verify text encoding stage inputs."""
         result = VerificationResult()
@@ -236,6 +329,9 @@ class TextEncodingStage(PipelineStage):
                     attn_masks_list.append(attention_mask)
                     return self.return_embeds(embeds_list, attn_masks_list, return_type, return_attention_mask, indices)
 
+            if len(processed_texts) > 1 and "padding" not in tok_kwargs:
+                tok_kwargs["padding"] = True
+
             # If tokenizer is a multimodal processor (e.g. Qwen2_5_VLProcessor),
             # use its inner tokenizer for text-only encoding.
             tok = getattr(tokenizer, "tokenizer", tokenizer)
@@ -351,7 +447,6 @@ class Cosmos25TextEncodingStage(PipelineStage):
     Cosmos 2.5 uses Reason1 (Qwen2.5-VL) and relies on the encoder's
     `compute_text_embeddings_online()`.
     """
-    performance_component_metric = "text_encoder_time_s"
 
     def __init__(self, text_encoder) -> None:
         super().__init__()

fastvideo/registry.py

@@ -20,7 +20,6 @@ from fastvideo.configs.pipelines.cosmos2_5 import (
     Cosmos25Config,
     Cosmos25_14BConfig,
 )
-from fastvideo.configs.pipelines.dreamx_world import DreamXWorld5BARPipelineConfig, DreamXWorld5BCamPipelineConfig
 from fastvideo.configs.pipelines.hunyuan import FastHunyuanConfig, HunyuanConfig
 from fastvideo.configs.pipelines.hunyuangamecraft import HunyuanGameCraftPipelineConfig
 from fastvideo.configs.pipelines.gen3c import Gen3CConfig
@@ -774,40 +773,6 @@ def _register_configs() -> None:
         model_family="wan",
         default_preset="wan_2_2_ti2v_5b",
     )
-
-    register_configs(
-        sampling_param_cls=None,
-        pipeline_config_cls=DreamXWorld5BCamPipelineConfig,
-        workload_types=(WorkloadType.I2V, ),
-        hf_model_paths=[
-            "FastVideo/DreamX-World-5B-Cam-Diffusers",
-        ],
-        model_detectors=[
-            # Mutually exclusive with the AR detector below: Cam requires an
-            # explicit "cam" marker so hyphenated AR local paths (e.g.
-            # /ckpts/dreamx-world-5b-converted) don't first-match here —
-            # detector resolution is first-match in registration order.
-            lambda path:
-            ("dreamx-world" in path.lower() and "cam" in path.lower()) or "dreamxworldpipeline" in path.lower()
-        ],
-        model_family="dreamx_world",
-        default_preset="dreamx_world_5b_cam",
-    )
-    register_configs(
-        sampling_param_cls=None,
-        pipeline_config_cls=DreamXWorld5BARPipelineConfig,
-        workload_types=(WorkloadType.I2V, ),
-        hf_model_paths=[
-            "FastVideo/DreamX-World-5B-Diffusers",
-        ],
-        model_detectors=[
-            lambda path:
-            ("dreamx-world-5b" in path.lower() and "cam" not in path.lower()) or "dreamxworldarpipeline" in path.lower(
-            )
-        ],
-        model_family="dreamx_world",
-        default_preset="dreamx_world_5b_ar",
-    )
     register_configs(
         sampling_param_cls=None,
         pipeline_config_cls=FastWan2_2_TI2V_5B_Config,
@@ -986,8 +951,6 @@ def _register_presets() -> None:
     from fastvideo.api.presets import register_preset
     from fastvideo.pipelines.basic.cosmos.presets import (
         ALL_PRESETS as COSMOS_PRESETS, )
-    from fastvideo.pipelines.basic.dreamx_world.presets import (
-        ALL_PRESETS as DREAMX_WORLD_PRESETS, )
     from fastvideo.pipelines.basic.gamecraft.presets import (
         ALL_PRESETS as GAMECRAFT_PRESETS, )
     from fastvideo.pipelines.basic.gen3c.presets import (
@@ -1021,7 +984,6 @@ def _register_presets() -> None:
 
     all_preset_groups = (
         COSMOS_PRESETS,
-        DREAMX_WORLD_PRESETS,
         FLUX2_PRESETS,
         GAMECRAFT_PRESETS,
         GEN3C_PRESETS,

fastvideo/tests/api/test_compat_translation.py

@@ -9,7 +9,7 @@ from fastvideo.api.compat import (
     generator_config_to_fastvideo_args,
     legacy_from_pretrained_to_config,
 )
-from fastvideo.api.schema import CompileConfig, GeneratorConfig
+from fastvideo.api.schema import BatchingConfig, CompileConfig, GeneratorConfig
 
 
 class TestLegacyTorchCompileKwargsTranslation:
@@ -17,6 +17,46 @@ class TestLegacyTorchCompileKwargsTranslation:
     first-class :class:`CompileConfig` fields and anything unknown falls
     into ``extras``."""
 
+    def test_all_typed_keys_promoted(self) -> None:
+        config = legacy_from_pretrained_to_config(
+            "/models/ltx2",
+            {
+                "enable_torch_compile": True,
+                "torch_compile_kwargs": {
+                    "backend": "inductor",
+                    "fullgraph": True,
+                    "mode": "max-autotune-no-cudagraphs",
+                    "dynamic": False,
+                },
+            },
+        )
+        compile_config = config.engine.compile
+        assert compile_config.enabled is True
+        assert compile_config.backend == "inductor"
+        assert compile_config.fullgraph is True
+        assert compile_config.mode == "max-autotune-no-cudagraphs"
+        assert compile_config.dynamic is False
+        assert compile_config.extras == {}
+
+    def test_unknown_keys_land_in_extras(self) -> None:
+        config = legacy_from_pretrained_to_config(
+            "/models/ltx2",
+            {
+                "enable_torch_compile": True,
+                "torch_compile_kwargs": {
+                    "backend": "inductor",
+                    "options": {"triton.cudagraphs": False},
+                    "disable": False,
+                },
+            },
+        )
+        compile_config = config.engine.compile
+        assert compile_config.backend == "inductor"
+        assert compile_config.extras == {
+            "options": {"triton.cudagraphs": False},
+            "disable": False,
+        }
+
     def test_empty_kwargs_produces_empty_extras(self) -> None:
         config = legacy_from_pretrained_to_config(
             "/models/ltx2",
@@ -77,6 +117,13 @@ class TestLegacyLtx2VaeTilingTranslation:
     ``generator.pipeline.vae_tiling``; reverse direction emits the
     legacy name back to FastVideoArgs."""
 
+    def test_forward_routes_to_pipeline_vae_tiling(self) -> None:
+        config = legacy_from_pretrained_to_config(
+            "/models/ltx2",
+            {"ltx2_vae_tiling": False},
+        )
+        assert config.pipeline.vae_tiling is False
+
     def test_true_round_trips(self) -> None:
         config = legacy_from_pretrained_to_config(
             "/models/ltx2",
@@ -153,6 +200,48 @@ class TestLegacyTextEncoderCompileTranslation:
         assert "enable_torch_compile_text_encoder" not in args.kwargs
 
 
+class TestBatchingTranslation:
+
+    def test_flat_kwargs_promote_to_engine_batching(self) -> None:
+        config = legacy_from_pretrained_to_config(
+            "/models/wan",
+            {
+                "batching_mode": "dynamic",
+                "batching_max_size": 4,
+                "batching_delay_ms": 25.0,
+                "batching_config": "/tmp/batching.json",
+                "enable_batching_metrics": True,
+            },
+        )
+
+        assert config.engine.batching.mode == "dynamic"
+        assert config.engine.batching.max_size == 4
+        assert config.engine.batching.delay_ms == 25.0
+        assert config.engine.batching.config_path == "/tmp/batching.json"
+        assert config.engine.batching.enable_metrics is True
+
+    def test_typed_batching_emits_fastvideo_args_kwargs(self, monkeypatch) -> None:
+        _stub_fastvideo_args_from_kwargs(monkeypatch)
+        config = GeneratorConfig(
+            model_path="/models/wan",
+            engine=_engine_with_batching(BatchingConfig(
+                mode="dynamic",
+                max_size=3,
+                delay_ms=10.0,
+                config_path="/tmp/batching.json",
+                enable_metrics=True,
+            )),
+        )
+
+        args = generator_config_to_fastvideo_args(config)
+
+        assert args.kwargs["batching_mode"] == "dynamic"
+        assert args.kwargs["batching_max_size"] == 3
+        assert args.kwargs["batching_delay_ms"] == 10.0
+        assert args.kwargs["batching_config"] == "/tmp/batching.json"
+        assert args.kwargs["enable_batching_metrics"] is True
+
+
 # -------------------------------------------------------------------
 # Helpers
 # -------------------------------------------------------------------
@@ -166,6 +255,13 @@ def _engine_with_compile(compile_config):
     return engine
 
 
+def _engine_with_batching(batching_config):
+    from fastvideo.api.schema import EngineConfig
+    engine = EngineConfig()
+    engine.batching = batching_config
+    return engine
+
+
 def _stub_fastvideo_args_from_kwargs(monkeypatch):
     """Swap ``FastVideoArgs.from_kwargs`` for a capture-only stub so
     translation tests don't need to construct a valid FastVideoArgs."""

fastvideo/tests/api/test_parser.py

@@ -130,6 +130,13 @@ def test_load_run_config_supports_yaml_roundtrip(tmp_path) -> None:
                 "use_fsdp_inference": False,
                 "disable_autocast": False,
                 "quantization": None,
+                "batching": {
+                    "mode": "disabled",
+                    "max_size": 1,
+                    "delay_ms": 0.0,
+                    "config_path": None,
+                    "enable_metrics": False,
+                },
             },
             "pipeline": {
                 "workload_type": None,

fastvideo/tests/batching/run_dynamic_batching_parity.py

@@ -0,0 +1,198 @@
+# SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import time
+from pathlib import Path
+from typing import Any
+
+import torch
+
+from fastvideo import VideoGenerator
+
+DEFAULT_PROMPTS = (
+    "A small robot sketches a city skyline at sunrise, cinematic lighting.",
+    "A glass teapot steams on a wooden table while rain falls outside.",
+)
+
+
+def _build_init_kwargs(args: argparse.Namespace, *, dynamic: bool) -> dict[str, Any]:
+    return {
+        "num_gpus": args.num_gpus,
+        "sp_size": args.sp_size,
+        "tp_size": args.tp_size,
+        "use_fsdp_inference": args.use_fsdp_inference,
+        "dit_cpu_offload": False,
+        "dit_layerwise_offload": False,
+        "flow_shift": args.flow_shift,
+        "text_encoder_precisions": ("fp32",),
+        "output_type": "latent",
+        "batching_mode": "dynamic" if dynamic else "disabled",
+        "batching_max_size": args.batch_size if dynamic else 1,
+        "batching_delay_ms": 0.0,
+    }
+
+
+def _request_kwargs(args: argparse.Namespace, prompt_index: int) -> dict[str, Any]:
+    return {
+        "prompt": args.prompts[prompt_index],
+        "height": args.height,
+        "width": args.width,
+        "num_frames": args.num_frames,
+        "num_inference_steps": args.num_inference_steps,
+        "guidance_scale": args.guidance_scale,
+        "embedded_cfg_scale": args.embedded_cfg_scale,
+        "seed": args.seed + prompt_index,
+        "fps": 24,
+        "save_video": False,
+        "return_frames": True,
+        "output_path": str(Path(args.output_dir) / f"request_{prompt_index}.mp4"),
+    }
+
+
+def _sync() -> None:
+    if torch.cuda.is_available():
+        torch.cuda.synchronize()
+
+
+def _run_sequential(generator: VideoGenerator, args: argparse.Namespace) -> tuple[list[dict[str, Any]], float]:
+    _sync()
+    start = time.perf_counter()
+    results = []
+    for index in range(args.batch_size):
+        kwargs = _request_kwargs(args, index)
+        prompt = kwargs.pop("prompt")
+        results.append(generator.generate_video(prompt=prompt, **kwargs))
+    _sync()
+    return results, time.perf_counter() - start
+
+
+def _run_dynamic(generator: VideoGenerator, args: argparse.Namespace) -> tuple[list[dict[str, Any]], float]:
+    if not hasattr(generator, "generate_video_batch"):
+        raise RuntimeError("VideoGenerator.generate_video_batch is unavailable in this checkout")
+    requests = [_request_kwargs(args, index) for index in range(args.batch_size)]
+    _sync()
+    start = time.perf_counter()
+    results = generator.generate_video_batch(requests)
+    _sync()
+    return results, time.perf_counter() - start
+
+
+def _tensor_metrics(sequential: list[dict[str, Any]], dynamic: list[dict[str, Any]]) -> dict[str, Any]:
+    per_request = []
+    for index, (seq_result, dyn_result) in enumerate(zip(sequential, dynamic, strict=True)):
+        seq = seq_result["samples"].detach().cpu().to(torch.float32)
+        dyn = dyn_result["samples"].detach().cpu().to(torch.float32)
+        diff = (seq - dyn).abs()
+        per_request.append({
+            "index": index,
+            "shape": list(seq.shape),
+            "max_abs_diff": float(diff.max().item()),
+            "mean_abs_diff": float(diff.mean().item()),
+            "allclose_atol_1e_5": bool(torch.allclose(seq, dyn, atol=1e-5, rtol=1e-5)),
+            "allclose_atol_1e_4": bool(torch.allclose(seq, dyn, atol=1e-4, rtol=1e-4)),
+        })
+    return {
+        "per_request": per_request,
+        "max_abs_diff": max(item["max_abs_diff"] for item in per_request),
+        "mean_abs_diff": sum(item["mean_abs_diff"] for item in per_request) / len(per_request),
+        "allclose_atol_1e_5": all(item["allclose_atol_1e_5"] for item in per_request),
+        "allclose_atol_1e_4": all(item["allclose_atol_1e_4"] for item in per_request),
+    }
+
+
+def run_parity(args: argparse.Namespace) -> dict[str, Any]:
+    generator = VideoGenerator.from_pretrained(args.model_path, **_build_init_kwargs(args, dynamic=True))
+    try:
+        sequential, sequential_s = _run_sequential(generator, args)
+        dynamic, dynamic_s = _run_dynamic(generator, args)
+        metrics = _tensor_metrics(sequential, dynamic)
+    finally:
+        generator.shutdown()
+    return {
+        "mode": "parity",
+        "model_path": args.model_path,
+        "num_gpus": args.num_gpus,
+        "shape": {
+            "height": args.height,
+            "width": args.width,
+            "num_frames": args.num_frames,
+            "num_inference_steps": args.num_inference_steps,
+        },
+        "batch_size": args.batch_size,
+        "sequential_time_s": sequential_s,
+        "dynamic_time_s": dynamic_s,
+        "speedup": sequential_s / dynamic_s if dynamic_s > 0 else None,
+        "tensor_metrics": metrics,
+    }
+
+
+def run_benchmark(args: argparse.Namespace, *, dynamic: bool) -> dict[str, Any]:
+    generator = VideoGenerator.from_pretrained(args.model_path, **_build_init_kwargs(args, dynamic=dynamic))
+    run = _run_dynamic if dynamic else _run_sequential
+    try:
+        for _ in range(args.warmup_runs):
+            run(generator, args)
+        times = []
+        for _ in range(args.measurement_runs):
+            _results, elapsed = run(generator, args)
+            times.append(elapsed)
+    finally:
+        generator.shutdown()
+    avg = sum(times) / len(times)
+    return {
+        "mode": "dynamic" if dynamic else "sequential",
+        "model_path": args.model_path,
+        "num_gpus": args.num_gpus,
+        "batch_size": args.batch_size,
+        "measurement_runs": args.measurement_runs,
+        "times_s": times,
+        "avg_time_s": avg,
+        "requests_per_second": args.batch_size / avg if avg > 0 else None,
+    }
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--mode", choices=("parity", "sequential", "dynamic"), required=True)
+    parser.add_argument("--model-path", default="Wan-AI/Wan2.1-T2V-1.3B-Diffusers")
+    parser.add_argument("--num-gpus", type=int, default=1)
+    parser.add_argument("--sp-size", type=int, default=1)
+    parser.add_argument("--tp-size", type=int, default=1)
+    parser.add_argument("--use-fsdp-inference", action="store_true")
+    parser.add_argument("--height", type=int, default=256)
+    parser.add_argument("--width", type=int, default=256)
+    parser.add_argument("--num-frames", type=int, default=9)
+    parser.add_argument("--num-inference-steps", type=int, default=2)
+    parser.add_argument("--guidance-scale", type=float, default=1.0)
+    parser.add_argument("--embedded-cfg-scale", type=float, default=6.0)
+    parser.add_argument("--flow-shift", type=float, default=7.0)
+    parser.add_argument("--seed", type=int, default=1024)
+    parser.add_argument("--batch-size", type=int, default=2)
+    parser.add_argument("--warmup-runs", type=int, default=1)
+    parser.add_argument("--measurement-runs", type=int, default=3)
+    parser.add_argument("--output-dir", default="/tmp/fastvideo_dynamic_batching")
+    parser.add_argument("--output-json", required=True)
+    parser.add_argument("--prompts", nargs="+", default=list(DEFAULT_PROMPTS))
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    if len(args.prompts) < args.batch_size:
+        raise ValueError("--prompts must contain at least --batch-size prompts")
+    os.makedirs(args.output_dir, exist_ok=True)
+    if args.mode == "parity":
+        result = run_parity(args)
+    else:
+        result = run_benchmark(args, dynamic=args.mode == "dynamic")
+    output_path = Path(args.output_json)
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    output_path.write_text(json.dumps(result, indent=2), encoding="utf-8")
+    print(json.dumps(result, indent=2))
+
+
+if __name__ == "__main__":
+    main()

fastvideo/tests/batching/test_admission.py

@@ -0,0 +1,89 @@
+# SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import pytest
+
+from fastvideo.batching.admission import (
+    AdmissionLimit,
+    BatchAdmissionController,
+    BatchingRule,
+    load_batching_config,
+)
+from fastvideo.configs.pipelines.base import PipelineConfig
+
+
+def test_admission_limit_rejects_batch_size_and_cost() -> None:
+    limit = AdmissionLimit(max_batch_size=2, max_cost=10.0)
+
+    assert limit.reject_reason(batch_size=3, batch_cost=1.0) == "config_cap:2"
+    assert limit.reject_reason(batch_size=2, batch_cost=11.0) == "cost_budget:11>10"
+    assert limit.reject_reason(batch_size=2, batch_cost=10.0) is None
+
+
+def test_batching_rule_validates_unknown_keys() -> None:
+    with pytest.raises(ValueError, match="did you mean 'max_batch_size'"):
+        BatchingRule.from_dict(
+            {
+                "model_contains": "wan",
+                "max_batch_siz": 2,
+            },
+            source="unit",
+        )
+
+
+@pytest.mark.parametrize(("value", "expected"), [(1, True), (0, False), (1.0, True), (0.0, False)])
+def test_batching_rule_parses_numeric_bool_values(value, expected) -> None:
+    rule = BatchingRule.from_dict(
+        {
+            "model_contains": "wan",
+            "offload": value,
+            "max_batch_size": 2,
+        },
+        source="unit",
+    )
+
+    assert rule.offload is expected
+
+
+def test_load_batching_config_supports_mapping_form(tmp_path) -> None:
+    path = tmp_path / "batching.json"
+    path.write_text(
+        '{"schema_version": 1, "wan|720x1280x81": {"max_batch_size": 3, "max_cost": 9}}',
+        encoding="utf-8",
+    )
+
+    rules = load_batching_config(str(path))
+
+    assert len(rules) == 1
+    assert rules[0].model == "wan"
+    assert rules[0].resolution == "720x1280x81"
+    assert rules[0].max_batch_size == 3
+    assert rules[0].max_cost == 9.0
+
+
+def test_admission_controller_applies_user_and_config_caps(tmp_path, monkeypatch) -> None:
+    path = tmp_path / "batching.json"
+    path.write_text(
+        '{"rules": [{"model_contains": "wan", "resolution": "720x1280x81", "max_batch_size": 3}]}',
+        encoding="utf-8",
+    )
+    monkeypatch.setattr(BatchAdmissionController, "_get_device_memory_gb", staticmethod(lambda gpu_id: 48.0))
+
+    args = SimpleNamespace(
+        batching_mode="dynamic",
+        batching_max_size=4,
+        batching_config=str(path),
+        model_path="/models/wan",
+        dit_cpu_offload=False,
+        dit_layerwise_offload=False,
+        pipeline_config=PipelineConfig(),
+    )
+    request = SimpleNamespace(height=720, width=1280, num_frames=81)
+
+    controller = BatchAdmissionController(args)
+
+    assert controller.enabled is True
+    assert controller.max_admissible_batch_size(request) == 3
+    assert controller.batch_is_full([request, request, request]) is True

fastvideo/tests/batching/test_signature.py

@@ -0,0 +1,57 @@
+# SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
+
+from fastvideo.api.sampling_param import SamplingParam
+from fastvideo.batching.signature import (
+    can_dynamic_batch,
+    dynamic_batch_signature,
+    resolution_key,
+)
+
+
+def _request(prompt: str = "a prompt", **overrides) -> SamplingParam:
+    request = SamplingParam(prompt=prompt, height=256, width=384, num_frames=17, num_inference_steps=4)
+    for key, value in overrides.items():
+        setattr(request, key, value)
+    return request
+
+
+def test_dynamic_batch_signature_excludes_request_local_fields() -> None:
+    first = _request(seed=1, output_path="/tmp/a.mp4", save_video=True, return_frames=False)
+    second = _request(seed=2, output_path="/tmp/b.mp4", save_video=False, return_frames=True)
+
+    assert dynamic_batch_signature(first) == dynamic_batch_signature(second)
+
+
+def test_can_dynamic_batch_accepts_matching_text_requests() -> None:
+    first = _request("first", seed=1)
+    second = _request("second", seed=2)
+
+    result = can_dynamic_batch(first, second)
+
+    assert result.can_batch is True
+    assert result.reason is None
+
+
+def test_can_dynamic_batch_rejects_sampling_mismatch() -> None:
+    first = _request(guidance_scale=1.0)
+    second = _request(guidance_scale=3.0)
+
+    result = can_dynamic_batch(first, second)
+
+    assert result.can_batch is False
+    assert result.reason == "sampling_params.guidance_scale"
+
+
+def test_can_dynamic_batch_rejects_image_conditioning() -> None:
+    first = _request()
+    second = _request(image_path="/tmp/image.png")
+
+    result = can_dynamic_batch(first, second)
+
+    assert result.can_batch is False
+    assert result.reason == "image_path"
+
+
+def test_resolution_key_uses_generation_shape() -> None:
+    assert resolution_key(_request(height=720, width=1280, num_frames=81)) == "720x1280x81"

fastvideo/tests/contract/test_ci_test_collection.py

@@ -1,84 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Guard: every test directory must be collected by some CI lane or be on
-the explicit allowlist below.
-
-Three separate incidents on 2026-07-05 found test files that no CI lane
-ever collects (fastvideo/tests/stages/, tests/local_tests/ additions in
-PR #1509, and this sweep found seven dark directories in total): the tests
-pass review, merge, and then silently never run. This test makes going
-dark an explicit, reviewed decision instead of an accident: adding a new
-test directory fails CI until it is either wired into a lane or
-allowlisted here with a reason.
-
-Pure text analysis — no fastvideo imports, no GPU, no torch.
-"""
-from pathlib import Path
-
-REPO_ROOT = Path(__file__).resolve().parents[3]
-TESTS_ROOT = REPO_ROOT / "fastvideo" / "tests"
-
-# Files whose text constitutes "a CI lane references this directory".
-CI_SOURCES = [
-    TESTS_ROOT / "modal" / "pr_test.py",
-    TESTS_ROOT / "modal" / "ssim_test.py",
-    *sorted((REPO_ROOT / ".buildkite").rglob("*.yml")),
-    *sorted((REPO_ROOT / ".buildkite").rglob("*.sh")),
-]
-
-# Directories that intentionally have no CI lane today. Every entry needs a
-# reason; remove the entry when the directory gets wired into a lane.
-# State as found on 2026-07-05 — these SHOULD shrink over time, not grow.
-ALLOWLIST = {
-    "attention": "no lane yet — GPU attention-backend tests, run manually",
-    "audio": "no lane yet — audio encoder tests, run manually",
-    "distributed": "no lane yet — multi-GPU torchrun tests, run manually",
-    "hooks": "no lane yet — run manually",
-    "layers": "no lane yet — torchrun FSDP dispatch tests, run manually",
-    "nightly": "by design: nightly cadence, not per-PR",
-    "modal": "CI infrastructure itself, not a test suite",
-}
-
-
-def _dirs_with_tests() -> list[str]:
-    dirs = []
-    for child in sorted(TESTS_ROOT.iterdir()):
-        if child.is_dir() and any(child.rglob("test_*.py")):
-            dirs.append(child.name)
-    return dirs
-
-
-def _ci_text() -> str:
-    return "\n".join(
-        src.read_text(errors="replace") for src in CI_SOURCES if src.exists())
-
-
-def test_every_test_directory_is_collected_or_allowlisted():
-    ci_text = _ci_text()
-    dark = [
-        name for name in _dirs_with_tests()
-        if f"tests/{name}" not in ci_text and name not in ALLOWLIST
-    ]
-    assert not dark, (
-        f"Test directories not referenced by any CI lane and not "
-        f"allowlisted: {dark}. Wire them into a lane in "
-        f"fastvideo/tests/modal/pr_test.py (or a Buildkite step), or add an "
-        f"allowlist entry with a reason in {__file__}.")
-
-
-def test_local_tests_stays_out_of_ci():
-    # tests/local_tests/ (repo root) is developer-local by design (author
-    # decision, 2026-07-05): parity scaffolds and machine-specific checks
-    # that must never gate CI. Fail if any CI source starts collecting it.
-    assert "tests/local_tests" not in _ci_text(), (
-        "tests/local_tests/ is local-only by design; remove the CI "
-        "reference or move the tests into a fastvideo/tests/ lane.")
-
-
-def test_allowlist_entries_are_still_real_directories():
-    # A stale allowlist hides regressions; entries must track reality.
-    missing = [
-        name for name in ALLOWLIST
-        if name != "modal" and not (TESTS_ROOT / name).is_dir()
-    ]
-    assert not missing, (
-        f"Allowlisted directories no longer exist — remove them: {missing}")

fastvideo/tests/entrypoints/test_openai_api.py

@@ -1,10 +1,15 @@
 """Unit tests for the OpenAI-compatible API server helpers (no GPU needed)."""
 
+import asyncio
 import os
+import time
+from types import SimpleNamespace
 from unittest.mock import patch
 
 import pytest
 
+from fastvideo.configs.pipelines.base import PipelineConfig
+from fastvideo.entrypoints.openai.batching import VideoBatchScheduler, _VideoBatchJob
 from fastvideo.api.parser import parse_config
 from fastvideo.api.schema import GenerationRequest
 from fastvideo.entrypoints.openai.protocol import (
@@ -21,6 +26,172 @@ from fastvideo.entrypoints.openai.utils import (
     parse_size,
 )
 
+
+class _FakeBatchGenerator:
+
+    def __init__(self):
+        self.calls = []
+
+    def generate_video_batch(self, request_kwargs):
+        self.calls.append([dict(item) for item in request_kwargs])
+        return [{"prompts": item["prompt"], "video_path": item["output_path"]} for item in request_kwargs]
+
+
+def _make_batch_job(request_id, kwargs):
+    loop = asyncio.get_running_loop()
+    return _VideoBatchJob(
+        request_id=request_id,
+        kwargs=dict(kwargs),
+        future=loop.create_future(),
+        enqueue_time=time.perf_counter(),
+    )
+
+
+def _batch_scheduler_args(**overrides):
+    defaults = dict(
+        model_path="test-model",
+        batching_mode="dynamic",
+        batching_max_size=2,
+        batching_delay_ms=25.0,
+        enable_batching_metrics=False,
+        pipeline_config=PipelineConfig(),
+    )
+    defaults.update(overrides)
+    return SimpleNamespace(**defaults)
+
+
+def test_video_batch_scheduler_groups_compatible_requests(tmp_path):
+    async def run():
+        generator = _FakeBatchGenerator()
+        scheduler = VideoBatchScheduler(generator, _batch_scheduler_args())
+        await scheduler.start()
+        try:
+            first = {
+                "prompt": "first",
+                "height": 256,
+                "width": 256,
+                "num_frames": 1,
+                "num_inference_steps": 2,
+                "seed": 1,
+                "output_path": str(tmp_path / "first.mp4"),
+                "save_video": False,
+            }
+            second = {
+                "prompt": "second",
+                "height": 256,
+                "width": 256,
+                "num_frames": 1,
+                "num_inference_steps": 2,
+                "seed": 2,
+                "output_path": str(tmp_path / "second.mp4"),
+                "save_video": False,
+            }
+            results = await asyncio.gather(
+                scheduler.submit("req-1", first),
+                scheduler.submit("req-2", second),
+            )
+        finally:
+            await scheduler.stop()
+        return generator.calls, results
+
+    calls, results = asyncio.run(run())
+
+    assert len(calls) == 1
+    assert [item["prompt"] for item in calls[0]] == ["first", "second"]
+    assert [result["prompts"] for result in results] == ["first", "second"]
+
+
+def test_video_batch_scheduler_keeps_incompatible_requests_separate(tmp_path):
+    async def run():
+        generator = _FakeBatchGenerator()
+        scheduler = VideoBatchScheduler(generator, _batch_scheduler_args())
+        await scheduler.start()
+        try:
+            text_only = {
+                "prompt": "first",
+                "height": 256,
+                "width": 256,
+                "num_frames": 1,
+                "num_inference_steps": 2,
+                "seed": 1,
+                "output_path": str(tmp_path / "first.mp4"),
+                "save_video": False,
+            }
+            image_conditioned = {
+                "prompt": "second",
+                "height": 256,
+                "width": 256,
+                "num_frames": 1,
+                "num_inference_steps": 2,
+                "seed": 2,
+                "image_path": str(tmp_path / "input.png"),
+                "output_path": str(tmp_path / "second.mp4"),
+                "save_video": False,
+            }
+            results = await asyncio.gather(
+                scheduler.submit("req-1", text_only),
+                scheduler.submit("req-2", image_conditioned),
+            )
+        finally:
+            await scheduler.stop()
+        return generator.calls, results
+
+    calls, results = asyncio.run(run())
+
+    assert len(calls) == 2
+    assert [[item["prompt"] for item in call] for call in calls] == [["first"], ["second"]]
+    assert [result["prompts"] for result in results] == ["first", "second"]
+
+
+def test_video_batch_scheduler_requeues_incompatible_pending_job_at_front(tmp_path):
+    async def run():
+        generator = _FakeBatchGenerator()
+        scheduler = VideoBatchScheduler(generator, _batch_scheduler_args())
+        first = {
+            "prompt": "first",
+            "height": 256,
+            "width": 256,
+            "num_frames": 1,
+            "num_inference_steps": 2,
+            "seed": 1,
+            "output_path": str(tmp_path / "first.mp4"),
+            "save_video": False,
+        }
+        incompatible = {
+            "prompt": "second",
+            "height": 256,
+            "width": 256,
+            "num_frames": 1,
+            "num_inference_steps": 2,
+            "seed": 2,
+            "image_path": str(tmp_path / "input.png"),
+            "output_path": str(tmp_path / "second.mp4"),
+            "save_video": False,
+        }
+        newer = {
+            "prompt": "third",
+            "height": 256,
+            "width": 256,
+            "num_frames": 1,
+            "num_inference_steps": 2,
+            "seed": 3,
+            "output_path": str(tmp_path / "third.mp4"),
+            "save_video": False,
+        }
+
+        scheduler._pending.extend([
+            _make_batch_job("req-2", incompatible),
+            _make_batch_job("req-3", newer),
+        ])
+        batch = await scheduler._collect_batch(_make_batch_job("req-1", first))
+        return [job.request_id for job in batch], [job.request_id for job in scheduler._pending]
+
+    batch_ids, pending_ids = asyncio.run(run())
+
+    assert batch_ids == ["req-1"]
+    assert pending_ids == ["req-2", "req-3"]
+
+
 # ---------------------------------------------------------------------------
 # parse_size
 # ---------------------------------------------------------------------------

fastvideo/tests/entrypoints/test_video_generator.py

@@ -3,6 +3,7 @@ from types import SimpleNamespace
 import warnings
 
 import pytest
+import torch
 
 from fastvideo.api import (
     GenerationRequest,
@@ -13,8 +14,10 @@ from fastvideo.api import (
     load_run_config,
 )
 from fastvideo.api.sampling_param import SamplingParam
+from fastvideo.configs.pipelines.base import PipelineConfig
 from fastvideo.entrypoints.video_generator import VideoGenerator
 from fastvideo.fastvideo_args import WorkloadType
+from fastvideo.pipelines import ForwardBatch
 
 
 def _new_video_generator() -> VideoGenerator:
@@ -37,6 +40,25 @@ def _new_runtime_video_generator() -> VideoGenerator:
     return generator
 
 
+def _batching_fastvideo_args(**overrides):
+    defaults = dict(
+        model_path="test-model",
+        prompt_txt=None,
+        workload_type=SimpleNamespace(value="t2v"),
+        batching_mode="dynamic",
+        batching_max_size=4,
+        batching_config=None,
+        dit_cpu_offload=False,
+        dit_layerwise_offload=False,
+        output_type="latent",
+        pin_cpu_memory=False,
+        VSA_sparsity=0.0,
+        pipeline_config=PipelineConfig(),
+    )
+    defaults.update(overrides)
+    return SimpleNamespace(**defaults)
+
+
 def _patch_from_fastvideo_args(monkeypatch):
     captured = {}
 
@@ -151,6 +173,117 @@ def test_prepare_output_path_empty_prompt_fallback(tmp_path):
     assert os.path.basename(result) == "output.mp4"
 
 
+def test_generate_prepared_work_items_merges_compatible_latent_requests(monkeypatch, tmp_path):
+    vg = _new_video_generator()
+    vg.fastvideo_args = _batching_fastvideo_args()
+    calls = []
+
+    def fake_device_memory(gpu_id):
+        return 48.0
+
+    def fake_run_forward(batch, fastvideo_args):
+        calls.append(batch)
+        batch_size = len(batch.prompt) if isinstance(batch.prompt, list) else 1
+        output = torch.arange(batch_size * 4, dtype=torch.float32).reshape(batch_size, 4, 1, 1, 1)
+        return ForwardBatch(data_type=batch.data_type, output=output, extra={"peak_memory_mb": 1.0}), 0.5, 10.0
+
+    monkeypatch.setattr(
+        "fastvideo.batching.admission.BatchAdmissionController._get_device_memory_gb",
+        staticmethod(fake_device_memory),
+    )
+    monkeypatch.setattr(vg, "_run_forward_batch", fake_run_forward)
+
+    first = SamplingParam(prompt="one", height=8, width=8, num_frames=1, seed=11, return_frames=True, save_video=False)
+    second = SamplingParam(prompt="two", height=8, width=8, num_frames=1, seed=22, return_frames=True, save_video=False)
+    work_items = [
+        vg._prepare_generation_work_item("one", first, vg.fastvideo_args, output_path=str(tmp_path / "one.mp4")),
+        vg._prepare_generation_work_item("two", second, vg.fastvideo_args, output_path=str(tmp_path / "two.mp4")),
+    ]
+
+    results = vg._generate_prepared_work_items(work_items)
+
+    assert len(calls) == 1
+    assert calls[0].prompt == ["one", "two"]
+    assert calls[0].seeds == [11, 22]
+    assert [result["prompts"] for result in results] == ["one", "two"]
+    assert [result["samples"].shape for result in results] == [(1, 4, 1, 1, 1), (1, 4, 1, 1, 1)]
+
+
+def test_generate_prepared_work_items_falls_back_for_incompatible_requests(monkeypatch, tmp_path):
+    vg = _new_video_generator()
+    vg.fastvideo_args = _batching_fastvideo_args()
+    calls = []
+
+    def fake_run_forward(batch, fastvideo_args):
+        calls.append(batch)
+        output = torch.zeros((1, 4, 1, 1, 1), dtype=torch.float32)
+        return ForwardBatch(data_type=batch.data_type, output=output), 0.5, 10.0
+
+    monkeypatch.setattr(vg, "_run_forward_batch", fake_run_forward)
+
+    first = SamplingParam(prompt="one", height=8, width=8, num_frames=1, guidance_scale=1.0, save_video=False)
+    second = SamplingParam(prompt="two", height=8, width=8, num_frames=1, guidance_scale=3.0, save_video=False)
+    work_items = [
+        vg._prepare_generation_work_item("one", first, vg.fastvideo_args, output_path=str(tmp_path / "one.mp4")),
+        vg._prepare_generation_work_item("two", second, vg.fastvideo_args, output_path=str(tmp_path / "two.mp4")),
+    ]
+
+    results = vg._generate_prepared_work_items(work_items)
+
+    assert len(calls) == 2
+    assert all(isinstance(call.prompt, str) for call in calls)
+    assert [result["prompts"] for result in results] == ["one", "two"]
+
+
+def test_generate_video_batch_routes_compat_kwargs(monkeypatch, tmp_path):
+    vg = _new_video_generator()
+    vg.fastvideo_args = _batching_fastvideo_args()
+    calls = []
+
+    def fake_device_memory(gpu_id):
+        return 48.0
+
+    def fake_run_forward(batch, fastvideo_args):
+        calls.append((batch, fastvideo_args))
+        output = torch.zeros((len(batch.prompt), 4, 1, 1, 1), dtype=torch.float32)
+        return ForwardBatch(data_type=batch.data_type, output=output), 0.5, 10.0
+
+    monkeypatch.setattr(
+        "fastvideo.batching.admission.BatchAdmissionController._get_device_memory_gb",
+        staticmethod(fake_device_memory),
+    )
+    monkeypatch.setattr(vg, "_run_forward_batch", fake_run_forward)
+
+    results = vg.generate_video_batch([
+        {
+            "prompt": "one",
+            "height": 8,
+            "width": 8,
+            "num_frames": 1,
+            "embedded_cfg_scale": 7.5,
+            "save_video": False,
+            "return_frames": True,
+            "output_path": str(tmp_path / "one.mp4"),
+        },
+        {
+            "prompt": "two",
+            "height": 8,
+            "width": 8,
+            "num_frames": 1,
+            "embedded_cfg_scale": 7.5,
+            "save_video": False,
+            "return_frames": True,
+            "output_path": str(tmp_path / "two.mp4"),
+        },
+    ])
+
+    assert len(calls) == 1
+    batch, fastvideo_args = calls[0]
+    assert batch.prompt == ["one", "two"]
+    assert fastvideo_args.pipeline_config.embedded_cfg_scale == 7.5
+    assert [result["prompts"] for result in results] == ["one", "two"]
+
+
 def test_from_config_normalizes_and_translates(monkeypatch):
     captured = _patch_from_fastvideo_args(monkeypatch)
     _patch_fastvideo_args_from_kwargs(monkeypatch)

fastvideo/tests/modal/pr_test.py

@@ -187,8 +187,6 @@ def run_transformer_tests():
 
 
 @app.function(gpu="L40S:4",
-              cpu=8.0,
-              memory=32768,
               image=image,
               timeout=900,
               secrets=[
@@ -203,8 +201,6 @@ def run_training_tests():
 
 
 @app.function(gpu="L40S:2",
-              cpu=8.0,
-              memory=32768,
               image=image,
               timeout=900,
               secrets=[
@@ -280,7 +276,7 @@ def run_self_forcing_tests():
 @app.function(gpu="L40S:1", image=image, timeout=900)
 def run_unit_test():
     run_test(
-        "pytest ./fastvideo/tests/api/ ./fastvideo/tests/contract/ ./fastvideo/tests/dataset/ ./fastvideo/tests/workflow/ ./fastvideo/tests/entrypoints/ ./fastvideo/tests/train/ ./fastvideo/tests/stages/ ./fastvideo/tests/ops/ ./fastvideo/tests/training/test_trackers.py --ignore=./fastvideo/tests/entrypoints/test_openai_api_integration.py --ignore=./fastvideo/tests/train/models --ignore=./fastvideo/tests/train/methods -vs"
+        "pytest ./fastvideo/tests/api/ ./fastvideo/tests/batching/ ./fastvideo/tests/contract/ ./fastvideo/tests/dataset/ ./fastvideo/tests/workflow/ ./fastvideo/tests/entrypoints/ ./fastvideo/tests/train/ ./fastvideo/tests/stages/ --ignore=./fastvideo/tests/entrypoints/test_openai_api_integration.py --ignore=./fastvideo/tests/train/models --ignore=./fastvideo/tests/train/methods -vs"
     )
 
 
@@ -322,8 +318,6 @@ def run_dreamverse_app_tests():
 
 
 @app.function(gpu="L40S:1",
-              cpu=8.0,
-              memory=32768,
               image=image,
               timeout=1800,
               secrets=[
@@ -406,8 +400,6 @@ def run_lora_extraction_tests():
 
 
 @app.function(gpu="L40S:2",
-              cpu=8.0,
-              memory=32768,
               image=image,
               timeout=1800,
               secrets=[
@@ -435,8 +427,6 @@ def run_performance_tests():
         "export PERF_RUN_SOURCE='unknown'; "
         "export PERF_UPLOAD_POLICY='never'; "
         "fi; "
-        "(nvidia-smi --query-gpu=index,timestamp,clocks.sm,clocks.max.sm,power.draw,power.limit,temperature.gpu "
-        "--format=csv -l 10 > /tmp/gpu_telemetry.csv 2>/dev/null &); "
         "pytest ./fastvideo/tests/performance -vs; "
         "PYTEST_RC=$?; "
         "PERF_RC=0; "
@@ -445,8 +435,6 @@ def run_performance_tests():
         "PERF_RC=$?; "
         "fi; "
         "python ./fastvideo/tests/performance/dashboard.py || true; "
-        "echo '--- GPU telemetry (clocks.sm vs clocks.max.sm reveals capped hosts) ---'; "
-        "cat /tmp/gpu_telemetry.csv || true; "
         "FINAL_RC=$PYTEST_RC; "
         "if [ $FINAL_RC -eq 0 ]; then FINAL_RC=$PERF_RC; fi; "
         "exit $FINAL_RC")

fastvideo/tests/ops/quantization/test_absmax_fp8.py

@@ -34,9 +34,6 @@ class TestAbsMaxFP8LinearMethod(unittest.TestCase):
             method._convert_scale(scale)
 
     def test_create_weights_rejects_invalid_dtype(self):
-        # float64 is outside the supported set (bf16/fp16/fp32). The test
-        # originally passed float32, which create_weights has accepted since
-        # PR #981 — it never ran in CI, so the mismatch went unnoticed.
         method = AbsMaxFP8LinearMethod()
         layer = nn.Module()
         with self.assertRaisesRegex(AssertionError, "only supports"):
@@ -46,7 +43,7 @@ class TestAbsMaxFP8LinearMethod(unittest.TestCase):
                 output_partition_sizes=[3],
                 input_size=2,
                 output_size=3,
-                params_dtype=torch.float64,
+                params_dtype=torch.float32,
             )
 
     def test_absmax_fp8_parameter_weight_loader(self):

fastvideo/tests/performance/compare_baseline.py

@@ -5,8 +5,7 @@ This script:
 1) reads current benchmark results from fastvideo/tests/performance/results,
 2) syncs the canonical baseline from the configured HF dataset repo,
 3) compares each current record against the median of up to 5 prior
-   baseline-eligible successful records in the same workload/variant/version,
-   GPU, recipe, hardware, and software cohort,
+   baseline-eligible successful records (filtered by gpu_type),
 4) writes normalized records back to the HF dataset repo according to
    PERF_UPLOAD_POLICY,
 5) exits non-zero if any gated metric exceeds both its percent and absolute
@@ -65,27 +64,6 @@ PERF_REPORTS_DIR = os.environ.get("PERF_REPORTS_DIR", "/root/data/perf_reports")
 UPLOAD_POLICY = os.environ.get("PERF_UPLOAD_POLICY", "never").strip().lower()
 VALID_UPLOAD_POLICIES = {"never", "pass", "always"}
 VALID_RUN_SOURCES = {"pr", "local", "scheduled_main", "unknown"}
-IDENTITY_KEYS = (
-    "workload_id",
-    "variant_id",
-    "benchmark_version",
-    "recipe",
-    "recipe_fingerprint",
-    "hardware_profile",
-    "hardware_profile_id",
-    "software_profile",
-    "software_profile_id",
-    "environment_metadata",
-    "environment_fingerprint",
-)
-COMPARISON_IDENTITY_KEYS = (
-    "workload_id",
-    "variant_id",
-    "benchmark_version",
-    "recipe_fingerprint",
-    "hardware_profile_id",
-    "software_profile_id",
-)
 
 
 def _should_persist_tracking() -> bool:
@@ -158,46 +136,6 @@ def _record_metadata(run_source: str, result: dict[str, Any]) -> dict[str, Any]:
     }
 
 
-def _identity_metadata(result: dict[str, Any]) -> dict[str, Any]:
-    metadata = {
-        key: result[key]
-        for key in IDENTITY_KEYS
-        if key in result and result[key] is not None
-    }
-    recipe = result.get("recipe")
-    benchmark = recipe.get("benchmark") if isinstance(recipe, dict) else None
-    if isinstance(benchmark, dict):
-        for key in ("workload_id", "variant_id", "benchmark_version"):
-            if key not in metadata and benchmark.get(key) is not None:
-                metadata[key] = benchmark[key]
-    return metadata
-
-
-def _comparison_identity_filters(record: dict[str, Any]) -> dict[str, str]:
-    missing = [
-        key for key in COMPARISON_IDENTITY_KEYS
-        if key not in record or record[key] is None or (isinstance(record[key], str) and not record[key].strip())
-    ]
-    if len(missing) == len(COMPARISON_IDENTITY_KEYS):
-        # Legacy v1 record: no identity fields at all. Keep the pre-v2
-        # (model_id, gpu_type) rolling-baseline comparison instead of
-        # crashing — v1 configs are documented as loadable.
-        return {}
-    if missing:
-        raise ValueError("Performance record missing required comparison identity fields: " + ", ".join(missing))
-    return {
-        key: str(record[key])
-        for key in COMPARISON_IDENTITY_KEYS
-    }
-
-
-def _format_identity_filters(filters: dict[str, str]) -> str:
-    if not filters:
-        return ""
-    compact = ", ".join(f"{key}={value}" for key, value in filters.items())
-    return f" ({compact})"
-
-
 def _load_current_results() -> list[dict[str, Any]]:
     pattern = os.path.join(RESULTS_DIR, "perf_*.json")
     records: list[dict[str, Any]] = []
@@ -231,7 +169,7 @@ def normalize_performance_result(result: dict[str, Any]) -> dict[str, Any]:
     vae_decode_time = safe_float(result.get("vae_decode_time_s"))
     metric_policies = resolve_metric_policies(result.get("regression_thresholds"))
 
-    record = {
+    return {
         "model_id": model_id,
         "timestamp": timestamp,
         "commit_sha": commit_sha,
@@ -246,8 +184,6 @@ def normalize_performance_result(result: dict[str, Any]) -> dict[str, Any]:
         "success": True,
         **_record_metadata(_detect_run_source(), result),
     }
-    record.update(_identity_metadata(result))
-    return record
 
 
 def _normalize_record(result: dict[str, Any]) -> dict[str, Any]:
@@ -486,31 +422,17 @@ def main() -> int:
             TRACKING_ROOT,
             record["model_id"],
             record["gpu_type"],
-            **_comparison_identity_filters(record),
             last_n=5,
             successful_only=True,
             baseline_eligible_only=True,
         )
 
         if not baseline_records:
-            # A brand-new cohort has NO regression gating until history
-            # accumulates — make that loud and machine-readable instead of
-            # an indistinguishable pass, so a cohort shift (intended or
-            # accidental, e.g. an identity-field change) never silently
-            # blinds the comparison.
-            record["baseline_status"] = "initialized_new_cohort"
-            print("=" * 72)
-            print(f"WARNING: NO BASELINE — initializing a NEW cohort for "
-                  f"{record['model_id']} on {record['gpu_type']}"
-                  f"{_format_identity_filters(_comparison_identity_filters(record))}")
-            print("Regression gating is INACTIVE for this cohort until "
-                  "baseline history accumulates. If this cohort shift is "
-                  "unexpected, check the identity fields above.")
-            print("=" * 72)
+            print(f"No baseline for {record['model_id']} on "
+                  f"{record['gpu_type']}. Initializing...")
             failures: list[str] = []
             record["success"] = True
         else:
-            record["baseline_status"] = "compared"
             failures = _check_regressions(record, baseline_records, metric_policies)
         if static_threshold_failed:
             failures.append(f"{record['model_id']} fixed-threshold phase failed "

fastvideo/tests/performance/dashboard.py

@@ -29,76 +29,15 @@ METRICS = (
     "dit_time_s",
     "vae_decode_time_s",
 )
-COMPARISON_COHORT_KEYS = (
-    "workload_id",
-    "variant_id",
-    "benchmark_version",
-    "recipe_fingerprint",
-    "hardware_profile_id",
-    "software_profile_id",
-)
-GROUP_KEYS = ("model_id", "gpu_type", *COMPARISON_COHORT_KEYS)
-
-
-def _cohort_value(value: object) -> str:
-    if value is None:
-        return ""
-    try:
-        if pd.isna(value):
-            return ""
-    except (TypeError, ValueError):
-        pass
-    return str(value)
-
-
-def _display_value(value: object) -> str:
-    return _cohort_value(value) or "legacy"
-
-
-def _short_value(value: object) -> str:
-    text = _display_value(value)
-    if text == "legacy" or len(text) <= 14:
-        return text
-    return text[:12]
-
-
-def _cohort_title(group_key: tuple[object, ...]) -> str:
-    workload = _display_value(group_key[2])
-    variant = _display_value(group_key[3])
-    version = _display_value(group_key[4])
-    version_label = version if version == "legacy" else f"v{version}"
-    return f"{workload} / {variant} / {version_label}"
-
-
-def _cohort_detail(group_key: tuple[object, ...]) -> str:
-    return " | ".join((
-        f"recipe {_short_value(group_key[5])}",
-        _short_value(group_key[6]),
-        _short_value(group_key[7]),
-    ))
-
-
-def _group_record(group_key: tuple[object, ...]) -> dict[str, str]:
-    return {
-        key: _cohort_value(value)
-        for key, value in zip(GROUP_KEYS, group_key)
-    }
-
-
-def _dashboard_frame(df: pd.DataFrame) -> pd.DataFrame:
-    dashboard_df = df.copy()
-    for key in GROUP_KEYS:
-        if key not in dashboard_df.columns:
-            dashboard_df[key] = ""
-        dashboard_df[key] = dashboard_df[key].map(_cohort_value)
-    return dashboard_df
 
 # -----------------------------
 # 1. Grouping
 # -----------------------------
 def group_data(df: pd.DataFrame):
-    # Group by the same comparison cohort used by baseline gating.
-    return _dashboard_frame(df).groupby(list(GROUP_KEYS), dropna=False)
+    # Group only by model+GPU so each group produces a time-series line.
+    # config_id (commit SHA) is carried as a column for hover/color use.
+    keys = ["model_id", "gpu_type"]
+    return df.groupby(keys, dropna=False)
 
 # -----------------------------
 # 2. Plot builder
@@ -107,20 +46,15 @@ def build_plots(df: pd.DataFrame) -> tuple[list, list[dict[str, object]]]:
     figs = []
     skipped_metrics: list[dict[str, object]] = []
 
-    for group_key, g in group_data(df):
-        model_id, gpu_type = group_key[:2]
-        group_record = _group_record(group_key)
-        cohort_title = _cohort_title(group_key)
-        cohort_detail = _cohort_detail(group_key)
+    for (model_id, gpu_type), g in group_data(df):
         g = g.sort_values("timestamp")
 
         # One chart per metric so the y-axes aren't on wildly different scales
         for metric in METRICS:
             if metric not in g.columns:
                 skipped_metrics.append({
-                    **group_record,
-                    "cohort": cohort_title,
-                    "cohort_detail": cohort_detail,
+                    "model_id": model_id,
+                    "gpu_type": gpu_type,
                     "metric": metric,
                     "reason": "column missing from loaded records",
                     "records": len(g),
@@ -131,9 +65,8 @@ def build_plots(df: pd.DataFrame) -> tuple[list, list[dict[str, object]]]:
             non_null = int(g[metric].notna().sum())
             if non_null == 0:
                 skipped_metrics.append({
-                    **group_record,
-                    "cohort": cohort_title,
-                    "cohort_detail": cohort_detail,
+                    "model_id": model_id,
+                    "gpu_type": gpu_type,
                     "metric": metric,
                     "reason": "no non-null values in loaded records",
                     "records": len(g),
@@ -146,8 +79,8 @@ def build_plots(df: pd.DataFrame) -> tuple[list, list[dict[str, object]]]:
                 x="timestamp",
                 y=metric,
                 markers=True,
-                hover_data=["config_id", "commit_sha", *COMPARISON_COHORT_KEYS],
-                title=f"{model_id} | {gpu_type} | {cohort_title} | {cohort_detail} | {metric}",
+                hover_data=["config_id", "commit_sha"],
+                title=f"{model_id} | {gpu_type} | {metric}",
                 labels={"timestamp": "Time", metric: metric},
             )
             figs.append(fig)
@@ -162,7 +95,7 @@ def render_skipped_metrics(skipped_metrics: list[dict[str, object]]) -> str:
     rows = [
         "<h3>Skipped Metric Plots</h3>",
         "<table>",
-        ("<thead><tr><th>Model</th><th>GPU</th><th>Cohort</th><th>Metric</th>"
+        ("<thead><tr><th>Model</th><th>GPU</th><th>Metric</th>"
          "<th>Records</th><th>Non-null</th><th>Reason</th></tr></thead>"),
         "<tbody>",
     ]
@@ -171,7 +104,6 @@ def render_skipped_metrics(skipped_metrics: list[dict[str, object]]) -> str:
             "<tr>"
             f"<td>{escape(str(item['model_id']))}</td>"
             f"<td>{escape(str(item['gpu_type']))}</td>"
-            f"<td>{escape(str(item['cohort']))}<br><code>{escape(str(item['cohort_detail']))}</code></td>"
             f"<td>{escape(str(item['metric']))}</td>"
             f"<td>{item['records']}</td>"
             f"<td>{item['non_null']}</td>"
@@ -233,7 +165,6 @@ def main() -> None:
         print("Skipped metric plots:")
         for item in skipped_metrics:
             print(f"  - {item['model_id']} | {item['gpu_type']} | "
-                  f"{item['cohort']} | {item['cohort_detail']} | "
                   f"{item['metric']}: {item['reason']} "
                   f"({item['non_null']}/{item['records']} non-null)")
     print(f"Generated {len(figs)} metric plot(s)")

fastvideo/tests/performance/identity.py

@@ -1,435 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Comparable identity helpers for performance benchmark records."""
-
-from __future__ import annotations
-
-import dataclasses
-import enum
-import hashlib
-import importlib.metadata
-import json
-import os
-import platform
-import re
-from collections.abc import Mapping, Sequence
-from pathlib import Path
-from typing import Any
-
-import torch
-
-RECIPE_SCHEMA_VERSION = 1
-PROFILE_ID_LENGTH = 16
-_PATH_EXCLUDED_GENERATION_KEYS = {"output_path", "output_video_name"}
-_PROMPT_KEYS = {"prompt", "negative_prompt", "neg_prompt"}
-REQUIRED_BENCHMARK_IDENTITY_KEYS = (
-    "benchmark_id",
-    "workload_id",
-    "variant_id",
-    "benchmark_version",
-)
-_PROFILE_ENV_VARS = (
-    "CUDA_VISIBLE_DEVICES",
-    "FASTVIDEO_ATTENTION_BACKEND",
-    "IMAGE_VERSION",
-    "UV_TORCH_BACKEND",
-)
-_PACKAGE_DISTRIBUTIONS = {
-    "fastvideo_kernel": ("fastvideo-kernel", "fastvideo_kernel"),
-    "flash_attn": ("flash-attn", "flash_attn"),
-    "flash_attn_4": ("flash-attn-4",),
-    "flash_attention_fp4": ("flash-attention-fp4",),
-    "flashinfer": ("flashinfer-python", "flashinfer"),
-    "nvidia_cutlass_dsl": ("nvidia-cutlass-dsl", "cutlass-dsl"),
-    "sageattention": ("sageattention",),
-    "sageattn3": ("sageattn3",),
-    "triton": ("triton",),
-    "xformers": ("xformers",),
-}
-
-
-def canonical_json(payload: Any) -> str:
-    """Return deterministic, compact JSON for a recipe/profile payload."""
-
-    return json.dumps(
-        _canonicalize(payload),
-        sort_keys=True,
-        separators=(",", ":"),
-        ensure_ascii=True,
-    )
-
-
-def sha256_hexdigest(payload: str | bytes) -> str:
-    if isinstance(payload, str):
-        payload = payload.encode("utf-8")
-    return hashlib.sha256(payload).hexdigest()
-
-
-def recipe_fingerprint(recipe: Mapping[str, Any]) -> str:
-    # Runtime-resolved values are audit metadata, not identity: hashing the
-    # resolved HF snapshot revision would churn the cohort (and via the
-    # no-baseline init path, silently reset the rolling baseline) on ANY
-    # upstream repo commit, including model-card edits with unchanged
-    # weights. The fingerprint is pinned to declared inputs only; the
-    # resolved values stay in the stored recipe for auditability.
-    pruned = {
-        key: (dict(value) if isinstance(value, Mapping) else value)
-        for key, value in recipe.items()
-    }
-    model = pruned.get("model")
-    if isinstance(model, dict):
-        model.pop("resolved_revision", None)
-    attention = pruned.get("attention")
-    if isinstance(attention, dict):
-        # Same reasoning as resolved_revision: with requested_backend="auto",
-        # a silent runtime fallback (e.g. a broken flash-attn wheel) must NOT
-        # open a fresh ungated cohort and seed degraded numbers as the new
-        # baseline — it should fail the regression gate in the old cohort.
-        # Intentional backend changes are declared via requested_backend,
-        # which stays in the hash.
-        attention.pop("resolved_backend", None)
-    return sha256_hexdigest(canonical_json(pruned))
-
-
-def profile_id(prefix: str, profile: Mapping[str, Any]) -> str:
-    return f"{prefix}-{sha256_hexdigest(canonical_json(profile))[:PROFILE_ID_LENGTH]}"
-
-
-def hardware_profile_id(profile: Mapping[str, Any]) -> str:
-    return profile_id("hw", profile)
-
-
-def software_profile_id(profile: Mapping[str, Any]) -> str:
-    return profile_id("sw", profile)
-
-
-def environment_fingerprint(metadata: Mapping[str, Any]) -> str:
-    return profile_id("env", metadata)
-
-
-def benchmark_identity_from_config(cfg: Mapping[str, Any]) -> dict[str, Any]:
-    missing = [
-        key for key in REQUIRED_BENCHMARK_IDENTITY_KEYS
-        if _none_if_empty(cfg.get(key)) is None
-    ]
-    if missing:
-        raise ValueError("Benchmark config missing required identity fields: " + ", ".join(missing))
-    return {
-        key: cfg[key]
-        for key in REQUIRED_BENCHMARK_IDENTITY_KEYS
-    }
-
-
-def build_recipe_from_benchmark_config(
-    cfg: Mapping[str, Any],
-    *,
-    attention_backend: str | None = None,
-    resolved_attention_backend: Any | None = None,
-    resolved_model_revision: Any | None = None,
-    measured_prompts: Sequence[Any] | None = None,
-) -> dict[str, Any]:
-    """Build a deterministic recipe document from a benchmark config.
-
-    The recipe captures fields that describe what benchmark workload was run. It
-    intentionally excludes timings, timestamps, commit metadata, and output paths.
-    """
-
-    model = dict(cfg.get("model") or {})
-    init_kwargs = dict(cfg.get("init_kwargs") or {})
-    generation_kwargs = dict(cfg.get("generation_kwargs") or {})
-    benchmark_identity = benchmark_identity_from_config(cfg)
-    prompts = list(measured_prompts) if measured_prompts is not None else list(
-        cfg.get("test_prompts") or ["A cinematic video."])
-    if attention_backend is None:
-        attention_backend = os.environ.get("FASTVIDEO_ATTENTION_BACKEND")
-
-    negative_prompt = generation_kwargs.pop("negative_prompt", generation_kwargs.pop("neg_prompt", None))
-    generation_recipe = {
-        key: value
-        for key, value in generation_kwargs.items()
-        if key not in _PATH_EXCLUDED_GENERATION_KEYS and key not in _PROMPT_KEYS
-    }
-
-    return {
-        "recipe_schema_version": RECIPE_SCHEMA_VERSION,
-        "benchmark": benchmark_identity,
-        "model": {
-            "model_path": normalize_model_path(model.get("model_path")),
-            "revision": _none_if_empty(model.get("revision") or cfg.get("revision")),
-            "resolved_revision": _none_if_empty(resolved_model_revision),
-        },
-        "init_kwargs": init_kwargs,
-        "generation_kwargs": generation_recipe,
-        "inputs": {
-            "prompt_count": len(prompts),
-            "prompt_sha256": [_optional_digest(prompt) for prompt in prompts],
-            "negative_prompt_sha256": _optional_digest(negative_prompt),
-        },
-        "attention": {
-            "requested_backend": _none_if_empty(attention_backend) or "auto",
-            "resolved_backend": _none_if_empty(resolved_attention_backend),
-        },
-    }
-
-
-def normalize_model_path(value: Any) -> str | None:
-    if value is None:
-        return None
-    text = os.fspath(value) if isinstance(value, os.PathLike) else str(value)
-    text = text.strip()
-    if not text:
-        return None
-    if _looks_like_local_path(text):
-        return Path(text).expanduser().as_posix()
-    return re.sub(r"/+", "/", text)
-
-
-def resolved_revision_from_model_path(value: Any) -> str | None:
-    normalized = normalize_model_path(value)
-    if normalized is None:
-        return None
-    parts = Path(normalized).parts
-    for index, part in enumerate(parts[:-1]):
-        if part == "snapshots":
-            revision = parts[index + 1]
-            if re.fullmatch(r"[0-9a-f]{40}", revision):
-                return revision
-    return None
-
-
-def hardware_profile(
-    *,
-    num_gpus: int | None = None,
-    gpu_devices: Sequence[Mapping[str, Any]] | None = None,
-    device_type: str | None = None,
-    interconnect: str | None = None,
-) -> dict[str, Any]:
-    """Return normalized hardware cohort metadata.
-
-    ``gpu_devices`` is an injection point for tests and non-CUDA callers.
-    """
-
-    if gpu_devices is not None:
-        devices = [_normalize_gpu_device(device) for device in gpu_devices]
-        return {
-            "device_type": device_type or "cuda",
-            "gpu_count": num_gpus if num_gpus is not None else len(devices),
-            "gpus": devices,
-            "interconnect": interconnect or "unknown",
-        }
-
-    if not torch.cuda.is_available():
-        return {
-            "device_type": device_type or "cpu",
-            "gpu_count": 0,
-            "gpus": [],
-            "interconnect": interconnect or "none",
-        }
-
-    gpu_count = int(num_gpus if num_gpus is not None else torch.cuda.device_count())
-    devices = []
-    cuda_device_count = torch.cuda.device_count()
-    for device_id in range(gpu_count):
-        if device_id >= cuda_device_count:
-            devices.append(_normalize_gpu_device({}))
-            continue
-        props = torch.cuda.get_device_properties(device_id)
-        capability = torch.cuda.get_device_capability(device_id)
-        devices.append(
-            _normalize_gpu_device({
-                "name": props.name,
-                "memory_bytes": props.total_memory,
-                "compute_capability": f"{capability[0]}.{capability[1]}",
-            }))
-
-    return {
-        "device_type": device_type or "cuda",
-        "gpu_count": gpu_count,
-        "gpus": devices,
-        "interconnect": interconnect or _detect_cuda_interconnect(gpu_count),
-    }
-
-
-def software_profile(
-    *,
-    python_version: str | None = None,
-    torch_version: str | None = None,
-    cuda_version: str | None = None,
-    package_versions: Mapping[str, str | None] | None = None,
-) -> dict[str, Any]:
-    """Return normalized software cohort metadata.
-
-    Python, PyTorch, and CUDA use major/minor cohorts. Attention and kernel
-    packages keep exact versions because patch releases can change performance.
-    """
-
-    versions = dict(package_versions) if package_versions is not None else _installed_package_versions()
-    return {
-        "python": _major_minor(python_version or platform.python_version()),
-        "pytorch": _major_minor(torch_version or torch.__version__),
-        "cuda": _major_minor(cuda_version or torch.version.cuda),
-        "packages": {
-            name: str(version)
-            for name, version in sorted(versions.items())
-            if version
-        },
-    }
-
-
-def environment_metadata(
-    *,
-    env: Mapping[str, str] | None = None,
-    package_versions: Mapping[str, str | None] | None = None,
-    hardware: Mapping[str, Any] | None = None,
-    software: Mapping[str, Any] | None = None,
-) -> dict[str, Any]:
-    """Return audit metadata kept separate from comparable identity keys."""
-
-    source_env = env if env is not None else os.environ
-    full_package_versions = dict(package_versions) if package_versions is not None else _installed_package_versions()
-    return {
-        "python": {
-            "version": platform.python_version(),
-            "implementation": platform.python_implementation(),
-        },
-        "platform": {
-            "system": platform.system(),
-            "machine": platform.machine(),
-            "release": platform.release(),
-        },
-        "torch": {
-            "version": torch.__version__,
-            "cuda": torch.version.cuda,
-            "cudnn": torch.backends.cudnn.version(),
-        },
-        "env": {
-            key: source_env.get(key)
-            for key in _PROFILE_ENV_VARS
-            if source_env.get(key) is not None
-        },
-        "packages": {
-            key: value
-            for key, value in sorted(full_package_versions.items())
-            if value
-        },
-        "hardware_profile": dict(hardware) if hardware is not None else hardware_profile(),
-        "software_profile": dict(software) if software is not None else software_profile(
-            package_versions=full_package_versions),
-    }
-
-
-def _canonicalize(value: Any) -> Any:
-    if dataclasses.is_dataclass(value) and not isinstance(value, type):
-        return _canonicalize(dataclasses.asdict(value))
-    if isinstance(value, Mapping):
-        return {
-            str(key): _canonicalize(value[key])
-            for key in sorted(value, key=lambda item: str(item))
-        }
-    if isinstance(value, (set, frozenset)):
-        return [_canonicalize(item) for item in sorted(value, key=lambda item: str(item))]
-    if isinstance(value, tuple):
-        return [_canonicalize(item) for item in value]
-    if isinstance(value, list):
-        return [_canonicalize(item) for item in value]
-    if isinstance(value, enum.Enum):
-        return value.name
-    if isinstance(value, Path):
-        return value.expanduser().as_posix()
-    if isinstance(value, os.PathLike):
-        return os.fspath(value)
-    if isinstance(value, type):
-        return f"{value.__module__}.{value.__qualname__}"
-    if isinstance(value, torch.dtype):
-        return str(value).removeprefix("torch.")
-    return value
-
-
-def _none_if_empty(value: Any) -> Any | None:
-    if value is None:
-        return None
-    if isinstance(value, str) and not value.strip():
-        return None
-    return value
-
-
-def _optional_digest(value: Any) -> str | None:
-    if value is None:
-        return None
-    return sha256_hexdigest(str(value))
-
-
-def _looks_like_local_path(value: str) -> bool:
-    if value.startswith(("/", "./", "../", "~")):
-        return True
-    looks_like_hf_repo = re.match(r"^[A-Za-z0-9_.-]+/[A-Za-z0-9_.-]+$", value)
-    has_separator = "/" in value or "\\" in value or os.sep in value
-    return has_separator and looks_like_hf_repo is None
-
-
-def _normalize_gpu_device(device: Mapping[str, Any]) -> dict[str, Any]:
-    memory_gb = device.get("memory_gb")
-    if memory_gb is None and device.get("memory_bytes") is not None:
-        memory_gb = round(float(device["memory_bytes"]) / (1024**3))
-    return {
-        "name": _none_if_empty(device.get("name")) or "unknown",
-        "memory_gb": memory_gb,
-        "compute_capability": _none_if_empty(device.get("compute_capability")),
-    }
-
-
-def _detect_cuda_interconnect(gpu_count: int) -> str:
-    if gpu_count <= 1:
-        return "single_gpu"
-    try:
-        from fastvideo.platforms import current_platform
-
-        if current_platform.is_cuda() and current_platform.is_full_nvlink(list(range(gpu_count))):
-            return "full_nvlink"
-        return "none_or_partial"
-    except Exception:
-        return "unknown"
-
-
-def _installed_package_versions() -> dict[str, str | None]:
-    return {
-        name: _distribution_version(*distribution_names)
-        for name, distribution_names in _PACKAGE_DISTRIBUTIONS.items()
-    }
-
-
-def _distribution_version(*names: str) -> str | None:
-    for name in names:
-        try:
-            return importlib.metadata.version(name)
-        except importlib.metadata.PackageNotFoundError:
-            continue
-    return None
-
-
-def _major_minor(version: str | None) -> str | None:
-    if not version:
-        return None
-    match = re.search(r"(\d+)\.(\d+)", version)
-    if match is None:
-        return version
-    return f"{match.group(1)}.{match.group(2)}"
-
-
-__all__ = [
-    "RECIPE_SCHEMA_VERSION",
-    "REQUIRED_BENCHMARK_IDENTITY_KEYS",
-    "benchmark_identity_from_config",
-    "canonical_json",
-    "environment_fingerprint",
-    "environment_metadata",
-    "hardware_profile",
-    "hardware_profile_id",
-    "normalize_model_path",
-    "profile_id",
-    "recipe_fingerprint",
-    "resolved_revision_from_model_path",
-    "build_recipe_from_benchmark_config",
-    "sha256_hexdigest",
-    "software_profile",
-    "software_profile_id",
-]

fastvideo/tests/performance/test_compare_baseline_policy.py

@@ -1,7 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
 
-import pytest
-
 from fastvideo.tests.performance import compare_baseline
 from fastvideo.performance.metric_policy import resolve_metric_policies
 
@@ -121,110 +119,6 @@ def test_boolean_regression_threshold_values_are_ignored():
     assert latency.threshold_percent == 0.08
     assert latency.threshold_absolute == 0.5
     assert latency.gated is False
-def test_normalized_record_preserves_identity_metadata(monkeypatch):
-    monkeypatch.setenv("PERF_RUN_SOURCE", "pr")
-    raw = _raw_result()
-    raw.update({
-        "workload_id": "wan-t2v",
-        "variant_id": "1.3b-sp2",
-        "benchmark_version": 2,
-        "recipe": {
-            "recipe_schema_version": 1,
-        },
-        "recipe_fingerprint": "recipe-1",
-        "hardware_profile": {
-            "gpu_count": 1,
-        },
-        "hardware_profile_id": "hw-1",
-        "software_profile": {
-            "python": "3.12",
-        },
-        "software_profile_id": "sw-1",
-        "environment_metadata": {
-            "env": {
-                "IMAGE_VERSION": "latest",
-            },
-        },
-        "environment_fingerprint": "env-1",
-    })
-
-    record = compare_baseline.normalize_performance_result(raw)
-
-    assert record["workload_id"] == "wan-t2v"
-    assert record["variant_id"] == "1.3b-sp2"
-    assert record["benchmark_version"] == 2
-    assert record["recipe"] == {"recipe_schema_version": 1}
-    assert record["recipe_fingerprint"] == "recipe-1"
-    assert record["hardware_profile"] == {"gpu_count": 1}
-    assert record["hardware_profile_id"] == "hw-1"
-    assert record["software_profile"] == {"python": "3.12"}
-    assert record["software_profile_id"] == "sw-1"
-    assert record["environment_metadata"] == {"env": {"IMAGE_VERSION": "latest"}}
-    assert record["environment_fingerprint"] == "env-1"
-
-
-def test_normalized_record_reads_identity_labels_from_recipe(monkeypatch):
-    monkeypatch.setenv("PERF_RUN_SOURCE", "pr")
-    raw = _raw_result()
-    raw["recipe"] = {
-        "benchmark": {
-            "workload_id": "wan-t2v",
-            "variant_id": "1.3b-sp2",
-            "benchmark_version": 2,
-        },
-    }
-
-    record = compare_baseline.normalize_performance_result(raw)
-
-    assert record["workload_id"] == "wan-t2v"
-    assert record["variant_id"] == "1.3b-sp2"
-    assert record["benchmark_version"] == 2
-
-
-def test_comparison_identity_filters_use_full_issue_key():
-    record = {
-        "workload_id": "wan-t2v",
-        "variant_id": "1.3b-sp2",
-        "benchmark_version": 2,
-        "recipe_fingerprint": "recipe-1",
-        "hardware_profile_id": "hw-1",
-        "software_profile_id": "sw-1",
-    }
-
-    assert compare_baseline._comparison_identity_filters(record) == {
-        "workload_id": "wan-t2v",
-        "variant_id": "1.3b-sp2",
-        "benchmark_version": "2",
-        "recipe_fingerprint": "recipe-1",
-        "hardware_profile_id": "hw-1",
-        "software_profile_id": "sw-1",
-    }
-
-
-def test_comparison_identity_filters_require_full_issue_key():
-    record = {
-        "workload_id": "wan-t2v",
-        "variant_id": "1.3b-sp2",
-        "benchmark_version": 0,
-        "recipe_fingerprint": "recipe-1",
-        "hardware_profile_id": "hw-1",
-    }
-
-    with pytest.raises(ValueError, match="software_profile_id"):
-        compare_baseline._comparison_identity_filters(record)
-
-
-def test_comparison_identity_filters_keep_zero_version():
-    record = {
-        "workload_id": "wan-t2v",
-        "variant_id": "1.3b-sp2",
-        "benchmark_version": 0,
-        "recipe_fingerprint": "recipe-1",
-        "hardware_profile_id": "hw-1",
-        "software_profile_id": "sw-1",
-    }
-
-    assert compare_baseline._comparison_identity_filters(record)["benchmark_version"] == "0"
 
 
 def test_baseline_eligibility_only_for_successful_scheduled_main():

fastvideo/tests/performance/test_dashboard_api.py

@@ -56,34 +56,8 @@ def _record(model_id, gpu_type, ts, commit, latency, throughput, success=True, *
 
 def test_summary_endpoint_returns_latest_group_status():
     app = create_app(FakeStore([
-        _record(
-            "wan",
-            "NVIDIA L40S",
-            "2026-01-01T00:00:00+00:00",
-            "a" * 40,
-            10.0,
-            10.0,
-            workload_id="wan-t2v",
-            variant_id="1.3b-sp2",
-            benchmark_version=2,
-            recipe_fingerprint="recipe-a",
-            hardware_profile_id="hw-l40s",
-            software_profile_id="sw-cu130",
-        ),
-        _record(
-            "wan",
-            "NVIDIA L40S",
-            "2026-01-02T00:00:00+00:00",
-            "b" * 40,
-            11.0,
-            9.0,
-            workload_id="wan-t2v",
-            variant_id="1.3b-sp2",
-            benchmark_version=2,
-            recipe_fingerprint="recipe-a",
-            hardware_profile_id="hw-l40s",
-            software_profile_id="sw-cu130",
-        ),
+        _record("wan", "NVIDIA L40S", "2026-01-01T00:00:00+00:00", "a" * 40, 10.0, 10.0),
+        _record("wan", "NVIDIA L40S", "2026-01-02T00:00:00+00:00", "b" * 40, 11.0, 9.0),
     ]))
     client = TestClient(app)
 
@@ -97,10 +71,6 @@ def test_summary_endpoint_returns_latest_group_status():
     assert body["rows"][0]["metrics"]["latency"]["threshold_exceeded"] is True
     assert body["rows"][0]["threshold_exceeded_metrics"] == ["latency", "throughput"]
     assert body["rows"][0]["computed_regression_status"] == "fail"
-    assert body["rows"][0]["workload_id"] == "wan-t2v"
-    assert body["rows"][0]["variant_id"] == "1.3b-sp2"
-    assert body["rows"][0]["benchmark_version"] == 2
-    assert body["rows"][0]["recipe_fingerprint"] == "recipe-a"
 
 
 def test_summary_status_is_independent_of_days_window():

fastvideo/tests/performance/test_dashboard_plotly.py

@@ -1,98 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-
-import pandas as pd
-
-from fastvideo.tests.performance import dashboard
-
-
-def _record(recipe_fingerprint="recipe-a", software_profile_id="sw-a", **overrides):
-    record = {
-        "model_id": "wan-t2v-1.3b-2gpu",
-        "gpu_type": "NVIDIA L40S",
-        "workload_id": "wan-t2v",
-        "variant_id": "1.3b-sp2",
-        "benchmark_version": 2,
-        "recipe_fingerprint": recipe_fingerprint,
-        "hardware_profile_id": "hw-l40s",
-        "software_profile_id": software_profile_id,
-        "timestamp": "2026-01-01T00:00:00+00:00",
-        "commit_sha": "a" * 40,
-        "config_id": "aaaaaaa",
-        "latency": 10.0,
-        "throughput": 4.5,
-        "memory": 10000.0,
-        "text_encoder_time_s": 2.0,
-        "dit_time_s": 8.0,
-        "vae_decode_time_s": 3.0,
-    }
-    record.update(overrides)
-    return record
-
-
-def test_group_data_uses_full_comparison_cohort():
-    df = pd.DataFrame([
-        _record(recipe_fingerprint="recipe-a", software_profile_id="sw-a"),
-        _record(recipe_fingerprint="recipe-b", software_profile_id="sw-b"),
-    ])
-
-    groups = list(dashboard.group_data(df))
-
-    assert len(groups) == 2
-    assert {key[5] for key, _group in groups} == {"recipe-a", "recipe-b"}
-    assert {key[7] for key, _group in groups} == {"sw-a", "sw-b"}
-
-
-def test_group_data_fills_legacy_cohort_columns():
-    df = pd.DataFrame([{
-        "model_id": "legacy-wan",
-        "gpu_type": "NVIDIA L40S",
-        "timestamp": "2026-01-01T00:00:00+00:00",
-        "commit_sha": "a" * 40,
-        "config_id": "aaaaaaa",
-        "latency": 10.0,
-    }])
-
-    groups = list(dashboard.group_data(df))
-
-    assert len(groups) == 1
-    assert groups[0][0][2:] == ("", "", "", "", "", "")
-
-
-def test_build_plots_labels_distinct_cohorts():
-    df = pd.DataFrame([
-        _record(recipe_fingerprint="recipe-a", software_profile_id="sw-a"),
-        _record(recipe_fingerprint="recipe-b", software_profile_id="sw-b"),
-    ])
-
-    figs, skipped_metrics = dashboard.build_plots(df)
-    titles = [fig.layout.title.text for fig in figs]
-
-    assert len(figs) == len(dashboard.METRICS) * 2
-    assert skipped_metrics == []
-    assert any("recipe recipe-a | hw-l40s | sw-a" in title for title in titles)
-    assert any("recipe recipe-b | hw-l40s | sw-b" in title for title in titles)
-
-
-def test_skipped_metric_table_includes_cohort_identity():
-    df = pd.DataFrame([{
-        "model_id": "wan-t2v-1.3b-2gpu",
-        "gpu_type": "NVIDIA L40S",
-        "workload_id": "wan-t2v",
-        "variant_id": "1.3b-sp2",
-        "benchmark_version": 2,
-        "recipe_fingerprint": "recipe-a",
-        "hardware_profile_id": "hw-l40s",
-        "software_profile_id": "sw-a",
-        "timestamp": "2026-01-01T00:00:00+00:00",
-        "commit_sha": "a" * 40,
-        "config_id": "aaaaaaa",
-        "latency": 10.0,
-    }])
-
-    _figs, skipped_metrics = dashboard.build_plots(df)
-    html = dashboard.render_skipped_metrics(skipped_metrics[:1])
-
-    assert skipped_metrics[0]["cohort"] == "wan-t2v / 1.3b-sp2 / v2"
-    assert skipped_metrics[0]["cohort_detail"] == "recipe recipe-a | hw-l40s | sw-a"
-    assert "wan-t2v / 1.3b-sp2 / v2" in html
-    assert "recipe recipe-a | hw-l40s | sw-a" in html

fastvideo/tests/performance/test_dashboard_service.py

@@ -142,111 +142,6 @@ def test_build_latest_summary_separates_informational_threshold_crossing():
     assert rows[0]["threshold_exceeded_metrics"] == ["latency"]
     assert rows[0]["failing_metrics"] == []
     assert rows[0]["computed_regression_status"] == "pass"
-def test_build_latest_summary_keeps_identity_cohorts_separate():
-    records = [
-        _record(
-            "2026-01-01T00:00:00+00:00",
-            "a" * 40,
-            10.0,
-            10.0,
-            recipe_fingerprint="recipe-a",
-            workload_id="wan-t2v",
-            variant_id="1.3b-sp2",
-            benchmark_version=2,
-            hardware_profile_id="hw-l40s",
-            software_profile_id="sw-cu130",
-            run_source="scheduled_main",
-            baseline_eligible=True,
-        ),
-        _record(
-            "2026-01-02T00:00:00+00:00",
-            "b" * 40,
-            20.0,
-            5.0,
-            recipe_fingerprint="recipe-b",
-            workload_id="wan-t2v",
-            variant_id="1.3b-sp2",
-            benchmark_version=2,
-            hardware_profile_id="hw-l40s",
-            software_profile_id="sw-cu130",
-            run_source="scheduled_main",
-            baseline_eligible=True,
-        ),
-        _record(
-            "2026-01-03T00:00:00+00:00",
-            "c" * 40,
-            22.0,
-            4.5,
-            recipe_fingerprint="recipe-b",
-            workload_id="wan-t2v",
-            variant_id="1.3b-sp2",
-            benchmark_version=2,
-            hardware_profile_id="hw-l40s",
-            software_profile_id="sw-cu130",
-        ),
-    ]
-
-    rows = build_latest_summary(records)
-    recipe_b_row = next(row for row in rows if row["recipe_fingerprint"] == "recipe-b")
-
-    assert len(rows) == 2
-    assert recipe_b_row["baseline_n"] == 1
-    assert recipe_b_row["metrics"]["latency"]["baseline"] == 20.0
-    assert recipe_b_row["workload_id"] == "wan-t2v"
-    assert recipe_b_row["variant_id"] == "1.3b-sp2"
-    assert recipe_b_row["benchmark_version"] == 2
-
-
-def test_build_latest_summary_keeps_variant_versions_separate():
-    records = [
-        _record(
-            "2026-01-01T00:00:00+00:00",
-            "a" * 40,
-            10.0,
-            10.0,
-            workload_id="wan-t2v",
-            variant_id="1.3b-sp2",
-            benchmark_version=1,
-            recipe_fingerprint="recipe-a",
-            hardware_profile_id="hw-l40s",
-            software_profile_id="sw-cu130",
-            run_source="scheduled_main",
-            baseline_eligible=True,
-        ),
-        _record(
-            "2026-01-02T00:00:00+00:00",
-            "b" * 40,
-            20.0,
-            5.0,
-            workload_id="wan-t2v",
-            variant_id="1.3b-sp2",
-            benchmark_version=2,
-            recipe_fingerprint="recipe-a",
-            hardware_profile_id="hw-l40s",
-            software_profile_id="sw-cu130",
-            run_source="scheduled_main",
-            baseline_eligible=True,
-        ),
-        _record(
-            "2026-01-03T00:00:00+00:00",
-            "c" * 40,
-            22.0,
-            4.5,
-            workload_id="wan-t2v",
-            variant_id="1.3b-sp2",
-            benchmark_version=2,
-            recipe_fingerprint="recipe-a",
-            hardware_profile_id="hw-l40s",
-            software_profile_id="sw-cu130",
-        ),
-    ]
-
-    rows = build_latest_summary(records)
-    version_2_row = next(row for row in rows if row["benchmark_version"] == 2)
-
-    assert len(rows) == 2
-    assert version_2_row["baseline_n"] == 1
-    assert version_2_row["metrics"]["latency"]["baseline"] == 20.0
 
 
 def test_filter_records_and_trends_preserve_metric_points():
@@ -323,78 +218,3 @@ def test_load_records_can_filter_baseline_eligible_records(tmp_path):
         "2026-01-02T00:00:00+00:00",
         "2026-01-03T00:00:00+00:00",
     }
-
-
-def test_load_records_for_model_filters_identity_cohort(tmp_path):
-    model_dir = tmp_path / "wan"
-    model_dir.mkdir()
-    (model_dir / "matching.json").write_text(
-        """
-        {
-          "model_id": "wan",
-          "gpu_type": "NVIDIA L40S",
-          "timestamp": "2026-01-01T00:00:00+00:00",
-          "success": true,
-          "baseline_eligible": true,
-          "workload_id": "wan-t2v",
-          "variant_id": "1.3b-sp2",
-          "benchmark_version": 2,
-          "recipe_fingerprint": "recipe-a",
-          "hardware_profile_id": "hw-l40s",
-          "software_profile_id": "sw-cu130"
-        }
-        """,
-        encoding="utf-8",
-    )
-    (model_dir / "other_recipe.json").write_text(
-        """
-        {
-          "model_id": "wan",
-          "gpu_type": "NVIDIA L40S",
-          "timestamp": "2026-01-02T00:00:00+00:00",
-          "success": true,
-          "baseline_eligible": true,
-          "workload_id": "wan-t2v",
-          "variant_id": "1.3b-sp2",
-          "benchmark_version": 2,
-          "recipe_fingerprint": "recipe-b",
-          "hardware_profile_id": "hw-l40s",
-          "software_profile_id": "sw-cu130"
-        }
-        """,
-        encoding="utf-8",
-    )
-    (model_dir / "other_version.json").write_text(
-        """
-        {
-          "model_id": "wan",
-          "gpu_type": "NVIDIA L40S",
-          "timestamp": "2026-01-03T00:00:00+00:00",
-          "success": true,
-          "baseline_eligible": true,
-          "workload_id": "wan-t2v",
-          "variant_id": "1.3b-sp2",
-          "benchmark_version": 3,
-          "recipe_fingerprint": "recipe-a",
-          "hardware_profile_id": "hw-l40s",
-          "software_profile_id": "sw-cu130"
-        }
-        """,
-        encoding="utf-8",
-    )
-
-    records = hf_store.load_records_for_model(
-        str(tmp_path),
-        "wan",
-        "NVIDIA L40S",
-        workload_id="wan-t2v",
-        variant_id="1.3b-sp2",
-        benchmark_version="2",
-        recipe_fingerprint="recipe-a",
-        hardware_profile_id="hw-l40s",
-        software_profile_id="sw-cu130",
-        baseline_eligible_only=True,
-    )
-
-    assert len(records) == 1
-    assert records[0]["recipe_fingerprint"] == "recipe-a"

fastvideo/tests/performance/test_identity.py

@@ -1,422 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-
-from copy import deepcopy
-
-import pytest
-import torch
-
-from fastvideo.tests.performance import identity as identity_module
-from fastvideo.tests.performance.identity import (
-    build_recipe_from_benchmark_config,
-    canonical_json,
-    environment_fingerprint,
-    environment_metadata,
-    hardware_profile,
-    hardware_profile_id,
-    recipe_fingerprint,
-    resolved_revision_from_model_path,
-    software_profile,
-    software_profile_id,
-)
-
-
-def _benchmark_config():
-    return {
-        "benchmark_id": "wan-t2v-1.3b-2gpu",
-        "workload_id": "wan-t2v",
-        "variant_id": "1.3b-sp2",
-        "benchmark_version": 2,
-        "model": {
-            "model_path": "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
-            "model_short_name": "Wan2.1-T2V-1.3B",
-        },
-        "init_kwargs": {
-            "num_gpus": 2,
-            "sp_size": 2,
-            "tp_size": 1,
-            "vae_sp": True,
-            "vae_tiling": True,
-            "text_encoder_precisions": ["fp32"],
-        },
-        "generation_kwargs": {
-            "height": 480,
-            "width": 832,
-            "num_frames": 45,
-            "num_inference_steps": 4,
-            "guidance_scale": 3,
-            "embedded_cfg_scale": 6,
-            "seed": 1024,
-            "fps": 24,
-            "neg_prompt": "low quality",
-        },
-        "test_prompts": ["A cinematic video."],
-    }
-
-
-def _fingerprint(cfg, attention_backend="FLASH_ATTN"):
-    recipe = build_recipe_from_benchmark_config(cfg, attention_backend=attention_backend)
-    return recipe_fingerprint(recipe)
-
-
-def test_canonical_json_is_stable_for_mapping_order():
-    left = {"b": [2, 1], "a": {"z": "last", "m": "middle"}}
-    right = {"a": {"m": "middle", "z": "last"}, "b": [2, 1]}
-
-    assert canonical_json(left) == canonical_json(right)
-
-
-def test_canonical_json_handles_sets_deterministically():
-    assert canonical_json({"values": {"b", "a"}}) == '{"values":["a","b"]}'
-    assert canonical_json({"values": frozenset(("b", "a"))}) == '{"values":["a","b"]}'
-
-
-def test_same_config_produces_same_recipe_fingerprint():
-    cfg = _benchmark_config()
-
-    assert _fingerprint(cfg) == _fingerprint(deepcopy(cfg))
-
-
-def test_recipe_includes_first_class_benchmark_identity():
-    recipe = build_recipe_from_benchmark_config(_benchmark_config())
-
-    assert recipe["benchmark"] == {
-        "benchmark_id": "wan-t2v-1.3b-2gpu",
-        "workload_id": "wan-t2v",
-        "variant_id": "1.3b-sp2",
-        "benchmark_version": 2,
-    }
-
-
-def test_recipe_requires_first_class_benchmark_identity():
-    cfg = _benchmark_config()
-    del cfg["workload_id"]
-
-    with pytest.raises(ValueError, match="workload_id"):
-        build_recipe_from_benchmark_config(cfg)
-
-
-def test_semantically_equivalent_sequence_forms_match():
-    cfg = _benchmark_config()
-    equivalent = deepcopy(cfg)
-    equivalent["init_kwargs"]["text_encoder_precisions"] = tuple(
-        equivalent["init_kwargs"]["text_encoder_precisions"])
-
-    assert _fingerprint(cfg) == _fingerprint(equivalent)
-
-
-def test_output_path_does_not_change_recipe_fingerprint():
-    cfg = _benchmark_config()
-    with_output_path = deepcopy(cfg)
-    with_output_path["generation_kwargs"]["output_path"] = "/tmp/generated"
-
-    assert _fingerprint(cfg) == _fingerprint(with_output_path)
-
-
-def test_runtime_resolved_values_do_not_change_recipe_fingerprint():
-    # Runtime-resolved values are audit metadata: an upstream repo commit
-    # (even a README-only edit) or a silent attention-backend fallback must
-    # not churn the cohort — that would open a fresh ungated cohort and let
-    # degraded numbers seed a new baseline. Declared inputs (including
-    # requested_backend) stay in the hash.
-    cfg = _benchmark_config()
-    recipe = build_recipe_from_benchmark_config(
-        cfg, attention_backend="FLASH_ATTN",
-        resolved_attention_backend="FLASH_ATTN",
-        resolved_model_revision="aaaa1111")
-    changed = build_recipe_from_benchmark_config(
-        cfg, attention_backend="FLASH_ATTN",
-        resolved_attention_backend="TORCH_SDPA",
-        resolved_model_revision="bbbb2222")
-
-    assert recipe["model"]["resolved_revision"] == "aaaa1111"  # kept for audit
-    assert recipe["attention"]["resolved_backend"] == "FLASH_ATTN"  # kept for audit
-    assert recipe_fingerprint(recipe) == recipe_fingerprint(changed)
-
-
-def test_measured_prompt_override_ignores_unused_configured_prompts():
-    cfg = _benchmark_config()
-    changed = deepcopy(cfg)
-    changed["test_prompts"] = [
-        "Unused prompt that should not affect this measured workload.",
-        cfg["test_prompts"][0],
-    ]
-
-    recipe = build_recipe_from_benchmark_config(cfg, measured_prompts=[cfg["test_prompts"][0]])
-    changed_recipe = build_recipe_from_benchmark_config(changed, measured_prompts=[cfg["test_prompts"][0]])
-
-    assert recipe_fingerprint(recipe) == recipe_fingerprint(changed_recipe)
-
-
-def test_num_inference_steps_changes_recipe_fingerprint():
-    cfg = _benchmark_config()
-    changed = deepcopy(cfg)
-    changed["generation_kwargs"]["num_inference_steps"] = 8
-
-    assert _fingerprint(cfg) != _fingerprint(changed)
-
-
-def test_benchmark_version_changes_recipe_fingerprint():
-    cfg = _benchmark_config()
-    changed = deepcopy(cfg)
-    changed["benchmark_version"] = 3
-
-    assert _fingerprint(cfg) != _fingerprint(changed)
-
-
-def test_attention_backend_changes_recipe_fingerprint():
-    cfg = _benchmark_config()
-
-    assert _fingerprint(cfg, attention_backend="FLASH_ATTN") != _fingerprint(
-        cfg, attention_backend="TORCH_SDPA")
-
-
-def test_precision_changes_recipe_fingerprint():
-    cfg = _benchmark_config()
-    changed = deepcopy(cfg)
-    changed["init_kwargs"]["text_encoder_precisions"] = ["bf16"]
-
-    assert _fingerprint(cfg) != _fingerprint(changed)
-
-
-def test_dimensions_and_seed_change_recipe_fingerprint():
-    cfg = _benchmark_config()
-    different_dimensions = deepcopy(cfg)
-    different_dimensions["generation_kwargs"]["width"] = 1024
-    different_seed = deepcopy(cfg)
-    different_seed["generation_kwargs"]["seed"] = 2048
-
-    assert _fingerprint(cfg) != _fingerprint(different_dimensions)
-    assert _fingerprint(cfg) != _fingerprint(different_seed)
-
-
-def test_distributed_layout_changes_recipe_fingerprint():
-    cfg = _benchmark_config()
-    changed = deepcopy(cfg)
-    changed["init_kwargs"]["sp_size"] = 1
-    changed["init_kwargs"]["tp_size"] = 2
-
-    assert _fingerprint(cfg) != _fingerprint(changed)
-
-
-def test_software_profile_uses_exact_attention_kernel_package_versions_for_id():
-    profile_a = software_profile(
-        python_version="3.12.4",
-        torch_version="2.12.0+cu130",
-        cuda_version="13.0",
-        package_versions={
-            "triton": "3.4.1",
-            "fastvideo_kernel": "0.3.2",
-            "flash_attn_4": "4.0.0.dev0",
-            "flashinfer": "0.2.11",
-            "nvidia_cutlass_dsl": "4.5.0",
-        },
-    )
-    profile_b = software_profile(
-        python_version="3.12.5",
-        torch_version="2.12.1+cu130",
-        cuda_version="13.0",
-        package_versions={
-            "triton": "3.4.9",
-            "fastvideo_kernel": "0.3.7",
-            "flash_attn_4": "4.0.0.dev1",
-            "flashinfer": "0.2.12",
-            "nvidia_cutlass_dsl": "4.5.1",
-        },
-    )
-
-    assert profile_a == {
-        "python": "3.12",
-        "pytorch": "2.12",
-        "cuda": "13.0",
-        "packages": {
-            "fastvideo_kernel": "0.3.2",
-            "flash_attn_4": "4.0.0.dev0",
-            "flashinfer": "0.2.11",
-            "nvidia_cutlass_dsl": "4.5.0",
-            "triton": "3.4.1",
-        },
-    }
-    assert profile_b["python"] == "3.12"
-    assert profile_b["pytorch"] == "2.12"
-    assert profile_b["cuda"] == "13.0"
-    assert software_profile_id(profile_a) != software_profile_id(profile_b)
-
-
-def test_installed_package_versions_checks_attention_kernel_distributions(monkeypatch):
-    versions = {
-        "fastvideo-kernel": "0.3.2",
-        "flash-attn": "2.8.1",
-        "flash-attn-4": "4.0.0.dev0",
-        "flash-attention-fp4": "0.1.0",
-        "flashinfer-python": "0.2.11",
-        "nvidia-cutlass-dsl": "4.5.0",
-    }
-
-    def fake_version(name):
-        if name not in versions:
-            raise identity_module.importlib.metadata.PackageNotFoundError(name)
-        return versions[name]
-
-    monkeypatch.setattr(identity_module.importlib.metadata, "version", fake_version)
-
-    installed = identity_module._installed_package_versions()
-
-    assert installed["fastvideo_kernel"] == "0.3.2"
-    assert installed["flash_attn"] == "2.8.1"
-    assert installed["flash_attn_4"] == "4.0.0.dev0"
-    assert installed["flash_attention_fp4"] == "0.1.0"
-    assert installed["flashinfer"] == "0.2.11"
-    assert installed["nvidia_cutlass_dsl"] == "4.5.0"
-
-
-def test_hardware_profile_id_uses_normalized_gpu_cohort():
-    profile = hardware_profile(
-        num_gpus=2,
-        gpu_devices=[
-            {
-                "name": "NVIDIA L40S",
-                "memory_bytes": 48 * 1024**3,
-                "compute_capability": "8.9",
-            },
-            {
-                "name": "NVIDIA L40S",
-                "memory_gb": 48,
-                "compute_capability": "8.9",
-            },
-        ],
-        interconnect="none_or_partial",
-    )
-
-    assert profile == {
-        "device_type": "cuda",
-        "gpu_count": 2,
-        "gpus": [
-            {
-                "name": "NVIDIA L40S",
-                "memory_gb": 48,
-                "compute_capability": "8.9",
-            },
-            {
-                "name": "NVIDIA L40S",
-                "memory_gb": 48,
-                "compute_capability": "8.9",
-            },
-        ],
-        "interconnect": "none_or_partial",
-    }
-    assert hardware_profile_id(profile).startswith("hw-")
-
-
-def test_hardware_profile_pads_missing_requested_cuda_devices(monkeypatch):
-    class Props:
-        name = "NVIDIA L40S"
-        total_memory = 48 * 1024**3
-
-    monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
-    monkeypatch.setattr(torch.cuda, "device_count", lambda: 1)
-    monkeypatch.setattr(torch.cuda, "get_device_properties", lambda device_id: Props())
-    monkeypatch.setattr(torch.cuda, "get_device_capability", lambda device_id: (8, 9))
-
-    profile = hardware_profile(num_gpus=2, interconnect="unknown")
-
-    assert profile["gpu_count"] == 2
-    assert profile["gpus"] == [
-        {
-            "name": "NVIDIA L40S",
-            "memory_gb": 48,
-            "compute_capability": "8.9",
-        },
-        {
-            "name": "unknown",
-            "memory_gb": None,
-            "compute_capability": None,
-        },
-    ]
-
-
-def test_local_path_detection_checks_both_path_separators(monkeypatch):
-    monkeypatch.setattr(identity_module.os, "sep", "\\")
-
-    assert identity_module._looks_like_local_path("models/local/checkpoint") is True
-    assert identity_module._looks_like_local_path(r"C:\models\checkpoint") is True
-    assert identity_module._looks_like_local_path("Wan-AI/Wan2.1-T2V-1.3B-Diffusers") is False
-
-
-def test_environment_metadata_is_separate_audit_fingerprint():
-    cfg = _benchmark_config()
-    recipe_hash = _fingerprint(cfg)
-    audit_a = environment_metadata(
-        env={"IMAGE_VERSION": "py3.12-cuda13.0.0"},
-        package_versions={"triton": "3.4.1"},
-        hardware={"device_type": "cuda", "gpu_count": 2},
-        software={"python": "3.12", "pytorch": "2.12", "cuda": "13.0"},
-    )
-    audit_b = environment_metadata(
-        env={"IMAGE_VERSION": "py3.12-cuda13.0.1"},
-        package_versions={"triton": "3.4.1"},
-        hardware={"device_type": "cuda", "gpu_count": 2},
-        software={"python": "3.12", "pytorch": "2.12", "cuda": "13.0"},
-    )
-
-    assert recipe_hash == _fingerprint(cfg)
-    assert environment_fingerprint(audit_a) != environment_fingerprint(audit_b)
-
-
-def test_resolved_revision_from_hf_snapshot_path():
-    revision = "a" * 40
-
-    assert resolved_revision_from_model_path(
-        f"/root/.cache/huggingface/hub/models--org--repo/snapshots/{revision}") == revision
-    assert resolved_revision_from_model_path("/models/local-repo") is None
-
-
-def test_runtime_identity_from_generator_summarizes_worker_records():
-    from fastvideo.tests.performance.test_inference_performance import _runtime_identity_from_generator
-
-    revision = "b" * 40
-
-    class FakeExecutor:
-        def collective_rpc(self, _method):
-            return [
-                {
-                    "resolved_attention_backends": ["FLASH_ATTN"],
-                    "resolved_model_path": f"/root/.cache/huggingface/hub/models--org--repo/snapshots/{revision}",
-                },
-                {
-                    "resolved_attention_backends": ["FLASH_ATTN"],
-                    "resolved_model_path": f"/root/.cache/huggingface/hub/models--org--repo/snapshots/{revision}",
-                },
-            ]
-
-    class FakeGenerator:
-        executor = FakeExecutor()
-
-    assert _runtime_identity_from_generator(FakeGenerator()) == {
-        "resolved_attention_backend": "FLASH_ATTN",
-        "resolved_model_revision": revision,
-    }
-
-
-def test_collect_worker_identity_handles_torch_module_pipeline():
-    from fastvideo.tests.performance.test_inference_performance import _collect_worker_identity
-
-    class Backend:
-        name = "FLASH_ATTN"
-
-    class Leaf(torch.nn.Module):
-        backend = Backend()
-
-    class Pipeline(torch.nn.Module):
-        model_path = "/models/local"
-
-        def __init__(self):
-            super().__init__()
-            self.leaf = Leaf()
-
-    class Worker:
-        pipeline = Pipeline()
-
-    assert _collect_worker_identity(Worker()) == {
-        "resolved_attention_backends": ["FLASH_ATTN"],
-        "resolved_model_path": "/models/local",
-    }

fastvideo/tests/performance/test_inference_performance.py

@@ -12,25 +12,12 @@ import os
 import time
 from collections.abc import Mapping
 from datetime import datetime, timezone
-from typing import Any
 
 import torch
 import pytest
 
 from fastvideo import VideoGenerator
 from fastvideo.logger import init_logger
-from fastvideo.tests.performance.identity import (
-    benchmark_identity_from_config,
-    build_recipe_from_benchmark_config,
-    environment_fingerprint,
-    environment_metadata,
-    hardware_profile,
-    hardware_profile_id,
-    recipe_fingerprint,
-    resolved_revision_from_model_path,
-    software_profile,
-    software_profile_id,
-)
 from fastvideo.worker.multiproc_executor import MultiprocExecutor
 
 logger = init_logger(__name__)
@@ -47,11 +34,8 @@ V2_REQUIRED_IDENTITY_FIELDS = (
     "variant_id",
     "benchmark_version",
 )
-# "recipe" is no longer config-declarable: the fingerprint follow-up landed,
-# and the generated recipe document owns that key in emitted records. A config
-# declaring it now fails validation loudly instead of being silently
-# overwritten by the generated one.
 V2_OPTIONAL_METADATA_FIELDS = (
+    "recipe",
     "metric_threshold_policy",
     "quality_metadata",
 )
@@ -195,11 +179,7 @@ def _extract_component_times(result: dict) -> dict[str, float | None]:
             logger.debug("Skipping malformed stage '%s' data: %r", stage_name, stage_data)
             continue
         stage_class = stage_data.get("stage_class", stage_name)
-        component_metric = stage_data.get("component_metric")
-        if isinstance(component_metric, str) and component_metric in component_times:
-            metric_key = component_metric
-        else:
-            metric_key = STAGE_METRIC_MAP.get(stage_class)
+        metric_key = STAGE_METRIC_MAP.get(stage_class)
         if metric_key is None:
             logger.debug("Unmapped stage '%s' class '%s' (%.3fs)",
                          stage_name,
@@ -247,110 +227,6 @@ def _write_results(results):
     logger.info("Performance results written to %s", filepath)
 
 
-def _backend_name(value) -> str:
-    if hasattr(value, "name"):
-        return str(value.name)
-    return str(value)
-
-
-def _collect_worker_identity(worker) -> dict[str, Any]:
-    pipeline = getattr(worker, "pipeline", None)
-    model_path = getattr(pipeline, "model_path", None)
-    backends: set[str] = set()
-
-    modules = getattr(pipeline, "modules", None)
-    if isinstance(modules, Mapping):
-        module_iter = modules.values()
-    elif isinstance(pipeline, torch.nn.Module):
-        module_iter = (pipeline,)
-    else:
-        module_iter = ()
-
-    for module in module_iter:
-        if isinstance(module, torch.nn.Module):
-            for submodule in module.modules():
-                backend = getattr(submodule, "backend", None)
-                if backend is not None:
-                    backends.add(_backend_name(backend))
-        else:
-            backend = getattr(module, "backend", None)
-            if backend is not None:
-                backends.add(_backend_name(backend))
-
-    return {
-        "resolved_attention_backends": sorted(backends),
-        "resolved_model_path": model_path,
-    }
-
-
-def _single_or_list(values: set[str]) -> str | list[str] | None:
-    ordered = sorted(values)
-    if not ordered:
-        return None
-    if len(ordered) == 1:
-        return ordered[0]
-    return ordered
-
-
-def _runtime_identity_from_generator(generator) -> dict[str, Any]:
-    worker_records = generator.executor.collective_rpc(_collect_worker_identity)
-    resolved_backends: set[str] = set()
-    resolved_revisions: set[str] = set()
-
-    for record in worker_records:
-        resolved_backends.update(record.get("resolved_attention_backends") or [])
-        revision = resolved_revision_from_model_path(record.get("resolved_model_path"))
-        if revision is not None:
-            resolved_revisions.add(revision)
-
-    return {
-        "resolved_attention_backend": _single_or_list(resolved_backends),
-        "resolved_model_revision": _single_or_list(resolved_revisions),
-    }
-
-
-def _benchmark_identity_fields(cfg):
-    identity = benchmark_identity_from_config(cfg)
-    return {
-        key: identity[key]
-        for key in ("workload_id", "variant_id", "benchmark_version")
-    }
-
-
-def _build_identity_fields(cfg, init_kwargs, prompt, runtime_identity):
-    # Legacy v1 configs (no config_schema_version) stay on the legacy
-    # (model_id, gpu_type) cohort: they lack the required identity fields,
-    # and building a recipe for them would raise AFTER the GPU measurement
-    # completed. The validator already enforces that v2 configs carry the
-    # identity fields, so v2 records always get the full identity block.
-    if cfg.get("config_schema_version") is None:
-        return {}
-    recipe = build_recipe_from_benchmark_config(
-        cfg,
-        resolved_attention_backend=runtime_identity.get("resolved_attention_backend"),
-        resolved_model_revision=runtime_identity.get("resolved_model_revision"),
-        measured_prompts=[prompt],
-    )
-    num_gpus = init_kwargs.get("num_gpus", cfg.get("run_config", {}).get("required_gpus", 1))
-    hw_profile = hardware_profile(num_gpus=num_gpus)
-    sw_profile = software_profile()
-    env_metadata = environment_metadata(
-        hardware=hw_profile,
-        software=sw_profile,
-    )
-    return {
-        "recipe": recipe,
-        "recipe_fingerprint": recipe_fingerprint(recipe),
-        "hardware_profile": hw_profile,
-        "hardware_profile_id": hardware_profile_id(hw_profile),
-        "software_profile": sw_profile,
-        "software_profile_id": software_profile_id(sw_profile),
-        "environment_metadata": env_metadata,
-        "environment_fingerprint": environment_fingerprint(env_metadata),
-        **_benchmark_identity_fields(cfg),
-    }
-
-
 # -- Test -------------------------------------------------------------------
 
 def _run_benchmark(cfg):
@@ -388,7 +264,6 @@ def _run_benchmark(cfg):
             model_path=model_info["model_path"],
             **init_kwargs,
         )
-        runtime_identity = _runtime_identity_from_generator(generator)
 
         for i in range(num_warmup):
             logger.info("Warmup run %d/%d", i + 1, num_warmup)
@@ -443,7 +318,6 @@ def _run_benchmark(cfg):
         "dit_time_s": _avg_component(all_component_times, "dit_time_s"),
         "vae_decode_time_s": _avg_component(all_component_times,
                                             "vae_decode_time_s"),
-        **_build_identity_fields(cfg, init_kwargs, prompt, runtime_identity),
     }
 
     logger.info(

fastvideo/tests/performance/test_inference_performance_component_times.py

@@ -1,20 +1,13 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from fastvideo.pipelines.pipeline_batch_info import PipelineLoggingInfo
-from fastvideo.pipelines.stages.denoising import Cosmos25AutoDenoisingStage, DenoisingStage
-from fastvideo.pipelines.stages.text_encoding import Cosmos25TextEncodingStage
 from fastvideo.tests.performance.test_inference_performance import _extract_component_times
 
 
-class SubclassStyleDenoisingStage(DenoisingStage):
-    pass
-
-
 def test_extract_component_times_handles_pipeline_logging_info_object():
     logging_info = PipelineLoggingInfo()
     logging_info.add_stage_execution_time("prompt_encoding_stage", 1.25)
     logging_info.add_stage_metric("prompt_encoding_stage", "stage_class", "TextEncodingStage")
-    logging_info.add_stage_metric("prompt_encoding_stage", "component_metric", "text_encoder_time_s")
 
     assert _extract_component_times({"logging_info": logging_info}) == {
         "text_encoder_time_s": 1.25,
@@ -52,35 +45,6 @@ def test_extract_component_times_uses_stage_class_for_pipeline_stage_keys():
     }
 
 
-def test_denoising_stage_subclasses_inherit_component_metric():
-    assert SubclassStyleDenoisingStage.performance_component_metric == "dit_time_s"
-
-
-def test_cosmos25_direct_pipeline_stages_define_component_metrics():
-    assert Cosmos25TextEncodingStage.performance_component_metric == "text_encoder_time_s"
-    assert Cosmos25AutoDenoisingStage.performance_component_metric == "dit_time_s"
-
-
-def test_extract_component_times_uses_component_metric_for_stage_subclasses():
-    result = {
-        "logging_info": {
-            "stages": {
-                "denoising_stage": {
-                    "execution_time": 4.2,
-                    "stage_class": "CosmosDenoisingStage",
-                    "component_metric": "dit_time_s",
-                },
-            },
-        },
-    }
-
-    assert _extract_component_times(result) == {
-        "text_encoder_time_s": None,
-        "dit_time_s": 4.2,
-        "vae_decode_time_s": None,
-    }
-
-
 def test_extract_component_times_keeps_legacy_class_name_keys():
     # Backward-compatibility check for logs produced before pipeline-unique
     # stage keys carried a separate stage_class field.

fastvideo/tests/ssim/inference_similarity_utils.py

@@ -3,9 +3,9 @@
 from __future__ import annotations
 
 import os
-from collections.abc import Iterator
 from contextlib import contextmanager
 from logging import Logger
+from typing import Iterator
 
 from fastvideo import VideoGenerator
 from fastvideo.tests.ssim.reference_utils import (
@@ -67,12 +67,6 @@ def _find_reference_video(reference_folder: str, prompt: str) -> str:
     raise FileNotFoundError("Reference video missing")
 
 
-def _remove_stale_generated_video(output_dir: str, output_video_name: str) -> None:
-    stale_path = os.path.join(output_dir, output_video_name)
-    if os.path.exists(stale_path):
-        os.remove(stale_path)
-
-
 def _assert_similarity(
     *,
     logger: Logger,
@@ -220,7 +214,6 @@ def run_text_to_video_similarity_test(
         )
         output_video_name = f"{prompt[:100].strip()}.mp4"
         os.makedirs(output_dir, exist_ok=True)
-        _remove_stale_generated_video(output_dir, output_video_name)
 
         params_map = select_ssim_params(
             default_params_map,
@@ -296,7 +289,6 @@ def run_image_to_video_similarity_test(
         )
         output_video_name = f"{prompt[:100].strip()}.mp4"
         os.makedirs(output_dir, exist_ok=True)
-        _remove_stale_generated_video(output_dir, output_video_name)
 
         params_map = select_ssim_params(
             default_params_map,

fastvideo/tests/ssim/test_dreamx_world_similarity.py

@@ -1,191 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-from __future__ import annotations
-
-import os
-from pathlib import Path
-
-import pytest
-from PIL import Image, ImageDraw
-
-from fastvideo.logger import init_logger
-from fastvideo.tests.ssim.inference_similarity_utils import (
-    resolve_inference_device_reference_folder,
-    run_image_to_video_similarity_test,
-    run_text_to_video_similarity_test,
-)
-
-logger = init_logger(__name__)
-
-REQUIRED_GPUS = 1
-
-device_reference_folder = resolve_inference_device_reference_folder(logger)
-
-_MODEL_PATH = os.getenv(
-    "DREAMX_WORLD_SSIM_MODEL_PATH",
-    "FastVideo/DreamX-World-5B-Cam-Diffusers",
-)
-_AR_MODEL_PATH = os.getenv(
-    "DREAMX_WORLD_AR_SSIM_MODEL_PATH",
-    "FastVideo/DreamX-World-5B-Diffusers",
-)
-
-DREAMX_WORLD_PARAMS = {
-    "num_gpus": 1,
-    "model_path": _MODEL_PATH,
-    "height": 64,
-    "width": 64,
-    "num_frames": 9,
-    "num_inference_steps": 1,
-    "guidance_scale": 1.0,
-    "seed": 1024,
-    "fps": 16,
-}
-
-DREAMX_WORLD_FULL_QUALITY_PARAMS = {
-    **DREAMX_WORLD_PARAMS,
-    "height": 480,
-    "width": 832,
-    "num_frames": 161,
-    "num_inference_steps": 30,
-    "guidance_scale": 5.0,
-}
-
-
-DREAMX_WORLD_AR_PARAMS = {
-    "num_gpus": 1,
-    "model_path": _AR_MODEL_PATH,
-    "height": 192,
-    "width": 192,
-    "num_frames": 81,
-    "num_inference_steps": 4,
-    "guidance_scale": 1.0,
-    "seed": 2048,
-    "fps": 16,
-}
-
-DREAMX_WORLD_AR_FULL_QUALITY_PARAMS = {
-    **DREAMX_WORLD_AR_PARAMS,
-    "height": 704,
-    "width": 1280,
-    "num_frames": 1005,
-}
-
-DREAMX_WORLD_MODEL_TO_PARAMS = {
-    "DreamX-World-5B-Cam": DREAMX_WORLD_PARAMS,
-}
-
-DREAMX_WORLD_AR_MODEL_TO_PARAMS = {
-    "DreamX-World-5B": DREAMX_WORLD_AR_PARAMS,
-}
-
-FULL_QUALITY_DREAMX_WORLD_MODEL_TO_PARAMS = {
-    "DreamX-World-5B-Cam": DREAMX_WORLD_FULL_QUALITY_PARAMS,
-}
-
-FULL_QUALITY_DREAMX_WORLD_AR_MODEL_TO_PARAMS = {
-    "DreamX-World-5B": DREAMX_WORLD_AR_FULL_QUALITY_PARAMS,
-}
-
-DREAMX_WORLD_TEST_CASES = [
-    (
-        "A cinematic first-person drive through a futuristic coastal city at sunrise, "
-        "reflective glass towers, clean streets, soft volumetric light.",
-        ("w", "d", "w"),
-        (4.0, 2.0, 4.0),
-    ),
-]
-
-
-DREAMX_WORLD_AR_TEST_CASES = [
-    (
-        "A long autonomous drive through a futuristic coastal city at sunrise, "
-        "smooth forward camera motion, reflective glass towers, clean streets.",
-        ("w", "d", "w", "a"),
-        (2.0, 1.0, 2.0, 1.0),
-    ),
-]
-
-
-def _write_deterministic_reference_image(path: Path) -> None:
-    path.parent.mkdir(parents=True, exist_ok=True)
-    image = Image.new("RGB", (96, 96), color=(42, 76, 112))
-    draw = ImageDraw.Draw(image)
-    draw.rectangle((0, 56, 96, 96), fill=(36, 44, 52))
-    draw.polygon([(0, 56), (48, 30), (96, 56)], fill=(120, 142, 158))
-    draw.rectangle((34, 42, 62, 72), fill=(178, 198, 212))
-    draw.line((0, 80, 96, 66), fill=(238, 209, 124), width=3)
-    image.save(path)
-
-
-@pytest.mark.parametrize(("prompt", "action_list", "action_speed_list"), DREAMX_WORLD_TEST_CASES)
-@pytest.mark.parametrize("attention_backend_name", ["TORCH_SDPA"])
-@pytest.mark.parametrize("model_id", list(DREAMX_WORLD_MODEL_TO_PARAMS.keys()))
-def test_dreamx_world_inference_similarity(
-    prompt: str,
-    action_list: tuple[str, ...],
-    action_speed_list: tuple[float, ...],
-    attention_backend_name: str,
-    model_id: str,
-    tmp_path: Path,
-) -> None:
-    image_path = tmp_path / "dreamx_world_ssim_input.png"
-    _write_deterministic_reference_image(image_path)
-
-    run_image_to_video_similarity_test(
-        logger=logger,
-        script_dir=os.path.dirname(os.path.abspath(__file__)),
-        device_reference_folder=device_reference_folder,
-        prompt=prompt,
-        image_path=str(image_path),
-        attention_backend_name=attention_backend_name,
-        model_id=model_id,
-        default_params_map=DREAMX_WORLD_MODEL_TO_PARAMS,
-        full_quality_params_map=FULL_QUALITY_DREAMX_WORLD_MODEL_TO_PARAMS,
-        min_acceptable_ssim=0.98,
-        init_kwargs_override={
-            "use_fsdp_inference": False,
-            "dit_cpu_offload": False,
-            "vae_cpu_offload": True,
-            "text_encoder_cpu_offload": True,
-            "pin_cpu_memory": False,
-            "override_pipeline_cls_name": "DreamXWorldPipeline",
-        },
-        generation_kwargs_override={
-            "action_list": list(action_list),
-            "action_speed_list": list(action_speed_list),
-        },
-    )
-
-@pytest.mark.parametrize(("prompt", "action_list", "action_speed_list"), DREAMX_WORLD_AR_TEST_CASES)
-@pytest.mark.parametrize("attention_backend_name", ["TORCH_SDPA"])
-@pytest.mark.parametrize("model_id", list(DREAMX_WORLD_AR_MODEL_TO_PARAMS.keys()))
-def test_dreamx_world_ar_inference_similarity(
-    prompt: str,
-    action_list: tuple[str, ...],
-    action_speed_list: tuple[float, ...],
-    attention_backend_name: str,
-    model_id: str,
-) -> None:
-    run_text_to_video_similarity_test(
-        logger=logger,
-        script_dir=os.path.dirname(os.path.abspath(__file__)),
-        device_reference_folder=device_reference_folder,
-        prompt=prompt,
-        attention_backend_name=attention_backend_name,
-        model_id=model_id,
-        default_params_map=DREAMX_WORLD_AR_MODEL_TO_PARAMS,
-        full_quality_params_map=FULL_QUALITY_DREAMX_WORLD_AR_MODEL_TO_PARAMS,
-        min_acceptable_ssim=0.98,
-        init_kwargs_override={
-            "use_fsdp_inference": False,
-            "dit_cpu_offload": False,
-            "vae_cpu_offload": True,
-            "text_encoder_cpu_offload": True,
-            "pin_cpu_memory": False,
-            "override_pipeline_cls_name": "DreamXWorldARPipeline",
-        },
-        generation_kwargs_override={
-            "action_list": list(action_list),
-            "action_speed_list": list(action_speed_list),
-        },
-    )

fastvideo/tests/stages/test_input_validation_batching.py

@@ -0,0 +1,40 @@
+# SPDX-License-Identifier: Apache-2.0
+from types import SimpleNamespace
+
+from fastvideo.pipelines import ForwardBatch
+from fastvideo.pipelines.stages.input_validation import InputValidationStage
+
+
+def test_input_validation_preserves_explicit_dynamic_batch_seeds() -> None:
+    batch = ForwardBatch(
+        data_type="video",
+        prompt=["one", "two"],
+        seed=100,
+        seeds=[17, 23],
+        height=8,
+        width=8,
+        num_frames=1,
+        num_inference_steps=1,
+    )
+
+    InputValidationStage()._generate_seeds(batch, SimpleNamespace())
+
+    assert batch.seeds == [17, 23]
+    assert [generator.initial_seed() for generator in batch.generator] == [17, 23]
+
+
+def test_input_validation_generates_one_seed_per_prompt() -> None:
+    batch = ForwardBatch(
+        data_type="video",
+        prompt=["one", "two"],
+        seed=100,
+        height=8,
+        width=8,
+        num_frames=1,
+        num_inference_steps=1,
+    )
+
+    InputValidationStage()._generate_seeds(batch, SimpleNamespace())
+
+    assert batch.seeds == [100, 101]
+    assert [generator.initial_seed() for generator in batch.generator] == [100, 101]

fastvideo/tests/stages/test_text_encoding.py

@@ -13,7 +13,13 @@ class TensorDict(dict):
         return TensorDict({k: v.to(device) for k, v in self.items()})
 
 class FakeTokenizer:
+    def __init__(self):
+        self.calls = []
+        self.texts = []
+
     def __call__(self, texts, **kwargs):
+        self.calls.append(kwargs)
+        self.texts.append(list(texts))
         B = len(texts)
         seq_len = int(kwargs.get("max_length", 4))
         return TensorDict({
@@ -159,6 +165,66 @@ def test_forward_integration_cfg_off_and_on():
     assert len(out2.prompt_attention_mask) == 2
     assert len(out2.negative_attention_mask) == 2
 
+def test_encode_text_adds_padding_for_prompt_lists():
+    fastvideo_args, hidden = make_args(num_encoders=1, text_len=4, hidden_size=8)
+    stage = make_stage(num_encoders=1, hidden_size=hidden)
+
+    stage.encode_text(["short", "a longer prompt"], fastvideo_args, encoder_index=[0])
+
+    assert stage.tokenizers[0].calls[-1]["padding"] is True
+
+
+def test_forward_prompt_list_preserves_single_prompt_text_encoding_path():
+    fastvideo_args, hidden = make_args(num_encoders=1, text_len=4, hidden_size=8)
+    stage = make_stage(num_encoders=1, hidden_size=hidden)
+    batch = ForwardBatch(
+        data_type="video",
+        prompt=["short", "a longer prompt"],
+        negative_prompt="",
+        do_classifier_free_guidance=False,
+        prompt_embeds=[],
+        negative_prompt_embeds=None,
+        prompt_attention_mask=[],
+        negative_attention_mask=None,
+    )
+
+    out = stage.forward(batch, fastvideo_args)
+
+    assert stage.tokenizers[0].texts == [["short"], ["a longer prompt"]]
+    assert out.prompt_embeds[0].shape == (2, hidden)
+    assert out.prompt_attention_mask[0].shape == (2, 4)
+
+
+def test_encode_prompt_list_individually_pads_variable_length_embeds_and_audio():
+    fastvideo_args, _hidden = make_args(num_encoders=1, text_len=4, hidden_size=8)
+    stage = TextEncodingStage(text_encoders=[], tokenizers=[])
+    lengths = {"short": 2, "a longer prompt": 4}
+
+    def fake_encode_text(text, *_args, **_kwargs):
+        length = lengths[text]
+        embeds = [torch.full((1, length, 3), fill_value=float(length))]
+        masks = [torch.ones((1, length), dtype=torch.long)]
+        stage._last_audio_embeds = [torch.full((1, length, 5), fill_value=float(length))]
+        return embeds, masks
+
+    stage.encode_text = fake_encode_text
+
+    embeds, masks = stage._encode_prompt_list_individually(
+        ["short", "a longer prompt"],
+        fastvideo_args,
+        encoder_index=[0],
+        return_attention_mask=True,
+    )
+
+    assert embeds[0].shape == (2, 4, 3)
+    assert masks[0].shape == (2, 4)
+    assert stage._last_audio_embeds is not None
+    assert stage._last_audio_embeds[0].shape == (2, 4, 5)
+    assert torch.equal(embeds[0][0, :2], torch.full((2, 3), 2.0))
+    assert torch.equal(embeds[0][0, 2:], torch.zeros((2, 3)))
+    assert torch.equal(stage._last_audio_embeds[0][0, :2], torch.full((2, 5), 2.0))
+    assert torch.equal(stage._last_audio_embeds[0][0, 2:], torch.zeros((2, 5)))
+
 
 def test_encode_text_hidden_state_flag_follows_encoder_config():
     fastvideo_args, hidden = make_args(num_encoders=1, text_len=4, hidden_size=8)

fastvideo/tests/train/fixtures/wan_causal_t2v_causal_cd_min.yaml

@@ -1,56 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# Minimum config to run a single training step of
-# CausalConsistencyDistillationMethod on WanCausalModel for the
-# per-method smoke test. Uses the real Wan 2.1 1.3B checkpoint for both
-# the trainable student and the frozen teacher (AR Euler-step target),
-# with tiny synthetic latents.
-
-models:
-  student:
-    _target_: fastvideo.train.models.wan.WanCausalModel
-    init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers
-    trainable: true
-  teacher:
-    _target_: fastvideo.train.models.wan.WanCausalModel
-    init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers
-    trainable: false
-  ema:
-    _target_: fastvideo.train.models.wan.WanCausalModel
-    init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers
-    trainable: false
-
-method:
-  _target_: fastvideo.train.methods.consistency_model.causal_cd.CausalConsistencyDistillationMethod
-  discrete_cd_N: 12
-  guidance_scale: 3.0
-  ema_decay: 0.95
-
-training:
-  dit_precision: bf16
-
-  distributed:
-    num_gpus: 1
-    sp_size: 1
-    tp_size: 1
-
-  data:
-    seed: 42
-    train_batch_size: 1
-    training_cfg_rate: 0.0
-    num_latent_t: 6
-    num_height: 64
-    num_width: 64
-    num_frames: 21
-
-  optimizer:
-    learning_rate: 2.0e-6
-    betas: [0.9, 0.999]
-    weight_decay: 0.01
-    lr_scheduler: constant
-    lr_warmup_steps: 0
-
-  loop:
-    max_train_steps: 1
-    gradient_accumulation_steps: 1
-
-pipeline: {}

fastvideo/tests/train/fixtures/wan_causal_t2v_dfsft_framewise_min.yaml

@@ -1,46 +0,0 @@
-# Minimum config to run a single training step of
-# DiffusionForcingSFTMethod on a frame-wise WanCausalModel
-# (num_frames_per_block=1, chunk_size=1) for the per-method smoke
-# test.  Uses the real Wan 2.1 1.3B checkpoint with tiny synthetic
-# latents.
-
-models:
-  student:
-    _target_: fastvideo.train.models.wan.WanCausalModel
-    init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers
-    trainable: true
-    num_frames_per_block: 1
-
-method:
-  _target_: fastvideo.train.methods.fine_tuning.dfsft.DiffusionForcingSFTMethod
-  chunk_size: 1
-
-training:
-  dit_precision: bf16
-
-  distributed:
-    num_gpus: 1
-    sp_size: 1
-    tp_size: 1
-
-  data:
-    seed: 42
-    train_batch_size: 1
-    training_cfg_rate: 0.0
-    num_latent_t: 6
-    num_height: 64
-    num_width: 64
-    num_frames: 21
-
-  optimizer:
-    learning_rate: 2.0e-6
-    betas: [0.9, 0.999]
-    weight_decay: 0.01
-    lr_scheduler: constant
-    lr_warmup_steps: 0
-
-  loop:
-    max_train_steps: 1
-    gradient_accumulation_steps: 1
-
-pipeline: {}

fastvideo/tests/train/fixtures/wan_causal_t2v_tfsft_min.yaml

@@ -1,45 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# Minimum config to run a single training step of
-# TeacherForcingSFTMethod on WanCausalModel for the per-method smoke
-# test. Identical to wan_causal_t2v_dfsft_min.yaml except the method,
-# which feeds clean history to the causal transformer (teacher forcing).
-
-models:
-  student:
-    _target_: fastvideo.train.models.wan.WanCausalModel
-    init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers
-    trainable: true
-
-method:
-  _target_: fastvideo.train.methods.fine_tuning.tfsft.TeacherForcingSFTMethod
-  chunk_size: 3
-
-training:
-  dit_precision: bf16
-
-  distributed:
-    num_gpus: 1
-    sp_size: 1
-    tp_size: 1
-
-  data:
-    seed: 42
-    train_batch_size: 1
-    training_cfg_rate: 0.0
-    num_latent_t: 6
-    num_height: 64
-    num_width: 64
-    num_frames: 21
-
-  optimizer:
-    learning_rate: 2.0e-6
-    betas: [0.9, 0.999]
-    weight_decay: 0.01
-    lr_scheduler: constant
-    lr_warmup_steps: 0
-
-  loop:
-    max_train_steps: 1
-    gradient_accumulation_steps: 1
-
-pipeline: {}

fastvideo/tests/train/methods/test_wan_causal_cd.py

@@ -1,147 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Per-method GPU smoke test: ``WanCausalModel`` + ``CausalConsistencyDistillationMethod``.
-
-Mirrors ``test_wan_causal_dfsft.py``. Causal consistency distillation
-bootstraps a consistency MSE between the student's ``x0`` at ``t`` and an EMA
-copy of the student at ``t_next``, where ``t_next`` is produced online by a
-single CFG Euler step of a frozen teacher (all under clean-history teacher
-forcing). This test exercises the full step: finite loss, nonzero student
-gradients, frozen teacher, and a post-step EMA update.
-"""
-
-from __future__ import annotations
-
-import os
-
-os.environ.setdefault("MASTER_ADDR", "localhost")
-os.environ.setdefault("MASTER_PORT", "29519")
-
-from pathlib import Path
-
-import pytest
-import torch
-
-from fastvideo.train.methods.consistency_model.causal_cd import (
-    CausalConsistencyDistillationMethod, )
-from fastvideo.train.models.wan import WanCausalModel
-from fastvideo.train.utils.config import load_run_config
-
-
-_FIXTURE = str(
-    Path(__file__).resolve().parent.parent / "fixtures"
-    / "wan_causal_t2v_causal_cd_min.yaml")
-
-
-def _build_synthetic_batch(
-    device: torch.device,
-    dtype: torch.dtype,
-) -> dict[str, torch.Tensor]:
-    batch_size = 1
-    return {
-        "text_embedding":
-        torch.randn(batch_size, 16, 4096, device=device, dtype=dtype),
-        "text_attention_mask":
-        torch.ones(batch_size, 16, device=device),
-        "vae_latent":
-        torch.randn(batch_size, 16, 6, 8, 8, device=device, dtype=dtype),
-    }
-
-
-@pytest.mark.usefixtures("distributed_setup")
-def test_wan_causal_cd_single_train_step(
-        monkeypatch: pytest.MonkeyPatch) -> None:
-    if not torch.cuda.is_available():
-        pytest.skip("requires CUDA")
-
-    cfg = load_run_config(_FIXTURE)
-
-    device = torch.device("cuda:0")
-    dtype = torch.bfloat16
-
-    monkeypatch.setattr(
-        "fastvideo.train.utils.dataloader."
-        "build_parquet_t2v_train_dataloader",
-        lambda *args, **kwargs: None,
-    )
-
-    student = WanCausalModel(
-        init_from=cfg.models["student"]["init_from"],
-        training_config=cfg.training,
-        trainable=True,
-    )
-    student.transformer = student.transformer.to(device=device, dtype=dtype)
-
-    teacher = WanCausalModel(
-        init_from=cfg.models["teacher"]["init_from"],
-        training_config=cfg.training,
-        trainable=False,
-    )
-    teacher.transformer = teacher.transformer.to(device=device, dtype=dtype)
-
-    ema = WanCausalModel(
-        init_from=cfg.models["ema"]["init_from"],
-        training_config=cfg.training,
-        trainable=False,
-    )
-    ema.transformer = ema.transformer.to(device=device, dtype=dtype)
-
-    method = CausalConsistencyDistillationMethod(
-        cfg=cfg,
-        role_models={"student": student, "teacher": teacher, "ema": ema},
-    )
-    method.on_train_start()
-
-    batch = _build_synthetic_batch(device, dtype)
-    loss_map, outputs, _metrics = method.single_train_step(batch, iteration=0)
-
-    loss = loss_map["total_loss"]
-    assert torch.is_tensor(loss), "total_loss must be a torch.Tensor"
-    assert torch.isfinite(loss).item(), (
-        f"total_loss is not finite: {loss.item()}")
-
-    method.backward(loss_map, outputs, grad_accum_rounds=1)
-
-    blocks = student.transformer.blocks
-    assert blocks is not None and len(blocks) > 0
-    layer0 = blocks[0]
-    trainable = [p for p in layer0.parameters() if p.requires_grad]
-    assert len(trainable) > 0, "student layer 0 has no trainable parameters"
-    for i, p in enumerate(trainable):
-        assert p.grad is not None, f"student layer 0 param[{i}] has None grad"
-        assert torch.isfinite(p.grad).all().item(), (
-            f"student layer 0 param[{i}] grad contains NaN/Inf")
-    assert any(p.grad.detach().float().norm().item() > 0.0 for p in trainable), (
-        "all student layer-0 grads are exactly zero; consistency loss "
-        "did not reach the first transformer block")
-
-    # Teacher must stay frozen.
-    assert all(not p.requires_grad for p in teacher.transformer.parameters()), (
-        "teacher must be frozen for Causal-CD")
-
-    # The EMA model and student start from the same checkpoint, so the first
-    # parameter must match before any update. FSDP fully_shard params are
-    # DTensors; compare the local shards (torch.equal is unsupported on
-    # DTensor).
-    def _local(p: torch.Tensor) -> torch.Tensor:
-        return p.to_local() if hasattr(p, "to_local") else p
-
-    ema_param = next(ema.transformer.parameters())
-    student_param = next(student.transformer.parameters())
-    assert torch.equal(_local(ema_param), _local(student_param)), (
-        "EMA model should start identical to the student (same checkpoint)")
-
-    # The EMA update must move EMA toward the student. Apply a visibly large
-    # perturbation so the bf16 lerp is well above rounding noise (the real
-    # optimizer step at lr=2e-6 would be sub-ULP in bf16).
-    with torch.no_grad():
-        student_param.add_(1.0)
-    before = _local(ema_param).detach().float().clone()
-    method._update_ema()
-    after = _local(ema_param).detach().float()
-    assert not torch.equal(before, after), (
-        "EMA weights did not move after _update_ema")
-    # EMA = decay*ema + (1-decay)*student moves ~ (1-decay) of the gap.
-    expected = before + (1.0 - method._ema_decay) * (
-        _local(student_param).detach().float() - before)
-    assert torch.allclose(after, expected, atol=1e-2), (
-        "EMA update did not follow the expected lerp")

fastvideo/tests/train/methods/test_wan_causal_dfsft_framewise.py

@@ -1,109 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Per-method GPU smoke test: frame-wise ``WanCausalModel`` + ``DiffusionForcingSFTMethod``.
-
-Mirrors ``test_wan_causal_dfsft.py`` but with a block size of 1 frame
-(``num_frames_per_block=1`` on the model, ``chunk_size=1`` on the method),
-so each frame gets its own independent noise level. The test asserts the
-override took effect and runs one train step: forward, finite loss, and
-nonzero gradients reaching the first transformer block.
-"""
-
-from __future__ import annotations
-
-import os
-
-os.environ.setdefault("MASTER_ADDR", "localhost")
-os.environ.setdefault("MASTER_PORT", "29520")
-
-from pathlib import Path
-
-import pytest
-import torch
-
-from fastvideo.train.methods.fine_tuning.dfsft import (
-    DiffusionForcingSFTMethod, )
-from fastvideo.train.models.wan import WanCausalModel
-from fastvideo.train.utils.config import load_run_config
-
-
-_FIXTURE = str(
-    Path(__file__).resolve().parent.parent / "fixtures"
-    / "wan_causal_t2v_dfsft_framewise_min.yaml")
-
-
-def _build_synthetic_batch(
-    device: torch.device,
-    dtype: torch.dtype,
-) -> dict[str, torch.Tensor]:
-    batch_size = 1
-    return {
-        "text_embedding":
-        torch.randn(batch_size, 16, 4096, device=device, dtype=dtype),
-        "text_attention_mask":
-        torch.ones(batch_size, 16, device=device),
-        "vae_latent":
-        torch.randn(batch_size, 16, 6, 8, 8, device=device, dtype=dtype),
-    }
-
-
-@pytest.mark.usefixtures("distributed_setup")
-def test_wan_causal_dfsft_framewise_single_train_step(
-        monkeypatch: pytest.MonkeyPatch) -> None:
-    if not torch.cuda.is_available():
-        pytest.skip("requires CUDA")
-
-    cfg = load_run_config(_FIXTURE)
-
-    device = torch.device("cuda:0")
-    dtype = torch.bfloat16
-
-    monkeypatch.setattr(
-        "fastvideo.train.utils.dataloader."
-        "build_parquet_t2v_train_dataloader",
-        lambda *args, **kwargs: None,
-    )
-
-    student_cfg = cfg.models["student"]
-    model = WanCausalModel(
-        init_from=student_cfg["init_from"],
-        training_config=cfg.training,
-        trainable=True,
-        num_frames_per_block=student_cfg.get("num_frames_per_block"),
-    )
-    assert model.transformer.num_frame_per_block == 1, (
-        "frame-wise override did not reach the transformer")
-    model.transformer = model.transformer.to(device=device, dtype=dtype)
-
-    method = DiffusionForcingSFTMethod(
-        cfg=cfg,
-        role_models={"student": model},
-    )
-    method.on_train_start()
-
-    batch = _build_synthetic_batch(device, dtype)
-    loss_map, outputs, _metrics = method.single_train_step(batch, iteration=0)
-
-    loss = loss_map["total_loss"]
-    assert torch.is_tensor(loss), "total_loss must be a torch.Tensor"
-    assert torch.isfinite(loss).item(), (
-        f"total_loss is not finite: {loss.item()}")
-
-    method.backward(loss_map, outputs, grad_accum_rounds=1)
-
-    blocks = getattr(model.transformer, "blocks", None)
-    assert blocks is not None and len(blocks) > 0
-    layer0 = blocks[0]
-
-    trainable = [p for p in layer0.parameters() if p.requires_grad]
-    assert len(trainable) > 0, "layer 0 has no trainable parameters"
-
-    for i, p in enumerate(trainable):
-        assert p.grad is not None, f"layer 0 param[{i}] has None grad"
-        assert torch.isfinite(p.grad).all().item(), (
-            f"layer 0 param[{i}] grad contains NaN/Inf")
-
-    any_nonzero = any(
-        p.grad.detach().float().norm().item() > 0.0 for p in trainable)
-    assert any_nonzero, (
-        "all layer-0 grads are exactly zero; backward did not "
-        "reach the first transformer block")

fastvideo/tests/train/methods/test_wan_causal_tfsft.py

@@ -1,113 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Per-method GPU smoke test: ``WanCausalModel`` + ``TeacherForcingSFTMethod``.
-
-Mirrors ``test_wan_causal_dfsft.py``. Teacher forcing concatenates a clean
-context copy of every frame inside the causal transformer (``clean_x``) and
-denoises the current block while attending to *clean* history. This test
-exercises the ``clean_x`` path end-to-end: forward, finite loss, and nonzero
-gradients reaching the first transformer block.
-"""
-
-from __future__ import annotations
-
-import os
-
-os.environ.setdefault("MASTER_ADDR", "localhost")
-os.environ.setdefault("MASTER_PORT", "29518")
-
-from pathlib import Path
-
-import pytest
-import torch
-
-from fastvideo.train.methods.fine_tuning.tfsft import (
-    TeacherForcingSFTMethod, )
-from fastvideo.train.models.wan import WanCausalModel
-from fastvideo.train.utils.config import load_run_config
-
-
-_FIXTURE = str(
-    Path(__file__).resolve().parent.parent / "fixtures"
-    / "wan_causal_t2v_tfsft_min.yaml")
-
-
-def _build_synthetic_batch(
-    device: torch.device,
-    dtype: torch.dtype,
-) -> dict[str, torch.Tensor]:
-    batch_size = 1
-    return {
-        "text_embedding":
-        torch.randn(batch_size, 16, 4096, device=device, dtype=dtype),
-        "text_attention_mask":
-        torch.ones(batch_size, 16, device=device),
-        "vae_latent":
-        torch.randn(batch_size, 16, 6, 8, 8, device=device, dtype=dtype),
-    }
-
-
-@pytest.mark.usefixtures("distributed_setup")
-def test_wan_causal_tfsft_single_train_step(
-        monkeypatch: pytest.MonkeyPatch) -> None:
-    if not torch.cuda.is_available():
-        pytest.skip("requires CUDA")
-
-    cfg = load_run_config(_FIXTURE)
-
-    device = torch.device("cuda:0")
-    dtype = torch.bfloat16
-
-    monkeypatch.setattr(
-        "fastvideo.train.utils.dataloader."
-        "build_parquet_t2v_train_dataloader",
-        lambda *args, **kwargs: None,
-    )
-
-    model = WanCausalModel(
-        init_from=cfg.models["student"]["init_from"],
-        training_config=cfg.training,
-        trainable=True,
-    )
-    model.transformer = model.transformer.to(device=device, dtype=dtype)
-
-    method = TeacherForcingSFTMethod(
-        cfg=cfg,
-        role_models={"student": model},
-    )
-    method.on_train_start()
-
-    batch = _build_synthetic_batch(device, dtype)
-    loss_map, outputs, _metrics = method.single_train_step(batch, iteration=0)
-
-    loss = loss_map["total_loss"]
-    assert torch.is_tensor(loss), "total_loss must be a torch.Tensor"
-    assert torch.isfinite(loss).item(), (
-        f"total_loss is not finite: {loss.item()}")
-
-    method.backward(loss_map, outputs, grad_accum_rounds=1)
-
-    blocks = getattr(model.transformer, "blocks", None)
-    assert blocks is not None and len(blocks) > 0, (
-        "CausalWanTransformer is expected to expose ``.blocks``")
-    layer0 = blocks[0]
-
-    trainable = [p for p in layer0.parameters() if p.requires_grad]
-    assert len(trainable) > 0, "layer 0 has no trainable parameters"
-
-    for i, p in enumerate(trainable):
-        assert p.grad is not None, f"layer 0 param[{i}] has None grad"
-        assert torch.isfinite(p.grad).all().item(), (
-            f"layer 0 param[{i}] grad contains NaN/Inf")
-
-    any_nonzero = any(
-        p.grad.detach().float().norm().item() > 0.0 for p in trainable)
-    assert any_nonzero, (
-        "all layer-0 grads are exactly zero; backward did not "
-        "reach the first transformer block")
-
-    # Teacher forcing must build its own (concatenated) attention mask and
-    # must not have constructed the diffusion-forcing mask.
-    assert model.transformer.teacher_forcing_block_mask is not None, (
-        "teacher-forcing mask was not constructed")
-    assert model.transformer.block_mask is None, (
-        "diffusion-forcing mask should not be built on the TF path")

fastvideo/tests/training/test_trackers.py

@@ -1,225 +0,0 @@
-from __future__ import annotations
-
-import builtins
-from io import BytesIO
-from pathlib import Path
-import sys
-from types import ModuleType
-from typing import Any
-
-import numpy as np
-from PIL import Image
-import pytest
-
-from fastvideo.training import trackers as trackers_module
-from fastvideo.training.trackers import (
-    _coerce_video_fps,
-    _prepare_video_array,
-    BaseTracker,
-    SequentialTracker,
-    SwanlabTracker,
-    WandbTracker,
-)
-
-
-class _FakeSwanlab(ModuleType):
-    def __init__(self) -> None:
-        super().__init__("swanlab")
-        self.init_kwargs: dict[str, Any] | None = None
-        self.logged: list[tuple[dict[str, Any], int]] = []
-        self.video_paths: list[str] = []
-        self.finished = False
-
-    def init(self, **kwargs: Any) -> object:
-        self.init_kwargs = kwargs
-        return object()
-
-    def log(self, metrics: dict[str, Any], step: int) -> None:
-        self.logged.append((metrics, step))
-
-    def finish(self) -> None:
-        self.finished = True
-
-    def Video(self, file_path: str, caption: str | None = None) -> dict[str, Any]:  # noqa: N802
-        self.video_paths.append(file_path)
-        return {
-            "content": Path(file_path).read_bytes(),
-            "caption": caption,
-        }
-
-
-def _install_fake_swanlab(monkeypatch: pytest.MonkeyPatch) -> _FakeSwanlab:
-    swanlab = _FakeSwanlab()
-    monkeypatch.setitem(sys.modules, "swanlab", swanlab)
-    return swanlab
-
-
-def test_prepare_video_array_converts_tchw_to_thwc() -> None:
-    video = np.arange(2 * 3 * 2 * 4, dtype=np.uint8).reshape(2, 3, 2, 4)
-
-    prepared = _prepare_video_array(video)
-
-    np.testing.assert_array_equal(prepared, video.transpose(0, 2, 3, 1))
-
-
-def test_default_video_fps_is_16() -> None:
-    assert _coerce_video_fps(None) == 16
-
-
-def test_wandb_video_uses_shared_default_fps(
-    monkeypatch: pytest.MonkeyPatch,
-    tmp_path: Path,
-) -> None:
-    wandb = ModuleType("wandb")
-    video_calls: list[dict[str, Any]] = []
-    wandb.init = lambda **_: object()
-
-    def video(data: Any, **kwargs: Any) -> dict[str, Any]:
-        video_calls.append(kwargs)
-        return {"data": data, **kwargs}
-
-    wandb.Video = video
-    monkeypatch.setitem(sys.modules, "wandb", wandb)
-    tracker = WandbTracker("project", str(tmp_path))
-
-    tracker.video(np.zeros((2, 3, 2, 2), dtype=np.uint8))
-
-    assert video_calls == [{"fps": 16, "format": "mp4"}]
-
-
-def test_prepare_video_array_tiles_and_pads_batches() -> None:
-    video = np.zeros((3, 1, 3, 1, 1), dtype=np.uint8)
-    video[0] = 1
-    video[1] = 2
-    video[2] = 3
-
-    prepared = _prepare_video_array(video)
-
-    assert prepared.shape == (1, 1, 4, 3)
-    np.testing.assert_array_equal(
-        prepared[0, 0],
-        np.array(
-            [
-                [1, 1, 1],
-                [2, 2, 2],
-                [3, 3, 3],
-                [0, 0, 0],
-            ],
-            dtype=np.uint8,
-        ),
-    )
-
-
-def test_swanlab_video_converts_array_to_gif_and_logs_artifact(
-    monkeypatch: pytest.MonkeyPatch,
-    tmp_path: Path,
-) -> None:
-    swanlab = _install_fake_swanlab(monkeypatch)
-    tracker = SwanlabTracker("project", str(tmp_path), run_name="run")
-    video = np.zeros((2, 3, 2, 2), dtype=np.uint8)
-    video[0, 0] = 255
-    video[1, 2] = 255
-
-    artifact = tracker.video(video, caption="validation sample", fps=5, format="mp4")
-    tracker.log_artifacts({"validation/videos": [artifact]}, step=7)
-
-    assert artifact["content"].startswith((b"GIF87a", b"GIF89a"))
-    assert artifact["caption"] == "validation sample"
-    assert not Path(swanlab.video_paths[0]).exists()
-    with Image.open(BytesIO(artifact["content"])) as image:
-        assert image.n_frames == 2
-        assert image.info["duration"] == 200
-    assert swanlab.logged == [({"validation/videos": [artifact]}, 7)]
-
-
-@pytest.mark.parametrize(("requested_fps", "expected_fps"), [(None, 12.0), (2, 2.0)])
-def test_swanlab_video_uses_requested_or_source_fps_during_gif_conversion(
-    monkeypatch: pytest.MonkeyPatch,
-    tmp_path: Path,
-    requested_fps: int | None,
-    expected_fps: float,
-) -> None:
-    _install_fake_swanlab(monkeypatch)
-    tracker = SwanlabTracker("project", str(tmp_path))
-    frames = [np.zeros((2, 2, 3), dtype=np.uint8), np.ones((2, 2, 3), dtype=np.uint8)]
-    written_fps: list[float] = []
-    write_gif = trackers_module._write_gif
-
-    monkeypatch.setattr(trackers_module, "_read_video_file", lambda _: (frames, 12.0))
-
-    def capture_fps(file_path: str, video_frames: Any, fps: float) -> None:
-        written_fps.append(fps)
-        write_gif(file_path, video_frames, fps)
-
-    monkeypatch.setattr(trackers_module, "_write_gif", capture_fps)
-
-    artifact = tracker.video(tmp_path / "validation.mp4", fps=requested_fps)
-
-    assert artifact["content"].startswith((b"GIF87a", b"GIF89a"))
-    assert written_fps == [expected_fps]
-
-
-def test_swanlab_video_passes_existing_gif_through(
-    monkeypatch: pytest.MonkeyPatch,
-    tmp_path: Path,
-) -> None:
-    swanlab = _install_fake_swanlab(monkeypatch)
-    tracker = SwanlabTracker("project", str(tmp_path))
-    gif_path = tmp_path / "validation.gif"
-    Image.new("RGB", (2, 2)).save(gif_path, format="GIF")
-
-    artifact = tracker.video(gif_path, caption="existing")
-
-    assert swanlab.video_paths == [str(gif_path)]
-    assert gif_path.exists()
-    assert artifact["caption"] == "existing"
-
-
-def test_swanlab_missing_dependency_error_is_actionable(
-    monkeypatch: pytest.MonkeyPatch,
-    tmp_path: Path,
-) -> None:
-    monkeypatch.delitem(sys.modules, "swanlab", raising=False)
-    real_import = builtins.__import__
-
-    def missing_swanlab(name: str, *args: Any, **kwargs: Any) -> Any:
-        if name == "swanlab":
-            raise ModuleNotFoundError("No module named 'swanlab'", name="swanlab")
-        return real_import(name, *args, **kwargs)
-
-    monkeypatch.setattr(builtins, "__import__", missing_swanlab)
-
-    with pytest.raises(ModuleNotFoundError, match=r"fastvideo\[swanlab\]"):
-        SwanlabTracker("project", str(tmp_path))
-
-
-class _RecordingTracker(BaseTracker):
-    def __init__(self, name: str) -> None:
-        super().__init__()
-        self.name = name
-        self.artifact_logs: list[dict[str, Any]] = []
-
-    def video(
-        self,
-        data: Any,
-        *,
-        caption: str | None = None,
-        fps: int | None = None,
-        format: str | None = None,
-    ) -> str:
-        return f"{self.name}:{data}"
-
-    def log_artifacts(self, artifacts: dict[str, Any], step: int) -> None:
-        self.artifact_logs.append(artifacts)
-
-
-def test_sequential_tracker_logs_backend_specific_video_artifacts() -> None:
-    first = _RecordingTracker("first")
-    second = _RecordingTracker("second")
-    tracker = SequentialTracker([first, second])
-
-    artifact = tracker.video("sample.mp4")
-    tracker.log_artifacts({"validation/videos": [artifact]}, step=3)
-
-    assert first.artifact_logs == [{"validation/videos": ["first:sample.mp4"]}]
-    assert second.artifact_logs == [{"validation/videos": ["second:sample.mp4"]}]

fastvideo/tests/transformers/test_hunyuanvideo.py

@@ -0,0 +1,104 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import glob
+import json
+import os
+import pytest
+import torch
+
+from fastvideo.configs.pipelines.base import PipelineConfig
+from fastvideo.distributed.parallel_state import get_sp_parallel_rank, get_sp_world_size
+from fastvideo.logger import init_logger
+from fastvideo.models.loader.component_loader import TransformerLoader
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.utils import maybe_download_model
+from fastvideo.forward_context import set_forward_context
+from fastvideo.configs.models.dits import HunyuanVideoConfig
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+
+logger = init_logger(__name__)
+
+os.environ["MASTER_ADDR"] = "localhost"
+os.environ["MASTER_PORT"] = "29503"
+
+BASE_MODEL_PATH = "hunyuanvideo-community/HunyuanVideo"
+MODEL_PATH = maybe_download_model(BASE_MODEL_PATH, local_dir=os.path.join("data", BASE_MODEL_PATH))
+TRANSFORMER_PATH = os.path.join(MODEL_PATH, "transformer")
+CONFIG_PATH = os.path.join(TRANSFORMER_PATH, "config.json")
+
+LOCAL_RANK = 0
+RANK = 0
+WORLD_SIZE = 1
+
+# Latent generated on commit c021e8a27cf437ac22827f2bc58b7f006561317f with 1 x L40S
+REFERENCE_LATENT = 93.23263549804688
+
+
+@pytest.mark.usefixtures("distributed_setup")
+@pytest.mark.skip(reason="Skipping test_hunyuanvideo_distributed")
+def test_hunyuanvideo_distributed():
+    logger.info(f"Initializing process: rank={RANK}, local_rank={LOCAL_RANK}, world_size={WORLD_SIZE}")
+
+    torch.cuda.set_device(f"cuda:{LOCAL_RANK}")
+
+    # Get tensor parallel info
+    sp_rank = get_sp_parallel_rank()
+    sp_world_size = get_sp_world_size()
+
+    logger.info(f"Process rank {RANK} initialized with SP rank {sp_rank} in SP world size {sp_world_size}")
+
+    config = json.load(open(CONFIG_PATH))
+    # remove   "_class_name": "HunyuanVideoTransformer3DModel",   "_diffusers_version": "0.32.0.dev0",
+    # TODO: write normalize config function
+    config.pop("_class_name")
+    config.pop("_diffusers_version")
+
+    precision_str = "bf16"
+    args = FastVideoArgs(
+        model_path=TRANSFORMER_PATH,
+        dit_cpu_offload=True,
+        pipeline_config=PipelineConfig(dit_config=HunyuanVideoConfig(), dit_precision=precision_str),
+    )
+    args.device = torch.device(f"cuda:{LOCAL_RANK}")
+
+    loader = TransformerLoader()
+    model = loader.load(TRANSFORMER_PATH, args)
+
+    model.eval()
+
+    # Create random inputs for testing
+    batch_size = 1
+    seq_len = 3
+    device = torch.device(f"cuda:{LOCAL_RANK}")
+
+    # Video latents [B, C, T, H, W]
+    hidden_states = torch.randn(batch_size, 16, 8, 16, 16, device=device, dtype=torch.bfloat16)
+    chunk_per_rank = hidden_states.shape[2] // sp_world_size
+    hidden_states = hidden_states[:, :, sp_rank * chunk_per_rank : (sp_rank + 1) * chunk_per_rank]
+
+    # Text embeddings [B, L, D] (including global token)
+    encoder_hidden_states = torch.randn(batch_size, seq_len + 1, 4096, device=device, dtype=torch.bfloat16)
+
+    # Timestep
+    timestep = torch.tensor([500], device=device, dtype=torch.bfloat16)
+    forward_batch = ForwardBatch(
+        data_type="dummy",
+    )
+
+    # Disable gradients for inference
+    with torch.no_grad():
+        # Run inference on model
+        with torch.amp.autocast(device_type="cuda", dtype=torch.bfloat16):
+            with set_forward_context(current_timestep=0, attn_metadata=None, forward_batch=forward_batch):
+                output = model(
+                    hidden_states=hidden_states,
+                    encoder_hidden_states=encoder_hidden_states,
+                    timestep=timestep,
+                )
+
+    latent = output.double().sum().item()
+
+    # Check if latents are similar
+    diff_output_latents = abs(REFERENCE_LATENT - latent)
+    logger.info(f"Reference latent: {REFERENCE_LATENT}, Current latent: {latent}")
+    assert diff_output_latents < 1e-4, f"Output latents differ significantly: max diff = {diff_output_latents}"

fastvideo/tests/transformers/test_matrixgame2.py

@@ -0,0 +1,133 @@
+# SPDX-License-Identifier: Apache-2.0
+import os
+
+import pytest
+import torch
+
+from fastvideo.configs.models.dits.matrixgame2 import MatrixGame2WanVideoConfig
+from fastvideo.configs.pipelines import PipelineConfig
+from fastvideo.distributed.parallel_state import (
+    get_sp_parallel_rank,
+    get_sp_world_size,
+)
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.forward_context import set_forward_context
+from fastvideo.logger import init_logger
+from fastvideo.models.loader.component_loader import TransformerLoader
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+from fastvideo.utils import maybe_download_model
+
+logger = init_logger(__name__)
+
+os.environ["MASTER_ADDR"] = "localhost"
+os.environ["MASTER_PORT"] = "29505"
+os.environ.setdefault("TORCHDYNAMO_DISABLE", "1")
+
+MODEL_PATH = maybe_download_model("FastVideo/Matrix-Game-2.0-Base-Distilled-Diffusers")
+TRANSFORMER_PATH = os.path.join(MODEL_PATH, "transformer")
+REFERENCE_LATENT = -111835.58895772696
+
+
+@pytest.mark.skip(reason="Not reliably reproducible")
+@pytest.mark.usefixtures("distributed_setup")
+def test_matrixgame2_transformer():
+    transformer_path = TRANSFORMER_PATH
+
+    sp_rank = get_sp_parallel_rank()
+    sp_world_size = get_sp_world_size()
+
+    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+    precision = torch.bfloat16
+    precision_str = "bf16"
+
+    args = FastVideoArgs(
+        model_path=transformer_path,
+        dit_cpu_offload=False,
+        use_fsdp_inference=False,
+        pipeline_config=PipelineConfig(
+            dit_config=MatrixGame2WanVideoConfig(), dit_precision=precision_str
+        ),
+    )
+    args.device = device
+
+    loader = TransformerLoader()
+    model = loader.load(transformer_path, args).to(device, dtype=precision)
+    model.eval()
+
+    total_params = sum(p.numel() for p in model.parameters())
+    weight_sum = sum(p.to(torch.float64).sum().item() for p in model.parameters())
+    weight_mean = weight_sum / total_params
+    logger.info("Total parameters: %s", total_params)
+    logger.info("Weight sum: %s", weight_sum)
+    logger.info("Weight mean: %s", weight_mean)
+
+    torch.manual_seed(42)
+
+    batch_size = 1
+    latent_frames = 15
+    latent_height = 44
+    latent_width = 80
+    image_dim = 1280
+    image_seq_len = 257
+
+    # Set num_frame_per_block to match actual latent_frames for action_model
+    model.num_frame_per_block = latent_frames
+    # Force full attention mode
+    model.local_attn_size = -1
+    # Clear cached block masks so they regenerate with correct settings
+    model.block_mask = None
+    model.block_mask_mouse = None
+    model.block_mask_keyboard = None
+
+    # Video latents [B, C, T, H, W]
+    x = torch.randn(batch_size, 16, latent_frames, latent_height, latent_width,
+                    device=device, dtype=precision)
+    cond_concat = torch.randn(batch_size, 20, latent_frames, latent_height, latent_width,
+                              device=device, dtype=precision)
+    hidden_states = torch.cat([x, cond_concat], dim=1)
+
+    if sp_world_size > 1:
+        chunk_per_rank = hidden_states.shape[2] // sp_world_size
+        hidden_states = hidden_states[:, :, sp_rank * chunk_per_rank:(sp_rank + 1) * chunk_per_rank]
+
+    # Image embeddings
+    encoder_hidden_states_image = [
+        torch.randn(batch_size, image_seq_len, image_dim, device=device, dtype=precision)
+    ]
+
+    # Timestep
+    timestep = torch.full((batch_size, latent_frames), 500, device=device, dtype=precision)
+
+    # Action conditions
+    num_pixel_frames = (latent_frames - 1) * 4 + 1
+    mouse_cond = torch.randn(batch_size, num_pixel_frames, 2, device=device, dtype=precision)
+    keyboard_cond = torch.randn(batch_size, num_pixel_frames, 4, device=device, dtype=precision)
+
+    # Text embeddings (empty for i2v)
+    encoder_hidden_states = torch.zeros(batch_size, 0, model.config.hidden_size,
+                                        device=device, dtype=precision)
+
+    forward_batch = ForwardBatch(data_type="dummy")
+    model.num_frame_per_block = latent_frames
+
+    with torch.no_grad():
+        with torch.amp.autocast('cuda', dtype=precision):
+            with set_forward_context(current_timestep=0, attn_metadata=None, forward_batch=forward_batch):
+                output = model(
+                    hidden_states=hidden_states,
+                    encoder_hidden_states=encoder_hidden_states,
+                    encoder_hidden_states_image=encoder_hidden_states_image,
+                    timestep=timestep,
+                    mouse_cond=mouse_cond,
+                    keyboard_cond=keyboard_cond,
+                )
+
+    latent = output.double().sum().item()
+
+    diff = abs(REFERENCE_LATENT - latent)
+    relative_diff = diff / abs(REFERENCE_LATENT)
+    logger.info(f"Reference latent: {REFERENCE_LATENT}, Current latent: {latent}")
+    logger.info(f"Absolute diff: {diff}, Relative diff: {relative_diff * 100:.4f}%")
+
+    # The 0.19% diff now
+    assert relative_diff < 0.005, f"Output latents differ significantly: relative diff = {relative_diff * 100:.4f}% (max allowed: 0.5%)"

fastvideo/tests/vaes/test_hunyuan_vae.py

@@ -0,0 +1,82 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import json
+import os
+
+import pytest
+import torch
+from safetensors.torch import load_file
+
+from fastvideo.configs.pipelines import PipelineConfig
+from fastvideo.logger import init_logger
+
+# from fastvideo.models.vaes.hunyuanvae import (
+#     AutoencoderKLHunyuanVideo as MyHunyuanVAE)
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.models.loader.component_loader import VAELoader
+from fastvideo.configs.models.vaes import HunyuanVAEConfig
+from fastvideo.utils import maybe_download_model
+
+logger = init_logger(__name__)
+
+os.environ["MASTER_ADDR"] = "localhost"
+os.environ["MASTER_PORT"] = "29503"
+
+BASE_MODEL_PATH = "hunyuanvideo-community/HunyuanVideo"
+MODEL_PATH = maybe_download_model(BASE_MODEL_PATH, local_dir=os.path.join("data", BASE_MODEL_PATH))
+VAE_PATH = os.path.join(MODEL_PATH, "vae")
+CONFIG_PATH = os.path.join(VAE_PATH, "config.json")
+
+# Latent generated on commit d71a4ebffc2034922fc379568b6a6aa722f3744c with 1 x A40
+# torch 2.7.1
+A40_REFERENCE_LATENT = -106.22467041015625
+
+# Latent generated on commit 2b54068960c41d42221e8b8719a374b499855029 with 1 x L40S
+L40S_REFERENCE_LATENT = -157.09130859375
+
+
+@pytest.mark.usefixtures("distributed_setup")
+@pytest.mark.skip(reason="Skipping test_hunyuan_vae")
+def test_hunyuan_vae():
+    device = torch.device("cuda:0")
+    precision = torch.bfloat16
+    precision_str = "bf16"
+    args = FastVideoArgs(
+        model_path=VAE_PATH, pipeline_config=PipelineConfig(vae_config=HunyuanVAEConfig(), vae_precision=precision_str)
+    )
+    args.device = device
+    args.vae_cpu_offload = False
+
+    loader = VAELoader()
+    model = loader.load(VAE_PATH, args)
+
+    model.enable_tiling(
+        tile_sample_min_height=32,
+        tile_sample_min_width=32,
+        tile_sample_min_num_frames=8,
+        tile_sample_stride_height=16,
+        tile_sample_stride_width=16,
+        tile_sample_stride_num_frames=4,
+    )
+
+    batch_size = 1
+
+    # Video input [B, C, T, H, W]
+    input_tensor = torch.randn(batch_size, 3, 21, 64, 64, device=device, dtype=torch.bfloat16)
+
+    # Disable gradients for inference
+    with torch.no_grad():
+        latent = model.encode(input_tensor).mean.double().sum().item()
+
+    # Check if latents are similar
+    device_name = torch.cuda.get_device_name()
+    if "A40" in device_name:
+        REFERENCE_LATENT = A40_REFERENCE_LATENT
+    elif "L40S" in device_name:
+        REFERENCE_LATENT = L40S_REFERENCE_LATENT
+    else:
+        raise ValueError(f"Unknown device: {device_name}")
+
+    diff_encoded_latents = abs(REFERENCE_LATENT - latent)
+    logger.info(f"Reference latent: {REFERENCE_LATENT}, Current latent: {latent}")
+    assert diff_encoded_latents < 1e-4, f"Encoded latents differ significantly: max diff = {diff_encoded_latents}"

fastvideo/tests/vaes/test_official_wan_vae.py

@@ -0,0 +1,113 @@
+# SPDX-License-Identifier: Apache-2.0
+import os
+
+import numpy as np
+import pytest
+import torch
+from diffusers import AutoencoderKLWan
+
+from fastvideo.configs.pipelines import PipelineConfig
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.logger import init_logger
+from fastvideo.models.loader.component_loader import VAELoader
+from fastvideo.configs.models.vaes import WanVAEConfig
+from fastvideo.utils import maybe_download_model
+
+try:
+    from wan.modules.vae2_2 import Wan2_2_VAE
+except ImportError:
+    Wan2_2_VAE = None
+
+logger = init_logger(__name__)
+
+os.environ["MASTER_ADDR"] = "localhost"
+os.environ["MASTER_PORT"] = "29503"
+
+BASE_MODEL_PATH = "Wan-AI/Wan2.2-TI2V-5B-Diffusers"
+MODEL_PATH = maybe_download_model(BASE_MODEL_PATH)
+VAE_PATH = os.path.join(MODEL_PATH, "vae")
+
+
+@pytest.mark.skip(reason="disable test")
+@pytest.mark.usefixtures("distributed_setup")
+def test_official_wan_vae():
+    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+    precision = torch.bfloat16
+    precision_str = "bf16"
+    args = FastVideoArgs(model_path=VAE_PATH, pipeline_config=PipelineConfig(vae_config=WanVAEConfig(), vae_precision=precision_str))
+    args.device = device
+    args.vae_cpu_offload = False
+
+    loader = VAELoader()
+    model2 = loader.load(VAE_PATH, args)
+    assert model2.use_feature_cache # Default to use the original WanVAE algorithm
+
+    model1 = Wan2_2_VAE(vae_pth="/mnt/weka/home/hao.zhang/wei/Wan2.2_VAE.pth", device=device, dtype=precision)
+
+    # Create identical inputs for both models
+    batch_size = 1
+
+    # Video input [B, C, T, H, W]
+    input_tensor = torch.randn(batch_size,
+                               3,
+                               81,
+                               32,
+                               32,
+                               device=device,
+                               dtype=precision)
+    # latent_tensor = torch.randn(batch_size,
+    #                             16,
+    #                             21,
+    #                             32,
+    #                             32,
+    #                             device=device,
+    #                             dtype=precision)
+
+    # Disable gradients for inference
+    with torch.no_grad():
+        # Test encoding
+        logger.info("Testing encoding...")
+        latent1 = model1.encode([input_tensor.squeeze(0)])[0].unsqueeze(0)
+        print("--------------------------------")
+        latent2 = model2.encode(input_tensor)
+        # Check if latents have the same shape
+        assert latent1.shape == latent2.mean.shape, f"Latent shapes don't match: {latent1.mean.shape} vs {latent2.mean.shape}"
+        # Check if latents are similar
+        max_diff_encode = torch.max(torch.abs(latent1 - latent2.mean))
+        mean_diff_encode = torch.mean(torch.abs(latent1 - latent2.mean))
+        logger.info("Maximum difference between encoded latents: %s",
+                    max_diff_encode.item())
+        logger.info("Mean difference between encoded latents: %s",
+                    mean_diff_encode.item())
+        assert max_diff_encode < 1e-5, f"Encoded latents differ significantly: max diff = {mean_diff_encode.item()}"
+        # Test decoding
+        logger.info("Testing decoding...")
+        # latent1_tensor = latent1.mode()
+        # mean1 = (torch.tensor(model1.config.latents_mean).view(
+        #     1, model1.config.z_dim, 1, 1, 1).to(input_tensor.device,
+        #                                         input_tensor.dtype))
+        # std1 = (1.0 / torch.tensor(model1.config.latents_std).view(
+        #     1, model1.config.z_dim, 1, 1, 1)).to(input_tensor.device,
+        #                                         input_tensor.dtype)
+        # latent1_tensor = latent1_tensor / std1 + mean1
+        # output1 = model1.decode(latent1_tensor).sample
+
+        output1 = model1.decode([latent1.squeeze(0)])
+        output1 = output1[0].unsqueeze(0)
+
+        mean2 = model2.config.arch_config.shift_factor.to(input_tensor.device, input_tensor.dtype)
+        std2 = model2.config.arch_config.scaling_factor.to(input_tensor.device, input_tensor.dtype)
+        latent2_tensor = latent2.mode()
+        latent2_tensor = latent2_tensor / std2 + mean2
+        output2 = model2.decode(latent2_tensor)
+        # Check if outputs have the same shape
+        assert output1.shape == output2.shape, f"Output shapes don't match: {output1.shape} vs {output2.shape}"
+
+        # Check if outputs are similar
+        max_diff_decode = torch.max(torch.abs(output1.float() - output2.float()))
+        mean_diff_decode = torch.mean(torch.abs(output1.float() - output2.float()))
+        logger.info("Maximum difference between decoded outputs: %s",
+                    max_diff_decode.item())
+        logger.info("Mean difference between decoded outputs: %s",
+                    mean_diff_decode.item())
+        assert max_diff_decode < 1e-5, f"Decoded outputs differ significantly: max diff = {max_diff_decode.item()}"

fastvideo/train/callbacks/validation.py

@@ -404,7 +404,6 @@ class ValidationCallback(Callback):
                         art = self.tracker.video(
                             fname,
                             caption=cap,
-                            fps=sp.fps,
                         )
                         if art is not None:
                             video_logs.append(art)
@@ -457,16 +456,11 @@ class ValidationCallback(Callback):
             None,
         )
 
-        loaded_modules: dict[str, Any] = {"transformer": transformer}
-        # Distillation methods build the flow-match scheduler their few-step DMD
-        # sampler needs; inject it so the pipeline doesn't fall back to UniPC.
-        method_scheduler = getattr(self.method, "_sf_scheduler", None)
-        if method_scheduler is not None:
-            loaded_modules["scheduler"] = method_scheduler
-
         kwargs: dict[str, Any] = {
             "inference_mode": True,
-            "loaded_modules": loaded_modules,
+            "loaded_modules": {
+                "transformer": transformer,
+            },
             "tp_size": tc.distributed.tp_size,
             "sp_size": tc.distributed.sp_size,
             "num_gpus": tc.distributed.num_gpus,
@@ -483,6 +477,12 @@ class ValidationCallback(Callback):
             **kwargs,
         )
 
+        scheduler = self._pipeline.get_module("scheduler")
+        if (scheduler is not None and type(scheduler).__name__ == "SelfForcingFlowMatchScheduler"):
+            scheduler.sigma_min = 0.0
+            scheduler.extra_one_step = True
+            scheduler.set_timesteps(num_inference_steps=1000, training=True)
+
         self._pipeline_key = key
         return self._pipeline
 

fastvideo/train/methods/__init__.py

@@ -9,8 +9,6 @@ __all__ = [
     "KDMethod",
     "SelfForcingMethod",
     "DiffusionForcingSFTMethod",
-    "TeacherForcingSFTMethod",
-    "CausalConsistencyDistillationMethod",
 ]
 
 
@@ -30,11 +28,4 @@ def __getattr__(name: str) -> object:
     if name == "DiffusionForcingSFTMethod":
         from fastvideo.train.methods.fine_tuning.dfsft import DiffusionForcingSFTMethod
         return DiffusionForcingSFTMethod
-    if name == "TeacherForcingSFTMethod":
-        from fastvideo.train.methods.fine_tuning.tfsft import TeacherForcingSFTMethod
-        return TeacherForcingSFTMethod
-    if name == "CausalConsistencyDistillationMethod":
-        from fastvideo.train.methods.consistency_model.causal_cd import (
-            CausalConsistencyDistillationMethod, )
-        return CausalConsistencyDistillationMethod
     raise AttributeError(name)

fastvideo/train/methods/consistency_model/__init__.py

@@ -1,22 +1,3 @@
 # SPDX-License-Identifier: Apache-2.0
 
-from __future__ import annotations
-
-from typing import TYPE_CHECKING
-
-if TYPE_CHECKING:
-    from fastvideo.train.methods.consistency_model.causal_cd import (
-        CausalConsistencyDistillationMethod, )
-
-__all__ = [
-    "CausalConsistencyDistillationMethod",
-]
-
-
-def __getattr__(name: str) -> object:
-    if name == "CausalConsistencyDistillationMethod":
-        from fastvideo.train.methods.consistency_model.causal_cd import (
-            CausalConsistencyDistillationMethod, )
-
-        return CausalConsistencyDistillationMethod
-    raise AttributeError(name)
+__all__: list[str] = []

fastvideo/train/methods/consistency_model/causal_cd.py

@@ -1,237 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Causal consistency distillation method (algorithm layer)."""
-
-from __future__ import annotations
-
-from typing import Any
-
-import torch
-import torch.nn.functional as F
-
-from fastvideo.models.schedulers.scheduling_self_forcing_flow_match import (SelfForcingFlowMatchScheduler)
-from fastvideo.train.methods.base import LogScalar, TrainingMethod
-from fastvideo.train.models.base import ModelBase
-from fastvideo.train.utils.checkpoint import _FullModelState
-from fastvideo.train.utils.optimizer import build_optimizer_and_scheduler
-
-
-class CausalConsistencyDistillationMethod(TrainingMethod):
-
-    def __init__(
-        self,
-        *,
-        cfg: Any,
-        role_models: dict[str, ModelBase],
-    ) -> None:
-        super().__init__(cfg=cfg, role_models=role_models)
-
-        for role in ("student", "teacher", "ema"):
-            if role not in role_models:
-                raise ValueError(f"Causal-CD requires role {role!r} "
-                                 "(student trainable; teacher + ema frozen, "
-                                 "both initialized from the student's "
-                                 "checkpoint)")
-        if not self.student._trainable:
-            raise ValueError("Causal-CD requires student to be trainable")
-        self.teacher = role_models["teacher"]
-        self.ema_model = role_models["ema"]
-
-        self._attn_kind = self._infer_attn_kind()
-        self._guidance_scale = float(self.method_config.get("guidance_scale", 3.0))
-        self._discrete_cd_n = int(self.method_config.get("discrete_cd_N", 48))
-        if self._discrete_cd_n < 2:
-            raise ValueError("method.discrete_cd_N must be >= 2")
-        self._ema_decay = float(self.method_config.get("ema_decay", 0.99))
-        self._ema_start_step = int(self.method_config.get("ema_start_step", 200))
-        shift = getattr(self.training_config.pipeline_config, "flow_shift", None)
-        self._flow_shift = float(shift) if shift else 5.0
-
-        self.student.init_preprocessors(self.training_config)
-        self._sf_scheduler = SelfForcingFlowMatchScheduler(
-            num_inference_steps=self._discrete_cd_n,
-            num_train_timesteps=int(self.student.num_train_timesteps),
-            shift=self._flow_shift,
-            sigma_min=0.0,
-            sigma_max=1.0,
-            extra_one_step=True,
-            training=False,
-        )
-        self._init_optimizers_and_schedulers()
-
-    # ------------------------------------------------------------------
-
-    @property
-    def _optimizer_dict(self) -> dict[str, Any]:
-        return {"student": self._student_optimizer}
-
-    @property
-    def _lr_scheduler_dict(self) -> dict[str, Any]:
-        return {"student": self._student_lr_scheduler}
-
-    def get_optimizers(self, iteration: int) -> list[torch.optim.Optimizer]:
-        del iteration
-        return [self._student_optimizer]
-
-    def get_lr_schedulers(self, iteration: int) -> list[Any]:
-        del iteration
-        return [self._student_lr_scheduler]
-
-    def checkpoint_state(self) -> dict[str, Any]:
-        # The EMA role is frozen (so the base class skips it) but mutated by
-        # _update_ema every step; without persisting it a resume reloads the
-        # EMA from init_from and the consistency target snaps back to the
-        # base checkpoint. Mirrors DiffusionNFT's frozen "old" role.
-        states = super().checkpoint_state()
-        states["roles.ema.transformer"] = _FullModelState(self.ema_model.transformer)
-        return states
-
-    # ------------------------------------------------------------------
-
-    def single_train_step(
-        self,
-        batch: dict[str, Any],
-        iteration: int,
-    ) -> tuple[dict[str, torch.Tensor], dict[str, Any], dict[str, LogScalar]]:
-        del iteration
-        training_batch = self.student.prepare_batch(
-            batch,
-            generator=self.cuda_generator,
-            latents_source="data",
-        )
-        clean_latents = training_batch.latents
-        if not torch.is_tensor(clean_latents) or clean_latents.ndim != 5:
-            raise ValueError("Causal-CD expects [B, T, C, H, W] latents")
-
-        batch_size, num_latents = int(clean_latents.shape[0]), int(clean_latents.shape[1])
-        device = clean_latents.device
-
-        sigmas = self._sf_scheduler.sigmas.to(device)
-        timesteps = self._sf_scheduler.timesteps.to(device)
-        idx = torch.randint(0, self._discrete_cd_n - 1, (1, ), generator=self.cuda_generator, device=device).squeeze(0)
-        t, t_next = timesteps[idx], timesteps[idx + 1]
-        sigma_t, sigma_t_next = sigmas[idx], sigmas[idx + 1]
-        t_pf = t * torch.ones(batch_size, num_latents, device=device)
-        t_next_pf = t_next * torch.ones(batch_size, num_latents, device=device)
-
-        noise = torch.randn(
-            clean_latents.shape,
-            generator=self.cuda_generator,
-            device=device,
-            dtype=clean_latents.dtype,
-        )
-        latent_t = (1.0 - sigma_t) * clean_latents + sigma_t * noise
-
-        # Set before any forward: predict_noise feeds batch.timesteps into
-        # set_forward_context (VSA sparsity gating), so the teacher CFG
-        # passes below must not see the stale timesteps from prepare_batch.
-        training_batch.timesteps = t_pf
-
-        with torch.no_grad():
-            v_cond = self._predict_flow(self.teacher,
-                                        latent_t,
-                                        t_pf,
-                                        training_batch,
-                                        conditional=True,
-                                        clean_x=clean_latents)
-            v_uncond = self._predict_flow(self.teacher,
-                                          latent_t,
-                                          t_pf,
-                                          training_batch,
-                                          conditional=False,
-                                          clean_x=clean_latents)
-            v_pred = v_uncond + self._guidance_scale * (v_cond - v_uncond)
-            dt = ((t - t_next) / float(self.student.num_train_timesteps))
-            latent_t_next = latent_t - dt * v_pred
-
-        flow_student = self._predict_flow(self.student,
-                                          latent_t,
-                                          t_pf,
-                                          training_batch,
-                                          conditional=True,
-                                          clean_x=clean_latents)
-        x0_t = latent_t - sigma_t * flow_student
-
-        with torch.no_grad():
-            flow_ema = self._predict_flow(self.ema_model,
-                                          latent_t_next,
-                                          t_next_pf,
-                                          training_batch,
-                                          conditional=True,
-                                          clean_x=clean_latents)
-            x0_t_next = latent_t_next - sigma_t_next * flow_ema
-
-        loss = F.mse_loss(x0_t.float(), x0_t_next.float())
-
-        loss_map = {"total_loss": loss, "causal_cd_loss": loss}
-        attn_metadata = (training_batch.attn_metadata_vsa if self._attn_kind == "vsa" else training_batch.attn_metadata)
-        outputs: dict[str, Any] = {"_fv_backward": (t_pf, attn_metadata)}
-        metrics: dict[str, LogScalar] = {}
-        return loss_map, outputs, metrics
-
-    # ------------------------------------------------------------------
-
-    def backward(
-        self,
-        loss_map: dict[str, torch.Tensor],
-        outputs: dict[str, Any],
-        *,
-        grad_accum_rounds: int = 1,
-    ) -> None:
-        grad_accum_rounds = max(1, int(grad_accum_rounds))
-        ctx = outputs.get("_fv_backward")
-        if ctx is None:
-            super().backward(loss_map, outputs, grad_accum_rounds=grad_accum_rounds)
-            return
-        self.student.backward(loss_map["total_loss"], ctx, grad_accum_rounds=grad_accum_rounds)
-
-    def optimizers_schedulers_step(self, iteration: int) -> None:
-        super().optimizers_schedulers_step(iteration)
-        if iteration >= self._ema_start_step:
-            self._update_ema()
-
-    # ------------------------------------------------------------------
-    # Internal helpers
-    # ------------------------------------------------------------------
-
-    def _predict_flow(
-        self,
-        model: ModelBase,
-        latents: torch.Tensor,
-        timestep: torch.Tensor,
-        batch: Any,
-        *,
-        conditional: bool,
-        clean_x: torch.Tensor,
-    ) -> torch.Tensor:
-        return model.predict_noise(latents,
-                                   timestep,
-                                   batch,
-                                   conditional=conditional,
-                                   cfg_uncond=None,
-                                   attn_kind=self._attn_kind,
-                                   clean_x=clean_x)
-
-    @torch.no_grad()
-    def _update_ema(self) -> None:
-        decay = self._ema_decay
-        for ema_p, p in zip(self.ema_model.transformer.parameters(), self.student.transformer.parameters(),
-                            strict=True):
-            ema_p.mul_(decay).add_(p.detach().to(ema_p.dtype), alpha=1.0 - decay)
-
-    def _init_optimizers_and_schedulers(self) -> None:
-        tc = self.training_config
-        student_lr = float(tc.optimizer.learning_rate)
-        if student_lr <= 0.0:
-            raise ValueError("training.learning_rate must be > 0 for causal-cd")
-        student_params = [p for p in self.student.transformer.parameters() if p.requires_grad]
-        (
-            self._student_optimizer,
-            self._student_lr_scheduler,
-        ) = build_optimizer_and_scheduler(
-            params=student_params,
-            optimizer_config=tc.optimizer,
-            loop_config=tc.loop,
-            learning_rate=student_lr,
-            betas=tc.optimizer.betas,
-            scheduler_name=str(tc.optimizer.lr_scheduler),
-        )

fastvideo/train/methods/fine_tuning/__init__.py

@@ -7,12 +7,10 @@ from typing import TYPE_CHECKING
 if TYPE_CHECKING:
     from fastvideo.train.methods.fine_tuning.dfsft import DiffusionForcingSFTMethod
     from fastvideo.train.methods.fine_tuning.finetune import FineTuneMethod
-    from fastvideo.train.methods.fine_tuning.tfsft import TeacherForcingSFTMethod
 
 __all__ = [
     "DiffusionForcingSFTMethod",
     "FineTuneMethod",
-    "TeacherForcingSFTMethod",
 ]
 
 
@@ -27,9 +25,4 @@ def __getattr__(name: str) -> object:
         from fastvideo.train.methods.fine_tuning.finetune import FineTuneMethod
 
         return FineTuneMethod
-    if name == "TeacherForcingSFTMethod":
-        from fastvideo.train.methods.fine_tuning.tfsft import (
-            TeacherForcingSFTMethod, )
-
-        return TeacherForcingSFTMethod
     raise AttributeError(name)

fastvideo/train/methods/fine_tuning/dfsft.py

@@ -135,11 +135,12 @@ class DiffusionForcingSFTMethod(TrainingMethod):
             t_inhom.flatten(),
         )
 
-        pred = self._predict_noise(
+        pred = self.student.predict_noise(
             noisy_latents,
             t_inhom,
             training_batch,
-            clean_latents,
+            conditional=True,
+            attn_kind=self._attn_kind,
         )
 
         if bool(self.training_config.model.precondition_outputs):
@@ -177,23 +178,6 @@ class DiffusionForcingSFTMethod(TrainingMethod):
         metrics: dict[str, LogScalar] = {}
         return loss_map, outputs, metrics
 
-    def _predict_noise(
-        self,
-        noisy_latents: torch.Tensor,
-        timestep: torch.Tensor,
-        training_batch: Any,
-        clean_latents: torch.Tensor,
-    ) -> torch.Tensor:
-        # Unused here; the teacher-forcing subclass overrides this to pass clean_x.
-        del clean_latents
-        return self.student.predict_noise(
-            noisy_latents,
-            timestep,
-            training_batch,
-            conditional=True,
-            attn_kind=self._attn_kind,
-        )
-
     # TrainingMethod override: backward
     def backward(
         self,

fastvideo/train/methods/fine_tuning/tfsft.py

@@ -1,30 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Teacher-forcing SFT method (TFSFT; algorithm layer)."""
-
-from __future__ import annotations
-
-from typing import Any
-
-import torch
-
-from fastvideo.train.methods.fine_tuning.dfsft import (
-    DiffusionForcingSFTMethod, )
-
-
-class TeacherForcingSFTMethod(DiffusionForcingSFTMethod):
-
-    def _predict_noise(
-        self,
-        noisy_latents: torch.Tensor,
-        timestep: torch.Tensor,
-        training_batch: Any,
-        clean_latents: torch.Tensor,
-    ) -> torch.Tensor:
-        return self.student.predict_noise(
-            noisy_latents,
-            timestep,
-            training_batch,
-            conditional=True,
-            attn_kind=self._attn_kind,
-            clean_x=clean_latents,
-        )

fastvideo/train/methods/rl/diffusion_nft.py

@@ -10,6 +10,12 @@ from typing import Any
 
 import torch
 import torch.distributed as dist
+from torch.distributed.checkpoint.state_dict import (
+    StateDictOptions,
+    get_model_state_dict,
+    set_model_state_dict,
+)
+from torch.distributed.checkpoint.stateful import Stateful
 from tqdm.auto import tqdm
 
 from fastvideo.dataset.parquet_dataset_map_style import (
@@ -31,7 +37,6 @@ from fastvideo.train.methods.rl.common import (
     validation_caption,
     validation_shard_indices,
 )
-from fastvideo.train.utils.checkpoint import _FullModelState
 from fastvideo.train.utils.config import (
     get_optional_float,
     get_optional_int,
@@ -73,6 +78,31 @@ class _DiffusionNFTEMAState:
         self._method._ema_update_count = int(update_count)
 
 
+class _FullModelState(Stateful):
+    """DCP wrapper that saves frozen model parameters too.
+
+    The shared ``ModelWrapper`` intentionally filters to ``requires_grad``
+    parameters. DiffusionNFT's old policy is frozen but must be restored on
+    resume, so it needs full model state.
+    """
+
+    def __init__(self, model: torch.nn.Module) -> None:
+        self.model = model
+
+    def state_dict(self) -> dict[str, Any]:
+        return get_model_state_dict(self.model)  # type: ignore[no-any-return]
+
+    def load_state_dict(
+        self,
+        state_dict: dict[str, Any],
+    ) -> None:
+        set_model_state_dict(
+            self.model,
+            model_state_dict=state_dict,
+            options=StateDictOptions(strict=False),
+        )
+
+
 class DiffusionNFTMethod(TrainingMethod):
     """DiffusionNFT-style RL for diffusion models.
 

fastvideo/train/models/wan/wan.py

@@ -320,8 +320,6 @@ class WanModel(ModelBase):
         conditional: bool,
         cfg_uncond: dict[str, Any] | None = None,
         attn_kind: Literal["dense", "vsa"] = "dense",
-        clean_x: torch.Tensor | None = None,
-        aug_t: torch.Tensor | None = None,
     ) -> torch.Tensor:
         device_type = self.device.type
         dtype = self._get_training_dtype()
@@ -349,11 +347,7 @@ class WanModel(ModelBase):
                 current_timestep=batch.timesteps,
                 attn_metadata=attn_metadata,
         ):
-            input_kwargs = (self._build_distill_input_kwargs(noisy_latents,
-                                                             timestep,
-                                                             text_dict,
-                                                             clean_x=clean_x,
-                                                             aug_t=aug_t))
+            input_kwargs = (self._build_distill_input_kwargs(noisy_latents, timestep, text_dict))
             transformer = self._get_transformer(timestep)
             pred_noise = transformer(**input_kwargs).permute(0, 2, 1, 3, 4)
         return pred_noise
@@ -536,24 +530,17 @@ class WanModel(ModelBase):
         noise_input: torch.Tensor,
         timestep: torch.Tensor,
         text_dict: dict[str, torch.Tensor] | None,
-        clean_x: torch.Tensor | None = None,
-        aug_t: torch.Tensor | None = None,
     ) -> dict[str, Any]:
         if text_dict is None:
             raise ValueError("text_dict cannot be None for "
                              "Wan distillation")
-        kwargs: dict[str, Any] = {
+        return {
             "hidden_states": noise_input.permute(0, 2, 1, 3, 4),
             "encoder_hidden_states": text_dict["encoder_hidden_states"],
             "encoder_attention_mask": text_dict["encoder_attention_mask"],
             "timestep": timestep,
             "return_dict": False,
         }
-        if clean_x is not None:
-            # Teacher forcing: clean context latents (+ optional aug timestep).
-            kwargs["clean_x"] = clean_x.permute(0, 2, 1, 3, 4)
-            kwargs["aug_t"] = aug_t
-        return kwargs
 
     def _get_transformer(self, timestep: torch.Tensor) -> torch.nn.Module:
         return self.transformer

fastvideo/train/models/wan/wan_causal.py

@@ -49,7 +49,6 @@ class WanCausalModel(WanModel, CausalModelBase):
         transformer_override_safetensor: str
         | None = None,
         lora: LoraConfig | dict[str, Any] | None = None,
-        num_frames_per_block: int | None = None,
     ) -> None:
         super().__init__(
             init_from=init_from,
@@ -63,16 +62,6 @@ class WanCausalModel(WanModel, CausalModelBase):
         )
         self._streaming_caches: (dict[tuple[int, str], _StreamingCaches]) = {}
 
-        if num_frames_per_block is not None:
-            num_frames_per_block = int(num_frames_per_block)
-            if not 1 <= num_frames_per_block <= 3:
-                # Same bound as CausalWanTransformer3DModel's config path
-                # (assert num_frame_per_block <= 3); this override must not
-                # bypass it.
-                raise ValueError("num_frames_per_block must be between 1 and 3, "
-                                 f"got {num_frames_per_block}")
-            self.transformer.num_frame_per_block = num_frames_per_block
-
     # --- CausalModelBase override: clear_caches ---
     def clear_caches(
         self,