[bugfix]: classify supports_dynamic_batching in schema parity inventory

Typed generate(request) prompt lists previously always expanded to
sequential single-request calls, so dynamic batching only served the
deprecated legacy-kwargs surface. When batching is enabled
(batching_mode='dynamic', batching_max_size>1) a typed prompt list now
expands into the same _GenerationWorkItem path the legacy prompt-file
and server scheduler routes use: shared FastVideoArgs (merge identity),
reserved output paths, compatibility-gated merging, and typed
GenerationResult wrapping with prompt_index.

Default (batching disabled) typed behavior is byte-identical; the new
route is only taken under the existing opt-in flags. Adds one CPU
regression test asserting a typed prompt list merges into a single
forward. Batching-touched suite: 140 -> 141 passed, zero skips.

.buildkite/pipeline.yml

@@ -436,7 +436,7 @@ steps:
                 - "pyproject.toml"
                 - "docker/Dockerfile"
               config:
-                command: "timeout 25m .buildkite/scripts/pr_test.sh"
+                command: "timeout 15m .buildkite/scripts/pr_test.sh"
                 label: ":test_tube: LoRA Training Tests"
                 env:
                   - TEST_TYPE=training_lora

.gitignore

@@ -72,7 +72,8 @@ docs/distillation/examples/
 # Python pickle files
 *.pkl
 
-# Reference videos (negations must come after the catch-all on line below)
+# Reference videos
+!fastvideo/tests/ssim/reference_videos/**/*.mp4
 
 # Static images
 !docs/assets/images/**/*.png
@@ -126,8 +127,6 @@ apps/dreamverse/web/.env.production.local
 .sisyphus/
 openspec/
 fastvideo/tests/ssim/reference_videos/**
-!fastvideo/tests/ssim/reference_videos/**/*.mp4
-!fastvideo/tests/ssim/reference_videos/**/*.png
 
 # Editor logs and local Python version pins (accidentally committed)
 *.nvimlog

docs/contributing/performance_benchmarks.md

@@ -232,25 +232,13 @@ by the harness and is not config-declarable.)
 
 Recipe fingerprinting, hardware/software profile IDs, exact-identity
 comparison, and dashboard cohort grouping land with this change: v2 records
-compare only within their identity cohort and are never compared against v1
-records (or vice versa). Legacy records missing any comparison identity field
-still run and are normalized for reporting, but they skip rolling-baseline
-comparison entirely (`baseline_status: "skipped_missing_identity"`, comparator
-verdict `PASS`, never baseline eligible); only static thresholds gate them.
-Each comparison reports an explicit `comparator_status` on the normalized
-record and in the Markdown summary:
-
-| Status | Meaning | CI |
-|---|---|---|
-| `PASS` | Comparable baseline found with no gated regression, or comparison was skipped for a record without the v2 identity block. | passes |
-| `REGRESSION` | A gated metric regressed past both its percent and absolute floors. | fails |
-| `CALIBRATION_NEEDED` | No comparable baseline exists. Gating is inactive and the record does **not** seed a baseline — seed new cohorts explicitly via the reseed workflow. The record also carries `baseline_status: "initialized_new_cohort"`. | passes |
-| `RECIPE_MISMATCH` | Baseline history exists for the same variant/hardware/software cohort but under a different `recipe_fingerprint`. Represent recipe changes as a new `variant_id`, or reseed. | fails |
-| `INFRA_ERROR` | The comparison itself failed (for example, baseline history could not be loaded). | fails |
-| `QUALITY_BLOCKED` | Reserved for the promoted-baseline workflow; never emitted by the comparator. | n/a |
-
-Metric-specific threshold policies and promoted baselines remain separate
-follow-ups.
+compare only within their identity cohort, and a record that opens a NEW
+cohort is marked `baseline_status: "initialized_new_cohort"` (regression
+gating starts once that cohort accumulates history). Legacy v1 configs still
+run and are normalized for reporting, but their records skip rolling-baseline
+comparison entirely (`baseline_status: "skipped_missing_identity"`, never
+baseline eligible); only static thresholds gate them. Metric-specific
+threshold policies and promoted baselines remain separate follow-ups.
 
 ### Raw record (`results/perf_*.json`)
 
@@ -379,8 +367,6 @@ result, used as the rolling-baseline source of truth.
   "build_id": "<buildkite-build-id>",
   "job_id": "<buildkite-job-id>",
   "quality_metadata": { "quality_status": "canonical" },
-  "baseline_status": "compared",
-  "comparator_status": "PASS",
   "success": true
 }
 ```
@@ -501,10 +487,10 @@ When the rolling-baseline phase runs, it emits:
 2. The pytest test auto-discovers all configs — no test code needed. CI
    picks it up on the next `/test performance` run.
 
-3. The first run with no HF history reports `CALIBRATION_NEEDED`: it passes,
-   but it does not seed the shared baseline. Seed the cohort explicitly with
-   the `reseed-performance-baseline` skill; subsequent scheduled-main runs
-   then compare against it and keep the rolling baseline advancing.
+3. The first persisted main-branch run with no HF history initializes the
+   baseline (passes automatically). Subsequent runs compare against it. Local
+   and pull-request runs with no HF history also pass, but they do not seed the
+   shared baseline.
 
 4. If the benchmark targets a GPU not currently in `thresholds`, either add
    that GPU as a key or rely on the `default` block. Note that `default` is
@@ -524,13 +510,8 @@ When the rolling-baseline phase runs, it emits:
 
 ## Troubleshooting
 
-**`CALIBRATION_NEEDED: no comparable baseline`** — first run for this cohort.
-The run passes but does not seed a baseline; seed it explicitly with the
-`reseed-performance-baseline` skill.
-
-**`RECIPE_MISMATCH`** — the benchmark recipe changed without a new
-`variant_id`. Either bump `variant_id` to open a new cohort, or reseed the
-baseline if the existing variant should adopt the new recipe.
+**"No baseline for ... Initializing"** — first run for this comparison cohort.
+Run will pass and (if persisting) seed the first record.
 
 **Persistent failure right after a torch / kernel / image upgrade** —
 genuine regression *or* baseline drift. Compare the failing normalized record

docs/design/inference_schema_parity_inventory.yaml

@@ -41,6 +41,11 @@ surfaces:
       disable_autocast: generator.engine.disable_autocast
       enable_stage_verification: generator.engine.enable_stage_verification
       prompt_txt: request.inputs.prompt_path
+      batching_mode: generator.engine.batching.mode
+      batching_max_size: generator.engine.batching.max_size
+      batching_delay_ms: generator.engine.batching.delay_ms
+      batching_config: generator.engine.batching.config_path
+      enable_batching_metrics: generator.engine.batching.enable_metrics
       override_text_encoder_safetensors: generator.pipeline.components.text_encoder_weights
       override_text_encoder_quant: generator.engine.quantization.text_encoder_quant
       transformer_quant: generator.engine.quantization.transformer_quant
@@ -128,6 +133,7 @@ surfaces:
       preprocess_text_funcs: "Internal text preprocessing hooks."
       postprocess_text_funcs: "Internal text postprocessing hooks."
       scheduler_step_in_fp32: "Runtime scheduler precision toggle; not part of the public typed inference API."
+      supports_dynamic_batching: "Per-pipeline-family capability flag for dynamic batching opt-in; not part of the public typed inference API."
 
   pipeline_config_extensions:
     preset_owned:
@@ -458,8 +464,6 @@ surfaces:
       guidance_scale: request.sampling.guidance_scale
       guidance_scale_2: request.sampling.guidance_scale_2
       guidance_rescale: request.sampling.guidance_rescale
-      use_embedded_guidance: request.sampling.use_embedded_guidance
-      true_cfg_scale: request.sampling.true_cfg_scale
       boundary_ratio: request.sampling.boundary_ratio
       sigmas: request.sampling.sigmas
       enable_teacache: request.runtime.enable_teacache

docs/distillation/anyflow.md

@@ -1,116 +0,0 @@
-# 🌊 AnyFlow Any-Step Video Distillation
-
-**AnyFlow** ([paper](https://arxiv.org/abs/2605.13724), [project page](https://nvlabs.github.io/AnyFlow/), [official code](https://github.com/NVlabs/AnyFlow), [model weights](https://huggingface.co/collections/nvidia/anyflow)) is an any-step video diffusion framework built on flow maps. A single distilled checkpoint can be evaluated at NFE ∈ {1, 2, 4, 8, 16, 32} without retraining, and quality scales **monotonically** with steps — unlike consistency-based distillation, which often degrades as NFE grows.
-
-The student network ``u_θ(x_t, t, r)`` predicts the *average velocity* from time ``t`` back to time ``r``, so one Euler step is
-
-```
-x_r = x_t - ((t - r) / N) · u_θ(x_t, t, r)
-```
-
-for any ``t > r``.
-
-## 📊 Model Overview
-
-NVIDIA publishes four checkpoints under [`nvidia/anyflow`](https://huggingface.co/collections/nvidia/anyflow):
-
-- `nvidia/AnyFlow-Wan2.1-T2V-1.3B-Diffusers` — bidirectional T2V, Wan2.1 1.3B base
-- `nvidia/AnyFlow-Wan2.1-T2V-14B-Diffusers` — bidirectional T2V, Wan2.1 14B base
-- `nvidia/AnyFlow-FAR-Wan2.1-1.3B-Diffusers` — frame-autoregressive variant, 1.3B
-- `nvidia/AnyFlow-FAR-Wan2.1-14B-Diffusers` — frame-autoregressive variant, 14B
-
-FastVideo currently supports the bidirectional T2V variants for training; the FAR variants can be loaded for inference through the diffusers integration.
-
-## ⚙️ Inference
-
-For inference, load the published checkpoint directly through diffusers; FastVideo's training-side ``WanModel`` config maps the HF AnyFlow ``delta_embedder`` weights onto its internal layout via ``param_names_mapping`` so the same checkpoint can be used as the ``init_from`` for the on-policy YAML below.
-
-## 🧠 Algorithm
-
-Training runs in two stages. Both use the dual-timestep Wan backbone — enabled by ``pipeline.dit_config.r_embedder: true`` in the YAML, which allocates a sibling ``condition_embedder.delta_embedder`` and fuses its embedding with the standard timestep embedding via either an additive or a gated mixer.
-
-### Stage 1 — Pretrain (flow-map central-difference)
-
-Method: ``AnyFlowPretrainMethod`` (``fastvideo/train/methods/distribution_matching/anyflow_pretrain.py``)
-
-For each batch, sample ``(t, r) ∈ [0, 1]`` as ``(max, min)`` of two uniform draws, then:
-
-- a ``diffusion_ratio`` fraction (default 0.5) gets ``r = t`` — recovers plain flow matching;
-- a ``consistency_ratio`` fraction (default 0.25) gets ``r = 0`` — forces consistency to clean data;
-- the remainder is free.
-
-The student forward at ``(t, r)`` is trained against the central-difference target
-
-```
-target = (eps - x_0) - (t - r) · dF/dt
-```
-
-where ``dF/dt`` is estimated from the student's own forward at ``(t ± δ, r)`` with the sample also moved along the flow trajectory by ``v_pred · (δ / N)``. Per-timestep weighting uses ``beta08`` (``w(t) = t · sqrt(1 - t)``, renormalized). A stop-gradient scale-balance keeps the non-diffusion branches' loss magnitude aligned with the diffusion branch.
-
-### Stage 2 — On-policy DMD
-
-Method: ``AnyFlowMethod`` (``fastvideo/train/methods/distribution_matching/anyflow.py``)
-
-Inherits ``DMD2Method``. The student is rolled out for ``student_sample_steps`` Euler-flow steps from pure noise; one randomly-chosen step is gradient-enabled (broadcast from rank 0 so every worker agrees), the rest run under ``torch.no_grad``. With ``use_mean_velocity: true`` (default) the rollout uses ``r = t_next`` at each step, matching AnyFlow's ``WanAnyFlowPipeline.training_rollout``.
-
-The inherited ``_dmd_loss`` (VSD with fake-score critic) consumes the rollout output and the teacher's CFG prediction. The optional pinned ``t_list_override`` lets configs reproduce the paper's hand-tuned 4-step schedule ``[999, 937, 833, 624, 0]``.
-
-## 🚀 Training Scripts
-
-### Stage 1 — pretrain
-
-```bash
-bash examples/train/run.sh \
-    examples/train/configs/distribution_matching/wan/anyflow_pretrain_t2v.yaml
-```
-
-**Key configuration** (in ``examples/train/configs/distribution_matching/wan/anyflow_pretrain_t2v.yaml``):
-
-- Global batch size: 32 (8 GPUs × 4 per-GPU)
-- Learning rate: 5e-5
-- Flow shift: 5.0
-- ``diffusion_ratio`` / ``consistency_ratio``: 0.5 / 0.25
-- ``epsilon`` (finite-difference step): 5 (absolute train-timestep units)
-- ``weight_type``: ``beta08``
-- ``fuse_guidance_scale``: 3.0
-- Training steps: 6000
-
-### Stage 2 — on-policy
-
-```bash
-bash examples/train/run.sh \
-    examples/train/configs/distribution_matching/wan/anyflow_onpolicy_t2v.yaml \
-    --models.student.init_from outputs/wan2.1_anyflow_pretrain/checkpoint-final
-```
-
-(Or point ``models.student.init_from`` directly at ``nvidia/AnyFlow-Wan2.1-T2V-1.3B-Diffusers`` to bootstrap from the paper weights and skip Stage 1.)
-
-**Key configuration**:
-
-- Global batch size: 8 (8 GPUs × 1 per-GPU)
-- Learning rate: 2e-6
-- Flow shift: 5.0
-- ``student_sample_steps``: 4
-- ``t_list_override``: ``[999, 937, 833, 624, 0]``
-- ``use_mean_velocity``: ``true`` (i.e. ``r = t_next`` during rollout)
-- ``real_score_guidance_scale``: 3.0
-- ``generator_update_interval``: 5 (DMD2 alternation)
-- Training steps: 4000
-
-## 🔌 Loading published AnyFlow checkpoints
-
-The HF AnyFlow checkpoints expose ``condition_embedder.delta_embedder.*`` weights that FastVideo internally maps onto its ``condition_embedder.delta_embedder.mlp.*`` layout. This rename happens automatically through the regex in ``WanVideoArchConfig.param_names_mapping`` — no separate adapter is needed. The same regex is a no-op on plain Wan checkpoints (which don't contain any ``delta_embedder`` keys).
-
-Set the YAML's ``pipeline.dit_config.r_embedder: true`` to allocate the ``delta_embedder`` module on the FastVideo side; when initializing from a plain Wan checkpoint the delta weights are deep-copied from ``time_embedder`` (matching AnyFlow's ``setup_flowmap_model()`` behavior).
-
-## 🧭 Note on ``fuse_guidance_scale``
-
-Stage 1 optionally fuses classifier-free guidance into the training target so the resulting checkpoint can be sampled at ``guidance_scale=1.0`` (no extra forward pass at inference time). The transformation is
-
-```
-noise_pred ← (noise_pred - (1 - g) · noise_pred_uncond) / g
-```
-
-with ``g = fuse_guidance_scale``. The negative prompt embedding comes from ``WanModel``'s ``ensure_negative_conditioning()`` — i.e. the dataset's configured ``sampling_param.negative_prompt``. Setting ``fuse_guidance_scale: 1.0`` skips the extra unconditional forward entirely.
-
-The on-policy stage's ``real_score_guidance_scale`` (inherited from DMD2) follows the same parameterization conventions documented in [``dmd.md``](dmd.md#-note-on-real_score_guidance_scale).

examples/inference/basic/basic_flux_dev.py

@@ -1,140 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-
-from __future__ import annotations
-
-import argparse
-import contextlib
-import os
-import re
-
-DEFAULT_PROMPTS = [
-    "a photo of a cat",
-    (
-        "a cinematic photo of a red panda wearing a tiny backpack, standing on a "
-        "rainy neon-lit street at night, shallow depth of field, sharp focus, "
-        "35mm, bokeh"
-    ),
-]
-
-
-def _safe_filename(text: str, max_len: int = 100) -> str:
-    """Make a stable, filesystem-friendly filename base."""
-    s = text[:max_len].strip()
-    s = s.replace(os.sep, "_")
-    if os.altsep:
-        s = s.replace(os.altsep, "_")
-    s = re.sub(r"\s+", " ", s)
-    s = re.sub(r"[^A-Za-z0-9 .,_-]", "_", s)
-    s = s.strip(" .")
-    return s or "prompt"
-
-
-def _remove_existing_outputs(out_dir: str, filename_base: str) -> None:
-    """Delete prior outputs so reruns do not get _1, _2 suffixes."""
-    if not os.path.isdir(out_dir):
-        return
-
-    pattern = re.compile(rf"^{re.escape(filename_base)}(_\d+)?\.(mp4|png)$")
-    for fn in os.listdir(out_dir):
-        if pattern.match(fn):
-            with contextlib.suppress(FileNotFoundError):
-                os.remove(os.path.join(out_dir, fn))
-
-
-def parse_args() -> argparse.Namespace:
-    p = argparse.ArgumentParser(
-        description="Run FLUX.1-dev text-to-image with FastVideo VideoGenerator.",
-    )
-    p.add_argument(
-        "--model-path",
-        default="official_weights/FLUX.1-dev",
-        help="Local Diffusers checkpoint dir or HF repo id.",
-    )
-    p.add_argument(
-        "--out-dir",
-        "--outdir",
-        default="outputs/flux_dev/samples",
-        help="Directory for saved PNG outputs.",
-    )
-    p.add_argument(
-        "--prompt",
-        action="append",
-        default=None,
-        help="Prompt. Repeat for multiple images.",
-    )
-    p.add_argument(
-        "--backend",
-        default=None,
-        help="Set FASTVIDEO_ATTENTION_BACKEND (e.g. TORCH_SDPA).",
-    )
-    p.add_argument("--seed", type=int, default=42, help="Base seed; each prompt uses seed + index.")
-    p.add_argument("--height", type=int, default=1024, help="Output height.")
-    p.add_argument("--width", type=int, default=1024, help="Output width.")
-    p.add_argument("--steps", type=int, default=28, help="Number of inference steps.")
-    p.add_argument("--guidance", type=float, default=3.5, help="Guidance scale.")
-    p.add_argument("--num-gpus", type=int, default=1, help="GPU count.")
-    return p.parse_args()
-
-
-def main() -> None:
-    args = parse_args()
-    prompts: list[str] = args.prompt if args.prompt else DEFAULT_PROMPTS
-
-    if args.backend:
-        os.environ["FASTVIDEO_ATTENTION_BACKEND"] = args.backend
-
-    from fastvideo import VideoGenerator
-
-    os.makedirs(args.out_dir, exist_ok=True)
-
-    init_kwargs = {
-        "num_gpus": args.num_gpus,
-        "workload_type": "t2i",
-        "sp_size": 1,
-        "tp_size": 1,
-        "dit_cpu_offload": False,
-        "dit_layerwise_offload": False,
-        "text_encoder_cpu_offload": False,
-        "vae_cpu_offload": False,
-        "image_encoder_cpu_offload": False,
-        "pin_cpu_memory": False,
-        "use_fsdp_inference": False,
-    }
-
-    generator = VideoGenerator.from_pretrained(
-        model_path=args.model_path,
-        **init_kwargs,
-    )
-    try:
-        for i, prompt in enumerate(prompts):
-            seed = args.seed + i
-            filename_base = (
-                f"flux_dev_{i:02d}_seed{seed}_{_safe_filename(prompt, max_len=80)}"
-            )
-            _remove_existing_outputs(args.out_dir, filename_base)
-            output_path = os.path.join(args.out_dir, f"{filename_base}.png")
-            print(f"[flux] prompt_idx={i} seed={seed} output_path={output_path}")
-
-            generation_kwargs = {
-                "output_path": output_path,
-                "height": args.height,
-                "width": args.width,
-                "num_frames": 1,
-                "fps": 1,
-                "num_inference_steps": args.steps,
-                "guidance_scale": args.guidance,
-                "use_embedded_guidance": True,
-                "true_cfg_scale": 1.0,
-                "seed": seed,
-                "save_video": True,
-            }
-
-            generator.generate_video(prompt, **generation_kwargs)
-
-        print(f"[flux] done. outputs written to: {args.out_dir}")
-    finally:
-        generator.shutdown()
-
-
-if __name__ == "__main__":
-    main()

examples/inference/basic/basic_glm_image.py

@@ -1,107 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Run GLM-Image text-to-image generation through FastVideo.
-
-User story:
-    "I have the HF `zai-org/GLM-Image` checkpoint and want a minimal
-    text-to-image generation command, saved as a PNG."
-"""
-import argparse
-from pathlib import Path
-
-from PIL import Image
-
-from fastvideo import VideoGenerator
-from fastvideo.api import (
-    EngineConfig,
-    GenerationRequest,
-    GeneratorConfig,
-    OutputConfig,
-    ParallelismConfig,
-    PipelineSelection,
-    SamplingConfig,
-)
-
-
-def parse_args() -> argparse.Namespace:
-    parser = argparse.ArgumentParser(description="Run GLM-Image text-to-image generation.")
-    parser.add_argument(
-        "--model-path",
-        default="zai-org/GLM-Image",
-        help="HF id or local diffusers-format GLM-Image weights directory.",
-    )
-    parser.add_argument(
-        "--output",
-        default="image_output/landscape.png",
-        help="Output PNG path.",
-    )
-    parser.add_argument(
-        "--prompt",
-        default=("A beautiful landscape photography with rolling hills, "
-                 "a winding river, and a vibrant sunset in the background. "
-                 "Warm golden light, photorealistic style."),
-        help="Text prompt.",
-    )
-    parser.add_argument("--height", type=int, default=1024)
-    parser.add_argument("--width", type=int, default=1024)
-    parser.add_argument("--steps", type=int, default=50)
-    parser.add_argument("--guidance-scale", type=float, default=1.5)
-    parser.add_argument("--seed", type=int, default=1024)
-    parser.add_argument("--num-gpus", type=int, default=1)
-    parser.add_argument("--tp-size", type=int, default=None)
-    parser.add_argument("--sp-size", type=int, default=None)
-    return parser.parse_args()
-
-
-def main() -> None:
-    args = parse_args()
-
-    output = Path(args.output)
-    output.parent.mkdir(parents=True, exist_ok=True)
-    tp_size = args.tp_size if args.tp_size is not None else (args.num_gpus if args.num_gpus > 1 else 1)
-    sp_size = args.sp_size if args.sp_size is not None else (1 if args.num_gpus > 1 else args.num_gpus)
-
-    # GLM-Image needs trust_remote_code for its AR encoder; offload and the
-    # pipeline class come from the model's registered defaults — don't override.
-    generator_config = GeneratorConfig(
-        model_path=args.model_path,
-        trust_remote_code=True,
-        engine=EngineConfig(
-            num_gpus=args.num_gpus,
-            parallelism=ParallelismConfig(tp_size=tp_size, sp_size=sp_size),
-        ),
-        pipeline=PipelineSelection(workload_type="t2i"),
-    )
-
-    generator = VideoGenerator.from_config(generator_config)
-    try:
-        request = GenerationRequest(
-            prompt=args.prompt,
-            sampling=SamplingConfig(
-                height=args.height,
-                width=args.width,
-                num_frames=1,
-                fps=1,
-                num_inference_steps=args.steps,
-                guidance_scale=args.guidance_scale,
-                seed=args.seed,
-            ),
-            output=OutputConfig(
-                output_path=str(output.parent),
-                save_video=False,
-                return_frames=True,
-            ),
-        )
-        result = generator.generate(request)
-        if isinstance(result, list):
-            result = result[0]
-
-        frames = result.frames
-        if frames is not None and len(frames):
-            Image.fromarray(frames[0]).save(output)
-            print(f"Saved image to {output}")
-    finally:
-        generator.shutdown()
-
-
-if __name__ == "__main__":
-    main()

examples/inference/basic/edit_glm_image.py

@@ -1,120 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Run GLM-Image image-to-image (edit) generation through FastVideo.
-
-User story:
-    "I have the HF `zai-org/GLM-Image` checkpoint and a condition image, and
-    want a minimal edit command (text + image -> edited image), saved as a PNG."
-
-GLM-Image is a single unified pipeline: passing a condition image switches it
-from text-to-image to the edit path (the condition enters the DiT via a KV-cache
-write pass), so the generator config is identical to `basic_glm_image.py` — the
-`inputs.pil_image` on the request is what selects the edit mode.
-"""
-import argparse
-from pathlib import Path
-
-from PIL import Image
-
-from fastvideo import VideoGenerator
-from fastvideo.api import (
-    EngineConfig,
-    GenerationRequest,
-    GeneratorConfig,
-    InputConfig,
-    OutputConfig,
-    ParallelismConfig,
-    PipelineSelection,
-    SamplingConfig,
-)
-
-
-def parse_args() -> argparse.Namespace:
-    parser = argparse.ArgumentParser(description="Run GLM-Image image-to-image (edit) generation.")
-    parser.add_argument(
-        "--model-path",
-        default="zai-org/GLM-Image",
-        help="HF id or local diffusers-format GLM-Image weights directory.",
-    )
-    parser.add_argument(
-        "--image",
-        default="assets/images/couple.jpg",
-        help="Condition image to edit.",
-    )
-    parser.add_argument(
-        "--output",
-        default="image_output/edited.png",
-        help="Output PNG path.",
-    )
-    parser.add_argument(
-        "--prompt",
-        default="Change the background to a snowy mountain landscape at golden hour.",
-        help="Edit instruction.",
-    )
-    parser.add_argument("--height", type=int, default=1024)
-    parser.add_argument("--width", type=int, default=1024)
-    parser.add_argument("--steps", type=int, default=50)
-    parser.add_argument("--guidance-scale", type=float, default=1.5)
-    parser.add_argument("--seed", type=int, default=1024)
-    parser.add_argument("--num-gpus", type=int, default=1)
-    parser.add_argument("--tp-size", type=int, default=None)
-    parser.add_argument("--sp-size", type=int, default=None)
-    return parser.parse_args()
-
-
-def main() -> None:
-    args = parse_args()
-
-    output = Path(args.output)
-    output.parent.mkdir(parents=True, exist_ok=True)
-    condition = Image.open(args.image).convert("RGB")
-    tp_size = args.tp_size if args.tp_size is not None else (args.num_gpus if args.num_gpus > 1 else 1)
-    sp_size = args.sp_size if args.sp_size is not None else (1 if args.num_gpus > 1 else args.num_gpus)
-
-    # GLM-Image needs trust_remote_code for its AR encoder; offload and the
-    # pipeline class come from the model's registered defaults — don't override.
-    # The pipeline is registered as t2i; passing inputs.pil_image below switches
-    # it to the edit path.
-    generator_config = GeneratorConfig(
-        model_path=args.model_path,
-        trust_remote_code=True,
-        engine=EngineConfig(
-            num_gpus=args.num_gpus,
-            parallelism=ParallelismConfig(tp_size=tp_size, sp_size=sp_size),
-        ),
-        pipeline=PipelineSelection(workload_type="t2i"),
-    )
-
-    generator = VideoGenerator.from_config(generator_config)
-    try:
-        request = GenerationRequest(
-            prompt=args.prompt,
-            inputs=InputConfig(pil_image=condition),
-            sampling=SamplingConfig(
-                height=args.height,
-                width=args.width,
-                num_frames=1,
-                fps=1,
-                num_inference_steps=args.steps,
-                guidance_scale=args.guidance_scale,
-                seed=args.seed,
-            ),
-            output=OutputConfig(
-                output_path=str(output.parent),
-                save_video=False,
-                return_frames=True,
-            ),
-        )
-        result = generator.generate(request)
-        if isinstance(result, list):
-            result = result[0]
-
-        frames = result.frames
-        if frames is not None and len(frames):
-            Image.fromarray(frames[0]).save(output)
-            print(f"Saved image to {output}")
-    finally:
-        generator.shutdown()
-
-
-if __name__ == "__main__":
-    main()

examples/train/configs/distribution_matching/wan/anyflow_onpolicy_t2v.yaml

@@ -1,105 +0,0 @@
-# AnyFlow on-policy DMD — Wan 2.1 T2V 1.3B.
-#
-# Stage 2 of the AnyFlow two-stage recipe. Continues from the pretrain
-# checkpoint; refines the student via DMD2 with a multi-step Euler-flow
-# rollout from pure noise. Teacher provides the real score, critic
-# learns the fake score; both inherited from DMD2Method.
-#
-# Replace <PATH_TO_PRETRAIN_CKPT> with the output of the pretrain stage,
-# or with the NVIDIA-released checkpoint
-# nvidia/AnyFlow-Wan2.1-T2V-1.3B-Diffusers to bootstrap directly from
-# the paper weights (the delta_embedder rename is handled by the
-# param_names_mapping in WanVideoArchConfig).
-
-models:
-  student:
-    _target_: fastvideo.train.models.wan.WanModel
-    init_from: <PATH_TO_PRETRAIN_CKPT>
-    trainable: true
-  teacher:
-    _target_: fastvideo.train.models.wan.WanModel
-    init_from: Wan-AI/Wan2.1-T2V-14B-Diffusers
-    trainable: false
-    disable_custom_init_weights: true
-  critic:
-    _target_: fastvideo.train.models.wan.WanModel
-    init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers
-    trainable: true
-    disable_custom_init_weights: true
-
-method:
-  _target_: fastvideo.train.methods.distribution_matching.anyflow.AnyFlowMethod
-  rollout_mode: simulate
-  generator_update_interval: 5
-  real_score_guidance_scale: 3.0
-  dmd_denoising_steps: [999, 937, 833, 624]
-  warp_denoising_step: false
-
-  # AnyFlow rollout knobs.
-  student_sample_steps: 4
-  use_mean_velocity: true
-  t_list_override: [999.0, 937.0, 833.0, 624.0, 0.0]
-  dmd_score_r_value: 0.0  # DMD scoring conditioning is at r=0 (consistency target).
-
-  # Critic optimizer (DMD2 inherited).
-  fake_score_learning_rate: 8.0e-6
-  fake_score_betas: [0.0, 0.999]
-  fake_score_lr_scheduler: constant
-
-  attn_kind: vsa
-
-training:
-  distributed:
-    num_gpus: 8
-    sp_size: 1
-    tp_size: 1
-    hsdp_replicate_dim: 1
-    hsdp_shard_dim: 8
-
-  data:
-    data_path: data/preprocessed
-    dataloader_num_workers: 4
-    train_batch_size: 1
-    training_cfg_rate: 0.0
-    seed: 1000
-    num_latent_t: 21
-    num_height: 480
-    num_width: 832
-    num_frames: 81
-
-  optimizer:
-    learning_rate: 2.0e-6
-    betas: [0.0, 0.999]
-    weight_decay: 0.01
-    lr_scheduler: constant
-    lr_warmup_steps: 0
-
-  loop:
-    max_train_steps: 4000
-    gradient_accumulation_steps: 1
-
-  checkpoint:
-    output_dir: outputs/wan2.1_anyflow_onpolicy
-    training_state_checkpointing_steps: 500
-    checkpoints_total_limit: 3
-    resume_from_checkpoint: latest
-
-  tracker:
-    project_name: anyflow-wan
-    run_name: wan2.1_t2v_anyflow_onpolicy
-
-  model:
-    enable_gradient_checkpointing_type: full
-
-callbacks:
-  grad_clip:
-    _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback
-    max_grad_norm: 1.0
-
-pipeline:
-  flow_shift: 5.0
-  dit_config:
-    r_embedder: true
-    r_embedder_fusion: gated
-    r_embedder_gate_value: 0.25
-    r_embedder_deltatime_type: r

examples/train/configs/distribution_matching/wan/anyflow_pretrain_t2v.yaml

@@ -1,83 +0,0 @@
-# AnyFlow pretrain (flow-map central-difference) — Wan 2.1 T2V 1.3B.
-#
-# Stage 1 of the AnyFlow two-stage recipe. Trains the dual-timestep
-# u_θ(x_t, t, r) on the central-difference target so the same checkpoint
-# can be sampled at arbitrary NFE in the on-policy stage.
-#
-# Initialize from base Wan 2.1 T2V 1.3B. No teacher or critic at this
-# stage; AnyFlowPretrainMethod owns a single student + one optimizer.
-
-models:
-  student:
-    _target_: fastvideo.train.models.wan.WanModel
-    init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers
-    trainable: true
-
-method:
-  _target_: fastvideo.train.methods.distribution_matching.anyflow_pretrain.AnyFlowPretrainMethod
-  diffusion_ratio: 0.5
-  consistency_ratio: 0.25
-  epsilon: 5            # finite-difference step in absolute train-timestep units
-  weight_type: beta08   # per-timestep loss weight = t * sqrt(1 - t), renormalized
-  fuse_guidance_scale: 3.0
-  # shift is taken from pipeline.flow_shift below.
-
-training:
-  distributed:
-    num_gpus: 8
-    sp_size: 1
-    tp_size: 1
-    hsdp_replicate_dim: 1
-    hsdp_shard_dim: 8
-
-  data:
-    data_path: data/preprocessed
-    dataloader_num_workers: 4
-    train_batch_size: 4
-    training_cfg_rate: 0.0
-    seed: 1000
-    num_latent_t: 21
-    num_height: 480
-    num_width: 832
-    num_frames: 81
-
-  optimizer:
-    learning_rate: 5.0e-5
-    betas: [0.9, 0.999]
-    weight_decay: 0.0
-    lr_scheduler: constant
-    lr_warmup_steps: 0
-
-  loop:
-    max_train_steps: 6000
-    gradient_accumulation_steps: 1
-
-  checkpoint:
-    output_dir: outputs/wan2.1_anyflow_pretrain
-    training_state_checkpointing_steps: 500
-    checkpoints_total_limit: 3
-    resume_from_checkpoint: latest
-
-  tracker:
-    project_name: anyflow-wan
-    run_name: wan2.1_t2v_anyflow_pretrain
-
-  model:
-    enable_gradient_checkpointing_type: full
-
-callbacks:
-  grad_clip:
-    _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback
-    max_grad_norm: 1.0
-
-pipeline:
-  flow_shift: 5.0
-  dit_config:
-    # Enable AnyFlow dual-timestep conditioning. The student loads from
-    # base Wan 2.1 — its checkpoint has no delta_embedder weights, so they
-    # get initialized identically to time_embedder via deep-copy in
-    # WanTimeTextImageEmbedding.__init__.
-    r_embedder: true
-    r_embedder_fusion: gated
-    r_embedder_gate_value: 0.25
-    r_embedder_deltatime_type: r

fastvideo/api/compat.py

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

fastvideo/api/flux.py

@@ -1,27 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-
-from __future__ import annotations
-
-from dataclasses import dataclass
-
-from fastvideo.api.sampling_param import SamplingParam
-
-
-@dataclass
-class FluxSamplingParam(SamplingParam):
-
-    prompt: str | None = "a photo of a cat"
-    negative_prompt: str = ""
-
-    num_videos_per_prompt: int = 1
-    seed: int = 0
-
-    num_frames: int = 1
-    height: int = 1024
-    width: int = 1024
-    fps: int = 1
-
-    num_inference_steps: int = 28
-    guidance_scale: float = 3.5
-    use_embedded_guidance: bool = True
-    true_cfg_scale: float = 1.0

fastvideo/api/sampling_param.py

@@ -90,10 +90,6 @@ class SamplingParam:
     num_inference_steps_sr: int = 50
     guidance_scale: float = 1.0
     guidance_scale_2: float | None = None
-    # Embedded guidance (FLUX): do not treat ``guidance_scale > 1`` as classic CFG.
-    use_embedded_guidance: bool = False
-    # Diffusers-style true CFG for FLUX when > 1 (requires negative prompt encoding).
-    true_cfg_scale: float = 1.0
     guidance_rescale: float = 0.0
     boundary_ratio: float | None = None
     sigmas: list[float] | None = None
@@ -329,18 +325,6 @@ class SamplingParam:
             default=SamplingParam.guidance_rescale,
             help="Guidance rescale factor",
         )
-        parser.add_argument(
-            "--use-embedded-guidance",
-            action="store_true",
-            default=SamplingParam.use_embedded_guidance,
-            help="Use embedded guidance scale (FLUX-style) instead of classic CFG",
-        )
-        parser.add_argument(
-            "--true-cfg-scale",
-            type=float,
-            default=SamplingParam.true_cfg_scale,
-            help="True CFG scale for FLUX when > 1 (requires negative prompt encoding)",
-        )
         parser.add_argument(
             "--boundary-ratio",
             type=float,

fastvideo/api/schema.py

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

fastvideo/attention/backends/sdpa.py

@@ -15,14 +15,8 @@ class SDPABackend(AttentionBackend):
     accept_output_buffer: bool = True
 
     @staticmethod
-    def get_supported_head_sizes() -> list[int] | None:
-        # torch.nn.functional.scaled_dot_product_attention is head-size
-        # agnostic: the math backend handles any size and the fused kernels
-        # fall back internally when they cannot. The previous list
-        # ([32, 64, ..., 256]) was copied from FlashAttentionBackend in the
-        # v1 refactor (#270) and wrongly excluded sizes such as 80 (e.g.
-        # CLIP vision encoders). None means "no restriction".
-        return None
+    def get_supported_head_sizes() -> list[int]:
+        return [32, 64, 96, 128, 160, 192, 224, 256]
 
     @staticmethod
     def get_name() -> str:

fastvideo/batching/__init__.py

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

fastvideo/batching/admission.py

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

fastvideo/batching/signature.py

@@ -0,0 +1,183 @@
+# 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",
+    "latents",
+    "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)
+    # Unrecognized types (tensors, PIL images, ...) cannot be reliably compared
+    # by value — repr() truncates large tensors, so two different payloads can
+    # collide. Treat them as equal only when both requests hold the exact same
+    # object; anything else is conservatively incompatible.
+    return (type(value).__qualname__, id(value))
+
+
+def _is_present(value: Any) -> bool:
+    if value is None:
+        return False
+    if value is False:
+        return False
+    return not (isinstance(value, list | tuple | dict | set) and not value)
+
+
+def _first_scalar(value: Any) -> Any:
+    if isinstance(value, list | tuple):
+        return value[0] if value else None
+    return value
+
+
+def _first_mismatch(
+    base_sig: tuple[tuple[str, Any], ...],
+    candidate_sig: tuple[tuple[str, Any], ...],
+) -> str | None:
+    if len(base_sig) != len(candidate_sig):
+        return "sampling_params"
+    for (name, base_value), (candidate_name, candidate_value) in zip(base_sig, candidate_sig, strict=True):
+        if name != candidate_name:
+            return "sampling_params"
+        if base_value != candidate_value:
+            return f"sampling_params.{name}"
+    return None

fastvideo/configs/models/dits/__init__.py

@@ -1,9 +1,7 @@
 from fastvideo.configs.models.dits.cosmos import CosmosVideoConfig
 from fastvideo.configs.models.dits.cosmos2_5 import Cosmos25VideoConfig
 from fastvideo.configs.models.dits.dreamx_world import DreamXWorldARConfig, DreamXWorldConfig
-from fastvideo.configs.models.dits.flux import FluxDiTConfig
 from fastvideo.configs.models.dits.flux_2 import Flux2Config
-from fastvideo.configs.models.dits.glm_image import GlmImageDiTConfig
 from fastvideo.configs.models.dits.hunyuangamecraft import HunyuanGameCraftConfig
 from fastvideo.configs.models.dits.hunyuanvideo import HunyuanVideoConfig
 from fastvideo.configs.models.dits.hunyuanvideo15 import HunyuanVideo15Config
@@ -17,7 +15,6 @@ from fastvideo.configs.models.dits.kandinsky5 import Kandinsky5VideoConfig
 
 __all__ = [
     "HunyuanVideoConfig", "HunyuanVideo15Config", "HunyuanGameCraftConfig", "WanVideoConfig", "DreamXWorldConfig",
-    "DreamXWorldARConfig", "CosmosVideoConfig", "Cosmos25VideoConfig", "FluxDiTConfig", "Flux2Config",
-    "LongCatVideoConfig", "LTX2VideoConfig", "HYWorldConfig", "Kandinsky5VideoConfig", "MagiHumanVideoConfig",
-    "StableAudioConfig", "GlmImageDiTConfig"
+    "DreamXWorldARConfig", "CosmosVideoConfig", "Cosmos25VideoConfig", "LongCatVideoConfig", "LTX2VideoConfig",
+    "HYWorldConfig", "Kandinsky5VideoConfig", "MagiHumanVideoConfig", "StableAudioConfig", "Flux2Config"
 ]

fastvideo/configs/models/dits/flux.py

@@ -1,27 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-
-from dataclasses import dataclass, field
-
-from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig
-
-
-@dataclass
-class FluxTransformer2DArchConfig(DiTArchConfig):
-
-    patch_size: int = 1
-    in_channels: int = 64
-    out_channels: int | None = None
-    num_layers: int = 19
-    num_single_layers: int = 38
-    attention_head_dim: int = 128
-    num_attention_heads: int = 24
-    joint_attention_dim: int = 4096
-    pooled_projection_dim: int = 768
-    guidance_embeds: bool = True
-    axes_dims_rope: tuple[int, int, int] = (16, 56, 56)
-
-
-@dataclass
-class FluxDiTConfig(DiTConfig):
-    arch_config: DiTArchConfig = field(default_factory=FluxTransformer2DArchConfig)
-    prefix: str = "flux"

fastvideo/configs/models/dits/glm_image.py

@@ -1,61 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-from dataclasses import dataclass, field
-
-from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig
-
-
-def is_blocks(n: str, m) -> bool:
-    return "transformer_blocks" in n and str.isdigit(n.split(".")[-1])
-
-
-@dataclass
-class GlmImageDiTArchConfig(DiTArchConfig):
-
-    _fsdp_shard_conditions: list = field(default_factory=lambda: [is_blocks])
-
-    hidden_size: int = 4096
-    num_attention_heads: int = 32
-    attention_head_dim: int = 128
-    in_channels: int = 16
-    out_channels: int = 16
-    num_layers: int = 30
-
-    text_embed_dim: int = 1472
-    time_embed_dim: int = 512
-    condition_dim: int = 256
-
-    prior_vq_quantizer_codebook_size: int = 16384
-
-    patch_size: int = 2
-
-    max_height: int = 2048
-    max_width: int = 2048
-
-    qk_norm: str = "layer_norm"
-    eps: float = 1e-5
-
-    exclude_lora_layers: list[str] = field(
-        default_factory=lambda: ["image_projector", "glyph_projector", "prior_token_embedding"])
-
-    param_names_mapping: dict = field(
-        default_factory=lambda: {
-            r"^glyph_projector\.net\.0\.proj\.(.*)$": r"glyph_projector.fc_in.\1",
-            r"^glyph_projector\.net\.2\.(.*)$": r"glyph_projector.fc_out.\1",
-            r"^prior_projector\.net\.0\.proj\.(.*)$": r"prior_projector.fc_in.\1",
-            r"^prior_projector\.net\.2\.(.*)$": r"prior_projector.fc_out.\1",
-            r"^transformer_blocks\.(\d+)\.ff\.net\.0\.proj\.(.*)$": r"transformer_blocks.\1.ff.fc_in.\2",
-            r"^transformer_blocks\.(\d+)\.ff\.net\.2\.(.*)$": r"transformer_blocks.\1.ff.fc_out.\2",
-        })
-
-    reverse_param_names_mapping: dict = field(default_factory=dict)
-    lora_param_names_mapping: dict = field(default_factory=dict)
-
-    def __post_init__(self):
-        super().__post_init__()
-        self.num_channels_latents = self.out_channels
-
-
-@dataclass
-class GlmImageDiTConfig(DiTConfig):
-    arch_config: DiTArchConfig = field(default_factory=GlmImageDiTArchConfig)
-    prefix: str = "GlmImage"

fastvideo/configs/models/dits/wanvideo.py

@@ -1,6 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass, field
-from typing import Literal
 
 from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig
 
@@ -20,11 +19,6 @@ class WanVideoArchConfig(DiTArchConfig):
             r"^condition_embedder\.text_embedder\.linear_2\.(.*)$": r"condition_embedder.text_embedder.fc_out.\1",
             r"^condition_embedder\.time_embedder\.linear_1\.(.*)$": r"condition_embedder.time_embedder.mlp.fc_in.\1",
             r"^condition_embedder\.time_embedder\.linear_2\.(.*)$": r"condition_embedder.time_embedder.mlp.fc_out.\1",
-            # AnyFlow dual-timestep checkpoints expose delta_embedder weights with the
-            # same internal layout as time_embedder. The regex is harmless on plain
-            # Wan checkpoints (no delta_embedder keys to match).
-            r"^condition_embedder\.delta_embedder\.linear_1\.(.*)$": r"condition_embedder.delta_embedder.mlp.fc_in.\1",
-            r"^condition_embedder\.delta_embedder\.linear_2\.(.*)$": r"condition_embedder.delta_embedder.mlp.fc_out.\1",
             r"^condition_embedder\.time_proj\.(.*)$": r"condition_embedder.time_modulation.linear.\1",
             r"^condition_embedder\.image_embedder\.ff\.net\.0\.proj\.(.*)$":
             r"condition_embedder.image_embedder.ff.fc_in.\1",
@@ -92,14 +86,6 @@ class WanVideoArchConfig(DiTArchConfig):
     # "relativistic" keeps long rollouts in-distribution; a no-op unless sink_size > 0 and local_attn_size > 0.
     rope_cache_policy: str = "absolute"
 
-    # AnyFlow dual-timestep conditioning. Defaults preserve bit-identity with
-    # the legacy single-timestep forward (no delta_embedder allocated, no
-    # extra computation on the embedder forward path).
-    r_embedder: bool = False
-    r_embedder_fusion: Literal["additive", "gated"] = "additive"
-    r_embedder_gate_value: float = 0.25
-    r_embedder_deltatime_type: Literal["r", "t-r"] = "r"
-
     def __post_init__(self):
         super().__post_init__()
         self.out_channels = self.out_channels or self.in_channels

fastvideo/configs/models/vaes/__init__.py

@@ -2,7 +2,6 @@ from fastvideo.configs.models.vaes.cosmosvae import CosmosVAEConfig
 from fastvideo.configs.models.vaes.cosmos2_5vae import Cosmos25VAEConfig
 from fastvideo.configs.models.vaes.gamecraftvae import GameCraftVAEConfig
 from fastvideo.configs.models.vaes.gen3cvae import Gen3CVAEConfig
-from fastvideo.configs.models.vaes.glm_image import GlmImageVAEConfig
 from fastvideo.configs.models.vaes.hunyuanvae import HunyuanVAEConfig
 from fastvideo.configs.models.vaes.hunyuan15vae import Hunyuan15VAEConfig
 from fastvideo.configs.models.vaes.ltx2vae import LTX2VAEConfig
@@ -22,5 +21,4 @@ __all__ = [
     "OobleckVAEArchConfig",
     "OobleckVAEConfig",
     "Flux2VAEConfig",
-    "GlmImageVAEConfig",
 ]

fastvideo/configs/models/vaes/glm_image.py

@@ -1,95 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-from __future__ import annotations
-
-from dataclasses import dataclass, field
-
-from fastvideo.configs.models.vaes.autoencoder_kl import (AutoencoderKLArchConfig, AutoencoderKLVAEConfig)
-
-_GLM_IMAGE_LATENTS_MEAN: tuple[float, ...] = (
-    -0.2080078125,
-    1.875,
-    -0.470703125,
-    -1.265625,
-    -1.421875,
-    0.77734375,
-    -0.3671875,
-    -0.9453125,
-    0.318359375,
-    0.7734375,
-    -0.1884765625,
-    -0.022216796875,
-    -0.220703125,
-    -1.59375,
-    -0.81640625,
-    -0.255859375,
-)
-_GLM_IMAGE_LATENTS_STD: tuple[float, ...] = (
-    3.0625,
-    2.203125,
-    2.265625,
-    4.84375,
-    2.5,
-    3.9375,
-    2.203125,
-    3.03125,
-    2.1875,
-    2.046875,
-    2.71875,
-    2.390625,
-    2.390625,
-    2.453125,
-    2.25,
-    2.15625,
-)
-
-
-@dataclass
-class GlmImageVAEArchConfig(AutoencoderKLArchConfig):
-    act_fn: str = "silu"
-    block_out_channels: tuple[int, ...] = (128, 512, 1024, 1024)
-    down_block_types: tuple[str, ...] = (
-        "DownEncoderBlock2D",
-        "DownEncoderBlock2D",
-        "DownEncoderBlock2D",
-        "DownEncoderBlock2D",
-    )
-    up_block_types: tuple[str, ...] = (
-        "UpDecoderBlock2D",
-        "UpDecoderBlock2D",
-        "UpDecoderBlock2D",
-        "UpDecoderBlock2D",
-    )
-    force_upcast: bool = True
-    in_channels: int = 3
-    out_channels: int = 3
-    latent_channels: int = 16
-    latents_mean: tuple[float, ...] = _GLM_IMAGE_LATENTS_MEAN
-    latents_std: tuple[float, ...] = _GLM_IMAGE_LATENTS_STD
-    layers_per_block: int = 3
-    mid_block_add_attention: bool = False
-    norm_num_groups: int = 32
-    sample_size: int = 1024
-    scaling_factor: float = 0.18215
-    shift_factor: float | None = None
-    use_quant_conv: bool = False
-    use_post_quant_conv: bool = False
-
-    temporal_compression_ratio: int = 1
-    spatial_compression_ratio: int = 8
-
-
-@dataclass
-class GlmImageVAEConfig(AutoencoderKLVAEConfig):
-    arch_config: GlmImageVAEArchConfig = field(default_factory=GlmImageVAEArchConfig)
-
-    use_tiling: bool = True
-    use_temporal_tiling: bool = False
-    use_parallel_tiling: bool = False
-
-    tile_sample_min_height: int = 512
-    tile_sample_min_width: int = 512
-    tile_sample_stride_height: int = 384
-    tile_sample_stride_width: int = 384
-
-    load_encoder: bool = True
-    load_decoder: bool = True

fastvideo/configs/pipelines/base.py

@@ -75,6 +75,12 @@ class PipelineConfig:
     # DMD parameters
     dmd_denoising_steps: list[int] | None = field(default=None)
 
+    # Whether this pipeline family is verified to merge multiple requests into
+    # one forward pass (dynamic batching). Conservative default: several
+    # pipelines are structurally batch-size-1, so each family must opt in
+    # explicitly after verifying output parity.
+    supports_dynamic_batching: bool = False
+
     # Wan2.2 task modifiers
     ti2v_task: bool = False
     lucy_edit_task: bool = False
@@ -276,6 +282,37 @@ 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 dynamic_batching_supported(self) -> bool:
+        """Whether the loaded pipeline may merge requests via dynamic batching.
+
+        Conservative gate: a pipeline family must opt in explicitly, and DMD /
+        causal denoising variants are always excluded because those denoising
+        stages do not consume per-item generators, so merged requests would
+        share RNG.
+        """
+        return (self.supports_dynamic_batching and self.dmd_denoising_steps is None and not self.is_causal)
+
+    def estimate_request_cost(self, request: Any) -> float:
+        """Estimate relative memory/compute cost for batching admission.
+
+        The default is intentionally simple and model-agnostic: pixel count
+        times frame count. Pipeline subclasses can override this when they have
+        calibrated costs.
+        """
+        height = getattr(request, "height", None)
+        width = getattr(request, "width", None)
+        num_frames = getattr(request, "num_frames", None)
+        if isinstance(height, list):
+            height = height[0] if height else None
+        if isinstance(width, list):
+            width = width[0] if width else None
+        if isinstance(num_frames, list):
+            num_frames = num_frames[0] if num_frames else None
+        height = int(height or 1)
+        width = int(width or 1)
+        num_frames = int(num_frames or 1)
+        return float(max(1, height) * max(1, width) * max(1, num_frames))
+
     def dump_to_json(self, file_path: str):
         output_dict = shallow_asdict(self)
         del_keys = []

fastvideo/configs/pipelines/flux.py

@@ -1,74 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-
-from __future__ import annotations
-
-from collections.abc import Callable
-from dataclasses import dataclass, field
-
-import torch
-
-from fastvideo.configs.models import EncoderConfig
-from fastvideo.configs.models.dits.flux import FluxDiTConfig
-from fastvideo.configs.models.encoders import (
-    BaseEncoderOutput,
-    CLIPTextConfig,
-    T5LargeConfig,
-)
-from fastvideo.configs.models.vaes.autoencoder_kl import AutoencoderKLVAEConfig
-from fastvideo.configs.pipelines.base import PipelineConfig, preprocess_text
-
-
-def _flux_clip_pooled_postprocess(outputs: BaseEncoderOutput) -> torch.Tensor:
-    """CLIP branch for FLUX: Diffusers uses pooled prompt embeddings only."""
-    if outputs.pooler_output is None:
-        raise RuntimeError(
-            "FLUX CLIP conditioning requires pooler_output. Ensure the CLIP text encoder returns pooled features.")
-    return outputs.pooler_output
-
-
-def _flux_t5_sequence_postprocess(outputs: BaseEncoderOutput) -> torch.Tensor:
-    if outputs.last_hidden_state is None:
-        raise RuntimeError("FLUX T5 conditioning requires last_hidden_state.")
-    return outputs.last_hidden_state
-
-
-@dataclass
-class FluxPipelineConfig(PipelineConfig):
-    """Pipeline layout for Diffusers FLUX.1-dev (CLIP + T5 + packed DiT + FlowMatch)."""
-
-    scheduler_arch: str = "FlowMatchEulerDiscreteScheduler"
-    transformer_arch: str = "FluxTransformer2DModel"
-    vae_arch: str = "AutoencoderKL"
-    text_encoder_archs: tuple[str, ...] = ("CLIPTextModel", "T5EncoderModel")
-    tokenizer_archs: tuple[str, ...] = ("CLIPTokenizer", "T5TokenizerFast")
-
-    dit_config: FluxDiTConfig = field(default_factory=FluxDiTConfig)
-    vae_config: AutoencoderKLVAEConfig = field(default_factory=AutoencoderKLVAEConfig)
-
-    embedded_cfg_scale: float = 3.5
-    flow_shift: float | None = None
-
-    text_encoder_configs: tuple[EncoderConfig, ...] = field(default_factory=lambda: (CLIPTextConfig(), T5LargeConfig()))
-    preprocess_text_funcs: tuple[Callable[[str], str],
-                                 ...] = field(default_factory=lambda: (preprocess_text, preprocess_text))
-    postprocess_text_funcs: tuple[Callable[[BaseEncoderOutput], torch.Tensor], ...] = field(
-        default_factory=lambda: (_flux_clip_pooled_postprocess, _flux_t5_sequence_postprocess))
-
-    dit_precision: str = "bf16"
-    vae_precision: str = "fp32"
-    text_encoder_precisions: tuple[str, ...] = field(default_factory=lambda: ("fp32", "bf16"))
-
-    def __post_init__(self) -> None:
-        te_cfgs = list(self.text_encoder_configs)
-        if len(te_cfgs) >= 1:
-            te_cfgs[0].tokenizer_kwargs.setdefault("padding", "max_length")
-            te_cfgs[0].tokenizer_kwargs.setdefault("max_length", 77)
-            te_cfgs[0].tokenizer_kwargs.setdefault("truncation", True)
-            te_cfgs[0].tokenizer_kwargs.setdefault("return_tensors", "pt")
-        if len(te_cfgs) >= 2:
-            cap = 512
-            te_cfgs[1].tokenizer_kwargs["max_length"] = min(int(te_cfgs[1].tokenizer_kwargs.get("max_length", cap)),
-                                                            cap)
-            te_cfgs[1].tokenizer_kwargs.setdefault("padding", "max_length")
-            te_cfgs[1].tokenizer_kwargs.setdefault("truncation", True)
-            te_cfgs[1].tokenizer_kwargs.setdefault("return_tensors", "pt")

fastvideo/configs/pipelines/glm_image.py

@@ -1,46 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-from collections.abc import Callable
-from dataclasses import dataclass, field
-
-import torch
-
-from fastvideo.configs.models import DiTConfig, EncoderConfig, VAEConfig
-from fastvideo.configs.models.dits.glm_image import GlmImageDiTConfig
-from fastvideo.configs.models.encoders import BaseEncoderOutput, T5Config
-from fastvideo.configs.models.vaes.glm_image import GlmImageVAEConfig
-from fastvideo.configs.pipelines.base import PipelineConfig
-
-
-def glm_image_t5_postprocess(outputs: BaseEncoderOutput) -> torch.Tensor:
-    mask: torch.Tensor = outputs.attention_mask
-    hidden_state: torch.Tensor = outputs.last_hidden_state
-    seq_lens = mask.gt(0).sum(dim=1).long()
-
-    assert torch.isnan(hidden_state).sum() == 0, "T5 hidden states contain NaN"
-
-    max_len = 512
-    prompt_embeds = [u[:min(v, max_len)] for u, v in zip(hidden_state, seq_lens, strict=True)]
-    prompt_embeds_tensor: torch.Tensor = torch.stack(
-        [torch.cat([u, u.new_zeros(max_len - u.size(0), u.size(1))]) for u in prompt_embeds], dim=0)
-
-    return prompt_embeds_tensor
-
-
-@dataclass
-class GlmImageConfig(PipelineConfig):
-
-    dit_config: DiTConfig = field(default_factory=GlmImageDiTConfig)
-    dit_precision: str = "bf16"
-
-    vae_config: VAEConfig = field(default_factory=GlmImageVAEConfig)
-    vae_precision: str = "fp32"
-    vae_tiling: bool = True
-    vae_sp: bool = False
-
-    text_encoder_configs: tuple[EncoderConfig, ...] = field(default_factory=lambda: (T5Config(), ))
-    text_encoder_precisions: tuple[str, ...] = field(default_factory=lambda: ("fp32", ))
-    postprocess_text_funcs: tuple[Callable[[BaseEncoderOutput], torch.Tensor],
-                                  ...] = field(default_factory=lambda: (glm_image_t5_postprocess, ))
-
-    flow_shift: float | None = 1.0
-    embedded_cfg_scale: float = 7.5

fastvideo/configs/pipelines/wan.py

@@ -59,6 +59,12 @@ class WanT2V480PConfig(PipelineConfig):
     vae_decode_precision: str = "bf16"
     text_encoder_precisions: tuple[str, ...] = field(default_factory=lambda: ("fp32", ))
 
+    # The generic Wan denoising path is the only one verified for dynamic
+    # request batching (parity runs on Wan2.1-T2V). DMD/causal Wan variants
+    # inherit this flag but are vetoed by
+    # PipelineConfig.dynamic_batching_supported().
+    supports_dynamic_batching: bool = True
+
     # self-forcing params
     warp_denoising_step: bool = True
 

fastvideo/entrypoints/openai/api_server.py

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

fastvideo/entrypoints/openai/batching.py

@@ -0,0 +1,227 @@
+# 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
+        pipeline_config = self._fastvideo_args.pipeline_config
+        if self.enabled and not pipeline_config.dynamic_batching_supported():
+            raise RuntimeError(
+                f"Pipeline config {type(pipeline_config).__name__} does not support dynamic request batching; "
+                "run with --batching-mode disabled, or use a pipeline that opts in via "
+                "supports_dynamic_batching=True (DMD/causal denoising variants are always excluded)")
+        self._task = asyncio.create_task(self._run(), name="fastvideo-video-batch-scheduler")
+        self._task.add_done_callback(self._on_run_done)
+
+    def _on_run_done(self, task: asyncio.Task) -> None:
+        if task.cancelled():
+            exc: BaseException | None = asyncio.CancelledError()
+        else:
+            exc = task.exception()
+        if exc is None:
+            return
+        logger.error("Video batch scheduler task died; failing all pending requests", exc_info=exc)
+        self._stopped = True
+        self._drain_and_fail_waiting(RuntimeError(f"Video batch scheduler crashed and cannot dispatch: {exc!r}"))
+
+    def _drain_and_fail_waiting(self, error: BaseException) -> None:
+        """Fail every job still waiting in ``_pending`` or ``_queue``."""
+        while self._pending:
+            job = self._pending.popleft()
+            if not job.future.done():
+                job.future.set_exception(error)
+        while True:
+            try:
+                queued = self._queue.get_nowait()
+            except asyncio.QueueEmpty:
+                return
+            if queued is not None and not queued.future.done():
+                queued.future.set_exception(error)
+
+    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
+            try:
+                batch = await self._collect_batch(job)
+                await self._dispatch(batch)
+            except BaseException:
+                # The job was already popped off the queue; park it back in
+                # _pending so the crash drain in _on_run_done fails its future
+                # instead of leaving it hanging forever.
+                self._pending.append(job)
+                raise
+
+        # Fail jobs stranded in _pending AND still sitting in _queue (e.g.
+        # submitted while the final dispatch was in flight when stop() landed).
+        self._drain_and_fail_waiting(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:
+            timeout = deadline - time.perf_counter()
+            if timeout > 0:
+                try:
+                    candidate = await asyncio.wait_for(self._get_next_job(), timeout=timeout)
+                except (TimeoutError, asyncio.TimeoutError):
+                    # asyncio.TimeoutError is only aliased to the builtin on
+                    # Python 3.11+; catch both so 3.10 does not kill _run.
+                    break
+            else:
+                # delay=0 means "don't wait", and an already-expired deadline
+                # means the first job waited behind an in-flight generation:
+                # in both cases greedily drain whatever is already queued so
+                # max_size still coalesces instead of dispatching batches of 1.
+                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 as exc:
+            # Deduplicated per message so a systematically broken probe (which
+            # silently disables batching) is visible without log spam.
+            logger.warning_once("Dynamic batch compatibility probe failed; treating requests as "
+                                f"incompatible (batching disabled for these requests): {exc!r}")
+            return False
+        return can_dynamic_batch(
+            base_sampling,
+            candidate_sampling,
+            base_extra=base_extra,
+            candidate_extra=candidate_extra,
+        ).can_batch
+
+    def _sampling_param_from_kwargs(self, kwargs: dict[str, Any]) -> tuple[SamplingParam, dict[str, Any]]:
+        sampling_param = SamplingParam.from_pretrained(self._fastvideo_args.model_path)
+        updates = dict(kwargs)
+        prompt = updates.pop("prompt", None)
+        extra: dict[str, Any] = {}
+        for key in (
+                "ltx2_audio_latents",
+                "ltx2_audio_clean_latent",
+                "ltx2_audio_denoise_mask",
+                "audio_num_frames",
+                "video_position_offset_sec",
+        ):
+            if key in updates:
+                extra[key] = updates.pop(key)
+        sampling_param.update(updates)
+        sampling_param.prompt = prompt
+        return sampling_param, extra

fastvideo/entrypoints/openai/state.py

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

fastvideo/entrypoints/openai/video_api.py

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

fastvideo/entrypoints/video_generator.py

@@ -18,6 +18,7 @@ import warnings
 from collections.abc import Mapping
 from contextlib import suppress
 from copy import deepcopy
+from dataclasses import dataclass
 from typing import Any
 
 import imageio
@@ -26,6 +27,8 @@ import torch
 import torchvision
 from einops import rearrange
 
+from fastvideo.batching.admission import BatchAdmissionController
+from fastvideo.batching.signature import can_dynamic_batch
 from fastvideo.api.compat import (
     expand_request_prompt_batch,
     generator_config_to_fastvideo_args,
@@ -94,6 +97,11 @@ _FROM_PRETRAINED_CONVENIENCE_KWARGS = frozenset({
     "pin_cpu_memory",
     "enable_torch_compile",
     "torch_compile_kwargs",
+    "batching_mode",
+    "batching_max_size",
+    "batching_delay_ms",
+    "batching_config",
+    "enable_batching_metrics",
     "output_type",
     "nvfp4_fa4",
 })
@@ -112,6 +120,17 @@ def _infer_latent_batch_size(batch: ForwardBatch) -> int:
     return latent_batch_size
 
 
+@dataclass
+class _GenerationWorkItem:
+    prompt: str
+    sampling_param: SamplingParam
+    fastvideo_args: FastVideoArgs
+    batch: ForwardBatch
+    output_path: str
+    target_height: int
+    target_width: int
+
+
 class VideoGenerator:
     """
     A unified class for generating videos using diffusion models.
@@ -416,14 +435,7 @@ class VideoGenerator:
                 legacy_kwargs=kwargs,
             )
 
-            fastvideo_args = self.fastvideo_args
-            pipeline_overrides = request_to_pipeline_overrides(request)
-            if pipeline_overrides:
-                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 = self._resolve_request_fastvideo_args(request)
 
             resolved_sampling_param = request_to_sampling_param(
                 request,
@@ -439,6 +451,70 @@ class VideoGenerator:
             if log_queue:
                 self.executor.clear_log_queue()
 
+    def generate_video_batch(self, request_kwargs: list[dict[str, Any]]) -> list[dict[str, Any]]:
+        """Generate multiple legacy video requests, batching compatible items."""
+        work_items: list[_GenerationWorkItem] = []
+        reserved_output_paths: set[str] = set()
+        fastvideo_args_by_pipeline_override: dict[tuple[tuple[str, str], ...], FastVideoArgs] = {
+            (): self.fastvideo_args
+        }
+        for raw_kwargs in request_kwargs:
+            kwargs = dict(raw_kwargs)
+            prompt = kwargs.pop("prompt", None)
+            if prompt is None:
+                raise ValueError("Each batched generation request must include prompt")
+            if not isinstance(prompt, str):
+                raise TypeError(f"`prompt` must be a string, but got {type(prompt)}")
+
+            sampling_param = kwargs.pop("sampling_param", None)
+            if sampling_param is None:
+                sampling_param = SamplingParam.from_pretrained(self.fastvideo_args.model_path)
+            else:
+                sampling_param = deepcopy(sampling_param)
+
+            extra_overrides: dict[str, Any] = {}
+            for _ek in _BATCH_EXTRA_PASSTHROUGH_KEYS:
+                if _ek in kwargs:
+                    extra_overrides[_ek] = kwargs.pop(_ek)
+
+            request = legacy_generate_call_to_request(
+                prompt,
+                sampling_param,
+                legacy_kwargs=kwargs,
+            )
+            if not isinstance(request.prompt, str):
+                raise TypeError(f"`prompt` must be a string, but got {type(request.prompt)}")
+
+            fastvideo_args = self.fastvideo_args
+            pipeline_overrides = request_to_pipeline_overrides(request)
+            if pipeline_overrides:
+                override_key = tuple((key, repr(value)) for key, value in sorted(pipeline_overrides.items()))
+                fastvideo_args = fastvideo_args_by_pipeline_override.get(override_key)
+                if fastvideo_args is None:
+                    fastvideo_args = deepcopy(self.fastvideo_args)
+                    for key, value in pipeline_overrides.items():
+                        if not hasattr(fastvideo_args.pipeline_config, key):
+                            raise ValueError(f"Request field {key!r} is not supported by pipeline config overrides")
+                        setattr(fastvideo_args.pipeline_config, key, deepcopy(value))
+                    fastvideo_args_by_pipeline_override[override_key] = fastvideo_args
+
+            resolved_sampling_param = request_to_sampling_param(
+                request,
+                model_path=self.fastvideo_args.model_path,
+            )
+            output_path = self._prepare_output_path(resolved_sampling_param.output_path, request.prompt,
+                                                    reserved_output_paths)
+            work_items.append(
+                self._prepare_generation_work_item(
+                    prompt=request.prompt,
+                    sampling_param=resolved_sampling_param,
+                    fastvideo_args=fastvideo_args,
+                    output_path=output_path,
+                    _extra_overrides=extra_overrides,
+                ))
+
+        return self._generate_prepared_work_items(work_items)
+
     def _generate_request_impl(
         self,
         request: GenerationRequest,
@@ -446,6 +522,8 @@ class VideoGenerator:
         if isinstance(request.prompt, list):
             if request.inputs.prompt_path is not None:
                 raise ValueError("request.prompt list cannot be combined with request.inputs.prompt_path")
+            if (self._dynamic_batching_enabled(self.fastvideo_args) and self.fastvideo_args.prompt_txt is None):
+                return self._generate_batched_typed_requests(request)
             results: list[GenerationResult] = []
             for index, single_request in enumerate(expand_request_prompt_batch(request)):
                 prompt = single_request.prompt
@@ -461,18 +539,58 @@ class VideoGenerator:
 
         return self._generate_single_request(request)
 
+    def _generate_batched_typed_requests(self, request: GenerationRequest) -> list[GenerationResult]:
+        """Route a typed prompt-list request through the dynamic-batching work-item path.
+
+        Expanded requests share every field except prompt/image_path/video_path
+        (none of which are pipeline overrides), so one resolved FastVideoArgs is
+        shared across items — merge compatibility requires identity.
+        """
+        single_requests = expand_request_prompt_batch(request)
+        fastvideo_args = self._resolve_request_fastvideo_args(request)
+        work_items: list[_GenerationWorkItem] = []
+        reserved_output_paths: set[str] = set()
+        for single_request in single_requests:
+            sampling_param = request_to_sampling_param(
+                single_request,
+                model_path=self.fastvideo_args.model_path,
+            )
+            output_path = self._prepare_output_path(sampling_param.output_path, single_request.prompt,
+                                                    reserved_output_paths)
+            work_items.append(
+                self._prepare_generation_work_item(
+                    prompt=single_request.prompt,
+                    sampling_param=sampling_param,
+                    fastvideo_args=fastvideo_args,
+                    output_path=output_path,
+                ))
+
+        results: list[GenerationResult] = []
+        for index, legacy_result in enumerate(self._generate_prepared_work_items(work_items)):
+            wrapped = GenerationResult.from_legacy_result(legacy_result)
+            wrapped.prompt_index = index
+            if wrapped.prompt is None:
+                wrapped.prompt = single_requests[index].prompt
+            results.append(wrapped)
+        return results
+
+    def _resolve_request_fastvideo_args(self, request: GenerationRequest) -> FastVideoArgs:
+        """Return ``self.fastvideo_args``, deep-copied with the request's pipeline overrides applied (if any)."""
+        pipeline_overrides = request_to_pipeline_overrides(request)
+        if not pipeline_overrides:
+            return self.fastvideo_args
+        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))
+        return fastvideo_args
+
     def _generate_single_request(
         self,
         request: GenerationRequest,
     ) -> GenerationResult | list[GenerationResult]:
-        fastvideo_args = self.fastvideo_args
-        pipeline_overrides = request_to_pipeline_overrides(request)
-        if pipeline_overrides:
-            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 = self._resolve_request_fastvideo_args(request)
 
         sampling_param = request_to_sampling_param(
             request,
@@ -535,6 +653,29 @@ class VideoGenerator:
 
             logger.info("Found %d prompts in %s", len(prompts), prompt_txt_path)
 
+            if self._dynamic_batching_enabled(fastvideo_args):
+                work_items: list[_GenerationWorkItem] = []
+                reserved_output_paths: set[str] = set()
+                for batch_prompt in prompts:
+                    item_kwargs = dict(kwargs)
+                    item_kwargs["output_path"] = self._prepare_output_path(sampling_param.output_path, batch_prompt,
+                                                                           reserved_output_paths)
+                    work_items.append(
+                        self._prepare_generation_work_item(
+                            prompt=batch_prompt,
+                            sampling_param=sampling_param,
+                            fastvideo_args=fastvideo_args,
+                            **item_kwargs,
+                        ))
+
+                results = self._generate_prepared_work_items(work_items, tolerate_failures=True)
+                for i, (result, batch_prompt) in enumerate(zip(results, prompts, strict=True)):
+                    result["prompt_index"] = i
+                    result["prompt"] = batch_prompt
+                logger.info("Completed batch processing. Generated %d videos successfully.",
+                            sum(1 for result in results if "error" not in result))
+                return results
+
             results = []
             for i, batch_prompt in enumerate(prompts):
                 logger.info("Processing prompt %d/%d: %s...", i + 1, len(prompts), batch_prompt[:100])
@@ -588,6 +729,7 @@ class VideoGenerator:
         self,
         output_path: str,
         prompt: str,
+        reserved_paths: set[str] | None = None,
     ) -> str:
         """Build a unique, sanitized output file path.
 
@@ -602,6 +744,9 @@ class VideoGenerator:
         - Invalid filename characters are removed; if the name changes, a
           warning is logged.
         - If the target path already exists, a numeric suffix is appended.
+        - ``reserved_paths`` lets batch callers resolve every path before any
+          file is written: paths in the set are treated as taken, and the
+          chosen path is added to the set.
         """
         target_ext = ".png" if self._is_image_workload() else ".mp4"
 
@@ -646,53 +791,57 @@ class VideoGenerator:
         if output_dir:
             os.makedirs(output_dir, exist_ok=True)
 
+        def _is_taken(path: str) -> bool:
+            return os.path.exists(path) or (reserved_paths is not None and path in reserved_paths)
+
         new_output_path = os.path.join(output_dir, out_name)
         counter = 1
-        while os.path.exists(new_output_path):
+        while _is_taken(new_output_path):
             name_part, ext_part = os.path.splitext(out_name)
             new_name = f"{name_part}_{counter}{ext_part}"
             new_output_path = os.path.join(output_dir, new_name)
             counter += 1
+        if reserved_paths is not None:
+            reserved_paths.add(new_output_path)
         return new_output_path
 
-    def _generate_single_video(
+    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,
-        sampling_param: SamplingParam | None = None,
-        fastvideo_args: FastVideoArgs | None = None,
+        prompt: str | list[str],
+        sampling_param: SamplingParam,
+        fastvideo_args: FastVideoArgs,
         **kwargs,
-    ) -> dict[str, Any]:
-        """Internal method for single video generation"""
-        if fastvideo_args is None:
-            fastvideo_args = self.fastvideo_args
+    ) -> _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)}")
 
-        # Validate inputs
-        if not isinstance(prompt, str):
-            raise TypeError(f"`prompt` must be a string, but got {type(prompt)}")
-        prompt = prompt.strip()
         sampling_param = deepcopy(sampling_param)
         output_path = kwargs["output_path"]
-        prompt_embeds = kwargs.get("prompt_embeds")
-        sampling_param.prompt = prompt
-        # Process negative prompt
+        sampling_param.prompt = prompt_value
         if sampling_param.negative_prompt is not None:
             sampling_param.negative_prompt = sampling_param.negative_prompt.strip()
 
-        # Validate dimensions
-        if (sampling_param.height <= 0 or sampling_param.width <= 0 or sampling_param.num_frames <= 0):
+        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}")
 
-        # Calculate sizes
         target_height = align_to(sampling_param.height, 16)
         target_width = align_to(sampling_param.width, 16)
-
-        # Calculate latent sizes
         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]
 
-        # Log parameters
         debug_str = f"""
                       height: {target_height}
                        width: {target_width}
@@ -712,7 +861,6 @@ class VideoGenerator:
         """ # type: ignore[attr-defined]
         logger.info(debug_str)
 
-        # Prepare batch
         batch = ForwardBatch(
             **shallow_asdict(sampling_param),
             eta=0.0,
@@ -720,19 +868,30 @@ class VideoGenerator:
             VSA_sparsity=fastvideo_args.VSA_sparsity,
         )
         # Allow precomputed prompt_embeds (e.g. from diffusers) to skip text encoding
+        prompt_embeds = kwargs.get("prompt_embeds")
         if prompt_embeds is not None:
             batch.prompt_embeds = (list(prompt_embeds) if isinstance(prompt_embeds, list | tuple) else [prompt_embeds])
 
-        extra_overrides = kwargs.pop("_extra_overrides", {})
+        extra_overrides = kwargs.get("_extra_overrides", {})
         for _ek, _ev in extra_overrides.items():
             batch.extra[_ek] = _ev
 
-        # Run inference
-        start_time = time.perf_counter()
+        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,
+        )
 
-        # Execute forward pass in a new thread for non-blocking tensor
-        # allocation. Capture thread exceptions so we can surface the true
-        # failure in the main thread instead of later hitting None outputs.
+    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": ""}
@@ -747,21 +906,6 @@ class VideoGenerator:
 
         thread = threading.Thread(target=execute_forward_thread)
         thread.start()
-        latent_batch_size = _infer_latent_batch_size(batch)
-        # When ``output_type == "latent"`` the forward output has latent
-        # shape (e.g. ``[B, C_latent, T_latent, H_latent, W_latent]``)
-        # rather than the pre-allocation's pixel shape. Skip the pinned
-        # ~50 MB buffer entirely; we always fall through to the
-        # ``samples = output_batch.output.cpu()`` branch below in that
-        # mode. ``skip_pixel_prealloc`` also gates the slow-path warning.
-        skip_pixel_prealloc = fastvideo_args.output_type == "latent"
-        if skip_pixel_prealloc:
-            samples = torch.empty(0, device='cpu')
-        else:
-            samples = torch.empty(
-                (latent_batch_size, 3, sampling_param.num_frames, sampling_param.height, sampling_param.width),
-                device='cpu',
-                pin_memory=fastvideo_args.pin_cpu_memory)
         thread.join()
 
         if thread_error["error"] is not None:
@@ -773,29 +917,51 @@ class VideoGenerator:
             raise RuntimeError("Forward execution returned no output tensor. "
                                "This usually means the executor/pipeline failed earlier.")
 
-        if output_batch.output.shape == samples.shape:
-            samples.copy_(output_batch.output)
-        else:
-            if not skip_pixel_prealloc:
-                logger.warning("Output shape %s does not match expected shape %s; use slow path",
-                               output_batch.output.shape, samples.shape)
-            samples = output_batch.output.cpu()
-        logging_info = output_batch.logging_info
-
         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
 
-        # Three mutually-exclusive output modes determine whether (a) we
-        # build an RGB frame buffer and (b) what file we write to disk:
-        #
-        #   1. `output_type == "latent"` — VAE is bypassed in DecodingStage
-        #      and `samples` holds raw latents (arbitrary channel count).
-        #      The RGB grid / uint8 / mp4 / png pipeline below cannot
-        #      consume those, so we skip it entirely and let callers work
-        #      with the latent tensor directly via `result["samples"]`.
-        #   2. Audio-only workload — `samples` is a 1×3×1×8×8 placeholder
-        #      no caller will use; skip the grid loop and save a `.wav`.
-        #   3. Pixel video / image — the historical happy path.
         is_latent_output = fastvideo_args.output_type == "latent"
         audio_only = bool(output_batch.extra.get("audio_only"))
 
@@ -821,9 +987,6 @@ class VideoGenerator:
         audio_mux_time = 0.0
         if save_to_disk:
             if audio_only:
-                # Audio-only workload: write a standalone .wav rather than
-                # muxing the audio into a placeholder mp4 (which forces
-                # ffmpeg to round 8x8 placeholder frames up to 16x16).
                 output_path = self._rewrite_extension(output_path, ".wav")
                 save_start = time.perf_counter()
                 self._write_pcm_wav(
@@ -834,20 +997,16 @@ class VideoGenerator:
                 save_video_time = time.perf_counter() - save_start
                 logger.info("Saved audio to %s", output_path)
             elif self._is_image_workload():
-                # Image workloads (t2i, i2i, …): save the first frame as PNG.
-                assert frames is not None  # implied by save_to_disk and not audio_only
+                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  # implied by save_to_disk and not audio_only
+                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:
-                    # Single-pass save path: encode video+audio once, avoiding
-                    # second-pass remux overhead. AudioMuxStage remains 0.0 on
-                    # success because audio is already present in the saved MP4.
                     save_start = time.perf_counter()
                     save_ok = self._save_video_with_audio_ffmpeg_pipe(
                         output_path=output_path,
@@ -895,17 +1054,14 @@ class VideoGenerator:
         e2e_time = time.perf_counter() - start_time
         logger.info("End-to-end latency: %.2f seconds", e2e_time)
 
-        result: dict[str, Any] = {
-            "prompts": prompt,
+        return {
+            "prompts": work_item.prompt,
             "samples": samples if batch.return_frames else None,
             "frames": frames if batch.return_frames else None,
-            # Audio is the primary output for audio workloads — return it
-            # whenever the pipeline produced one, regardless of
-            # `return_frames` (which gates the video-shaped buffers).
             "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": (target_height, target_width, batch.num_frames),
+            "size": (work_item.target_height, work_item.target_width, batch.num_frames),
             "generation_time": gen_time,
             "e2e_latency": e2e_time,
             "logging_info": logging_info,
@@ -916,8 +1072,168 @@ class VideoGenerator:
             "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])]
+
+        pipeline_config = fastvideo_args.pipeline_config
+        if not pipeline_config.dynamic_batching_supported():
+            raise ValueError(
+                f"Dynamic batching is enabled but pipeline config {type(pipeline_config).__name__} does not "
+                "support merging requests; run with batching_mode='disabled', or use a pipeline that opts in "
+                "via supports_dynamic_batching=True (DMD/causal denoising variants are always excluded)")
+
+        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,
+        sampling_param: SamplingParam | None = None,
+        fastvideo_args: FastVideoArgs | None = None,
+        **kwargs,
+    ) -> dict[str, Any]:
+        """Internal method for single video generation"""
+        if fastvideo_args is None:
+            fastvideo_args = self.fastvideo_args
+
+        if not isinstance(prompt, str):
+            raise TypeError(f"`prompt` must be a string, but got {type(prompt)}")
+
+        work_item = self._prepare_generation_work_item(
+            prompt=prompt,
+            sampling_param=sampling_param,
+            fastvideo_args=fastvideo_args,
+            **kwargs,
+        )
+        return self._execute_single_work_item(work_item)
+
     @staticmethod
     def _wrap_legacy_result(
         result: dict[str, Any] | list[dict[str, Any]], ) -> GenerationResult | list[GenerationResult]:

fastvideo/fastvideo_args.py

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

fastvideo/layers/rotary_embedding.py

@@ -295,7 +295,6 @@ def get_1d_rotary_pos_embed(
     interpolation_factor: float = 1.0,
     dtype: torch.dtype = torch.float32,
     use_real: bool = True,
-    freqs_dtype: torch.dtype | None = None,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """
     Precompute the frequency tensor for complex exponential (cis) with given dimensions.
@@ -320,16 +319,12 @@ def get_1d_rotary_pos_embed(
     if isinstance(pos, int):
         pos = torch.arange(pos).float()
 
-    # freqs_dtype is an alias for dtype (Diffusers-compatible calling convention).
-    if freqs_dtype is not None:
-        dtype = freqs_dtype
-
     # proposed by reddit user bloc97, to rescale rotary embeddings to longer sequence length without fine-tuning
     # has some connection to NTK literature
     if theta_rescale_factor != 1.0:
         theta *= theta_rescale_factor**(dim / (dim - 2))
 
-    freqs = 1.0 / (theta**(torch.arange(0, dim, 2, device=pos.device)[:(dim // 2)].to(dtype) / dim))  # [D/2]
+    freqs = 1.0 / (theta**(torch.arange(0, dim, 2)[:(dim // 2)].to(dtype) / dim))  # [D/2]
     freqs = torch.outer(pos * interpolation_factor, freqs)  # [S, D/2]
     freqs_cos = freqs.cos()  # [S, D/2]
     freqs_sin = freqs.sin()  # [S, D/2]
@@ -450,21 +445,6 @@ def get_nd_rotary_pos_embed(
     return cos, sin
 
 
-_ROTARY_POS_EMBED_CACHE: dict[tuple, tuple[torch.Tensor, torch.Tensor]] = {}
-# Bound the table cache so long-running servers / causal models (which vary
-# start_frame per frame) cannot grow it without limit; entries are large float64
-# tensors. Least-recently-used eviction keeps the active resolution(s) hot while
-# capping memory.
-_ROTARY_POS_EMBED_CACHE_MAXSIZE = 16
-
-
-def _hashable(value: Any) -> Any:
-    """Return a hashable view of a scalar or sequence for use in a cache key."""
-    if isinstance(value, list | tuple):
-        return tuple(value)
-    return value
-
-
 def get_rotary_pos_embed(
     rope_sizes,
     hidden_size,
@@ -515,31 +495,6 @@ def get_rotary_pos_embed(
         sp_rank = 0
         sp_world_size = 1
 
-    # Memoize on every output-affecting argument; the table is constant across
-    # denoising steps, so this avoids recomputing the float64 cos/sin tables.
-    cache_key = (
-        _hashable(rope_sizes),
-        tuple(rope_dim_list),
-        rope_theta,
-        _hashable(theta_rescale_factor),
-        _hashable(interpolation_factor),
-        shard_dim,
-        sp_rank,
-        sp_world_size,
-        dtype,
-        start_frame,
-        use_real,
-    )
-    cached = _ROTARY_POS_EMBED_CACHE.get(cache_key)
-    if cached is not None:
-        # Move to most-recently-used position so the active table is not evicted
-        # when several resolutions / buckets share the process (LRU recency).
-        # Pop with a default: a concurrent eviction between the get() above and
-        # here would otherwise raise KeyError on the hit path.
-        if _ROTARY_POS_EMBED_CACHE.pop(cache_key, None) is not None:
-            _ROTARY_POS_EMBED_CACHE[cache_key] = cached
-        return cached
-
     freqs_cos, freqs_sin = get_nd_rotary_pos_embed(
         rope_dim_list,
         rope_sizes,
@@ -553,14 +508,6 @@ def get_rotary_pos_embed(
         start_frame=start_frame,
         use_real=use_real,
     )
-    # The returned tensors are shared cache entries: callers must never mutate
-    # them in place. Note .to(device) is an identity alias when the tensor is
-    # already on the target device (e.g. CPU runs), so it does NOT guarantee a
-    # copy — treat the tables as read-only and copy before any in-place op.
-    # Reached only on a miss, so evict the least-recently-used entry at capacity.
-    if len(_ROTARY_POS_EMBED_CACHE) >= _ROTARY_POS_EMBED_CACHE_MAXSIZE:
-        _ROTARY_POS_EMBED_CACHE.pop(next(iter(_ROTARY_POS_EMBED_CACHE)))
-    _ROTARY_POS_EMBED_CACHE[cache_key] = (freqs_cos, freqs_sin)
     return freqs_cos, freqs_sin
 
 

fastvideo/models/dits/flux.py

@@ -1,578 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-
-from __future__ import annotations
-
-from contextlib import nullcontext
-from dataclasses import dataclass
-import math
-from typing import Any
-
-import torch
-import torch.nn as nn
-
-from fastvideo.layers.rotary_embedding import apply_rotary_emb, get_1d_rotary_pos_embed
-
-from fastvideo.attention import DistributedAttention
-from fastvideo.configs.models import DiTConfig
-from fastvideo.forward_context import get_forward_context, set_forward_context
-from fastvideo.layers.linear import ReplicatedLinear
-from fastvideo.layers.visual_embedding import Timesteps
-from fastvideo.models.dits.base import BaseDiT
-from fastvideo.models.dits.sd3 import (
-    CombinedTimestepTextProjEmbeddings,
-    SD3AdaLayerNormContinuous,
-    SD3AdaLayerNormZero,
-    SD3FeedForward,
-    SD3TextProjection,
-    SD3TimestepEmbedding,
-)
-from fastvideo.platforms import AttentionBackendEnum
-
-
-@dataclass
-class FluxTransformer2DModelOutput:
-    sample: torch.Tensor
-
-
-class FluxPosEmbed(nn.Module):
-    """1D RoPE axes concatenated per Diffusers `FluxPosEmbed`."""
-
-    def __init__(self, theta: int, axes_dim: list[int]) -> None:
-        super().__init__()
-        self.theta = theta
-        self.axes_dim = axes_dim
-
-    def forward(self, ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
-        n_axes = ids.shape[-1]
-        cos_out: list[torch.Tensor] = []
-        sin_out: list[torch.Tensor] = []
-        pos = ids.float()
-        is_mps = ids.device.type == "mps"
-        is_npu = ids.device.type == "npu"
-        freqs_dtype = torch.float32 if (is_mps or is_npu) else torch.float64
-        for i in range(n_axes):
-            cos, sin = get_1d_rotary_pos_embed(
-                self.axes_dim[i],
-                pos[:, i],
-                theta=self.theta,
-                use_real=True,
-                freqs_dtype=freqs_dtype,
-            )
-            cos_out.append(cos)
-            sin_out.append(sin)
-        freqs_cos = torch.cat(cos_out, dim=-1).to(ids.device)
-        freqs_sin = torch.cat(sin_out, dim=-1).to(ids.device)
-        return freqs_cos, freqs_sin
-
-
-class FluxCombinedTimestepGuidanceTextProjEmbeddings(nn.Module):
-    def __init__(self, embedding_dim: int, pooled_projection_dim: int) -> None:
-        super().__init__()
-        self.time_proj = Timesteps(
-            num_channels=256,
-            flip_sin_to_cos=True,
-            downscale_freq_shift=0,
-        )
-        self.timestep_embedder = SD3TimestepEmbedding(
-            in_channels=256,
-            time_embed_dim=embedding_dim,
-            act_fn="silu",
-        )
-        self.guidance_embedder = SD3TimestepEmbedding(
-            in_channels=256,
-            time_embed_dim=embedding_dim,
-            act_fn="silu",
-        )
-        self.text_embedder = SD3TextProjection(
-            pooled_projection_dim,
-            embedding_dim,
-            act_fn="silu",
-        )
-
-    def forward(
-        self,
-        timestep: torch.Tensor,
-        guidance: torch.Tensor,
-        pooled_projection: torch.Tensor,
-    ) -> torch.Tensor:
-        timesteps_proj = self.time_proj(timestep)
-        timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=pooled_projection.dtype))
-        guidance_proj = self.time_proj(guidance)
-        guidance_emb = self.guidance_embedder(guidance_proj.to(dtype=pooled_projection.dtype))
-        time_guidance_emb = timesteps_emb + guidance_emb
-        pooled_projections = self.text_embedder(pooled_projection)
-        return time_guidance_emb + pooled_projections
-
-
-class FluxAdaLayerNormZeroSingle(nn.Module):
-    def __init__(self, embedding_dim: int, bias: bool = True) -> None:
-        super().__init__()
-        self.silu = nn.SiLU()
-        self.linear = ReplicatedLinear(embedding_dim, 3 * embedding_dim, bias=bias)
-        self.norm = nn.LayerNorm(embedding_dim, elementwise_affine=False, eps=1e-6)
-
-    def forward(
-        self,
-        x: torch.Tensor,
-        emb: torch.Tensor,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        emb, _ = self.linear(self.silu(emb))
-        shift_msa, scale_msa, gate_msa = emb.chunk(3, dim=1)
-        x = self.norm(x) * (1 + scale_msa[:, None]) + shift_msa[:, None]
-        return x, gate_msa
-
-
-class FluxJointAttention(nn.Module):
-    """Joint attention: text tokens precede image tokens (Diffusers order)."""
-
-    def __init__(
-        self,
-        dim: int,
-        num_attention_heads: int,
-        attention_head_dim: int,
-        supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None,
-    ) -> None:
-        super().__init__()
-        self.heads = num_attention_heads
-        self.head_dim = attention_head_dim
-        self.inner_dim = num_attention_heads * attention_head_dim
-
-        self.norm_q = nn.RMSNorm(attention_head_dim, eps=1e-6)
-        self.norm_k = nn.RMSNorm(attention_head_dim, eps=1e-6)
-        self.norm_added_q = nn.RMSNorm(attention_head_dim, eps=1e-6)
-        self.norm_added_k = nn.RMSNorm(attention_head_dim, eps=1e-6)
-
-        self.to_q = ReplicatedLinear(dim, self.inner_dim, bias=True)
-        self.to_k = ReplicatedLinear(dim, self.inner_dim, bias=True)
-        self.to_v = ReplicatedLinear(dim, self.inner_dim, bias=True)
-        self.add_q_proj = ReplicatedLinear(dim, self.inner_dim, bias=True)
-        self.add_k_proj = ReplicatedLinear(dim, self.inner_dim, bias=True)
-        self.add_v_proj = ReplicatedLinear(dim, self.inner_dim, bias=True)
-
-        self.to_out = nn.ModuleList(
-            [
-                ReplicatedLinear(self.inner_dim, dim, bias=True),
-                nn.Dropout(0.0),
-            ]
-        )
-        self.to_add_out = ReplicatedLinear(self.inner_dim, dim, bias=True)
-
-        self.attn = DistributedAttention(
-            num_heads=num_attention_heads,
-            head_size=attention_head_dim,
-            causal=False,
-            supported_attention_backends=supported_attention_backends,
-        )
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        encoder_hidden_states: torch.Tensor,
-        image_rotary_emb: tuple[torch.Tensor, torch.Tensor],
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        batch_size = hidden_states.shape[0]
-        text_seq_len = encoder_hidden_states.shape[1]
-        img_seq_len = hidden_states.shape[1]
-
-        q, _ = self.to_q(hidden_states)
-        k, _ = self.to_k(hidden_states)
-        v, _ = self.to_v(hidden_states)
-        q = q.view(batch_size, img_seq_len, self.heads, self.head_dim)
-        k = k.view(batch_size, img_seq_len, self.heads, self.head_dim)
-        v = v.view(batch_size, img_seq_len, self.heads, self.head_dim)
-        q = self.norm_q(q)
-        k = self.norm_k(k)
-
-        enc_q, _ = self.add_q_proj(encoder_hidden_states)
-        enc_k, _ = self.add_k_proj(encoder_hidden_states)
-        enc_v, _ = self.add_v_proj(encoder_hidden_states)
-        enc_q = enc_q.view(batch_size, text_seq_len, self.heads, self.head_dim)
-        enc_k = enc_k.view(batch_size, text_seq_len, self.heads, self.head_dim)
-        enc_v = enc_v.view(batch_size, text_seq_len, self.heads, self.head_dim)
-        enc_q = self.norm_added_q(enc_q)
-        enc_k = self.norm_added_k(enc_k)
-
-        q = torch.cat([enc_q, q], dim=1)
-        k = torch.cat([enc_k, k], dim=1)
-        v = torch.cat([enc_v, v], dim=1)
-
-        q = apply_rotary_emb(q, image_rotary_emb, sequence_dim=1)
-        k = apply_rotary_emb(k, image_rotary_emb, sequence_dim=1)
-
-        joint_out, _ = self.attn(q, k, v)
-        joint_out = joint_out.reshape(batch_size, text_seq_len + img_seq_len, self.inner_dim)
-
-        enc_out = joint_out[:, :text_seq_len]
-        img_out = joint_out[:, text_seq_len:]
-
-        img_out, _ = self.to_out[0](img_out)
-        img_out = self.to_out[1](img_out)
-        enc_out, _ = self.to_add_out(enc_out)
-        return img_out, enc_out
-
-
-class FluxSingleStreamAttention(nn.Module):
-    """Self-attention on concatenated text+image sequence (single blocks)."""
-
-    def __init__(
-        self,
-        dim: int,
-        num_attention_heads: int,
-        attention_head_dim: int,
-        supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None,
-    ) -> None:
-        super().__init__()
-        self.heads = num_attention_heads
-        self.head_dim = attention_head_dim
-        self.inner_dim = num_attention_heads * attention_head_dim
-
-        self.norm_q = nn.RMSNorm(attention_head_dim, eps=1e-6)
-        self.norm_k = nn.RMSNorm(attention_head_dim, eps=1e-6)
-        self.to_q = ReplicatedLinear(dim, self.inner_dim, bias=True)
-        self.to_k = ReplicatedLinear(dim, self.inner_dim, bias=True)
-        self.to_v = ReplicatedLinear(dim, self.inner_dim, bias=True)
-        self.attn = DistributedAttention(
-            num_heads=num_attention_heads,
-            head_size=attention_head_dim,
-            causal=False,
-            supported_attention_backends=supported_attention_backends,
-        )
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        image_rotary_emb: tuple[torch.Tensor, torch.Tensor],
-    ) -> torch.Tensor:
-        batch_size, seq_len, _ = hidden_states.shape
-        q, _ = self.to_q(hidden_states)
-        k, _ = self.to_k(hidden_states)
-        v, _ = self.to_v(hidden_states)
-        q = q.view(batch_size, seq_len, self.heads, self.head_dim)
-        k = k.view(batch_size, seq_len, self.heads, self.head_dim)
-        v = v.view(batch_size, seq_len, self.heads, self.head_dim)
-        q = self.norm_q(q)
-        k = self.norm_k(k)
-        q = apply_rotary_emb(q, image_rotary_emb, sequence_dim=1)
-        k = apply_rotary_emb(k, image_rotary_emb, sequence_dim=1)
-        out, _ = self.attn(q, k, v)
-        return out.reshape(batch_size, seq_len, self.inner_dim)
-
-
-class FluxTransformerBlock(nn.Module):
-    def __init__(
-        self,
-        dim: int,
-        num_attention_heads: int,
-        attention_head_dim: int,
-        supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None,
-    ) -> None:
-        super().__init__()
-        self.norm1 = SD3AdaLayerNormZero(dim)
-        self.norm1_context = SD3AdaLayerNormZero(dim)
-        self.attn = FluxJointAttention(
-            dim=dim,
-            num_attention_heads=num_attention_heads,
-            attention_head_dim=attention_head_dim,
-            supported_attention_backends=supported_attention_backends,
-        )
-        self.norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
-        self.ff = SD3FeedForward(
-            dim=dim,
-            dim_out=dim,
-            activation_fn="gelu-approximate",
-        )
-        self.norm2_context = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
-        self.ff_context = SD3FeedForward(
-            dim=dim,
-            dim_out=dim,
-            activation_fn="gelu-approximate",
-        )
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        encoder_hidden_states: torch.Tensor,
-        temb: torch.Tensor,
-        image_rotary_emb: tuple[torch.Tensor, torch.Tensor],
-        joint_attention_kwargs: dict[str, Any] | None = None,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        del joint_attention_kwargs
-        norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(hidden_states, emb=temb)
-        (norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp) = self.norm1_context(
-            encoder_hidden_states, emb=temb
-        )
-
-        attn_output, context_attn_output = self.attn(
-            hidden_states=norm_hidden_states,
-            encoder_hidden_states=norm_encoder_hidden_states,
-            image_rotary_emb=image_rotary_emb,
-        )
-
-        attn_output = gate_msa.unsqueeze(1) * attn_output
-        hidden_states = hidden_states + attn_output
-
-        norm_hidden_states = self.norm2(hidden_states)
-        norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
-        ff_output = self.ff(norm_hidden_states)
-        hidden_states = hidden_states + gate_mlp.unsqueeze(1) * ff_output
-
-        context_attn_output = c_gate_msa.unsqueeze(1) * context_attn_output
-        encoder_hidden_states = encoder_hidden_states + context_attn_output
-
-        norm_encoder_hidden_states = self.norm2_context(encoder_hidden_states)
-        norm_encoder_hidden_states = norm_encoder_hidden_states * (1 + c_scale_mlp[:, None]) + c_shift_mlp[:, None]
-        context_ff_output = self.ff_context(norm_encoder_hidden_states)
-        encoder_hidden_states = encoder_hidden_states + (c_gate_mlp.unsqueeze(1) * context_ff_output)
-        if encoder_hidden_states.dtype == torch.float16:
-            encoder_hidden_states = encoder_hidden_states.clip(-65504, 65504)
-
-        return encoder_hidden_states, hidden_states
-
-
-class FluxSingleTransformerBlock(nn.Module):
-    def __init__(
-        self,
-        dim: int,
-        num_attention_heads: int,
-        attention_head_dim: int,
-        mlp_ratio: float = 4.0,
-        supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None,
-    ) -> None:
-        super().__init__()
-        mlp_hidden_dim = int(dim * mlp_ratio)
-        self.norm = FluxAdaLayerNormZeroSingle(dim)
-        self.proj_mlp = ReplicatedLinear(dim, mlp_hidden_dim, bias=True)
-        self.act_mlp = nn.GELU(approximate="tanh")
-        self.proj_out = ReplicatedLinear(dim + mlp_hidden_dim, dim, bias=True)
-        self.attn = FluxSingleStreamAttention(
-            dim=dim,
-            num_attention_heads=num_attention_heads,
-            attention_head_dim=attention_head_dim,
-            supported_attention_backends=supported_attention_backends,
-        )
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        encoder_hidden_states: torch.Tensor,
-        temb: torch.Tensor,
-        image_rotary_emb: tuple[torch.Tensor, torch.Tensor],
-        joint_attention_kwargs: dict[str, Any] | None = None,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        del joint_attention_kwargs
-        text_seq_len = encoder_hidden_states.shape[1]
-        hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
-        residual = hidden_states
-        norm_hidden_states, gate = self.norm(hidden_states, emb=temb)
-        mlp_hidden_states = self.act_mlp(self.proj_mlp(norm_hidden_states)[0])
-        attn_output = self.attn(
-            hidden_states=norm_hidden_states,
-            image_rotary_emb=image_rotary_emb,
-        )
-        hidden_states = torch.cat([attn_output, mlp_hidden_states], dim=2)
-        gate = gate.unsqueeze(1)
-        hidden_states = gate * self.proj_out(hidden_states)[0]
-        hidden_states = residual + hidden_states
-        if hidden_states.dtype == torch.float16:
-            hidden_states = hidden_states.clip(-65504, 65504)
-        encoder_hidden_states = hidden_states[:, :text_seq_len]
-        hidden_states = hidden_states[:, text_seq_len:]
-        return encoder_hidden_states, hidden_states
-
-
-class FluxTransformer2DModel(BaseDiT):
-    """FastVideo FLUX transformer; load Diffusers FLUX safetensors 1:1."""
-
-    _fsdp_shard_conditions = [
-        lambda n, m: (n.startswith("transformer_blocks.") or n.startswith("single_transformer_blocks."))
-        and n.split(".")[-1].isdigit(),
-    ]
-    _compile_conditions = _fsdp_shard_conditions
-    # HF weight names already match this module layout (cf. SGLang regex maps).
-    param_names_mapping: dict[str, Any] = {}
-    reverse_param_names_mapping: dict[str, Any] = {}
-    lora_param_names_mapping: dict[str, Any] = {}
-    _supported_attention_backends = (
-        AttentionBackendEnum.FLASH_ATTN,
-        AttentionBackendEnum.TORCH_SDPA,
-    )
-
-    def __init__(self, config: DiTConfig, hf_config: dict[str, Any], **kwargs) -> None:
-        del kwargs
-        super().__init__(config=config, hf_config=hf_config)
-        self.fastvideo_config = config
-        self.hf_config = hf_config
-        arch = config.arch_config
-
-        out_ch = arch.out_channels
-        self.out_channels = out_ch if out_ch is not None else arch.in_channels
-        self.inner_dim = arch.num_attention_heads * arch.attention_head_dim
-        self.hidden_size = self.inner_dim
-        self.num_attention_heads = arch.num_attention_heads
-        self.num_channels_latents = arch.in_channels
-
-        axes_list = list(arch.axes_dims_rope)
-        self.pos_embed = FluxPosEmbed(theta=10000, axes_dim=axes_list)
-        if arch.guidance_embeds:
-            self.time_text_embed = FluxCombinedTimestepGuidanceTextProjEmbeddings(
-                embedding_dim=self.inner_dim,
-                pooled_projection_dim=arch.pooled_projection_dim,
-            )
-        else:
-            self.time_text_embed = CombinedTimestepTextProjEmbeddings(
-                embedding_dim=self.inner_dim,
-                pooled_projection_dim=arch.pooled_projection_dim,
-            )
-        self.context_embedder = ReplicatedLinear(arch.joint_attention_dim, self.inner_dim)
-        self.x_embedder = ReplicatedLinear(arch.in_channels, self.inner_dim)
-
-        self.transformer_blocks = nn.ModuleList(
-            [
-                FluxTransformerBlock(
-                    dim=self.inner_dim,
-                    num_attention_heads=arch.num_attention_heads,
-                    attention_head_dim=arch.attention_head_dim,
-                    supported_attention_backends=self._supported_attention_backends,
-                )
-                for _ in range(arch.num_layers)
-            ]
-        )
-        self.single_transformer_blocks = nn.ModuleList(
-            [
-                FluxSingleTransformerBlock(
-                    dim=self.inner_dim,
-                    num_attention_heads=arch.num_attention_heads,
-                    attention_head_dim=arch.attention_head_dim,
-                    supported_attention_backends=self._supported_attention_backends,
-                )
-                for _ in range(arch.num_single_layers)
-            ]
-        )
-
-        self.norm_out = SD3AdaLayerNormContinuous(
-            self.inner_dim,
-            self.inner_dim,
-            elementwise_affine=False,
-            eps=1e-6,
-            bias=True,
-            norm_type="layer_norm",
-        )
-        self.proj_out = ReplicatedLinear(
-            self.inner_dim,
-            arch.patch_size * arch.patch_size * self.out_channels,
-            bias=True,
-        )
-        self.gradient_checkpointing = False
-        self.__post_init__()
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        encoder_hidden_states: torch.Tensor | None = None,
-        pooled_projections: torch.Tensor | None = None,
-        timestep: torch.LongTensor | torch.Tensor | None = None,
-        img_ids: torch.Tensor | None = None,
-        txt_ids: torch.Tensor | None = None,
-        guidance: torch.Tensor | None = None,
-        joint_attention_kwargs: dict[str, Any] | None = None,
-        return_dict: bool = True,
-        controlnet_block_samples: Any | None = None,
-        controlnet_single_block_samples: Any | None = None,
-        controlnet_blocks_repeat: bool = False,
-        **kwargs: Any,
-    ) -> FluxTransformer2DModelOutput | tuple[torch.Tensor, ...]:
-        del kwargs
-        if encoder_hidden_states is None:
-            raise ValueError("encoder_hidden_states must be provided")
-        if pooled_projections is None:
-            raise ValueError("pooled_projections must be provided")
-        if timestep is None:
-            raise ValueError("timestep must be provided")
-        if img_ids is None or txt_ids is None:
-            raise ValueError("img_ids and txt_ids must be provided")
-
-        arch = self.fastvideo_config.arch_config
-        if arch.guidance_embeds and guidance is None:
-            raise ValueError("guidance must be provided when guidance_embeds=True")
-
-        if timestep.dim() == 0:
-            timestep = timestep[None]
-        if timestep.dim() > 1:
-            timestep = timestep.reshape(-1)
-        if timestep.shape[0] == 1 and hidden_states.shape[0] > 1:
-            timestep = timestep.expand(hidden_states.shape[0])
-
-        try:
-            get_forward_context()
-            forward_context = nullcontext()
-        except AssertionError:
-            if timestep.numel() == 0:
-                ts0 = 0
-            elif torch.is_floating_point(timestep):
-                ts0 = int(round(timestep[0].item() * 1000))
-            else:
-                ts0 = int(timestep[0].item())
-            forward_context = set_forward_context(current_timestep=ts0, attn_metadata=None)
-
-        with forward_context:
-            hidden_states, _ = self.x_embedder(hidden_states)
-
-            ts = timestep.to(hidden_states.dtype) * 1000
-            g = None if guidance is None else guidance.to(hidden_states.dtype) * 1000
-
-            if arch.guidance_embeds:
-                assert g is not None
-                temb = self.time_text_embed(ts, g, pooled_projections)
-            else:
-                temb = self.time_text_embed(timestep=ts, pooled_projection=pooled_projections)
-
-            encoder_hidden_states, _ = self.context_embedder(encoder_hidden_states)
-
-            if txt_ids.ndim == 3:
-                txt_ids = txt_ids[0]
-            if img_ids.ndim == 3:
-                img_ids = img_ids[0]
-
-            ids = torch.cat((txt_ids, img_ids), dim=0)
-            image_rotary_emb = self.pos_embed(ids)
-
-            jkwargs = joint_attention_kwargs or {}
-
-            for idx, block in enumerate(self.transformer_blocks):
-                encoder_hidden_states, hidden_states = block(
-                    hidden_states=hidden_states,
-                    encoder_hidden_states=encoder_hidden_states,
-                    temb=temb,
-                    image_rotary_emb=image_rotary_emb,
-                    joint_attention_kwargs=jkwargs,
-                )
-                if controlnet_block_samples:
-                    interval = len(self.transformer_blocks) / len(controlnet_block_samples)
-                    interval = int(math.ceil(interval))
-                    if controlnet_blocks_repeat:
-                        hidden_states = hidden_states + controlnet_block_samples[idx % len(controlnet_block_samples)]
-                    else:
-                        hidden_states = hidden_states + controlnet_block_samples[idx // interval]
-
-            for idx, block in enumerate(self.single_transformer_blocks):
-                encoder_hidden_states, hidden_states = block(
-                    hidden_states=hidden_states,
-                    encoder_hidden_states=encoder_hidden_states,
-                    temb=temb,
-                    image_rotary_emb=image_rotary_emb,
-                    joint_attention_kwargs=jkwargs,
-                )
-                if controlnet_single_block_samples:
-                    interval = len(self.single_transformer_blocks) / len(controlnet_single_block_samples)
-                    interval = int(math.ceil(interval))
-                    hidden_states = hidden_states + controlnet_single_block_samples[idx // interval]
-
-            hidden_states = self.norm_out(hidden_states, temb)
-            output, _ = self.proj_out(hidden_states)
-
-            if not return_dict:
-                return (output,)
-            return FluxTransformer2DModelOutput(sample=output)
-
-
-EntryClass = FluxTransformer2DModel

fastvideo/models/dits/glm_image.py

@@ -1,776 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-from typing import Any, Dict, List, Optional, Tuple, Union
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from fastvideo.layers.mlp import MLP
-
-from fastvideo.attention import LocalAttention
-from fastvideo.configs.models.dits.glm_image import GlmImageDiTConfig
-from fastvideo.layers.layernorm import ScaleResidualLayerNormScaleShift
-from fastvideo.layers.linear import ReplicatedLinear
-from fastvideo.layers.rotary_embedding import _apply_rotary_emb
-from fastvideo.layers.visual_embedding import Timesteps
-from fastvideo.models.dits.base import BaseDiT
-from fastvideo.platforms import AttentionBackendEnum
-
-
-class GlmImageLayerKVCache:
-
-    def __init__(self):
-        self.k_cache = None
-        self.v_cache = None
-        self.mode: Optional[str] = None
-
-    def store(self, k: torch.Tensor, v: torch.Tensor):
-        # Append along seq (dim=1).
-        if self.k_cache is None:
-            self.k_cache = k
-            self.v_cache = v
-        else:
-            self.k_cache = torch.cat([self.k_cache, k], dim=1)
-            self.v_cache = torch.cat([self.v_cache, v], dim=1)
-
-    def get(self):
-        return self.k_cache, self.v_cache
-
-    def clear(self):
-        self.k_cache = None
-        self.v_cache = None
-        self.mode = None
-
-
-class GlmImageKVCache:
-
-    def __init__(self, num_layers: int):
-        self.num_layers = num_layers
-        self.caches = [GlmImageLayerKVCache() for _ in range(num_layers)]
-
-    def __getitem__(self, layer_idx: int) -> GlmImageLayerKVCache:
-        return self.caches[layer_idx]
-
-    def set_mode(self, mode: Optional[str]):
-        if mode is not None and mode not in ["write", "read", "skip"]:
-            raise ValueError(
-                f"Invalid mode: {mode}, must be one of 'write', 'read', 'skip'"
-            )
-        for cache in self.caches:
-            cache.mode = mode
-
-    def clear(self):
-        for cache in self.caches:
-            cache.clear()
-
-
-# =============================================================================
-# Timestep and Text Projection
-# =============================================================================
-class GlmImageTimestepEmbedding(nn.Module):
-
-    def __init__(
-        self,
-        in_channels: int,
-        time_embed_dim: int,
-        act_fn: str = "silu",
-        out_dim: int = None,
-    ):
-        super().__init__()
-        if out_dim is None:
-            out_dim = time_embed_dim
-        self.linear_1 = ReplicatedLinear(in_channels, time_embed_dim, bias=True)
-        if act_fn == "silu":
-            self.act = nn.SiLU()
-        elif act_fn == "gelu":
-            self.act = nn.GELU(approximate="tanh")
-        else:
-            self.act = nn.SiLU()
-        self.linear_2 = ReplicatedLinear(time_embed_dim, out_dim, bias=True)
-
-    def forward(self, sample: torch.Tensor) -> torch.Tensor:
-        sample, _ = self.linear_1(sample)
-        sample = self.act(sample)
-        sample, _ = self.linear_2(sample)
-        return sample
-
-
-class GlmImageTextProjection(nn.Module):
-
-    def __init__(
-        self,
-        in_features: int,
-        hidden_size: int,
-        out_features: int = None,
-        act_fn: str = "silu",
-    ):
-        super().__init__()
-        if out_features is None:
-            out_features = hidden_size
-        self.linear_1 = ReplicatedLinear(in_features, hidden_size, bias=True)
-        if act_fn == "silu":
-            self.act_1 = nn.SiLU()
-        elif act_fn == "gelu_tanh":
-            self.act_1 = nn.GELU(approximate="tanh")
-        else:
-            self.act_1 = nn.SiLU()
-        self.linear_2 = ReplicatedLinear(hidden_size, out_features, bias=True)
-
-    def forward(self, caption: torch.Tensor) -> torch.Tensor:
-        hidden_states, _ = self.linear_1(caption)
-        hidden_states = self.act_1(hidden_states)
-        hidden_states, _ = self.linear_2(hidden_states)
-        return hidden_states
-
-
-class GlmImageCombinedTimestepSizeEmbeddings(nn.Module):
-
-    def __init__(
-        self,
-        embedding_dim: int,
-        condition_dim: int,
-        pooled_projection_dim: int,
-        timesteps_dim: int = 256,
-    ):
-        super().__init__()
-
-        self.time_proj = Timesteps(
-            num_channels=timesteps_dim, flip_sin_to_cos=True, downscale_freq_shift=0
-        )
-        self.condition_proj = Timesteps(
-            num_channels=condition_dim, flip_sin_to_cos=True, downscale_freq_shift=0
-        )
-        self.timestep_embedder = GlmImageTimestepEmbedding(
-            in_channels=timesteps_dim, time_embed_dim=embedding_dim
-        )
-        self.condition_embedder = GlmImageTextProjection(
-            pooled_projection_dim, embedding_dim, act_fn="silu"
-        )
-
-    def forward(
-        self,
-        timestep: torch.Tensor,
-        target_size: torch.Tensor,
-        crop_coords: torch.Tensor,
-        hidden_dtype: torch.dtype,
-    ) -> torch.Tensor:
-        timesteps_proj = self.time_proj(timestep)
-
-        crop_coords_proj = self.condition_proj(crop_coords.flatten()).view(
-            crop_coords.size(0), -1
-        )
-        target_size_proj = self.condition_proj(target_size.flatten()).view(
-            target_size.size(0), -1
-        )
-
-        condition_proj = torch.cat([crop_coords_proj, target_size_proj], dim=1)
-
-        timesteps_emb = self.timestep_embedder(
-            timesteps_proj.to(dtype=hidden_dtype)
-        )
-        condition_emb = self.condition_embedder(
-            condition_proj.to(dtype=hidden_dtype)
-        )
-
-        conditioning = timesteps_emb + condition_emb
-        return conditioning
-
-
-# =============================================================================
-# Image Projector
-# =============================================================================
-class GlmImageImageProjector(nn.Module):
-
-    def __init__(
-        self,
-        in_channels: int = 16,
-        hidden_size: int = 2560,
-        patch_size: int = 2,
-    ):
-        super().__init__()
-        self.patch_size = patch_size
-        self.proj = nn.Linear(in_channels * patch_size**2, hidden_size)
-
-    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        batch_size, channel, height, width = hidden_states.shape
-        post_patch_height = height // self.patch_size
-        post_patch_width = width // self.patch_size
-
-        hidden_states = hidden_states.reshape(
-            batch_size,
-            channel,
-            post_patch_height,
-            self.patch_size,
-            post_patch_width,
-            self.patch_size,
-        )
-        hidden_states = (
-            hidden_states.permute(0, 2, 4, 1, 3, 5).flatten(3, 5).flatten(1, 2)
-        )
-        hidden_states = self.proj(hidden_states)
-
-        return hidden_states
-
-
-# =============================================================================
-# AdaLayerNorm
-# =============================================================================
-class GlmImageAdaLayerNormZero(nn.Module):
-
-    def __init__(self, embedding_dim: int, dim: int) -> None:
-        super().__init__()
-
-        self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-5)
-        self.norm_context = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-5)
-        self.linear = ReplicatedLinear(embedding_dim, 12 * dim, bias=True)
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        encoder_hidden_states: torch.Tensor,
-        temb: torch.Tensor,
-    ) -> Tuple[torch.Tensor, ...]:
-        dtype = hidden_states.dtype
-        norm_hidden_states = self.norm(hidden_states).to(dtype=dtype)
-        norm_encoder_hidden_states = self.norm_context(encoder_hidden_states).to(
-            dtype=dtype
-        )
-
-        emb, _ = self.linear(temb)
-        (
-            shift_msa,
-            c_shift_msa,
-            scale_msa,
-            c_scale_msa,
-            gate_msa,
-            c_gate_msa,
-            shift_mlp,
-            c_shift_mlp,
-            scale_mlp,
-            c_scale_mlp,
-            gate_mlp,
-            c_gate_mlp,
-        ) = emb.chunk(12, dim=1)
-
-        hidden_states = norm_hidden_states * (
-            1 + scale_msa.unsqueeze(1)
-        ) + shift_msa.unsqueeze(1)
-        encoder_hidden_states = norm_encoder_hidden_states * (
-            1 + c_scale_msa.unsqueeze(1)
-        ) + c_shift_msa.unsqueeze(1)
-
-        return (
-            hidden_states,
-            gate_msa,
-            shift_mlp,
-            scale_mlp,
-            gate_mlp,
-            encoder_hidden_states,
-            c_gate_msa,
-            c_shift_mlp,
-            c_scale_mlp,
-            c_gate_mlp,
-        )
-
-
-# =============================================================================
-# Attention
-# =============================================================================
-class GlmImageAttention(nn.Module):
-
-    def __init__(
-        self,
-        query_dim: int,
-        heads: int,
-        dim_head: int,
-        out_dim: int,
-        bias: bool = True,
-        qk_norm: str = "layer_norm",
-        elementwise_affine: bool = False,
-        eps: float = 1e-5,
-        supported_attention_backends: set[AttentionBackendEnum] | None = None,
-        prefix: str = "",
-    ):
-        super().__init__()
-
-        self.heads = out_dim // dim_head if out_dim is not None else heads
-        self.num_kv_heads = self.heads
-        self.dim_head = dim_head
-        self.inner_dim = out_dim if out_dim is not None else dim_head * heads
-        self.inner_kv_dim = self.inner_dim
-        self.out_dim = out_dim if out_dim is not None else query_dim
-
-        self.to_q = ReplicatedLinear(query_dim, self.inner_dim, bias=bias)
-        self.to_k = ReplicatedLinear(query_dim, self.inner_kv_dim, bias=bias)
-        self.to_v = ReplicatedLinear(query_dim, self.inner_kv_dim, bias=bias)
-
-        self.to_out = nn.ModuleList(
-            [ReplicatedLinear(self.inner_dim, self.out_dim, bias=True)]
-        )
-
-        if qk_norm is None:
-            self.norm_q = None
-            self.norm_k = None
-        elif qk_norm == "layer_norm":
-            self.norm_q = nn.LayerNorm(
-                dim_head, eps=eps, elementwise_affine=elementwise_affine
-            )
-            self.norm_k = nn.LayerNorm(
-                dim_head, eps=eps, elementwise_affine=elementwise_affine
-            )
-        else:
-            raise ValueError(f"unknown qk_norm: {qk_norm}")
-
-        self.attn = LocalAttention(
-            num_heads=self.heads,
-            head_size=dim_head,
-            num_kv_heads=self.heads,
-            softmax_scale=None,
-            causal=False,
-            supported_attention_backends=supported_attention_backends,
-        )
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        encoder_hidden_states: torch.Tensor,
-        image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
-        kv_cache: Optional[GlmImageLayerKVCache] = None,
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
-        dtype = encoder_hidden_states.dtype
-
-        batch_size, text_seq_length, embed_dim = encoder_hidden_states.shape
-        batch_size, image_seq_length, embed_dim = hidden_states.shape
-        hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
-
-        # 1. QKV projections
-        query, _ = self.to_q(hidden_states)
-        key, _ = self.to_k(hidden_states)
-        value, _ = self.to_v(hidden_states)
-
-        query = query.unflatten(2, (self.heads, -1))
-        key = key.unflatten(2, (self.heads, -1))
-        value = value.unflatten(2, (self.heads, -1))
-
-        # 2. QK normalization
-        if self.norm_q is not None:
-            query = self.norm_q(query).to(dtype=dtype)
-        if self.norm_k is not None:
-            key = self.norm_k(key).to(dtype=dtype)
-
-        # 3. Rotational positional embeddings applied to latent stream
-        if image_rotary_emb is not None:
-            cos, sin = image_rotary_emb
-
-            query[:, text_seq_length:, :, :] = _apply_rotary_emb(
-                query[:, text_seq_length:, :, :], cos, sin, is_neox_style=True
-            )
-            key[:, text_seq_length:, :, :] = _apply_rotary_emb(
-                key[:, text_seq_length:, :, :], cos, sin, is_neox_style=True
-            )
-
-        # 4. KV Cache handling
-        if kv_cache is not None:
-            if kv_cache.mode == "write":
-                kv_cache.store(key, value)
-            elif kv_cache.mode == "read":
-                # Prepend cached condition k/v along seq (dim=1).
-                k_cache, v_cache = kv_cache.get()
-                key = torch.cat([k_cache, key], dim=1) if k_cache is not None else key
-                value = (
-                    torch.cat([v_cache, value], dim=1) if v_cache is not None else value
-                )
-            elif kv_cache.mode == "skip":
-                pass
-
-        hidden_states = self.attn(query, key, value)
-        hidden_states = hidden_states.flatten(2, 3)
-        hidden_states = hidden_states.to(query.dtype)
-
-        # 6. Output projection
-        hidden_states, _ = self.to_out[0](hidden_states)
-
-        encoder_hidden_states, hidden_states = hidden_states.split(
-            [text_seq_length, hidden_states.size(1) - text_seq_length], dim=1
-        )
-        return hidden_states, encoder_hidden_states
-
-
-# =============================================================================
-# Transformer Block
-# =============================================================================
-class GlmImageTransformerBlock(nn.Module):
-
-    def __init__(
-        self,
-        dim: int = 2560,
-        num_attention_heads: int = 64,
-        attention_head_dim: int = 40,
-        time_embed_dim: int = 512,
-        supported_attention_backends: set[AttentionBackendEnum] | None = None,
-        prefix: str = "",
-    ) -> None:
-        super().__init__()
-
-        # 1. Attention
-        self.norm1 = GlmImageAdaLayerNormZero(time_embed_dim, dim)
-
-        self.attn1 = GlmImageAttention(
-            query_dim=dim,
-            heads=num_attention_heads,
-            dim_head=attention_head_dim,
-            out_dim=dim,
-            bias=True,
-            qk_norm="layer_norm",
-            elementwise_affine=False,
-            eps=1e-5,
-            supported_attention_backends=supported_attention_backends,
-            prefix=f"{prefix}.attn1",
-        )
-
-        # 2. Feedforward with fused ScaleResidualLayerNorm
-        self.norm2 = ScaleResidualLayerNormScaleShift(
-            dim, norm_type="layer", eps=1e-5, elementwise_affine=False
-        )
-        self.norm2_context = ScaleResidualLayerNormScaleShift(
-            dim, norm_type="layer", eps=1e-5, elementwise_affine=False
-        )
-        self.ff = MLP(input_dim=dim, mlp_hidden_dim=dim * 4, output_dim=dim, act_type="gelu_pytorch_tanh")
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        encoder_hidden_states: torch.Tensor,
-        temb: Optional[torch.Tensor] = None,
-        image_rotary_emb: Optional[
-            Union[
-                Tuple[torch.Tensor, torch.Tensor],
-                List[Tuple[torch.Tensor, torch.Tensor]],
-            ]
-        ] = None,
-        attention_kwargs: Optional[Dict[str, Any]] = None,
-        kv_cache: Optional[GlmImageLayerKVCache] = None,
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
-        # 1. Timestep conditioning
-        (
-            norm_hidden_states,
-            gate_msa,
-            shift_mlp,
-            scale_mlp,
-            gate_mlp,
-            norm_encoder_hidden_states,
-            c_gate_msa,
-            c_shift_mlp,
-            c_scale_mlp,
-            c_gate_mlp,
-        ) = self.norm1(hidden_states, encoder_hidden_states, temb)
-
-        # 2. Attention
-        if attention_kwargs is None:
-            attention_kwargs = {}
-
-        attn_hidden_states, attn_encoder_hidden_states = self.attn1(
-            hidden_states=norm_hidden_states,
-            encoder_hidden_states=norm_encoder_hidden_states,
-            image_rotary_emb=image_rotary_emb,
-            kv_cache=kv_cache,
-            **attention_kwargs,
-        )
-
-        # 3. Feedforward (fused residual + norm + scale/shift)
-        norm_hidden_states, hidden_states = self.norm2(
-            hidden_states,
-            attn_hidden_states,
-            gate_msa.unsqueeze(1),
-            shift_mlp.unsqueeze(1),
-            scale_mlp.unsqueeze(1),
-        )
-        norm_encoder_hidden_states, encoder_hidden_states = self.norm2_context(
-            encoder_hidden_states,
-            attn_encoder_hidden_states,
-            c_gate_msa.unsqueeze(1),
-            c_shift_mlp.unsqueeze(1),
-            c_scale_mlp.unsqueeze(1),
-        )
-
-        ff_output = self.ff(norm_hidden_states)
-        ff_output_context = self.ff(norm_encoder_hidden_states)
-        hidden_states = hidden_states + ff_output * gate_mlp.unsqueeze(1)
-        encoder_hidden_states = (
-            encoder_hidden_states + ff_output_context * c_gate_mlp.unsqueeze(1)
-        )
-
-        return hidden_states, encoder_hidden_states
-
-
-# =============================================================================
-# Rotary Positional Embedding
-# =============================================================================
-class GlmImageRotaryPosEmbed(nn.Module):
-
-    def __init__(self, dim: int, patch_size: int, theta: float = 10000.0) -> None:
-        super().__init__()
-        self.dim = dim
-        self.patch_size = patch_size
-        self.theta = theta
-        self._cache_key: tuple | None = None
-        self._cache_value: tuple[torch.Tensor, torch.Tensor] | None = None
-
-    def forward(self, hidden_states: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
-        _, _, raw_h, raw_w = hidden_states.shape
-        height = raw_h // self.patch_size
-        width = raw_w // self.patch_size
-        device = hidden_states.device
-
-        cache_key = (height, width, device.type,
-                     device.index if device.index is not None else -1)
-        if self._cache_key == cache_key and self._cache_value is not None:
-            return self._cache_value
-
-        dim_h, dim_w = self.dim // 2, self.dim // 2
-        h_inv_freq = 1.0 / (
-            self.theta
-            ** (
-                torch.arange(0, dim_h, 2, dtype=torch.float32, device=device)[
-                    : (dim_h // 2)
-                ]
-                / dim_h
-            )
-        )
-        w_inv_freq = 1.0 / (
-            self.theta
-            ** (
-                torch.arange(0, dim_w, 2, dtype=torch.float32, device=device)[
-                    : (dim_w // 2)
-                ]
-                / dim_w
-            )
-        )
-        h_seq = torch.arange(height, device=device)
-        w_seq = torch.arange(width, device=device)
-        freqs_h = torch.outer(h_seq, h_inv_freq).unsqueeze(1).expand(height, width, -1)
-        freqs_w = torch.outer(w_seq, w_inv_freq).unsqueeze(0).expand(height, width, -1)
-
-        freqs = torch.cat([freqs_h, freqs_w], dim=-1).reshape(height * width, -1)
-        result = (freqs.cos(), freqs.sin())
-        self._cache_key = cache_key
-        self._cache_value = result
-        return result
-
-
-# =============================================================================
-# Final AdaLayerNorm
-# =============================================================================
-class GlmImageAdaLayerNormContinuous(nn.Module):
-
-    def __init__(
-        self,
-        embedding_dim: int,
-        conditioning_embedding_dim: int,
-        elementwise_affine: bool = True,
-        eps: float = 1e-5,
-        bias: bool = True,
-        norm_type: str = "layer_norm",
-    ):
-        super().__init__()
-        self.linear = nn.Linear(
-            conditioning_embedding_dim, embedding_dim * 2, bias=bias
-        )
-        if norm_type == "layer_norm":
-            self.norm = nn.LayerNorm(embedding_dim, eps, elementwise_affine, bias)
-        elif norm_type == "rms_norm":
-            self.norm = nn.RMSNorm(embedding_dim, eps, elementwise_affine)
-        else:
-            raise ValueError(f"unknown norm_type {norm_type}")
-
-    def forward(
-        self, x: torch.Tensor, conditioning_embedding: torch.Tensor
-    ) -> torch.Tensor:
-        emb = self.linear(conditioning_embedding.to(x.dtype))
-        scale, shift = torch.chunk(emb, 2, dim=1)
-        x = self.norm(x) * (1 + scale)[:, None, :] + shift[:, None, :]
-        return x
-
-
-# =============================================================================
-# Main Model
-# =============================================================================
-class GlmImageTransformer2DModel(BaseDiT):
-
-    _fsdp_shard_conditions = GlmImageDiTConfig().arch_config._fsdp_shard_conditions
-    _compile_conditions = GlmImageDiTConfig().arch_config._compile_conditions
-    _supported_attention_backends = (
-        AttentionBackendEnum.FLASH_ATTN,
-        AttentionBackendEnum.TORCH_SDPA,
-    )
-
-    param_names_mapping = GlmImageDiTConfig().arch_config.param_names_mapping
-    reverse_param_names_mapping = {}
-    lora_param_names_mapping = {}
-
-    def __init__(
-        self,
-        config: GlmImageDiTConfig,
-        hf_config: dict[str, Any],
-    ):
-        super().__init__(config=config, hf_config=hf_config)
-
-        arch_config = config.arch_config
-
-        self.in_channels = arch_config.in_channels
-        self.out_channels = arch_config.out_channels
-        self.patch_size = arch_config.patch_size
-        self.num_layers = arch_config.num_layers
-        self.attention_head_dim = arch_config.attention_head_dim
-        self.num_attention_heads = arch_config.num_attention_heads
-        self.text_embed_dim = arch_config.text_embed_dim
-        self.time_embed_dim = arch_config.time_embed_dim
-
-        # GlmImage uses 2 additional SDXL-like conditions - target_size, crop_coords
-        # Each of these are sincos embeddings of shape 2 * condition_dim
-        pooled_projection_dim = 2 * 2 * arch_config.condition_dim
-        inner_dim = arch_config.num_attention_heads * arch_config.attention_head_dim
-
-        self.hidden_size = inner_dim
-        self.num_channels_latents = arch_config.out_channels
-
-        # 1. RoPE
-        self.rotary_emb = GlmImageRotaryPosEmbed(
-            arch_config.attention_head_dim, arch_config.patch_size, theta=10000.0
-        )
-
-        # 2. Patch & Text-timestep embedding
-        self.image_projector = GlmImageImageProjector(
-            arch_config.in_channels, inner_dim, arch_config.patch_size
-        )
-        self.glyph_projector = MLP(
-            input_dim=arch_config.text_embed_dim,
-            mlp_hidden_dim=inner_dim,
-            output_dim=inner_dim,
-            act_type="gelu",
-        )
-        self.prior_token_embedding = nn.Embedding(
-            arch_config.prior_vq_quantizer_codebook_size, inner_dim
-        )
-        self.prior_projector = MLP(
-            input_dim=inner_dim,
-            mlp_hidden_dim=inner_dim,
-            output_dim=inner_dim,
-            act_type="silu",
-        )
-
-        self.time_condition_embed = GlmImageCombinedTimestepSizeEmbeddings(
-            embedding_dim=arch_config.time_embed_dim,
-            condition_dim=arch_config.condition_dim,
-            pooled_projection_dim=pooled_projection_dim,
-            timesteps_dim=arch_config.time_embed_dim,
-        )
-
-        # 3. Transformer blocks
-        self.transformer_blocks = nn.ModuleList(
-            [
-                GlmImageTransformerBlock(
-                    inner_dim,
-                    arch_config.num_attention_heads,
-                    arch_config.attention_head_dim,
-                    arch_config.time_embed_dim,
-                    supported_attention_backends=self._supported_attention_backends,
-                    prefix=f"transformer_blocks.{i}",
-                )
-                for i in range(arch_config.num_layers)
-            ]
-        )
-
-        # 4. Output projection
-        self.norm_out = GlmImageAdaLayerNormContinuous(
-            inner_dim, arch_config.time_embed_dim, elementwise_affine=False
-        )
-        self.proj_out = nn.Linear(
-            inner_dim,
-            arch_config.patch_size * arch_config.patch_size * arch_config.out_channels,
-            bias=True,
-        )
-
-        self.gradient_checkpointing = False
-
-        self.__post_init__()
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        encoder_hidden_states: torch.Tensor,
-        prior_token_id: torch.Tensor,
-        prior_token_drop: torch.Tensor,
-        timestep: torch.LongTensor,
-        target_size: torch.Tensor,
-        crop_coords: torch.Tensor,
-        attention_kwargs: Optional[Dict[str, Any]] = None,
-        kv_caches: Optional[GlmImageKVCache] = None,
-        kv_caches_mode: Optional[str] = None,
-        freqs_cis: Optional[
-            Union[
-                Tuple[torch.Tensor, torch.Tensor],
-                List[Tuple[torch.Tensor, torch.Tensor]],
-            ]
-        ] = None,
-        guidance: torch.Tensor = None,
-        **kwargs,
-    ) -> torch.Tensor:
-        if kv_caches is not None:
-            kv_caches.set_mode(kv_caches_mode)
-
-        batch_size, num_channels, height, width = hidden_states.shape
-
-        if isinstance(encoder_hidden_states, list):
-            encoder_hidden_states = encoder_hidden_states[0]
-
-        # 1. RoPE
-        image_rotary_emb = freqs_cis
-        if image_rotary_emb is None:
-            image_rotary_emb = self.rotary_emb(hidden_states)
-
-        # 2. Patch & Timestep embeddings
-        p = self.patch_size
-        post_patch_height = height // p
-        post_patch_width = width // p
-
-        hidden_states = self.image_projector(hidden_states)
-        encoder_hidden_states = self.glyph_projector(encoder_hidden_states)
-        prior_embedding = self.prior_token_embedding(prior_token_id)
-        # Zero dropped priors by multiply: boolean indexing + .any() syncs each step.
-        keep = (~prior_token_drop).to(device=prior_embedding.device, dtype=prior_embedding.dtype)
-        while keep.dim() < prior_embedding.dim():
-            keep = keep.unsqueeze(-1)
-        prior_embedding = prior_embedding * keep
-        prior_hidden_states = self.prior_projector(prior_embedding)
-        hidden_states = hidden_states + prior_hidden_states
-
-        temb = self.time_condition_embed(
-            timestep, target_size, crop_coords, hidden_states.dtype
-        )
-        temb = F.silu(temb)
-
-        # 3. Transformer blocks
-        for idx, block in enumerate(self.transformer_blocks):
-            hidden_states, encoder_hidden_states = block(
-                hidden_states,
-                encoder_hidden_states,
-                temb,
-                image_rotary_emb,
-                attention_kwargs,
-                kv_cache=kv_caches[idx] if kv_caches is not None else None,
-            )
-
-        # 4. Output norm & projection
-        hidden_states = self.norm_out(hidden_states, temb)
-        hidden_states = self.proj_out(hidden_states)
-
-        # 5. Unpatchify
-        hidden_states = hidden_states.reshape(
-            batch_size, post_patch_height, post_patch_width, -1, p, p
-        )
-        output = hidden_states.permute(0, 3, 1, 4, 2, 5).flatten(4, 5).flatten(2, 3)
-
-        return output.float()
-
-
-EntryClass = GlmImageTransformer2DModel

fastvideo/models/dits/ltx2.py

@@ -72,13 +72,8 @@ def _rms_norm_dispatch(
     """Use QuACK RMSNorm only for LTX-2 refine stage, else Torch RMSNorm."""
     if _is_ltx2_refine_stage():
         return _quack_rmsnorm(x, weight=weight, eps=eps)
-    # Official LTX-2 runs the DiT natively in bf16 with no autocast, so its
-    # rms_norm outputs stay bf16 on every torch version. torch 2.12 added
-    # rms_norm to autocast's float32 cast policy, which under our
-    # autocast(bfloat16) forward returns float32 instead. Restore the input
-    # dtype to match the official numerics (and torch <= 2.11 behavior).
     return torch.nn.functional.rms_norm(
-        x, (x.shape[-1], ), weight=weight, eps=eps).to(x.dtype)
+        x, (x.shape[-1], ), weight=weight, eps=eps)
 
 
 class StageAwareRMSNorm(nn.RMSNorm):

fastvideo/models/dits/wanvideo.py

@@ -1,6 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
 
-import copy
 import math
 from typing import Any
 
@@ -60,11 +59,6 @@ class WanTimeTextImageEmbedding(nn.Module):
         time_freq_dim: int,
         text_embed_dim: int,
         image_embed_dim: int | None = None,
-        *,
-        r_embedder: bool = False,
-        r_embedder_fusion: str = "additive",
-        r_embedder_gate_value: float = 0.25,
-        r_embedder_deltatime_type: str = "r",
     ):
         super().__init__()
 
@@ -83,57 +77,14 @@ class WanTimeTextImageEmbedding(nn.Module):
         if image_embed_dim is not None:
             self.image_embedder = WanImageEmbedding(image_embed_dim, dim)
 
-        # AnyFlow dual-timestep support. When r_embedder is False the forward
-        # path bypasses delta_embedder entirely and the output is byte-identical
-        # to the legacy single-timestep implementation.
-        self._r_embedder_enabled = bool(r_embedder)
-        self._r_embedder_fusion = r_embedder_fusion
-        self._r_embedder_deltatime_type = r_embedder_deltatime_type
-        if self._r_embedder_enabled:
-            if r_embedder_fusion not in ("additive", "gated"):
-                raise ValueError(
-                    "r_embedder_fusion must be one of {additive, gated}, "
-                    f"got {r_embedder_fusion!r}")
-            if r_embedder_deltatime_type not in ("r", "t-r"):
-                raise ValueError(
-                    "r_embedder_deltatime_type must be one of {r, t-r}, "
-                    f"got {r_embedder_deltatime_type!r}")
-            # Deep-copy preserves identical initialization with time_embedder,
-            # matching AnyFlow reference setup_flowmap_model() behavior.
-            self.delta_embedder = copy.deepcopy(self.time_embedder)
-            # Non-persistent buffer — gate is a hyperparameter, not learned.
-            self.register_buffer(
-                "_r_embedder_gate",
-                torch.tensor(float(r_embedder_gate_value)),
-                persistent=False,
-            )
-        else:
-            self.delta_embedder = None
-
     def forward(
         self,
         timestep: torch.Tensor,
         encoder_hidden_states: torch.Tensor,
         encoder_hidden_states_image: torch.Tensor | None = None,
         timestep_seq_len: int | None = None,
-        r_timestep: torch.Tensor | None = None,
     ):
         temb = self.time_embedder(timestep, timestep_seq_len)
-
-        if self._r_embedder_enabled and r_timestep is not None:
-            assert self.delta_embedder is not None
-            if self._r_embedder_deltatime_type == "r":
-                delta_input = r_timestep
-            else:
-                delta_input = timestep - r_timestep
-            delta_emb = self.delta_embedder(delta_input, timestep_seq_len)
-            gate = self._r_embedder_gate
-            if self._r_embedder_fusion == "gated":
-                temb = (1.0 - gate) * temb + gate * delta_emb
-            else:
-                # Additive (additional channel, no convex blend).
-                temb = temb + gate * delta_emb
-
         timestep_proj = self.time_modulation(temb)
 
         if self.text_embedder is not None:
@@ -644,10 +595,6 @@ class WanTransformer3DModel(BaseDiT):
             time_freq_dim=config.freq_dim,
             text_embed_dim=config.text_dim,
             image_embed_dim=config.image_dim,
-            r_embedder=config.r_embedder,
-            r_embedder_fusion=config.r_embedder_fusion,
-            r_embedder_gate_value=config.r_embedder_gate_value,
-            r_embedder_deltatime_type=config.r_embedder_deltatime_type,
         )
 
         # 3. Transformer blocks
@@ -689,7 +636,6 @@ class WanTransformer3DModel(BaseDiT):
                 encoder_hidden_states_image: torch.Tensor | list[torch.Tensor]
                 | None = None,
                 guidance=None,
-                r_timestep: torch.Tensor | None = None,
                 **kwargs) -> torch.Tensor:
         orig_dtype = hidden_states.dtype
         if encoder_hidden_states is not None and not isinstance(encoder_hidden_states, torch.Tensor):
@@ -738,17 +684,8 @@ class WanTransformer3DModel(BaseDiT):
         else:
             ts_seq_len = None
 
-        # AnyFlow dual-timestep — match timestep's flattening so embedder
-        # sees aligned shapes.
-        if r_timestep is not None and r_timestep.dim() == 2:
-            r_timestep = r_timestep.flatten()
-
         temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = self.condition_embedder(
-            timestep,
-            encoder_hidden_states,
-            encoder_hidden_states_image,
-            timestep_seq_len=ts_seq_len,
-            r_timestep=r_timestep)
+            timestep, encoder_hidden_states, encoder_hidden_states_image, timestep_seq_len=ts_seq_len)
         if ts_seq_len is not None:
             # batch_size, seq_len, 6, inner_dim
             timestep_proj = timestep_proj.unflatten(2, (6, -1))

fastvideo/models/encoders/glm_image_ar_loader.py

@@ -1,64 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Sole HF-import boundary for GLM-Image's AR encoder (lazy-wrapper exception E001)."""
-from __future__ import annotations
-
-from typing import Any
-
-import torch
-import torch.nn as nn
-
-from fastvideo.logger import init_logger
-
-logger = init_logger(__name__)
-
-
-class GlmImageARLoader(nn.Module):
-
-    def __init__(self, model_path: str, processor_path: str | None = None,
-                 *, torch_dtype: torch.dtype = torch.bfloat16,
-                 trust_remote_code: bool = True) -> None:
-        super().__init__()
-        from transformers import (AutoProcessor,
-                                  GlmImageForConditionalGeneration)
-        logger.info("Loading GLM-Image AR encoder from %s", model_path)
-        self._model = GlmImageForConditionalGeneration.from_pretrained(
-            model_path,
-            torch_dtype=torch_dtype,
-            trust_remote_code=trust_remote_code,
-        )
-        if processor_path is not None:
-            logger.info("Loading GLM-Image processor from %s", processor_path)
-            self.processor = AutoProcessor.from_pretrained(
-                processor_path, trust_remote_code=trust_remote_code)
-        else:
-            self.processor = None
-
-    @torch.no_grad()
-    def generate(self, *args: Any, **kwargs: Any) -> torch.Tensor:
-        return self._model.generate(*args, **kwargs)
-
-    @torch.no_grad()
-    def get_image_features(self, pixel_values: torch.Tensor,
-                           image_grid_thw: torch.Tensor) -> Any:
-        return self._model.get_image_features(pixel_values, image_grid_thw)
-
-    @torch.no_grad()
-    def get_image_tokens(self, image_embeds: torch.Tensor,
-                         image_grid_thw: torch.Tensor) -> torch.Tensor:
-        return self._model.get_image_tokens(image_embeds, image_grid_thw)
-
-    @property
-    def config(self):  # type: ignore[no-untyped-def]
-        return self._model.config
-
-    @property
-    def generation_config(self):  # type: ignore[no-untyped-def]
-        return self._model.generation_config
-
-    def to(self, *args, **kwargs):  # type: ignore[override]
-        self._model = self._model.to(*args, **kwargs)
-        return super().to(*args, **kwargs)
-
-    def eval(self):  # type: ignore[override]
-        self._model = self._model.eval()
-        return super().eval()

fastvideo/models/loader/benchmarks/benchmark_weight_loading_comparison.py

@@ -42,7 +42,7 @@ def load_independent(files: list[str], device: str):
     """Before-PR behavior: every rank reads every tensor from disk to GPU."""
     for st_file in files:
         with safe_open(st_file, framework="pt", device=device) as f:
-            for name in f.keys():  # noqa: SIM118
+            for name in f:
                 param = f.get_tensor(name)
                 yield name, param
 
@@ -54,7 +54,7 @@ def load_broadcast(files: list[str], device: str, node_group,
     handles = []
     for st_file in files:
         with safe_open(st_file, framework="pt", device=device) as f:
-            for name in f.keys():  # noqa: SIM118
+            for name in f:
                 if local_rank == 0:
                     param = f.get_tensor(name)
                 else:

fastvideo/models/loader/component_loader.py

@@ -95,8 +95,6 @@ class ComponentLoader(ABC):
             "image_processor": (ImageProcessorLoader, "transformers"),
             "feature_extractor": (ImageProcessorLoader, "transformers"),
             "image_encoder": (ImageEncoderLoader, "transformers"),
-            "vision_language_encoder": (VisionLanguageEncoderLoader, "transformers"),
-            "processor": (ProcessorLoader, "transformers"),
             "upsampler": (UpsamplerLoader, "diffusers"),
             "upsampler_2": (UpsamplerLoader, "diffusers"),
             # Stable Audio's `StableAudioMultiConditioner` bundles T5 +
@@ -525,39 +523,6 @@ class ImageEncoderLoader(TextEncoderLoader):
         )
 
 
-class VisionLanguageEncoderLoader(ComponentLoader):
-    """Loader for vision-language autoregressive encoders."""
-
-    def load(self, model_path: str, fastvideo_args: FastVideoArgs):
-        from fastvideo.distributed.parallel_state import get_local_torch_device
-        from fastvideo.models.encoders.glm_image_ar_loader import (
-            GlmImageARLoader)
-
-        logger.info("Loading vision-language encoder from %s", model_path)
-        target_device = get_local_torch_device()
-        loader = GlmImageARLoader(
-            model_path,
-            torch_dtype=torch.bfloat16,
-            trust_remote_code=fastvideo_args.trust_remote_code,
-        ).to(target_device).eval()
-        return loader
-
-
-class ProcessorLoader(ComponentLoader):
-    """Loader for HF processors that pair with vision-language encoders."""
-
-    def load(self, model_path: str, fastvideo_args: FastVideoArgs):
-        from transformers import AutoProcessor
-
-        logger.info("Loading processor from %s", model_path)
-        processor = AutoProcessor.from_pretrained(
-            model_path,
-            trust_remote_code=fastvideo_args.trust_remote_code,
-        )
-        logger.info("Loaded processor: %s", processor.__class__.__name__)
-        return processor
-
-
 class ImageProcessorLoader(ComponentLoader):
     """Loader for image processor."""
 

fastvideo/models/registry.py

@@ -61,11 +61,6 @@ _IMAGE_TO_VIDEO_DIT_MODELS = {
     "MatrixGame3WanModel": ("dits", "matrixgame3", "MatrixGame3WanModel"),
 }
 
-# Text-to-image DiT models (2D image generation)
-_TEXT_TO_IMAGE_DIT_MODELS = {
-    "GlmImageTransformer2DModel": ("dits", "glm_image", "GlmImageTransformer2DModel"),
-}
-
 _TEXT_ENCODER_MODELS = {
     "CLIPTextModel": ("encoders", "clip", "CLIPTextModel"),
     "CLIPTextModelWithProjection":
@@ -139,7 +134,6 @@ _UPSAMPLERS = {
 _LEGACY_FAST_VIDEO_MODELS = {
     **_TEXT_TO_VIDEO_DIT_MODELS,
     **_IMAGE_TO_VIDEO_DIT_MODELS,
-    **_TEXT_TO_IMAGE_DIT_MODELS,
     **_TEXT_ENCODER_MODELS,
     **_IMAGE_ENCODER_MODELS,
     **_VAE_MODELS,

fastvideo/models/schedulers/scheduling_flow_map_euler_discrete.py

@@ -1,202 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Flow-map any-step Euler scheduler for AnyFlow.
-
-The model predicts the *average* velocity ``u_θ(x_t, t, r)`` from time
-``t`` back to time ``r``, so one Euler step is
-
-    x_r = x_t - ((t - r) / num_train_timesteps) * u_θ(x_t, t, r)
-
-regardless of how far apart ``t`` and ``r`` are. The scheduler also
-provides the AnyFlow training-time helpers ``apply_shift`` (flow-matching
-shift transform) and ``get_train_weight`` (per-timestep loss weight,
-including ``beta08``).
-
-Standalone — does not depend on diffusers' ConfigMixin/SchedulerMixin.
-"""
-
-from __future__ import annotations
-
-from typing import Literal
-
-import torch
-
-from fastvideo.models.schedulers.base import BaseScheduler
-
-
-WeightType = Literal["uniform", "gaussian", "beta08"]
-
-
-class FlowMapEulerDiscreteScheduler(BaseScheduler):
-    """Minimal flow-map scheduler.
-
-    Parameters
-    ----------
-    num_train_timesteps:
-        Discretization granularity for training. ``t`` is expressed in
-        absolute units in ``[0, num_train_timesteps]``.
-    shift:
-        Flow-matching shift parameter (Wan video default: ``5.0``). Set
-        to ``1.0`` for an identity shift.
-    """
-
-    order: int = 1
-
-    def __init__(
-        self,
-        *,
-        num_train_timesteps: int = 1000,
-        shift: float = 1.0,
-    ) -> None:
-        self.num_train_timesteps = int(num_train_timesteps)
-        self.shift = float(shift)
-        self.timesteps: torch.Tensor = torch.empty(0)
-        self.sigmas: torch.Tensor = torch.empty(0)
-        super().__init__()
-
-    # ------------------------------------------------------------------
-    # BaseScheduler abstract surface
-
-    def set_shift(self, shift: float) -> None:
-        self.shift = float(shift)
-
-    def scale_model_input(
-        self,
-        sample: torch.Tensor,
-        timestep: int | None = None,
-    ) -> torch.Tensor:
-        # Flow-matching has no per-step input scaling; pass through.
-        del timestep
-        return sample
-
-    # ------------------------------------------------------------------
-    # Public helpers used by the AnyFlow pretrain method.
-
-    def apply_shift(
-        self,
-        t: torch.Tensor,
-        *,
-        shift: float | None = None,
-    ) -> torch.Tensor:
-        """Apply the flow-matching shift: ``t' = s * t / (1 + (s - 1) * t)``.
-
-        Operates in the normalized ``[0, 1]`` domain — callers should pass
-        ``t / num_train_timesteps`` (or sample ``t`` directly from
-        ``[0, 1]``).
-        """
-        s = self.shift if shift is None else float(shift)
-        if s == 1.0:
-            return t
-        return s * t / (1.0 + (s - 1.0) * t)
-
-    def get_train_weight(
-        self,
-        t: torch.Tensor,
-        *,
-        weight_type: WeightType = "beta08",
-    ) -> torch.Tensor:
-        """Per-timestep training weight, renormalized so the total weight
-        mass equals ``num_train_timesteps`` (matching AnyFlow reference's
-        ``scheduling_flowmap_euler_discrete.py``).
-
-        ``beta08``: ``w(t) = t * sqrt(1 - t)`` (in normalized t-space).
-        """
-        # Auto-detect domain: if t was given in absolute units, normalize.
-        t_f = t.float()
-        max_val = t_f.max() if t_f.numel() > 0 else torch.tensor(0.0)
-        if max_val > 1.0 + 1e-6:
-            t_norm = t_f / self.num_train_timesteps
-        else:
-            t_norm = t_f
-        t_norm = t_norm.clamp(min=0.0, max=1.0)
-
-        if weight_type == "uniform":
-            w = torch.ones_like(t_norm)
-        elif weight_type == "gaussian":
-            w = torch.exp(-0.5 * ((t_norm - 0.5) / 0.2) ** 2)
-        elif weight_type == "beta08":
-            w = t_norm.pow(1.0) * (1.0 - t_norm).clamp_min(0.0).pow(0.5)
-        else:
-            raise ValueError(f"Unknown weight_type: {weight_type!r}")
-
-        denom = w.sum().clamp_min(1e-8)
-        return w * (float(self.num_train_timesteps) / denom)
-
-    # ------------------------------------------------------------------
-
-    def set_timesteps(
-        self,
-        *,
-        num_inference_steps: int,
-        device: torch.device | str = "cpu",
-        custom_timesteps: list[float] | torch.Tensor | None = None,
-    ) -> None:
-        """Build a descending timestep schedule ending at 0.
-
-        With ``num_inference_steps=N`` the schedule has ``N + 1`` entries
-        ``[T_max, ..., 0]`` so a rollout consumes ``N`` Euler steps.
-
-        ``custom_timesteps`` overrides the linspace+shift schedule with
-        a pinned list (in absolute train-timestep units), useful for the
-        AnyFlow paper's hand-tuned ``[999, 937, 833, 624, 0]`` schedule.
-        """
-        if num_inference_steps <= 0:
-            raise ValueError(
-                "num_inference_steps must be positive, "
-                f"got {num_inference_steps}")
-        device = torch.device(device)
-
-        if custom_timesteps is not None:
-            ts = torch.as_tensor(
-                custom_timesteps, dtype=torch.float32, device=device)
-            if ts.ndim != 1:
-                raise ValueError(
-                    "custom_timesteps must be 1-D, got shape "
-                    f"{tuple(ts.shape)}")
-            if not torch.all(ts[:-1] >= ts[1:]):
-                raise ValueError(
-                    "custom_timesteps must be descending (largest first)")
-        else:
-            ts_norm = torch.linspace(
-                1.0, 0.0, num_inference_steps + 1, device=device)
-            ts_norm = self.apply_shift(ts_norm)
-            ts = ts_norm * self.num_train_timesteps
-
-        self.timesteps = ts
-        self.sigmas = ts / self.num_train_timesteps
-
-    def step(
-        self,
-        model_output: torch.Tensor,
-        *,
-        sample: torch.Tensor,
-        timestep: torch.Tensor,
-        r_timestep: torch.Tensor,
-    ) -> torch.Tensor:
-        """One Euler step from ``t`` to ``r``.
-
-        ``model_output`` is the average-velocity prediction
-        ``u_θ(x_t, t, r)``. Both ``timestep`` and ``r_timestep`` are in
-        absolute train-timestep units (``[0, num_train_timesteps]``).
-        """
-        t = timestep.to(sample.device, dtype=sample.dtype)
-        r = r_timestep.to(sample.device, dtype=sample.dtype)
-        dt_norm = (t - r) / float(self.num_train_timesteps)
-        # Broadcast dt over channel/spatial dims.
-        view: list[int] = [-1] + [1] * (sample.ndim - 1)
-        return sample - dt_norm.view(*view) * model_output
-
-    def add_noise(
-        self,
-        original_samples: torch.Tensor,
-        noise: torch.Tensor,
-        timestep: torch.Tensor,
-    ) -> torch.Tensor:
-        """Linear flow-matching interpolation: ``x_t = (1 - σ) * x_0 + σ * ε``,
-        where ``σ = t / num_train_timesteps``.
-        """
-        sigma = (timestep.to(original_samples.device,
-                             dtype=original_samples.dtype)
-                 / float(self.num_train_timesteps))
-        view: list[int] = [-1] + [1] * (original_samples.ndim - 1)
-        sigma = sigma.view(*view)
-        return (1.0 - sigma) * original_samples + sigma * noise

fastvideo/pipelines/basic/flux/__init__.py

@@ -1 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0

fastvideo/pipelines/basic/flux/flux_pipeline.py

@@ -1,74 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-
-from __future__ import annotations
-
-from fastvideo.fastvideo_args import FastVideoArgs
-from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase
-from fastvideo.pipelines.stages.flux_stages import (
-    FluxConditioningStage,
-    FluxDecodingStage,
-    FluxDenoisingStage,
-    FluxInputValidationStage,
-    FluxLatentPreparationStage,
-    FluxTimestepPreparationStage,
-)
-from fastvideo.pipelines.stages.text_encoding import TextEncodingStage
-
-
-class FluxPipeline(ComposedPipelineBase):
-    """FLUX.1-dev T2I (Diffusers module layout, packed latents, embedded guidance)."""
-
-    _required_config_modules = [
-        "scheduler",
-        "transformer",
-        "vae",
-        "text_encoder",
-        "text_encoder_2",
-        "tokenizer",
-        "tokenizer_2",
-    ]
-
-    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None:
-        self.add_stage(stage_name="input_validation_stage", stage=FluxInputValidationStage())
-
-        self.add_stage(
-            stage_name="text_encoding_stage",
-            stage=TextEncodingStage(
-                text_encoders=[
-                    self.get_module("text_encoder"),
-                    self.get_module("text_encoder_2"),
-                ],
-                tokenizers=[
-                    self.get_module("tokenizer"),
-                    self.get_module("tokenizer_2"),
-                ],
-            ),
-        )
-
-        self.add_stage(stage_name="flux_conditioning_stage", stage=FluxConditioningStage())
-
-        self.add_stage(
-            stage_name="timestep_preparation_stage",
-            stage=FluxTimestepPreparationStage(scheduler=self.get_module("scheduler")),
-        )
-
-        self.add_stage(
-            stage_name="latent_preparation_stage",
-            stage=FluxLatentPreparationStage(scheduler=self.get_module("scheduler")),
-        )
-
-        self.add_stage(
-            stage_name="denoising_stage",
-            stage=FluxDenoisingStage(
-                transformer=self.get_module("transformer"),
-                scheduler=self.get_module("scheduler"),
-            ),
-        )
-
-        self.add_stage(
-            stage_name="decoding_stage",
-            stage=FluxDecodingStage(vae=self.get_module("vae")),
-        )
-
-
-EntryClass = FluxPipeline

fastvideo/pipelines/basic/glm_image/__init__.py

@@ -1,7 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""GLM-Image pipeline package."""
-
-from fastvideo.pipelines.basic.glm_image.glm_image_pipeline import (
-    GlmImagePipeline, )
-
-__all__ = ["GlmImagePipeline"]

fastvideo/pipelines/basic/glm_image/glm_image_pipeline.py

@@ -1,82 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-from fastvideo.fastvideo_args import FastVideoArgs
-from fastvideo.logger import init_logger
-from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import (
-    FlowMatchEulerDiscreteScheduler, )
-from fastvideo.pipelines import ComposedPipelineBase, LoRAPipeline
-from fastvideo.pipelines.basic.glm_image.stages import (
-    GlmImageBeforeDenoisingStage,
-    GlmImageConditionEncodingStage,
-    GlmImageDecodingStage,
-    GlmImageDenoisingStage,
-)
-from fastvideo.pipelines.stages import InputValidationStage
-
-logger = init_logger(__name__)
-
-
-class GlmImagePipeline(LoRAPipeline, ComposedPipelineBase):
-    pipeline_name = "GlmImagePipeline"
-
-    _required_config_modules = [
-        "text_encoder",
-        "tokenizer",
-        "vae",
-        "transformer",
-        "scheduler",
-        "vision_language_encoder",
-        "processor",
-    ]
-
-    _optional_config_modules: list[str] = []
-
-    def initialize_pipeline(self, fastvideo_args: FastVideoArgs) -> None:
-        self.modules["scheduler"] = FlowMatchEulerDiscreteScheduler(shift=1.0)
-
-    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None:
-        self.add_stage(
-            stage_name="input_validation_stage",
-            stage=InputValidationStage(),
-        )
-
-        self.add_stage(
-            stage_name="glm_image_before_denoising_stage",
-            stage=GlmImageBeforeDenoisingStage(
-                vae=self.get_module("vae"),
-                text_encoder=self.get_module("text_encoder"),
-                tokenizer=self.get_module("tokenizer"),
-                processor=self.get_module("processor"),
-                transformer=self.get_module("transformer"),
-                scheduler=self.get_module("scheduler"),
-                vision_language_encoder=self.get_module("vision_language_encoder"),
-            ),
-        )
-
-        self.add_stage(
-            stage_name="glm_image_condition_encoding_stage",
-            stage=GlmImageConditionEncodingStage(
-                vae=self.get_module("vae"),
-                transformer=self.get_module("transformer"),
-            ),
-        )
-
-        self.add_stage(
-            stage_name="denoising_stage",
-            stage=GlmImageDenoisingStage(
-                transformer=self.get_module("transformer"),
-                scheduler=self.get_module("scheduler"),
-                vae=self.get_module("vae"),
-                pipeline=self,
-            ),
-        )
-
-        self.add_stage(
-            stage_name="decoding_stage",
-            stage=GlmImageDecodingStage(
-                vae=self.get_module("vae"),
-                pipeline=self,
-            ),
-        )
-
-
-EntryClass = GlmImagePipeline

fastvideo/pipelines/basic/glm_image/stages/__init__.py

@@ -1,14 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""GLM-Image pipeline stages."""
-
-from fastvideo.pipelines.basic.glm_image.stages.before_denoising import (GlmImageBeforeDenoisingStage)
-from fastvideo.pipelines.basic.glm_image.stages.condition_encoding import (GlmImageConditionEncodingStage)
-from fastvideo.pipelines.basic.glm_image.stages.decoding import (GlmImageDecodingStage)
-from fastvideo.pipelines.basic.glm_image.stages.denoising import (GlmImageDenoisingStage)
-
-__all__ = [
-    "GlmImageBeforeDenoisingStage",
-    "GlmImageConditionEncodingStage",
-    "GlmImageDecodingStage",
-    "GlmImageDenoisingStage",
-]

fastvideo/pipelines/basic/glm_image/stages/before_denoising.py

@@ -1,269 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-from __future__ import annotations
-
-import re
-from math import sqrt
-
-import numpy as np
-import torch
-
-from fastvideo.distributed import get_local_torch_device
-from fastvideo.fastvideo_args import FastVideoArgs
-from fastvideo.logger import init_logger
-from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
-from fastvideo.pipelines.stages.base import PipelineStage
-from fastvideo.pipelines.stages.validators import StageValidators as V
-from fastvideo.pipelines.stages.validators import VerificationResult
-
-logger = init_logger(__name__)
-
-
-def calculate_shift(
-    image_seq_len: int,
-    base_seq_len: int = 256,
-    base_shift: float = 0.25,
-    max_shift: float = 0.75,
-) -> float:
-    return (image_seq_len / base_seq_len)**0.5 * max_shift + base_shift
-
-
-def get_glyph_texts(prompt: str | list[str]) -> list[str] | list[list[str]]:
-    if isinstance(prompt, str):
-        prompts: list[str] = [prompt]
-        is_batch = False
-    else:
-        prompts = prompt
-        is_batch = True
-    out: list[list[str]] = []
-    for p in prompts:
-        out.append(
-            re.findall(r"'([^']*)'", p) + re.findall(r"“([^“”]*)”", p) + re.findall(r'"([^"]*)"', p) +
-            re.findall(r"「([^「」]*)」", p))
-    return out if is_batch else out[0]
-
-
-def compute_glyph_embeds(
-    prompts: list[str],
-    tokenizer,
-    text_encoder,
-    device: torch.device,
-    dtype: torch.dtype,
-    max_sequence_length: int = 2048,
-) -> torch.Tensor:
-    all_glyph_texts = get_glyph_texts(prompts)
-    all_glyph_embeds = []
-    for glyph_texts in all_glyph_texts:
-        if len(glyph_texts) == 0:
-            glyph_texts = [""]
-        input_ids = tokenizer(
-            glyph_texts,
-            max_length=max_sequence_length,
-            truncation=True,
-        ).input_ids
-        input_ids = [[tokenizer.pad_token_id] * ((len(input_ids) + 1) % 2) + ids for ids in input_ids]
-        max_length = max(len(ids) for ids in input_ids)
-        attention_mask = torch.tensor(
-            [[1] * len(ids) + [0] * (max_length - len(ids)) for ids in input_ids],
-            device=device,
-        )
-        input_ids_t = torch.tensor(
-            [ids + [tokenizer.pad_token_id] * (max_length - len(ids)) for ids in input_ids],
-            device=device,
-        )
-        outputs = text_encoder(input_ids_t, attention_mask=attention_mask)
-        glyph_embeds = outputs.last_hidden_state[attention_mask.bool()].unsqueeze(0)
-        all_glyph_embeds.append(glyph_embeds)
-
-    max_seq_len = max(emb.size(1) for emb in all_glyph_embeds)
-    padded = []
-    for emb in all_glyph_embeds:
-        if emb.size(1) < max_seq_len:
-            pad = torch.zeros(emb.size(0), max_seq_len - emb.size(1), emb.size(2), device=device, dtype=emb.dtype)
-            emb = torch.cat([pad, emb], dim=1)
-        padded.append(emb)
-    return torch.cat(padded, dim=0).to(device=device, dtype=dtype)
-
-
-def _grid_dims(height: int, width: int) -> tuple[int, int, int, int]:
-    th, tw = height // 32, width // 32
-    ratio = th / tw
-    pth = int(sqrt(ratio) * 16)
-    ptw = int(sqrt(1 / ratio) * 16)
-    return th, tw, pth, ptw
-
-
-def _upsample_d32_to_d16(tokens: torch.Tensor, th: int, tw: int) -> torch.Tensor:
-    tokens = tokens.view(1, 1, th, tw).float()
-    tokens = torch.nn.functional.interpolate(tokens, scale_factor=2, mode="nearest").long()
-    return tokens.view(1, -1)
-
-
-class GlmImageBeforeDenoisingStage(PipelineStage):
-
-    def __init__(self,
-                 vae,
-                 text_encoder,
-                 tokenizer,
-                 processor,
-                 transformer,
-                 scheduler,
-                 vision_language_encoder=None) -> None:
-        super().__init__()
-        self.vae = vae
-        self.text_encoder = text_encoder
-        self.tokenizer = tokenizer
-        self.processor = processor
-        self.transformer = transformer
-        self.scheduler = scheduler
-        if isinstance(vision_language_encoder, tuple):
-            self.vision_language_encoder, self.vl_processor = (vision_language_encoder[0], vision_language_encoder[1]
-                                                               or processor)
-        else:
-            self.vision_language_encoder = vision_language_encoder
-            self.vl_processor = processor
-
-    @torch.no_grad()
-    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
-        device = get_local_torch_device()
-        dtype = torch.bfloat16
-        th, tw, pth, ptw = _grid_dims(batch.height, batch.width)
-
-        if batch.seed is not None:
-            torch.manual_seed(batch.seed)
-            if torch.cuda.is_available():
-                torch.cuda.manual_seed(batch.seed)
-
-        # 1-3. AR token generation. I2I prepends the condition image and uses a
-        # single-scale target grid; T2I is multi-scale.
-        is_t2i = batch.pil_image is None
-        if self.vision_language_encoder is not None:
-            content = [{"type": "text", "text": batch.prompt}]
-            if not is_t2i:
-                content.insert(0, {"type": "image", "image": batch.pil_image})
-            messages = [{"role": "user", "content": content}]
-            inputs = self.vl_processor.apply_chat_template(messages,
-                                                           tokenize=True,
-                                                           target_h=batch.height,
-                                                           target_w=batch.width,
-                                                           return_dict=True,
-                                                           return_tensors="pt").to(device)
-
-            if is_t2i:
-                up_h, up_w = th, tw
-                large_start, large_count = pth * ptw, th * tw
-                max_new = large_count + (pth * ptw) + 1
-            else:
-                # Condition grid(s) first, target grid last.
-                _, t_h, t_w = inputs["image_grid_thw"][-1].tolist()
-                up_h, up_w = int(t_h), int(t_w)
-                large_start, large_count = 0, up_h * up_w
-                max_new = large_count + 1
-
-            outputs = self.vision_language_encoder.generate(**inputs, max_new_tokens=max_new, do_sample=True)
-            gen_tokens = outputs[0][inputs.input_ids.shape[-1]:]
-            if gen_tokens.shape[0] >= large_start + large_count:
-                large_tokens = gen_tokens[large_start:large_start + large_count]
-            else:
-                available = gen_tokens[large_start:]
-                large_tokens = torch.zeros(large_count, dtype=gen_tokens.dtype, device=gen_tokens.device)
-                if available.shape[0] > 0:
-                    large_tokens[:min(available.shape[0], large_count)] = available[:large_count]
-                logger.warning("AR generated %d tokens, expected %d. Padding with zeros.", gen_tokens.shape[0],
-                               large_start + large_count)
-            batch.prior_token_id = _upsample_d32_to_d16(large_tokens, up_h, up_w)
-            batch.prior_token_drop = torch.zeros(batch.prior_token_id.shape, dtype=torch.bool, device=device)
-
-            if not is_t2i:
-                self._compute_source_prior_tokens(batch, inputs)
-        else:
-            num_prior_tokens = 4 * th * tw
-            logger.warning("No vision_language_encoder provided; using random dropped priors.")
-            batch.prior_token_id = torch.randint(0, 16384, (1, num_prior_tokens), device=device)
-            batch.prior_token_drop = torch.ones(batch.prior_token_id.shape, dtype=torch.bool, device=device)
-
-        # 4. Glyph T5 encoding.
-        prompts = [batch.prompt] if isinstance(batch.prompt, str) else list(batch.prompt)
-        prompt_embeds = compute_glyph_embeds(prompts, self.tokenizer, self.text_encoder, device, dtype)
-
-        # 5. CFG-side negative encoding.
-        if batch.do_classifier_free_guidance:
-            neg_prompts = [batch.negative_prompt or ""] * len(prompts)
-            neg_embeds = compute_glyph_embeds(neg_prompts, self.tokenizer, self.text_encoder, device, dtype)
-            L_pos, L_neg = prompt_embeds.shape[1], neg_embeds.shape[1]
-            max_L = max(L_pos, L_neg)
-            if L_pos < max_L:
-                pad = torch.zeros(prompt_embeds.shape[0],
-                                  max_L - L_pos,
-                                  prompt_embeds.shape[2],
-                                  device=device,
-                                  dtype=dtype)
-                prompt_embeds = torch.cat([pad, prompt_embeds], dim=1)
-            if L_neg < max_L:
-                pad = torch.zeros(neg_embeds.shape[0], max_L - L_neg, neg_embeds.shape[2], device=device, dtype=dtype)
-                neg_embeds = torch.cat([pad, neg_embeds], dim=1)
-            # Row 0 conditional (positive), row 1 unconditional (negative).
-            prompt_embeds = torch.cat([prompt_embeds, neg_embeds], dim=0)
-            att_pos = torch.ones((1, max_L), device=device)
-            att_neg = torch.ones((1, max_L), device=device)
-            if L_pos < max_L:
-                att_pos[:, :max_L - L_pos] = 0
-            if L_neg < max_L:
-                att_neg[:, :max_L - L_neg] = 0
-            attention_mask = torch.cat([att_pos, att_neg], dim=0)
-        else:
-            attention_mask = torch.ones((1, prompt_embeds.shape[1]), device=device)
-
-        batch.prompt_embeds = [prompt_embeds]
-        batch.attention_mask = attention_mask
-
-        # 6. Latents + dynamic flow shift.
-        if batch.seed is not None:
-            torch.manual_seed(batch.seed)
-            if torch.cuda.is_available():
-                torch.cuda.manual_seed(batch.seed)
-        batch.latents = torch.randn((1, 16, 1, batch.height // 8, batch.width // 8), device=device, dtype=dtype)
-
-        # Integer-cast linspace timesteps with resolution-dependent shift applied to
-        # sigmas only; the DiT is conditioned on the unshifted integer timesteps.
-        ntt = self.scheduler.config.num_train_timesteps
-        patch_size = self.transformer.patch_size
-        image_seq_len = ((batch.height // 8) * (batch.width // 8)) // (patch_size**2)
-        sched_timesteps = np.linspace(ntt, 1.0, batch.num_inference_steps + 1)[:-1].astype(np.int64).astype(np.float32)
-        sched_sigmas = sched_timesteps / ntt
-        self.scheduler.set_shift(calculate_shift(image_seq_len))
-        self.scheduler.set_timesteps(batch.num_inference_steps,
-                                     device=device,
-                                     sigmas=sched_sigmas.tolist(),
-                                     timesteps=sched_timesteps.tolist())
-        batch.timesteps = self.scheduler.timesteps
-        return batch
-
-    @torch.no_grad()
-    def _compute_source_prior_tokens(self, batch: ForwardBatch, inputs) -> None:
-        image_grid_thw = inputs["image_grid_thw"]
-        num_condition_images = image_grid_thw.shape[0] - 1
-        source_grids = image_grid_thw[:num_condition_images]
-        image_features = self.vision_language_encoder.get_image_features(inputs["pixel_values"], source_grids)
-        image_feature_parts = getattr(image_features, "pooler_output", image_features)
-        embed = torch.cat(image_feature_parts, dim=0)
-        src_ids_d32 = self.vision_language_encoder.get_image_tokens(embed, source_grids)
-        split_sizes = source_grids.prod(dim=-1).tolist()
-        upsampled = [
-            _upsample_d32_to_d16(ids, int(grid[1]), int(grid[2])).squeeze(0)
-            for ids, grid in zip(torch.split(src_ids_d32, split_sizes), source_grids, strict=False)
-        ]
-        src_grids_up = source_grids.clone()
-        src_grids_up[:, 1] *= 2
-        src_grids_up[:, 2] *= 2
-        batch.extra["glm_prior_token_image_ids"] = torch.cat(upsampled, dim=0)
-        batch.extra["glm_source_image_grid_thw"] = src_grids_up
-
-    def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult:
-        result = VerificationResult()
-        result.add_check("prompt", batch.prompt, V.string_not_empty)
-        return result
-
-    def verify_output(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult:
-        result = VerificationResult()
-        result.add_check("prior_token_id", batch.prior_token_id, V.is_tensor)
-        return result

fastvideo/pipelines/basic/glm_image/stages/condition_encoding.py

@@ -1,67 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-from __future__ import annotations
-
-import torch
-
-from fastvideo.distributed import get_local_torch_device
-from fastvideo.fastvideo_args import FastVideoArgs
-from fastvideo.forward_context import set_forward_context
-from fastvideo.image_processor import ImageProcessor
-from fastvideo.models.dits.glm_image import GlmImageKVCache
-from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
-from fastvideo.pipelines.stages.base import PipelineStage
-
-_CONDITION_MULTIPLE_OF = 16  # vae_scale_factor (8) * DiT patch_size (2)
-
-
-class GlmImageConditionEncodingStage(PipelineStage):
-
-    def __init__(self, vae, transformer) -> None:
-        super().__init__()
-        self.vae = vae
-        self.transformer = transformer
-        self.image_processor = ImageProcessor(vae_scale_factor=_CONDITION_MULTIPLE_OF)
-
-    @torch.no_grad()
-    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
-        if batch.pil_image is None:
-            return batch
-
-        device = get_local_torch_device()
-        dtype = torch.bfloat16
-        self.vae.to(device)
-
-        prior_ids = batch.extra["glm_prior_token_image_ids"].to(device)
-        if prior_ids.dim() == 1:
-            prior_ids = prior_ids.unsqueeze(0)
-
-        # Latent patch count must match the source prior tokens; mismatch is fatal.
-        src_grid = batch.extra["glm_source_image_grid_thw"][0]
-        cond_h = int(src_grid[1]) * _CONDITION_MULTIPLE_OF
-        cond_w = int(src_grid[2]) * _CONDITION_MULTIPLE_OF
-        cond_img = self.image_processor.preprocess(batch.pil_image, cond_h, cond_w).to(device=device,
-                                                                                       dtype=torch.float32)
-        latent = self.vae.encode(cond_img).latent_dist.mode()
-        # NOTE: at runtime self.vae.config is a diffusers FrozenDict with flat
-        # latents_mean/latents_std fields. Access the flat fields directly.
-        cfg = self.vae.config
-        mean = torch.tensor(cfg.latents_mean, device=device, dtype=torch.float32).view(1, -1, 1, 1)
-        std = torch.tensor(cfg.latents_std, device=device, dtype=torch.float32).view(1, -1, 1, 1)
-        latent = ((latent - mean) / std).to(dtype)
-
-        kv_caches = GlmImageKVCache(num_layers=self.transformer.num_layers)
-        empty_text = batch.prompt_embeds[0][:1, :0, :].to(device=device, dtype=dtype)
-        with set_forward_context(current_timestep=0, attn_metadata=None, forward_batch=batch):
-            self.transformer(
-                hidden_states=latent,
-                encoder_hidden_states=empty_text,
-                prior_token_id=prior_ids,
-                prior_token_drop=torch.zeros((prior_ids.shape[0], ), dtype=torch.bool, device=device),
-                timestep=torch.zeros((1, ), device=device),
-                target_size=torch.tensor([tuple(cond_img.shape[-2:])], device=device, dtype=torch.long),
-                crop_coords=torch.zeros((1, 2), device=device, dtype=torch.long),
-                kv_caches=kv_caches,
-                kv_caches_mode="write",
-            )
-        batch.extra["glm_kv_caches"] = kv_caches
-        return batch

fastvideo/pipelines/basic/glm_image/stages/decoding.py

@@ -1,35 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-from __future__ import annotations
-
-import torch
-
-from fastvideo.distributed import get_local_torch_device
-from fastvideo.fastvideo_args import FastVideoArgs
-from fastvideo.pipelines.stages.decoding import DecodingStage
-from fastvideo.utils import PRECISION_TO_TYPE
-
-
-class GlmImageDecodingStage(DecodingStage):
-
-    @torch.no_grad()
-    def decode(self, latents: torch.Tensor, fastvideo_args: FastVideoArgs) -> torch.Tensor:
-        self.vae.to(get_local_torch_device())
-        latents = latents.to(get_local_torch_device())
-
-        vae_dtype = PRECISION_TO_TYPE[fastvideo_args.pipeline_config.vae_precision]
-        vae_autocast = (vae_dtype != torch.float32 and not fastvideo_args.disable_autocast)
-
-        latents = self._denormalize_latents(latents)
-        if latents.dim() == 5:
-            latents = latents.squeeze(2)
-
-        with torch.autocast(device_type="cuda", dtype=vae_dtype, enabled=vae_autocast):
-            if fastvideo_args.pipeline_config.vae_tiling:
-                self.vae.enable_tiling()
-            if not vae_autocast:
-                latents = latents.to(vae_dtype)
-            decoded = self.vae.decode(latents)
-
-        image = decoded.sample if hasattr(decoded, "sample") else decoded
-        image = (image / 2 + 0.5).clamp(0, 1)
-        return image.unsqueeze(2)

fastvideo/pipelines/basic/glm_image/stages/denoising.py

@@ -1,192 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""CFG convention (both denoise paths): row 0 conditional (positive), row 1 unconditional."""
-from __future__ import annotations
-
-import torch
-
-from fastvideo.attention.backends.sdpa import SDPAMetadata
-from fastvideo.distributed import get_local_torch_device
-from fastvideo.fastvideo_args import FastVideoArgs
-from fastvideo.forward_context import set_forward_context
-from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
-from fastvideo.pipelines.stages.denoising import DenoisingStage
-from fastvideo.pipelines.stages.validators import StageValidators as V
-from fastvideo.pipelines.stages.validators import VerificationResult
-from fastvideo.platforms import AttentionBackendEnum
-
-
-class GlmImageDenoisingStage(DenoisingStage):
-
-    def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult:
-        result = VerificationResult()
-        result.add_check("timesteps", batch.timesteps, [V.is_tensor, V.min_dims(1)])
-        latents = getattr(batch, "latent", getattr(batch, "latents", None))
-        result.add_check("latents", latents, [V.is_tensor, V.with_dims(5)])
-        result.add_check("num_inference_steps", batch.num_inference_steps, V.positive_int)
-        result.add_check("prompt_embeds", batch.prompt_embeds, V.list_not_empty)
-        return result
-
-    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
-        device = get_local_torch_device()
-        dtype = torch.bfloat16
-        guidance_scale = batch.guidance_scale
-        do_cfg = guidance_scale > 1.0
-
-        latents = getattr(batch, "latent", getattr(batch, "latents", None))
-        if latents is None:
-            raise ValueError("No latents found in batch.")
-        if latents.dim() == 5:
-            latents = latents.squeeze(2)
-
-        prompt_embeds = batch.prompt_embeds[0]
-        text_attention_mask = getattr(batch, "attention_mask", None)
-        timesteps = batch.timesteps
-
-        patch_size = self.transformer.patch_size
-        _, _, h, w = latents.shape
-        image_seq_length = (h // patch_size) * (w // patch_size)
-        text_seq_length = prompt_embeds.shape[1] if prompt_embeds.dim() >= 2 else 0
-
-        first_block = self.transformer.transformer_blocks[0]
-        backend = getattr(first_block.attn1.attn, "backend", None)
-        sdpa = backend == AttentionBackendEnum.TORCH_SDPA and text_attention_mask is not None
-
-        kv_caches = batch.extra.get("glm_kv_caches")
-        if kv_caches is None:
-            self._denoise_t2i(batch, latents, prompt_embeds, text_attention_mask, timesteps, do_cfg, guidance_scale,
-                              text_seq_length, image_seq_length, sdpa, device, dtype)
-        else:
-            self._denoise_i2i(batch, latents, prompt_embeds, text_attention_mask, timesteps, do_cfg, guidance_scale,
-                              text_seq_length, image_seq_length, sdpa, kv_caches, device, dtype)
-        return batch
-
-    def _denoise_t2i(self, batch, latents, prompt_embeds, text_attention_mask, timesteps, do_cfg, guidance_scale,
-                     text_seq_length, image_seq_length, sdpa, device, dtype) -> None:
-        num_inference_steps = batch.num_inference_steps
-        bs = 2 if do_cfg else 1
-        target_size = torch.tensor([[batch.height, batch.width]], device=device, dtype=torch.long).repeat(bs, 1)
-        crop_coords = torch.zeros((bs, 2), device=device, dtype=torch.long)
-
-        prior_token_id = batch.prior_token_id
-        if do_cfg and prior_token_id.shape[0] == 1:
-            prior_token_id = prior_token_id.repeat(2, 1)
-        if do_cfg:
-            prior_token_drop = torch.tensor([False, True], device=device)
-        else:
-            prior_token_drop = getattr(batch, "prior_token_drop", torch.tensor([False], device=device))
-
-        attention_mask_kv = None
-        if sdpa:
-            if (text_attention_mask.shape[0] == 1 and bs > 1):
-                text_attention_mask = text_attention_mask.repeat(bs, 1)
-            mix_attn_mask = torch.ones((bs, text_seq_length + image_seq_length), device=device, dtype=torch.float32)
-            mix_attn_mask[:, :text_seq_length] = (text_attention_mask.float().to(device))
-            attention_mask_kv = (mix_attn_mask > 0).unsqueeze(1).unsqueeze(2)
-
-        with self.progress_bar(total=num_inference_steps) as progress_bar:
-            for i, t in enumerate(timesteps):
-                latent_model_input = torch.cat([latents] * 2) if do_cfg else latents
-                latent_model_input = self.scheduler.scale_model_input(latent_model_input, t).to(dtype)
-                t_expand = t.expand(latent_model_input.shape[0]) - 1
-
-                attn_metadata = (SDPAMetadata(current_timestep=i, attn_mask=attention_mask_kv)
-                                 if attention_mask_kv is not None else None)
-
-                with torch.no_grad(), set_forward_context(current_timestep=i,
-                                                          attn_metadata=attn_metadata,
-                                                          forward_batch=batch):
-                    noise_pred = self.transformer(
-                        latent_model_input,
-                        prompt_embeds,
-                        prior_token_id,
-                        prior_token_drop,
-                        t_expand,
-                        target_size,
-                        crop_coords,
-                    )
-
-                    if do_cfg:
-                        noise_pred_cond, noise_pred_uncond = noise_pred.float().chunk(2)
-                        noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_cond - noise_pred_uncond)
-
-                        guidance_rescale = getattr(batch, "guidance_rescale", 0.0)
-                        if guidance_rescale > 0.0:
-                            dims = list(range(1, noise_pred_cond.ndim))
-                            std_text = noise_pred_cond.std(dim=dims, keepdim=True)
-                            std_cfg = noise_pred.std(dim=dims, keepdim=True)
-                            rescaled = noise_pred * (std_text / std_cfg)
-                            noise_pred = (guidance_rescale * rescaled + (1 - guidance_rescale) * noise_pred)
-
-                    latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
-                progress_bar.update()
-
-        batch.latents = latents.unsqueeze(2)
-
-    def _denoise_i2i(self, batch, latents, prompt_embeds, text_attention_mask, timesteps, do_cfg, guidance_scale,
-                     text_seq_length, image_seq_length, sdpa, kv_caches, device, dtype) -> None:
-        """Two separate transformer calls (cond reads the cache, uncond skips it):
-        the cache mode is one global flag with batch-1 k/v, so a 2-row CFG call
-        cannot express both."""
-        num_inference_steps = batch.num_inference_steps
-        target_size = torch.tensor([[batch.height, batch.width]], device=device, dtype=torch.long)
-        crop_coords = torch.zeros((1, 2), device=device, dtype=torch.long)
-        prior_token_id = batch.prior_token_id[:1]
-        drop_keep = torch.zeros((1, ), dtype=torch.bool, device=device)
-        drop_all = torch.ones((1, ), dtype=torch.bool, device=device)
-        cache_len = kv_caches[0].k_cache.shape[1] if kv_caches[0].k_cache is not None else 0
-
-        def _mask(row: int, with_cache: bool):
-            if not sdpa:
-                return None
-            prefix = cache_len if with_cache else 0
-            m = torch.ones((1, prefix + text_seq_length + image_seq_length), device=device, dtype=torch.float32)
-            m[:, prefix:prefix + text_seq_length] = text_attention_mask[row:row + 1].float().to(device)
-            return (m > 0).unsqueeze(1).unsqueeze(2)
-
-        cond_mask = _mask(0, with_cache=True)
-        uncond_mask = _mask(1, with_cache=False) if do_cfg else None
-        pos_embeds = prompt_embeds[:1]
-        neg_embeds = prompt_embeds[1:2] if do_cfg else None
-
-        with self.progress_bar(total=num_inference_steps) as progress_bar:
-            for i, t in enumerate(timesteps):
-                latent_model_input = self.scheduler.scale_model_input(latents, t).to(dtype)
-                t_expand = t.expand(1) - 1
-
-                cond_meta = (SDPAMetadata(current_timestep=i, attn_mask=cond_mask) if cond_mask is not None else None)
-                with torch.no_grad(), set_forward_context(current_timestep=i,
-                                                          attn_metadata=cond_meta,
-                                                          forward_batch=batch):
-                    noise_pred = self.transformer(latent_model_input,
-                                                  pos_embeds,
-                                                  prior_token_id,
-                                                  drop_keep,
-                                                  t_expand,
-                                                  target_size,
-                                                  crop_coords,
-                                                  kv_caches=kv_caches,
-                                                  kv_caches_mode="read")
-
-                if do_cfg:
-                    uncond_meta = (SDPAMetadata(current_timestep=i, attn_mask=uncond_mask)
-                                   if uncond_mask is not None else None)
-                    with torch.no_grad(), set_forward_context(current_timestep=i,
-                                                              attn_metadata=uncond_meta,
-                                                              forward_batch=batch):
-                        noise_pred_uncond = self.transformer(latent_model_input,
-                                                             neg_embeds,
-                                                             prior_token_id,
-                                                             drop_all,
-                                                             t_expand,
-                                                             target_size,
-                                                             crop_coords,
-                                                             kv_caches=kv_caches,
-                                                             kv_caches_mode="skip")
-                    noise_pred = noise_pred_uncond.float() + guidance_scale * (noise_pred.float() -
-                                                                               noise_pred_uncond.float())
-
-                latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
-                progress_bar.update()
-
-        kv_caches.clear()
-        batch.latents = latents.unsqueeze(2)

fastvideo/pipelines/pipeline_batch_info.py

@@ -109,10 +109,6 @@ class ForwardBatch:
     max_sequence_length: int | None = None
     prompt_template: dict[str, Any] | None = None
     do_classifier_free_guidance: bool = False
-    # When True, ``guidance_scale`` is passed into models that use embedded guidance (e.g. FLUX)
-    # and must not imply classic dual-forward CFG. Use ``true_cfg_scale > 1`` for true CFG.
-    use_embedded_guidance: bool = False
-    true_cfg_scale: float = 1.0
 
     # Batch info
     batch_size: int | None = None
@@ -256,12 +252,9 @@ class ForwardBatch:
     def __post_init__(self):
         """Initialize dependent fields after dataclass initialization."""
 
-        # LTX-2 text CFG scales; FLUX uses ``use_embedded_guidance`` so ``guidance_scale > 1`` alone
-        # does not enable classifier-free guidance.
+        # Enable CFG for standard guidance_scale and LTX-2 text CFG scales.
         ltx2_text_cfg_enabled = (self.ltx2_cfg_scale_video != 1.0 or self.ltx2_cfg_scale_audio != 1.0)
-        if self.use_embedded_guidance:
-            self.do_classifier_free_guidance = (self.true_cfg_scale > 1.0) or ltx2_text_cfg_enabled
-        elif self.guidance_scale > 1.0 or ltx2_text_cfg_enabled:
+        if self.guidance_scale > 1.0 or ltx2_text_cfg_enabled:
             self.do_classifier_free_guidance = True
         if self.negative_prompt_embeds is None:
             self.negative_prompt_embeds = []

fastvideo/pipelines/stages/denoising.py

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

fastvideo/pipelines/stages/flux_stages.py

@@ -1,424 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-
-from __future__ import annotations
-
-import inspect
-from typing import Any
-
-import torch
-from diffusers.utils.torch_utils import randn_tensor
-
-from fastvideo.distributed import get_local_torch_device
-from fastvideo.fastvideo_args import FastVideoArgs
-from fastvideo.forward_context import set_forward_context
-from fastvideo.logger import init_logger
-from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
-from fastvideo.pipelines.stages.base import PipelineStage
-from fastvideo.pipelines.stages.input_validation import InputValidationStage
-from fastvideo.pipelines.stages.timestep_preparation import TimestepPreparationStage
-from fastvideo.utils import PRECISION_TO_TYPE
-
-logger = init_logger(__name__)
-
-
-def _pack_latents(
-    latents: torch.Tensor,
-    batch_size: int,
-    num_channels_latents: int,
-    height: int,
-    width: int,
-) -> torch.Tensor:
-    """Diffusers ``_pack_latents`` for FLUX (2×2 spatial pack in latent space)."""
-    latents = latents.view(batch_size, num_channels_latents, height // 2, 2, width // 2, 2)
-    latents = latents.permute(0, 2, 4, 1, 3, 5)
-    return latents.reshape(batch_size, (height // 2) * (width // 2), num_channels_latents * 4)
-
-
-def _unpack_latents(
-    latents: torch.Tensor,
-    batch_size: int,
-    num_channels_latents: int,
-    height: int,
-    width: int,
-) -> torch.Tensor:
-    """Inverse of ``_pack_latents``."""
-    latents = latents.reshape(batch_size, height // 2, width // 2, num_channels_latents, 2, 2)
-    latents = latents.permute(0, 3, 1, 4, 2, 5)
-    return latents.reshape(batch_size, num_channels_latents, height, width)
-
-
-def _prepare_latent_image_ids(
-    patch_height: int,
-    patch_width: int,
-    device: torch.device,
-    dtype: torch.dtype = torch.long,
-) -> torch.Tensor:
-    """Match Diffusers ``FluxPipeline._prepare_latent_image_ids`` (no batch dim)."""
-    latent_image_ids = torch.zeros(patch_height, patch_width, 3, device=device, dtype=torch.float32)
-    latent_image_ids[..., 1] = latent_image_ids[..., 1] + torch.arange(patch_height, device=device)[:, None]
-    latent_image_ids[..., 2] = latent_image_ids[..., 2] + torch.arange(patch_width, device=device)[None, :]
-    h, w, c = latent_image_ids.shape
-    return latent_image_ids.reshape(h * w, c).to(dtype=dtype)
-
-
-class FluxInputValidationStage(InputValidationStage):
-    """Require height/width divisible by 16 (VAE scale × 2 for FLUX packing)."""
-
-    def forward(
-        self,
-        batch: ForwardBatch,
-        fastvideo_args: FastVideoArgs,
-    ) -> ForwardBatch:
-        if (batch.height is not None and batch.width is not None and (batch.height % 16 != 0 or batch.width % 16 != 0)):
-            raise ValueError("FLUX expects height and width divisible by 16 "
-                             f"(VAE latent grid × 2× packing); got {batch.height}×{batch.width}.")
-        return super().forward(batch, fastvideo_args)
-
-
-class FluxConditioningStage(PipelineStage):
-    """Build CLIP pooled + T5 sequence + ``text_ids`` (and optional negative for true CFG)."""
-
-    @torch.no_grad()
-    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
-        if len(batch.prompt_embeds) < 2:
-            raise ValueError("FluxConditioningStage expects 2 prompt_embeds (CLIP pooled, T5 sequence), "
-                             f"got {len(batch.prompt_embeds)}")
-
-        device = get_local_torch_device()
-        target_dtype = PRECISION_TO_TYPE[fastvideo_args.pipeline_config.dit_precision]
-
-        pooled = batch.prompt_embeds[0].to(device=device, dtype=target_dtype)
-        enc = batch.prompt_embeds[1].to(device=device, dtype=target_dtype)
-        seq_len = enc.shape[1]
-        text_ids = torch.zeros(seq_len, 3, device=device, dtype=torch.long)
-
-        batch.extra["flux_pooled_projections"] = pooled
-        batch.extra["flux_encoder_hidden_states"] = enc
-        batch.extra["flux_text_ids"] = text_ids
-
-        if batch.do_classifier_free_guidance:
-            if not batch.negative_prompt_embeds or len(batch.negative_prompt_embeds) < 2:
-                raise ValueError("True CFG requires two negative_prompt_embeds (CLIP, T5).")
-            neg_pooled = batch.negative_prompt_embeds[0].to(device=device, dtype=target_dtype)
-            neg_enc = batch.negative_prompt_embeds[1].to(device=device, dtype=target_dtype)
-            batch.extra["flux_negative_pooled_projections"] = neg_pooled
-            batch.extra["flux_negative_encoder_hidden_states"] = neg_enc
-
-        return batch
-
-
-class FluxTimestepPreparationStage(TimestepPreparationStage):
-    """Flow Match with resolution-dependent ``mu`` from packed image sequence length."""
-
-    @staticmethod
-    def _calculate_mu(
-        image_seq_len: int,
-        base_seq_len: int = 256,
-        max_seq_len: int = 4096,
-        base_shift: float = 0.5,
-        max_shift: float = 1.15,
-    ) -> float:
-        m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
-        b = base_shift - m * base_seq_len
-        return float(image_seq_len) * m + b
-
-    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
-        sig = inspect.signature(self.scheduler.set_timesteps)
-        if "mu" not in sig.parameters:
-            logger.warning(
-                "FLUX timestep prep: scheduler %s.set_timesteps does not accept 'mu'; falling back to the base "
-                "timestep schedule. FLUX expects a FlowMatchEulerDiscreteScheduler with resolution-dependent "
-                "dynamic shifting — output quality may degrade.",
-                type(self.scheduler).__name__)
-            return super().forward(batch, fastvideo_args)
-
-        cfg = getattr(self.scheduler, "config", None)
-        use_dynamic = bool(getattr(cfg, "use_dynamic_shifting", False)) if cfg is not None else False
-        if not use_dynamic:
-            logger.warning(
-                "FLUX timestep prep: scheduler has use_dynamic_shifting=False; falling back to the base timestep "
-                "schedule and skipping the resolution-dependent 'mu' shift. FLUX requires dynamic shifting for "
-                "correct timesteps — output quality may degrade.")
-            return super().forward(batch, fastvideo_args)
-
-        if batch.height is None or batch.width is None:
-            raise ValueError("height/width must be set before FluxTimestepPreparationStage")
-
-        vae_arch = fastvideo_args.pipeline_config.vae_config.arch_config
-        spatial_ratio = int(getattr(vae_arch, "spatial_compression_ratio", 8))
-        h_lat = batch.height // spatial_ratio
-        w_lat = batch.width // spatial_ratio
-        if h_lat % 2 != 0 or w_lat % 2 != 0:
-            raise ValueError(
-                f"Latent spatial dims must be even for FLUX packing; got {h_lat}×{w_lat} from {batch.height}×{batch.width}."
-            )
-        image_seq_len = (h_lat // 2) * (w_lat // 2)
-
-        base_seq_len = int(getattr(cfg, "base_image_seq_len", 256))
-        max_seq_len = int(getattr(cfg, "max_image_seq_len", 4096))
-        base_shift = float(getattr(cfg, "base_shift", 0.5))
-        max_shift = float(getattr(cfg, "max_shift", 1.15))
-
-        device = get_local_torch_device()
-        mu = self._calculate_mu(
-            image_seq_len=image_seq_len,
-            base_seq_len=base_seq_len,
-            max_seq_len=max_seq_len,
-            base_shift=base_shift,
-            max_shift=max_shift,
-        )
-        self.scheduler.set_timesteps(batch.num_inference_steps, device=device, mu=mu)
-        batch.timesteps = self.scheduler.timesteps
-        return batch
-
-
-class FluxLatentPreparationStage(PipelineStage):
-
-    def __init__(self, scheduler) -> None:
-        super().__init__()
-        self.scheduler = scheduler
-
-    @torch.no_grad()
-    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
-        if batch.height is None or batch.width is None:
-            raise ValueError("height/width required for FluxLatentPreparationStage")
-
-        if isinstance(batch.prompt, list):
-            batch_size = len(batch.prompt)
-        elif batch.prompt is not None:
-            batch_size = 1
-        else:
-            if not batch.prompt_embeds:
-                raise ValueError("prompt or prompt_embeds must be provided")
-            batch_size = batch.prompt_embeds[0].shape[0]
-
-        batch_size *= batch.num_videos_per_prompt
-
-        if isinstance(batch.generator, list) and len(batch.generator) != batch_size:
-            raise ValueError(f"generator list length {len(batch.generator)} does not match batch_size {batch_size}")
-
-        arch = fastvideo_args.pipeline_config.dit_config.arch_config
-        in_channels = int(getattr(arch, "in_channels", 64))
-        num_channels_latents = in_channels // 4
-
-        vae_arch = fastvideo_args.pipeline_config.vae_config.arch_config
-        spatial_ratio = int(getattr(vae_arch, "spatial_compression_ratio", 8))
-
-        h_lat = batch.height // spatial_ratio
-        w_lat = batch.width // spatial_ratio
-
-        dtype = PRECISION_TO_TYPE[fastvideo_args.pipeline_config.dit_precision]
-        device = get_local_torch_device()
-
-        shape = (batch_size, num_channels_latents, h_lat, w_lat)
-        latents = batch.latents
-        if latents is None:
-            latents = randn_tensor(shape, generator=batch.generator, device=device, dtype=dtype)
-            if hasattr(self.scheduler, "init_noise_sigma"):
-                latents = latents * self.scheduler.init_noise_sigma
-        else:
-            latents = latents.to(device=device, dtype=dtype)
-            if latents.shape != shape:
-                raise ValueError(f"Expected latents shape {shape}, got {tuple(latents.shape)}")
-            if hasattr(self.scheduler, "init_noise_sigma"):
-                latents = latents * self.scheduler.init_noise_sigma
-
-        packed = _pack_latents(latents, batch_size, num_channels_latents, h_lat, w_lat)
-
-        patch_h, patch_w = h_lat // 2, w_lat // 2
-        img_ids = _prepare_latent_image_ids(patch_h, patch_w, device, dtype=torch.long)
-
-        batch.latents = packed
-        batch.raw_latent_shape = shape
-        batch.extra["flux_h_lat"] = h_lat
-        batch.extra["flux_w_lat"] = w_lat
-        batch.extra["flux_num_channels_latents"] = num_channels_latents
-        batch.extra["flux_latent_image_ids"] = img_ids
-
-        return batch
-
-
-class FluxDenoisingStage(PipelineStage):
-
-    def __init__(self, transformer, scheduler) -> None:
-        super().__init__()
-        self.transformer = transformer
-        self.scheduler = scheduler
-
-    @staticmethod
-    def _step_kwargs(scheduler_step, batch: ForwardBatch) -> dict[str, Any]:
-        kwargs: dict[str, Any] = {}
-        sig = inspect.signature(scheduler_step)
-        if "generator" in sig.parameters:
-            gen = batch.generator[0] if isinstance(batch.generator, list) else batch.generator
-            kwargs["generator"] = gen
-        return kwargs
-
-    @torch.no_grad()
-    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
-        if batch.timesteps is None:
-            raise ValueError("timesteps must be set before FluxDenoisingStage")
-        if batch.latents is None:
-            raise ValueError("latents must be set before FluxDenoisingStage")
-
-        packed = batch.latents
-        timesteps = batch.timesteps
-
-        pooled = batch.extra["flux_pooled_projections"]
-        enc = batch.extra["flux_encoder_hidden_states"]
-        txt_ids = batch.extra["flux_text_ids"]
-        img_ids = batch.extra["flux_latent_image_ids"]
-
-        neg_pooled = batch.extra.get("flux_negative_pooled_projections")
-        neg_enc = batch.extra.get("flux_negative_encoder_hidden_states")
-
-        true_cfg_scale = float(batch.true_cfg_scale)
-        use_true_cfg = batch.do_classifier_free_guidance and true_cfg_scale > 1.0
-
-        # Prefer the loaded transformer's arch (HF ``guidance_embeds``), not static pipeline defaults.
-        tr_arch = self.transformer.fastvideo_config.arch_config
-        guidance_embeds = bool(getattr(tr_arch, "guidance_embeds", False))
-
-        device = get_local_torch_device()
-        target_dtype = PRECISION_TO_TYPE[fastvideo_args.pipeline_config.dit_precision]
-        autocast_enabled = (target_dtype != torch.float32) and not fastvideo_args.disable_autocast
-
-        bs = packed.shape[0]
-        if guidance_embeds:
-            guidance = torch.full((bs, ), float(batch.guidance_scale), device=device, dtype=torch.float32)
-        else:
-            guidance = None
-
-        step_extras = self._step_kwargs(self.scheduler.step, batch)
-
-        for t in timesteps:
-            t_scalar = t
-            if not isinstance(t_scalar, torch.Tensor):
-                t_scalar = torch.tensor([t_scalar], device=device, dtype=torch.float32)
-            t_scalar = t_scalar.to(device=device, dtype=torch.float32)
-
-            timestep_model = t_scalar.expand(bs).float() / 1000.0
-            timestep_model = timestep_model.to(dtype=target_dtype)
-
-            ts_ctx = int(t_scalar.reshape(-1)[0].item())
-            with (
-                    torch.autocast(
-                        device_type="cuda",
-                        dtype=target_dtype,
-                        enabled=autocast_enabled and device.type == "cuda",
-                    ),
-                    set_forward_context(
-                        current_timestep=ts_ctx,
-                        attn_metadata=None,
-                        forward_batch=batch,
-                    ),
-            ):
-                if use_true_cfg:
-                    assert neg_enc is not None and neg_pooled is not None
-                    n_neg = self.transformer(
-                        hidden_states=packed,
-                        encoder_hidden_states=neg_enc,
-                        pooled_projections=neg_pooled,
-                        timestep=timestep_model,
-                        guidance=guidance,
-                        txt_ids=txt_ids,
-                        img_ids=img_ids,
-                        return_dict=False,
-                    )[0]
-                    n_pos = self.transformer(
-                        hidden_states=packed,
-                        encoder_hidden_states=enc,
-                        pooled_projections=pooled,
-                        timestep=timestep_model,
-                        guidance=guidance,
-                        txt_ids=txt_ids,
-                        img_ids=img_ids,
-                        return_dict=False,
-                    )[0]
-                    noise_pred = n_neg + true_cfg_scale * (n_pos - n_neg)
-                else:
-                    noise_pred = self.transformer(
-                        hidden_states=packed,
-                        encoder_hidden_states=enc,
-                        pooled_projections=pooled,
-                        timestep=timestep_model,
-                        guidance=guidance,
-                        txt_ids=txt_ids,
-                        img_ids=img_ids,
-                        return_dict=False,
-                    )[0]
-
-            packed = self.scheduler.step(
-                noise_pred,
-                t_scalar,
-                packed,
-                return_dict=False,
-                **step_extras,
-            )[0]
-
-        batch.latents = packed
-        return batch
-
-
-class FluxDecodingStage(PipelineStage):
-    """Unpack latents, apply VAE scaling/shift, decode to pixels (5D output ``B×3×1×H×W``)."""
-
-    def __init__(self, vae) -> None:
-        super().__init__()
-        self.vae = vae
-
-    @staticmethod
-    def _denormalize_latents(latents: torch.Tensor, vae: Any) -> torch.Tensor:
-        cfg = getattr(vae, "config", None)
-        sf = getattr(cfg, "scaling_factor", None) if cfg is not None else None
-        sh = getattr(cfg, "shift_factor", None) if cfg is not None else None
-        if sf is None and hasattr(vae, "scaling_factor"):
-            sf = vae.scaling_factor
-        if sh is None and hasattr(vae, "shift_factor"):
-            sh = vae.shift_factor
-        if sf is not None:
-            latents = latents / (sf.to(latents.device, latents.dtype) if isinstance(sf, torch.Tensor) else sf)
-            if sh is not None:
-                latents = latents + (sh.to(latents.device, latents.dtype) if isinstance(sh, torch.Tensor) else sh)
-        return latents
-
-    @torch.no_grad()
-    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
-        packed = batch.latents
-        if packed is None:
-            raise ValueError("latents must be set before FluxDecodingStage")
-
-        h_lat = int(batch.extra["flux_h_lat"])
-        w_lat = int(batch.extra["flux_w_lat"])
-        num_ch = int(batch.extra["flux_num_channels_latents"])
-        raw_shape = batch.raw_latent_shape
-        if raw_shape is None:
-            raise ValueError("raw_latent_shape missing; FluxLatentPreparationStage must run first.")
-        batch_size = int(raw_shape[0])
-
-        infer_device = get_local_torch_device()
-        packed = packed.to(infer_device)
-
-        latents_4d = _unpack_latents(packed, batch_size, num_ch, h_lat, w_lat)
-        latents_4d = self._denormalize_latents(latents_4d, self.vae)
-
-        vae_device = next(self.vae.parameters()).device
-        latents_4d = latents_4d.to(device=vae_device)
-
-        vae_dtype = PRECISION_TO_TYPE[fastvideo_args.pipeline_config.vae_precision]
-        autocast_enabled = (vae_dtype != torch.float32) and not fastvideo_args.disable_autocast
-        use_cuda_autocast = autocast_enabled and vae_device.type == "cuda"
-
-        with torch.autocast(
-                device_type="cuda",
-                dtype=vae_dtype,
-                enabled=use_cuda_autocast,
-        ):
-            if not autocast_enabled:
-                latents_4d = latents_4d.to(dtype=vae_dtype)
-            dec = self.vae.decode(latents_4d)
-            image = dec.sample if hasattr(dec, "sample") else dec[0]
-
-        image = (image / 2 + 0.5).clamp(0, 1)
-        batch.output = image.unsqueeze(2).detach().float().cpu()
-        return batch

fastvideo/pipelines/stages/input_validation.py

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

fastvideo/pipelines/stages/text_encoding.py

@@ -68,12 +68,20 @@ class TextEncodingStage(PipelineStage):
         assert batch.prompt is not None
         prompt_text: str | list[str] = batch.prompt
         all_indices: list[int] = list(range(len(self.text_encoders)))
-        prompt_embeds_list, prompt_masks_list = self.encode_text(
-            prompt_text,
-            fastvideo_args,
-            encoder_index=all_indices,
-            return_attention_mask=True,
-        )
+        if isinstance(prompt_text, list):
+            prompt_embeds_list, prompt_masks_list = self._encode_prompt_list_individually(
+                prompt_text,
+                fastvideo_args,
+                encoder_index=all_indices,
+                return_attention_mask=True,
+            )
+        else:
+            prompt_embeds_list, prompt_masks_list = self.encode_text(
+                prompt_text,
+                fastvideo_args,
+                encoder_index=all_indices,
+                return_attention_mask=True,
+            )
         if self._last_audio_embeds is not None:
             batch.extra["ltx2_audio_prompt_embeds"] = self._last_audio_embeds
 
@@ -85,13 +93,24 @@ class TextEncodingStage(PipelineStage):
 
         # Encode negative prompt if CFG is enabled
         if batch.do_classifier_free_guidance:
-            assert isinstance(batch.negative_prompt, str)
-            neg_embeds_list, neg_masks_list = self.encode_text(
-                batch.negative_prompt,
-                fastvideo_args,
-                encoder_index=all_indices,
-                return_attention_mask=True,
-            )
+            assert isinstance(batch.negative_prompt, str | list)
+            negative_prompt: str | list[str] = batch.negative_prompt
+            if isinstance(batch.prompt, list) and isinstance(negative_prompt, str):
+                negative_prompt = [negative_prompt] * len(batch.prompt)
+            if isinstance(negative_prompt, list):
+                neg_embeds_list, neg_masks_list = self._encode_prompt_list_individually(
+                    negative_prompt,
+                    fastvideo_args,
+                    encoder_index=all_indices,
+                    return_attention_mask=True,
+                )
+            else:
+                neg_embeds_list, neg_masks_list = self.encode_text(
+                    negative_prompt,
+                    fastvideo_args,
+                    encoder_index=all_indices,
+                    return_attention_mask=True,
+                )
             if self._last_audio_embeds is not None:
                 batch.extra["ltx2_audio_negative_embeds"] = self._last_audio_embeds
 
@@ -104,6 +123,81 @@ class TextEncodingStage(PipelineStage):
 
         return batch
 
+    def _encode_prompt_list_individually(
+        self,
+        texts: list[str],
+        fastvideo_args: FastVideoArgs,
+        *,
+        encoder_index: list[int],
+        return_attention_mask: bool,
+    ) -> tuple[list[torch.Tensor], list[torch.Tensor]]:
+        per_prompt_embeds: list[list[torch.Tensor]] = []
+        per_prompt_masks: list[list[torch.Tensor]] = []
+        per_prompt_audio_embeds: list[list[torch.Tensor] | None] = []
+
+        for text in texts:
+            embeds, masks = self.encode_text(
+                text,
+                fastvideo_args,
+                encoder_index=encoder_index,
+                return_attention_mask=return_attention_mask,
+            )
+            per_prompt_embeds.append(embeds)
+            per_prompt_masks.append(masks)
+            per_prompt_audio_embeds.append(self._last_audio_embeds)
+
+        merged_embeds = [
+            self._cat_tensors([prompt_embeds[encoder_pos] for prompt_embeds in per_prompt_embeds])
+            for encoder_pos in range(len(per_prompt_embeds[0]))
+        ]
+        merged_masks = [
+            self._cat_attention_masks([prompt_masks[encoder_pos] for prompt_masks in per_prompt_masks])
+            for encoder_pos in range(len(per_prompt_masks[0]))
+        ]
+        if per_prompt_audio_embeds and all(audio_embeds is not None for audio_embeds in per_prompt_audio_embeds):
+            audio_embed_lists = [audio_embeds for audio_embeds in per_prompt_audio_embeds if audio_embeds is not None]
+            self._last_audio_embeds = [
+                self._cat_tensors([audio_embeds[encoder_pos] for audio_embeds in audio_embed_lists])
+                for encoder_pos in range(len(audio_embed_lists[0]))
+            ]
+        else:
+            self._last_audio_embeds = None
+        return merged_embeds, merged_masks
+
+    @staticmethod
+    def _cat_tensors(tensors: list[torch.Tensor]) -> torch.Tensor:
+        base_shape = tensors[0].shape[1:]
+        if all(tensor.shape[1:] == base_shape for tensor in tensors):
+            return torch.cat(tensors, dim=0)
+        if all(tensor.ndim == 3 for tensor in tensors):
+            base_trailing_shape = tensors[0].shape[2:]
+            if all(tensor.shape[2:] == base_trailing_shape for tensor in tensors):
+                max_length = max(tensor.shape[1] for tensor in tensors)
+                padded_tensors = []
+                for tensor in tensors:
+                    pad_width = max_length - tensor.shape[1]
+                    if pad_width > 0:
+                        tensor = torch.nn.functional.pad(tensor, (0, 0, 0, pad_width), value=0.0)
+                    padded_tensors.append(tensor)
+                return torch.cat(padded_tensors, dim=0)
+        raise ValueError(f"Cannot concatenate tensors with shapes: {[list(tensor.shape) for tensor in tensors]}")
+
+    @staticmethod
+    def _cat_attention_masks(masks: list[torch.Tensor]) -> torch.Tensor:
+        base_shape = masks[0].shape[1:]
+        if all(mask.shape[1:] == base_shape for mask in masks):
+            return torch.cat(masks, dim=0)
+        if all(mask.ndim == 2 for mask in masks):
+            max_length = max(mask.shape[1] for mask in masks)
+            padded_masks = []
+            for mask in masks:
+                pad_width = max_length - mask.shape[1]
+                if pad_width > 0:
+                    mask = torch.nn.functional.pad(mask, (0, pad_width), value=0)
+                padded_masks.append(mask)
+            return torch.cat(padded_masks, dim=0)
+        raise ValueError(f"Cannot concatenate attention masks with shapes: {[list(mask.shape) for mask in masks]}")
+
     def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult:
         """Verify text encoding stage inputs."""
         result = VerificationResult()
@@ -253,6 +347,9 @@ class TextEncodingStage(PipelineStage):
                     attn_masks_list.append(attention_mask)
                     return self.return_embeds(embeds_list, attn_masks_list, return_type, return_attention_mask, indices)
 
+            if len(processed_texts) > 1 and "padding" not in tok_kwargs:
+                tok_kwargs["padding"] = True
+
             # If tokenizer is a multimodal processor (e.g. Qwen2_5_VLProcessor),
             # use its inner tokenizer for text-only encoding.
             tok = getattr(tokenizer, "tokenizer", tokenizer)

fastvideo/registry.py

@@ -58,14 +58,11 @@ from fastvideo.configs.pipelines.wan import (
     WanT2V480PConfig,
     WanT2V720PConfig,
 )
-from fastvideo.configs.pipelines.glm_image import GlmImageConfig
-from fastvideo.configs.pipelines.flux import FluxPipelineConfig
 from fastvideo.configs.pipelines.sd35 import SD35Config
 from fastvideo.configs.pipelines.stable_audio import (StableAudioOpenSmallConfig, StableAudioT2AConfig)
 from fastvideo.api.sampling_param import SamplingParam
 from fastvideo.api.matrixgame2 import MatrixGame2SamplingParam
 from fastvideo.api.matrixgame3 import MatrixGame3SamplingParam
-from fastvideo.api.flux import FluxSamplingParam
 
 from fastvideo.fastvideo_args import WorkloadType
 from fastvideo.logger import init_logger
@@ -1077,33 +1074,6 @@ def _register_configs() -> None:
         default_preset="sd35_medium",
     )
 
-    # GLM-Image
-    register_configs(
-        sampling_param_cls=None,
-        pipeline_config_cls=GlmImageConfig,
-        hf_model_paths=[
-            "zai-org/GLM-Image",
-        ],
-        model_detectors=[lambda path: "glmimage" in path.lower() or "glm-image" in path.lower()],
-        workload_types=(WorkloadType.T2I, ),
-        model_family="glm_image",
-    )
-
-    # FLUX.1-dev (Diffusers)
-    register_configs(
-        sampling_param_cls=FluxSamplingParam,
-        pipeline_config_cls=FluxPipelineConfig,
-        workload_types=(WorkloadType.T2I, ),
-        hf_model_paths=[
-            "black-forest-labs/FLUX.1-dev",
-        ],
-        model_detectors=[
-            lambda path: "fluxpipeline" in path,
-            lambda path: "flux.1-dev" in path or "flux_1_dev" in path,
-            lambda path: "/flux/" in path or path.endswith("/flux"),
-        ],
-    )
-
 
 # --- Part 3: Main Resolver ---
 

fastvideo/tests/api/test_compat_translation.py

@@ -9,7 +9,7 @@ from fastvideo.api.compat import (
     generator_config_to_fastvideo_args,
     legacy_from_pretrained_to_config,
 )
-from fastvideo.api.schema import CompileConfig, GeneratorConfig
+from fastvideo.api.schema import BatchingConfig, CompileConfig, GeneratorConfig
 
 
 class TestLegacyTorchCompileKwargsTranslation:
@@ -153,6 +153,48 @@ class TestLegacyTextEncoderCompileTranslation:
         assert "enable_torch_compile_text_encoder" not in args.kwargs
 
 
+class TestBatchingTranslation:
+
+    def test_flat_kwargs_promote_to_engine_batching(self) -> None:
+        config = legacy_from_pretrained_to_config(
+            "/models/wan",
+            {
+                "batching_mode": "dynamic",
+                "batching_max_size": 4,
+                "batching_delay_ms": 25.0,
+                "batching_config": "/tmp/batching.json",
+                "enable_batching_metrics": True,
+            },
+        )
+
+        assert config.engine.batching.mode == "dynamic"
+        assert config.engine.batching.max_size == 4
+        assert config.engine.batching.delay_ms == 25.0
+        assert config.engine.batching.config_path == "/tmp/batching.json"
+        assert config.engine.batching.enable_metrics is True
+
+    def test_typed_batching_emits_fastvideo_args_kwargs(self, monkeypatch) -> None:
+        _stub_fastvideo_args_from_kwargs(monkeypatch)
+        config = GeneratorConfig(
+            model_path="/models/wan",
+            engine=_engine_with_batching(BatchingConfig(
+                mode="dynamic",
+                max_size=3,
+                delay_ms=10.0,
+                config_path="/tmp/batching.json",
+                enable_metrics=True,
+            )),
+        )
+
+        args = generator_config_to_fastvideo_args(config)
+
+        assert args.kwargs["batching_mode"] == "dynamic"
+        assert args.kwargs["batching_max_size"] == 3
+        assert args.kwargs["batching_delay_ms"] == 10.0
+        assert args.kwargs["batching_config"] == "/tmp/batching.json"
+        assert args.kwargs["enable_batching_metrics"] is True
+
+
 # -------------------------------------------------------------------
 # Helpers
 # -------------------------------------------------------------------
@@ -166,6 +208,13 @@ def _engine_with_compile(compile_config):
     return engine
 
 
+def _engine_with_batching(batching_config):
+    from fastvideo.api.schema import EngineConfig
+    engine = EngineConfig()
+    engine.batching = batching_config
+    return engine
+
+
 def _stub_fastvideo_args_from_kwargs(monkeypatch):
     """Swap ``FastVideoArgs.from_kwargs`` for a capture-only stub so
     translation tests don't need to construct a valid FastVideoArgs."""

fastvideo/tests/api/test_parser.py

@@ -130,6 +130,13 @@ def test_load_run_config_supports_yaml_roundtrip(tmp_path) -> None:
                 "use_fsdp_inference": False,
                 "disable_autocast": False,
                 "quantization": None,
+                "batching": {
+                    "mode": "disabled",
+                    "max_size": 1,
+                    "delay_ms": 0.0,
+                    "config_path": None,
+                    "enable_metrics": False,
+                },
             },
             "pipeline": {
                 "workload_type": None,
@@ -183,7 +190,6 @@ def test_load_run_config_supports_yaml_roundtrip(tmp_path) -> None:
                 "guidance_scale_2": None,
                 "guidance_rescale": 0.0,
                 "true_cfg_scale": None,
-                "use_embedded_guidance": None,
                 "boundary_ratio": None,
                 "sigmas": None,
             },

fastvideo/tests/attention/test_sdpa_head_sizes.py

@@ -1,49 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""CPU tests: SDPA must not restrict attention head sizes.
-
-``torch.nn.functional.scaled_dot_product_attention`` is head-size
-agnostic, but ``SDPABackend.get_supported_head_sizes()`` used to declare
-a list copied from FlashAttentionBackend that omitted e.g. 80 (the head
-size of CLIP-L vision encoders such as matrixgame2's). That made
-capability validation flag SDPA as incompatible with head-80 layers even
-though it handles them fine.
-
-CPU-only; no GPU or flash-attn dependency.
-"""
-
-import torch
-
-from fastvideo.attention.backends.sdpa import SDPABackend, SDPAImpl, SDPAMetadata
-from fastvideo.platforms.cuda import CudaPlatformBase
-from fastvideo.platforms.interface import AttentionBackendEnum
-
-
-def test_sdpa_head_sizes_unrestricted() -> None:
-    """SDPA declares no head-size restriction (None), so 80 is allowed."""
-    assert SDPABackend.get_supported_head_sizes() is None
-    # Once the capability-validation API (#1494) lands, head-80 must pass.
-    validate = getattr(SDPABackend, "validate_compatibility", None)
-    if validate is not None:
-        assert validate(head_size=80, dtype=torch.bfloat16) is None
-
-
-def test_explicit_torch_sdpa_pin_head_80_resolves() -> None:
-    """An explicit TORCH_SDPA pin on a head-80 layer resolves to SDPABackend."""
-    cls_str = CudaPlatformBase.get_attn_backend_cls(AttentionBackendEnum.TORCH_SDPA, head_size=80, dtype=torch.bfloat16)
-    assert cls_str == "fastvideo.attention.backends.sdpa.SDPABackend"
-
-
-def test_sdpa_impl_computes_head_80_attention() -> None:
-    """SDPAImpl produces correct output for head_size=80 (not in the old list)."""
-    bs, seq, heads, head_dim = 2, 8, 2, 80
-    softmax_scale = head_dim**-0.5
-    impl = SDPAImpl(num_heads=heads, head_size=head_dim, causal=False, softmax_scale=softmax_scale)
-    torch.manual_seed(0)
-    q, k, v = (torch.randn(bs, seq, heads, head_dim) for _ in range(3))
-
-    out = impl.forward(q, k, v, SDPAMetadata(current_timestep=0))
-
-    assert out.shape == (bs, seq, heads, head_dim)
-    qt, kt, vt = (t.transpose(1, 2) for t in (q, k, v))
-    ref = torch.softmax(qt @ kt.transpose(-2, -1) * softmax_scale, dim=-1) @ vt
-    torch.testing.assert_close(out, ref.transpose(1, 2), rtol=1e-4, atol=1e-4)

fastvideo/tests/batching/run_dynamic_batching_parity.py

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

fastvideo/tests/batching/test_admission.py

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

fastvideo/tests/batching/test_scheduler.py

@@ -0,0 +1,226 @@
+# SPDX-License-Identifier: Apache-2.0
+"""CPU regression tests for the dynamic VideoBatchScheduler (no GPU needed)."""
+
+from __future__ import annotations
+
+import asyncio
+import threading
+from types import SimpleNamespace
+
+import pytest
+
+from fastvideo.configs.pipelines.base import PipelineConfig
+from fastvideo.configs.pipelines.wan import (
+    FastWan2_1_T2V_480P_Config,
+    SelfForcingWanT2V480PConfig,
+    WanT2V480PConfig,
+)
+from fastvideo.entrypoints.openai.batching import VideoBatchScheduler
+
+
+class _FakeBatchGenerator:
+
+    def __init__(self):
+        self.calls: list[list[dict]] = []
+
+    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]
+
+
+class _BlockingFirstCallGenerator(_FakeBatchGenerator):
+    """Blocks the first generate_video_batch call until released.
+
+    Lets a test deterministically hold one generation "in flight" while more
+    requests queue up behind it.
+    """
+
+    def __init__(self):
+        super().__init__()
+        self.started = threading.Event()
+        self.release = threading.Event()
+        self._first_call = True
+
+    def generate_video_batch(self, request_kwargs):
+        if self._first_call:
+            self._first_call = False
+            self.started.set()
+            assert self.release.wait(timeout=10)
+        return super().generate_video_batch(request_kwargs)
+
+
+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(supports_dynamic_batching=True),
+    )
+    defaults.update(overrides)
+    return SimpleNamespace(**defaults)
+
+
+def _job_kwargs(tmp_path, index: int) -> dict:
+    return {
+        "prompt": f"p{index}",
+        "height": 256,
+        "width": 256,
+        "num_frames": 1,
+        "num_inference_steps": 2,
+        "seed": index,
+        "output_path": str(tmp_path / f"video_{index}.mp4"),
+        "save_video": False,
+    }
+
+
+def test_expired_deadline_coalesces_jobs_queued_behind_inflight_generation(tmp_path):
+    """Jobs whose delay deadline expired while a generation was in flight must
+    still coalesce (greedy drain), not dispatch as batches of 1."""
+
+    async def run():
+        generator = _BlockingFirstCallGenerator()
+        scheduler = VideoBatchScheduler(generator, _batch_scheduler_args(batching_max_size=4, batching_delay_ms=25.0))
+        await scheduler.start()
+        try:
+            loop = asyncio.get_running_loop()
+            first = asyncio.ensure_future(scheduler.submit("req-1", _job_kwargs(tmp_path, 1)))
+            await asyncio.wait_for(loop.run_in_executor(None, generator.started.wait, 5), timeout=10)
+            rest = [asyncio.ensure_future(scheduler.submit(f"req-{i}", _job_kwargs(tmp_path, i))) for i in (2, 3, 4)]
+            # Let the in-queue jobs outlive the 25 ms batching window while the
+            # first generation is still blocking the scheduler.
+            await asyncio.sleep(0.1)
+            generator.release.set()
+            await asyncio.wait_for(asyncio.gather(first, *rest), timeout=10)
+        finally:
+            await scheduler.stop()
+        return [[item["prompt"] for item in call] for call in generator.calls]
+
+    assert asyncio.run(run()) == [["p1"], ["p2", "p3", "p4"]]
+
+
+def test_submit_after_stop_raises(tmp_path):
+
+    async def run():
+        scheduler = VideoBatchScheduler(_FakeBatchGenerator(), _batch_scheduler_args())
+        await scheduler.start()
+        await scheduler.stop()
+        with pytest.raises(RuntimeError, match="stopped"):
+            await scheduler.submit("req-1", _job_kwargs(tmp_path, 1))
+
+    asyncio.run(run())
+
+
+def test_stop_fails_jobs_still_queued_without_hanging(tmp_path):
+    """stop() during an in-flight generation must fail the still-queued jobs
+    instead of stranding their futures forever."""
+
+    async def run():
+        generator = _BlockingFirstCallGenerator()
+        scheduler = VideoBatchScheduler(generator, _batch_scheduler_args())
+        await scheduler.start()
+        loop = asyncio.get_running_loop()
+        first = asyncio.ensure_future(scheduler.submit("req-1", _job_kwargs(tmp_path, 1)))
+        await asyncio.wait_for(loop.run_in_executor(None, generator.started.wait, 5), timeout=10)
+        queued = [asyncio.ensure_future(scheduler.submit(f"req-{i}", _job_kwargs(tmp_path, i))) for i in (2, 3)]
+        await asyncio.sleep(0)  # let the queued submits enqueue
+        stop_task = asyncio.ensure_future(scheduler.stop())
+        await asyncio.sleep(0)  # let stop() mark the scheduler stopped
+        generator.release.set()
+
+        result = await asyncio.wait_for(first, timeout=10)
+        assert result["prompts"] == "p1"
+        for task in queued:
+            with pytest.raises(RuntimeError, match="stopped before dispatch"):
+                await asyncio.wait_for(task, timeout=10)
+        await asyncio.wait_for(stop_task, timeout=10)
+
+    asyncio.run(run())
+
+
+def test_dispatch_exception_fails_all_group_futures(tmp_path):
+
+    class _ExplodingGenerator:
+
+        def generate_video_batch(self, request_kwargs):
+            raise ValueError("generation exploded")
+
+    async def run():
+        scheduler = VideoBatchScheduler(_ExplodingGenerator(), _batch_scheduler_args())
+        await scheduler.start()
+        try:
+            return await asyncio.wait_for(
+                asyncio.gather(
+                    scheduler.submit("req-1", _job_kwargs(tmp_path, 1)),
+                    scheduler.submit("req-2", _job_kwargs(tmp_path, 2)),
+                    return_exceptions=True,
+                ),
+                timeout=10,
+            )
+        finally:
+            await scheduler.stop()
+
+    results = asyncio.run(run())
+
+    assert len(results) == 2
+    for result in results:
+        assert isinstance(result, ValueError)
+        assert "generation exploded" in str(result)
+
+
+def test_run_crash_fails_inflight_and_queued_futures(tmp_path):
+    """If the _run task ever dies, every waiting future must fail instead of
+    hanging, and later submits must be rejected."""
+
+    async def run():
+        scheduler = VideoBatchScheduler(_FakeBatchGenerator(), _batch_scheduler_args())
+
+        async def boom(first):
+            raise RuntimeError("kaboom")
+
+        scheduler._collect_batch = boom  # type: ignore[method-assign]
+        await scheduler.start()
+        results = await asyncio.wait_for(
+            asyncio.gather(
+                scheduler.submit("req-1", _job_kwargs(tmp_path, 1)),
+                scheduler.submit("req-2", _job_kwargs(tmp_path, 2)),
+                return_exceptions=True,
+            ),
+            timeout=10,
+        )
+        with pytest.raises(RuntimeError, match="stopped"):
+            await scheduler.submit("req-3", _job_kwargs(tmp_path, 3))
+        return results
+
+    results = asyncio.run(run())
+
+    assert len(results) == 2
+    for result in results:
+        assert isinstance(result, RuntimeError)
+        assert "crashed" in str(result)
+
+
+def test_scheduler_start_rejects_pipeline_without_capability(tmp_path):
+
+    async def run():
+        scheduler = VideoBatchScheduler(
+            _FakeBatchGenerator(),
+            _batch_scheduler_args(pipeline_config=PipelineConfig()),
+        )
+        with pytest.raises(RuntimeError, match="does not support dynamic request batching"):
+            await scheduler.start()
+        assert scheduler._task is None
+
+    asyncio.run(run())
+
+
+def test_dynamic_batching_capability_gate_defaults():
+    # Default: every pipeline is excluded until its family opts in.
+    assert PipelineConfig().dynamic_batching_supported() is False
+    # The verified generic Wan denoising path opts in.
+    assert WanT2V480PConfig().dynamic_batching_supported() is True
+    # DMD and causal denoising variants are always vetoed (shared-RNG risk),
+    # even though they inherit the Wan opt-in flag.
+    assert FastWan2_1_T2V_480P_Config().dynamic_batching_supported() is False
+    assert SelfForcingWanT2V480PConfig().dynamic_batching_supported() is False

fastvideo/tests/batching/test_signature.py

@@ -0,0 +1,83 @@
+# 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_can_dynamic_batch_rejects_preset_latents() -> None:
+    first = _request()
+    second = _request(latents=object())
+
+    result = can_dynamic_batch(first, second)
+
+    assert result.can_batch is False
+    assert result.reason == "latents"
+
+
+def test_signature_treats_distinct_opaque_objects_as_incompatible() -> None:
+
+    class Opaque:
+        pass
+
+    first = _request("first")
+    second = _request("second")
+
+    distinct = can_dynamic_batch(first, second, base_extra={"foo": Opaque()}, candidate_extra={"foo": Opaque()})
+    assert distinct.can_batch is False
+
+    shared = Opaque()
+    same_object = can_dynamic_batch(first, second, base_extra={"foo": shared}, candidate_extra={"foo": shared})
+    assert same_object.can_batch is True
+
+
+def test_resolution_key_uses_generation_shape() -> None:
+    assert resolution_key(_request(height=720, width=1280, num_frames=81)) == "720x1280x81"

fastvideo/tests/entrypoints/test_openai_api.py

@@ -1,10 +1,15 @@
 """Unit tests for the OpenAI-compatible API server helpers (no GPU needed)."""
 
+import asyncio
 import os
+import time
+from types import SimpleNamespace
 from unittest.mock import patch
 
 import pytest
 
+from fastvideo.configs.pipelines.base import PipelineConfig
+from fastvideo.entrypoints.openai.batching import VideoBatchScheduler, _VideoBatchJob
 from fastvideo.api.parser import parse_config
 from fastvideo.api.schema import GenerationRequest
 from fastvideo.entrypoints.openai.protocol import (
@@ -21,6 +26,172 @@ from fastvideo.entrypoints.openai.utils import (
     parse_size,
 )
 
+
+class _FakeBatchGenerator:
+
+    def __init__(self):
+        self.calls = []
+
+    def generate_video_batch(self, request_kwargs):
+        self.calls.append([dict(item) for item in request_kwargs])
+        return [{"prompts": item["prompt"], "video_path": item["output_path"]} for item in request_kwargs]
+
+
+def _make_batch_job(request_id, kwargs):
+    loop = asyncio.get_running_loop()
+    return _VideoBatchJob(
+        request_id=request_id,
+        kwargs=dict(kwargs),
+        future=loop.create_future(),
+        enqueue_time=time.perf_counter(),
+    )
+
+
+def _batch_scheduler_args(**overrides):
+    defaults = dict(
+        model_path="test-model",
+        batching_mode="dynamic",
+        batching_max_size=2,
+        batching_delay_ms=25.0,
+        enable_batching_metrics=False,
+        pipeline_config=PipelineConfig(supports_dynamic_batching=True),
+    )
+    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,18 +3,22 @@ from types import SimpleNamespace
 import warnings
 
 import pytest
+import torch
 
 from fastvideo.api import (
     GenerationRequest,
     GenerationResult,
     GeneratorConfig,
     InputConfig,
+    OutputConfig,
     SamplingConfig,
     load_run_config,
 )
 from fastvideo.api.sampling_param import SamplingParam
+from fastvideo.configs.pipelines.base import PipelineConfig
 from fastvideo.entrypoints.video_generator import VideoGenerator
 from fastvideo.fastvideo_args import WorkloadType
+from fastvideo.pipelines import ForwardBatch
 
 
 def _new_video_generator() -> VideoGenerator:
@@ -37,6 +41,25 @@ def _new_runtime_video_generator() -> VideoGenerator:
     return generator
 
 
+def _batching_fastvideo_args(**overrides):
+    defaults = dict(
+        model_path="test-model",
+        prompt_txt=None,
+        workload_type=SimpleNamespace(value="t2v"),
+        batching_mode="dynamic",
+        batching_max_size=4,
+        batching_config=None,
+        dit_cpu_offload=False,
+        dit_layerwise_offload=False,
+        output_type="latent",
+        pin_cpu_memory=False,
+        VSA_sparsity=0.0,
+        pipeline_config=PipelineConfig(supports_dynamic_batching=True),
+    )
+    defaults.update(overrides)
+    return SimpleNamespace(**defaults)
+
+
 def _patch_from_fastvideo_args(monkeypatch):
     captured = {}
 
@@ -151,6 +174,117 @@ def test_prepare_output_path_empty_prompt_fallback(tmp_path):
     assert os.path.basename(result) == "output.mp4"
 
 
+def test_generate_prepared_work_items_merges_compatible_latent_requests(monkeypatch, tmp_path):
+    vg = _new_video_generator()
+    vg.fastvideo_args = _batching_fastvideo_args()
+    calls = []
+
+    def fake_device_memory(gpu_id):
+        return 48.0
+
+    def fake_run_forward(batch, fastvideo_args):
+        calls.append(batch)
+        batch_size = len(batch.prompt) if isinstance(batch.prompt, list) else 1
+        output = torch.arange(batch_size * 4, dtype=torch.float32).reshape(batch_size, 4, 1, 1, 1)
+        return ForwardBatch(data_type=batch.data_type, output=output, extra={"peak_memory_mb": 1.0}), 0.5, 10.0
+
+    monkeypatch.setattr(
+        "fastvideo.batching.admission.BatchAdmissionController._get_device_memory_gb",
+        staticmethod(fake_device_memory),
+    )
+    monkeypatch.setattr(vg, "_run_forward_batch", fake_run_forward)
+
+    first = SamplingParam(prompt="one", height=8, width=8, num_frames=1, seed=11, return_frames=True, save_video=False)
+    second = SamplingParam(prompt="two", height=8, width=8, num_frames=1, seed=22, return_frames=True, save_video=False)
+    work_items = [
+        vg._prepare_generation_work_item("one", first, vg.fastvideo_args, output_path=str(tmp_path / "one.mp4")),
+        vg._prepare_generation_work_item("two", second, vg.fastvideo_args, output_path=str(tmp_path / "two.mp4")),
+    ]
+
+    results = vg._generate_prepared_work_items(work_items)
+
+    assert len(calls) == 1
+    assert calls[0].prompt == ["one", "two"]
+    assert calls[0].seeds == [11, 22]
+    assert [result["prompts"] for result in results] == ["one", "two"]
+    assert [result["samples"].shape for result in results] == [(1, 4, 1, 1, 1), (1, 4, 1, 1, 1)]
+
+
+def test_generate_prepared_work_items_falls_back_for_incompatible_requests(monkeypatch, tmp_path):
+    vg = _new_video_generator()
+    vg.fastvideo_args = _batching_fastvideo_args()
+    calls = []
+
+    def fake_run_forward(batch, fastvideo_args):
+        calls.append(batch)
+        output = torch.zeros((1, 4, 1, 1, 1), dtype=torch.float32)
+        return ForwardBatch(data_type=batch.data_type, output=output), 0.5, 10.0
+
+    monkeypatch.setattr(vg, "_run_forward_batch", fake_run_forward)
+
+    first = SamplingParam(prompt="one", height=8, width=8, num_frames=1, guidance_scale=1.0, save_video=False)
+    second = SamplingParam(prompt="two", height=8, width=8, num_frames=1, guidance_scale=3.0, save_video=False)
+    work_items = [
+        vg._prepare_generation_work_item("one", first, vg.fastvideo_args, output_path=str(tmp_path / "one.mp4")),
+        vg._prepare_generation_work_item("two", second, vg.fastvideo_args, output_path=str(tmp_path / "two.mp4")),
+    ]
+
+    results = vg._generate_prepared_work_items(work_items)
+
+    assert len(calls) == 2
+    assert all(isinstance(call.prompt, str) for call in calls)
+    assert [result["prompts"] for result in results] == ["one", "two"]
+
+
+def test_generate_video_batch_routes_compat_kwargs(monkeypatch, tmp_path):
+    vg = _new_video_generator()
+    vg.fastvideo_args = _batching_fastvideo_args()
+    calls = []
+
+    def fake_device_memory(gpu_id):
+        return 48.0
+
+    def fake_run_forward(batch, fastvideo_args):
+        calls.append((batch, fastvideo_args))
+        output = torch.zeros((len(batch.prompt), 4, 1, 1, 1), dtype=torch.float32)
+        return ForwardBatch(data_type=batch.data_type, output=output), 0.5, 10.0
+
+    monkeypatch.setattr(
+        "fastvideo.batching.admission.BatchAdmissionController._get_device_memory_gb",
+        staticmethod(fake_device_memory),
+    )
+    monkeypatch.setattr(vg, "_run_forward_batch", fake_run_forward)
+
+    results = vg.generate_video_batch([
+        {
+            "prompt": "one",
+            "height": 8,
+            "width": 8,
+            "num_frames": 1,
+            "embedded_cfg_scale": 7.5,
+            "save_video": False,
+            "return_frames": True,
+            "output_path": str(tmp_path / "one.mp4"),
+        },
+        {
+            "prompt": "two",
+            "height": 8,
+            "width": 8,
+            "num_frames": 1,
+            "embedded_cfg_scale": 7.5,
+            "save_video": False,
+            "return_frames": True,
+            "output_path": str(tmp_path / "two.mp4"),
+        },
+    ])
+
+    assert len(calls) == 1
+    batch, fastvideo_args = calls[0]
+    assert batch.prompt == ["one", "two"]
+    assert fastvideo_args.pipeline_config.embedded_cfg_scale == 7.5
+    assert [result["prompts"] for result in results] == ["one", "two"]
+
+
 def test_from_config_normalizes_and_translates(monkeypatch):
     captured = _patch_from_fastvideo_args(monkeypatch)
     _patch_fastvideo_args_from_kwargs(monkeypatch)
@@ -558,6 +692,46 @@ def test_generate_batched_request_fans_out_media_inputs(monkeypatch):
     ]
 
 
+def test_generate_typed_prompt_list_routes_through_dynamic_batching(monkeypatch, tmp_path):
+    generator = _new_video_generator()
+    generator.fastvideo_args = _batching_fastvideo_args()
+    _patch_sampling_param_from_pretrained(monkeypatch)
+    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), 0.5, 10.0
+
+    monkeypatch.setattr(
+        "fastvideo.batching.admission.BatchAdmissionController._get_device_memory_gb",
+        staticmethod(fake_device_memory),
+    )
+    monkeypatch.setattr(generator, "_run_forward_batch", fake_run_forward)
+
+    results = generator.generate(
+        GenerationRequest(
+            prompt=["first prompt", "second prompt"],
+            sampling=SamplingConfig(num_frames=1, height=8, width=8),
+            output=OutputConfig(
+                output_path=str(tmp_path),
+                save_video=False,
+                return_frames=True,
+            ),
+        )
+    )
+
+    assert len(calls) == 1
+    assert calls[0].prompt == ["first prompt", "second prompt"]
+    assert [result.prompt for result in results] == ["first prompt", "second prompt"]
+    assert [result.prompt_index for result in results] == [0, 1]
+    assert all(isinstance(result, GenerationResult) for result in results)
+
+
 def test_generate_batched_request_rejects_mismatched_media_inputs(monkeypatch):
     generator = _new_runtime_video_generator()
     _patch_sampling_param_from_pretrained(monkeypatch)

fastvideo/tests/modal/pr_test.py

@@ -276,7 +276,7 @@ def run_self_forcing_tests():
 @app.function(gpu="L40S:1", image=image, timeout=900, secrets=[ci_env_secret])
 def run_unit_test():
     run_test(
-        "pytest ./fastvideo/tests/api/ ./fastvideo/tests/contract/ ./fastvideo/tests/dataset/ ./fastvideo/tests/workflow/ ./fastvideo/tests/entrypoints/ ./fastvideo/tests/train/ ./fastvideo/tests/stages/ ./fastvideo/tests/ops/ ./fastvideo/tests/training/test_trackers.py ./fastvideo/tests/attention/test_sdpa_metadata_mask_contract.py --ignore=./fastvideo/tests/entrypoints/test_openai_api_integration.py --ignore=./fastvideo/tests/train/models --ignore=./fastvideo/tests/train/methods -vs"
+        "pytest ./fastvideo/tests/api/ ./fastvideo/tests/batching/ ./fastvideo/tests/contract/ ./fastvideo/tests/dataset/ ./fastvideo/tests/workflow/ ./fastvideo/tests/entrypoints/ ./fastvideo/tests/train/ ./fastvideo/tests/stages/ ./fastvideo/tests/ops/ ./fastvideo/tests/training/test_trackers.py ./fastvideo/tests/attention/test_sdpa_metadata_mask_contract.py --ignore=./fastvideo/tests/entrypoints/test_openai_api_integration.py --ignore=./fastvideo/tests/train/models --ignore=./fastvideo/tests/train/methods -vs"
     )
 
 

fastvideo/tests/ops/test_rotary_pos_embed_cache.py

@@ -1,190 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Unit tests for the get_rotary_pos_embed memoization cache."""
-
-import pytest
-import torch
-
-from fastvideo.layers.rotary_embedding import (
-    _ROTARY_POS_EMBED_CACHE,
-    _ROTARY_POS_EMBED_CACHE_MAXSIZE,
-    get_rotary_pos_embed,
-)
-
-
-def _rope_dim_list(hidden_size: int, heads_num: int) -> list[int]:
-    """Return the default 3-axis rope_dim_list used by the video DiTs."""
-    d = hidden_size // heads_num
-    return [d - 4 * (d // 6), 2 * (d // 6), 2 * (d // 6)]
-
-
-def _call(
-    rope_sizes=(21, 30, 52),
-    hidden_size=1536,
-    heads_num=12,
-    rope_dim_list="default",
-    rope_theta=10000.0,
-    dtype=torch.float64,
-    start_frame=0,
-    use_real=True,
-    **kwargs,
-):
-    """Thin wrapper around get_rotary_pos_embed with DiT-like defaults."""
-    if rope_dim_list == "default":
-        rope_dim_list = _rope_dim_list(hidden_size, heads_num)
-    return get_rotary_pos_embed(
-        rope_sizes,
-        hidden_size,
-        heads_num,
-        rope_dim_list,
-        rope_theta,
-        dtype=dtype,
-        start_frame=start_frame,
-        use_real=use_real,
-        **kwargs,
-    )
-
-
-@pytest.fixture(autouse=True)
-def _clear_cache():
-    """Isolate every test by clearing the module-level cache around it."""
-    _ROTARY_POS_EMBED_CACHE.clear()
-    yield
-    _ROTARY_POS_EMBED_CACHE.clear()
-
-
-def test_repeated_call_hits_cache():
-    """A second identical call returns the exact same tensor objects."""
-    cos1, sin1 = _call()
-    cos2, sin2 = _call()
-    assert cos1 is cos2 and sin1 is sin2
-    assert len(_ROTARY_POS_EMBED_CACHE) == 1
-
-
-def test_many_identical_calls_keep_single_entry():
-    """Many identical calls never grow the cache beyond one entry."""
-    for _ in range(10):
-        _call()
-    assert len(_ROTARY_POS_EMBED_CACHE) == 1
-
-
-@pytest.mark.parametrize(
-    "rope_sizes,hidden_size,heads_num,dtype,use_real",
-    [
-        ((21, 30, 52), 1536, 12, torch.float64, True),
-        ((21, 45, 80), 5120, 40, torch.float64, True),
-        ((1, 16, 16), 1536, 12, torch.float32, True),
-        ((4, 8, 8), 1536, 12, torch.float64, False),
-    ],
-)
-def test_cached_matches_fresh_recompute(rope_sizes, hidden_size, heads_num,
-                                        dtype, use_real):
-    """Cached tensors are bitwise-equal to a fresh uncached recompute."""
-    cos_cached, sin_cached = _call(rope_sizes=rope_sizes,
-                                   hidden_size=hidden_size,
-                                   heads_num=heads_num,
-                                   dtype=dtype,
-                                   use_real=use_real)
-    _ROTARY_POS_EMBED_CACHE.clear()
-    cos_fresh, sin_fresh = _call(rope_sizes=rope_sizes,
-                                 hidden_size=hidden_size,
-                                 heads_num=heads_num,
-                                 dtype=dtype,
-                                 use_real=use_real)
-    assert torch.equal(cos_cached, cos_fresh)
-    assert torch.equal(sin_cached, sin_fresh)
-
-
-@pytest.mark.parametrize(
-    "kwargs_a,kwargs_b",
-    [
-        ({"rope_sizes": (21, 30, 52)}, {"rope_sizes": (21, 45, 80)}),
-        ({"dtype": torch.float64}, {"dtype": torch.float32}),
-        ({"use_real": True}, {"use_real": False}),
-        ({"start_frame": 0}, {"start_frame": 3}),
-        ({"rope_theta": 10000.0}, {"rope_theta": 5000.0}),
-        ({"shard_dim": 0}, {"shard_dim": 1}),
-    ],
-)
-def test_distinct_args_create_distinct_entries(kwargs_a, kwargs_b):
-    """Any output-affecting argument difference yields a separate cache entry."""
-    _call(**kwargs_a)
-    _call(**kwargs_b)
-    assert len(_ROTARY_POS_EMBED_CACHE) == 2
-
-
-def test_none_rope_dim_list_shares_key_with_equivalent_list():
-    """None rope_dim_list and its derived explicit list map to one entry."""
-    # head_dim must be divisible by 3 for the None branch to stay valid.
-    hidden_size, heads_num = 1536, 16  # head_dim == 96 -> [32, 32, 32]
-    _call(rope_dim_list=None, hidden_size=hidden_size, heads_num=heads_num)
-    before = len(_ROTARY_POS_EMBED_CACHE)
-    _call(rope_dim_list=[32, 32, 32], hidden_size=hidden_size, heads_num=heads_num)
-    assert len(_ROTARY_POS_EMBED_CACHE) == before == 1
-
-
-def test_use_real_controls_last_dim():
-    """use_real=True spans full head_dim; use_real=False spans half."""
-    cos_full, _ = _call(use_real=True)
-    cos_half, _ = _call(use_real=False)
-    assert cos_full.shape[-1] == 128
-    assert cos_half.shape[-1] == 64
-
-
-@pytest.mark.parametrize("rope_sizes", [(1, 1, 1), (1, 30, 52), (21, 1, 1)])
-def test_degenerate_grid_shapes(rope_sizes):
-    """Degenerate single-element axes still produce a correctly sized table."""
-    cos, sin = _call(rope_sizes=rope_sizes)
-    expected = rope_sizes[0] * rope_sizes[1] * rope_sizes[2]
-    assert cos.shape[0] == expected
-    assert sin.shape[0] == expected
-
-
-def test_scalar_and_list_factors_are_hashable_and_distinct():
-    """List-valued rescale factors are hashable and keyed apart from scalars."""
-    _call(theta_rescale_factor=1.0)
-    _call(theta_rescale_factor=[1.0, 1.0, 1.0])
-    assert len(_ROTARY_POS_EMBED_CACHE) == 2
-
-
-def test_caller_device_copy_does_not_corrupt_cache():
-    """The .to()/.float() copy callers perform must not mutate cached tensors."""
-    cos, _ = _call()
-    snapshot = cos.clone()
-    _ = cos.to("cpu").float()
-    cos_again, _ = _call()
-    assert torch.equal(cos_again, snapshot)
-
-
-def test_start_frame_offsets_values():
-    """A non-zero start_frame shifts the temporal positions, changing output."""
-    cos0, _ = _call(start_frame=0)
-    cos3, _ = _call(start_frame=3)
-    assert not torch.equal(cos0, cos3)
-    assert len(_ROTARY_POS_EMBED_CACHE) == 2
-
-
-def test_cache_is_bounded_and_evicts_oldest():
-    """The cache caps at the max size and evicts the oldest entry first."""
-    # Tiny grids keep this lightweight; each start_frame is a distinct key.
-    overshoot = _ROTARY_POS_EMBED_CACHE_MAXSIZE + 4
-    for frame in range(overshoot):
-        _call(rope_sizes=(2, 2, 2), start_frame=frame)
-        assert len(_ROTARY_POS_EMBED_CACHE) <= _ROTARY_POS_EMBED_CACHE_MAXSIZE
-    assert len(_ROTARY_POS_EMBED_CACHE) == _ROTARY_POS_EMBED_CACHE_MAXSIZE
-    # The earliest-inserted frames must have been evicted; the latest survive.
-    surviving = {key[-2] for key in _ROTARY_POS_EMBED_CACHE}  # start_frame slot
-    assert overshoot - 1 in surviving
-    assert 0 not in surviving
-
-
-def test_cache_hit_refreshes_recency():
-    """Re-accessing an entry protects it from eviction over an untouched one."""
-    _call(rope_sizes=(2, 2, 2), start_frame=0)  # entry we will keep hot
-    for frame in range(1, _ROTARY_POS_EMBED_CACHE_MAXSIZE):
-        _call(rope_sizes=(2, 2, 2), start_frame=frame)
-    assert len(_ROTARY_POS_EMBED_CACHE) == _ROTARY_POS_EMBED_CACHE_MAXSIZE
-    _call(rope_sizes=(2, 2, 2), start_frame=0)   # hit -> frame 0 becomes most recent
-    _call(rope_sizes=(2, 2, 2), start_frame=99)  # miss -> evicts now-oldest (frame 1)
-    surviving = {key[-2] for key in _ROTARY_POS_EMBED_CACHE}
-    assert 0 in surviving
-    assert 1 not in surviving

fastvideo/tests/performance/compare_baseline.py

@@ -89,14 +89,6 @@ COMPARISON_IDENTITY_KEYS = (
     "hardware_profile_id",
     "software_profile_id",
 )
-# Comparator verdicts (issue #1532), stored as record["comparator_status"].
-STATUS_PASS = "PASS"
-STATUS_REGRESSION = "REGRESSION"
-STATUS_CALIBRATION_NEEDED = "CALIBRATION_NEEDED"
-STATUS_RECIPE_MISMATCH = "RECIPE_MISMATCH"
-STATUS_INFRA_ERROR = "INFRA_ERROR"
-# Reserved for the promoted-baseline workflow; never emitted by this comparator.
-STATUS_QUALITY_BLOCKED = "QUALITY_BLOCKED"
 
 
 def _should_persist_tracking() -> bool:
@@ -351,90 +343,6 @@ def _check_regressions(
     return failures
 
 
-def _recipe_mismatch_fingerprints(
-    record: dict[str, Any],
-    identity_filters: dict[str, str],
-) -> list[str]:
-    """Return baseline recipe fingerprints for the same variant cohort.
-
-    Non-empty means the same (workload, variant, version, hardware, software)
-    cohort has baseline history under a DIFFERENT recipe fingerprint: the
-    recipe changed without being represented as a new variant.
-    """
-    if not identity_filters:
-        return []
-    variant_filters = {key: value for key, value in identity_filters.items() if key != "recipe_fingerprint"}
-    same_variant = load_records_for_model(
-        TRACKING_ROOT,
-        record["model_id"],
-        record["gpu_type"],
-        **variant_filters,
-        successful_only=True,
-        baseline_eligible_only=True,
-    )
-    fingerprints = {str(r.get("recipe_fingerprint")) for r in same_variant}
-    return sorted(fingerprints - {identity_filters["recipe_fingerprint"]})
-
-
-def _compare_record(
-    record: dict[str, Any],
-    metric_policies: tuple[MetricPolicy, ...],
-) -> tuple[list[str], list[dict[str, Any]]]:
-    """Compare one normalized record against its exact-identity cohort.
-
-    Sets ``record["comparator_status"]`` (and ``record["baseline_status"]``
-    for cohort bookkeeping) and returns ``(failures, baseline_records)``.
-    """
-    model = record.get("model_id", "unknown")
-    try:
-        identity_filters = _comparison_identity_filters(record)
-        baseline_records = load_records_for_model(
-            TRACKING_ROOT,
-            record["model_id"],
-            record["gpu_type"],
-            **identity_filters,
-            last_n=5,
-            successful_only=True,
-            baseline_eligible_only=True,
-        )
-        if baseline_records:
-            record["baseline_status"] = "compared"
-            failures = _check_regressions(record, baseline_records, metric_policies)
-            record["comparator_status"] = STATUS_REGRESSION if failures else STATUS_PASS
-            return failures, baseline_records
-
-        mismatched = _recipe_mismatch_fingerprints(record, identity_filters)
-    except Exception as exc:
-        record["comparator_status"] = STATUS_INFRA_ERROR
-        return [f"{model} baseline comparison hit an infra error: {exc}"], []
-
-    if mismatched:
-        record["comparator_status"] = STATUS_RECIPE_MISMATCH
-        return [
-            f"{model} recipe fingerprint {record.get('recipe_fingerprint')} does not "
-            f"match baseline fingerprint(s) {', '.join(mismatched)} for variant "
-            f"{record.get('variant_id')}. Represent recipe changes as a new "
-            "variant_id, or reseed the baseline."
-        ], []
-
-    # No comparable baseline anywhere: a brand-new cohort. Make that loud and
-    # machine-readable instead of an indistinguishable pass, so a cohort shift
-    # (intended or accidental, e.g. an identity-field change) never silently
-    # blinds the comparison — and never silently seeds a passing baseline.
-    record["baseline_status"] = "initialized_new_cohort"
-    record["comparator_status"] = STATUS_CALIBRATION_NEEDED
-    print("=" * 72)
-    print(f"CALIBRATION_NEEDED: no comparable baseline for {model} on "
-          f"{record.get('gpu_type', 'unknown')}"
-          f"{_format_identity_filters(identity_filters)}")
-    print("Regression gating is INACTIVE for this record and it will NOT "
-          "seed a baseline. Seed the cohort explicitly via the reseed "
-          "workflow. If this cohort shift is unexpected, check the identity "
-          "fields above.")
-    print("=" * 72)
-    return [], []
-
-
 def _compact_value(value: float | None, precision: int = 3) -> str:
     if value is None:
         return "n/a"
@@ -486,7 +394,6 @@ def _build_summary_row(
     return {
         "model_id": record["model_id"],
         "gpu_type": record["gpu_type"],
-        "comparator_status": record.get("comparator_status", STATUS_PASS),
         "baseline_n": len(baseline_records),
         "metrics": metric_values,
         "worst_regression_pct": worst_regression_pct,
@@ -528,10 +435,7 @@ def _build_markdown_summary(
             else "none"
         )
         failing_metrics = ", ".join(row["failing_metrics"]) if row["failing_metrics"] else "none"
-        status = row["comparator_status"]
-        if row["failed"] and status == STATUS_PASS:
-            # Baseline comparison passed but the fixed-threshold phase failed.
-            status = "FAIL"
+        status = "FAIL" if row["failed"] else "PASS"
 
         lines.append(f"| {row['model_id']} | {row['gpu_type']} | "
                      f"{row['baseline_n']} | "
@@ -597,26 +501,46 @@ def main() -> int:
         if identity_filters is None:
             # Records without the full v2 identity block skip rolling-baseline
             # comparison entirely: only the static thresholds gate them and
-            # they never become baseline eligible. Nothing was compared and
-            # nothing failed, so the comparator verdict is PASS;
-            # baseline_status keeps the skip machine-readable.
+            # they never become baseline eligible.
             record["baseline_status"] = "skipped_missing_identity"
-            record["comparator_status"] = STATUS_PASS
             failures: list[str] = []
         else:
-            failures, baseline_records = _compare_record(record, metric_policies)
+            baseline_records = load_records_for_model(
+                TRACKING_ROOT,
+                record["model_id"],
+                record["gpu_type"],
+                **identity_filters,
+                last_n=5,
+                successful_only=True,
+                baseline_eligible_only=True,
+            )
+
+            if not baseline_records:
+                # A brand-new cohort has NO regression gating until history
+                # accumulates — make that loud and machine-readable instead of
+                # an indistinguishable pass, so a cohort shift (intended or
+                # accidental, e.g. an identity-field change) never silently
+                # blinds the comparison.
+                record["baseline_status"] = "initialized_new_cohort"
+                print("=" * 72)
+                print(f"WARNING: NO BASELINE — initializing a NEW cohort for "
+                      f"{record['model_id']} on {record['gpu_type']}"
+                      f"{_format_identity_filters(identity_filters)}")
+                print("Regression gating is INACTIVE for this cohort until "
+                      "baseline history accumulates. If this cohort shift is "
+                      "unexpected, check the identity fields above.")
+                print("=" * 72)
+                failures = []
+            else:
+                record["baseline_status"] = "compared"
+                failures = _check_regressions(record, baseline_records, metric_policies)
         if static_threshold_failed:
             failures.append(f"{record['model_id']} fixed-threshold phase failed "
                             f"(PERF_PYTEST_RC={os.environ.get('PERF_PYTEST_RC')})")
 
         record["success"] = not failures
-        # Only compared-and-PASS records may advance the rolling baseline: a
-        # CALIBRATION_NEEDED record must never silently seed a new cohort, and
-        # records without the v2 identity block never become eligible.
         record["baseline_eligible"] = (
-            identity_filters is not None
-            and record["comparator_status"] == STATUS_PASS
-            and _is_baseline_eligible(record["run_source"], record["success"])
+            identity_filters is not None and _is_baseline_eligible(record["run_source"], record["success"])
         )
         all_failures.extend(failures)
 

fastvideo/tests/performance/test_compare_baseline_policy.py

@@ -421,196 +421,3 @@ def test_informational_metric_remains_visible_without_failing():
     assert row["metrics"]["throughput"]["regressed"] is False
     assert row["threshold_exceeded_metrics"] == ["throughput"]
     assert row["failing_metrics"] == []
-
-
-def _v2_raw_result(**overrides):
-    raw = _raw_result()
-    raw.update({
-        "workload_id": "wan-t2v",
-        "variant_id": "1.3b-sp2",
-        "benchmark_version": 2,
-        "recipe_fingerprint": "recipe-1",
-        "hardware_profile_id": "hw-1",
-        "software_profile_id": "sw-1",
-    })
-    raw.update(overrides)
-    return raw
-
-
-def _v2_baseline_record(**overrides):
-    record = {
-        "gpu_type": "NVIDIA L40S",
-        "workload_id": "wan-t2v",
-        "variant_id": "1.3b-sp2",
-        "benchmark_version": 2,
-        "recipe_fingerprint": "recipe-1",
-        "hardware_profile_id": "hw-1",
-        "software_profile_id": "sw-1",
-        "latency": 10.0,
-        "throughput": 4.5,
-        "memory": 10000.0,
-        "timestamp": "2026-06-15T00:00:00+00:00",
-        "success": True,
-        "run_source": "scheduled_main",
-        "baseline_eligible": True,
-    }
-    record.update(overrides)
-    return record
-
-
-def _run_compare(monkeypatch, tmp_path, raw_result, baseline_records):
-    """Run compare_baseline.main() against a local tracking root.
-
-    Returns (exit_code, normalized_record, markdown_report).
-    """
-    results_dir = tmp_path / "results"
-    results_dir.mkdir()
-    (results_dir / "perf_current.json").write_text(json.dumps(raw_result))
-
-    model_dir = tmp_path / "tracking" / raw_result["benchmark_id"]
-    model_dir.mkdir(parents=True)
-    for index, record in enumerate(baseline_records):
-        (model_dir / f"rec{index}.json").write_text(json.dumps(record))
-
-    reports_dir = tmp_path / "reports"
-    monkeypatch.setattr(compare_baseline, "RESULTS_DIR", str(results_dir))
-    monkeypatch.setattr(compare_baseline, "TRACKING_ROOT", str(tmp_path / "tracking"))
-    monkeypatch.setattr(compare_baseline, "PERF_REPORTS_DIR", str(reports_dir))
-    monkeypatch.setattr(compare_baseline, "UPLOAD_POLICY", "never")
-    monkeypatch.setattr(compare_baseline, "sync_from_hf", lambda *args, **kwargs: None)
-    monkeypatch.setenv("PERF_RUN_SOURCE", "scheduled_main")
-    monkeypatch.delenv("PERF_PYTEST_RC", raising=False)
-    monkeypatch.delenv("GITHUB_STEP_SUMMARY", raising=False)
-
-    exit_code = compare_baseline.main()
-
-    normalized_paths = list((reports_dir / "results").glob("normalized_perf_*.json"))
-    assert len(normalized_paths) == 1
-    markdown = "\n".join(path.read_text() for path in reports_dir.glob("perf_*.md"))
-    return exit_code, json.loads(normalized_paths[0].read_text()), markdown
-
-
-def test_main_pass_with_comparable_baseline(monkeypatch, tmp_path):
-    exit_code, record, markdown = _run_compare(
-        monkeypatch, tmp_path, _v2_raw_result(), [_v2_baseline_record()])
-
-    assert exit_code == 0
-    assert record["comparator_status"] == "PASS"
-    assert record["baseline_status"] == "compared"
-    assert record["success"] is True
-    # Scheduled-main PASS records keep advancing the rolling baseline.
-    assert record["baseline_eligible"] is True
-    assert "| PASS |" in markdown
-
-
-def test_main_regression_on_gated_metric_fails_ci(monkeypatch, tmp_path):
-    exit_code, record, markdown = _run_compare(
-        monkeypatch, tmp_path, _v2_raw_result(avg_generation_time_s=20.0),
-        [_v2_baseline_record()])
-
-    assert exit_code == 1
-    assert record["comparator_status"] == "REGRESSION"
-    assert record["success"] is False
-    assert record["baseline_eligible"] is False
-    assert "| REGRESSION |" in markdown
-
-
-def test_main_missing_baseline_is_calibration_needed_and_does_not_seed(monkeypatch, tmp_path):
-    exit_code, record, markdown = _run_compare(
-        monkeypatch, tmp_path, _v2_raw_result(), [])
-
-    assert exit_code == 0
-    assert record["comparator_status"] == "CALIBRATION_NEEDED"
-    assert record["baseline_status"] == "initialized_new_cohort"
-    assert record["success"] is True
-    # Visible, but never silently seeds a passing baseline.
-    assert record["baseline_eligible"] is False
-    assert "| CALIBRATION_NEEDED |" in markdown
-
-
-def test_main_recipe_change_without_new_variant_is_recipe_mismatch(monkeypatch, tmp_path):
-    exit_code, record, markdown = _run_compare(
-        monkeypatch, tmp_path, _v2_raw_result(recipe_fingerprint="recipe-2"),
-        [_v2_baseline_record()])
-
-    assert exit_code == 1
-    assert record["comparator_status"] == "RECIPE_MISMATCH"
-    assert record["success"] is False
-    assert record["baseline_eligible"] is False
-    assert "| RECIPE_MISMATCH |" in markdown
-
-
-def test_main_recipe_change_as_new_variant_is_calibration_needed(monkeypatch, tmp_path):
-    exit_code, record, _ = _run_compare(
-        monkeypatch, tmp_path,
-        _v2_raw_result(variant_id="1.3b-sp2-r2", recipe_fingerprint="recipe-2"),
-        [_v2_baseline_record()])
-
-    assert exit_code == 0
-    assert record["comparator_status"] == "CALIBRATION_NEEDED"
-
-
-def test_main_v2_record_never_compares_against_v1_baselines(monkeypatch, tmp_path):
-    legacy_v1_baseline = {
-        "gpu_type": "NVIDIA L40S",
-        # Would be a >100% latency regression if the comparator (wrongly)
-        # matched the v2 record against v1 history.
-        "latency": 1.0,
-        "timestamp": "2026-06-15T00:00:00+00:00",
-        "success": True,
-    }
-
-    exit_code, record, _ = _run_compare(
-        monkeypatch, tmp_path, _v2_raw_result(), [legacy_v1_baseline])
-
-    assert exit_code == 0
-    assert record["comparator_status"] == "CALIBRATION_NEEDED"
-
-
-def test_main_legacy_v1_record_skips_comparison_with_pass_verdict(monkeypatch, tmp_path):
-    legacy_v1_baseline = {
-        "gpu_type": "NVIDIA L40S",
-        # Would be a >100% latency regression if the comparator (wrongly)
-        # compared the identity-less v1 record against this history.
-        "latency": 1.0,
-        "timestamp": "2026-06-15T00:00:00+00:00",
-        "success": True,
-    }
-
-    exit_code, record, markdown = _run_compare(
-        monkeypatch, tmp_path, _raw_result(), [legacy_v1_baseline])
-
-    assert exit_code == 0
-    assert record["comparator_status"] == "PASS"
-    assert record["baseline_status"] == "skipped_missing_identity"
-    assert record["baseline_eligible"] is False
-    assert "| PASS |" in markdown
-
-
-def test_main_partial_identity_skips_comparison(monkeypatch, tmp_path):
-    raw = _v2_raw_result()
-    del raw["software_profile_id"]
-
-    exit_code, record, _ = _run_compare(monkeypatch, tmp_path, raw, [])
-
-    assert exit_code == 0
-    assert record["comparator_status"] == "PASS"
-    assert record["baseline_status"] == "skipped_missing_identity"
-    assert record["success"] is True
-    assert record["baseline_eligible"] is False
-
-
-def test_main_baseline_load_failure_is_infra_error_and_fails_ci(monkeypatch, tmp_path):
-    def _boom(*_args, **_kwargs):
-        raise OSError("hf store unavailable")
-
-    monkeypatch.setattr(compare_baseline, "load_records_for_model", _boom)
-
-    exit_code, record, markdown = _run_compare(
-        monkeypatch, tmp_path, _v2_raw_result(), [])
-
-    assert exit_code == 1
-    assert record["comparator_status"] == "INFRA_ERROR"
-    assert record["success"] is False
-    assert record["baseline_eligible"] is False
-    assert "| INFRA_ERROR |" in markdown

fastvideo/tests/ssim/inference_similarity_utils.py

@@ -62,29 +62,12 @@ def resolve_inference_device_reference_folder(logger: Logger) -> str:
     return device_reference_folder
 
 
-def _find_reference_media(
-    reference_folder: str,
-    prompt: str,
-    *,
-    media_extension: str,
-) -> str:
-    """Pick a reference file whose basename contains the prompt prefix."""
+def _find_reference_video(reference_folder: str, prompt: str) -> str:
     prompt_prefix = prompt[:100].strip()
-    allowed = (media_extension.lower(), ".mp4", ".png", ".jpg", ".jpeg")
-    matches: list[str] = []
     for filename in os.listdir(reference_folder):
-        low = filename.lower()
-        if not any(low.endswith(ext) for ext in allowed):
-            continue
-        if prompt_prefix in filename:
-            matches.append(filename)
-    if not matches:
-        raise FileNotFoundError("Reference media missing")
-    preferred = media_extension.lower().lstrip(".")
-    for name in matches:
-        if name.lower().endswith(f".{preferred}"):
-            return os.path.join(reference_folder, name)
-    return os.path.join(reference_folder, matches[0])
+        if filename.endswith(".mp4") and prompt_prefix in filename:
+            return os.path.join(reference_folder, filename)
+    raise FileNotFoundError("Reference video missing")
 
 
 def _remove_stale_generated_video(output_dir: str, output_video_name: str) -> None:
@@ -97,23 +80,21 @@ def _assert_similarity(
     *,
     logger: Logger,
     output_dir: str,
-    output_media_name: str,
+    output_video_name: str,
     reference_folder: str,
     prompt: str,
     num_inference_steps: int,
     min_acceptable_ssim: float,
     model_id: str,
     attention_backend_name: str,
-    media_extension: str,
 ) -> None:
-    generated_media_path = os.path.join(output_dir, output_media_name)
-    artifact_kind = "image" if media_extension.lower() in (".png", ".jpg", ".jpeg") else "video"
+    generated_video_path = os.path.join(output_dir, output_video_name)
     if not os.path.exists(reference_folder):
         logger.error("Reference folder missing: %s", reference_folder)
         xfail_missing_reference_in_bootstrap_mode(
-            generated_artifact_path=generated_media_path,
+            generated_artifact_path=generated_video_path,
             reference_folder=reference_folder,
-            artifact_kind=artifact_kind,
+            artifact_kind="video",
         )
         error_msg = (
             f"Reference video folder does not exist: {reference_folder}\n"
@@ -123,32 +104,28 @@ def _assert_similarity(
         raise FileNotFoundError(error_msg)
 
     try:
-        reference_media_path = _find_reference_media(
-            reference_folder,
-            prompt,
-            media_extension=media_extension,
-        )
+        reference_video_path = _find_reference_video(reference_folder, prompt)
     except FileNotFoundError as error:
         logger.error(
-            "Reference media not found for prompt: %s with backend: %s",
+            "Reference video not found for prompt: %s with backend: %s",
             prompt,
             attention_backend_name,
         )
         xfail_missing_reference_in_bootstrap_mode(
-            generated_artifact_path=generated_media_path,
+            generated_artifact_path=generated_video_path,
             reference_folder=reference_folder,
-            artifact_kind=artifact_kind,
+            artifact_kind="video",
         )
         raise error
 
     logger.info(
         "Computing SSIM between %s and %s",
-        reference_media_path,
-        generated_media_path,
+        reference_video_path,
+        generated_video_path,
     )
     ssim_values = compute_video_ssim_torchvision(
-        reference_media_path,
-        generated_media_path,
+        reference_video_path,
+        generated_video_path,
         use_ms_ssim=True,
     )
 
@@ -159,8 +136,8 @@ def _assert_similarity(
     success = write_ssim_results(
         output_dir,
         ssim_values,
-        reference_media_path,
-        generated_media_path,
+        reference_video_path,
+        generated_video_path,
         num_inference_steps,
         prompt,
     )
@@ -245,7 +222,6 @@ def run_text_to_video_similarity_test(
     min_acceptable_ssim: float,
     init_kwargs_override: dict[str, object] | None = None,
     generation_kwargs_override: dict[str, object] | None = None,
-    media_extension: str = ".mp4",
 ) -> None:
     with attention_backend(attention_backend_name):
         output_dir = build_generated_output_dir(
@@ -254,9 +230,9 @@ def run_text_to_video_similarity_test(
             model_id,
             attention_backend_name,
         )
-        output_media_name = f"{prompt[:100].strip()}{media_extension}"
+        output_video_name = f"{prompt[:100].strip()}.mp4"
         os.makedirs(output_dir, exist_ok=True)
-        _remove_stale_generated_video(output_dir, output_media_name)
+        _remove_stale_generated_video(output_dir, output_video_name)
 
         params_map = select_ssim_params(
             default_params_map,
@@ -298,14 +274,13 @@ def run_text_to_video_similarity_test(
     _assert_similarity(
         logger=logger,
         output_dir=output_dir,
-        output_media_name=output_media_name,
+        output_video_name=output_video_name,
         reference_folder=reference_folder,
         prompt=prompt,
         num_inference_steps=num_inference_steps,
         min_acceptable_ssim=min_acceptable_ssim,
         model_id=model_id,
         attention_backend_name=attention_backend_name,
-        media_extension=media_extension,
     )
 
 
@@ -331,9 +306,9 @@ def run_image_to_video_similarity_test(
             model_id,
             attention_backend_name,
         )
-        output_media_name = f"{prompt[:100].strip()}.mp4"
+        output_video_name = f"{prompt[:100].strip()}.mp4"
         os.makedirs(output_dir, exist_ok=True)
-        _remove_stale_generated_video(output_dir, output_media_name)
+        _remove_stale_generated_video(output_dir, output_video_name)
 
         params_map = select_ssim_params(
             default_params_map,
@@ -376,12 +351,11 @@ def run_image_to_video_similarity_test(
     _assert_similarity(
         logger=logger,
         output_dir=output_dir,
-        output_media_name=output_media_name,
+        output_video_name=output_video_name,
         reference_folder=reference_folder,
         prompt=prompt,
         num_inference_steps=num_inference_steps,
         min_acceptable_ssim=min_acceptable_ssim,
         model_id=model_id,
         attention_backend_name=attention_backend_name,
-        media_extension=".mp4",
     )

fastvideo/tests/ssim/test_flux_t2i_similarity.py

@@ -1,134 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-
-from __future__ import annotations
-
-import os
-
-import pytest
-import torch
-
-from fastvideo.api.flux import FluxSamplingParam
-from fastvideo.logger import init_logger
-from fastvideo.tests.ssim.inference_similarity_utils import (
-    run_text_to_video_similarity_test,
-)
-from fastvideo.tests.ssim.reference_utils import (
-    get_cuda_device_name,
-    resolve_device_reference_folder,
-)
-
-logger = init_logger(__name__)
-
-REQUIRED_GPUS = 1
-
-# MS-SSIM gate (see module docstring).
-FLUX_T2I_MIN_SSIM = 0.98
-
-FLUX_MODEL_PATH = os.getenv(
-    "FLUX_T2I_MODEL_DIR",
-    "black-forest-labs/FLUX.1-dev",
-)
-
-device_reference_folder = resolve_device_reference_folder(
-    (
-        ("A40", "A40"),
-        ("L40S", "L40S"),
-        ("H100", "H100"),
-        ("H200", "H200"),
-        ("RTX 4090", "RTX4090"),
-        ("4090", "RTX4090"),
-    ),
-    device_name=get_cuda_device_name(),
-    fallback_device_prefix="L40S",
-    logger=logger,
-)
-
-# Folder token must match Hub path with slashes → double underscore (SD3.5
-# pattern in ``test_sd35_similarity.py``).
-MODEL_ID = "black-forest-labs__FLUX.1-dev"
-
-TEST_PROMPTS = [
-    "a photo of a cat",
-]
-
-FLUX_DEFAULT_PARAMS: dict[str, object] = {
-    "num_gpus": 1,
-    "model_path": FLUX_MODEL_PATH,
-    "sp_size": 1,
-    "tp_size": 1,
-    "height": 256,
-    "width": 256,
-    "num_frames": 1,
-    "fps": 1,
-    "num_inference_steps": 8,
-    "guidance_scale": 3.5,
-    "seed": 0,
-}
-
-_flux_full_defaults = FluxSamplingParam()
-FLUX_FULL_QUALITY_PARAMS: dict[str, object] = {
-    "num_gpus": 1,
-    "model_path": FLUX_MODEL_PATH,
-    "sp_size": 1,
-    "tp_size": 1,
-    "height": _flux_full_defaults.height,
-    "width": _flux_full_defaults.width,
-    "num_frames": 1,
-    "fps": _flux_full_defaults.fps,
-    "num_inference_steps": _flux_full_defaults.num_inference_steps,
-    "guidance_scale": _flux_full_defaults.guidance_scale,
-    "seed": _flux_full_defaults.seed,
-}
-
-FLUX_MODEL_TO_PARAMS = {
-    MODEL_ID: FLUX_DEFAULT_PARAMS,
-}
-FLUX_FULL_QUALITY_MODEL_TO_PARAMS = {
-    MODEL_ID: FLUX_FULL_QUALITY_PARAMS,
-}
-
-
-@pytest.mark.skipif(
-    not torch.cuda.is_available(),
-    reason="FLUX T2I SSIM test requires CUDA",
-)
-@pytest.mark.parametrize("prompt", TEST_PROMPTS)
-@pytest.mark.parametrize("attention_backend_name", ["TORCH_SDPA"])
-@pytest.mark.parametrize("model_id", list(FLUX_MODEL_TO_PARAMS.keys()))
-def test_flux_t2i_similarity(
-    prompt: str,
-    attention_backend_name: str,
-    model_id: str,
-) -> None:
-    is_hf_repo = "/" in FLUX_MODEL_PATH and not FLUX_MODEL_PATH.startswith("/")
-    if not is_hf_repo and not os.path.isdir(FLUX_MODEL_PATH):
-        pytest.skip(
-            f"FLUX weights not found at {FLUX_MODEL_PATH} "
-            f"(set FLUX_T2I_MODEL_DIR to override)"
-        )
-
-    run_text_to_video_similarity_test(
-        logger=logger,
-        script_dir=os.path.dirname(os.path.abspath(__file__)),
-        device_reference_folder=device_reference_folder,
-        prompt=prompt,
-        attention_backend_name=attention_backend_name,
-        model_id=model_id,
-        default_params_map=FLUX_MODEL_TO_PARAMS,
-        full_quality_params_map=FLUX_FULL_QUALITY_MODEL_TO_PARAMS,
-        min_acceptable_ssim=FLUX_T2I_MIN_SSIM,
-        media_extension=".png",
-        init_kwargs_override={
-            "workload_type": "t2i",
-            "use_fsdp_inference": False,
-            "text_encoder_cpu_offload": False,
-            "vae_cpu_offload": False,
-            "image_encoder_cpu_offload": False,
-            "pin_cpu_memory": False,
-        },
-        generation_kwargs_override={
-            "save_video": True,
-            "use_embedded_guidance": True,
-            "true_cfg_scale": 1.0,
-        },
-    )

fastvideo/tests/ssim/test_glm_image_similarity.py

@@ -1,141 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""SSIM-based regression test for GLM-Image generation."""
-from __future__ import annotations
-
-import os
-from pathlib import Path
-
-import pytest
-import torch
-
-from fastvideo.logger import init_logger
-from fastvideo.tests.ssim.inference_similarity_utils import (
-    run_text_to_video_similarity_test,
-)
-from fastvideo.tests.ssim.reference_utils import (
-    get_cuda_device_name,
-    resolve_device_reference_folder,
-)
-
-logger = init_logger(__name__)
-
-REQUIRED_GPUS = 1
-
-REPO_ROOT = Path(__file__).resolve().parents[3]
-LOCAL_WEIGHTS_DIR = Path(
-    os.getenv("GLM_IMAGE_LOCAL_WEIGHTS_DIR",
-              REPO_ROOT / "official_weights" / "glm_image"))
-GLM_IMAGE_MODEL_PATH = os.getenv("GLM_IMAGE_MODEL_DIR", str(LOCAL_WEIGHTS_DIR))
-
-device_reference_folder = resolve_device_reference_folder(
-    (
-        ("A40", "A40"),
-        ("L40S", "L40S"),
-        ("H100", "H100"),
-        ("H200", "H200"),
-        ("B200", "B200"),
-    ),
-    device_name=get_cuda_device_name(),
-    fallback_device_prefix="L40S",
-    logger=logger,
-)
-
-MODEL_ID = "zai-org__GLM-Image"
-
-TEST_PROMPTS = [
-    "A beautiful landscape photography with rolling hills, "
-    "a winding river, and a vibrant sunset in the background. "
-    "Warm golden light, photorealistic style.",
-]
-
-GLM_IMAGE_PARAMS = {
-    "num_gpus": 1,
-    "model_path": GLM_IMAGE_MODEL_PATH,
-    "sp_size": 1,
-    "tp_size": 1,
-    "height": 256,
-    "width": 256,
-    "num_frames": 1,
-    "fps": 1,
-    "num_inference_steps": 4,
-    "guidance_scale": 1.5,
-    "seed": 0,
-    "neg_prompt": "",
-}
-
-GLM_IMAGE_FULL_QUALITY_PARAMS = {
-    "num_gpus": 1,
-    "model_path": GLM_IMAGE_MODEL_PATH,
-    "sp_size": 1,
-    "tp_size": 1,
-    "height": 1024,
-    "width": 1024,
-    "num_frames": 1,
-    "fps": 1,
-    "num_inference_steps": 50,
-    "guidance_scale": 1.5,
-    "seed": 0,
-    "neg_prompt": "",
-}
-
-GLM_IMAGE_MODEL_TO_PARAMS = {
-    MODEL_ID: GLM_IMAGE_PARAMS,
-}
-GLM_IMAGE_FULL_QUALITY_MODEL_TO_PARAMS = {
-    MODEL_ID: GLM_IMAGE_FULL_QUALITY_PARAMS,
-}
-
-
-def _has_weights() -> bool:
-    required = ["transformer", "vae", "text_encoder",
-                "vision_language_encoder", "processor", "tokenizer",
-                "scheduler"]
-    return all((LOCAL_WEIGHTS_DIR / r).exists() for r in required)
-
-
-def _upstream_glm_image_available() -> bool:
-    try:
-        import transformers
-    except ImportError:
-        return False
-    return hasattr(transformers, "GlmImageForConditionalGeneration")
-
-
-@pytest.mark.skipif(
-    not torch.cuda.is_available(),
-    reason="GLM-Image SSIM test requires CUDA",
-)
-@pytest.mark.skipif(
-    not _has_weights(),
-    reason=f"GLM-Image full weights not found at {LOCAL_WEIGHTS_DIR}.",
-)
-@pytest.mark.skipif(
-    not _upstream_glm_image_available(),
-    reason="GLM-Image needs transformers>=5.0.0rc0 (ships the AR encoder).",
-)
-@pytest.mark.parametrize("prompt", TEST_PROMPTS)
-@pytest.mark.parametrize("attention_backend_name", ["TORCH_SDPA"])
-@pytest.mark.parametrize("model_id", list(GLM_IMAGE_MODEL_TO_PARAMS.keys()))
-def test_glm_image_similarity(
-    prompt: str,
-    attention_backend_name: str,
-    model_id: str,
-) -> None:
-    run_text_to_video_similarity_test(
-        logger=logger,
-        script_dir=os.path.dirname(os.path.abspath(__file__)),
-        device_reference_folder=device_reference_folder,
-        prompt=prompt,
-        attention_backend_name=attention_backend_name,
-        model_id=model_id,
-        default_params_map=GLM_IMAGE_MODEL_TO_PARAMS,
-        full_quality_params_map=GLM_IMAGE_FULL_QUALITY_MODEL_TO_PARAMS,
-        min_acceptable_ssim=0.98,
-        init_kwargs_override={
-            "trust_remote_code": True,
-            "use_fsdp_inference": False,
-        },
-        generation_kwargs_override={
-            "save_video": True,
-        },
-    )

fastvideo/tests/stages/test_input_validation_batching.py

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

fastvideo/tests/stages/test_text_encoding.py

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

fastvideo/tests/training/distill/test_anyflow_onpolicy.py

@@ -1,231 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""AnyFlow on-policy method tests (CPU-only, no Wan instantiation).
-
-The full AnyFlowMethod requires a real student/teacher/critic trio plus
-DMD2's optimizer wiring — too heavyweight for a unit test. These tests
-exercise the rollout-shape helpers and source-level invariants via
-``object.__new__`` bypassing of ``__init__``.
-"""
-
-from __future__ import annotations
-
-import inspect
-from types import SimpleNamespace
-from typing import Any
-
-import pytest
-import torch
-
-from fastvideo.train.methods.distribution_matching.anyflow import AnyFlowMethod
-
-
-# ---------------------------------------------------------------------------
-# Helpers: build a "naked" AnyFlowMethod that skips __init__.
-# ---------------------------------------------------------------------------
-
-
-def _naked_method(
-    *,
-    student_sample_steps: int = 4,
-    use_mean_velocity: bool = True,
-    t_list_override: list[float] | None = None,
-    denoising_step_list: list[float] | None = None,
-) -> AnyFlowMethod:
-    method = AnyFlowMethod.__new__(AnyFlowMethod)
-    method._student_sample_steps = int(student_sample_steps)  # type: ignore[attr-defined]
-    method._use_mean_velocity = bool(use_mean_velocity)  # type: ignore[attr-defined]
-    method._t_list_override = (  # type: ignore[attr-defined]
-        list(t_list_override) if t_list_override else None)
-    method._dmd_score_r = 0.0  # type: ignore[attr-defined]
-    method._real_score_guidance = 1.0  # type: ignore[attr-defined]
-    method.cuda_generator = None  # type: ignore[attr-defined]
-    method._cfg_uncond = None  # type: ignore[attr-defined]
-    method._denoising_step_list_cache = None  # type: ignore[attr-defined]
-
-    # Stub _get_denoising_step_list — DMD2 reads method_config but for these
-    # focused tests we want a deterministic schedule.
-    raw = denoising_step_list or [999, 750, 500, 250]
-    cached = torch.tensor(raw, dtype=torch.long)
-
-    def _stub(self, device: torch.device) -> torch.Tensor:
-        return cached.to(device=device)
-
-    bound = _stub.__get__(method, AnyFlowMethod)
-    method._get_denoising_step_list = bound  # type: ignore[assignment]
-    return method
-
-
-# ---------------------------------------------------------------------------
-# Schedule construction.
-# ---------------------------------------------------------------------------
-
-
-def test_get_rollout_schedule_uses_t_list_override_verbatim() -> None:
-    method = _naked_method(
-        t_list_override=[999.0, 937.0, 833.0, 624.0, 0.0])
-    schedule = method._get_rollout_schedule(device=torch.device("cpu"))
-    torch.testing.assert_close(
-        schedule,
-        torch.tensor([999.0, 937.0, 833.0, 624.0, 0.0],
-                     dtype=torch.float32),
-    )
-
-
-def test_get_rollout_schedule_falls_back_to_denoising_step_list() -> None:
-    method = _naked_method(
-        denoising_step_list=[999, 750, 500, 250])
-    schedule = method._get_rollout_schedule(device=torch.device("cpu"))
-    # Tail must be a 0 boundary so the final Euler step lands at t=0.
-    assert float(schedule[-1].item()) == 0.0
-    # Original step list preserved at the front.
-    torch.testing.assert_close(
-        schedule[:4],
-        torch.tensor([999.0, 750.0, 500.0, 250.0], dtype=torch.float32),
-    )
-
-
-def test_get_rollout_schedule_does_not_double_append_zero() -> None:
-    method = _naked_method(
-        denoising_step_list=[999, 750, 500, 0])
-    schedule = method._get_rollout_schedule(device=torch.device("cpu"))
-    assert schedule.numel() == 4  # no extra boundary inserted.
-
-
-# ---------------------------------------------------------------------------
-# t_list_override validation.
-# ---------------------------------------------------------------------------
-
-
-def test_anyflow_method_rejects_ascending_t_list_override() -> None:
-    """__init__ validates that t_list_override is descending. We test the
-    validation logic by stitching together a minimal cfg + role_models
-    path; if construction fails for an unrelated reason we still catch
-    the descending check via the explicit error message."""
-    src = inspect.getsource(AnyFlowMethod.__init__)
-    assert 't_list_override must be descending' in src
-    assert 'descending' in src
-
-
-def test_anyflow_method_rejects_non_positive_student_sample_steps() -> None:
-    src = inspect.getsource(AnyFlowMethod.__init__)
-    assert 'student_sample_steps must be positive' in src
-
-
-# ---------------------------------------------------------------------------
-# Rollout dynamics — stubbed student.
-# ---------------------------------------------------------------------------
-
-
-class _SpyStudent:
-    """Stand-in student that records every (t, r) pair seen during a
-    rollout and predicts a constant velocity field."""
-
-    def __init__(self, num_train_timesteps: int = 1000) -> None:
-        self.num_train_timesteps = num_train_timesteps
-        self.seen: list[tuple[float, float]] = []
-        # A single trainable parameter so callers can verify gradient flow.
-        self.param = torch.nn.Parameter(torch.zeros(1))
-
-    def predict_velocity_with_r(
-        self,
-        noisy: torch.Tensor,
-        t: torch.Tensor,
-        r: torch.Tensor,
-        batch: Any,
-        *,
-        conditional: bool = True,
-        cfg_uncond: Any = None,
-        attn_kind: str = "vsa",
-    ) -> torch.Tensor:
-        del batch, conditional, cfg_uncond, attn_kind
-        self.seen.append((float(t.flatten()[0].item()),
-                          float(r.flatten()[0].item())))
-        # Constant velocity field of magnitude param so we can backprop.
-        return self.param * torch.ones_like(noisy)
-
-
-def _make_batch(shape: tuple[int, ...]) -> SimpleNamespace:
-    batch = SimpleNamespace()
-    batch.latents = torch.randn(*shape)
-    batch.dmd_latent_vis_dict = {}
-    return batch
-
-
-def test_rollout_uses_mean_velocity_r_equals_t_next() -> None:
-    """With use_mean_velocity=True, r at step i must equal t at step i+1."""
-    method = _naked_method(
-        student_sample_steps=4,
-        use_mean_velocity=True,
-        t_list_override=[999.0, 750.0, 500.0, 250.0, 0.0],
-    )
-    student = _SpyStudent()
-    method.student = student  # type: ignore[assignment]
-
-    batch = _make_batch((1, 2, 4, 4, 4))
-    _ = method._student_rollout(batch, with_grad=False)
-
-    # 4 forwards = 4 (t, r) pairs.
-    assert len(student.seen) == 4
-    for i in range(3):
-        # r at step i must equal t at step i+1.
-        assert student.seen[i][1] == student.seen[i + 1][0]
-
-
-def test_rollout_use_mean_velocity_false_uses_r_equal_t() -> None:
-    method = _naked_method(
-        student_sample_steps=2,
-        use_mean_velocity=False,
-        t_list_override=[999.0, 500.0, 0.0],
-    )
-    student = _SpyStudent()
-    method.student = student  # type: ignore[assignment]
-    batch = _make_batch((1, 2, 4, 4, 4))
-    _ = method._student_rollout(batch, with_grad=False)
-    for t_seen, r_seen in student.seen:
-        assert t_seen == r_seen
-
-
-def test_rollout_with_grad_true_produces_differentiable_output() -> None:
-    method = _naked_method(
-        student_sample_steps=4,
-        use_mean_velocity=True,
-        t_list_override=[999.0, 750.0, 500.0, 250.0, 0.0],
-    )
-    student = _SpyStudent()
-    method.student = student  # type: ignore[assignment]
-    batch = _make_batch((1, 2, 4, 4, 4))
-    out = method._student_rollout(batch, with_grad=True)
-    assert out.requires_grad, (
-        "Rollout output must keep a gradient so the DMD loss can backprop "
-        "through the chosen step.")
-    out.sum().backward()
-    assert student.param.grad is not None
-    assert student.param.grad.abs().sum() > 0
-
-
-def test_rollout_with_grad_false_blocks_gradient_completely() -> None:
-    method = _naked_method(
-        student_sample_steps=4,
-        use_mean_velocity=True,
-        t_list_override=[999.0, 750.0, 500.0, 250.0, 0.0],
-    )
-    student = _SpyStudent()
-    method.student = student  # type: ignore[assignment]
-    batch = _make_batch((1, 2, 4, 4, 4))
-    out = method._student_rollout(batch, with_grad=False)
-    assert not out.requires_grad
-
-
-def test_broadcast_grad_step_index_in_range() -> None:
-    method = _naked_method(student_sample_steps=4)
-    for _ in range(20):
-        idx = method._broadcast_grad_step_index(
-            num_steps=4, device=torch.device("cpu"))
-        assert 0 <= idx < 4
-
-
-def test_broadcast_grad_step_index_rejects_non_positive_num_steps() -> None:
-    method = _naked_method()
-    with pytest.raises(ValueError, match="num_steps must be positive"):
-        method._broadcast_grad_step_index(
-            num_steps=0, device=torch.device("cpu"))

fastvideo/tests/training/distill/test_anyflow_pretrain.py

@@ -1,604 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""AnyFlow pretrain method tests.
-
-CPU-only unit tests covering:
-- Config flag defaults (bit-identity preserved on legacy paths).
-- ``WanTimeTextImageEmbedding`` dual-timestep forward (additive default = bit-identical
-  to legacy; gated mode reproduces AnyFlow's ``(1 - g) * temb + g * delta_emb`` fusion).
-- ``WanTransformer3DModel.forward`` accepts ``r_timestep``.
-- ``FlowMapEulerDiscreteScheduler`` numerics: ``apply_shift``, ``get_train_weight``,
-  ``step``.
-- ``(t, r)`` per-batch sampling distribution.
-- Central-difference target math.
-- AnyFlow HF checkpoint key remap (``remap_anyflow_keys``).
-"""
-
-from __future__ import annotations
-
-import copy
-
-import pytest
-import torch
-
-from fastvideo.configs.models.dits import WanVideoConfig
-
-
-# ---------------------------------------------------------------------------
-# Task 1: r_embedder config flags default to bit-identity preservation.
-# ---------------------------------------------------------------------------
-
-
-def test_wan_arch_defaults_preserve_bit_identity() -> None:
-    cfg = WanVideoConfig()
-    arch = cfg.arch_config
-    assert arch.r_embedder is False
-    assert arch.r_embedder_fusion == "additive"
-    assert arch.r_embedder_gate_value == 0.25
-    assert arch.r_embedder_deltatime_type == "r"
-
-
-# ---------------------------------------------------------------------------
-# Task 2: WanTimeTextImageEmbedding dual-timestep forward.
-# ---------------------------------------------------------------------------
-
-
-def _init_uninitialized_weights(module: torch.nn.Module, seed: int = 0) -> None:
-    """FastVideo's ``ReplicatedLinear`` allocates weights with
-    ``torch.empty`` and relies on a downstream ``load_weights`` pass to
-    populate them. Unit tests bypass that pass, so weights start as
-    uninitialized garbage (typically NaN/Inf). Manually init every
-    Linear / RMSNorm / LayerNorm parameter so the forward produces
-    deterministic finite outputs.
-    """
-    torch.manual_seed(seed)
-    with torch.no_grad():
-        for p in module.parameters():
-            if p.ndim >= 2:
-                # Xavier-uniform scaled by inverse fan-in for stable forwards.
-                torch.nn.init.xavier_uniform_(p)
-            else:
-                p.zero_()
-
-
-def _make_embedder(
-    *,
-    r_embedder: bool,
-    fusion: str = "additive",
-    gate: float = 0.25,
-    deltatime_type: str = "r",
-    init_seed: int = 0,
-):
-    from fastvideo.models.dits.wanvideo import WanTimeTextImageEmbedding
-
-    emb = WanTimeTextImageEmbedding(
-        dim=32,
-        time_freq_dim=64,
-        text_embed_dim=16,
-        image_embed_dim=None,
-        r_embedder=r_embedder,
-        r_embedder_fusion=fusion,
-        r_embedder_gate_value=gate,
-        r_embedder_deltatime_type=deltatime_type,
-    )
-    _init_uninitialized_weights(emb, seed=init_seed)
-    emb.eval()
-    return emb
-
-
-def test_embedder_default_path_no_delta_module() -> None:
-    """When r_embedder=False, delta_embedder must not be allocated."""
-    emb = _make_embedder(r_embedder=False)
-    assert emb.delta_embedder is None
-
-
-def test_embedder_default_path_is_bit_identical_to_legacy() -> None:
-    """With r_embedder=False, forward output must match the legacy single-t path
-    (no r_timestep kwarg, no extra computation)."""
-    torch.manual_seed(0)
-    emb = _make_embedder(r_embedder=False)
-    t = torch.randint(0, 1000, (2,), dtype=torch.long)
-    txt = torch.randn(2, 4, 16)
-    temb_a, proj_a, _, _ = emb(t, txt)
-    # Calling without r_timestep again must be deterministic-equal.
-    temb_b, proj_b, _, _ = emb(t, txt)
-    torch.testing.assert_close(temb_a, temb_b)
-    torch.testing.assert_close(proj_a, proj_b)
-    assert temb_a.shape == (2, 32)
-    assert proj_a.shape == (2, 32 * 6)
-
-
-def test_embedder_enabled_without_r_timestep_is_bit_identical_to_legacy() -> None:
-    """Even with r_embedder=True, if r_timestep is None at call time the
-    forward must skip the delta path entirely so existing call sites that
-    don't pass r_timestep stay byte-equal to the legacy result."""
-    torch.manual_seed(0)
-    emb_legacy = _make_embedder(r_embedder=False)
-    torch.manual_seed(0)
-    emb_dual = _make_embedder(r_embedder=True, fusion="additive")
-    t = torch.randint(0, 1000, (2,), dtype=torch.long)
-    txt = torch.randn(2, 4, 16)
-    temb_legacy, proj_legacy, _, _ = emb_legacy(t, txt)
-    temb_dual, proj_dual, _, _ = emb_dual(t, txt)  # No r_timestep.
-    torch.testing.assert_close(temb_legacy, temb_dual)
-    torch.testing.assert_close(proj_legacy, proj_dual)
-
-
-def test_embedder_gated_fusion_formula() -> None:
-    """Gated mode: rt_emb = (1 - g) * temb_t + g * delta_emb (with delta_input=r)."""
-    torch.manual_seed(0)
-    emb = _make_embedder(r_embedder=True, fusion="gated", gate=0.25)
-    t = torch.tensor([500, 500], dtype=torch.long)
-    r = torch.tensor([100, 100], dtype=torch.long)
-    txt = torch.randn(2, 4, 16)
-
-    temb_t = emb.time_embedder(t)
-    delta_emb = emb.delta_embedder(r)
-    expected = 0.75 * temb_t + 0.25 * delta_emb
-
-    rt_emb, _, _, _ = emb(t, txt, r_timestep=r)
-    torch.testing.assert_close(rt_emb, expected, rtol=1e-5, atol=1e-5)
-
-
-def test_embedder_additive_fusion_formula() -> None:
-    """Additive mode: rt_emb = temb_t + g * delta_emb."""
-    torch.manual_seed(0)
-    emb = _make_embedder(r_embedder=True, fusion="additive", gate=0.3)
-    t = torch.tensor([700, 700], dtype=torch.long)
-    r = torch.tensor([200, 200], dtype=torch.long)
-    txt = torch.randn(2, 4, 16)
-
-    temb_t = emb.time_embedder(t)
-    delta_emb = emb.delta_embedder(r)
-    expected = temb_t + 0.3 * delta_emb
-
-    rt_emb, _, _, _ = emb(t, txt, r_timestep=r)
-    torch.testing.assert_close(rt_emb, expected, rtol=1e-5, atol=1e-5)
-
-
-def test_embedder_deltatime_type_t_minus_r() -> None:
-    """When deltatime_type='t-r', delta_embedder consumes (t - r)."""
-    torch.manual_seed(0)
-    emb = _make_embedder(
-        r_embedder=True, fusion="gated", gate=0.5, deltatime_type="t-r")
-    t = torch.tensor([800, 800], dtype=torch.long)
-    r = torch.tensor([300, 300], dtype=torch.long)
-    txt = torch.randn(2, 4, 16)
-
-    temb_t = emb.time_embedder(t)
-    delta_emb = emb.delta_embedder(t - r)
-    expected = 0.5 * temb_t + 0.5 * delta_emb
-
-    rt_emb, _, _, _ = emb(t, txt, r_timestep=r)
-    torch.testing.assert_close(rt_emb, expected, rtol=1e-5, atol=1e-5)
-
-
-def test_embedder_invalid_fusion_raises() -> None:
-    with pytest.raises(ValueError, match="r_embedder_fusion"):
-        _make_embedder(r_embedder=True, fusion="bogus")
-
-
-def test_embedder_invalid_deltatime_type_raises() -> None:
-    with pytest.raises(ValueError, match="r_embedder_deltatime_type"):
-        _make_embedder(r_embedder=True, fusion="gated", deltatime_type="2t-r")
-
-
-def test_embedder_gate_not_in_state_dict() -> None:
-    """Gate is a non-persistent buffer; it must not appear in state_dict so
-    checkpoints stay portable across different gate hyperparameters."""
-    emb = _make_embedder(r_embedder=True, fusion="gated", gate=0.25)
-    keys = list(emb.state_dict().keys())
-    assert not any("_r_embedder_gate" in k for k in keys)
-
-
-# ---------------------------------------------------------------------------
-# Task 3: WanTransformer3DModel threads r_timestep through.
-# ---------------------------------------------------------------------------
-
-
-def test_wan_transformer_forward_signature_has_r_timestep() -> None:
-    """The forward signature must declare r_timestep explicitly (not
-    swallowed by **kwargs) so callers and type checkers can see it."""
-    import inspect
-
-    from fastvideo.models.dits.wanvideo import WanTransformer3DModel
-
-    sig = inspect.signature(WanTransformer3DModel.forward)
-    assert "r_timestep" in sig.parameters
-    param = sig.parameters["r_timestep"]
-    assert param.default is None
-
-
-def test_wan_transformer_init_propagates_r_embedder_config() -> None:
-    """When the arch config sets r_embedder=True the WanTransformer3DModel
-    constructor must instantiate the embedder with the delta path active.
-
-    We avoid full WanTransformer3DModel instantiation (which requires
-    distributed init) by reading the source's __init__ to confirm it
-    forwards the four arch flags to WanTimeTextImageEmbedding.
-    """
-    import inspect
-
-    from fastvideo.models.dits.wanvideo import WanTransformer3DModel
-
-    src = inspect.getsource(WanTransformer3DModel.__init__)
-    # All four arch config fields must be passed to WanTimeTextImageEmbedding.
-    assert "r_embedder=config.r_embedder" in src
-    assert "r_embedder_fusion=config.r_embedder_fusion" in src
-    assert "r_embedder_gate_value=config.r_embedder_gate_value" in src
-    assert "r_embedder_deltatime_type=config.r_embedder_deltatime_type" in src
-
-
-def test_wan_transformer_forward_threads_r_timestep_to_embedder() -> None:
-    """The forward must pass r_timestep into the embedder call (verified via
-    source inspection to avoid heavyweight distributed bring-up)."""
-    import inspect
-
-    from fastvideo.models.dits.wanvideo import WanTransformer3DModel
-
-    src = inspect.getsource(WanTransformer3DModel.forward)
-    assert "r_timestep=r_timestep" in src, (
-        "WanTransformer3DModel.forward must forward r_timestep into "
-        "self.condition_embedder")
-
-
-# ---------------------------------------------------------------------------
-# Task 4: FlowMapEulerDiscreteScheduler numerics.
-# ---------------------------------------------------------------------------
-
-
-def _scheduler(*, shift: float = 1.0, n_train: int = 1000):
-    from fastvideo.models.schedulers.scheduling_flow_map_euler_discrete import (
-        FlowMapEulerDiscreteScheduler, )
-    return FlowMapEulerDiscreteScheduler(
-        num_train_timesteps=n_train, shift=shift)
-
-
-def test_flow_map_scheduler_set_timesteps_descending() -> None:
-    sched = _scheduler(shift=5.0)
-    sched.set_timesteps(num_inference_steps=4, device=torch.device("cpu"))
-    ts = sched.timesteps
-    # N inference steps → N + 1 boundary entries.
-    assert ts.numel() == 5
-    assert torch.all(ts[:-1] >= ts[1:])  # descending
-    assert ts[-1].item() == 0.0
-    assert ts[0].item() == pytest.approx(1000.0, abs=1e-3)
-
-
-def test_flow_map_scheduler_set_timesteps_custom_overrides_schedule() -> None:
-    sched = _scheduler(shift=5.0)
-    custom = [999.0, 937.0, 833.0, 624.0, 0.0]
-    sched.set_timesteps(
-        num_inference_steps=4,
-        device=torch.device("cpu"),
-        custom_timesteps=custom,
-    )
-    torch.testing.assert_close(
-        sched.timesteps, torch.tensor(custom, dtype=torch.float32))
-
-
-def test_flow_map_scheduler_custom_timesteps_must_be_descending() -> None:
-    sched = _scheduler(shift=5.0)
-    with pytest.raises(ValueError, match="descending"):
-        sched.set_timesteps(
-            num_inference_steps=4,
-            device=torch.device("cpu"),
-            custom_timesteps=[100.0, 500.0, 900.0],
-        )
-
-
-def test_flow_map_scheduler_apply_shift_endpoints_invariant() -> None:
-    """apply_shift fixes the endpoints {0, 1} and produces non-trivial
-    motion in the interior for shift != 1."""
-    sched = _scheduler(shift=5.0)
-    t = torch.tensor([0.0, 0.5, 1.0])
-    shifted = sched.apply_shift(t)
-    torch.testing.assert_close(
-        shifted, torch.tensor([0.0, 5.0 / 6.0, 1.0]), rtol=1e-6, atol=1e-6)
-
-
-def test_flow_map_scheduler_apply_shift_shift_one_is_identity() -> None:
-    sched = _scheduler(shift=1.0)
-    t = torch.linspace(0.0, 1.0, 100)
-    torch.testing.assert_close(sched.apply_shift(t), t)
-
-
-def test_flow_map_scheduler_step_one_euler_iteration_matches_formula() -> None:
-    """One step: x_r = x_t - ((t - r) / N) * model_output."""
-    sched = _scheduler(shift=1.0)
-    sched.set_timesteps(num_inference_steps=4, device=torch.device("cpu"))
-
-    torch.manual_seed(0)
-    x_t = torch.randn(2, 4, 1, 8, 8)
-    v = torch.randn_like(x_t)
-    t = torch.tensor([750.0, 500.0])
-    r = torch.tensor([500.0, 250.0])
-
-    out = sched.step(v, sample=x_t, timestep=t, r_timestep=r)
-    expected = x_t - ((t - r) / 1000.0).view(-1, 1, 1, 1, 1) * v
-    torch.testing.assert_close(out, expected, rtol=1e-6, atol=1e-6)
-
-
-def test_flow_map_scheduler_get_train_weight_beta08_shape_and_renorm() -> None:
-    """beta08: t * sqrt(1-t), renormalized so sum equals num_train_timesteps.
-    The interior of the schedule must dominate the endpoints (monotone up
-    then monotone down)."""
-    sched = _scheduler()
-    t = torch.linspace(0.001, 0.999, 1000)
-    w = sched.get_train_weight(t, weight_type="beta08")
-    assert torch.allclose(w.sum(), torch.tensor(1000.0), rtol=1e-3)
-    assert torch.all(w >= 0.0)
-    # Endpoints smaller than the middle bump.
-    mid = len(w) // 2
-    assert w[0] < w[mid]
-    assert w[-1] < w[mid]
-
-
-def test_flow_map_scheduler_get_train_weight_uniform_is_constant_norm() -> None:
-    sched = _scheduler()
-    t = torch.linspace(0.0, 1.0, 1000)
-    w = sched.get_train_weight(t, weight_type="uniform")
-    assert torch.allclose(w.sum(), torch.tensor(1000.0), rtol=1e-3)
-    # All entries equal to 1.0 after renormalization.
-    torch.testing.assert_close(w, torch.ones_like(w), rtol=1e-6, atol=1e-6)
-
-
-def test_flow_map_scheduler_get_train_weight_accepts_absolute_units() -> None:
-    """When t is provided in [0, num_train_timesteps] the helper auto-
-    normalizes; result must match the [0, 1] call."""
-    sched = _scheduler()
-    t_norm = torch.linspace(0.001, 0.999, 1000)
-    t_abs = t_norm * 1000
-    w_norm = sched.get_train_weight(t_norm, weight_type="beta08")
-    w_abs = sched.get_train_weight(t_abs, weight_type="beta08")
-    torch.testing.assert_close(w_norm, w_abs, rtol=1e-5, atol=1e-5)
-
-
-def test_flow_map_scheduler_add_noise_matches_flow_matching_formula() -> None:
-    """Linear flow-matching: x_t = (1 - sigma) * x_0 + sigma * eps, where
-    sigma = t / num_train_timesteps."""
-    sched = _scheduler()
-    torch.manual_seed(0)
-    x0 = torch.randn(2, 4, 1, 4, 4)
-    eps = torch.randn_like(x0)
-    t = torch.tensor([250.0, 750.0])
-    out = sched.add_noise(x0, eps, t)
-    sigma = (t / 1000.0).view(-1, 1, 1, 1, 1)
-    expected = (1.0 - sigma) * x0 + sigma * eps
-    torch.testing.assert_close(out, expected, rtol=1e-6, atol=1e-6)
-
-
-# ---------------------------------------------------------------------------
-# Task 6: (t, r) per-batch sampling.
-# ---------------------------------------------------------------------------
-
-
-def test_sample_pair_timesteps_partitions_batch_correctly() -> None:
-    """For batch=8 with diffusion=0.5/consistency=0.25: 4 r=t, 2 r=0,
-    2 free entries."""
-    from fastvideo.train.methods.distribution_matching.anyflow_pretrain import (
-        _sample_pair_timesteps, )
-
-    torch.manual_seed(42)
-    t, r, is_diffusion, is_consistency = _sample_pair_timesteps(
-        batch_size=8,
-        diffusion_ratio=0.5,
-        consistency_ratio=0.25,
-        device=torch.device("cpu"),
-        generator=None,
-    )
-    assert t.shape == (8,)
-    assert r.shape == (8,)
-    assert int(is_diffusion.sum()) == 4
-    assert int(is_consistency.sum()) == 2
-    # The masks must be disjoint.
-    assert not torch.any(is_diffusion & is_consistency)
-
-    diff_idx = torch.nonzero(is_diffusion).flatten()
-    torch.testing.assert_close(r[diff_idx], t[diff_idx])
-
-    cons_idx = torch.nonzero(is_consistency).flatten()
-    torch.testing.assert_close(r[cons_idx], torch.zeros(2))
-
-    free_idx = torch.nonzero(~(is_diffusion | is_consistency)).flatten()
-    assert torch.all(r[free_idx] <= t[free_idx])
-    assert torch.all(r[free_idx] >= 0.0)
-    assert torch.all(t[free_idx] <= 1.0)
-
-
-def test_sample_pair_timesteps_t_max_r_min_ordering() -> None:
-    """For the free fraction (ratios=0), t and r come from max/min of two
-    uniform draws so r <= t holds."""
-    torch.manual_seed(0)
-    from fastvideo.train.methods.distribution_matching.anyflow_pretrain import (
-        _sample_pair_timesteps, )
-
-    for _ in range(50):
-        t, r, is_diff, is_cons = _sample_pair_timesteps(
-            batch_size=4,
-            diffusion_ratio=0.0,
-            consistency_ratio=0.0,
-            device=torch.device("cpu"),
-            generator=None,
-        )
-        assert torch.all(t >= r)
-        assert torch.all(r >= 0.0)
-        assert torch.all(t <= 1.0)
-        assert int(is_diff.sum()) == 0
-        assert int(is_cons.sum()) == 0
-
-
-def test_sample_pair_timesteps_rejects_ratios_summing_above_one() -> None:
-    from fastvideo.train.methods.distribution_matching.anyflow_pretrain import (
-        _sample_pair_timesteps, )
-
-    with pytest.raises(ValueError, match="must be <= 1"):
-        _sample_pair_timesteps(
-            batch_size=8,
-            diffusion_ratio=0.7,
-            consistency_ratio=0.4,
-            device=torch.device("cpu"),
-            generator=None,
-        )
-
-
-def test_sample_pair_timesteps_rejects_negative_ratios() -> None:
-    from fastvideo.train.methods.distribution_matching.anyflow_pretrain import (
-        _sample_pair_timesteps, )
-
-    with pytest.raises(ValueError, match="non-negative"):
-        _sample_pair_timesteps(
-            batch_size=8,
-            diffusion_ratio=-0.1,
-            consistency_ratio=0.25,
-            device=torch.device("cpu"),
-            generator=None,
-        )
-
-
-# ---------------------------------------------------------------------------
-# Task 7: central-difference target.
-# ---------------------------------------------------------------------------
-
-
-class _StubStudent:
-    """Stand-in student for unit-testing the central-difference helper.
-
-    The "velocity prediction" is a closed-form function of x and t
-    (no actual neural network) so we can compare against the analytical
-    derivative.
-    """
-
-    def __init__(self, alpha: float = 0.3) -> None:
-        self.alpha = float(alpha)
-
-    def predict_velocity_with_r(
-        self,
-        noisy: torch.Tensor,
-        t: torch.Tensor,
-        r: torch.Tensor,
-        batch,
-        *,
-        conditional: bool = True,
-        attn_kind: str = "dense",
-        cfg_uncond=None,
-    ) -> torch.Tensor:
-        del batch, conditional, attn_kind, cfg_uncond, r
-        # f(x, t) = x + alpha * (t / 1000) → dF/dt = alpha / 1000.
-        view = [-1] + [1] * (noisy.ndim - 1)
-        return noisy + self.alpha * (t.view(*view).float() / 1000.0)
-
-
-def test_central_difference_dF_dt_linear_function() -> None:
-    """For f(x, t) = x + alpha * (t / N), the central-difference estimate
-    of dF/dt must be alpha / N (in absolute t-units that's alpha / N
-    velocity per t-unit, and our helper divides the difference by
-    2*delta -> exactly alpha / N)."""
-    from fastvideo.train.methods.distribution_matching.anyflow_pretrain import (
-        _central_difference_dF_dt, )
-
-    student = _StubStudent(alpha=0.3)
-    x = torch.randn(2, 1, 4, 4, 4)
-    latents = torch.zeros_like(x)
-    noise = torch.zeros_like(x)  # v_pred = noise - latents = 0
-    t = torch.tensor([500.0, 250.0])
-    r = torch.tensor([100.0, 0.0])
-
-    dF = _central_difference_dF_dt(
-        student=student,
-        batch=None,
-        noisy=x,
-        latents=latents,
-        noise=noise,
-        t=t,
-        r=r,
-        delta=5.0,
-        num_train_timesteps=1000.0,
-    )
-
-    expected = torch.full_like(x, 0.3 / 1000.0)
-    torch.testing.assert_close(dF, expected, rtol=1e-5, atol=1e-5)
-
-
-def test_central_difference_dF_dt_guidance_scaling() -> None:
-    """With guidance_scale != 1, the result must be divided by guidance."""
-    from fastvideo.train.methods.distribution_matching.anyflow_pretrain import (
-        _central_difference_dF_dt, )
-
-    student = _StubStudent(alpha=0.6)
-    x = torch.zeros(1, 1, 4, 4, 4)
-    latents = torch.zeros_like(x)
-    noise = torch.zeros_like(x)
-    t = torch.tensor([500.0])
-    r = torch.tensor([100.0])
-
-    dF_g1 = _central_difference_dF_dt(
-        student=student, batch=None, noisy=x, latents=latents, noise=noise,
-        t=t, r=r, delta=5.0, num_train_timesteps=1000.0, guidance_scale=1.0)
-    dF_g3 = _central_difference_dF_dt(
-        student=student, batch=None, noisy=x, latents=latents, noise=noise,
-        t=t, r=r, delta=5.0, num_train_timesteps=1000.0, guidance_scale=3.0)
-
-    torch.testing.assert_close(dF_g1 / 3.0, dF_g3, rtol=1e-5, atol=1e-5)
-
-
-def test_central_difference_dF_dt_rejects_zero_delta() -> None:
-    from fastvideo.train.methods.distribution_matching.anyflow_pretrain import (
-        _central_difference_dF_dt, )
-
-    with pytest.raises(ValueError, match="delta must be positive"):
-        _central_difference_dF_dt(
-            student=_StubStudent(),
-            batch=None,
-            noisy=torch.zeros(1, 1, 4, 4, 4),
-            latents=torch.zeros(1, 1, 4, 4, 4),
-            noise=torch.zeros(1, 1, 4, 4, 4),
-            t=torch.tensor([500.0]),
-            r=torch.tensor([100.0]),
-            delta=0.0,
-            num_train_timesteps=1000.0,
-        )
-
-
-# ---------------------------------------------------------------------------
-# Task 9: param_names_mapping handles AnyFlow checkpoints + is a no-op on plain
-# Wan checkpoints. We don't ship a separate remap_anyflow_keys helper because
-# the existing param_names_mapping regex mechanism does the job (delta_embedder
-# rename is a no-op when the source state dict doesn't contain those keys).
-# ---------------------------------------------------------------------------
-
-
-def test_param_names_mapping_includes_delta_embedder_rename() -> None:
-    """The Wan arch config's param_names_mapping must rename HF AnyFlow
-    delta_embedder weights into FastVideo's internal mlp.fc_in/fc_out layout."""
-    arch = WanVideoConfig().arch_config
-    mapping_keys = list(arch.param_names_mapping.keys())
-    assert any("delta_embedder" in k for k in mapping_keys), (
-        "WanVideoArchConfig.param_names_mapping must include delta_embedder "
-        "rename so HF AnyFlow checkpoints load without a separate adapter")
-
-
-def test_param_names_mapping_default_doesnt_break_plain_wan_keys() -> None:
-    """The new delta_embedder regex must not match any key in a plain
-    pretrained Wan2.1 checkpoint (those don't have delta_embedder)."""
-    import re
-
-    plain_wan_keys = [
-        "patch_embedding.weight",
-        "condition_embedder.time_embedder.linear_1.weight",
-        "condition_embedder.time_embedder.linear_2.weight",
-        "condition_embedder.time_proj.weight",
-        "condition_embedder.text_embedder.linear_1.weight",
-        "blocks.0.attn1.to_q.weight",
-        "blocks.0.ffn.net.0.proj.weight",
-    ]
-    arch = WanVideoConfig().arch_config
-    delta_regexes = [
-        k for k in arch.param_names_mapping if "delta_embedder" in k
-    ]
-    assert delta_regexes, "expected at least one delta_embedder regex"
-
-    for plain in plain_wan_keys:
-        for rx in delta_regexes:
-            assert re.match(rx, plain) is None, (
-                f"delta_embedder regex {rx!r} unexpectedly matched plain "
-                f"Wan key {plain!r}")

fastvideo/tests/training/distill/test_anyflow_smoke.py

@@ -1,130 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""GPU smoke test for AnyFlow pretrain + on-policy.
-
-Mirrors ``test_distill_dmd.py`` — runs the new YAML-driven training
-entrypoint via ``torchrun`` for two iterations to verify end-to-end
-wiring (model load, optimizer build, forward, backward, step, save).
-
-CPU-only environments are skipped; the test is intended to fire on the
-Buildkite ``/test distillation`` lane and on local boxes with at least
-2 H100/H200 GPUs.
-"""
-
-from __future__ import annotations
-
-import os
-import subprocess
-import sys
-from pathlib import Path
-
-import pytest
-
-
-REPO_ROOT = Path(__file__).resolve().parents[4]
-PRETRAIN_YAML = (
-    REPO_ROOT
-    / "examples/train/configs/distribution_matching/wan/anyflow_pretrain_t2v.yaml")
-ONPOLICY_YAML = (
-    REPO_ROOT
-    / "examples/train/configs/distribution_matching/wan/anyflow_onpolicy_t2v.yaml")
-
-NUM_NODES = "1"
-NUM_GPUS_PER_NODE = "2"
-
-
-def _have_enough_gpus() -> bool:
-    """Return True iff at least 2 CUDA devices are visible. The smoke test
-    needs HSDP/FSDP with a non-trivial world size; single-GPU bring-up
-    races against the new framework's distributed barriers."""
-    try:
-        import torch
-    except Exception:
-        return False
-    if not torch.cuda.is_available():
-        return False
-    return torch.cuda.device_count() >= 2
-
-
-pytestmark = pytest.mark.skipif(
-    not _have_enough_gpus(),
-    reason="AnyFlow smoke test requires >= 2 CUDA devices")
-
-
-def _run_torchrun(config_path: Path, *, output_dir: Path) -> None:
-    if not config_path.exists():
-        pytest.fail(f"YAML config missing: {config_path}")
-
-    env = os.environ.copy()
-    env.setdefault("MASTER_ADDR", "127.0.0.1")
-    env.setdefault("MASTER_PORT", "29551")
-    env.setdefault("WANDB_MODE", "offline")
-    env.setdefault("TOKENIZERS_PARALLELISM", "false")
-
-    cmd = [
-        sys.executable,
-        "-m",
-        "torch.distributed.run",
-        "--nnodes", NUM_NODES,
-        "--nproc_per_node", NUM_GPUS_PER_NODE,
-        "--master_port", env["MASTER_PORT"],
-        "-m", "fastvideo.train.entrypoint.train",
-        "--config", str(config_path),
-        "--training.loop.max_train_steps", "2",
-        "--training.checkpoint.output_dir", str(output_dir),
-        "--training.distributed.num_gpus", NUM_GPUS_PER_NODE,
-        "--training.distributed.hsdp_shard_dim", NUM_GPUS_PER_NODE,
-        "--training.data.train_batch_size", "1",
-    ]
-    process = subprocess.run(cmd, capture_output=True, text=True, env=env)
-    if process.stdout:
-        print("STDOUT:", process.stdout)
-    if process.stderr:
-        print("STDERR:", process.stderr)
-    if process.returncode != 0:
-        raise subprocess.CalledProcessError(
-            process.returncode, cmd, process.stdout, process.stderr)
-
-
-def test_anyflow_pretrain_smoke(tmp_path: Path) -> None:
-    """Two-iteration pretrain — exercises (t, r) sampling, central-difference
-    target, scale balance, optimizer step, and checkpoint save path."""
-    _run_torchrun(PRETRAIN_YAML, output_dir=tmp_path / "pretrain")
-
-
-def test_anyflow_onpolicy_smoke(tmp_path: Path) -> None:
-    """Two-iteration on-policy DMD — exercises the multi-step Euler-flow
-    rollout, grad-step broadcast, DMD2 alternating updates."""
-    # Override init_from to the public Wan2.1-T2V-1.3B-Diffusers checkpoint
-    # for the smoke run; param_names_mapping handles the (non-existent)
-    # delta_embedder rename as a no-op.
-    env = os.environ.copy()
-    env.setdefault("MASTER_ADDR", "127.0.0.1")
-    env.setdefault("MASTER_PORT", "29552")
-    env.setdefault("WANDB_MODE", "offline")
-    env.setdefault("TOKENIZERS_PARALLELISM", "false")
-
-    cmd = [
-        sys.executable,
-        "-m",
-        "torch.distributed.run",
-        "--nnodes", NUM_NODES,
-        "--nproc_per_node", NUM_GPUS_PER_NODE,
-        "--master_port", env["MASTER_PORT"],
-        "-m", "fastvideo.train.entrypoint.train",
-        "--config", str(ONPOLICY_YAML),
-        "--training.loop.max_train_steps", "2",
-        "--training.checkpoint.output_dir", str(tmp_path / "onpolicy"),
-        "--training.distributed.num_gpus", NUM_GPUS_PER_NODE,
-        "--training.distributed.hsdp_shard_dim", NUM_GPUS_PER_NODE,
-        "--training.data.train_batch_size", "1",
-        "--models.student.init_from", "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
-        "--method.student_sample_steps", "2",
-    ]
-    process = subprocess.run(cmd, capture_output=True, text=True, env=env)
-    if process.stdout:
-        print("STDOUT:", process.stdout)
-    if process.stderr:
-        print("STDERR:", process.stderr)
-    if process.returncode != 0:
-        raise subprocess.CalledProcessError(
-            process.returncode, cmd, process.stdout, process.stderr)

fastvideo/tests/transformers/test_flux.py

@@ -1,185 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-
-from __future__ import annotations
-
-import glob
-import os
-
-import pytest
-import torch
-from diffusers import FluxTransformer2DModel as HFFluxTransformer2DModel
-from torch.testing import assert_close
-
-from fastvideo.configs.models.dits.flux import FluxDiTConfig
-from fastvideo.configs.pipelines.base import PipelineConfig
-from fastvideo.fastvideo_args import FastVideoArgs
-from fastvideo.forward_context import set_forward_context
-from fastvideo.models.loader.component_loader import TransformerLoader
-from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
-
-os.environ.setdefault("MASTER_ADDR", "localhost")
-os.environ.setdefault("MASTER_PORT", "29517")
-
-_REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", ".."))
-_DEFAULT_FLUX_TRANSFORMER = os.path.join(
-    _REPO_ROOT,
-    "official_weights",
-    "FLUX.1-dev",
-    "transformer",
-)
-
-
-def _flux_transformer_path() -> str:
-    return os.environ.get("FLUX_TRANSFORMER_PATH", _DEFAULT_FLUX_TRANSFORMER)
-
-
-def _prepare_latent_image_ids(
-    height: int,
-    width: int,
-    device: torch.device,
-    dtype: torch.dtype = torch.long,
-) -> torch.Tensor:
-    """Match Diffusers ``FluxPipeline._prepare_latent_image_ids`` (batch omitted)."""
-    latent_image_ids = torch.zeros(height, width, 3)
-    latent_image_ids[..., 1] = latent_image_ids[..., 1] + torch.arange(height)[:, None]
-    latent_image_ids[..., 2] = latent_image_ids[..., 2] + torch.arange(width)[None, :]
-    h, w, c = latent_image_ids.shape
-    latent_image_ids = latent_image_ids.reshape(h * w, c)
-    return latent_image_ids.to(device=device, dtype=dtype)
-
-
-@pytest.fixture
-def torch_sdpa_attention_backend(monkeypatch: pytest.MonkeyPatch) -> None:
-    monkeypatch.setenv("FASTVIDEO_ATTENTION_BACKEND", "TORCH_SDPA")
-
-
-requires_cuda = pytest.mark.skipif(
-    not torch.cuda.is_available(),
-    reason="FLUX DiT parity test requires CUDA",
-)
-
-requires_weights = pytest.mark.skipif(
-    not glob.glob(os.path.join(_flux_transformer_path(), "*.safetensors")),
-    reason=(
-        f"No safetensors under {_flux_transformer_path()} — download FLUX.1-dev "
-        "transformer or set FLUX_TRANSFORMER_PATH"
-    ),
-)
-
-
-@requires_cuda
-@requires_weights
-@pytest.mark.usefixtures("distributed_setup", "torch_sdpa_attention_backend")
-def test_flux_transformer_parity_vs_diffusers() -> None:
-    """Single forward: FastVideo DiT vs Diffusers ``FluxTransformer2DModel``."""
-    device = torch.device("cuda:0")
-    precision = torch.bfloat16
-    transformer_path = _flux_transformer_path()
-
-    args = FastVideoArgs(
-        model_path=transformer_path,
-        dit_cpu_offload=False,
-        dit_layerwise_offload=False,
-        pipeline_config=PipelineConfig(dit_config=FluxDiTConfig(), dit_precision="bf16"),
-    )
-    args.device = device
-
-    generator = torch.Generator(device=device).manual_seed(0)
-    torch.manual_seed(0)
-
-    batch_size = 1
-    latent_h, latent_w = 4, 4
-    img_seq = latent_h * latent_w
-    text_len = 32
-
-    hidden_states = torch.randn(
-        batch_size,
-        img_seq,
-        64,
-        device=device,
-        dtype=precision,
-        generator=generator,
-    )
-    encoder_hidden_states = torch.randn(
-        batch_size,
-        text_len,
-        4096,
-        device=device,
-        dtype=precision,
-        generator=generator,
-    )
-    pooled_projections = torch.randn(
-        batch_size,
-        768,
-        device=device,
-        dtype=precision,
-        generator=generator,
-    )
-
-    # Diffusers pipeline passes scheduler timesteps / 1000 (float, same dtype as latents).
-    timestep = torch.tensor([512.0], device=device, dtype=precision) / 1000.0
-    guidance = torch.full((batch_size,), 3.5, device=device, dtype=torch.float32)
-
-    txt_ids = torch.zeros(text_len, 3, device=device, dtype=torch.long)
-    img_ids = _prepare_latent_image_ids(latent_h, latent_w, device, dtype=torch.long)
-
-    forward_batch = ForwardBatch(data_type="dummy")
-
-    # One ~12B model at a time avoids peak VRAM from holding both checkpoints.
-    loader = TransformerLoader()
-    fv_model = loader.load(transformer_path, args).to(device=device, dtype=precision)
-    fv_model.eval()
-    with (
-        torch.no_grad(),
-        torch.amp.autocast("cuda", dtype=precision),
-        set_forward_context(
-            current_timestep=512,
-            attn_metadata=None,
-            forward_batch=forward_batch,
-        ),
-    ):
-        fv_out = fv_model(
-            hidden_states=hidden_states.clone(),
-            encoder_hidden_states=encoder_hidden_states.clone(),
-            pooled_projections=pooled_projections.clone(),
-            timestep=timestep.clone(),
-            guidance=guidance.clone(),
-            txt_ids=txt_ids,
-            img_ids=img_ids,
-            return_dict=False,
-        )[0]
-    fv_out_cpu = fv_out.detach().float().cpu()
-    del fv_model
-    if torch.cuda.is_available():
-        torch.cuda.empty_cache()
-
-    hf_model = (
-        HFFluxTransformer2DModel.from_pretrained(
-            transformer_path,
-            torch_dtype=precision,
-        )
-        .to(device)
-        .eval()
-    )
-    with torch.no_grad(), torch.amp.autocast("cuda", dtype=precision):
-        hf_out = hf_model(
-            hidden_states=hidden_states.clone(),
-            encoder_hidden_states=encoder_hidden_states.clone(),
-            pooled_projections=pooled_projections.clone(),
-            timestep=timestep.clone(),
-            guidance=guidance.clone(),
-            txt_ids=txt_ids,
-            img_ids=img_ids,
-            return_dict=False,
-        )[0]
-
-    assert hf_out.shape == fv_out_cpu.shape
-    hf_cpu = hf_out.float().cpu()
-    abs_diff = (hf_cpu - fv_out_cpu).abs()
-    print(f"[FLUX DiT parity] max_diff={abs_diff.max():.4f}  mean_diff={abs_diff.mean():.4f}  "
-          f"median_diff={abs_diff.median():.4f}  p99_diff="
-          f"{abs_diff.flatten().kthvalue(int(0.99 * abs_diff.numel())).values:.4f}")
-    # bfloat16 accumulation over 57 transformer layers produces tail errors up to ~0.5
-    # on isolated elements (median=0, mean~0.04 on L40S).  atol=0.5 catches real bugs
-    # (wrong weights / missing layers) which produce mean_diff >> 0.1.
-    assert_close(hf_cpu, fv_out_cpu, atol=0.5, rtol=0.0)

fastvideo/tests/utils.py

@@ -77,32 +77,13 @@ def _read_video_frames(path: str) -> torch.Tensor:
     return torch.stack(frames)
 
 
-def _read_image_as_single_frame_video(path: str) -> torch.Tensor:
-    """Read one image as a single-frame ``(1, C, H, W)`` uint8 tensor."""
-    from torchvision.io import read_image
-
-    img = read_image(path)
-    return img.unsqueeze(0)
-
-
-def _read_visual_frames(path: str) -> torch.Tensor:
-    """Read a video or a single image as ``(T, C, H, W)`` uint8."""
-    ext = os.path.splitext(path)[1].lower()
-    if ext in {".png", ".jpg", ".jpeg", ".webp"}:
-        return _read_image_as_single_frame_video(path)
-    return _read_video_frames(path)
-
-
 def compute_video_ssim_torchvision(video1_path, video2_path, use_ms_ssim=True):
     """
-    Compute SSIM between two videos or single-frame image files.
-
-    Image paths (``.png``, ``.jpg``, ``.jpeg``, ``.webp``) are treated as
-    one-frame clips so T2I SSIM can share the same MS-SSIM path as video.
+    Compute SSIM between two videos.
 
     Args:
-        video1_path: Path to the first video or image.
-        video2_path: Path to the second video or image.
+        video1_path: Path to the first video.
+        video2_path: Path to the second video.
         use_ms_ssim: Whether to use Multi-Scale Structural Similarity(MS-SSIM) instead of SSIM.
     """
     from pytorch_msssim import ms_ssim, ssim
@@ -113,8 +94,8 @@ def compute_video_ssim_torchvision(video1_path, video2_path, use_ms_ssim=True):
     if not os.path.exists(video2_path):
         raise FileNotFoundError(f"Video2 not found: {video2_path}")
 
-    frames1 = _read_visual_frames(video1_path)
-    frames2 = _read_visual_frames(video2_path)
+    frames1 = _read_video_frames(video1_path)
+    frames2 = _read_video_frames(video2_path)
 
     # Ensure same number of frames
     min_frames = min(frames1.shape[0], frames2.shape[0])

fastvideo/train/methods/distribution_matching/__init__.py

@@ -1,8 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
 
-from fastvideo.train.methods.distribution_matching.anyflow import AnyFlowMethod
-from fastvideo.train.methods.distribution_matching.anyflow_pretrain import (
-    AnyFlowPretrainMethod, )
 from fastvideo.train.methods.distribution_matching.dmd2 import DMD2Method
 from fastvideo.train.methods.distribution_matching.self_forcing import (
     SelfForcingMethod, )
@@ -10,8 +7,6 @@ from fastvideo.train.methods.distribution_matching.streaming_long_tuning import
     StreamingLongTuningMethod, )
 
 __all__ = [
-    "AnyFlowMethod",
-    "AnyFlowPretrainMethod",
     "DMD2Method",
     "SelfForcingMethod",
     "StreamingLongTuningMethod",

fastvideo/train/methods/distribution_matching/anyflow.py

@@ -1,209 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""AnyFlow on-policy distillation method.
-
-Stage 2 of the AnyFlow two-stage recipe. Continues from a pretrained
-flow-map student and refines it via distribution-matching distillation
-(DMD2) where the student is rolled out for ``student_sample_steps``
-Euler-flow steps from pure noise. One randomly-chosen step in the
-rollout is gradient-enabled (and broadcast across ranks so every worker
-agrees on which step to gradient-enable); the rest run under
-``torch.no_grad``.
-
-Inherits ``DMD2Method`` for the alternating student / critic update
-machinery and the existing DMD VSD-with-fake-score loss. Overrides
-``_student_rollout`` to drive the multi-step Euler-flow rollout with
-``r = t_next`` (mean-velocity sampling — matches the AnyFlow paper's
-``WanAnyFlowPipeline.training_rollout`` with ``use_mean_velocity=True``).
-
-Reference: ``pipeline_wan_anyflow.py::training_rollout`` in
-NVlabs/AnyFlow at commit ``549236a``.
-"""
-
-from __future__ import annotations
-
-from typing import Any, Literal
-
-import torch
-import torch.distributed as dist
-
-from fastvideo.train.methods.distribution_matching.dmd2 import DMD2Method
-from fastvideo.train.utils.config import (
-    get_optional_float,
-    get_optional_int,
-)
-
-
-class AnyFlowMethod(DMD2Method):
-    """AnyFlow on-policy distillation (multi-step rollout)."""
-
-    def __init__(
-        self,
-        *,
-        cfg: Any,
-        role_models: dict[str, Any],
-    ) -> None:
-        super().__init__(cfg=cfg, role_models=role_models)
-        mcfg = self.method_config
-
-        student_sample_steps = get_optional_int(mcfg, "student_sample_steps", where="method.student_sample_steps")
-        if student_sample_steps is None:
-            student_sample_steps = 4
-        if int(student_sample_steps) <= 0:
-            raise ValueError("method.student_sample_steps must be positive, "
-                             f"got {student_sample_steps}")
-        self._student_sample_steps = int(student_sample_steps)
-
-        use_mean_velocity_raw = mcfg.get("use_mean_velocity", True)
-        if not isinstance(use_mean_velocity_raw, bool):
-            raise ValueError("method.use_mean_velocity must be a bool, "
-                             f"got {type(use_mean_velocity_raw).__name__}")
-        self._use_mean_velocity = bool(use_mean_velocity_raw)
-
-        # Optional pinned rollout schedule (descending, absolute t-units).
-        # Falls back to dmd_denoising_steps when absent.
-        raw_t_list = mcfg.get("t_list_override", None)
-        if raw_t_list is None:
-            self._t_list_override: list[float] | None = None
-        else:
-            if not isinstance(raw_t_list, list) or not raw_t_list:
-                raise ValueError("method.t_list_override must be a non-empty list of "
-                                 f"floats when set, got {raw_t_list!r}")
-            t_list = [float(x) for x in raw_t_list]
-            for i in range(len(t_list) - 1):
-                if t_list[i] < t_list[i + 1]:
-                    raise ValueError("method.t_list_override must be descending, "
-                                     f"got {t_list!r}")
-            self._t_list_override = t_list
-
-        # Scoring conditioning: AnyFlow scores against r=0 for the DMD branch.
-        score_r_raw = mcfg.get("dmd_score_r_value", 0.0)
-        try:
-            self._dmd_score_r = float(score_r_raw)
-        except (TypeError, ValueError) as exc:
-            raise ValueError("method.dmd_score_r_value must be numeric, "
-                             f"got {score_r_raw!r}") from exc
-
-        # Optional teacher guidance scale for the DMD loss (carry over from
-        # DMD2Method's behavior; default 1.0).
-        guidance = get_optional_float(mcfg, "real_score_guidance_scale", where="method.real_score_guidance_scale")
-        self._real_score_guidance = float(guidance) if guidance is not None else 1.0
-
-    # ------------------------------------------------------------------
-    # Rollout schedule
-
-    def _get_rollout_schedule(self, *, device: torch.device) -> torch.Tensor:
-        """Build the descending timestep schedule used by the on-policy
-        rollout. Length is ``num_steps + 1`` so ``num_steps`` Euler steps
-        consume the full range.
-
-        Order of precedence:
-          1. ``method.t_list_override`` — used verbatim (absolute units).
-          2. ``method.dmd_denoising_steps`` (inherited from DMD2) appended
-             with a final 0 boundary if the last entry isn't already 0.
-        """
-        if self._t_list_override is not None:
-            return torch.tensor(self._t_list_override, device=device, dtype=torch.float32)
-
-        steps = self._get_denoising_step_list(device).to(dtype=torch.float32)
-        if float(steps[-1].item()) != 0.0:
-            zero = torch.zeros(1, device=device, dtype=torch.float32)
-            steps = torch.cat([steps, zero], dim=0)
-        return steps
-
-    def _broadcast_grad_step_index(
-        self,
-        num_steps: int,
-        *,
-        device: torch.device,
-    ) -> int:
-        """Pick the rollout step that gets gradient enabled. In distributed
-        runs the choice is broadcast from rank 0 so every worker agrees."""
-        if num_steps <= 0:
-            raise ValueError("num_steps must be positive")
-        if dist.is_initialized() and dist.get_rank() != 0:
-            idx_tensor = torch.empty((1, ), dtype=torch.long, device=device)
-        else:
-            idx_tensor = torch.randint(0,
-                                       num_steps, (1, ),
-                                       device=device,
-                                       dtype=torch.long,
-                                       generator=self.cuda_generator)
-        if dist.is_initialized():
-            dist.broadcast(idx_tensor, src=0)
-        return int(idx_tensor.item())
-
-    # ------------------------------------------------------------------
-    # Rollout
-
-    def _student_rollout(
-        self,
-        batch: Any,
-        *,
-        with_grad: bool,
-    ) -> torch.Tensor:
-        """Multi-step Euler-flow rollout from pure noise.
-
-        Returns the predicted clean latent ``x_0`` after the chosen
-        gradient step (or the final ``x`` after the last step if
-        ``with_grad`` is False — used by the critic path).
-        """
-        latents = batch.latents
-        if latents is None or latents.ndim != 5:
-            raise RuntimeError("AnyFlow on-policy rollout requires TrainingBatch.latents "
-                               "of shape [B, T, C, H, W] for shape templating")
-        device = latents.device
-        dtype = latents.dtype
-
-        schedule = self._get_rollout_schedule(device=device)
-        num_entries = int(schedule.numel())
-        num_steps = num_entries - 1
-        if num_steps <= 0:
-            raise RuntimeError("rollout schedule must have at least two entries "
-                               f"(got {num_entries})")
-        if num_steps > self._student_sample_steps:
-            # Trim to the configured cap, keeping the last (=0) boundary.
-            schedule = torch.cat([schedule[:self._student_sample_steps], schedule[-1:]], dim=0)
-            num_steps = self._student_sample_steps
-
-        grad_step = self._broadcast_grad_step_index(num_steps, device=device) if with_grad else -1
-
-        attn_kind: Literal["dense", "vsa"] = "vsa"
-        n_train = float(self.student.num_train_timesteps)
-
-        x = torch.randn(latents.shape, device=device, dtype=dtype, generator=self.cuda_generator)
-        last_pred_x0: torch.Tensor | None = None
-        batch_size = int(latents.shape[0])
-
-        for i in range(num_steps):
-            t_cur = schedule[i].expand(batch_size)
-            t_next = schedule[i + 1].expand(batch_size)
-            r = t_next if self._use_mean_velocity else t_cur
-
-            enable_grad = bool(with_grad) and (i == grad_step)
-            with torch.set_grad_enabled(enable_grad):
-                v = self.student.predict_velocity_with_r(
-                    x,
-                    t_cur,
-                    r,
-                    batch,
-                    conditional=True,
-                    cfg_uncond=self._cfg_uncond,
-                    attn_kind=attn_kind,
-                )
-            # Euler step in absolute units: x ← x - ((t_cur - t_next) / N) * v.
-            view = [-1] + [1] * (x.ndim - 1)
-            dt = ((t_cur - t_next) / n_train).view(*view)
-            x = x - dt * v
-
-            if enable_grad:
-                # We treat the rollout output (post-step) as a predicted
-                # clean latent — AnyFlow's last Euler step lands at t=0.
-                last_pred_x0 = x
-
-        if last_pred_x0 is None:
-            # No gradient step taken (with_grad=False path).
-            last_pred_x0 = x
-
-        if hasattr(batch, "dmd_latent_vis_dict"):
-            batch.dmd_latent_vis_dict["generator_timestep"] = (schedule[-1].detach().clone())
-        return last_pred_x0

fastvideo/train/methods/distribution_matching/anyflow_pretrain.py

@@ -1,404 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""AnyFlow pretrain (flow-map central-difference) training method.
-
-Stage 1 of the AnyFlow two-stage recipe. Trains a single student network
-``u_θ(x_t, t, r)`` to predict the average velocity from time ``t`` back
-to time ``r`` via the central-difference target
-
-    target = (eps - x_0) - ((t - r) / N) * dF/dt
-
-where ``N = num_train_timesteps`` and ``dF/dt`` is estimated from the
-student's own forward at ``(t ± δ, r)`` (with one-sided fallback near the
-schedule endpoints).
-
-Per-batch ``(t, r)`` sampling follows the AnyFlow paper:
-
-- ``diffusion_ratio`` fraction: ``r = t`` (recovers plain flow matching).
-- ``consistency_ratio`` fraction: ``r = 0`` (consistency to clean data).
-- Remaining fraction: ``(t, r) = (max, min)`` of two independent uniform
-  draws (full reconstruction range).
-
-Reference: ``trainer_wan_anyflow_pretrain.py`` in NVlabs/AnyFlow at
-commit ``549236a``.
-"""
-
-from __future__ import annotations
-
-from typing import Any
-from collections.abc import Sequence
-
-import torch
-
-from fastvideo.train.methods.base import LogScalar, TrainingMethod
-from fastvideo.train.models.base import ModelBase
-from fastvideo.train.utils.config import (
-    get_optional_float,
-    get_optional_int,
-)
-from fastvideo.train.utils.optimizer import build_optimizer_and_scheduler
-
-
-def _sample_pair_timesteps(
-    *,
-    batch_size: int,
-    diffusion_ratio: float,
-    consistency_ratio: float,
-    device: torch.device,
-    generator: torch.Generator | None,
-) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
-    """Sample ``(t, r)`` per the AnyFlow paper.
-
-    Two uniform draws ``u1, u2 ∈ [0, 1]`` per sample, then
-    ``t = max(u1, u2)`` and ``r = min(u1, u2)``. After the base sample,
-    the first ``diffusion_ratio * B`` entries get ``r = t`` (diffusion
-    branch, plain flow matching), and the next ``consistency_ratio * B``
-    entries get ``r = 0`` (consistency branch).
-
-    Returns
-    -------
-    t, r, is_diffusion, is_consistency
-        Each tensor has shape ``(batch_size,)``. ``t`` and ``r`` are in
-        ``[0, 1]`` (i.e. *not* yet shifted and *not* yet in absolute
-        train-timestep units). ``is_diffusion`` and ``is_consistency``
-        are bool masks that partition a subset of the batch — entries
-        outside both masks are the "free" reconstruction fraction.
-    """
-    if batch_size <= 0:
-        raise ValueError(f"batch_size must be positive, got {batch_size}")
-    if diffusion_ratio < 0.0 or consistency_ratio < 0.0:
-        raise ValueError("diffusion_ratio and consistency_ratio must be non-negative")
-    if diffusion_ratio + consistency_ratio > 1.0:
-        raise ValueError("diffusion_ratio + consistency_ratio must be <= 1, "
-                         f"got {diffusion_ratio} + {consistency_ratio}")
-
-    u1 = torch.rand(batch_size, device=device, generator=generator)
-    u2 = torch.rand(batch_size, device=device, generator=generator)
-    t = torch.maximum(u1, u2)
-    r = torch.minimum(u1, u2)
-
-    n_diff = int(diffusion_ratio * batch_size)
-    n_cons = int(consistency_ratio * batch_size)
-    is_diffusion = torch.zeros(batch_size, dtype=torch.bool, device=device)
-    is_consistency = torch.zeros(batch_size, dtype=torch.bool, device=device)
-    is_diffusion[:n_diff] = True
-    is_consistency[n_diff:n_diff + n_cons] = True
-
-    # Override per the AnyFlow paper:
-    # - diffusion entries: r = t (plain flow matching)
-    # - consistency entries: r = 0 (consistency to clean data)
-    r = torch.where(is_diffusion, t, r)
-    r = torch.where(is_consistency, torch.zeros_like(r), r)
-    return t, r, is_diffusion, is_consistency
-
-
-@torch.no_grad()
-def _central_difference_dF_dt(
-    *,
-    student: Any,
-    batch: Any,
-    noisy: torch.Tensor,
-    latents: torch.Tensor,
-    noise: torch.Tensor,
-    t: torch.Tensor,
-    r: torch.Tensor,
-    delta: float,
-    num_train_timesteps: float,
-    attn_kind: str = "dense",
-    guidance_scale: float = 1.0,
-) -> torch.Tensor:
-    """Estimate ``dF/dt`` for the AnyFlow central-difference target.
-
-    Computes a symmetric finite difference of the velocity prediction in
-    *absolute train-timestep units*:
-
-        dF/dt ≈ [u_θ(x_{t+δ}, t+δ, r) - u_θ(x_{t-δ}, t-δ, r)] / (2 * δ * guidance)
-
-    The sample is also moved along the flow trajectory by the same
-    finite step (``v_pred * (δ / N)``) to match AnyFlow's reference
-    formulation in ``trainer_wan_anyflow_pretrain.py::compute_central_difference``.
-    Wrapped in ``torch.no_grad`` so the two extra forwards never enter
-    the backward graph.
-    """
-    if delta <= 0.0:
-        raise ValueError(f"delta must be positive, got {delta}")
-    if guidance_scale <= 0.0:
-        raise ValueError(f"guidance_scale must be positive, got {guidance_scale}")
-
-    v_pred = noise - latents  # ground-truth flow velocity
-    delta_x = delta / float(num_train_timesteps)
-
-    t_plus = t + delta
-    noisy_plus = noisy + v_pred * delta_x
-    f_plus = student.predict_velocity_with_r(noisy_plus, t_plus, r, batch, conditional=True, attn_kind=attn_kind)
-
-    t_minus = t - delta
-    noisy_minus = noisy - v_pred * delta_x
-    f_minus = student.predict_velocity_with_r(noisy_minus, t_minus, r, batch, conditional=True, attn_kind=attn_kind)
-
-    return (f_plus - f_minus) / (2.0 * delta * guidance_scale)
-
-
-class AnyFlowPretrainMethod(TrainingMethod):
-    """AnyFlow flow-map pretrain method.
-
-    Single-student training; no teacher or critic. The student must
-    implement ``predict_velocity_with_r(noisy, t, r, batch, ...)`` —
-    typically a ``WanModel`` with ``r_embedder=True`` in its arch config.
-    """
-
-    def __init__(
-        self,
-        *,
-        cfg: Any,
-        role_models: dict[str, ModelBase],
-    ) -> None:
-        super().__init__(cfg=cfg, role_models=role_models)
-        if "student" not in role_models:
-            raise ValueError("AnyFlowPretrainMethod requires role 'student'")
-        if not self.student._trainable:
-            raise ValueError("AnyFlowPretrainMethod requires student to be trainable")
-
-        mcfg = self.method_config
-        self._diffusion_ratio = float(
-            get_optional_float(mcfg, "diffusion_ratio", where="method.diffusion_ratio") or 0.5)
-        self._consistency_ratio = float(
-            get_optional_float(mcfg, "consistency_ratio", where="method.consistency_ratio") or 0.25)
-        if self._diffusion_ratio + self._consistency_ratio > 1.0:
-            raise ValueError("method.diffusion_ratio + method.consistency_ratio must "
-                             f"be <= 1, got {self._diffusion_ratio} + "
-                             f"{self._consistency_ratio}")
-
-        # δ: finite-difference step in absolute train-timestep units.
-        epsilon = get_optional_int(mcfg, "epsilon", where="method.epsilon")
-        self._fd_epsilon = float(epsilon) if epsilon is not None else 5.0
-
-        # Loss weighting scheme (uniform / gaussian / beta08).
-        raw_weight_type = mcfg.get("weight_type", "beta08")
-        if not isinstance(raw_weight_type, str):
-            raise ValueError("method.weight_type must be a string, got "
-                             f"{type(raw_weight_type).__name__}")
-        weight_type = raw_weight_type.strip().lower()
-        if weight_type not in {"uniform", "gaussian", "beta08"}:
-            raise ValueError("method.weight_type must be one of "
-                             "{uniform, gaussian, beta08}, "
-                             f"got {raw_weight_type!r}")
-        self._weight_type = weight_type
-
-        # Guidance fused into the training target (default 1.0 = unused).
-        fg = get_optional_float(mcfg, "fuse_guidance_scale", where="method.fuse_guidance_scale")
-        self._fuse_guidance_scale = float(fg) if fg is not None else 1.0
-        if self._fuse_guidance_scale <= 0.0:
-            raise ValueError("method.fuse_guidance_scale must be positive, "
-                             f"got {self._fuse_guidance_scale}")
-
-        # Flow-map scheduler — uses pipeline_config.flow_shift if present
-        # and falls back to method.shift (and finally 1.0).
-        shift = float(getattr(self.training_config.pipeline_config, "flow_shift", 0.0) or 0.0)
-        if shift <= 0.0:
-            shift_override = get_optional_float(mcfg, "shift", where="method.shift")
-            shift = float(shift_override) if shift_override is not None else 1.0
-        self._shift = shift
-
-        # Lazy-imported to avoid circular imports on package load.
-        from fastvideo.models.schedulers.scheduling_flow_map_euler_discrete import (
-            FlowMapEulerDiscreteScheduler, )
-        self._flow_map_scheduler = FlowMapEulerDiscreteScheduler(
-            num_train_timesteps=int(self.student.num_train_timesteps),
-            shift=self._shift,
-        )
-
-        self.student.init_preprocessors(self.training_config)
-        self._init_optimizer_and_scheduler()
-
-    @property
-    def _optimizer_dict(self) -> dict[str, torch.optim.Optimizer]:
-        return {"student": self._student_optimizer}
-
-    @property
-    def _lr_scheduler_dict(self) -> dict[str, Any]:
-        return {"student": self._student_lr_scheduler}
-
-    def get_optimizers(
-        self,
-        iteration: int,
-    ) -> Sequence[torch.optim.Optimizer]:
-        del iteration
-        return [self._student_optimizer]
-
-    def get_lr_schedulers(self, iteration: int) -> Sequence[Any]:
-        del iteration
-        return [self._student_lr_scheduler]
-
-    def single_train_step(
-        self,
-        batch: dict[str, Any],
-        iteration: int,
-    ) -> tuple[
-            dict[str, torch.Tensor],
-            dict[str, Any],
-            dict[str, LogScalar],
-    ]:
-        del iteration  # AnyFlow pretrain has no iteration-dependent dispatch.
-
-        training_batch = self.student.prepare_batch(
-            batch,
-            generator=self.cuda_generator,
-            latents_source="data",
-        )
-        latents = training_batch.latents  # [B, T, C, H, W] (post-permute in prepare_batch).
-        if latents is None or latents.ndim != 5:
-            raise RuntimeError("AnyFlow pretrain expects TrainingBatch.latents of shape "
-                               "[B, T, C, H, W] after prepare_batch; got "
-                               f"{None if latents is None else tuple(latents.shape)}")
-        device = latents.device
-        dtype = latents.dtype
-        batch_size = int(latents.shape[0])
-
-        # AnyFlow (t, r) sampling — overrides the timestep drawn by
-        # WanModel._sample_timesteps inside prepare_batch.
-        t_norm, r_norm, is_diffusion, is_consistency = _sample_pair_timesteps(
-            batch_size=batch_size,
-            diffusion_ratio=self._diffusion_ratio,
-            consistency_ratio=self._consistency_ratio,
-            device=device,
-            generator=self.cuda_generator,
-        )
-
-        sched = self._flow_map_scheduler
-        n_train = float(self.student.num_train_timesteps)
-        t = (sched.apply_shift(t_norm) * n_train).to(device=device, dtype=dtype)
-        r = (sched.apply_shift(r_norm) * n_train).to(device=device, dtype=dtype)
-
-        # Fresh noise drawn from the method's RNG; ignore the noise that
-        # prepare_batch attached (it pairs with the discarded timestep).
-        noise = torch.randn(
-            latents.shape,
-            device=device,
-            dtype=dtype,
-            generator=self.cuda_generator,
-        )
-        noisy = sched.add_noise(latents, noise, t)
-
-        # Keep training_batch coherent with the new (t, noisy): downstream
-        # forward_context uses these fields.
-        training_batch.timesteps = t
-        training_batch.noise = noise
-        training_batch.noisy_model_input = noisy.permute(0, 2, 1, 3, 4)
-
-        # Student velocity prediction at (t, r).
-        noise_pred = self.student.predict_velocity_with_r(
-            noisy,
-            t,
-            r,
-            training_batch,
-            conditional=True,
-            attn_kind="dense",
-        )
-
-        # Optional guidance distillation — fuse CFG into the training target so
-        # the resulting checkpoint can be sampled at guidance_scale=1.0.
-        if self._fuse_guidance_scale != 1.0:
-            with torch.no_grad():
-                noise_pred_uncond = self.student.predict_velocity_with_r(
-                    noisy,
-                    t,
-                    r,
-                    training_batch,
-                    conditional=False,
-                    attn_kind="dense",
-                )
-            g = float(self._fuse_guidance_scale)
-            noise_pred = (noise_pred - (1.0 - g) * noise_pred_uncond) / g
-
-        dF_dt = _central_difference_dF_dt(
-            student=self.student,
-            batch=training_batch,
-            noisy=noisy,
-            latents=latents,
-            noise=noise,
-            t=t,
-            r=r,
-            delta=self._fd_epsilon,
-            num_train_timesteps=n_train,
-            attn_kind="dense",
-            guidance_scale=self._fuse_guidance_scale,
-        )
-
-        # AnyFlow target: target = (eps - x_0) - (t - r) * dF/dt
-        # dF/dt is in (velocity per absolute t-unit); (t - r) is in absolute units;
-        # the product cancels back to velocity units, matching noise_pred.
-        view = [batch_size] + [1] * (latents.ndim - 1)
-        target = (noise - latents) - (t - r).view(*view) * dF_dt
-
-        # Per-sample squared error, then per-timestep weight, then scale-balance
-        # so the non-diffusion branches stay on the same magnitude as the
-        # diffusion branch (matches AnyFlow's stop-grad rescaling).
-        per_sample = torch.mean(
-            ((noise_pred.float() - target.float())**2).reshape(batch_size, -1),
-            dim=-1,
-        )
-        weight = sched.get_train_weight(t, weight_type=self._weight_type)
-        per_sample = per_sample * weight
-
-        with torch.no_grad():
-            diff_mask = is_diffusion
-            diff_mean = per_sample[diff_mask].mean() if diff_mask.any() else per_sample.mean()
-            non_diff_mask = ~diff_mask
-            if non_diff_mask.any():
-                scale = diff_mean / (per_sample[non_diff_mask] + 1e-5)
-            else:
-                scale = torch.tensor(1.0, device=device, dtype=per_sample.dtype)
-        non_diff_idx = torch.nonzero(non_diff_mask, as_tuple=False).flatten()
-        if non_diff_idx.numel() > 0:
-            per_sample = per_sample.clone()
-            per_sample[non_diff_idx] = per_sample[non_diff_idx] * scale
-
-        total_loss = per_sample.mean()
-
-        loss_map = {"total_loss": total_loss}
-        metrics: dict[str, LogScalar] = {
-            "diffusion_fraction": float(is_diffusion.float().mean()),
-            "consistency_fraction": float(is_consistency.float().mean()),
-            "scale_weight_mean": float(scale.mean()) if isinstance(scale, torch.Tensor) else float(scale),
-        }
-        outputs = {
-            "student_ctx": (
-                training_batch.timesteps,
-                training_batch.attn_metadata,
-            ),
-        }
-        return loss_map, outputs, metrics
-
-    def backward(
-        self,
-        loss_map: dict[str, torch.Tensor],
-        outputs: dict[str, Any],
-        *,
-        grad_accum_rounds: int = 1,
-    ) -> None:
-        """Route the loss backward through the student's forward_context
-        so attn metadata stays attached during gradient computation."""
-        student_ctx = outputs.get("student_ctx")
-        if student_ctx is None:
-            super().backward(loss_map, outputs, grad_accum_rounds=grad_accum_rounds)
-            return
-        self.student.backward(
-            loss_map["total_loss"],
-            student_ctx,
-            grad_accum_rounds=grad_accum_rounds,
-        )
-
-    def _init_optimizer_and_scheduler(self) -> None:
-        tc = self.training_config
-        params = [p for p in self.student.transformer.parameters() if p.requires_grad]
-        (
-            self._student_optimizer,
-            self._student_lr_scheduler,
-        ) = build_optimizer_and_scheduler(
-            params=params,
-            optimizer_config=tc.optimizer,
-            loop_config=tc.loop,
-            learning_rate=float(tc.optimizer.learning_rate),
-            betas=tc.optimizer.betas,
-            scheduler_name=str(tc.optimizer.lr_scheduler),
-        )

fastvideo/train/models/wan/wan.py

@@ -358,53 +358,6 @@ class WanModel(ModelBase):
             pred_noise = transformer(**input_kwargs).permute(0, 2, 1, 3, 4)
         return pred_noise
 
-    def predict_velocity_with_r(
-        self,
-        noisy_latents: torch.Tensor,
-        timestep: torch.Tensor,
-        r_timestep: torch.Tensor,
-        batch: TrainingBatch,
-        *,
-        conditional: bool,
-        cfg_uncond: dict[str, Any] | None = None,
-        attn_kind: Literal["dense", "vsa"] = "dense",
-    ) -> torch.Tensor:
-        """AnyFlow forward: predict average velocity from ``t`` back to ``r``.
-
-        Same plumbing as :meth:`predict_noise` but injects ``r_timestep``
-        into the transformer kwargs. The transformer must have been
-        constructed with an arch config that sets ``r_embedder=True`` for
-        the dual-timestep branch to be active — otherwise ``r_timestep``
-        is silently ignored by the embedder and the forward reduces to
-        the single-timestep path.
-        """
-        device_type = self.device.type
-        dtype = noisy_latents.dtype
-        if conditional:
-            text_dict = batch.conditional_dict
-            if text_dict is None:
-                raise RuntimeError("Missing conditional_dict in "
-                                   "TrainingBatch")
-        else:
-            text_dict = self._get_uncond_text_dict(batch, cfg_uncond=cfg_uncond)
-
-        if attn_kind == "dense":
-            attn_metadata = batch.attn_metadata
-        elif attn_kind == "vsa":
-            attn_metadata = batch.attn_metadata_vsa
-        else:
-            raise ValueError(f"Unknown attn_kind: {attn_kind!r}")
-
-        with torch.autocast(device_type, dtype=dtype), set_forward_context(
-                current_timestep=batch.timesteps,
-                attn_metadata=attn_metadata,
-        ):
-            input_kwargs = (self._build_distill_input_kwargs(noisy_latents, timestep, text_dict))
-            input_kwargs["r_timestep"] = r_timestep
-            transformer = self._get_transformer(timestep)
-            pred_velocity = transformer(**input_kwargs).permute(0, 2, 1, 3, 4)
-        return pred_velocity
-
     def backward(
         self,
         loss: torch.Tensor,

mkdocs.yml

@@ -182,7 +182,6 @@ nav:
   - Distillation:
     - Data Preprocessing: distillation/data_preprocess.md
     - DMD: distillation/dmd.md
-    - AnyFlow: distillation/anyflow.md
   - Attention:
     - Overview: attention/index.md
     - Video Sparse Attention: attention/vsa/index.md

pyproject.toml

@@ -21,9 +21,7 @@ dependencies = [
     "requests>=2.32.2",
 
     # Machine Learning & Transformers
-    # GLM-Image's AR encoder (GlmImageForConditionalGeneration) first ships in
-    # transformers 5.0.0; floor bumped from >=4.57.3 to >=5.0.0 (stable, not rc).
-    "transformers>=5.0.0",
+    "transformers>=4.57.3",
     # <0.23: tokenizers 0.23 renamed RobertaProcessing's binding args, so
     # transformers' CLIP-style tokenizer loading dies with
     # "RobertaProcessing.__new__() got an unexpected keyword argument 'cls'".

pyproject_other.toml

@@ -21,9 +21,7 @@ dependencies = [
     "requests>=2.32.2",
 
     # Machine Learning & Transformers
-    # GLM-Image's AR encoder (GlmImageForConditionalGeneration) first ships in
-    # transformers 5.0.0; floor bumped from >=4.57.3 to >=5.0.0 (stable, not rc).
-    "transformers>=5.0.0",
+    "transformers>=4.57.3",
     "tokenizers>=0.20.1",
     "sentencepiece>=0.2.0",
     "timm>=1.0.11",

scripts/demo_anyflow_14b.py

@@ -1,282 +0,0 @@
-#!/usr/bin/env python3
-# SPDX-License-Identifier: Apache-2.0
-"""FastVideo-side AnyFlow 14B T2V demo at NFE=4 and NFE=50.
-
-Loads ``nvidia/AnyFlow-Wan2.1-T2V-14B-Diffusers`` into FastVideo's
-``WanTransformer3DModel`` (with the ``param_names_mapping`` regex
-handling the ``delta_embedder`` rename), runs the
-``FlowMapEulerDiscreteScheduler`` for the requested NFE schedule, and
-saves the decoded video as MP4. Matches the prompt / shift / guidance
-recipe used by the parallel FastGen demo so the videos are directly
-comparable.
-
-Memory tactics (single H200, 141 GB HBM):
-- Encode prompts with UMT5, free the encoder.
-- Build the FastVideo Wan-14B transformer, load AnyFlow safetensor
-  shards via param_names_mapping translation.
-- Sample at both NFEs without re-loading the transformer.
-- Free the transformer, then load the Wan VAE with tiling for decode.
-
-Configure local checkout paths via env vars (defaults assume a sibling
-layout next to this repo):
-
-  ANYFLOW_LOCAL     — path to ``nvidia/AnyFlow-Wan2.1-T2V-14B-Diffusers``
-                      (default ``./anyflow-14b``)
-  ANYFLOW_DEMO_OUT  — output directory for the rendered MP4s
-                      (default ``./demo_videos``)
-
-Run via::
-
-    PYTHONPATH=$PWD python scripts/demo_anyflow_14b.py
-"""
-
-from __future__ import annotations
-
-import gc
-import os
-import re
-import sys
-import time
-from pathlib import Path
-
-import torch
-from safetensors.torch import load_file
-
-
-ANYFLOW_LOCAL = Path(os.environ.get("ANYFLOW_LOCAL", "./anyflow-14b")).expanduser()
-OUT_DIR = Path(os.environ.get("ANYFLOW_DEMO_OUT", "./demo_videos")).expanduser()
-OUT_DIR.mkdir(parents=True, exist_ok=True)
-
-SEED = 0
-DEVICE = torch.device("cuda")
-DTYPE = torch.bfloat16
-NUM_FRAMES = 81
-HEIGHT, WIDTH = 480, 832
-PROMPT = (
-    "CG game concept digital art, a majestic elephant with a vibrant tusk and sleek fur "
-    "running swiftly towards a herd of its kind. The elephant has a calm yet determined "
-    "expression, with its ears flapping slightly as it moves at high speed. The herd consists "
-    "of several other elephants of various ages and sizes, all moving in unison. The landscape "
-    "is vast savanna with rolling hills, tall grasses, and scattered acacia trees. The sun "
-    "sets behind the horizon, casting a warm golden glow over the scene. Low-angle view, focus "
-    "on the elephant as it accelerates towards the herd."
-)
-NEG_PROMPT = "blurry, low quality, distorted"
-# The published nvidia/AnyFlow-* checkpoints are on-policy distilled with
-# fuse_guidance_scale=3.0 baked into the weights — inference uses 1.0
-# (single conditional forward, no CFG; matches AnyFlow's official demo.py).
-GUIDANCE = 1.0
-
-
-def banner(msg: str) -> None:
-    print("\n" + "=" * 80)
-    print(msg)
-    print("=" * 80)
-
-
-def free_gpu() -> None:
-    gc.collect()
-    torch.cuda.empty_cache()
-    torch.cuda.synchronize()
-    used = torch.cuda.memory_allocated() / 1e9
-    print(f"  [mem] allocated {used:.1f} GB after free")
-
-
-def init_single_rank() -> None:
-    os.environ.setdefault("MASTER_ADDR", "127.0.0.1")
-    os.environ.setdefault("MASTER_PORT", "29571")
-    os.environ.setdefault("WORLD_SIZE", "1")
-    os.environ.setdefault("RANK", "0")
-    os.environ.setdefault("LOCAL_RANK", "0")
-    from fastvideo.distributed import (
-        init_distributed_environment,
-        initialize_model_parallel,
-    )
-    init_distributed_environment(world_size=1, rank=0, local_rank=0, backend="nccl")
-    initialize_model_parallel(
-        tensor_model_parallel_size=1,
-        sequence_model_parallel_size=1,
-        data_parallel_size=1,
-    )
-
-
-def translate_keys(raw: dict, *, mapping: dict[str, str]) -> dict:
-    out: dict = {}
-    for k, v in raw.items():
-        new_k = k
-        for pat, repl in mapping.items():
-            new_k = re.sub(pat, repl, new_k)
-        out[new_k] = v
-    return out
-
-
-def encode_prompts():
-    banner("(1) Encode prompts via UMT5")
-    from transformers import AutoTokenizer, UMT5EncoderModel
-
-    tok = AutoTokenizer.from_pretrained(str(ANYFLOW_LOCAL), subfolder="tokenizer", use_fast=False)
-    enc = UMT5EncoderModel.from_pretrained(
-        str(ANYFLOW_LOCAL), subfolder="text_encoder", torch_dtype=DTYPE,
-    ).to(DEVICE).eval()
-
-    @torch.no_grad()
-    def encode_one(prompts):
-        out = tok(
-            prompts, padding="max_length", max_length=512, truncation=True,
-            return_attention_mask=True, return_tensors="pt")
-        ids = out.input_ids.to(DEVICE)
-        mask = out.attention_mask.to(DEVICE)
-        seq_lens = mask.gt(0).sum(dim=1).long()
-        embeds = enc(ids, mask).last_hidden_state
-        padded = []
-        for i, l in enumerate(seq_lens):
-            e = embeds[i, :l]
-            pad = torch.zeros(512 - e.size(0), e.size(1), device=DEVICE, dtype=embeds.dtype)
-            padded.append(torch.cat([e, pad], dim=0))
-        return torch.stack(padded, dim=0)
-
-    text_e = encode_one([PROMPT])
-    neg_e = encode_one([NEG_PROMPT])
-    print(f"  prompts encoded: text={tuple(text_e.shape)} neg={tuple(neg_e.shape)}")
-    del enc, tok
-    free_gpu()
-    return text_e, neg_e
-
-
-def load_transformer():
-    banner("(2) Build FastVideo Wan-14B + load AnyFlow weights")
-    from fastvideo.configs.models.dits import WanVideoConfig
-    from fastvideo.models.dits.wanvideo import WanTransformer3DModel
-
-    cfg = WanVideoConfig()
-    arch = cfg.arch_config
-    # Wan2.1-T2V-14B arch (per AnyFlow checkpoint config.json).
-    arch.num_attention_heads = 40
-    arch.attention_head_dim = 128
-    arch.num_layers = 40
-    arch.ffn_dim = 13824
-    arch.r_embedder = True
-    arch.r_embedder_fusion = "gated"
-    arch.r_embedder_gate_value = 0.25
-    arch.r_embedder_deltatime_type = "r"
-    arch.__post_init__()
-
-    t0 = time.time()
-    model = WanTransformer3DModel(config=cfg, hf_config={}).to(DEVICE, dtype=DTYPE).eval()
-    print(f"  transformer built in {time.time() - t0:.1f}s; "
-          f"params: {sum(p.numel() for p in model.parameters())/1e9:.2f}B")
-
-    ckpt_dir = ANYFLOW_LOCAL / "transformer"
-    sd: dict = {}
-    for shard in sorted(ckpt_dir.glob("diffusion_pytorch_model-*.safetensors")):
-        sd.update(load_file(str(shard), device="cpu"))
-    print(f"  AnyFlow state dict: {len(sd)} tensors loaded")
-    sd = translate_keys(sd, mapping=arch.param_names_mapping)
-    info = model.load_state_dict(sd, strict=False)
-    print(f"  load: missing={len(info.missing_keys)} unexpected={len(info.unexpected_keys)}")
-    del sd
-    free_gpu()
-    return model
-
-
-@torch.no_grad()
-def sample(model, text_e, neg_e, nfe: int) -> torch.Tensor:
-    banner(f"(3) Sample 14B NFE={nfe}")
-    from fastvideo.forward_context import set_forward_context
-    from fastvideo.models.schedulers.scheduling_flow_map_euler_discrete import (
-        FlowMapEulerDiscreteScheduler, )
-
-    scheduler = FlowMapEulerDiscreteScheduler(num_train_timesteps=1000, shift=5.0)
-    scheduler.set_timesteps(num_inference_steps=nfe, device=DEVICE)
-    timesteps = scheduler.timesteps.to(DEVICE, dtype=DTYPE)
-
-    B, C = 1, 16
-    F = (NUM_FRAMES - 1) // 4 + 1  # temporal VAE compression = 4 (81 → 21)
-    H_l, W_l = HEIGHT // 8, WIDTH // 8
-    g = torch.Generator(device=DEVICE).manual_seed(SEED)
-    x = torch.randn(B, C, F, H_l, W_l, device=DEVICE, dtype=DTYPE, generator=g)
-
-    t0 = time.time()
-    for i, (t_cur, t_next) in enumerate(zip(timesteps[:-1], timesteps[1:])):
-        t_in = t_cur.expand(B).to(DTYPE)
-        r_in = t_next.expand(B).to(DTYPE)
-        with set_forward_context(current_timestep=t_in, attn_metadata=None):
-            flow_cond = model(
-                hidden_states=x, encoder_hidden_states=text_e,
-                timestep=t_in, r_timestep=r_in)
-            if GUIDANCE != 1.0:
-                flow_uncond = model(
-                    hidden_states=x, encoder_hidden_states=neg_e,
-                    timestep=t_in, r_timestep=r_in)
-                flow = flow_uncond + GUIDANCE * (flow_cond - flow_uncond)
-            else:
-                flow = flow_cond
-        x = scheduler.step(
-            flow, sample=x,
-            timestep=t_cur.repeat(B), r_timestep=t_next.repeat(B))
-    print(f"  NFE={nfe} sample time: {time.time() - t0:.1f}s "
-          f"({(time.time() - t0) / nfe:.1f}s/step)")
-    xf = x.float()
-    print(f"  latents mean={xf.mean().item():+.3f} std={xf.std().item():.3f} "
-          f"range=[{xf.min().item():+.2f}, {xf.max().item():+.2f}] "
-          f"finite={torch.isfinite(xf).all().item()}")
-    return x.detach()
-
-
-@torch.no_grad()
-def decode_to_mp4(latents: torch.Tensor, out_path: Path) -> tuple[Path, tuple]:
-    from diffusers.models.autoencoders.autoencoder_kl_wan import AutoencoderKLWan
-    import imageio.v3 as iio
-
-    vae = AutoencoderKLWan.from_pretrained(
-        str(ANYFLOW_LOCAL), subfolder="vae", torch_dtype=DTYPE,
-    ).to(DEVICE).eval()
-    try:
-        vae.enable_tiling()
-        print(f"  VAE tiling enabled")
-    except Exception:
-        pass
-
-    mean = torch.tensor(vae.config.latents_mean, device=DEVICE, dtype=DTYPE).view(1, -1, 1, 1, 1)
-    std = torch.tensor(vae.config.latents_std, device=DEVICE, dtype=DTYPE).view(1, -1, 1, 1, 1)
-    latents_unscaled = latents * std + mean
-    t0 = time.time()
-    frames = vae.decode(latents_unscaled, return_dict=False)[0]
-    print(f"  VAE decode time: {time.time() - t0:.1f}s")
-    frames = (frames.clamp(-1, 1) + 1) / 2
-    frames = frames[0].permute(1, 2, 3, 0).float().cpu().numpy()
-    frames = (frames * 255).astype("uint8")
-    iio.imwrite(str(out_path), frames, fps=16, codec="libx264", quality=8)
-    del vae
-    free_gpu()
-    return out_path, frames.shape
-
-
-def main() -> None:
-    torch.manual_seed(SEED)
-    torch.cuda.manual_seed_all(SEED)
-    init_single_rank()
-
-    text_e, neg_e = encode_prompts()
-    model = load_transformer()
-
-    latents_list = []
-    for nfe in [4, 50]:
-        lat = sample(model, text_e, neg_e, nfe=nfe)
-        latents_list.append((nfe, lat))
-
-    del model, text_e, neg_e
-    free_gpu()
-
-    for nfe, lat in latents_list:
-        banner(f"(4) Decode 14B NFE={nfe}")
-        out_path = OUT_DIR / f"fastvideo_anyflow_14b_nfe{nfe}_seed{SEED}.mp4"
-        path, shape = decode_to_mp4(lat, out_path)
-        print(f"  decoded {shape}")
-        print(f"  saved: {path}  ({path.stat().st_size / 1e6:.2f} MB)")
-
-    banner("DONE")
-
-
-if __name__ == "__main__":
-    main()

scripts/verify_anyflow_fastvideo_parity.py

@@ -1,404 +0,0 @@
-#!/usr/bin/env python3
-# SPDX-License-Identifier: Apache-2.0
-"""FastVideo↔AnyFlow numerical parity verification.
-
-Two checks, both on a single H200:
-
-  (A) Forward parity — load nvidia/AnyFlow-Wan2.1-T2V-1.3B-Diffusers into
-      FastVideo's WanTransformer3DModel (with r_embedder enabled +
-      param_names_mapping handling the delta_embedder rename), forward
-      it on identical inputs against AnyFlow's reference loader, and
-      compare. Expectation: rel-mean diff < 10% in bf16 (bf16 kernel noise).
-
-  (B) Any-step end-to-end sampling — run the new
-      FlowMapEulerDiscreteScheduler through 4 Euler-flow steps on the
-      same loaded weights, confirm the final latent is finite and
-      well-scaled.
-
-Configure local checkout paths via env vars (defaults assume a sibling
-layout next to this repo):
-
-  ANYFLOW_LOCAL  — path to ``nvidia/AnyFlow-Wan2.1-T2V-1.3B-Diffusers``
-                   (default ``./anyflow-1.3b``)
-  ANYFLOW_REF    — path to the NVlabs/AnyFlow reference repo, used to
-                   import its loader (default ``./anyflow-ref``)
-
-Run via::
-
-    PYTHONPATH=$PWD python scripts/verify_anyflow_fastvideo_parity.py
-"""
-
-from __future__ import annotations
-
-import os
-import re
-import sys
-import time
-from pathlib import Path
-
-import torch
-from safetensors.torch import load_file
-
-
-ANYFLOW_LOCAL = Path(os.environ.get("ANYFLOW_LOCAL", "./anyflow-1.3b")).expanduser()
-ANYFLOW_REF = Path(os.environ.get("ANYFLOW_REF", "./anyflow-ref")).expanduser()
-sys.path.insert(0, str(ANYFLOW_REF))
-
-SEED = 1234
-DEVICE = torch.device("cuda")
-DTYPE = torch.bfloat16
-
-
-def banner(msg: str) -> None:
-    print("\n" + "=" * 80)
-    print(msg)
-    print("=" * 80)
-
-
-# ---------------------------------------------------------------------------
-# Distributed bootstrap (single-rank). Required by WanTransformer3DModel
-# which calls get_sp_world_size() and uses ReplicatedLinear.
-# ---------------------------------------------------------------------------
-
-
-def init_single_rank() -> None:
-    os.environ.setdefault("MASTER_ADDR", "127.0.0.1")
-    os.environ.setdefault("MASTER_PORT", "29551")
-    os.environ.setdefault("WORLD_SIZE", "1")
-    os.environ.setdefault("RANK", "0")
-    os.environ.setdefault("LOCAL_RANK", "0")
-    from fastvideo.distributed import (
-        init_distributed_environment,
-        initialize_model_parallel,
-    )
-    init_distributed_environment(world_size=1, rank=0, local_rank=0,
-                                 backend="nccl")
-    initialize_model_parallel(
-        tensor_model_parallel_size=1,
-        sequence_model_parallel_size=1,
-        data_parallel_size=1,
-    )
-    print("  single-rank distributed environment + TP/SP/DP groups initialized")
-
-
-# ---------------------------------------------------------------------------
-# Translate AnyFlow HF safetensor keys onto FastVideo's WanTransformer3DModel
-# internal layout, applying the regex from WanVideoArchConfig.param_names_mapping.
-# ---------------------------------------------------------------------------
-
-
-def translate_keys(
-    raw: dict[str, torch.Tensor],
-    *,
-    mapping: dict[str, str],
-) -> dict[str, torch.Tensor]:
-    out: dict[str, torch.Tensor] = {}
-    for k, v in raw.items():
-        new_k = k
-        for pat, repl in mapping.items():
-            new_k = re.sub(pat, repl, new_k)
-        out[new_k] = v
-    return out
-
-
-# ---------------------------------------------------------------------------
-# Build FastVideo WanTransformer3DModel with AnyFlow weights loaded.
-# ---------------------------------------------------------------------------
-
-
-def build_fastvideo_transformer():
-    banner("(1) Build FastVideo WanTransformer3DModel + load AnyFlow weights")
-    from fastvideo.configs.models.dits import WanVideoConfig
-    from fastvideo.models.dits.wanvideo import WanTransformer3DModel
-
-    cfg = WanVideoConfig()
-    arch = cfg.arch_config
-    # AnyFlow Wan2.1-T2V-1.3B arch dims.
-    arch.num_attention_heads = 12
-    arch.attention_head_dim = 128
-    arch.num_layers = 30
-    arch.ffn_dim = 8960
-    # AnyFlow dual-timestep.
-    arch.r_embedder = True
-    arch.r_embedder_fusion = "gated"
-    arch.r_embedder_gate_value = 0.25
-    arch.r_embedder_deltatime_type = "r"
-    arch.__post_init__()
-
-    # WanTransformer3DModel takes (config, hf_config). hf_config is
-    # diffusers-style — we provide a minimal dict; only fields the model
-    # actually reads matter.
-    hf_config: dict = {}
-    t0 = time.time()
-    model = WanTransformer3DModel(config=cfg, hf_config=hf_config)
-    model = model.to(DEVICE, dtype=DTYPE).eval()
-    print(f"  Wan transformer built in {time.time() - t0:.1f}s; "
-          f"params: {sum(p.numel() for p in model.parameters())/1e9:.2f}B")
-
-    # Load AnyFlow checkpoint.
-    af_path = ANYFLOW_LOCAL / "transformer" / "diffusion_pytorch_model.safetensors"
-    af_raw = load_file(str(af_path), device="cpu")
-    print(f"  AnyFlow checkpoint: {len(af_raw)} tensors")
-
-    translated = translate_keys(af_raw, mapping=arch.param_names_mapping)
-    info = model.load_state_dict(translated, strict=False)
-    miss, unex = info.missing_keys, info.unexpected_keys
-    print(f"  missing_keys    : {len(miss)} (first 5: {miss[:5]})")
-    print(f"  unexpected_keys : {len(unex)} (first 5: {unex[:5]})")
-    return model, len(miss), len(unex)
-
-
-# ---------------------------------------------------------------------------
-# Build AnyFlow reference net.
-# ---------------------------------------------------------------------------
-
-
-def build_anyflow_reference():
-    banner("(2) Build AnyFlow reference loader")
-    from far.models import build_model
-
-    af_net = build_model("FAR_Wan_Transformer3DModel").from_pretrained(
-        str(ANYFLOW_LOCAL),
-        subfolder="transformer",
-        chunk_partition=None,
-        full_chunk_limit=0,
-        compressed_patch_size=[1, 4, 4],
-    ).to(DEVICE, dtype=DTYPE).eval()
-    print(f"  AnyFlow {type(af_net).__name__} ready")
-    return af_net
-
-
-# ---------------------------------------------------------------------------
-# Forward parity test.
-# ---------------------------------------------------------------------------
-
-
-def forward_compare(fv_model, af_net):
-    banner("(3) Forward output comparison")
-    B, C, F, H, W = 1, 16, 21, 60, 104
-    SEQ, DIM = 32, 4096
-    g = torch.Generator(device=DEVICE).manual_seed(SEED)
-    x = torch.randn(B, C, F, H, W, device=DEVICE, dtype=DTYPE, generator=g)
-    enc = torch.randn(B, SEQ, DIM, device=DEVICE, dtype=DTYPE, generator=g)
-
-    # AnyFlow expects per-frame (t, r) [B, F]; FastVideo T2V expects a
-    # single scalar per sample [B] (timestep.dim()==2 is reserved for
-    # Wan2.2 ti2v's per-token timestep schedule). For shared-t T2V the two
-    # are semantically equivalent.
-    t_per_frame = torch.full((B, F), 500.0, device=DEVICE, dtype=DTYPE)
-    r_per_frame = torch.full((B, F), 200.0, device=DEVICE, dtype=DTYPE)
-    t_per_sample = torch.full((B,), 500.0, device=DEVICE, dtype=DTYPE)
-    r_per_sample = torch.full((B,), 200.0, device=DEVICE, dtype=DTYPE)
-
-    with torch.no_grad():
-        # AnyFlow native: takes [B, F, C, H, W] + [B, F] (t, r).
-        x_af = x.permute(0, 2, 1, 3, 4).contiguous()
-        af_out = af_net(
-            x_af,
-            timestep=t_per_frame,
-            r_timestep=r_per_frame,
-            encoder_hidden_states=enc,
-            return_dict=False,
-            is_causal=False,
-        )[0]
-        if af_out.shape[1] != C and af_out.shape[2] == C:
-            af_out = af_out.permute(0, 2, 1, 3, 4).contiguous()
-
-        # FastVideo: [B, C, F, H, W] + [B] (t, r). Must wrap in
-        # set_forward_context so the attention layer can pick up the
-        # current timestep / attn_metadata.
-        from fastvideo.forward_context import set_forward_context
-        with set_forward_context(
-            current_timestep=t_per_sample,
-            attn_metadata=None,
-        ):
-            fv_out = fv_model(
-                hidden_states=x,
-                encoder_hidden_states=enc,
-                timestep=t_per_sample,
-                r_timestep=r_per_sample,
-            )
-
-    print(f"  AnyFlow out: shape={tuple(af_out.shape)} dtype={af_out.dtype}")
-    print(f"  FastVideo  : shape={tuple(fv_out.shape)} dtype={fv_out.dtype}")
-    if af_out.shape != fv_out.shape:
-        print("  ❌ shape mismatch")
-        return False
-    diff = (af_out.float() - fv_out.float()).abs()
-    ref = af_out.float().abs().mean().item() + 1e-12
-    print(f"  max abs diff : {diff.max().item():.3e}")
-    print(f"  mean abs diff: {diff.mean().item():.3e}")
-    print(f"  rel mean diff: {diff.mean().item() / ref:.3e}")
-    ok = diff.mean().item() / ref < 0.10
-    print(f"  >>> {'PASS' if ok else 'FAIL'} (target rel diff < 10%, bf16 noise)")
-    return ok
-
-
-# ---------------------------------------------------------------------------
-# Any-step sampling smoke via FlowMapEulerDiscreteScheduler.
-# ---------------------------------------------------------------------------
-
-
-def sample_anystep(fv_model):
-    banner("(4) Any-step 4-step Euler-flow sampling")
-    from fastvideo.models.schedulers.scheduling_flow_map_euler_discrete import (
-        FlowMapEulerDiscreteScheduler, )
-
-    scheduler = FlowMapEulerDiscreteScheduler(num_train_timesteps=1000, shift=5.0)
-    scheduler.set_timesteps(num_inference_steps=4, device=DEVICE)
-    timesteps = scheduler.timesteps.to(dtype=DTYPE)
-
-    B, C, F, H, W = 1, 16, 21, 60, 104
-    SEQ, DIM = 32, 4096
-    g = torch.Generator(device=DEVICE).manual_seed(SEED)
-    x = torch.randn(B, C, F, H, W, device=DEVICE, dtype=DTYPE, generator=g)
-    enc = torch.randn(B, SEQ, DIM, device=DEVICE, dtype=DTYPE, generator=g)
-
-    from fastvideo.forward_context import set_forward_context
-    t0 = time.time()
-    with torch.no_grad():
-        for t_cur, t_next in zip(timesteps[:-1], timesteps[1:]):
-            t_in = t_cur.expand(B).to(DTYPE)
-            r_in = t_next.expand(B).to(DTYPE)
-            with set_forward_context(current_timestep=t_in, attn_metadata=None):
-                v = fv_model(
-                    hidden_states=x,
-                    encoder_hidden_states=enc,
-                    timestep=t_in,
-                    r_timestep=r_in,
-                )
-            x = scheduler.step(
-                v, sample=x,
-                timestep=t_cur.repeat(B),
-                r_timestep=t_next.repeat(B),
-            )
-    elapsed = time.time() - t0
-    xf = x.float()
-    print(f"  elapsed: {elapsed:.1f}s")
-    print(f"  final latent: mean={xf.mean().item():+.3f} std={xf.std().item():.3f} "
-          f"range=[{xf.min().item():+.2f}, {xf.max().item():+.2f}] "
-          f"finite={torch.isfinite(xf).all().item()}")
-    ok = torch.isfinite(xf).all().item() and 0.01 < xf.std().item() < 30
-    print(f"  >>> {'PASS' if ok else 'FAIL'}")
-    return ok
-
-
-NUM_TRAIN_TIMESTEPS = 1000
-EPSILON = 5.0  # AnyFlow paper default
-
-
-@torch.no_grad()
-def training_step_compare(fv_model, af_net) -> bool:
-    """Inline replica of AnyFlow's train_bidirection central-difference loss
-    on both code paths with identical synthetic (real, noise, t, r) inputs.
-
-    Compares scalar loss + intermediate flow_pred / target tensors.
-    """
-    banner("(5) Training-step loss comparison (central-difference target)")
-    from fastvideo.forward_context import set_forward_context
-
-    B, C, F, H, W = 1, 16, 21, 60, 104
-    SEQ, DIM = 32, 4096
-    g = torch.Generator(device=DEVICE).manual_seed(SEED)
-    real = torch.randn(B, C, F, H, W, device=DEVICE, dtype=DTYPE, generator=g)
-    enc = torch.randn(B, SEQ, DIM, device=DEVICE, dtype=DTYPE, generator=g)
-    noise = torch.randn_like(real)
-    t_abs_pf = torch.full((B, F), 500.0, device=DEVICE, dtype=DTYPE)
-    r_abs_pf = torch.full((B, F), 200.0, device=DEVICE, dtype=DTYPE)
-    t_abs_ps = torch.full((B,), 500.0, device=DEVICE, dtype=DTYPE)
-    r_abs_ps = torch.full((B,), 200.0, device=DEVICE, dtype=DTYPE)
-
-    # AnyFlow inline replica: uses [B, F, C, H, W] layout.
-    real_btchw = real.permute(0, 2, 1, 3, 4).contiguous()
-    noise_btchw = noise.permute(0, 2, 1, 3, 4).contiguous()
-    t_norm_pf = (t_abs_pf / NUM_TRAIN_TIMESTEPS).view(B, F, 1, 1, 1).to(DTYPE)
-    noisy_btchw = t_norm_pf * noise_btchw + (1 - t_norm_pf) * real_btchw
-
-    def u_func_af(x_in, t_in, r_in):
-        return af_net(
-            x_in, timestep=t_in, r_timestep=r_in,
-            encoder_hidden_states=enc, return_dict=False, is_causal=False)[0]
-
-    v_pred = noise_btchw - real_btchw
-    eps = EPSILON
-    F_plus = u_func_af(noisy_btchw + v_pred * (eps / NUM_TRAIN_TIMESTEPS),
-                       t_abs_pf + eps, r_abs_pf)
-    F_minus = u_func_af(noisy_btchw - v_pred * (eps / NUM_TRAIN_TIMESTEPS),
-                        t_abs_pf - eps, r_abs_pf)
-    dF_dt_af = (F_plus - F_minus) / (2 * eps)
-    target_af = ((noise_btchw - real_btchw)
-                 - (t_abs_pf - r_abs_pf).view(B, F, 1, 1, 1) * dF_dt_af)
-    flow_af = u_func_af(noisy_btchw, t_abs_pf, r_abs_pf)
-    loss_af = (flow_af.float() - target_af.float()).pow(2).reshape(B, -1).mean(-1)
-
-    # FastVideo inline replica: uses [B, C, F, H, W] layout + [B] t/r.
-    real_bcfhw = real
-    noise_bcfhw = noise
-    t_norm_ps = (t_abs_ps / NUM_TRAIN_TIMESTEPS).view(B, 1, 1, 1, 1).to(DTYPE)
-    noisy_bcfhw = t_norm_ps * noise_bcfhw + (1 - t_norm_ps) * real_bcfhw
-
-    def u_func_fv(x_in, t_in, r_in):
-        with set_forward_context(current_timestep=t_in, attn_metadata=None):
-            return fv_model(
-                hidden_states=x_in,
-                encoder_hidden_states=enc,
-                timestep=t_in,
-                r_timestep=r_in,
-            )
-
-    v_pred_fv = noise_bcfhw - real_bcfhw
-    F_plus_fv = u_func_fv(
-        noisy_bcfhw + v_pred_fv * (eps / NUM_TRAIN_TIMESTEPS),
-        t_abs_ps + eps, r_abs_ps)
-    F_minus_fv = u_func_fv(
-        noisy_bcfhw - v_pred_fv * (eps / NUM_TRAIN_TIMESTEPS),
-        t_abs_ps - eps, r_abs_ps)
-    dF_dt_fv = (F_plus_fv - F_minus_fv) / (2 * eps)
-    target_fv = ((noise_bcfhw - real_bcfhw)
-                 - (t_abs_ps - r_abs_ps).view(B, 1, 1, 1, 1) * dF_dt_fv)
-    flow_fv = u_func_fv(noisy_bcfhw, t_abs_ps, r_abs_ps)
-    loss_fv = (flow_fv.float() - target_fv.float()).pow(2).reshape(B, -1).mean(-1)
-
-    # Compare. Align AnyFlow's [B, F, C, H, W] → [B, C, F, H, W].
-    flow_af_aligned = flow_af.permute(0, 2, 1, 3, 4)
-    target_af_aligned = target_af.permute(0, 2, 1, 3, 4)
-    flow_diff = (flow_fv.float() - flow_af_aligned.float()).abs()
-    target_diff = (target_fv.float() - target_af_aligned.float()).abs()
-
-    af_loss_v = loss_af.mean().item()
-    fv_loss_v = loss_fv.mean().item()
-    abs_diff = abs(af_loss_v - fv_loss_v)
-    rel_diff = abs_diff / abs(af_loss_v + 1e-12)
-    print(f"  AnyFlow loss : {af_loss_v:.6f}")
-    print(f"  FastVideo loss: {fv_loss_v:.6f}")
-    print(f"  abs diff     : {abs_diff:.3e}")
-    print(f"  rel diff     : {rel_diff:.3e}")
-    print(f"  flow_pred  : max abs {flow_diff.max().item():.3e}  "
-          f"mean {flow_diff.mean().item():.3e}")
-    print(f"  target     : max abs {target_diff.max().item():.3e}  "
-          f"mean {target_diff.mean().item():.3e}")
-    ok = rel_diff < 0.20
-    print(f"  >>> {'PASS' if ok else 'FAIL'} (target rel loss diff < 20%)")
-    return ok
-
-
-def main() -> None:
-    torch.manual_seed(SEED)
-    torch.cuda.manual_seed_all(SEED)
-
-    init_single_rank()
-    fv_model, n_miss, n_unex = build_fastvideo_transformer()
-    af_net = build_anyflow_reference()
-    forward_ok = forward_compare(fv_model, af_net)
-    sample_ok = sample_anystep(fv_model)
-    train_ok = training_step_compare(fv_model, af_net)
-    banner(
-        f"SUMMARY: missing_keys={n_miss} unexpected_keys={n_unex} "
-        f"forward_parity={forward_ok} sample_smoke={sample_ok} "
-        f"training_parity={train_ok}"
-    )
-    sys.exit(0 if (forward_ok and sample_ok and train_ok) else 1)
-
-
-if __name__ == "__main__":
-    main()

tests/local_tests/flux/PORT_STATUS.md

@@ -1,59 +0,0 @@
-# FLUX.1-dev Port Status
-
-Model family: flux
-Official ref: black-forest-labs/FLUX.1-dev (Diffusers layout)
-Workload: T2I
-Last updated: 2026-06-20
-
-## Component Status
-
-| Component | Type | Parity test | Status | Notes |
-|---|---|---|---|---|
-| FluxTransformer2DModel | DiT (ported) | fastvideo/tests/transformers/test_flux.py | PASS (A40, 2026-05-11) | max_diff=0.5, mean_diff=0.04, median=0 — bf16 tail error, see notes |
-| AutoencoderKL | VAE (reused) | tests/local_tests/flux/test_flux_dev_component_loaders.py | PASS (A40, 2026-05-11) | Loader smoke: 54.14s |
-| CLIPTextModel | encoder (reused) | tests/local_tests/flux/test_flux_dev_component_loaders.py | PASS (A40, 2026-05-11) | Shared component |
-| T5EncoderModel | encoder (reused) | tests/local_tests/flux/test_flux_dev_component_loaders.py | PASS (A40, 2026-05-11) | Shared component |
-| FlowMatchEulerDiscreteScheduler | scheduler (reused) | tests/local_tests/flux/test_flux_dev_component_loaders.py | PASS (A40, 2026-05-11) | Shared component |
-
-## Conversion
-
-No conversion script needed — FLUX.1-dev uses native Diffusers checkpoint layout.
-Components load via production loaders with no key remapping.
-Strict-load evidence: PASS — component loader test runs strict load of all components (2026-05-11).
-
-## Pipeline
-
-| Test | Status | Notes |
-|---|---|---|
-| Pipeline smoke (2-step, 256×256) | PASS (A40, 2026-05-11) | 91.14s — output finite, non-null |
-| Pipeline parity vs Diffusers FluxPipeline | PASS (A40, 2026-05-11) | Output finite, in [0,1], correct shape — pixel parity not enforced (pipelines sample noise independently) |
-
-## Quality
-
-| Item | Status |
-|---|---|
-| SSIM test | written (fastvideo/tests/ssim/test_flux_t2i_similarity.py) |
-| Reference images committed | TORCH_SDPA committed (A40, 2026-05-11); FLASH_ATTN pending |
-
-## DiT Parity Notes
-
-DiT parity test compares FastVideo `FluxTransformer2DModel` vs Diffusers under identical inputs (bfloat16, L40S).
-- max_diff=0.5000, mean_diff=0.0382, median_diff=0.0000, p99_diff=0.2500
-- Median=0 and mean=0.04 confirm implementations are equivalent; tail errors up to 0.5 are expected
-  bfloat16 accumulation over 57 transformer layers (eps~7.8e-3, accumulated per-GEMM error ~0.4).
-- Test uses atol=0.5 which catches real bugs (wrong weights/layers produce mean_diff >> 0.1).
-
-## Known Blockers
-
-None. (Both prior blockers resolved — see below.)
-
-## Resolved
-
-1. SSIM reference image — RESOLVED. The TORCH_SDPA reference (A40, 256×256, 8 steps,
-   seed=0) is committed, so the SSIM gate runs against it. A FLASH_ATTN reference is
-   still pending (see Quality table) but is not required for the default backend.
-
-2. Registry registration — RESOLVED. FLUX is registered unconditionally in
-   `_register_configs()` (fastvideo/registry.py); there is no early-return guard that
-   could silently skip it. The earlier `_CONFIG_REGISTRY` pre-population concern no
-   longer applies after the registry refactor.

tests/local_tests/flux/README.md

@@ -1,73 +0,0 @@
-# FLUX.1-dev Local Tests
-
-## Prerequisites
-
-- CUDA GPU
-- `official_weights/FLUX.1-dev` (Diffusers layout) or set `FLUX_DEV_ROOT`
-- `FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA` for parity tests
-
-Download weights:
-```bash
-huggingface-cli download black-forest-labs/FLUX.1-dev --local-dir official_weights/FLUX.1-dev
-```
-
-## Component Loader Smoke Test
-
-Verifies CLIP/T5 encoders, tokenizers, VAE, and FlowMatch scheduler all load
-from the Diffusers checkpoint layout.
-
-```bash
-pytest tests/local_tests/flux/test_flux_dev_component_loaders.py -vs
-```
-
-Status: requires weights — not run in CI.
-
-## Pipeline Smoke Test
-
-Short denoise + decode (2 steps, 256×256) end-to-end.
-
-```bash
-FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA \
-pytest tests/local_tests/pipelines/test_flux_dev_pipeline_smoke.py -vs
-```
-
-Status: requires weights — not run in CI.
-
-## Pipeline Parity Test
-
-Compares FastVideo FluxPipeline decoded image output against Diffusers
-FluxPipeline under identical prompt, seed, and inference parameters (4
-steps, 256×256, seed=42). Tolerance: atol=5e-2.
-
-```bash
-FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA \
-pytest tests/local_tests/pipelines/test_flux_dev_pipeline_parity.py -vs
-```
-
-Status: PASS on A40 (2026-05-11) — requires weights; not run in CI.
-
-## DiT Parity Test
-
-Single-forward comparison of FastVideo `FluxTransformer2DModel` against
-Diffusers `FluxTransformer2DModel` under identical random inputs.
-
-```bash
-FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA \
-pytest fastvideo/tests/transformers/test_flux.py -vs
-```
-
-Status: requires weights — not run in CI.
-Pass evidence: PASS recorded in PORT_STATUS.md (A40, 2026-05-11).
-
-## SSIM Regression Test
-
-Full-quality image generation compared against per-device reference images.
-Reference images must be seeded first via `reference_videos_cli.py`.
-
-```bash
-FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA \
-FLUX_T2I_MODEL_DIR=official_weights/FLUX.1-dev \
-pytest fastvideo/tests/ssim/test_flux_t2i_similarity.py -vs
-```
-
-Status: seeded reference images committed (TORCH_SDPA, A40, 2026-05-11) — see PORT_STATUS.md.

tests/local_tests/flux/test_flux_dev_component_loaders.py

@@ -1,131 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Loader smoke tests for FLUX.1-dev component loading from a local checkpoint.
-
-Verifies that tokenizers, CLIP/T5 text encoders, VAE, and FlowMatch scheduler
-all load correctly from a Diffusers-layout ``FLUX.1-dev`` directory.
-
-Requires ``official_weights/FLUX.1-dev`` (or set ``FLUX_DEV_ROOT``) and CUDA.
-
-Run from repo root::
-
-    pytest tests/local_tests/flux/test_flux_dev_component_loaders.py -vs
-"""
-from __future__ import annotations
-
-import os
-from dataclasses import dataclass, field
-from pathlib import Path
-
-import pytest
-import torch
-
-from fastvideo.configs.models.encoders import (
-    BaseEncoderOutput,
-    CLIPTextConfig,
-    T5LargeConfig,
-)
-from fastvideo.configs.models.vaes.autoencoder_kl import AutoencoderKLVAEConfig
-from fastvideo.configs.pipelines.base import PipelineConfig, preprocess_text
-
-os.environ.setdefault("MASTER_ADDR", "localhost")
-os.environ.setdefault("MASTER_PORT", "29517")
-
-_REPO_ROOT = Path(__file__).resolve().parents[3]
-_FLUX_DEV_ROOT = os.environ.get(
-    "FLUX_DEV_ROOT",
-    str(_REPO_ROOT / "official_weights" / "FLUX.1-dev"),
-)
-
-
-def _flux_clip_post(outputs: BaseEncoderOutput) -> torch.Tensor:
-    if outputs.pooler_output is None:
-        raise RuntimeError("CLIP pooled output required")
-    return outputs.pooler_output
-
-
-def _flux_t5_post(outputs: BaseEncoderOutput) -> torch.Tensor:
-    if outputs.last_hidden_state is None:
-        raise RuntimeError("T5 last_hidden_state required")
-    return outputs.last_hidden_state
-
-
-@dataclass
-class _FluxDevLoaderPipelineConfig(PipelineConfig):
-    """Two text encoders (CLIP + T5) matching FLUX.1-dev model_index.json."""
-
-    vae_config: AutoencoderKLVAEConfig = field(
-        default_factory=AutoencoderKLVAEConfig
-    )
-    text_encoder_configs: tuple[CLIPTextConfig, T5LargeConfig] = field(
-        default_factory=lambda: (CLIPTextConfig(), T5LargeConfig())
-    )
-    text_encoder_precisions: tuple[str, ...] = ("fp32", "bf16")
-    preprocess_text_funcs: tuple = field(
-        default_factory=lambda: (preprocess_text, preprocess_text)
-    )
-    postprocess_text_funcs: tuple = field(
-        default_factory=lambda: (_flux_clip_post, _flux_t5_post)
-    )
-
-
-def test_flux_dev_model_index_components_load() -> None:
-    """CLIP/T5 encoders, tokenizers, VAE, and FlowMatch scheduler load."""
-    if not torch.cuda.is_available():
-        pytest.skip("FLUX component loader test requires CUDA")
-    if not Path(_FLUX_DEV_ROOT, "model_index.json").is_file():
-        pytest.skip(
-            "official_weights/FLUX.1-dev missing "
-            "(set FLUX_DEV_ROOT or download black-forest-labs/FLUX.1-dev)"
-        )
-
-    from fastvideo.distributed import (
-        cleanup_dist_env_and_memory,
-        maybe_init_distributed_environment_and_model_parallel,
-    )
-    from fastvideo.fastvideo_args import FastVideoArgs
-    from fastvideo.models.loader.component_loader import (
-        SchedulerLoader,
-        TextEncoderLoader,
-        TokenizerLoader,
-        VAELoader,
-    )
-
-    maybe_init_distributed_environment_and_model_parallel(1, 1)
-    try:
-        args = FastVideoArgs(
-            model_path=_FLUX_DEV_ROOT,
-            pipeline_config=_FluxDevLoaderPipelineConfig(),
-            hsdp_shard_dim=1,
-            pin_cpu_memory=False,
-        )
-
-        tok_clip = TokenizerLoader().load(
-            os.path.join(_FLUX_DEV_ROOT, "tokenizer"), args
-        )
-        tok_t5 = TokenizerLoader().load(
-            os.path.join(_FLUX_DEV_ROOT, "tokenizer_2"), args
-        )
-        assert "CLIP" in tok_clip.__class__.__name__
-        assert "T5" in tok_t5.__class__.__name__
-
-        te_clip = TextEncoderLoader().load(
-            os.path.join(_FLUX_DEV_ROOT, "text_encoder"), args
-        )
-        te_t5 = TextEncoderLoader().load(
-            os.path.join(_FLUX_DEV_ROOT, "text_encoder_2"), args
-        )
-        # Distributed init may wrap encoders in an FSDP shell.
-        assert te_clip.__class__.__name__.endswith("CLIPTextModel")
-        assert te_t5.__class__.__name__.endswith("T5EncoderModel")
-
-        vae = VAELoader().load(os.path.join(_FLUX_DEV_ROOT, "vae"), args)
-        assert vae.__class__.__name__ == "AutoencoderKL"
-
-        scheduler = SchedulerLoader().load(
-            os.path.join(_FLUX_DEV_ROOT, "scheduler"), args
-        )
-        assert scheduler.__class__.__name__ == "FlowMatchEulerDiscreteScheduler"
-        scheduler.set_timesteps(4, mu=0.7)
-        assert len(scheduler.timesteps) == 4
-    finally:
-        cleanup_dist_env_and_memory()

tests/local_tests/glm_image/README.md

@@ -1,74 +0,0 @@
-# GLM-Image Local Tests
-
-Local-only parity and smoke tests for the `glm_image` port: FastVideo vs the
-official HF `transformers` + `diffusers` GLM-Image implementation. These are
-**not run in CI** — they need the full weights, CUDA, and (for the reference
-side) the diffusers GLM-Image classes; each test skips with an actionable
-message when a prerequisite is missing.
-
-## What you need
-
-| | |
-|---|---|
-| Weights | `zai-org/GLM-Image` → `official_weights/glm_image` (~30 GB: AR encoder 4 shards, transformer 3 shards, vae + text_encoder) |
-| Runtime deps | `transformers>=5.0.0` (first release with the AR encoder `GlmImageForConditionalGeneration`) and `diffusers>=0.38.0` — both committed in `pyproject.toml` |
-| Test-only deps | the diffusers `GlmImageTransformer2DModel` / `GlmImagePipeline` reference classes need `diffusers>=0.37.0.dev0` (used **only** by the parity tests) |
-
-Download weights (from the repo root; cache via `HF_HOME` if `/` is small):
-
-```bash
-python ".agents/skills/add-model-01-prep/scripts/download_hf_weights.py" \
-    "zai-org/GLM-Image" "official_weights/glm_image"
-```
-
-Optionally install the diffusers reference classes to run the parity side:
-
-```bash
-uv pip install -U "git+https://github.com/huggingface/diffusers.git@ff3b86b4"
-```
-
-Parity was checked against transformers `a65bf6c9` + diffusers `ff3b86b4`.
-
-## Parity results (GB200)
-
-| Component | Test | Result |
-|---|---|---|
-| `transformer` (DiT) | `test_transformer_parity.py` | fp32 cosine `1.000000` (100% within 5e-3); bf16 cosine `0.999969` (element-wise rounding over 30 layers vs a different SDPA kernel — not a bug; RoPE proven exactly equivalent). fp32 + bf16 + strict-load |
-| `vae` (`AutoencoderKL`) | `test_vae_parity.py` | decode + encode + config, 3/3 |
-| `text_encoder` (T5) + **ByT5** tokenizer | `test_t5_parity.py` | tokenizer + glyph extract + embeds, 3/3 |
-| `vision_language_encoder` (GLM-4-9B AR) | `test_ar_parity.py` | surface + to/eval + greedy determinism, 3/3 |
-| `pipeline` (T2I) | `test_pipeline_parity.py` | injected-input determinism: latent cosine `0.999997`, decoded-image MAE `0.159/255`. validity + deterministic parity, 2/2 |
-| `pipeline` (I2I / edit) | `test_edit_pipeline_parity.py` | validity / wiring gate (condition image enters via a KV-cache write pass, not noise-the-latent). diffusers distribution parity deferred (stochastic AR) |
-
-`processor` (`GlmImageProcessor`) and `scheduler` (`FlowMatchEulerDiscreteScheduler`,
-shared flow-matching) are exercised via the AR and pipeline tests.
-
-Run:
-
-```bash
-pytest tests/local_tests/glm_image/test_transformer_parity.py -v -s
-pytest tests/local_tests/glm_image/test_vae_parity.py -v -s
-pytest tests/local_tests/glm_image/test_t5_parity.py -v -s
-pytest tests/local_tests/glm_image/test_ar_parity.py -v -s
-pytest tests/local_tests/glm_image/test_pipeline_parity.py -v -s
-pytest tests/local_tests/glm_image/test_edit_pipeline_parity.py -v -s
-```
-
-## Design notes (load-bearing — don't regress)
-
-- **Pipeline parity is by injection.** Full-pipeline pixel parity is ill-posed
-  (stochastic AR + independent latent RNG), so the pipeline tests inject matched
-  `(prompt_embeds, prior_token_ids, latents)` into both the diffusers pipeline
-  and the real `GlmImageDenoisingStage`, then compare denoised latents + decoded
-  image — not raw end-to-end pixels.
-- **HF imports are confined** to the lazy-wrapper loader
-  `fastvideo/models/encoders/glm_image_ar_loader.py`, never the production
-  loaders (production-boundary rule).
-- **Tokenizer is ByT5**, matching the diffusers reference (not plain T5);
-  `compute_glyph_embeds` mirrors diffusers `_get_glyph_embeds`.
-- **Model-specific stages** live under
-  `fastvideo/pipelines/basic/glm_image/stages/` per the `add-model` Files Map.
-- **Param mapping** is handled in-place at load via `param_names_mapping`
-  (`fastvideo/configs/models/dits/glm_image.py`) + the native VAE/encoder
-  configs — no conversion script; the loader raises on any unmatched param and
-  `test_transformer_strict_load.py` asserts completeness with `strict=True`.

tests/local_tests/glm_image/test_ar_parity.py

@@ -1,109 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-from __future__ import annotations
-
-import os
-from pathlib import Path
-
-import pytest
-import torch
-
-os.environ.setdefault("MASTER_ADDR", "localhost")
-os.environ.setdefault("MASTER_PORT", "29519")
-os.environ.setdefault("DISABLE_SP", "1")
-
-REPO_ROOT = Path(__file__).resolve().parents[3]
-FAMILY = "glm_image"
-LOCAL_WEIGHTS_DIR = Path(
-    os.getenv("GLM_IMAGE_LOCAL_WEIGHTS_DIR",
-              REPO_ROOT / "official_weights" / FAMILY))
-AR_DIR = LOCAL_WEIGHTS_DIR / "vision_language_encoder"
-PROCESSOR_DIR = LOCAL_WEIGHTS_DIR / "processor"
-
-
-def _has_weights() -> bool:
-    return AR_DIR.exists() and any(AR_DIR.glob("*.safetensors"))
-
-
-def _has_glm_image_transformers() -> bool:
-    import importlib.util
-    spec = importlib.util.find_spec("transformers")
-    if spec is None:
-        return False
-    import transformers
-    return hasattr(transformers, "GlmImageForConditionalGeneration")
-
-
-pytestmark = [
-    pytest.mark.skipif(
-        not _has_weights(),
-        reason=f"GLM-Image AR encoder weights not found at {AR_DIR}.",
-    ),
-    pytest.mark.skipif(
-        not _has_glm_image_transformers(),
-        reason=("transformers in this env lacks "
-                "`GlmImageForConditionalGeneration` (main pins 4.57.3; needs "
-                "5.0.0rc0+). Bump transformers locally to exercise the AR "
-                "lazy-wrapper."),
-    ),
-]
-
-
-SAMPLE_PROMPT = (
-    "A beautiful landscape photography with rolling hills, a winding river, "
-    "and a vibrant sunset in the background. Photorealistic style."
-)
-
-
-@pytest.fixture(scope="module")
-def device() -> torch.device:
-    if not torch.cuda.is_available():
-        pytest.skip("CUDA required for GLM-Image AR encoder parity.")
-    return torch.device("cuda")
-
-
-@pytest.fixture(scope="module")
-def fastvideo_ar(device):
-    pytest.importorskip("fastvideo")
-    try:
-        from fastvideo.models.encoders.glm_image_ar_loader import (
-            GlmImageARLoader)
-    except ImportError as e:
-        pytest.skip(f"FastVideo lazy-wrapper not yet at target path: {e}")
-    loader = GlmImageARLoader(str(AR_DIR), str(PROCESSOR_DIR),
-                              torch_dtype=torch.bfloat16)
-    loader.to(device).eval()
-    return loader
-
-
-def test_ar_wrapper_exposes_expected_surface(fastvideo_ar):
-    assert callable(getattr(fastvideo_ar, "generate", None))
-    assert getattr(fastvideo_ar, "processor", None) is not None
-    assert getattr(fastvideo_ar, "config", None) is not None
-    assert getattr(fastvideo_ar, "generation_config", None) is not None
-    from transformers import GlmImageForConditionalGeneration
-    assert isinstance(fastvideo_ar._model, GlmImageForConditionalGeneration)
-
-
-def test_ar_to_and_eval_propagate(fastvideo_ar, device):
-    assert next(fastvideo_ar._model.parameters()).device.type == device.type
-    assert not fastvideo_ar._model.training
-
-
-def test_ar_generate_is_deterministic_under_do_sample_false(fastvideo_ar,
-                                                            device):
-    processor = fastvideo_ar.processor
-    messages = [{
-        "role": "user",
-        "content": [{"type": "text", "text": SAMPLE_PROMPT}],
-    }]
-    inputs = processor.apply_chat_template(messages, tokenize=True,
-                                           target_h=1024, target_w=1024,
-                                           return_dict=True,
-                                           return_tensors="pt").to(device)
-    with torch.no_grad():
-        out_a = fastvideo_ar.generate(**inputs, max_new_tokens=32,
-                                      do_sample=False)
-        out_b = fastvideo_ar.generate(**inputs, max_new_tokens=32,
-                                      do_sample=False)
-    assert out_a.shape == out_b.shape
-    assert torch.equal(out_a, out_b)

tests/local_tests/glm_image/test_edit_pipeline_parity.py

@@ -1,126 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""End-to-end scaffold for the GLM-Image image-to-image (edit) pipeline.
-
-Drives FastVideo `VideoGenerator` in edit mode (a condition image is passed) and
-checks that the unified pipeline routes through the condition-encoding stage +
-KV-cache denoising path and emits a well-formed image. Like the T2I parity test,
-pixel-exact parity vs diffusers is NOT asserted: the AR prior is sampled
-(`do_sample=True`) and the diffusion latents draw from independent RNG streams,
-so the codebases produce different valid samples of the same edit. We gate on
-validity (and, when diffusers is present, a comparable brightness regime).
-
-Skips cleanly until weights are available; GPU-heavy.
-"""
-from __future__ import annotations
-
-import os
-from pathlib import Path
-
-import numpy as np
-import pytest
-import torch
-from PIL import Image
-
-os.environ.setdefault("MASTER_ADDR", "localhost")
-os.environ.setdefault("MASTER_PORT", "29520")
-os.environ.setdefault("DISABLE_SP", "1")
-os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "TORCH_SDPA")
-
-REPO_ROOT = Path(__file__).resolve().parents[3]
-FAMILY = "glm_image"
-LOCAL_WEIGHTS_DIR = Path(
-    os.getenv("GLM_IMAGE_LOCAL_WEIGHTS_DIR",
-              REPO_ROOT / "official_weights" / FAMILY))
-CONDITION_IMAGE = REPO_ROOT / "assets" / "images" / "couple.jpg"
-
-
-def _has_weights() -> bool:
-    required = ["transformer", "vae", "text_encoder",
-                "vision_language_encoder", "processor", "tokenizer",
-                "scheduler"]
-    return all((LOCAL_WEIGHTS_DIR / r).exists() for r in required)
-
-
-def _upstream_glm_image_available() -> bool:
-    try:
-        import transformers
-    except ImportError:
-        return False
-    return hasattr(transformers, "GlmImageForConditionalGeneration")
-
-
-pytestmark = [
-    pytest.mark.skipif(
-        not _has_weights(),
-        reason=f"GLM-Image full weights not found at {LOCAL_WEIGHTS_DIR}.",
-    ),
-    pytest.mark.skipif(
-        not _upstream_glm_image_available(),
-        reason=("Edit pipeline needs transformers>=5.0.0rc0 "
-                "(GlmImageForConditionalGeneration); main pin predates it. "
-                "Bump locally to run."),
-    ),
-]
-
-
-EDIT_PROMPT = "Make the scene a snowy winter landscape."
-SEED = 0
-HEIGHT = 512
-WIDTH = 512
-STEPS = 8  # low step count keeps the test under 5 min on a single GPU
-
-
-@pytest.fixture(scope="module")
-def device():
-    if not torch.cuda.is_available():
-        pytest.skip("CUDA required for GLM-Image edit pipeline.")
-    return torch.device("cuda")
-
-
-def _to_uint8_hwc(a) -> np.ndarray:
-    """Coerce a frame buffer (np or torch, [0,1] float or uint8, with optional
-    leading batch/time dims, HWC or CHW) to a single (H, W, 3) uint8 image."""
-    if torch.is_tensor(a):
-        a = a.float().cpu().numpy()
-    a = np.asarray(a)
-    while a.ndim > 3:
-        a = a[0]
-    if a.ndim == 3 and a.shape[0] in (1, 3) and a.shape[-1] not in (1, 3):
-        a = np.transpose(a, (1, 2, 0))
-    if a.dtype != np.uint8:
-        scale = 255.0 if float(a.max()) <= 1.5 else 1.0
-        a = np.clip(a * scale, 0, 255).astype(np.uint8)
-    return a
-
-
-def _fastvideo_edit_image(device) -> np.ndarray:
-    pytest.importorskip("fastvideo")
-    try:
-        from fastvideo import VideoGenerator
-    except ImportError as e:
-        pytest.skip(f"FastVideo VideoGenerator unavailable: {e}")
-    condition = Image.open(CONDITION_IMAGE).convert("RGB")
-    gen = VideoGenerator.from_pretrained(str(LOCAL_WEIGHTS_DIR), num_gpus=1,
-                                         trust_remote_code=True)
-    result = gen.generate_video(prompt=EDIT_PROMPT,
-                                pil_image=condition,
-                                save_video=False,
-                                return_frames=True,
-                                height=HEIGHT, width=WIDTH,
-                                num_inference_steps=STEPS,
-                                guidance_scale=1.5,
-                                seed=SEED)
-    gen.shutdown()
-    return _to_uint8_hwc(result["frames"][0])
-
-
-def test_edit_pipeline_produces_valid_image(device):
-    """Wiring gate for the edit path: passing a condition image routes through
-    the condition-encoding (KV-cache write) stage and the read/skip denoising
-    path, and the pipeline emits a well-formed, non-degenerate image of the
-    requested size. Component numerical correctness is covered by the (green)
-    DiT / VAE / T5 / AR component-parity tests."""
-    fv = _fastvideo_edit_image(device)
-    assert fv.shape == (HEIGHT, WIDTH, 3), f"unexpected shape {fv.shape}"
-    assert fv.dtype == np.uint8 and np.isfinite(fv).all()
-    assert fv.std() > 10.0, f"image is near-constant (std={fv.std():.2f})"

tests/local_tests/glm_image/test_pipeline_parity.py

@@ -1,331 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-from __future__ import annotations
-
-import os
-from pathlib import Path
-
-import numpy as np
-import pytest
-import torch
-
-os.environ.setdefault("MASTER_ADDR", "localhost")
-os.environ.setdefault("MASTER_PORT", "29519")
-os.environ.setdefault("DISABLE_SP", "1")
-os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "TORCH_SDPA")
-
-REPO_ROOT = Path(__file__).resolve().parents[3]
-FAMILY = "glm_image"
-LOCAL_WEIGHTS_DIR = Path(
-    os.getenv("GLM_IMAGE_LOCAL_WEIGHTS_DIR",
-              REPO_ROOT / "official_weights" / FAMILY))
-TRANSFORMER_DIR = LOCAL_WEIGHTS_DIR / "transformer"
-
-
-def _has_weights() -> bool:
-    required = ["transformer", "vae", "text_encoder",
-                "vision_language_encoder", "processor", "tokenizer",
-                "scheduler"]
-    return all((LOCAL_WEIGHTS_DIR / r).exists() for r in required)
-
-
-def _upstream_glm_image_available() -> bool:
-    try:
-        import transformers
-        import diffusers
-    except ImportError:
-        return False
-    return (hasattr(transformers, "GlmImageForConditionalGeneration")
-            and hasattr(diffusers, "GlmImagePipeline"))
-
-
-pytestmark = [
-    pytest.mark.skipif(
-        not _has_weights(),
-        reason=f"GLM-Image full weights not found at {LOCAL_WEIGHTS_DIR}.",
-    ),
-    pytest.mark.skipif(
-        not _upstream_glm_image_available(),
-        reason=("Pipeline parity needs transformers>=5.0.0rc0 and "
-                "diffusers>=0.37.0.dev0; main pins predate both. Bump locally "
-                "to run."),
-    ),
-]
-
-
-SAMPLE_PROMPT = (
-    "A landscape photo with rolling green hills under a clear blue sky.")
-SEED = 0
-HEIGHT = 512
-WIDTH = 512
-STEPS = 8
-
-# bf16 + a different SDPA kernel across 30 DiT layers x STEPS steps leaves a
-# small residual; a real wiring/schedule bug is far larger than these bounds.
-LATENT_COSINE_MIN = 0.995
-IMAGE_MAE_MAX = 5.0  # /255
-
-
-@pytest.fixture(scope="module")
-def device():
-    if not torch.cuda.is_available():
-        pytest.skip("CUDA required for GLM-Image pipeline parity.")
-    return torch.device("cuda")
-
-
-def _to_uint8_hwc(a) -> np.ndarray:
-    if torch.is_tensor(a):
-        a = a.detach().float().cpu().numpy()
-    a = np.asarray(a)
-    while a.ndim > 3:
-        a = a[0]
-    if a.ndim == 3 and a.shape[0] in (1, 3) and a.shape[-1] not in (1, 3):
-        a = np.transpose(a, (1, 2, 0))
-    if a.dtype != np.uint8:
-        scale = 255.0 if float(a.max()) <= 1.5 else 1.0
-        a = np.clip(a * scale, 0, 255).astype(np.uint8)
-    return a
-
-
-# --------------------------------------------------------------------------- #
-# FastVideo native DiT loader (mirrors test_transformer_parity.py).
-# --------------------------------------------------------------------------- #
-def _load_state_dict(dir_path: Path) -> dict[str, torch.Tensor]:
-    safetensors = pytest.importorskip("safetensors.torch")
-    sd: dict[str, torch.Tensor] = {}
-    for shard in sorted(dir_path.glob("*.safetensors")):
-        sd.update(safetensors.load_file(str(shard)))
-    return sd
-
-
-def _apply_param_mapping(sd, mapping):
-    import re
-    out = {}
-    for k, v in sd.items():
-        new_k = k
-        for pat, repl in mapping.items():
-            if re.match(pat, k):
-                new_k = re.sub(pat, repl, k)
-                break
-        out[new_k] = v
-    return out
-
-
-def _ensure_distributed():
-    """The denoising stage calls get_local_torch_device(), which needs
-    FastVideo's world/TP groups. Initialize a single-rank (world_size=1) group
-    in-process (mirrors tests/local_tests/sd35/test_sd35_component_parity.py)."""
-    import torch.distributed as dist
-    from fastvideo.distributed.parallel_state import (
-        get_tp_group, init_distributed_environment, initialize_model_parallel)
-    try:
-        get_tp_group()
-        return
-    except Exception:
-        pass
-    if not dist.is_initialized():
-        os.environ.setdefault("RANK", "0")
-        os.environ.setdefault("WORLD_SIZE", "1")
-        os.environ.setdefault("LOCAL_RANK", "0")
-        torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))
-        store_path = f"/tmp/fastvideo_glm_pg_{os.getpid()}.store"
-        dist.init_process_group(backend="nccl",
-                                init_method=f"file://{store_path}",
-                                rank=0, world_size=1)
-    init_distributed_environment(world_size=1, rank=0, local_rank=0,
-                                 distributed_init_method="env://")
-    try:
-        get_tp_group()
-    except Exception:
-        initialize_model_parallel(tensor_model_parallel_size=1,
-                                  sequence_model_parallel_size=1,
-                                  data_parallel_size=1)
-
-
-def _load_fastvideo_transformer(device, dtype):
-    from fastvideo.configs.models.dits.glm_image import GlmImageDiTConfig
-    from fastvideo.models.dits.glm_image import GlmImageTransformer2DModel
-    cfg = GlmImageDiTConfig()
-    model = GlmImageTransformer2DModel(
-        cfg, {"_class_name": "GlmImageTransformer2DModel"})
-    sd = _load_state_dict(TRANSFORMER_DIR)
-    sd = _apply_param_mapping(sd, cfg.arch_config.param_names_mapping)
-    missing, unexpected = model.load_state_dict(sd, strict=False)
-    assert not missing, f"FastVideo DiT missing keys: {missing[:10]}"
-    assert not unexpected, f"FastVideo DiT unexpected keys: {unexpected[:10]}"
-    return model.to(device, dtype=dtype).eval()
-
-
-def _fastvideo_denoise_latents(device, dtype, *, prompt_embeds, prior_token_ids,
-                               init_latents):
-    """Drive the real GlmImageDenoisingStage with injected, matched inputs and
-    return the denoised latents (1, 16, H/8, W/8)."""
-    from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import (
-        FlowMatchEulerDiscreteScheduler)
-    from fastvideo.pipelines.basic.glm_image.stages.before_denoising import (
-        calculate_shift)
-    from fastvideo.pipelines.basic.glm_image.stages.denoising import (
-        GlmImageDenoisingStage)
-    from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
-
-    _ensure_distributed()
-    transformer = _load_fastvideo_transformer(device, dtype)
-    scheduler = FlowMatchEulerDiscreteScheduler(shift=1.0)
-    stage = GlmImageDenoisingStage(transformer=transformer, scheduler=scheduler)
-
-    # Reproduce the production timestep/sigma schedule (before_denoising.py):
-    # integer-cast linspace timesteps, resolution-dependent shift on sigmas.
-    ntt = scheduler.config.num_train_timesteps
-    patch = transformer.patch_size
-    image_seq_len = ((HEIGHT // 8) * (WIDTH // 8)) // (patch**2)
-    sched_t = np.linspace(ntt, 1.0, STEPS + 1)[:-1].astype(
-        np.int64).astype(np.float32)
-    scheduler.set_shift(calculate_shift(image_seq_len))
-    scheduler.set_timesteps(STEPS, device=device,
-                            sigmas=(sched_t / ntt).tolist(),
-                            timesteps=sched_t.tolist())
-
-    batch = ForwardBatch(data_type="image")
-    batch.prompt_embeds = [prompt_embeds.to(device, dtype)]
-    batch.attention_mask = None  # match diffusers (no text padding mask)
-    batch.prior_token_id = prior_token_ids.to(device)
-    batch.prior_token_drop = torch.zeros_like(prior_token_ids,
-                                              dtype=torch.bool, device=device)
-    batch.latents = init_latents.clone().unsqueeze(2).to(device, dtype)
-    batch.timesteps = scheduler.timesteps
-    batch.height, batch.width = HEIGHT, WIDTH
-    batch.num_inference_steps = STEPS
-    batch.guidance_scale = 1.5
-    batch.do_classifier_free_guidance = True
-    batch.seed = SEED
-    batch.extra = {}  # no glm_kv_caches -> T2I denoise path
-
-    out = stage.forward(batch, fastvideo_args=None)
-    latents = out.latents
-    if latents.dim() == 5:
-        latents = latents.squeeze(2)
-    latents = latents.detach().float().cpu()
-    del transformer, stage
-    torch.cuda.empty_cache()
-    import torch.distributed as dist
-    if dist.is_initialized():
-        dist.destroy_process_group()
-    return latents
-
-
-def _fastvideo_image(device) -> np.ndarray:
-    pytest.importorskip("fastvideo")
-    try:
-        from fastvideo import VideoGenerator
-    except ImportError as e:
-        pytest.skip(f"FastVideo VideoGenerator unavailable: {e}")
-    gen = VideoGenerator.from_pretrained(str(LOCAL_WEIGHTS_DIR), num_gpus=1,
-                                         trust_remote_code=True)
-    result = gen.generate_video(prompt=SAMPLE_PROMPT,
-                                save_video=False,
-                                return_frames=True,
-                                height=HEIGHT, width=WIDTH,
-                                num_inference_steps=STEPS,
-                                guidance_scale=1.5,
-                                seed=SEED)
-    gen.shutdown()
-    return _to_uint8_hwc(result["frames"][0])
-
-
-def test_fastvideo_pipeline_produces_valid_image(device):
-    fv = _fastvideo_image(device)
-    assert fv.shape == (HEIGHT, WIDTH, 3), f"unexpected shape {fv.shape}"
-    assert fv.dtype == np.uint8 and np.isfinite(fv).all()
-    assert fv.std() > 10.0, f"image is near-constant (std={fv.std():.2f})"
-
-
-def test_pipeline_denoise_parity_deterministic(device):
-    """Inject identical (prompt_embeds, prior_token_ids, init_latents) into the
-    diffusers pipeline and FastVideo's GlmImageDenoisingStage, then compare the
-    denoised latents and the decoded images. With the stochastic AR and the
-    independent latent RNG removed, the two implementations must agree to within
-    bf16/SDPA-kernel noise."""
-    # Inference only: the T5 glyph encoder leaves a grad graph on the embeds,
-    # which later breaks image_processor.postprocess()'s .numpy() on the decoded
-    # image. Disabling grad for the whole test also trims memory.
-    torch.set_grad_enabled(False)
-    diffusers = pytest.importorskip("diffusers")
-    from diffusers.utils.torch_utils import randn_tensor
-
-    dtype = torch.bfloat16
-    pipe = diffusers.GlmImagePipeline.from_pretrained(
-        str(LOCAL_WEIGHTS_DIR), torch_dtype=dtype).to(device)
-
-    gen = torch.Generator(device=device).manual_seed(SEED)
-
-    # --- shared, deterministic inputs (computed once, fed to both) ---------- #
-    pos, neg = pipe.encode_prompt(SAMPLE_PROMPT,
-                                  do_classifier_free_guidance=True,
-                                  device=device, dtype=dtype)
-    assert pos.shape[1] == neg.shape[1], (
-        f"quote-free prompt should give equal pos/neg glyph lengths, got "
-        f"{pos.shape[1]} vs {neg.shape[1]}")
-    prior_token_ids, _, _ = pipe.generate_prior_tokens(
-        SAMPLE_PROMPT, HEIGHT, WIDTH, image=None, device=device, generator=gen)
-    latent_ch = pipe.transformer.config.in_channels
-    init_latents = randn_tensor((1, latent_ch, HEIGHT // 8, WIDTH // 8),
-                                generator=gen, device=device, dtype=dtype)
-
-    # --- official denoised latents ----------------------------------------- #
-    # prompt=None: diffusers check_inputs rejects passing prompt + prompt_embeds
-    # together; the embeds (and prior tokens) fully specify the run.
-    official = pipe(prompt=None,
-                    prompt_embeds=pos, negative_prompt_embeds=neg,
-                    prior_token_ids=prior_token_ids,
-                    latents=init_latents.clone(),
-                    height=HEIGHT, width=WIDTH,
-                    num_inference_steps=STEPS, guidance_scale=1.5,
-                    output_type="latent").images.float().cpu()
-
-    # FastVideo packs CFG as a single [pos; neg] 2-row tensor (no padding here
-    # because L_pos == L_neg, asserted above).
-    packed = torch.cat([pos, neg], dim=0)
-
-    # --- FastVideo denoised latents (real GlmImageDenoisingStage) ---------- #
-    fv = _fastvideo_denoise_latents(device, dtype,
-                                    prompt_embeds=packed,
-                                    prior_token_ids=prior_token_ids,
-                                    init_latents=init_latents)
-
-    assert official.shape == fv.shape, (
-        f"latent shape mismatch: {official.shape} vs {fv.shape}")
-
-    # --- latent-level agreement -------------------------------------------- #
-    cos = torch.nn.functional.cosine_similarity(
-        official.flatten(), fv.flatten(), dim=0).item()
-    lat_mae = (official - fv).abs().mean().item()
-    lat_scale = official.abs().mean().item()
-    print(f"[glm-image denoise parity] latent cosine={cos:.6f} "
-          f"MAE={lat_mae:.5f} (|latent| mean={lat_scale:.5f}, "
-          f"rel={lat_mae / max(lat_scale, 1e-6):.4f})")
-
-    # --- decoded-image agreement (same VAE both sides isolates denoise) ----- #
-    def _decode(lat):
-        lat = lat.to(device, dtype)
-        mean = torch.tensor(pipe.vae.config.latents_mean).view(
-            1, -1, 1, 1).to(device, dtype)
-        std = torch.tensor(pipe.vae.config.latents_std).view(
-            1, -1, 1, 1).to(device, dtype)
-        img = pipe.vae.decode((lat * std + mean), return_dict=False)[0]
-        return _to_uint8_hwc(pipe.image_processor.postprocess(
-            img, output_type="np")[0])
-
-    off_img, fv_img = _decode(official), _decode(fv)
-    img_mae = np.abs(off_img.astype(np.float32)
-                     - fv_img.astype(np.float32)).mean()
-    print(f"[glm-image denoise parity] decoded image MAE={img_mae:.3f}/255 "
-          f"(diffusers mean={off_img.mean():.1f}, fv mean={fv_img.mean():.1f})")
-
-    del pipe
-    torch.cuda.empty_cache()
-
-    assert cos > LATENT_COSINE_MIN, (
-        f"denoised-latent cosine {cos:.6f} < {LATENT_COSINE_MIN}: the FastVideo "
-        "and official denoise diverge well beyond kernel noise (wiring bug).")
-    assert img_mae < IMAGE_MAE_MAX, (
-        f"decoded-image MAE {img_mae:.3f}/255 >= {IMAGE_MAE_MAX}: outputs are "
-        "not equivalent under matched inputs (wiring/schedule bug).")