[ci]: collect the tracker unit tests in the unit lane

fastvideo/tests/training/test_trackers.py (CPU-only, fake-swanlab
fixtures) was not collected by any lane - the training lanes pick
specific subdirs and run_unit_test skipped tests/training entirely.

.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: `PERF_MAX_REGRESSION` if set, otherwise `0.05` (5%). |
+| `max_intra_batch_regression` | No | Maximum allowed regression of any source JSON against the source batch median. Default: `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,8 +148,7 @@ 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 `PERF_MAX_REGRESSION` when set,
-otherwise `0.05`. Print a table with per-source values, batch median, and
+Default `max_intra_batch_regression` to `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
@@ -183,7 +182,7 @@ present, that run is not a valid source for baseline reseeding.
 
 ### 2. Sync and back up existing HF records under /tmp
 
-Use `fastvideo/tests/performance/hf_store.py` helpers directly. Do **not** use
+Use `fastvideo/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.
 
@@ -192,7 +191,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}"
-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)'
+python -c 'from fastvideo.performance.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`:
@@ -200,8 +199,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=$(PYTHONPATH=fastvideo/tests/performance python - <<'PY'
-from hf_store import sanitize
+MODEL_SAFE=$(python - <<'PY'
+from fastvideo.performance.hf_store import sanitize
 print(sanitize("<model_id>"))
 PY
 )
@@ -235,7 +234,7 @@ first baseline seed. Continue, but report that baseline history was empty.
 Load the last 5 successful records for the target:
 
 ```python
-from hf_store import load_records_for_model
+from fastvideo.performance.hf_store import load_records_for_model
 
 records = load_records_for_model(
     "/tmp/perf-tracking",
@@ -372,7 +371,7 @@ prepared records plus backup on disk.
 Use the shared storage helper so the path and repo type match CI:
 
 ```python
-from hf_store import upload_record
+from fastvideo.performance.hf_store import upload_record
 
 upload_record("<local_record_path>", record, strict=True)
 ```
@@ -460,7 +459,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/tests/performance/hf_store.py` — HF sync, record loading,
+- `fastvideo/performance/hf_store.py` — HF sync, record loading,
   `sanitize()`, and `upload_record()`.
 - `fastvideo/tests/performance/test_inference_performance.py` — source result
   JSON schema.

.github/workflows/ci-precommit.yml

@@ -1,7 +1,11 @@
 name: pre-commit
 
 on:
-  pull_request:
+  # 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:
     branches: [main]
   workflow_call:
     inputs:
@@ -15,12 +19,25 @@ permissions:
 
 jobs:
   pre-commit:
-    if: github.event_name == 'workflow_call' || github.event.pull_request.draft != true
+    if: 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"

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=940cd9680f3315f2f06b43ab5bea2c2cf2d96806
+ARG FA4_CUTE_REF=82d6441eec5d4dfec120153db2c0145ae855a083
 
 # 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 cutlass-4.5-safe upstream FA4 cute (FA4_CUTE_REF) over the
+# Overlay the CuTe-DSL-4.6-compatible 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,10 +112,11 @@ 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, 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
+`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
 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
@@ -432,5 +433,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 name is listed in `STAGE_METRIC_MAP` in
-`test_inference_performance.py`.
+and that the stage emits `component_metric` or is covered by the legacy
+`STAGE_METRIC_MAP` fallback in `test_inference_performance.py`.

docs/design/inference_schema_parity_inventory.yaml

@@ -41,11 +41,6 @@ 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

docs/training/trackers.md

@@ -0,0 +1,91 @@
+# 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,6 +161,9 @@ 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

fastvideo/api/compat.py

@@ -159,16 +159,6 @@ 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":
@@ -254,11 +244,6 @@ 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,15 +71,6 @@ 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
@@ -91,7 +82,6 @@ class EngineConfig:
     use_fsdp_inference: bool = False
     disable_autocast: bool = False
     quantization: QuantizationConfig | None = None
-    batching: BatchingConfig = field(default_factory=BatchingConfig)
 
 
 @dataclass
@@ -290,7 +280,6 @@ class ServeConfig:
 
 
 __all__ = [
-    "BatchingConfig",
     "CompileConfig",
     "ComponentConfig",
     "ContinuationState",

fastvideo/batching/__init__.py

@@ -1,26 +0,0 @@
-# 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

@@ -1,297 +0,0 @@
-# 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

@@ -1,178 +0,0 @@
-# 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/pipelines/base.py

@@ -276,27 +276,6 @@ 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/distributed/parallel_state.py

@@ -23,6 +23,7 @@ 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
@@ -1006,6 +1007,13 @@ 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,7 +10,6 @@ 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,
@@ -60,29 +59,11 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]:
     generator = VideoGenerator.from_fastvideo_args(args)
     logger.info("Model loaded successfully.")
 
-    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,
-    )
+    set_state(generator, args, output_dir, default_request=default_request)
 
     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

@@ -1,183 +0,0 @@
-# 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,7 +11,6 @@ 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
 
@@ -21,7 +20,6 @@ _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:
@@ -46,31 +44,23 @@ 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, _video_batch_scheduler
+    global _generator, _fastvideo_args, _output_dir, _default_request
     _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, _video_batch_scheduler
+    global _generator, _fastvideo_args, _default_request
     _generator = None
     _fastvideo_args = None
     _default_request = None
-    _video_batch_scheduler = None

fastvideo/entrypoints/openai/video_api.py

@@ -26,7 +26,6 @@ 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,
@@ -152,19 +151,15 @@ 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()
 
-        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),
-            )
+        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,7 +18,6 @@ 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
@@ -27,8 +26,6 @@ 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,
@@ -97,11 +94,6 @@ _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",
 })
@@ -120,17 +112,6 @@ 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.
@@ -458,70 +439,6 @@ 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,
@@ -618,29 +535,6 @@ 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])
@@ -694,7 +588,6 @@ class VideoGenerator:
         self,
         output_path: str,
         prompt: str,
-        reserved_paths: set[str] | None = None,
     ) -> str:
         """Build a unique, sanitized output file path.
 
@@ -709,9 +602,6 @@ 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"
 
@@ -756,416 +646,15 @@ 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 _is_taken(new_output_path):
+        while os.path.exists(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,14 +169,6 @@ 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
@@ -454,37 +446,6 @@ 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(
@@ -792,13 +753,6 @@ 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/performance/hf_store.py

@@ -310,8 +310,12 @@ 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,6 +36,7 @@ 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
@@ -57,9 +58,7 @@ 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):
@@ -80,25 +79,15 @@ def resolve_metric_policies(
                 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,
@@ -118,10 +107,7 @@ 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/pipelines/stages/base.py

@@ -34,6 +34,7 @@ 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:
         """
@@ -155,6 +156,9 @@ 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,6 +28,7 @@ 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,6 +51,7 @@ 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__()
@@ -234,6 +235,7 @@ 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
@@ -1186,6 +1188,7 @@ 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,11 +35,7 @@ class InputValidationStage(PipelineStage):
         num_videos_per_prompt = batch.num_videos_per_prompt
 
         assert seed is not None
-        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)]
+        seeds = [seed + i for i in range(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,6 +24,7 @@ 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:
         """
@@ -65,20 +66,12 @@ 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)))
-        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,
-            )
+        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
 
@@ -90,24 +83,13 @@ class TextEncodingStage(PipelineStage):
 
         # Encode negative prompt if CFG is enabled
         if batch.do_classifier_free_guidance:
-            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,
-                )
+            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,
+            )
             if self._last_audio_embeds is not None:
                 batch.extra["ltx2_audio_negative_embeds"] = self._last_audio_embeds
 
@@ -120,81 +102,6 @@ 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()
@@ -329,9 +236,6 @@ 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)
@@ -447,6 +351,7 @@ 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/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 BatchingConfig, CompileConfig, GeneratorConfig
+from fastvideo.api.schema import CompileConfig, GeneratorConfig
 
 
 class TestLegacyTorchCompileKwargsTranslation:
@@ -17,46 +17,6 @@ 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",
@@ -117,13 +77,6 @@ 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",
@@ -200,48 +153,6 @@ 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
 # -------------------------------------------------------------------
@@ -255,13 +166,6 @@ 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,13 +130,6 @@ 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

@@ -1,198 +0,0 @@
-# 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

@@ -1,89 +0,0 @@
-# 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

@@ -1,57 +0,0 @@
-# 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

@@ -0,0 +1,84 @@
+# 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,15 +1,10 @@
 """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 (
@@ -26,172 +21,6 @@ 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,7 +3,6 @@ from types import SimpleNamespace
 import warnings
 
 import pytest
-import torch
 
 from fastvideo.api import (
     GenerationRequest,
@@ -14,10 +13,8 @@ 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:
@@ -40,25 +37,6 @@ 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 = {}
 
@@ -173,117 +151,6 @@ 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,6 +187,8 @@ def run_transformer_tests():
 
 
 @app.function(gpu="L40S:4",
+              cpu=8.0,
+              memory=32768,
               image=image,
               timeout=900,
               secrets=[
@@ -201,6 +203,8 @@ def run_training_tests():
 
 
 @app.function(gpu="L40S:2",
+              cpu=8.0,
+              memory=32768,
               image=image,
               timeout=900,
               secrets=[
@@ -276,7 +280,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/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"
+        "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"
     )
 
 
@@ -318,6 +322,8 @@ def run_dreamverse_app_tests():
 
 
 @app.function(gpu="L40S:1",
+              cpu=8.0,
+              memory=32768,
               image=image,
               timeout=1800,
               secrets=[
@@ -400,6 +406,8 @@ def run_lora_extraction_tests():
 
 
 @app.function(gpu="L40S:2",
+              cpu=8.0,
+              memory=32768,
               image=image,
               timeout=1800,
               secrets=[

fastvideo/tests/ops/quantization/test_absmax_fp8.py

@@ -34,6 +34,9 @@ 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"):
@@ -43,7 +46,7 @@ class TestAbsMaxFP8LinearMethod(unittest.TestCase):
                 output_partition_sizes=[3],
                 input_size=2,
                 output_size=3,
-                params_dtype=torch.float32,
+                params_dtype=torch.float64,
             )
 
     def test_absmax_fp8_parameter_weight_loader(self):

fastvideo/tests/performance/test_inference_performance.py

@@ -179,7 +179,11 @@ 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)
-        metric_key = STAGE_METRIC_MAP.get(stage_class)
+        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)
         if metric_key is None:
             logger.debug("Unmapped stage '%s' class '%s' (%.3fs)",
                          stage_name,

fastvideo/tests/performance/test_inference_performance_component_times.py

@@ -1,13 +1,20 @@
 # 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,
@@ -45,6 +52,35 @@ 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/stages/test_input_validation_batching.py

@@ -1,40 +0,0 @@
-# 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,13 +13,7 @@ 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({
@@ -165,66 +159,6 @@ 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/training/test_trackers.py

@@ -0,0 +1,225 @@
+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

@@ -1,104 +0,0 @@
-# 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

@@ -1,133 +0,0 @@
-# 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

@@ -1,82 +0,0 @@
-# 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

@@ -1,113 +0,0 @@
-# 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,6 +404,7 @@ class ValidationCallback(Callback):
                         art = self.tracker.video(
                             fname,
                             caption=cap,
+                            fps=sp.fps,
                         )
                         if art is not None:
                             video_logs.append(art)

fastvideo/training/distillation_pipeline.py

@@ -1180,7 +1180,7 @@ class DistillationPipeline(TrainingPipeline):
 
                     artifacts = []
                     for filename, caption in zip(video_filenames, all_captions, strict=True):
-                        video_artifact = self.tracker.video(filename, caption=caption)
+                        video_artifact = self.tracker.video(filename, caption=caption, fps=sampling_param.fps)
                         if video_artifact is not None:
                             artifacts.append(video_artifact)
                     if artifacts:

fastvideo/training/trackers.py

@@ -7,22 +7,28 @@ interface that can be used across all FastVideo training pipelines.
 
 from __future__ import annotations
 
+from collections.abc import Iterable, Iterator
 import contextlib
 import copy
+import math
 import os
 import pathlib
+import tempfile
 import time
 from dataclasses import dataclass
 from enum import Enum
 from typing import Any
-from collections.abc import Iterable, Iterator
 
+import numpy as np
 import torch
 
 from fastvideo.logger import init_logger
 
 logger = init_logger(__name__)
 
+_DEFAULT_VIDEO_FPS = 16
+_MISSING_ARTIFACT = object()
+
 
 def _sanitize_wandb_config(value: Any) -> Any:
     """Best-effort conversion of nested config objects to W&B-safe values."""
@@ -56,6 +62,140 @@ def _sanitize_wandb_config(value: Any) -> Any:
     return repr(value)
 
 
+def _prepare_video_array(data: Any) -> np.ndarray:
+    """Convert W&B-style TCHW/BTCHW video data into GIF-ready frames."""
+    if isinstance(data, torch.Tensor):
+        data = data.detach().cpu().numpy()
+
+    video = np.asarray(data)
+    if video.ndim == 4:
+        video = video.reshape(1, *video.shape)
+    elif video.ndim != 5:
+        raise ValueError("Video data must have shape [T, C, H, W] or [B, T, C, H, W]")
+
+    batch_size, num_frames, channels, height, width = video.shape
+    if batch_size == 0 or num_frames == 0:
+        raise ValueError("Video data must contain at least one batch item and one frame")
+    if channels not in (1, 3, 4):
+        raise ValueError(f"Video data must have 1, 3, or 4 channels; got {channels}")
+    if video.dtype != np.uint8:
+        logger.warning("Converting video data to uint8 for SwanLab")
+        video = video.astype(np.uint8)
+
+    # Match wandb.Video's batch tiling so the same input has a familiar layout
+    # in either tracker.
+    if batch_size & (batch_size - 1):
+        padded_batch_size = 1 << batch_size.bit_length()
+        padding = np.zeros(
+            (padded_batch_size - batch_size, num_frames, channels, height, width),
+            dtype=video.dtype,
+        )
+        video = np.concatenate((video, padding), axis=0)
+
+    num_rows = 1 << ((batch_size.bit_length() - 1) // 2)
+    num_columns = video.shape[0] // num_rows
+    video = video.reshape(num_rows, num_columns, num_frames, channels, height, width)
+    video = np.transpose(video, axes=(2, 0, 4, 1, 5, 3))
+    video = video.reshape(num_frames, num_rows * height, num_columns * width, channels)
+    if channels == 1:
+        video = video[..., 0]
+    return np.ascontiguousarray(video)
+
+
+def _read_video_file(file_path: str) -> tuple[list[np.ndarray], float | None]:
+    """Decode a video file and return its frames and encoded frame rate."""
+    import imageio.v2 as imageio
+
+    reader = imageio.get_reader(file_path)
+    try:
+        metadata = reader.get_meta_data() or {}
+        frames = [np.asarray(frame) for frame in reader]
+    finally:
+        reader.close()
+
+    if not frames:
+        raise ValueError(f"Video file contains no frames: {file_path}")
+
+    raw_source_fps = metadata.get("fps")
+    try:
+        source_fps = float(str(raw_source_fps))
+    except ValueError:
+        source_fps = None
+    if source_fps is not None and (not math.isfinite(source_fps) or source_fps <= 0):
+        source_fps = None
+    return frames, source_fps
+
+
+def _coerce_video_fps(fps: int | float | None) -> float:
+    if fps is None:
+        return _DEFAULT_VIDEO_FPS
+    value = float(fps)
+    if not math.isfinite(value) or value <= 0:
+        raise ValueError(f"Video fps must be a positive finite number; got {fps!r}")
+    return value
+
+
+def _write_gif(file_path: str, frames: Any, fps: float) -> None:
+    """Encode frames as an animated GIF for SwanLab's video API."""
+    from PIL import Image
+
+    images = []
+    for frame in frames:
+        array = np.asarray(frame)
+        if array.dtype != np.uint8:
+            array = array.astype(np.uint8)
+        if array.ndim == 3 and array.shape[-1] == 1:
+            array = array[..., 0]
+        if array.ndim not in (2, 3) or (array.ndim == 3 and array.shape[-1] not in (3, 4)):
+            raise ValueError(f"GIF frames must be grayscale, RGB, or RGBA; got shape {array.shape}")
+        images.append(Image.fromarray(array))
+
+    if not images:
+        raise ValueError("Video data must contain at least one frame")
+
+    duration_ms = max(1, round(1000 / fps))
+    images[0].save(
+        file_path,
+        format="GIF",
+        save_all=True,
+        append_images=images[1:],
+        duration=duration_ms,
+        loop=0,
+    )
+
+
+@dataclass(frozen=True)
+class _SequentialArtifact:
+    """Backend-specific versions of one artifact created by a tracker group."""
+
+    values: tuple[Any | None, ...]
+
+
+def _select_tracker_artifact(value: Any, tracker_index: int) -> Any:
+    """Resolve nested sequential artifacts for one child tracker."""
+    if isinstance(value, _SequentialArtifact):
+        selected = value.values[tracker_index]
+        return _MISSING_ARTIFACT if selected is None else selected
+    if isinstance(value, dict):
+        selected_dict = {}
+        for key, item in value.items():
+            selected = _select_tracker_artifact(item, tracker_index)
+            if selected is not _MISSING_ARTIFACT:
+                selected_dict[key] = selected
+        return selected_dict if selected_dict else _MISSING_ARTIFACT
+    if isinstance(value, list):
+        selected_list = [
+            selected for item in value
+            if (selected := _select_tracker_artifact(item, tracker_index)) is not _MISSING_ARTIFACT
+        ]
+        return selected_list if selected_list else _MISSING_ARTIFACT
+    if isinstance(value, tuple):
+        selected_tuple = tuple(selected for item in value
+                               if (selected := _select_tracker_artifact(item, tracker_index)) is not _MISSING_ARTIFACT)
+        return selected_tuple if selected_tuple else _MISSING_ARTIFACT
+    return value
+
+
 @dataclass
 class Timer:
     """Simple timer utility used by the trackers."""
@@ -222,8 +362,7 @@ class WandbTracker(BaseTracker):
         kwargs: dict[str, Any] = {}
         if caption is not None:
             kwargs["caption"] = caption
-        if fps is not None:
-            kwargs["fps"] = fps
+        kwargs["fps"] = fps if fps is not None else _DEFAULT_VIDEO_FPS
         if format is not None:
             kwargs["format"] = format
         else:
@@ -254,8 +393,10 @@ class SequentialTracker(BaseTracker):
         self._timed_metrics = {}
 
     def log_artifacts(self, artifacts: dict[str, Any], step: int) -> None:
-        for tracker in self._trackers:
-            tracker.log_artifacts(artifacts, step)
+        for tracker_index, tracker in enumerate(self._trackers):
+            tracker_artifacts = _select_tracker_artifact(artifacts, tracker_index)
+            if tracker_artifacts is not _MISSING_ARTIFACT:
+                tracker.log_artifacts(tracker_artifacts, step)
         self._timed_metrics = {}
 
     def log_file(
@@ -278,16 +419,91 @@ class SequentialTracker(BaseTracker):
         fps: int | None = None,
         format: str | None = None,
     ) -> Any | None:
-        for tracker in self._trackers:
-            video = tracker.video(data, caption=caption, fps=fps, format=format)
-            if video is not None:
-                return video
-        return None
+        videos = tuple(tracker.video(data, caption=caption, fps=fps, format=format) for tracker in self._trackers)
+        if all(video is None for video in videos):
+            return None
+        return _SequentialArtifact(videos)
+
+
+class SwanlabTracker(BaseTracker):
+    """Tracker implementation for SwanLab."""
+
+    def __init__(
+        self,
+        experiment_name: str,
+        log_dir: str,
+        *,
+        config: dict[str, Any] | None = None,
+        run_name: str | None = None,
+    ) -> None:
+        super().__init__()
+
+        try:
+            import swanlab
+        except ModuleNotFoundError as error:
+            if error.name != "swanlab":
+                raise
+            raise ModuleNotFoundError("SwanLab tracking requires the optional 'swanlab' dependency. "
+                                      "Install it with `uv pip install 'fastvideo[swanlab]'` (or "
+                                      "`uv pip install -e '.[swanlab]'` from a source checkout).") from error
+
+        pathlib.Path(log_dir).mkdir(parents=True, exist_ok=True)
+
+        self._swanlab = swanlab
+        self._run = swanlab.init(
+            project=experiment_name,
+            experiment_name=run_name,
+            config=(_sanitize_wandb_config(config) if config is not None else None),
+            logdir=log_dir,
+        )
+        logger.info("Initialized SwanLab tracker")
+
+    def log(self, metrics: dict[str, Any], step: int) -> None:
+        metrics = {**self._timed_metrics, **metrics}
+        if metrics:
+            self._swanlab.log(metrics, step=step)
+        self._timed_metrics = {}
+
+    def finish(self) -> None:
+        self._swanlab.finish()
+
+    def video(
+        self,
+        data: Any,
+        *,
+        caption: str | None = None,
+        fps: int | None = None,
+        format: str | None = None,
+    ) -> Any:
+        """Create a SwanLab GIF artifact from a file or W&B-style array."""
+        del format  # SwanLab currently supports GIF artifacts only.
+
+        if isinstance(data, str | os.PathLike):
+            source_path = os.fspath(data)
+            if pathlib.Path(source_path).suffix.lower() == ".gif":
+                return self._swanlab.Video(source_path, caption=caption)
+            frames, source_fps = _read_video_file(source_path)
+            if fps is not None:
+                video_fps = _coerce_video_fps(fps)
+            else:
+                video_fps = source_fps if source_fps is not None else _coerce_video_fps(None)
+        else:
+            frames = _prepare_video_array(data)
+            video_fps = _coerce_video_fps(fps)
+
+        file_descriptor, gif_path = tempfile.mkstemp(suffix=".gif")
+        os.close(file_descriptor)
+        try:
+            _write_gif(gif_path, frames, video_fps)
+            return self._swanlab.Video(gif_path, caption=caption)
+        finally:
+            pathlib.Path(gif_path).unlink(missing_ok=True)
 
 
 class Trackers(str, Enum):
     NONE = "none"
     WANDB = "wandb"
+    SWANLAB = "swanlab"
 
 
 SUPPORTED_TRACKERS = {tracker.value for tracker in Trackers}
@@ -324,6 +540,14 @@ def initialize_trackers(
                     config=config,
                     run_name=run_name,
                 ))
+        elif tracker_name == Trackers.SWANLAB.value:
+            tracker_instances.append(
+                SwanlabTracker(
+                    experiment_name,
+                    os.path.abspath(log_dir),
+                    config=config,
+                    run_name=run_name,
+                ))
 
     if not tracker_instances:
         return DummyTracker()

fastvideo/training/training_pipeline.py

@@ -831,7 +831,7 @@ class TrainingPipeline(LoRAPipeline, ABC):
 
                 artifacts = []
                 for filename, caption in zip(video_filenames, all_captions, strict=True):
-                    video_artifact = self.tracker.video(filename, caption=caption)
+                    video_artifact = self.tracker.video(filename, caption=caption, fps=sampling_param.fps)
                     if video_artifact is not None:
                         artifacts.append(video_artifact)
                 if artifacts:

mkdocs.yml

@@ -169,6 +169,7 @@ nav:
   - Training:
     - Overview: training/overview.md
     - Training Infrastructure: training/train_infra.md
+    - Experiment Tracking: training/trackers.md
     - Data Preprocessing: training/data_preprocess.md
     - Fine-tuning: training/finetune.md
     - Examples:

pyproject.toml

@@ -114,12 +114,15 @@ fastvideo-kernel = [
 imagebind = { git = "https://github.com/facebookresearch/ImageBind.git", rev = "53680b02d7e37b19b124fa37bae4b6c98c38f5be" }
 # FA4 cute, pinned to a cutlass-4.5-compatible revision. torch.compile support
 # comes from FastVideo's own custom_op wrappers.
-flash-attn-4 = { git = "https://github.com/Dao-AILab/flash-attention.git", rev = "940cd9680f3315f2f06b43ab5bea2c2cf2d96806", subdirectory = "flash_attn/cute" }
+flash-attn-4 = { git = "https://github.com/Dao-AILab/flash-attention.git", rev = "82d6441eec5d4dfec120153db2c0145ae855a083", subdirectory = "flash_attn/cute" }
 
 [project.optional-dependencies]
 
 # flash-attn: uv pip install flash-attn==2.8.1 --no-cache-dir --no-build-isolation
 
+swanlab = [
+    "swanlab>=0.6.7",
+]
 
 lint = [
     "pre-commit==4.0.1",

tests/local_tests/test_nvfp4_fa4.py

@@ -10,15 +10,6 @@ import pytest
 
 os.environ["CUTE_DSL_ENABLE_TVM_FFI"] = "1"
 
-# Warm up flashinfer JIT before any FA4 imports
-from flashinfer.quantization import nvfp4_quantize, SfLayout
-_warmup = nvfp4_quantize(
-    torch.randn(4, 128, device="cuda", dtype=torch.bfloat16),
-    torch.ones(1, device="cuda", dtype=torch.float32),
-    sfLayout=SfLayout.layout_128x4, do_shuffle=False,
-)
-del _warmup
-
 
 def _make_impl(nheads, headdim, nvfp4=False):
     from fastvideo.attention.backends.flash_attn import FlashAttentionImpl
@@ -45,7 +36,8 @@ def _cuda_timer(fn, warmup=5, iters=20):
 
 
 @pytest.mark.skipif(
-    torch.cuda.get_device_capability() not in [(10, 0), (10, 3)],
+    not torch.cuda.is_available()
+    or torch.cuda.get_device_capability() not in [(10, 0), (10, 3)],
     reason="Requires Blackwell GPU (sm100a or sm103a)"
 )
 class TestNVFP4FA4:
@@ -55,6 +47,17 @@ class TestNVFP4FA4:
     MODEL_HEADDIM = 128
     MODEL_SEQLEN = 32760  # 480x832 video, 81 frames
 
+    @pytest.fixture(scope="class", autouse=True)
+    def _flashinfer_warmup(self):
+        # Warm up flashinfer JIT before any FA4 imports
+        from flashinfer.quantization import SfLayout, nvfp4_quantize
+        _warmup = nvfp4_quantize(
+            torch.randn(4, 128, device="cuda", dtype=torch.bfloat16),
+            torch.ones(1, device="cuda", dtype=torch.float32),
+            sfLayout=SfLayout.layout_128x4, do_shuffle=False,
+        )
+        del _warmup
+
     def test_quantize_fn_shapes(self):
         """_nvfp4_quantize_for_fa4 produces correct shapes and strides."""
         from fastvideo.attention.backends.flash_attn import _nvfp4_quantize_for_fa4