Merge branch 'main' into test-hf-sync

Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

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

@@ -29,8 +29,8 @@
     "Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of a bustling street food market. The scene captures a mix of humor and authenticity. Mid-shot framing, vibrant lighting."
   ],
   "run_config": {
-    "num_warmup_runs": 1,
-    "num_measurement_runs": 3,
+    "num_warmup_runs": 2,
+    "num_measurement_runs": 5,
     "required_gpus": 2
   },
   "thresholds": {

.buildkite/scripts/pr_test.sh

@@ -63,7 +63,72 @@ EFFECTIVE_PR=${BUILDKITE_PULL_REQUEST:-false}
 if [ "$EFFECTIVE_PR" = "false" ] && [ -n "${PR_NUMBER:-}" ]; then
     EFFECTIVE_PR=$PR_NUMBER
 fi
-MODAL_ENV="BUILDKITE_REPO=$BUILDKITE_REPO BUILDKITE_COMMIT=$BUILDKITE_COMMIT BUILDKITE_PULL_REQUEST=$EFFECTIVE_PR IMAGE_VERSION=$IMAGE_VERSION"
+MODAL_ENV="BUILDKITE_REPO=$BUILDKITE_REPO BUILDKITE_COMMIT=$BUILDKITE_COMMIT BUILDKITE_PULL_REQUEST=$EFFECTIVE_PR BUILDKITE_BRANCH=${BUILDKITE_BRANCH:-} TEST_SCOPE=${TEST_SCOPE:-} IMAGE_VERSION=$IMAGE_VERSION"
+
+POST_RUN_HOOK=""
+
+upload_performance_artifacts() {
+    SHORT_SHA=${BUILDKITE_COMMIT:0:7}
+    LOCAL_DIR="downloaded_reports"
+
+    _download_reports() {
+        log "Downloading perf_reports/ from Modal Volume..."
+        mkdir -p "$LOCAL_DIR"
+        if ! modal volume get hf-model-weights "perf_reports/" "$LOCAL_DIR"; then
+            log "Error: Failed to download perf_reports/ from Modal Volume."
+            return 1
+        fi
+    }
+
+    _upload_dashboard() {
+        local target
+        target=$(find "$LOCAL_DIR" -name "dashboard_*${SHORT_SHA}*" | head -n 1)
+        log "TARGET dashboard: '$target'"
+
+        if [ -n "$target" ]; then
+            log "Found dashboard: $target. Uploading to Buildkite..."
+            buildkite-agent artifact upload "$target"
+            buildkite-agent annotate --style info --context "perf-dashboard" < "$target"
+        else
+            log "Warning: Could not find a dashboard file matching $SHORT_SHA"
+        fi
+    }
+
+    _upload_perf_summary() {
+        local target
+        target=$(find "$LOCAL_DIR" -name "perf_*${SHORT_SHA}*" | head -n 1)
+        log "TARGET perf summary: '$target'"
+
+        if [ -n "$target" ]; then
+            log "Found perf summary: $target. Uploading to Buildkite..."
+            buildkite-agent artifact upload "$target"
+            buildkite-agent annotate --style info --context "perf-summary" < "$target"
+        else
+            log "Warning: Could not find a perf summary file matching $SHORT_SHA"
+        fi
+    }
+
+    _cleanup_modal_volume() {
+        log "Cleaning up perf_reports/ from Modal Volume..."
+        if modal volume rm hf-model-weights "perf_reports/" --recursive; then
+            log "Successfully deleted perf_reports/ from Modal Volume."
+        else
+            log "Warning: Failed to delete perf_reports/ from Modal Volume. Manual cleanup may be required."
+        fi
+    }
+
+    _cleanup_local() {
+        log "Cleaning up local download directory..."
+        rm -rf "$LOCAL_DIR"
+    }
+
+    # --- Main flow ---
+    _download_reports || { _cleanup_local; return 1; }
+    _upload_dashboard
+    _upload_perf_summary
+    _cleanup_modal_volume
+    _cleanup_local
+}
 
 case "$TEST_TYPE" in
     "encoder")
@@ -124,8 +189,9 @@ case "$TEST_TYPE" in
         MODAL_COMMAND="$MODAL_ENV HF_API_KEY=$HF_API_KEY python3 -m modal run $MODAL_TEST_FILE::run_lora_extraction_tests"
         ;;
     "performance")
-        log "Running performance tests..."
+        log "Running performance tests on Modal..."
         MODAL_COMMAND="$MODAL_ENV HF_API_KEY=$HF_API_KEY python3 -m modal run $MODAL_TEST_FILE::run_performance_tests"
+        POST_RUN_HOOK="upload_performance_artifacts"
         ;;
     "api_server")
         log "Running API server integration tests..."
@@ -147,5 +213,10 @@ else
     log "Error: Modal test failed with exit code: $TEST_EXIT_CODE"
 fi
 
+if [ -n "$POST_RUN_HOOK" ]; then
+    log "Executing post-run hook: $POST_RUN_HOOK"
+    "$POST_RUN_HOOK"
+fi
+
 log "=== Test execution completed with exit code: $TEST_EXIT_CODE ==="
 exit $TEST_EXIT_CODE

.pre-commit-config.yaml

@@ -18,6 +18,7 @@ exclude: |
     fastvideo/train\.py|
     fastvideo/utils/.*|
     examples/.*|
+    \.agents/.*|
     .github/workflows/publish-fastvideo.yml|
     .github/workflows/_template-build-image.yml|
     docs/source/inference/support_matrix.md

docs/contributing/coding_agents.md

@@ -296,8 +296,10 @@ Action:
 
 - Add or reuse a numerical parity test that loads the official model and the
   FastVideo model and compares outputs.
-- See examples in `tests/local_tests/` (e.g., `tests/local_tests/upsamplers/`)
-  and the commands in `tests/local_tests/README.md`.
+- See examples in `tests/local_tests/` organized by model family
+  (e.g., `tests/local_tests/sd35/`, `tests/local_tests/ltx2/`,
+  `tests/local_tests/stable_audio/`) and the navigation index in
+  `tests/local_tests/README.md`.
 - If there are discrepancies, add opt‑in logging to both models and compare
   activation summaries (layer output sums, per‑stage logs).
 - First align the loaded weights (validate `param_names_mapping`).
@@ -348,7 +350,8 @@ Purpose:
 
 Action:
 
-- Add a pipeline parity test under `tests/local_tests/pipelines/`.
+- Add a pipeline parity test under `tests/local_tests/<family>/`
+  (e.g., `tests/local_tests/<family>/test_<family>_pipeline_parity.py`).
 - See the [Testing Guide](testing.md) for test conventions.
 
 ### 7) Add user‑facing examples

examples/inference/basic/basic_stable_audio.py

@@ -0,0 +1,77 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Stable Audio Open 1.0 — text-to-audio (baseline) example.
+
+User story (game-audio designer, prototyping):
+    "I'm prototyping a level and I need 6 seconds of background
+    ambience — gentle wind, distant thunder, a hint of birdsong. I
+    don't want to dig through a sound library; I want to type what I
+    hear in my head and get a wav back. If it's wrong I'll iterate
+    on the prompt. This is the first stop."
+
+User story (musician sketching ideas):
+    "I want to bounce a 30s lo-fi drum loop to use as a placeholder
+    bed while I build the rest of the track. Type prompt, get audio,
+    drop into the DAW. The actual production beat I'll record
+    myself, but I need *something* to write the chords against."
+
+User story (researcher exploring the model):
+    "First time touching Stable Audio Open — what does it sound
+    like at default settings? This is the smallest amount of code
+    that goes from prompt to mp4."
+
+How it works:
+    Pure text-to-audio (T2A). The pipeline runs:
+        T5 + NumberConditioner -> StableAudioDiT -> Oobleck VAE
+    via the `dpmpp-3m-sde` k-diffusion sampler. All components are
+    FastVideo-native — no diffusers / transformers model imports at
+    runtime (see REVIEW item 30). Mirrors upstream
+    `stable_audio_tools.inference.generation.generate_diffusion_cond`
+    bit-for-bit (~0.2% abs_mean drift on 25 steps).
+
+Tunable knobs (the "creative dials"):
+    audio_end_in_s
+        1–6   — quick ideation (sub-10s wall clock at 100 steps)
+        10–30 — full musical phrase / loop length (the README example
+                uses 30s)
+        47.5  — model maximum (full sample_size = 2097152 / 44100 Hz)
+    num_inference_steps
+        25  — fast preview, occasional artifacts
+        100 — preset default (matches the HF model card)
+        250 — diminishing returns past here
+    guidance_scale
+        3   — looser, more variation per seed
+        7   — preset default; matches README
+        12+ — sharper but can sound "fried"
+
+Prerequisites:
+  1. Accept the terms on https://huggingface.co/stabilityai/stable-audio-open-1.0
+     and export your HF token in the shell:
+         export HF_TOKEN=hf_...
+  2. Install optional inference deps (one-time):
+         pip install k_diffusion einops_exts alias_free_torch torchsde
+"""
+from fastvideo import VideoGenerator
+
+PROMPT = "Lo-fi hip hop instrumental with vinyl crackle and gentle piano."
+
+
+def main() -> None:
+    generator = VideoGenerator.from_pretrained(
+        "FastVideo/stable-audio-open-1.0-Diffusers",
+        num_gpus=1,
+    )
+    output_path = "outputs_audio/stable_audio_basic/output_stable_audio.wav"
+    generator.generate_video(
+        prompt=PROMPT,
+        output_path=output_path,
+        save_video=True,
+        # 6-second clip; the model max is ~47.5s.
+        audio_end_in_s=6.0,
+        # The registered preset gives 100 steps + CFG=7.0 by default;
+        # override num_inference_steps / guidance_scale here for QA.
+    )
+    generator.shutdown()
+
+
+if __name__ == "__main__":
+    main()

examples/inference/basic/basic_stable_audio_a2a.py

@@ -0,0 +1,77 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Stable Audio Open 1.0 — audio-to-audio variation example.
+
+User story (musician, late at night):
+    "I generated this 12-second lo-fi loop earlier and I love the chord
+    progression and overall vibe, but the snare hit at 0:08 sounds wrong
+    and the rhythm feels stiff. I don't want to start over from scratch
+    and lose what's working — I want the model to keep the harmony and
+    mood but reroll the percussion + groove."
+
+User story (sound designer, on a deadline):
+    "I have one good 'sword clang' SFX. The art director wants 8 sibling
+    variations that all feel like the same sword from different angles —
+    same metal, same weight, slightly different impact. I'd rather
+    refine my one good take than text-prompt my way through 50 misses."
+
+Pass `init_audio=path/to/clip` (any wav/mp3/mp4/m4a/flac the standard
+deps decode) and the model will use it as a starting point for the
+text prompt instead of pure noise.
+
+Picking `init_audio_strength` (0.0 to 1.0):
+
+    Higher = closer to the source clip. Lower = more transformation.
+    (Same convention as the "Input Audio Strength" slider in
+    Stability's commercial Stable Audio web UI, so values transfer
+    directly.)
+
+      | strength | what you get                                       |
+      |----------|----------------------------------------------------|
+      |   1.00   | Output ≈ reference. No transformation.             |
+      |   0.85   | Texture micro-variation only.                      |
+      |   0.70   | Light reroll, same instruments.                    |
+      |   0.60   | Default. Instrument identity is replaceable        |
+      |          | (cello can take over from piano on the same notes).|
+      |   0.50   | Heavy — only melody / chord progression survives.  |
+      |   0.30   | Reference acts as a loose mood prompt.             |
+      |   0.00   | Plain T2A — reference ignored.                     |
+
+    Rule of thumb by intent:
+      * "Fix one part of this clip"            -> 0.75 .. 0.85
+      * "Same notes, different instrument"     -> 0.55 .. 0.65
+      * "Same chord progression, new content"  -> 0.40 .. 0.55
+      * "Use this as a loose mood prompt"      -> 0.20 .. 0.35
+
+    If the reference timbre is bleeding through more than you want,
+    lower it; if the structure is gone, raise it.
+
+Prerequisites: same as `basic_stable_audio.py`.
+"""
+from fastvideo import VideoGenerator
+
+PROMPT = "Change the piano to a cello playing the same notes"
+# Path to any audio-bearing file (wav, mp3, mp4, m4a, flac, ...).
+# Set to `None` to skip A2A and run plain T2A.
+INIT_AUDIO_PATH: str | None = None
+# Reference fidelity in [0, 1] -- higher = closer to source.
+INIT_AUDIO_STRENGTH = 0.6
+
+
+def main() -> None:
+    generator = VideoGenerator.from_pretrained(
+        "FastVideo/stable-audio-open-1.0-Diffusers",
+        num_gpus=1,
+    )
+    generator.generate_video(
+        prompt=PROMPT,
+        output_path="outputs_audio/stable_audio_a2a/output_a2a.wav",
+        save_video=True,
+        audio_end_in_s=6.0,
+        init_audio=INIT_AUDIO_PATH,
+        init_audio_strength=INIT_AUDIO_STRENGTH,
+    )
+    generator.shutdown()
+
+
+if __name__ == "__main__":
+    main()

examples/inference/basic/basic_stable_audio_inpaint.py

@@ -0,0 +1,84 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Stable Audio Open 1.0 — inpainting / outpainting (loop extension) example.
+
+User story (loop extension — the killer app):
+    "I have a 6-second drum loop my client likes. They want it as
+    background bed for a 30-second ad. I need it to loop seamlessly,
+    but a hard cut every 6s sounds bad. Let me extend it to 30s,
+    keeping the first 6s exactly as-is and letting the model continue
+    the groove for the remaining 24s."
+
+User story (audio repair):
+    "There's a microphone bump at 0:14 in this 30-second field
+    recording — really obvious in headphones. Mask out 0:13 to 0:15
+    and let the model regenerate plausible ambience that blends in.
+    Everything else stays exactly as I recorded it."
+
+User story (transition smoothing):
+    "I have two 10-second clips I want to crossfade. Mask out a 1s
+    overlap region in the middle and let the model invent a coherent
+    transition between the two."
+
+How it works (RePaint-style blending):
+    Stable Audio Open 1.0 wasn't trained as an inpainting model
+    (`model_type=diffusion_cond`, not `diffusion_cond_inpaint`), so we
+    can't use the upstream's mask-conditioned approach directly. We
+    use the RePaint trick instead, which works on any v-prediction
+    diffusion model:
+
+      1. Encode the reference clip into latent space.
+      2. At every denoising step `i`, replace the kept region of the
+         in-flight latent (where mask == 1) with the reference
+         re-noised to the next timestep's sigma. Only the unkept
+         region (mask == 0) is freely denoised.
+      3. After the loop, the kept region is exactly the reference;
+         the unkept region is freshly generated content.
+
+    This is approximate compared to a properly trained inpainting
+    checkpoint — the seam between kept/unkept can have slight EQ
+    discontinuity — but it works on the existing public model.
+
+Tunable: the mask is a 1-D tensor in {0, 1} at the model's sample
+rate. Conventions:
+    1.0 = keep this sample from the reference
+    0.0 = regenerate this sample
+
+Prerequisites: same as `basic_stable_audio.py`.
+"""
+import os
+
+from fastvideo import VideoGenerator
+
+PROMPT = "Steady lo-fi hip hop drum loop with vinyl crackle."
+# Required: path to the reference audio file (wav, mp3, mp4, m4a, flac,
+# ...) you want to extend or repair. The pipeline raises if a mask is
+# passed without a reference, so this must be a real path.
+REFERENCE_AUDIO_PATH = "path/to/your/loop.wav"
+KEEP_SECONDS = 6.0       # first KEEP_SECONDS preserved exactly
+TOTAL_SECONDS = 12.0     # extend the loop to this duration
+
+
+def main() -> None:
+    if not os.path.isfile(REFERENCE_AUDIO_PATH):
+        raise FileNotFoundError(
+            f"REFERENCE_AUDIO_PATH={REFERENCE_AUDIO_PATH!r} does not exist. "
+            "Edit this script to point at a real audio file (wav/mp3/mp4/"
+            "m4a/flac) before running.")
+    generator = VideoGenerator.from_pretrained(
+        "FastVideo/stable-audio-open-1.0-Diffusers",
+        num_gpus=1,
+    )
+    generator.generate_video(
+        prompt=PROMPT,
+        output_path="outputs_audio/stable_audio_inpaint/output_inpaint.wav",
+        save_video=True,
+        audio_end_in_s=TOTAL_SECONDS,
+        inpaint_audio=REFERENCE_AUDIO_PATH,
+        # Tuple form: keep first KEEP_SECONDS, regenerate the rest.
+        inpaint_mask=(KEEP_SECONDS, TOTAL_SECONDS),
+    )
+    generator.shutdown()
+
+
+if __name__ == "__main__":
+    main()

examples/inference/basic/basic_stable_audio_small.py

@@ -0,0 +1,53 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Stable Audio Open Small — fast / lightweight T2A example.
+
+User story (interactive UI builder):
+    "I'm building a sound-design UI where the user types a prompt and
+    we want sub-2-second feedback so the experience feels like
+    autocomplete, not a render queue. The full Stable Audio Open 1.0
+    takes ~8s on a single GPU; the small variant takes a fraction of
+    that — quality is lower but completely usable for real-time
+    iteration."
+
+User story (overnight batch jobs):
+    "I'm generating 10,000 short SFX variants for a procedural game.
+    Wall-clock matters more than per-clip polish — give me the small
+    model so I can fit the run in one night instead of a week."
+
+How it works:
+    The small variant is a separate Stability AI checkpoint
+    (`stabilityai/stable-audio-open-small`) that ships the same Oobleck
+    VAE as the 1.0 base model but a smaller / faster DiT (`embed_dim=1024`,
+    `depth=16`, `qk_norm="ln"`) and only one duration conditioner
+    (`seconds_total`, no `seconds_start`). FastVideo loads from the
+    converted Diffusers-format repo `FastVideo/stable-audio-open-small-Diffusers`
+    via the standard component loader; per-variant arch fields come
+    from `transformer/config.json` and `conditioner/config.json`.
+
+Prerequisites: same as `basic_stable_audio.py`. The converted repo is
+public so no gated-access flow is required.
+"""
+from fastvideo import VideoGenerator
+
+PROMPT = "Lo-fi hip hop instrumental with vinyl crackle and gentle piano."
+
+
+def main() -> None:
+    generator = VideoGenerator.from_pretrained(
+        "FastVideo/stable-audio-open-small-Diffusers",
+        num_gpus=1,
+    )
+    output_path = "outputs_audio/stable_audio_small/output_stable_audio_small.wav"
+    generator.generate_video(
+        prompt=PROMPT,
+        output_path=output_path,
+        save_video=True,
+        # Small variant trains on a ~11.9s window — keep `audio_end_in_s`
+        # at or below that.
+        audio_end_in_s=6.0,
+    )
+    generator.shutdown()
+
+
+if __name__ == "__main__":
+    main()

examples/train/configs/fine_tuning/cosmos/t2v.yaml

@@ -0,0 +1,73 @@
+# Cosmos Predict2 2B T2V finetune config.
+#
+# Data must be preprocessed with Cosmos VAE + T5 text encoder
+# into parquet format before training.
+
+models:
+  student:
+    _target_: fastvideo.train.models.cosmos.CosmosModel
+    init_from: nvidia/Cosmos-Predict2-2B-Video2World
+    trainable: true
+
+method:
+  _target_: fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod
+
+training:
+  distributed:
+    num_gpus: 8
+    sp_size: 1
+    tp_size: 1
+    hsdp_replicate_dim: 8
+    hsdp_shard_dim: 1
+
+  data:
+    data_path: data/cosmos_preprocessed
+    dataloader_num_workers: 4
+    train_batch_size: 1
+    training_cfg_rate: 0.0
+    seed: 1000
+    # Cosmos VAE: 4x temporal, 8x spatial compression.
+    # 93 frames -> 24 latent frames, 480x832 -> 60x104
+    num_latent_t: 24
+    num_height: 480
+    num_width: 832
+    num_frames: 93
+
+  optimizer:
+    learning_rate: 1.0e-5
+    betas: [0.9, 0.999]
+    weight_decay: 0.01
+    lr_scheduler: constant
+    lr_warmup_steps: 0
+
+  loop:
+    max_train_steps: 5000
+    gradient_accumulation_steps: 1
+
+  checkpoint:
+    output_dir: outputs/cosmos_finetune
+    training_state_checkpointing_steps: 500
+    checkpoints_total_limit: 3
+    resume_from_checkpoint: latest
+
+  tracker:
+    project_name: fastvideo_cosmos
+    run_name: cosmos_finetune
+
+  model:
+    enable_gradient_checkpointing_type: full
+
+callbacks:
+  grad_clip:
+    _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback
+    max_grad_norm: 1.0
+  validation:
+    _target_: fastvideo.train.callbacks.validation.ValidationCallback
+    pipeline_target: fastvideo.pipelines.basic.cosmos.cosmos_pipeline.Cosmos2VideoToWorldPipeline
+    dataset_file: data/cosmos_preprocessed/validation_prompts.json
+    every_steps: 100
+    sampling_steps: [50]
+    guidance_scale: 6.0
+
+pipeline:
+  flow_shift: 1.0

examples/train/configs/overfit_cosmos25_t2w.yaml

@@ -0,0 +1,79 @@
+# Cosmos-Predict2.5-2B Text-to-World overfitting test config.
+#
+# Overfits on a few short videos (480x832, 93 frames) to verify the
+# Cosmos 2.5 training plugin works end-to-end.
+#
+# Preprocess data first:
+#   CUDA_VISIBLE_DEVICES=0 python fastvideo/pipelines/preprocess/preprocess_cosmos25_overfit.py
+#
+# Run:
+#   bash examples/train/run.sh examples/train/configs/overfit_cosmos25_t2w.yaml
+
+models:
+  student:
+    _target_: fastvideo.train.models.cosmos.CosmosModel
+    init_from: KyleShao/Cosmos-Predict2.5-2B-Diffusers
+    trainable: true
+    enable_gradient_checkpointing_type: full
+    flow_shift: 1.0
+
+method:
+  _target_: fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod
+
+training:
+  distributed:
+    num_gpus: 1
+    sp_size: 1
+    tp_size: 1
+    hsdp_replicate_dim: 1
+    hsdp_shard_dim: 1
+
+  data:
+    data_path: data/cosmos25_overfit_preprocessed
+    dataloader_num_workers: 0
+    train_batch_size: 1
+    training_cfg_rate: 0.0
+    seed: 42
+    num_latent_t: 24
+    num_height: 480
+    num_width: 832
+    num_frames: 93
+
+  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: 300
+    gradient_accumulation_steps: 1
+
+  checkpoint:
+    output_dir: outputs/cosmos25_overfit
+    training_state_checkpointing_steps: 50
+    checkpoints_total_limit: 2
+
+  tracker:
+    project_name: fastvideo_cosmos25
+    run_name: cosmos25_overfit
+
+  model:
+    precondition_outputs: false
+    enable_gradient_checkpointing_type: full
+
+callbacks:
+  grad_clip:
+    _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback
+    max_grad_norm: 1.0
+  validation:
+    _target_: fastvideo.train.callbacks.validation.ValidationCallback
+    pipeline_target: fastvideo.pipelines.basic.cosmos.cosmos2_5_pipeline.Cosmos2_5Pipeline
+    dataset_file: data/cosmos25_overfit_preprocessed/validation_prompts.json
+    every_steps: 150
+    sampling_steps: [35]
+    guidance_scale: 7.0
+
+pipeline:
+  flow_shift: 1.0

fastvideo/api/results.py

@@ -15,6 +15,7 @@ class GenerationResult:
     samples: Any | None = None
     frames: Any | None = None
     audio: Any | None = None
+    audio_sample_rate: int | None = None
     size: tuple[int, int, int] | None = None
     generation_time: float | None = None
     logging_info: Any | None = None
@@ -44,6 +45,7 @@ class GenerationResult:
                 "samples",
                 "frames",
                 "audio",
+                "audio_sample_rate",
                 "size",
                 "generation_time",
                 "logging_info",
@@ -62,6 +64,7 @@ class GenerationResult:
             samples=result.get("samples"),
             frames=result.get("frames"),
             audio=result.get("audio"),
+            audio_sample_rate=result.get("audio_sample_rate"),
             size=result.get("size"),
             generation_time=result.get("generation_time"),
             logging_info=result.get("logging_info"),
@@ -80,6 +83,7 @@ class GenerationResult:
             "samples": self.samples,
             "frames": self.frames,
             "audio": self.audio,
+            "audio_sample_rate": self.audio_sample_rate,
             "size": self.size,
             "generation_time": self.generation_time,
             "logging_info": self.logging_info,

fastvideo/api/sampling_param.py

@@ -112,6 +112,33 @@ class SamplingParam:
     ltx2_stg_blocks_video: list[int] = field(default_factory=lambda: [29])
     ltx2_stg_blocks_audio: list[int] = field(default_factory=lambda: [29])
 
+    # Stable Audio (T2A): clip start/end in seconds. Honored by
+    # `StableAudioConditioningStage` + `StableAudioDecodingStage`. Other
+    # families ignore them.
+    audio_start_in_s: float | None = None
+    audio_end_in_s: float | None = None
+
+    # Stable Audio audio-to-audio (variation):
+    #   `init_audio` -- a path or `[B, C, samples]` waveform at the model
+    #                   sample rate; the pipeline encodes it via the VAE
+    #                   and uses it as the starting latent.
+    #   `init_audio_strength` -- 0..1, higher = closer to the reference
+    #                            (matches the convention of Stability's
+    #                            commercial Stable Audio 2.0 UI). 1.0 ~=
+    #                            VAE round-trip, 0.0 ~= plain T2A.
+    #   `init_noise_level` -- legacy raw `sigma_max` override (0.3..500,
+    #                         higher = more freedom). Kept for callers
+    #                         that already use it; prefer `init_audio_strength`.
+    init_audio: Any = None
+    init_audio_strength: float | None = None
+    init_noise_level: float | None = None
+
+    # Stable Audio inpainting (RePaint-style): `inpaint_audio` is the
+    # reference clip, `inpaint_mask` is a [samples] tensor in {0, 1} where
+    # 1 means *keep the reference* and 0 means *regenerate*.
+    inpaint_audio: Any = None
+    inpaint_mask: Any = None
+
     # Continuation state carried across streaming/multi-segment calls.
     continuation_state: ContinuationState | None = None
     # When True, the pipeline returns a ContinuationState on the result so

fastvideo/configs/models/base.py

@@ -48,8 +48,6 @@ class ModelConfig:
         for key, value in source_model_dict.items():
             if key in valid_fields:
                 setattr(arch_config, key, value)
-            else:
-                raise AttributeError(f"{type(arch_config).__name__} has no field '{key}'")
 
         if hasattr(arch_config, "__post_init__"):
             arch_config.__post_init__()

fastvideo/configs/models/dits/__init__.py

@@ -5,11 +5,13 @@ from fastvideo.configs.models.dits.hunyuanvideo import HunyuanVideoConfig
 from fastvideo.configs.models.dits.hunyuanvideo15 import HunyuanVideo15Config
 from fastvideo.configs.models.dits.longcat import LongCatVideoConfig
 from fastvideo.configs.models.dits.ltx2 import LTX2VideoConfig
+from fastvideo.configs.models.dits.stable_audio import StableAudioConfig
 from fastvideo.configs.models.dits.wanvideo import WanVideoConfig
 from fastvideo.configs.models.dits.hyworld import HYWorldConfig
 from fastvideo.configs.models.dits.kandinsky5 import Kandinsky5VideoConfig
 
 __all__ = [
     "HunyuanVideoConfig", "HunyuanVideo15Config", "HunyuanGameCraftConfig", "WanVideoConfig", "CosmosVideoConfig",
-    "Cosmos25VideoConfig", "LongCatVideoConfig", "LTX2VideoConfig", "HYWorldConfig", "Kandinsky5VideoConfig"
+    "Cosmos25VideoConfig", "LongCatVideoConfig", "LTX2VideoConfig", "HYWorldConfig", "Kandinsky5VideoConfig",
+    "StableAudioConfig"
 ]

fastvideo/configs/models/dits/cosmos.py

@@ -50,7 +50,8 @@ class CosmosArchConfig(DiTArchConfig):
         })
 
     # Cosmos-specific config parameters based on transformer_cosmos.py
-    in_channels: int = 16
+    # in_channels includes the condition_mask channel (16 latent + 1 cond = 17)
+    in_channels: int = 17
     out_channels: int = 16
     num_attention_heads: int = 16
     attention_head_dim: int = 128

fastvideo/configs/models/dits/stable_audio.py

@@ -0,0 +1,76 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Config for the Stable Audio Open 1.0 DiT.
+
+Note: the SA pipeline bypasses the standard `ComposedPipelineBase`
+component loader because the published HF repo ships a single monolithic
+`model.safetensors` (no Diffusers-style `model_index.json` or
+per-subfolder layout). The arch fields and `param_names_mapping` here
+document the architecture and key remap so the same conventions used by
+the rest of the DiT family apply (FSDP shard conditions, supported
+attention backends, future loader integrations) — they are not currently
+consumed by `fastvideo/models/loader/fsdp_load.py` for SA.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig
+from fastvideo.platforms import AttentionBackendEnum
+
+
+def _is_transformer_layer(n: str, m) -> bool:
+    # Matches `transformer.layers.{i}` in the SA DiT module tree.
+    parts = n.split(".")
+    return (len(parts) >= 3 and parts[-3] == "transformer" and parts[-2] == "layers" and parts[-1].isdigit())
+
+
+@dataclass
+class StableAudioArchConfig(DiTArchConfig):
+    _fsdp_shard_conditions: list = field(default_factory=lambda: [_is_transformer_layer])
+
+    # SA's checkpoint is `stable_audio_tools` raw format (not Diffusers),
+    # so the only remaps are: strip the `model.model.` host-pipeline
+    # prefix, and rename `nn.LayerNorm`'s `gamma`/`beta` to torch's
+    # canonical `weight`/`bias`. Linear / cross-attention naming already
+    # matches FastVideo's conventions, so no further remap is needed.
+    param_names_mapping: dict = field(
+        default_factory=lambda: {
+            r"^model\.model\.(.*?)\.gamma$": r"\1.weight",
+            r"^model\.model\.(.*?)\.beta$": r"\1.bias",
+            r"^model\.model\.(.*)$": r"\1",
+        })
+
+    # SA only supports backends compatible with single-GPU LocalAttention.
+    _supported_attention_backends: tuple[AttentionBackendEnum, ...] = (
+        AttentionBackendEnum.FLASH_ATTN,
+        AttentionBackendEnum.TORCH_SDPA,
+    )
+
+    # Architecture constants (from the published `model_config.json` for
+    # `stabilityai/stable-audio-open-1.0`).
+    io_channels: int = 64
+    embed_dim: int = 1536
+    depth: int = 24
+    num_attention_heads: int = 24
+    cond_token_dim: int = 768
+    global_cond_dim: int = 1536
+    project_cond_tokens: bool = False
+    project_global_cond: bool = True
+    # Set to "ln" to wrap attention Q/K in LayerNorm (used by
+    # `stable-audio-open-small`; absent in the 1.0 base).
+    qk_norm: str | None = None
+
+    def __post_init__(self) -> None:
+        super().__post_init__()
+        self.hidden_size = self.embed_dim
+        self.in_channels = self.io_channels
+        self.out_channels = self.io_channels
+        self.num_channels_latents = self.io_channels
+        self.attention_head_dim = self.embed_dim // self.num_attention_heads
+
+
+@dataclass
+class StableAudioConfig(DiTConfig):
+    arch_config: DiTArchConfig = field(default_factory=StableAudioArchConfig)
+
+    prefix: str = "StableAudio"

fastvideo/configs/models/encoders/__init__.py

@@ -7,9 +7,12 @@ from fastvideo.configs.models.encoders.qwen2_5 import Qwen2_5_VLConfig
 from fastvideo.configs.models.encoders.siglip import SiglipVisionConfig
 from fastvideo.configs.models.encoders.reason1 import Reason1ArchConfig, Reason1Config
 from fastvideo.configs.models.encoders.gemma import LTX2GemmaConfig
+from fastvideo.configs.models.encoders.stable_audio_conditioner import (StableAudioConditionerArchConfig,
+                                                                        StableAudioConditionerConfig)
 
 __all__ = [
     "EncoderConfig", "TextEncoderConfig", "ImageEncoderConfig", "BaseEncoderOutput", "CLIPTextConfig",
     "CLIPVisionConfig", "WAN2_1ControlCLIPVisionConfig", "LlamaConfig", "T5Config", "T5LargeConfig", "Qwen2_5_VLConfig",
-    "Reason1ArchConfig", "Reason1Config", "LTX2GemmaConfig", "SiglipVisionConfig"
+    "Reason1ArchConfig", "Reason1Config", "LTX2GemmaConfig", "SiglipVisionConfig", "StableAudioConditionerArchConfig",
+    "StableAudioConditionerConfig"
 ]

fastvideo/configs/models/encoders/stable_audio_conditioner.py

@@ -0,0 +1,80 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Config for the Stable Audio Open 1.0 multi-conditioner.
+
+The conditioner bundles three sub-conditioners — a T5 text encoder
+(prompt) and two NumberConditioners (`seconds_start` / `seconds_total`)
+— into the (cross_attn_cond, cross_attn_mask, global_embed) triple the
+DiT consumes. The architecture is fully specified by the official
+`stable_audio_tools` `MultiConditioner` config; the constants here
+mirror that.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+from fastvideo.configs.models.base import ArchConfig
+from fastvideo.configs.models.encoders.base import (EncoderArchConfig, EncoderConfig)
+
+
+def _default_configs() -> list[dict]:
+    """Default = `stable-audio-open-1.0`'s three sub-conditioners."""
+    return [
+        {
+            "id": "prompt",
+            "type": "t5",
+            "config": {
+                "t5_model_name": "t5-base",
+                "max_length": 128
+            }
+        },
+        {
+            "id": "seconds_start",
+            "type": "number",
+            "config": {
+                "min_val": 0,
+                "max_val": 512
+            }
+        },
+        {
+            "id": "seconds_total",
+            "type": "number",
+            "config": {
+                "min_val": 0,
+                "max_val": 512
+            }
+        },
+    ]
+
+
+@dataclass
+class StableAudioConditionerArchConfig(EncoderArchConfig):
+    architectures: list[str] = field(default_factory=lambda: ["StableAudioMultiConditioner"])
+
+    # Shared embedding width across all sub-conditioners (T5 last-hidden
+    # dim and NumberEmbedder feature dim both = `cond_dim`).
+    cond_dim: int = 768
+
+    # Sub-conditioner identifiers. Order in `cross_attention_cond_ids`
+    # is the concat order for the cross-attn token sequence; order in
+    # `global_cond_ids` is the concat order for the global FiLM-style
+    # embedding.
+    cross_attention_cond_ids: tuple[str, ...] = ("prompt", "seconds_start", "seconds_total")
+    global_cond_ids: tuple[str, ...] = ("seconds_start", "seconds_total")
+
+    # Per-sub-conditioner spec list (mirrors upstream
+    # `model_config.json.model.conditioning.configs`). Each entry is
+    # `{"id": ..., "type": "t5"|"number", "config": {...}}`. The default
+    # matches `stable-audio-open-1.0`; SA-small overrides via the
+    # `conditioner/config.json` shipped in the converted repo.
+    configs: list = field(default_factory=_default_configs)
+
+    # Match official `stable_audio_tools/models/conditioners.py:334`:
+    # T5 is loaded directly in fp16.
+    t5_dtype: str = "float16"
+
+
+@dataclass
+class StableAudioConditionerConfig(EncoderConfig):
+    arch_config: ArchConfig = field(default_factory=StableAudioConditionerArchConfig)
+
+    prefix: str = "stable_audio_conditioner"

fastvideo/configs/models/encoders/t5.py

@@ -41,6 +41,14 @@ class T5ArchConfig(TextEncoderArchConfig):
     text_len: int = 512
     dtype: str | None = None
     gradient_checkpointing: bool = False
+    # Extra fields present in upstream HF T5Config but unused by FastVideo's
+    # encoder. Declared here so `update_model_arch` doesn't reject them when
+    # loading repos like `stabilityai/stable-audio-open-1.0` that ship the
+    # full HF config.
+    n_positions: int = 512
+    decoder_start_token_id: int = 0
+    output_past: bool = True
+    task_specific_params: dict | None = None
     stacked_params_mapping: list[tuple[str, str, str]] = field(default_factory=lambda: [
         # (param_name, shard_name, shard_id)
         (".qkv_proj", ".q", "q"),

fastvideo/configs/models/vaes/__init__.py

@@ -5,6 +5,7 @@ from fastvideo.configs.models.vaes.gen3cvae import Gen3CVAEConfig
 from fastvideo.configs.models.vaes.hunyuanvae import HunyuanVAEConfig
 from fastvideo.configs.models.vaes.hunyuan15vae import Hunyuan15VAEConfig
 from fastvideo.configs.models.vaes.ltx2vae import LTX2VAEConfig
+from fastvideo.configs.models.vaes.oobleck import OobleckVAEArchConfig, OobleckVAEConfig
 from fastvideo.configs.models.vaes.wanvae import WanVAEConfig
 
 __all__ = [
@@ -16,4 +17,6 @@ __all__ = [
     "Gen3CVAEConfig",
     "Hunyuan15VAEConfig",
     "LTX2VAEConfig",
+    "OobleckVAEArchConfig",
+    "OobleckVAEConfig",
 ]

fastvideo/configs/models/vaes/oobleck.py

@@ -0,0 +1,68 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Config for the Stable Audio Open 1.0 "Oobleck" VAE.
+
+Mirrors the per-channel `vae/config.json` shipped in
+`stabilityai/stable-audio-open-1.0` 1:1 (see
+`fastvideo/models/vaes/oobleck.py::OobleckVAE.from_pretrained`, which
+constructs the VAE from these fields). Inherits the FastVideo VAEConfig
+base so the standard `load_encoder` / `load_decoder` flags + tiling
+knobs apply.
+
+Naming: the VAE architecture is officially "Oobleck" (per Stability
+AI's stable-audio-tools) — the surrounding model family is "Stable
+Audio Open 1.0". This config is named after the architecture
+(`OobleckVAEConfig`) since the same VAE is shared across Stable Audio
+checkpoints; downstream pipelines reference it by its arch name, not
+by a host-pipeline name.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+from fastvideo.configs.models.vaes.base import VAEArchConfig, VAEConfig
+
+
+@dataclass
+class OobleckVAEArchConfig(VAEArchConfig):
+    """Stable Audio Open 1.0 VAE architecture constants."""
+
+    architectures: list[str] = field(default_factory=lambda: ["AutoencoderOobleck"])
+
+    # From stabilityai/stable-audio-open-1.0/vae/config.json.
+    encoder_hidden_size: int = 128
+    downsampling_ratios: list[int] = field(default_factory=lambda: [2, 4, 4, 8, 8])
+    channel_multiples: list[int] = field(default_factory=lambda: [1, 2, 4, 8, 16])
+    decoder_channels: int = 128
+    decoder_input_channels: int = 64
+    audio_channels: int = 2  # stereo
+    sampling_rate: int = 44100
+
+
+@dataclass
+class OobleckVAEConfig(VAEConfig):
+    """FastVideo VAE config wrapping the Oobleck arch.
+
+    Audio VAEs don't use the temporal/spatial tiling defaults that the
+    base VAEConfig is shaped for (those exist for video VAEs); they are
+    retained but irrelevant for audio.
+    """
+
+    arch_config: VAEArchConfig = field(default_factory=OobleckVAEArchConfig)
+
+    # Audio is 1-D, so the video-VAE tiling defaults are inert. Disable
+    # them so callers don't accidentally trip on tile-stride math built
+    # for spatial tensors.
+    use_tiling: bool = False
+    use_temporal_tiling: bool = False
+    use_parallel_tiling: bool = False
+
+    # Where the FastVideo loader / pipeline-glue wrapper should fetch
+    # weights from when no local path is supplied. Gated repo — caller's
+    # HF token must have accepted terms on
+    # https://huggingface.co/stabilityai/stable-audio-open-1.0.
+    pretrained_path: str = "stabilityai/stable-audio-open-1.0"
+    pretrained_subfolder: str = "vae"
+    # Match official `stable_audio_tools`: VAE runs in fp16 (the
+    # `pretransform.model_half` path in
+    # `stable_audio_tools/models/pretransforms.py`).
+    pretrained_dtype: str = "float16"

fastvideo/configs/pipelines/stable_audio.py

@@ -0,0 +1,67 @@
+# SPDX-License-Identifier: Apache-2.0
+"""`PipelineConfig` for Stable Audio Open 1.0."""
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+from fastvideo.configs.models import DiTConfig, VAEConfig
+from fastvideo.configs.models.dits import StableAudioConfig
+from fastvideo.configs.models.vaes import OobleckVAEConfig
+from fastvideo.configs.pipelines.base import PipelineConfig
+
+
+@dataclass
+class StableAudioT2AConfig(PipelineConfig):
+    """Stable Audio Open 1.0 pipeline config."""
+
+    dit_config: DiTConfig = field(default_factory=StableAudioConfig)
+    # Standard `TransformerLoader` reads `dit_precision`; default in
+    # `PipelineConfig` is bf16, but we want fp16 to match official.
+    dit_precision: str = "fp16"
+
+    vae_config: VAEConfig = field(default_factory=OobleckVAEConfig)
+    vae_tiling: bool = False
+    vae_sp: bool = False
+
+    # `StableAudioMultiConditioner` owns its own T5; zero out the
+    # parent's text-encoder slots so the length-equality validator passes.
+    text_encoder_configs: tuple = field(default_factory=tuple)
+    preprocess_text_funcs: tuple = field(default_factory=tuple)
+    postprocess_text_funcs: tuple = field(default_factory=tuple)
+
+    num_inference_steps: int = 100
+    guidance_scale: float = 7.0
+    audio_end_in_s: float = 10.0  # short-clip default
+    audio_start_in_s: float = 0.0
+    sampling_rate: int = 44100
+    audio_channels: int = 2
+    # Stable Audio Open 1.0 was trained at a fixed 2,097,152-sample
+    # window (= 2097152 / 44100 ≈ 47.55s). Anything past this is
+    # silently truncated by the post-decode slice — validate up-front.
+    sample_size: int = 2097152
+    max_audio_duration_s: float = 2097152 / 44100
+
+    # Match the official `stable_audio_tools` defaults (`model_half=True`
+    # in `run_gradio.py`), which loads the DiT, VAE, and T5 in fp16 and
+    # wraps T5 forward in `autocast(fp16)`. fp16 is also a hard
+    # requirement for FlashAttention-2 / FA-3.
+    precision: str = "fp16"
+    vae_precision: str = "fp16"
+    text_encoder_precisions: tuple[str, ...] = field(default_factory=tuple)
+
+    def __post_init__(self) -> None:
+        # A2A needs encode; load both halves for either path.
+        self.vae_config.load_encoder = True
+        self.vae_config.load_decoder = True
+
+
+@dataclass
+class StableAudioOpenSmallConfig(StableAudioT2AConfig):
+    """`stable-audio-open-small` overrides: shorter training window
+    (524288 samples ≈ 11.89s @ 44.1 kHz) and a faster default sampler
+    config carried by the small preset.
+    """
+
+    sample_size: int = 524288
+    max_audio_duration_s: float = 524288 / 44100
+    audio_end_in_s: float = 6.0  # short-clip default suitable for the small window

fastvideo/entrypoints/video_generator.py

@@ -627,18 +627,32 @@ class VideoGenerator:
         gen_time = time.perf_counter() - start_time
         logger.info("Generated successfully in %.2f seconds", gen_time)
 
-        # Process outputs
-        videos = rearrange(samples, "b c t h w -> t b c h w")
-        frames = []
-        for x in videos:
-            x = torchvision.utils.make_grid(x, nrow=6)
-            x = x.permute(1, 2, 0).squeeze(-1)
-            x = (x * 255).to(torch.uint8)
-            frames.append(x.cpu().numpy())
+        # Process outputs (skip the make_grid loop for audio-only, where
+        # `samples` is a 1×3×1×8×8 placeholder no caller will use).
+        audio_only = bool(output_batch.extra.get("audio_only"))
+        frames: list[np.ndarray] = []
+        if not audio_only:
+            videos = rearrange(samples, "b c t h w -> t b c h w")
+            for x in videos:
+                x = torchvision.utils.make_grid(x, nrow=6)
+                x = x.permute(1, 2, 0).squeeze(-1)
+                x = (x * 255).to(torch.uint8)
+                frames.append(x.cpu().numpy())
 
         # Save output if requested
         if batch.save_video:
-            if self._is_image_workload():
+            if output_batch.extra.get("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")
+                self._write_pcm_wav(
+                    output_path,
+                    output_batch.extra["audio"],
+                    int(output_batch.extra["audio_sample_rate"]),
+                )
+                logger.info("Saved audio to %s", output_path)
+            elif self._is_image_workload():
                 # Image workloads (t2i, i2i, …): save the first frame as PNG.
                 imageio.imwrite(output_path, frames[0])
                 logger.info("Saved image to %s", output_path)
@@ -655,7 +669,11 @@ class VideoGenerator:
             "prompts": prompt,
             "samples": samples if batch.return_frames else None,
             "frames": frames if batch.return_frames else None,
-            "audio": output_batch.extra.get("audio") 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"),
             "size": (target_height, target_width, batch.num_frames),
             "generation_time": gen_time,
             "logging_info": logging_info,
@@ -683,7 +701,55 @@ class VideoGenerator:
         return result.to_legacy_dict()
 
     @staticmethod
+    def _rewrite_extension(path: str, new_ext: str) -> str:
+        root, old_ext = os.path.splitext(path)
+        new_path = root + new_ext
+        if old_ext and old_ext.lower() != new_ext.lower():
+            logger.info("Rewriting output extension %s -> %s.", old_ext, new_ext)
+        return new_path
+
+    @staticmethod
+    def _audio_to_int16(audio: torch.Tensor | np.ndarray, ) -> tuple[np.ndarray, int]:
+        """Normalize `[samples]` / `[samples, channels]` / `[channels,
+        samples]` audio in roughly [-1, 1] to a `(int16 [samples,
+        channels], num_channels)` pair. Raises `ValueError` for shapes
+        we can't classify.
+        """
+        if torch.is_tensor(audio):
+            audio_np = audio.detach().cpu().float().numpy()
+        else:
+            audio_np = np.asarray(audio, dtype=np.float32)
+        if audio_np.ndim == 1:
+            audio_np = audio_np[:, None]
+        elif audio_np.ndim == 2:
+            if audio_np.shape[0] <= 8 and audio_np.shape[1] > audio_np.shape[0]:
+                audio_np = audio_np.T
+        else:
+            raise ValueError(f"Unexpected audio shape {audio_np.shape}.")
+        audio_np = np.clip(audio_np, -1.0, 1.0)
+        audio_int16 = (audio_np * 32767.0).astype(np.int16)
+        return audio_int16, audio_int16.shape[1]
+
+    @classmethod
+    def _write_pcm_wav(
+        cls,
+        wav_path: str,
+        audio: torch.Tensor | np.ndarray,
+        sample_rate: int,
+    ) -> int:
+        """Write 16-bit PCM WAV; returns the channel count."""
+        import wave
+        audio_int16, num_channels = cls._audio_to_int16(audio)
+        with wave.open(wav_path, "wb") as f:
+            f.setnchannels(num_channels)
+            f.setsampwidth(2)
+            f.setframerate(sample_rate)
+            f.writeframes(audio_int16.tobytes())
+        return num_channels
+
+    @classmethod
     def _mux_audio(
+        cls,
         video_path: str,
         audio: torch.Tensor | np.ndarray,
         sample_rate: int,
@@ -696,37 +762,13 @@ class VideoGenerator:
                            "Install with: pip install av")
             return False
 
-        if torch.is_tensor(audio):
-            audio_np = audio.detach().cpu().float().numpy()
-        else:
-            audio_np = np.asarray(audio, dtype=np.float32)
-
-        if audio_np.ndim == 1:
-            audio_np = audio_np[:, None]
-        elif audio_np.ndim == 2:
-            if audio_np.shape[0] <= 8 and audio_np.shape[1] > audio_np.shape[0]:
-                audio_np = audio_np.T
-        else:
-            logger.warning("Unexpected audio shape %s; skipping mux.", audio_np.shape)
-            return False
-
-        audio_np = np.clip(audio_np, -1.0, 1.0)
-        audio_int16 = (audio_np * 32767.0).astype(np.int16)
-        num_channels = audio_int16.shape[1]
-        layout = "stereo" if num_channels == 2 else "mono"
-
         try:
-            import wave
             with tempfile.TemporaryDirectory() as tmpdir:
                 out_path = os.path.join(tmpdir, "muxed.mp4")
                 wav_path = os.path.join(tmpdir, "audio.wav")
 
-                # Write audio to WAV file
-                with wave.open(wav_path, "wb") as wav_file:
-                    wav_file.setnchannels(num_channels)
-                    wav_file.setsampwidth(2)
-                    wav_file.setframerate(sample_rate)
-                    wav_file.writeframes(audio_int16.tobytes())
+                num_channels = cls._write_pcm_wav(wav_path, audio, sample_rate)
+                layout = "stereo" if num_channels == 2 else "mono"
 
                 # Open input video and audio
                 input_video = av.open(video_path)

fastvideo/models/dits/cosmos.py

@@ -556,7 +556,10 @@ class CosmosTransformer3DModel(BaseDiT):
         self.extra_pos_embed_type = config.extra_pos_embed_type
 
         # 1. Patch Embedding
-        patch_embed_in_channels = config.in_channels + 1 if config.concat_padding_mask else config.in_channels
+        # config.in_channels already includes the condition_mask channel
+        # (HF config: in_channels=17 = 16 latent + 1 condition_mask).
+        # Only add +1 for the padding_mask when concat_padding_mask=True.
+        patch_embed_in_channels = config.in_channels + (1 if config.concat_padding_mask else 0)
         self.patch_embed = CosmosPatchEmbed(patch_embed_in_channels,
                                             inner_dim,
                                             config.patch_size,
@@ -617,6 +620,28 @@ class CosmosTransformer3DModel(BaseDiT):
 
         batch_size, num_channels, num_frames, height, width = hidden_states.shape
 
+        # Defensive dtype alignment: the Cosmos checkpoint is bf16 but
+        # FSDP-wrapped training copies may report fp32 via
+        # `next(parameters()).dtype`, which disables autocast in the
+        # shared denoising stage and feeds fp32 tensors into bf16
+        # weights. Cast every external input to the patch_embed weight
+        # dtype so the model forward is robust regardless of caller.
+        _target_dtype = self.patch_embed.proj.weight.dtype
+        if hidden_states.dtype != _target_dtype:
+            hidden_states = hidden_states.to(_target_dtype)
+        if condition_mask is not None and condition_mask.dtype != _target_dtype:
+            condition_mask = condition_mask.to(_target_dtype)
+        if padding_mask is not None and padding_mask.dtype != _target_dtype:
+            padding_mask = padding_mask.to(_target_dtype)
+        if isinstance(encoder_hidden_states, torch.Tensor):
+            if encoder_hidden_states.dtype != _target_dtype:
+                encoder_hidden_states = encoder_hidden_states.to(_target_dtype)
+        else:
+            encoder_hidden_states = [
+                t.to(_target_dtype) if t.dtype != _target_dtype else t
+                for t in encoder_hidden_states
+            ]
+
         # 1. Concatenate padding mask if needed & prepare attention mask
         if condition_mask is not None:
             hidden_states = torch.cat([hidden_states, condition_mask], dim=1)
@@ -626,6 +651,10 @@ class CosmosTransformer3DModel(BaseDiT):
             padding_mask = transforms.functional.resize(
                 padding_mask, list(hidden_states.shape[-2:]), interpolation=transforms.InterpolationMode.NEAREST
             )
+            # torchvision.resize may upcast bf16/fp16 → fp32; restore
+            # hidden_states' dtype so the subsequent cat doesn't promote
+            # everything and break patch_embed (bf16 weights).
+            padding_mask = padding_mask.to(hidden_states.dtype)
             hidden_states = torch.cat(
                 [hidden_states, padding_mask.unsqueeze(2).repeat(batch_size, 1, num_frames, 1, 1)], dim=1
             )
@@ -703,6 +732,11 @@ class CosmosTransformer3DModel(BaseDiT):
         hidden_states = hidden_states.permute(0, 7, 1, 6, 2, 4, 3, 5)
         hidden_states = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3)
 
+        # Return as tuple for compatibility with callers that
+        # do `transformer(..., return_dict=False)[0]` (diffusers
+        # convention used by CosmosDenoisingStage).
+        if not kwargs.get("return_dict", True):
+            return (hidden_states,)
         return hidden_states
 
 # Entry point for model registry

fastvideo/models/dits/stable_audio.py

@@ -0,0 +1,389 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Stable Audio Open 1.0 DiT.
+
+Continuous transformer with rotary self-attention, GQA cross-attention,
+and prepend global conditioning. 24 layers, embed_dim=1536, head_dim=64.
+"""
+from __future__ import annotations
+
+import math
+from typing import Any
+
+import torch
+from einops import rearrange
+from torch import nn
+
+from fastvideo.attention import LocalAttention
+from fastvideo.configs.models.dits import StableAudioConfig
+from fastvideo.layers.layernorm import FP32LayerNorm
+from fastvideo.layers.linear import ReplicatedLinear
+from fastvideo.models.dits.base import BaseDiT
+from fastvideo.models.loader.utils import get_param_names_mapping
+
+# Single import-time snapshot — re-reading via `StableAudioConfig()` per
+# `Attention.__init__` would rebuild the nested dataclass + regex map ~48
+# times during a single DiT construction. Reused for the class-level
+# attribute defaults below.
+_DEFAULT_CONFIG = StableAudioConfig()
+_SUPPORTED_BACKENDS = _DEFAULT_CONFIG.arch_config._supported_attention_backends
+
+
+class FourierFeatures(nn.Module):
+    """Random-Fourier learned-frequency timestep encoder."""
+
+    def __init__(self, in_features: int, out_features: int, std: float = 1.0) -> None:
+        super().__init__()
+        assert out_features % 2 == 0
+        self.weight = nn.Parameter(torch.randn([out_features // 2, in_features]) * std)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        f = 2 * math.pi * x @ self.weight.T
+        return torch.cat([f.cos(), f.sin()], dim=-1)
+
+
+# Partial-rotary with halves-swap (`unbind(-2)`, `[-x2, x1]`). Different
+# from FastVideo's `_apply_rotary_emb` (interleaved pairs, `unbind(-1)`),
+# so kept local.
+
+
+class RotaryEmbedding(nn.Module):
+
+    def __init__(self, dim: int, base: float = 10000.0) -> None:
+        super().__init__()
+        inv_freq = 1.0 / (base**(torch.arange(0, dim, 2).float() / dim))
+        self.register_buffer("inv_freq", inv_freq)
+        self.register_buffer("scale", None)
+
+    def forward_from_seq_len(self, seq_len: int):
+        t = torch.arange(seq_len, device=self.inv_freq.device, dtype=torch.float32)
+        freqs = torch.einsum("i , j -> i j", t, self.inv_freq)
+        freqs = torch.cat((freqs, freqs), dim=-1)
+        return freqs, 1.0
+
+
+def _rotate_half(x: torch.Tensor) -> torch.Tensor:
+    x = rearrange(x, "... (j d) -> ... j d", j=2)
+    x1, x2 = x.unbind(dim=-2)
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def _apply_rotary_pos_emb(t: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor:
+    out_dtype = t.dtype
+    rot_dim, seq_len = freqs.shape[-1], t.shape[-2]
+    freqs = freqs.to(torch.float32)[-seq_len:, :]
+    t = t.to(torch.float32)
+    if t.ndim == 4 and freqs.ndim == 3:
+        freqs = rearrange(freqs, "b n d -> b 1 n d")
+    t_rot, t_unrot = t[..., :rot_dim], t[..., rot_dim:]
+    t_rot = (t_rot * freqs.cos()) + (_rotate_half(t_rot) * freqs.sin())
+    return torch.cat((t_rot.to(out_dtype), t_unrot.to(out_dtype)), dim=-1)
+
+
+# SwiGLU FF — local because `fastvideo.layers.mlp.MLP` is non-gated.
+
+
+class _GLU(nn.Module):
+
+    def __init__(self, dim_in: int, dim_out: int, activation: nn.Module) -> None:
+        super().__init__()
+        self.act = activation
+        self.proj = ReplicatedLinear(dim_in, dim_out * 2, bias=True)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x, _ = self.proj(x)
+        x, gate = x.chunk(2, dim=-1)
+        return x * self.act(gate)
+
+
+class FeedForward(nn.Module):
+    # Sequential layout `(GLU, Identity, Linear, Identity)` keeps the
+    # checkpoint keys at indices 0 and 2.
+
+    def __init__(self, dim: int, mult: int = 4, zero_init_output: bool = True) -> None:
+        super().__init__()
+        inner_dim = int(dim * mult)
+        linear_in = _GLU(dim, inner_dim, nn.SiLU())
+        linear_out = ReplicatedLinear(inner_dim, dim, bias=True)
+        if zero_init_output:
+            nn.init.zeros_(linear_out.weight)
+            nn.init.zeros_(linear_out.bias)
+        self.ff = nn.Sequential(linear_in, nn.Identity(), linear_out, nn.Identity())
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        for mod in self.ff:
+            if isinstance(mod, ReplicatedLinear):
+                x, _ = mod(x)
+            else:
+                x = mod(x)
+        return x
+
+
+# Cross-attention is GQA (24 query heads, 12 KV heads); both backends
+# (FlashAttn, SDPA with `enable_gqa=True`) handle it.
+
+
+class Attention(nn.Module):
+
+    def __init__(self, dim: int, dim_heads: int = 64, dim_context: int | None = None,
+                 zero_init_output: bool = True, qk_norm: str | None = None) -> None:
+        super().__init__()
+        self.dim = dim
+        self.dim_heads = dim_heads
+        dim_kv = dim_context if dim_context is not None else dim
+        self.num_heads = dim // dim_heads
+        self.kv_heads = dim_kv // dim_heads
+        if dim_context is not None:
+            self.to_q = ReplicatedLinear(dim, dim, bias=False)
+            self.to_kv = ReplicatedLinear(dim_kv, dim_kv * 2, bias=False)
+        else:
+            self.to_qkv = ReplicatedLinear(dim, dim * 3, bias=False)
+        self.to_out = ReplicatedLinear(dim, dim, bias=False)
+        if zero_init_output:
+            nn.init.zeros_(self.to_out.weight)
+
+        # `stable-audio-open-small` wraps Q/K in LayerNorm before attn
+        # (`attn_kwargs.qk_norm = "ln"` in its `model_config.json`); the
+        # 1.0 base does not. Names match upstream (`q_norm`/`k_norm`)
+        # so the converted state dict loads strict.
+        if qk_norm == "ln":
+            self.q_norm = nn.LayerNorm(dim_heads)
+            self.k_norm = nn.LayerNorm(dim_heads)
+        elif qk_norm is None:
+            self.q_norm = nn.Identity()
+            self.k_norm = nn.Identity()
+        else:
+            raise ValueError(f"Unsupported qk_norm={qk_norm!r}; expected 'ln' or None.")
+
+        self.attn = LocalAttention(num_heads=self.num_heads, head_size=dim_heads,
+                                   num_kv_heads=self.kv_heads, causal=False,
+                                   supported_attention_backends=_SUPPORTED_BACKENDS)
+
+    def forward(self, x: torch.Tensor, context: torch.Tensor | None = None,
+                rotary_pos_emb: tuple[torch.Tensor, float] | None = None) -> torch.Tensor:
+        h, kv_h, has_context = self.num_heads, self.kv_heads, context is not None
+        kv_input = context if has_context else x
+        if has_context:
+            q, _ = self.to_q(x)
+            kv, _ = self.to_kv(kv_input)
+            k, v = kv.chunk(2, dim=-1)
+        else:
+            qkv, _ = self.to_qkv(x)
+            q, k, v = qkv.chunk(3, dim=-1)
+        # LocalAttention expects [batch, seq_len, num_heads, head_dim].
+        q = rearrange(q, "b n (h d) -> b n h d", h=h)
+        k = rearrange(k, "b n (h d) -> b n h d", h=kv_h)
+        v = rearrange(v, "b n (h d) -> b n h d", h=kv_h)
+        q = self.q_norm(q)
+        k = self.k_norm(k)
+
+        if rotary_pos_emb is not None:
+            freqs, _ = rotary_pos_emb
+            v_dtype = v.dtype
+            # Partial rotary (rot_dim < head_dim) with halves-swap, so
+            # apply outside LocalAttention. q,k come in as [B, S, H, D];
+            # transpose to [B, H, S, D] for the helper.
+            q_t = q.transpose(1, 2)
+            k_t = k.transpose(1, 2)
+            if q_t.shape[-2] >= k_t.shape[-2]:
+                ratio = q_t.shape[-2] / k_t.shape[-2]
+                q_freqs, k_freqs = freqs, ratio * freqs
+            else:
+                ratio = k_t.shape[-2] / q_t.shape[-2]
+                q_freqs, k_freqs = ratio * freqs, freqs
+            q = _apply_rotary_pos_emb(q_t, q_freqs).to(v_dtype).transpose(1, 2)
+            k = _apply_rotary_pos_emb(k_t, k_freqs).to(v_dtype).transpose(1, 2)
+
+        out = self.attn(q, k, v)
+        out = rearrange(out, "b n h d -> b n (h d)")
+        out, _ = self.to_out(out)
+        return out
+
+
+class TransformerBlock(nn.Module):
+
+    def __init__(self, dim: int, dim_heads: int = 64, cross_attend: bool = False,
+                 dim_context: int | None = None, zero_init_branch_outputs: bool = True,
+                 qk_norm: str | None = None) -> None:
+        super().__init__()
+        self.dim = dim
+        self.dim_heads = min(dim_heads, dim)
+        self.cross_attend = cross_attend
+        self.pre_norm = FP32LayerNorm(dim, elementwise_affine=True)
+        self.self_attn = Attention(dim, dim_heads=self.dim_heads,
+                                   zero_init_output=zero_init_branch_outputs,
+                                   qk_norm=qk_norm)
+        if cross_attend:
+            self.cross_attend_norm = FP32LayerNorm(dim, elementwise_affine=True)
+            self.cross_attn = Attention(dim, dim_heads=self.dim_heads, dim_context=dim_context,
+                                        zero_init_output=zero_init_branch_outputs,
+                                        qk_norm=qk_norm)
+        self.ff_norm = FP32LayerNorm(dim, elementwise_affine=True)
+        self.ff = FeedForward(dim, zero_init_output=zero_init_branch_outputs)
+
+    def forward(self, x: torch.Tensor, context: torch.Tensor | None = None,
+                rotary_pos_emb: tuple[torch.Tensor, float] | None = None) -> torch.Tensor:
+        x = x + self.self_attn(self.pre_norm(x), rotary_pos_emb=rotary_pos_emb)
+        if context is not None and self.cross_attend:
+            x = x + self.cross_attn(self.cross_attend_norm(x), context=context)
+        x = x + self.ff(self.ff_norm(x))
+        return x
+
+
+class ContinuousTransformer(nn.Module):
+
+    def __init__(self, dim: int, depth: int, *, dim_heads: int = 64, dim_in: int | None = None,
+                 dim_out: int | None = None, cross_attend: bool = False,
+                 cond_token_dim: int | None = None, zero_init_branch_outputs: bool = True,
+                 qk_norm: str | None = None) -> None:
+        super().__init__()
+        self.dim = dim
+        self.depth = depth
+        self.project_in = (ReplicatedLinear(dim_in, dim, bias=False) if dim_in is not None
+                           else nn.Identity())
+        self.project_out = (ReplicatedLinear(dim, dim_out, bias=False) if dim_out is not None
+                            else nn.Identity())
+        self.rotary_pos_emb = RotaryEmbedding(max(dim_heads // 2, 32))
+        self.layers = nn.ModuleList([
+            TransformerBlock(dim, dim_heads=dim_heads, cross_attend=cross_attend,
+                             dim_context=cond_token_dim,
+                             zero_init_branch_outputs=zero_init_branch_outputs,
+                             qk_norm=qk_norm) for _ in range(depth)
+        ])
+
+    def forward(self, x: torch.Tensor, prepend_embeds: torch.Tensor | None = None,
+                context: torch.Tensor | None = None) -> torch.Tensor:
+        if isinstance(self.project_in, ReplicatedLinear):
+            x, _ = self.project_in(x)
+        if prepend_embeds is not None:
+            assert prepend_embeds.shape[-1] == x.shape[-1]
+            x = torch.cat((prepend_embeds, x), dim=-2)
+        rotary = self.rotary_pos_emb.forward_from_seq_len(x.shape[1])
+        for layer in self.layers:
+            x = layer(x, context=context, rotary_pos_emb=rotary)
+        if isinstance(self.project_out, ReplicatedLinear):
+            x, _ = self.project_out(x)
+        return x
+
+
+class StableAudioDiT(BaseDiT):
+    """Stable Audio Open 1.0 diffusion transformer."""
+
+    _fsdp_shard_conditions = _DEFAULT_CONFIG.arch_config._fsdp_shard_conditions
+    _compile_conditions = _DEFAULT_CONFIG.arch_config._compile_conditions
+    param_names_mapping = _DEFAULT_CONFIG.arch_config.param_names_mapping
+    reverse_param_names_mapping: dict = {}
+
+    def __init__(self, config: StableAudioConfig | None = None,
+                 hf_config: dict[str, Any] | None = None) -> None:
+        if config is None:
+            config = StableAudioConfig()
+        super().__init__(config=config, hf_config=hf_config or {})
+        arch = config.arch_config
+        self.hidden_size = arch.hidden_size
+        self.num_attention_heads = arch.num_attention_heads
+        self.num_channels_latents = arch.num_channels_latents
+        io_channels = arch.io_channels
+        embed_dim = arch.embed_dim
+        depth = arch.depth
+        num_heads = arch.num_attention_heads
+        cond_token_dim = arch.cond_token_dim
+        global_cond_dim = arch.global_cond_dim
+        project_cond_tokens = arch.project_cond_tokens
+        project_global_cond = arch.project_global_cond
+        qk_norm = arch.qk_norm
+
+        self.cond_token_dim = cond_token_dim
+        timestep_features_dim = 256
+        self.timestep_features = FourierFeatures(1, timestep_features_dim)
+        self.to_timestep_embed = nn.Sequential(
+            ReplicatedLinear(timestep_features_dim, embed_dim, bias=True),
+            nn.SiLU(),
+            ReplicatedLinear(embed_dim, embed_dim, bias=True),
+        )
+        self.diffusion_objective = "v"
+
+        cond_embed_dim = cond_token_dim if not project_cond_tokens else embed_dim
+        self.to_cond_embed = nn.Sequential(
+            ReplicatedLinear(cond_token_dim, cond_embed_dim, bias=False),
+            nn.SiLU(),
+            ReplicatedLinear(cond_embed_dim, cond_embed_dim, bias=False),
+        )
+
+        global_embed_dim = global_cond_dim if not project_global_cond else embed_dim
+        self.to_global_embed = nn.Sequential(
+            ReplicatedLinear(global_cond_dim, global_embed_dim, bias=False),
+            nn.SiLU(),
+            ReplicatedLinear(global_embed_dim, global_embed_dim, bias=False),
+        )
+
+        self.transformer = ContinuousTransformer(
+            dim=embed_dim, depth=depth, dim_heads=embed_dim // num_heads, dim_in=io_channels,
+            dim_out=io_channels, cross_attend=True, cond_token_dim=cond_embed_dim,
+            qk_norm=qk_norm,
+        )
+
+        self.preprocess_conv = nn.Conv1d(io_channels, io_channels, 1, bias=False)
+        nn.init.zeros_(self.preprocess_conv.weight)
+        self.postprocess_conv = nn.Conv1d(io_channels, io_channels, 1, bias=False)
+        nn.init.zeros_(self.postprocess_conv.weight)
+
+        self.io_channels = io_channels
+        self.embed_dim = embed_dim
+        self.depth = depth
+        self.num_heads = num_heads
+        self.__post_init__()
+
+    @staticmethod
+    def _seq_apply(seq: nn.Sequential, x: torch.Tensor) -> torch.Tensor:
+        for mod in seq:
+            if isinstance(mod, ReplicatedLinear):
+                x, _ = mod(x)
+            else:
+                x = mod(x)
+        return x
+
+    def forward(self, x: torch.Tensor, t: torch.Tensor, *, cross_attn_cond: torch.Tensor,
+                global_embed: torch.Tensor) -> torch.Tensor:
+        """Forward over a single batch. CFG batching is the caller's job."""
+        model_dtype = next(self.parameters()).dtype
+        x = x.to(model_dtype)
+        t = t.to(model_dtype)
+        cross_attn_cond = cross_attn_cond.to(model_dtype)
+        global_embed = global_embed.to(model_dtype)
+
+        cross_attn_cond = self._seq_apply(self.to_cond_embed, cross_attn_cond)
+        global_embed = self._seq_apply(self.to_global_embed, global_embed)
+        timestep_embed = self._seq_apply(self.to_timestep_embed, self.timestep_features(t[:, None]))
+        global_embed = global_embed + timestep_embed
+        prepend_inputs = global_embed.unsqueeze(1)
+
+        x = self.preprocess_conv(x) + x
+        x = rearrange(x, "b c t -> b t c")
+        out = self.transformer(x, prepend_embeds=prepend_inputs, context=cross_attn_cond)
+        out = rearrange(out, "b t c -> b c t")[:, :, prepend_inputs.shape[1]:]
+        return self.postprocess_conv(out) + out
+
+    @classmethod
+    def from_official_state_dict(cls, state_dict: dict[str, torch.Tensor],
+                                 prefix: str = "model.model.") -> "StableAudioDiT":
+        """Load from a raw `stable_audio_tools` monolithic state dict.
+        Kept for tests / older checkpoints; production loads go through
+        the standard `TransformerLoader` against the converted Diffusers
+        repo.
+        """
+        model = cls()
+        mapping_fn = get_param_names_mapping(model.config.arch_config.param_names_mapping)
+        remapped: dict[str, torch.Tensor] = {}
+        for k, v in state_dict.items():
+            if not k.startswith(prefix):
+                continue
+            new_key, _, _ = mapping_fn(k)
+            remapped[new_key] = v
+        missing, unexpected = model.load_state_dict(remapped, strict=True)
+        if missing or unexpected:
+            raise RuntimeError(
+                f"StableAudioDiT load mismatch — missing={missing[:5]} unexpected={unexpected[:5]}")
+        return model
+
+
+EntryClass = StableAudioDiT

fastvideo/models/encoders/stable_audio_conditioner.py

@@ -0,0 +1,214 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Stable Audio Open 1.0 conditioner.
+
+T5-base text encoder + two NumberConditioners (`seconds_start`,
+`seconds_total`), wrapped by `StableAudioMultiConditioner` which
+produces the cross-attention and global-conditioning tensors the DiT
+expects.
+"""
+from __future__ import annotations
+
+import math
+
+import torch
+import torch.nn as nn
+from einops import rearrange
+
+from fastvideo.configs.models.encoders import StableAudioConditionerConfig
+
+
+class _LearnedPositionalEmbedding(nn.Module):
+
+    def __init__(self, dim: int) -> None:
+        super().__init__()
+        assert (dim % 2) == 0
+        self.weights = nn.Parameter(torch.randn(dim // 2))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = rearrange(x, "b -> b 1")
+        freqs = x * rearrange(self.weights, "d -> 1 d") * 2 * math.pi
+        fouriered = torch.cat((freqs.sin(), freqs.cos()), dim=-1)
+        return torch.cat((x, fouriered), dim=-1)
+
+
+def _time_positional_embedding(dim: int, out_features: int) -> nn.Sequential:
+    return nn.Sequential(_LearnedPositionalEmbedding(dim),
+                         nn.Linear(in_features=dim + 1, out_features=out_features))
+
+
+class NumberEmbedder(nn.Module):
+
+    def __init__(self, features: int, dim: int = 256) -> None:
+        super().__init__()
+        self.features = features
+        self.embedding = _time_positional_embedding(dim=dim, out_features=features)
+
+    def forward(self, x: torch.Tensor | list[float]) -> torch.Tensor:
+        if not torch.is_tensor(x):
+            device = next(self.embedding.parameters()).device
+            x = torch.tensor(x, device=device)
+        shape = x.shape
+        x = rearrange(x, "... -> (...)")
+        out = self.embedding(x)
+        return out.view(*shape, self.features)
+
+
+class _Conditioner(nn.Module):
+
+    def __init__(self, dim: int, output_dim: int, project_out: bool = False) -> None:
+        super().__init__()
+        self.dim = dim
+        self.output_dim = output_dim
+        self.proj_out = (nn.Linear(dim, output_dim) if dim != output_dim or project_out
+                         else nn.Identity())
+
+
+class T5Conditioner(_Conditioner):
+    """T5 text conditioner. Pads to `model_max_length` (=128 for the SA
+    repo's tokenizer, NOT the standard 512) and emits a masked
+    last-hidden-state.
+    """
+
+    T5_MODEL_DIMS = {"t5-base": 768}
+
+    def __init__(self, output_dim: int, t5_model_name: str = "t5-base",
+                 max_length: int = 128, dtype: str = "float16") -> None:
+        super().__init__(self.T5_MODEL_DIMS[t5_model_name], output_dim, project_out=False)
+        from transformers import AutoTokenizer, T5EncoderModel
+        self.max_length = max_length
+        self.tokenizer = AutoTokenizer.from_pretrained(t5_model_name)
+        # T5 loaded directly in fp16 (config-driven) to match official
+        # `stable_audio_tools/models/conditioners.py:334`. Registered as
+        # a normal submodule so `.to(device)` / `torch.compile` track it;
+        # `from_official_state_dict` filters `conditioners.prompt.*` from
+        # the missing-key check (T5 weights are absent from the SA
+        # checkpoint by design).
+        # Explicit lookup so a typo (e.g. "fp16" instead of "float16") errors
+        # at load time rather than silently falling back to a wrong dtype.
+        torch_dtype = getattr(torch, dtype)
+        if not isinstance(torch_dtype, torch.dtype):
+            raise ValueError(f"T5Conditioner dtype={dtype!r} is not a torch.dtype.")
+        self._t5_dtype = torch_dtype
+        self.model = (T5EncoderModel.from_pretrained(t5_model_name).eval().requires_grad_(False).to(torch_dtype))
+
+    def forward(self, texts: list[str], device: torch.device | str) -> tuple[torch.Tensor, torch.Tensor]:
+        encoded = self.tokenizer(texts, truncation=True, max_length=self.max_length,
+                                 padding="max_length", return_tensors="pt")
+        input_ids = encoded["input_ids"].to(device)
+        attention_mask = encoded["attention_mask"].to(device).to(torch.bool)
+        # Mirror official's `autocast(fp16)` wrap on T5 forward.
+        with torch.no_grad(), torch.autocast(device_type="cuda", dtype=self._t5_dtype):
+            embeddings = self.model(input_ids=input_ids,
+                                    attention_mask=attention_mask)["last_hidden_state"]
+        embeddings = self.proj_out(embeddings) * attention_mask.unsqueeze(-1).float()
+        return embeddings, attention_mask
+
+
+class NumberConditioner(_Conditioner):
+    """Float-valued conditioner with min/max clamping + NumberEmbedder."""
+
+    def __init__(self, output_dim: int, min_val: float = 0, max_val: float = 1) -> None:
+        super().__init__(output_dim, output_dim)
+        self.min_val = min_val
+        self.max_val = max_val
+        self.embedder = NumberEmbedder(features=output_dim)
+
+    def forward(self, floats: list[float], device: torch.device | str) -> tuple[torch.Tensor, torch.Tensor]:
+        floats = [float(x) for x in floats]
+        floats_t = torch.tensor(floats, device=device).clamp(self.min_val, self.max_val)
+        normalized = (floats_t - self.min_val) / (self.max_val - self.min_val)
+        emb_dtype = next(self.embedder.parameters()).dtype
+        normalized = normalized.to(emb_dtype)
+        float_embeds = self.embedder(normalized).unsqueeze(1)
+        return float_embeds, torch.ones(float_embeds.shape[0], 1, device=device)
+
+
+class StableAudioMultiConditioner(nn.Module):
+    """SA-Open-1.0 conditioner: T5 prompt + duration NumberConditioners.
+
+    All hardcoded constants (cond_dim, sub-conditioner ids, T5 model
+    name + max_length, NumberConditioner ranges) live on
+    `StableAudioConditionerConfig` — see
+    `fastvideo/configs/models/encoders/stable_audio_conditioner.py`.
+    """
+
+    def __init__(self, config: StableAudioConditionerConfig | None = None) -> None:
+        super().__init__()
+        self.config = config or StableAudioConditionerConfig()
+        arch = self.config.arch_config
+        # Build sub-conditioners from the `configs` list (mirrors
+        # upstream's `MultiConditioner` factory).
+        sub: dict[str, nn.Module] = {}
+        for spec in arch.configs:
+            sid = spec["id"]
+            stype = spec["type"]
+            scfg = spec["config"]
+            if stype == "t5":
+                sub[sid] = T5Conditioner(output_dim=arch.cond_dim,
+                                         t5_model_name=scfg["t5_model_name"],
+                                         max_length=scfg["max_length"],
+                                         dtype=arch.t5_dtype)
+            elif stype == "number":
+                sub[sid] = NumberConditioner(output_dim=arch.cond_dim,
+                                             min_val=scfg["min_val"], max_val=scfg["max_val"])
+            else:
+                raise ValueError(f"Unknown sub-conditioner type {stype!r} for id {sid!r}.")
+        self.conditioners = nn.ModuleDict(sub)
+        self.cross_attention_cond_ids = tuple(arch.cross_attention_cond_ids)
+        self.global_cond_ids = tuple(arch.global_cond_ids)
+
+    def forward(self, batch_metadata: list[dict],
+                device: torch.device | str) -> dict[str, tuple[torch.Tensor, torch.Tensor]]:
+        out: dict[str, tuple[torch.Tensor, torch.Tensor]] = {}
+        for key, conditioner in self.conditioners.items():
+            inputs = [x[key] for x in batch_metadata]
+            out[key] = conditioner(inputs, device)
+        return out
+
+    def get_conditioning_inputs(
+        self, cond: dict[str, tuple[torch.Tensor, torch.Tensor]]
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Pack conditioner outputs into the (cross_attn_cond,
+        cross_attn_mask, global_embed) triple the DiT consumes. Order
+        is driven by `cross_attention_cond_ids` / `global_cond_ids`
+        from the config — SA-1.0 uses three sub-conditioners
+        (prompt + seconds_start + seconds_total); SA-small uses two
+        (prompt + seconds_total).
+        """
+        x_embs = [cond[i][0] for i in self.cross_attention_cond_ids]
+        x_masks = [cond[i][1] for i in self.cross_attention_cond_ids]
+        cross_attn_cond = torch.cat(x_embs, dim=1)
+        cross_attn_mask = torch.cat(x_masks, dim=1)
+        global_embed = torch.cat([cond[i][0][:, 0] for i in self.global_cond_ids], dim=-1)
+        return cross_attn_cond, cross_attn_mask, global_embed
+
+    @classmethod
+    def from_official_state_dict(cls, state_dict: dict[str, torch.Tensor],
+                                 prefix: str = "conditioner.") -> "StableAudioMultiConditioner":
+        """Load NumberConditioner weights from a raw `stable_audio_tools`
+        monolithic state dict. Kept for tests / older checkpoints;
+        production loads go through the standard `ConditionerLoader`
+        against the converted Diffusers repo.
+        """
+        mc = cls()
+        own_state = mc.state_dict()
+        loaded: dict[str, torch.Tensor] = {}
+        for k, v in state_dict.items():
+            if not k.startswith(prefix):
+                continue
+            stripped = k[len(prefix):]
+            if stripped in own_state:
+                loaded[stripped] = v
+        # T5 keys are intentionally absent from the checkpoint.
+        missing = [k for k in own_state.keys() if k not in loaded
+                   and not k.startswith("conditioners.prompt.")]
+        unexpected = [k for k in loaded.keys() if k not in own_state]
+        if missing or unexpected:
+            raise RuntimeError(
+                f"StableAudioMultiConditioner load mismatch — missing={missing[:5]} unexpected={unexpected[:5]}"
+            )
+        mc.load_state_dict(loaded, strict=False)
+        return mc
+
+
+EntryClass = StableAudioMultiConditioner

fastvideo/models/layerwise_offload.py

@@ -1,224 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# Inspired by SGLang's layerwise offload implementation:
-# https://github.com/sgl-project/sglang/pull/15511
-#
-# This implementation provides a lightweight layerwise CPU offload manager
-# with async H2D prefetch using a dedicated CUDA stream, following SGLang's design.
-
-import re
-from contextlib import contextmanager
-from typing import Dict, Set, Optional, Tuple
-
-import torch
-
-
-class LayerwiseOffloadManager:
-    """A lightweight layerwise CPU offload manager.
-
-    Offloads per-layer parameters/buffers from GPU to CPU, and supports async H2D
-    prefetch using a dedicated CUDA stream.
-    """
-
-    def __init__(
-        self,
-        model: torch.nn.Module,
-        *,
-        module_list_attr: str,
-        num_layers: int,
-        enabled: bool,
-        pin_cpu_memory: bool = True,
-        auto_initialize: bool = False,
-    ) -> None:
-        self.model = model
-        self.module_list_attr = module_list_attr
-        self.num_layers = int(num_layers)
-        self.pin_cpu_memory = bool(pin_cpu_memory)
-
-        self.enabled = bool(enabled and torch.cuda.is_available())
-        self.device = (
-            torch.device("cuda", torch.cuda.current_device()) if self.enabled else None
-        )
-        self.copy_stream = torch.cuda.Stream() if self.enabled else None
-
-        self._layer_name_re = re.compile(
-            rf"(^|\.){re.escape(module_list_attr)}\.(\d+)(\.|$)"
-        )
-
-        self._cpu_weights: Dict[int, Dict[str, torch.Tensor]] = {}
-        self._cpu_dtypes: Dict[int, Dict[str, torch.dtype]] = {}
-
-        self._gpu_layers: Dict[int, Set[str]] = {}
-
-        self._named_parameters: Dict[str, torch.nn.Parameter] = {}
-        self._named_buffers: Dict[str, torch.Tensor] = {}
-
-        self._meta: Dict[str, Tuple[int, torch.dtype]] = {}
-
-        if auto_initialize:
-            self.initialize()
-
-    def _match_layer_idx(self, name: str) -> Optional[int]:
-        m = self._layer_name_re.search(name)
-        if not m:
-            return None
-        try:
-            return int(m.group(2))
-        except Exception:
-            return None
-
-    def _record_meta(self, name: str, t: torch.Tensor) -> None:
-        if name not in self._meta:
-            self._meta[name] = (int(t.ndim), t.dtype)
-
-    def _make_placeholder(self, name: str) -> torch.Tensor:
-        """Rank-preserving empty placeholder on GPU."""
-        assert self.device is not None
-        ndim, dtype = self._meta[name]
-        shape = (0,) if ndim <= 0 else (0,) * ndim
-        return torch.empty(shape, device=self.device, dtype=dtype)
-
-    def _get_target(self, name: str) -> torch.Tensor:
-        if name in self._named_parameters:
-            return self._named_parameters[name]
-        return self._named_buffers[name]
-
-    def _offload_tensor(self, name: str, tensor: torch.Tensor, layer_idx: int) -> None:
-        if layer_idx not in self._cpu_weights:
-            self._cpu_weights[layer_idx] = {}
-            self._cpu_dtypes[layer_idx] = {}
-
-        self._record_meta(name, tensor)
-
-        cpu_weight = tensor.detach().to("cpu")
-        if self.pin_cpu_memory:
-            cpu_weight = cpu_weight.pin_memory()
-
-        self._cpu_weights[layer_idx][name] = cpu_weight
-        self._cpu_dtypes[layer_idx][name] = tensor.dtype
-
-        if self.device is not None:
-            tensor.data = self._make_placeholder(name)
-
-    @torch.compiler.disable
-    def initialize(self) -> None:
-        """Offload all matched layer tensors to CPU and prefetch layer 0 (sync)."""
-        if not self.enabled:
-            return
-
-        self._named_parameters = dict(self.model.named_parameters())
-        self._named_buffers = dict(self.model.named_buffers())
-
-        for name, param in self._named_parameters.items():
-            layer_idx = self._match_layer_idx(name)
-            if layer_idx is None or layer_idx >= self.num_layers:
-                continue
-            self._offload_tensor(name, param, layer_idx)
-
-        for name, buf in self._named_buffers.items():
-            layer_idx = self._match_layer_idx(name)
-            if layer_idx is None or layer_idx >= self.num_layers:
-                continue
-            self._offload_tensor(name, buf, layer_idx)
-
-        self.prefetch_layer(0, non_blocking=False)
-        if self.copy_stream is not None:
-            torch.cuda.current_stream().wait_stream(self.copy_stream)
-
-    @torch.compiler.disable
-    def prefetch_layer(self, layer_idx: int, non_blocking: bool = True) -> None:
-        """Prefetch a layer's tensors from CPU to GPU (async on copy_stream)."""
-        if not self.enabled or self.device is None or self.copy_stream is None:
-            return
-        if layer_idx < 0 or layer_idx >= self.num_layers:
-            return
-        if layer_idx in self._gpu_layers:
-            return
-        if layer_idx not in self._cpu_weights:
-            return
-
-        self.copy_stream.wait_stream(torch.cuda.current_stream())
-
-        param_names: Set[str] = set()
-        with torch.cuda.stream(self.copy_stream):
-            for name, cpu_weight in self._cpu_weights[layer_idx].items():
-                target = self._get_target(name)
-
-                gpu_weight = torch.empty(
-                    cpu_weight.shape,
-                    dtype=self._cpu_dtypes[layer_idx][name],
-                    device=self.device,
-                )
-                gpu_weight.copy_(cpu_weight, non_blocking=non_blocking)
-
-                target.data = gpu_weight
-                param_names.add(name)
-
-        self._gpu_layers[layer_idx] = param_names
-
-    @contextmanager
-    def layer_scope(
-        self,
-        *,
-        prefetch_layer_idx: Optional[int],
-        release_layer_idx: Optional[int],
-        non_blocking: bool = True,
-    ):
-        if self.enabled and release_layer_idx is not None:
-            cur = release_layer_idx
-            if (
-                cur not in self._gpu_layers
-                and cur in self._cpu_weights
-                and self.device is not None
-                and self.copy_stream is not None
-            ):
-                self.prefetch_layer(cur, non_blocking=False)
-                torch.cuda.current_stream().wait_stream(self.copy_stream)
-
-        if self.enabled and prefetch_layer_idx is not None:
-            self.prefetch_layer(prefetch_layer_idx, non_blocking=non_blocking)
-
-        try:
-            yield
-        finally:
-            if self.enabled and self.copy_stream is not None:
-                torch.cuda.current_stream().wait_stream(self.copy_stream)
-            if self.enabled and release_layer_idx is not None:
-                self.release_layer(release_layer_idx)
-
-
-    @torch.compiler.disable
-    def release_layer(self, layer_idx: int) -> None:
-        """Release a layer's tensors back to placeholders (free VRAM)."""
-        if not self.enabled or self.device is None:
-            return
-
-        if layer_idx < 0:
-            return
-
-        param_names = self._gpu_layers.pop(layer_idx, None)
-        if not param_names:
-            return
-
-        for name in param_names:
-            target = self._get_target(name)
-            # Ensure meta exists even if something unexpected happened
-            self._record_meta(name, target)
-            target.data = self._make_placeholder(name)
-
-    @torch.compiler.disable
-    def release_all(self) -> None:
-        """Release all currently-resident layers back to placeholders."""
-        if not self.enabled or self.device is None:
-            return
-
-        if self.copy_stream is not None:
-            torch.cuda.current_stream().wait_stream(self.copy_stream)
-
-        for layer_idx in list(self._gpu_layers.keys()):
-            param_names = self._gpu_layers.pop(layer_idx, None)
-            if not param_names:
-                continue
-            for name in param_names:
-                target = self._get_target(name)
-                self._record_meta(name, target)
-                target.data = self._make_placeholder(name)

fastvideo/models/loader/component_loader.py

@@ -95,6 +95,10 @@ class ComponentLoader(ABC):
             "image_encoder": (ImageEncoderLoader, "transformers"),
             "upsampler": (UpsamplerLoader, "diffusers"),
             "upsampler_2": (UpsamplerLoader, "diffusers"),
+            # Stable Audio's `StableAudioMultiConditioner` bundles T5 +
+            # NumberConditioners; not a pure text encoder, so it gets
+            # its own loader.
+            "conditioner": (ConditionerLoader, "fastvideo"),
         }
 
         if module_type in module_loaders:
@@ -998,6 +1002,61 @@ class SchedulerLoader(ComponentLoader):
         return scheduler
 
 
+class ConditionerLoader(ComponentLoader):
+    """Loader for multi-conditioner components (e.g. Stable Audio's
+    `StableAudioMultiConditioner`, which bundles T5 + NumberConditioners
+    and is neither a pure text encoder nor a Diffusers-shaped module).
+    Reads `<subfolder>/config.json` to resolve the class via
+    `ModelRegistry`, instantiates with no args (the class pulls its own
+    defaults from its FastVideo config), then loads
+    `diffusion_pytorch_model.safetensors` non-strictly so externally
+    fetched sub-encoders (T5) don't trip the missing-key check.
+    """
+
+    def load(self, model_path: str, fastvideo_args: FastVideoArgs):
+        config = get_diffusers_config(model=model_path)
+        class_name = config.pop("_class_name", None)
+        config.pop("_name_or_path", None)
+        if class_name is None:
+            raise ValueError(
+                f"Conditioner config at {model_path} is missing the "
+                f"`_class_name` attribute required to resolve a model class.")
+        model_cls, _ = ModelRegistry.resolve_model_cls(class_name)
+
+        target_device = get_local_torch_device()
+        precision = getattr(fastvideo_args.pipeline_config, "precision", "fp16")
+        target_dtype = PRECISION_TO_TYPE.get(precision, torch.float16)
+
+        # Without this merge the model falls back to its dataclass
+        # defaults (e.g. SA-1.0's 3-conditioner spec — wrong for SA-small).
+        from dataclasses import fields as _fields
+        from fastvideo.configs.models.encoders import (
+            StableAudioConditionerConfig, )
+        if model_cls.__name__ == "StableAudioMultiConditioner":
+            cond_config = StableAudioConditionerConfig()
+            # `update_model_arch` is strict (raises on unknown keys); the
+            # converter writes a few non-arch keys (`_class_name`,
+            # `_diffusers_version`, `_name_or_path`) that must be filtered
+            # out first.
+            valid = {f.name for f in _fields(cond_config.arch_config)}
+            cond_config.update_model_arch({k: v for k, v in config.items() if k in valid})
+            with set_default_torch_dtype(target_dtype):
+                model = model_cls(cond_config)
+        else:
+            with set_default_torch_dtype(target_dtype):
+                model = model_cls()
+
+        weights = os.path.join(str(model_path), "diffusion_pytorch_model.safetensors")
+        if not os.path.isfile(weights):
+            raise FileNotFoundError(
+                f"Conditioner weights not found: {weights}")
+        state = safetensors_load_file(weights)
+        # Non-strict: T5 weights live outside this checkpoint (fetched in
+        # the conditioner's `__init__` from the standard HF repo).
+        model.load_state_dict(state, strict=False)
+        return model.to(device=target_device, dtype=target_dtype).eval()
+
+
 class UpsamplerLoader(ComponentLoader):
     """Loader for upsamplers."""
 

fastvideo/models/registry.py

@@ -86,6 +86,9 @@ _VAE_MODELS = {
     ("vaes", "gen3c_tokenizer_vae", "AutoencoderKLGen3CTokenizer"),
     "AutoencoderKLStepvideo": ("vaes", "stepvideovae", "AutoencoderKLStepvideo"),
     "CausalVideoAutoencoder": ("vaes", "ltx2vae", "LTX2CausalVideoAutoencoder"),
+    # `stable-audio-open-1.0/vae/config.json` ships `_class_name="AutoencoderOobleck"`
+    # (Diffusers' name); FastVideo's class is `OobleckVAE`.
+    "AutoencoderOobleck": ("vaes", "oobleck", "OobleckVAE"),
 }
 
 _AUDIO_MODELS = {

fastvideo/models/vaes/oobleck.py

@@ -0,0 +1,376 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Stable Audio Open 1.0 "Oobleck" VAE.
+
+5-stage Conv1d autoencoder with Snake activations + diagonal-Gaussian
+bottleneck. Loads `stabilityai/stable-audio-open-1.0/vae/` weights
+directly via `OobleckVAE.from_pretrained(...)`.
+
+    vae = OobleckVAE.from_pretrained("stabilityai/stable-audio-open-1.0", subfolder="vae")
+    waveform = vae.decode(latent)            # (B, audio_channels, samples)
+    latent   = vae.encode(waveform).sample()  # or .mode()
+"""
+from __future__ import annotations
+
+import json
+import math
+import os
+from dataclasses import dataclass
+
+import numpy as np
+import torch
+import torch.nn as nn
+from torch.nn.utils import weight_norm
+
+from fastvideo.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+class Snake1d(nn.Module):
+    """A 1D Snake activation with learnable per-channel alpha/beta."""
+
+    def __init__(self, hidden_dim: int, logscale: bool = True):
+        super().__init__()
+        self.alpha = nn.Parameter(torch.zeros(1, hidden_dim, 1))
+        self.beta = nn.Parameter(torch.zeros(1, hidden_dim, 1))
+        self.alpha.requires_grad = True
+        self.beta.requires_grad = True
+        self.logscale = logscale
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        shape = x.shape
+        alpha = self.alpha if not self.logscale else torch.exp(self.alpha)
+        beta = self.beta if not self.logscale else torch.exp(self.beta)
+        x = x.reshape(shape[0], shape[1], -1)
+        x = x + (beta + 1e-9).reciprocal() * torch.sin(alpha * x).pow(2)
+        return x.reshape(shape)
+
+
+class OobleckResidualUnit(nn.Module):
+    def __init__(self, dimension: int = 16, dilation: int = 1):
+        super().__init__()
+        pad = ((7 - 1) * dilation) // 2
+        self.snake1 = Snake1d(dimension)
+        self.conv1 = weight_norm(nn.Conv1d(
+            dimension, dimension, kernel_size=7, dilation=dilation, padding=pad,
+        ))
+        self.snake2 = Snake1d(dimension)
+        self.conv2 = weight_norm(nn.Conv1d(dimension, dimension, kernel_size=1))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        out = self.conv1(self.snake1(x))
+        out = self.conv2(self.snake2(out))
+        pad = (x.shape[-1] - out.shape[-1]) // 2
+        if pad > 0:
+            x = x[..., pad:-pad]
+        return x + out
+
+
+class OobleckEncoderBlock(nn.Module):
+    def __init__(self, input_dim: int, output_dim: int, stride: int = 1):
+        super().__init__()
+        self.res_unit1 = OobleckResidualUnit(input_dim, dilation=1)
+        self.res_unit2 = OobleckResidualUnit(input_dim, dilation=3)
+        self.res_unit3 = OobleckResidualUnit(input_dim, dilation=9)
+        self.snake1 = Snake1d(input_dim)
+        self.conv1 = weight_norm(nn.Conv1d(
+            input_dim, output_dim,
+            kernel_size=2 * stride, stride=stride,
+            padding=math.ceil(stride / 2),
+        ))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.res_unit1(x)
+        x = self.res_unit2(x)
+        x = self.snake1(self.res_unit3(x))
+        return self.conv1(x)
+
+
+class OobleckDecoderBlock(nn.Module):
+    def __init__(self, input_dim: int, output_dim: int, stride: int = 1):
+        super().__init__()
+        self.snake1 = Snake1d(input_dim)
+        self.conv_t1 = weight_norm(nn.ConvTranspose1d(
+            input_dim, output_dim,
+            kernel_size=2 * stride, stride=stride,
+            padding=math.ceil(stride / 2),
+        ))
+        self.res_unit1 = OobleckResidualUnit(output_dim, dilation=1)
+        self.res_unit2 = OobleckResidualUnit(output_dim, dilation=3)
+        self.res_unit3 = OobleckResidualUnit(output_dim, dilation=9)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.snake1(x)
+        x = self.conv_t1(x)
+        x = self.res_unit1(x)
+        x = self.res_unit2(x)
+        return self.res_unit3(x)
+
+
+class OobleckDiagonalGaussianDistribution:
+    """Diagonal-Gaussian VAE posterior with `softplus(scale) + 1e-4` std."""
+
+    def __init__(self, parameters: torch.Tensor, deterministic: bool = False):
+        self.parameters = parameters
+        self.mean, self.scale = parameters.chunk(2, dim=1)
+        self.std = nn.functional.softplus(self.scale) + 1e-4
+        self.var = self.std * self.std
+        self.logvar = torch.log(self.var)
+        self.deterministic = deterministic
+
+    def sample(self, generator: torch.Generator | None = None) -> torch.Tensor:
+        noise = torch.randn(
+            self.mean.shape, generator=generator,
+            device=self.parameters.device, dtype=self.parameters.dtype,
+        )
+        return self.mean + self.std * noise
+
+    def mode(self) -> torch.Tensor:
+        return self.mean
+
+
+@dataclass
+class OobleckDecoderOutput:
+    sample: torch.Tensor
+
+
+class OobleckEncoder(nn.Module):
+    def __init__(
+        self,
+        encoder_hidden_size: int,
+        audio_channels: int,
+        downsampling_ratios: list[int],
+        channel_multiples: list[int],
+    ):
+        super().__init__()
+        strides = downsampling_ratios
+        channel_multiples = [1] + list(channel_multiples)
+        self.conv1 = weight_norm(nn.Conv1d(
+            audio_channels, encoder_hidden_size, kernel_size=7, padding=3,
+        ))
+        self.block = nn.ModuleList([
+            OobleckEncoderBlock(
+                input_dim=encoder_hidden_size * channel_multiples[i],
+                output_dim=encoder_hidden_size * channel_multiples[i + 1],
+                stride=s,
+            )
+            for i, s in enumerate(strides)
+        ])
+        d_model = encoder_hidden_size * channel_multiples[-1]
+        self.snake1 = Snake1d(d_model)
+        self.conv2 = weight_norm(nn.Conv1d(
+            d_model, encoder_hidden_size, kernel_size=3, padding=1,
+        ))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.conv1(x)
+        for m in self.block:
+            x = m(x)
+        x = self.snake1(x)
+        return self.conv2(x)
+
+
+class OobleckDecoder(nn.Module):
+    def __init__(
+        self,
+        channels: int,
+        input_channels: int,
+        audio_channels: int,
+        upsampling_ratios: list[int],
+        channel_multiples: list[int],
+    ):
+        super().__init__()
+        strides = upsampling_ratios
+        channel_multiples = [1] + list(channel_multiples)
+        self.conv1 = weight_norm(nn.Conv1d(
+            input_channels, channels * channel_multiples[-1],
+            kernel_size=7, padding=3,
+        ))
+        self.block = nn.ModuleList([
+            OobleckDecoderBlock(
+                input_dim=channels * channel_multiples[len(strides) - i],
+                output_dim=channels * channel_multiples[len(strides) - i - 1],
+                stride=s,
+            )
+            for i, s in enumerate(strides)
+        ])
+        self.snake1 = Snake1d(channels)
+        self.conv2 = weight_norm(nn.Conv1d(
+            channels, audio_channels, kernel_size=7, padding=3, bias=False,
+        ))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.conv1(x)
+        for layer in self.block:
+            x = layer(x)
+        x = self.snake1(x)
+        return self.conv2(x)
+
+
+# ---------------------------------------------------------------------------
+# Top-level VAE
+# ---------------------------------------------------------------------------
+
+
+class OobleckVAE(nn.Module):
+    """Stable Audio Open 1.0 VAE.
+
+    Constructed either from an `OobleckVAEConfig` (the standard
+    `VAELoader` path) or from explicit kwargs (back-compat for tests
+    and `from_pretrained` callers).
+    """
+
+    def __init__(
+        self,
+        config=None,  # type: OobleckVAEConfig | None
+        *,
+        encoder_hidden_size: int = 128,
+        downsampling_ratios: list[int] | None = None,
+        channel_multiples: list[int] | None = None,
+        decoder_channels: int = 128,
+        decoder_input_channels: int = 64,
+        audio_channels: int = 2,
+        sampling_rate: int = 44100,
+    ):
+        super().__init__()
+        if config is not None:
+            arch = config.arch_config
+            encoder_hidden_size = arch.encoder_hidden_size
+            downsampling_ratios = list(arch.downsampling_ratios)
+            channel_multiples = list(arch.channel_multiples)
+            decoder_channels = arch.decoder_channels
+            decoder_input_channels = arch.decoder_input_channels
+            audio_channels = arch.audio_channels
+            sampling_rate = arch.sampling_rate
+        if downsampling_ratios is None:
+            downsampling_ratios = [2, 4, 4, 8, 8]
+        if channel_multiples is None:
+            channel_multiples = [1, 2, 4, 8, 16]
+        self.encoder_hidden_size = encoder_hidden_size
+        self.downsampling_ratios = downsampling_ratios
+        self.decoder_channels = decoder_channels
+        self.upsampling_ratios = list(reversed(downsampling_ratios))
+        self.hop_length = int(np.prod(downsampling_ratios))
+        self.sampling_rate = sampling_rate
+        self.audio_channels = audio_channels
+        self.decoder_input_channels = decoder_input_channels
+
+        self.encoder = OobleckEncoder(
+            encoder_hidden_size=encoder_hidden_size,
+            audio_channels=audio_channels,
+            downsampling_ratios=downsampling_ratios,
+            channel_multiples=channel_multiples,
+        )
+        self.decoder = OobleckDecoder(
+            channels=decoder_channels,
+            input_channels=decoder_input_channels,
+            audio_channels=audio_channels,
+            upsampling_ratios=self.upsampling_ratios,
+            channel_multiples=channel_multiples,
+        )
+
+    def encode(
+        self, x: torch.Tensor,
+    ) -> OobleckDiagonalGaussianDistribution:
+        return OobleckDiagonalGaussianDistribution(self.encoder(x))
+
+    def decode(self, z: torch.Tensor) -> OobleckDecoderOutput:
+        return OobleckDecoderOutput(sample=self.decoder(z))
+
+    def forward(
+        self, sample: torch.Tensor, sample_posterior: bool = False,
+    ) -> OobleckDecoderOutput:
+        posterior = self.encode(sample)
+        z = posterior.sample() if sample_posterior else posterior.mode()
+        return self.decode(z)
+
+    # -------------------------------------------------------------------
+    # Loader
+    # -------------------------------------------------------------------
+    @classmethod
+    def from_pretrained(
+        cls,
+        model_path: str,
+        *,
+        subfolder: str | None = None,
+        torch_dtype: torch.dtype | None = None,
+    ) -> "OobleckVAE":
+        """Instantiate and load weights from a Stable Audio VAE dir.
+
+        `model_path` may be:
+          * a HF repo id (e.g. `stabilityai/stable-audio-open-1.0`),
+          * a local directory containing `config.json` + safetensors,
+          * a local directory whose `subfolder="vae"` holds those files.
+
+        For gated repos, the HF token is read from `HF_TOKEN` /
+        `HUGGINGFACE_HUB_TOKEN` / `HF_API_KEY` (see `resolve_hf_token`).
+        """
+        import inspect
+        from safetensors.torch import load_file
+
+        from fastvideo.utils import resolve_hf_token
+
+        # Resolve to a local directory.
+        if os.path.isdir(model_path):
+            root = model_path
+        else:
+            from huggingface_hub import snapshot_download
+            allow = ["vae/*"] if subfolder else ["*"]
+            root = snapshot_download(
+                repo_id=model_path, token=resolve_hf_token(), allow_patterns=allow,
+            )
+        if subfolder:
+            root = os.path.join(root, subfolder)
+        if not os.path.isdir(root):
+            raise FileNotFoundError(f"Not a directory: {root}")
+
+        cfg_path = os.path.join(root, "config.json")
+        if not os.path.isfile(cfg_path):
+            raise FileNotFoundError(
+                f"Expected config.json at {cfg_path}. If using a HF repo, "
+                f"pass subfolder='vae'."
+            )
+        with open(cfg_path) as f:
+            cfg = json.load(f)
+        cfg_fields = {k: v for k, v in cfg.items() if not k.startswith("_")}
+        # Diffusers configs commonly carry extra fields (`scaling_factor`,
+        # `_diffusers_version`, ...) the bare `OobleckVAE` ctor doesn't accept.
+        init_params = inspect.signature(cls.__init__).parameters
+        cfg_fields = {k: v for k, v in cfg_fields.items() if k in init_params}
+
+        model = cls(**cfg_fields)
+
+        weights_path = os.path.join(root, "diffusion_pytorch_model.safetensors")
+        if not os.path.isfile(weights_path):
+            # Allow `model.safetensors` as a fallback.
+            alt = os.path.join(root, "model.safetensors")
+            if os.path.isfile(alt):
+                weights_path = alt
+            else:
+                raise FileNotFoundError(
+                    f"No safetensors weights under {root}. Expected "
+                    f"diffusion_pytorch_model.safetensors."
+                )
+        state = load_file(weights_path)
+        missing, unexpected = model.load_state_dict(state, strict=False)
+        if missing:
+            raise RuntimeError(
+                f"OobleckVAE missing {len(missing)} keys from {weights_path}: "
+                f"{missing[:5]}"
+            )
+        if unexpected:
+            # Non-critical: some checkpoints embed the VAE inside a larger
+            # container (e.g. `pretransform.model.*`). Log the count so
+            # genuine loader regressions don't go unnoticed.
+            logger.debug(
+                "OobleckVAE: ignored %d unexpected keys from %s "
+                "(first 3: %s)", len(unexpected), weights_path, unexpected[:3],
+            )
+
+        if torch_dtype is not None:
+            model = model.to(dtype=torch_dtype)
+        model.eval()
+        return model
+
+
+EntryClass = OobleckVAE

fastvideo/models/vaes/sa_audio.py

@@ -0,0 +1,122 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Lazy-loading pipeline wrapper around `OobleckVAE`.
+
+Two reasons this exists rather than using `OobleckVAE` directly:
+
+  1. The underlying VAE is fetched on first `encode`/`decode` call, not
+     at construction — lets pipelines build the module tree on CPU
+     before knowing the target device.
+  2. The lazy VAE's params are hidden from `named_parameters()` so the
+     FastVideo pipeline-component loader doesn't try to match Oobleck's
+     safetensors against the host pipeline's converted-repo state dict.
+
+For standalone use prefer `OobleckVAE.from_pretrained(...)` directly.
+"""
+from __future__ import annotations
+
+import os
+
+import torch
+from torch import nn
+
+from fastvideo.configs.models.vaes import OobleckVAEConfig
+
+
+class SAAudioVAEModel(nn.Module):
+    """Pipeline-glue lazy loader around `OobleckVAE`."""
+
+    def __init__(self, config: OobleckVAEConfig) -> None:
+        super().__init__()
+        self.config = config
+        arch = config.arch_config
+        self.pretrained_path: str = config.pretrained_path
+        self.pretrained_subfolder: str | None = config.pretrained_subfolder
+        self.pretrained_dtype: str = config.pretrained_dtype
+        self.sampling_rate: int = arch.sampling_rate
+        self.audio_channels: int = arch.audio_channels
+        self.decoder_input_channels: int = arch.decoder_input_channels
+        self._oobleck_vae = None
+
+    def named_parameters(self, prefix: str = "", recurse: bool = True):
+        # Hide the lazy-loaded VAE — its weights are fetched separately
+        # and shouldn't appear in the host pipeline's loader sweep.
+        for name, param in super().named_parameters(prefix=prefix, recurse=recurse):
+            if name.startswith("_oobleck_vae.") or name == "_oobleck_vae":
+                continue
+            yield name, param
+
+    def _build(self, device: torch.device | None = None):
+        from fastvideo.models.vaes.oobleck import OobleckVAE
+
+        path = self.pretrained_path
+        if not path:
+            raise ValueError(
+                "OobleckVAEConfig.pretrained_path must be set; expected "
+                "`stabilityai/stable-audio-open-1.0` or a local path."
+            )
+        dtype = getattr(torch, self.pretrained_dtype, torch.float32)
+        # If the caller already pointed us at the VAE dir directly, drop
+        # the subfolder. Otherwise pass through (default "vae").
+        subfolder: str | None = self.pretrained_subfolder
+        if subfolder and os.path.isdir(path) and os.path.isfile(os.path.join(path, "config.json")):
+            subfolder = None
+        model = OobleckVAE.from_pretrained(path, subfolder=subfolder, torch_dtype=dtype)
+        if device is not None:
+            model = model.to(device=device)
+        model.eval()
+        return model
+
+    @property
+    def oobleck_vae(self):
+        if self._oobleck_vae is None:
+            self._oobleck_vae = self._build()
+        return self._oobleck_vae
+
+    # Back-compat alias: callers that imported this earlier referred to
+    # the underlying VAE as `sa_audio_vae_model`. Both names point at the
+    # same object.
+    @property
+    def sa_audio_vae_model(self):
+        return self.oobleck_vae
+
+    @property
+    def hop_length(self) -> int:
+        return int(self.oobleck_vae.hop_length)
+
+    def _move_to_input_device(self, model, ref: torch.Tensor):
+        if ref is None:
+            return model
+        first_param = next(model.parameters(), None)
+        if first_param is not None and first_param.device != ref.device:
+            model = model.to(device=ref.device)
+            self._oobleck_vae = model
+        return model
+
+    def decode(self, latent: torch.Tensor) -> torch.Tensor:
+        """Decode an audio latent (`[B, C_latent, L]`) -> waveform
+        (`[B, audio_channels, samples]`).
+        """
+        model = self.oobleck_vae
+        model = self._move_to_input_device(model, latent)
+        with torch.no_grad():
+            out = model.decode(latent.to(next(model.parameters()).dtype))
+        if hasattr(out, "sample"):
+            return out.sample
+        return out
+
+    def encode(self, waveform: torch.Tensor, sample_posterior: bool = False) -> torch.Tensor:
+        """Encode `[B, C_audio, samples]` -> latent `[B, C_latent, L]`.
+
+        `sample_posterior=False` (default): deterministic mean.
+        `sample_posterior=True`: stochastic sample (`mean + softplus(scale) * randn`).
+        """
+        model = self.oobleck_vae
+        model = self._move_to_input_device(model, waveform)
+        with torch.no_grad():
+            out = model.encode(waveform.to(next(model.parameters()).dtype))
+        if hasattr(out, "latent_dist"):
+            out = out.latent_dist
+        return out.sample() if sample_posterior else out.mode()
+
+
+EntryClass = SAAudioVAEModel

fastvideo/pipelines/basic/cosmos/cosmos_pipeline.py

@@ -22,23 +22,18 @@ class Cosmos2VideoToWorldPipeline(ComposedPipelineBase):
     _required_config_modules = ["text_encoder", "tokenizer", "vae", "transformer", "scheduler", "safety_checker"]
 
     def initialize_pipeline(self, fastvideo_args: FastVideoArgs):
-        self.modules["scheduler"] = FlowMatchEulerDiscreteScheduler(shift=fastvideo_args.pipeline_config.flow_shift,
-                                                                    use_karras_sigmas=True)
-
-        sigma_max = 80.0
-        sigma_min = 0.002
-        sigma_data = 1.0
-        final_sigmas_type = "sigma_min"
-
-        if self.modules["scheduler"] is not None:
-            scheduler = self.modules["scheduler"]
-            scheduler.config.sigma_max = sigma_max
-            scheduler.config.sigma_min = sigma_min
-            scheduler.config.sigma_data = sigma_data
-            scheduler.config.final_sigmas_type = final_sigmas_type
-            scheduler.sigma_max = sigma_max
-            scheduler.sigma_min = sigma_min
-            scheduler.sigma_data = sigma_data
+        scheduler = FlowMatchEulerDiscreteScheduler(
+            shift=fastvideo_args.pipeline_config.flow_shift,
+            use_karras_sigmas=True,
+        )
+        scheduler.config.sigma_max = 80.0
+        scheduler.config.sigma_min = 0.002
+        scheduler.config.sigma_data = 1.0
+        scheduler.config.final_sigmas_type = "sigma_min"
+        scheduler.sigma_max = 80.0
+        scheduler.sigma_min = 0.002
+        scheduler.sigma_data = 1.0
+        self.modules["scheduler"] = scheduler
 
     def create_pipeline_stages(self, fastvideo_args: FastVideoArgs):
         """Set up pipeline stages with proper dependency injection."""

fastvideo/pipelines/basic/stable_audio/__init__.py

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

fastvideo/pipelines/basic/stable_audio/presets.py

@@ -0,0 +1,63 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Stable Audio presets.
+
+Sampling defaults track the published HF model card
+(https://huggingface.co/stabilityai/stable-audio-open-1.0):
+100 steps, CFG=7, dpmpp-3m-sde, sigma_min=0.3, sigma_max=500, rho=1.0.
+"""
+from fastvideo.api.presets import InferencePreset, PresetStageSpec
+
+_DENOISE_STAGE = PresetStageSpec(
+    name="denoise",
+    kind="denoising",
+    description="Stable Audio Cosine-DPM++ denoising with text + duration CFG.",
+    allowed_overrides=frozenset({
+        "num_inference_steps",
+        "guidance_scale",
+    }),
+)
+
+# `audio_start_in_s` / `audio_end_in_s` are call-kwargs, kept off here.
+# `height`/`width` are pinned to 8 (the shared `InputValidationStage`
+# rejects values that aren't divisible by 8) and `num_frames` to 1 so
+# the video-shaped preallocation in `VideoGenerator` stays tiny — the
+# real output is the audio waveform on `result["audio"]`, not the
+# placeholder frame tensor.
+_SHARED_DEFAULTS = {
+    "seed": 0,
+    "guidance_scale": 7.0,
+    "num_inference_steps": 100,
+    "negative_prompt": "",
+    "height": 8,
+    "width": 8,
+    "num_frames": 1,
+}
+
+STABLE_AUDIO_OPEN_1_0_BASE = InferencePreset(
+    name="stable_audio_open_1_0_base",
+    version=1,
+    model_family="stable_audio",
+    description=("Stability AI Stable Audio Open 1.0 text-to-audio. Generates up "
+                 "to ~47.5s of stereo 44.1 kHz audio per call. Default duration "
+                 "is 10s; raise via `audio_end_in_s` up to the model max."),
+    workload_type="t2v",  # NOTE: WorkloadType has no T2A variant yet (REVIEW item 28)
+    stage_schemas=(_DENOISE_STAGE, ),
+    defaults=dict(_SHARED_DEFAULTS),
+)
+
+# Smaller / faster checkpoint with the same Oobleck VAE but a 1024-dim
+# 16-layer DiT with `qk_norm="ln"`. Sampling defaults match the official
+# `stable-audio-open-small` model card.
+STABLE_AUDIO_OPEN_SMALL = InferencePreset(
+    name="stable_audio_open_small",
+    version=1,
+    model_family="stable_audio",
+    description=("Stability AI Stable Audio Open Small text-to-audio. Faster than "
+                 "the 1.0 base; supports up to ~11.9s of stereo 44.1 kHz audio per "
+                 "call (smaller training window)."),
+    workload_type="t2v",
+    stage_schemas=(_DENOISE_STAGE, ),
+    defaults=dict(_SHARED_DEFAULTS),
+)
+
+ALL_PRESETS = (STABLE_AUDIO_OPEN_1_0_BASE, STABLE_AUDIO_OPEN_SMALL)

fastvideo/pipelines/basic/stable_audio/stable_audio_pipeline.py

@@ -0,0 +1,125 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Stable Audio Open 1.0 pipeline (T2A + A2A + RePaint inpainting).
+
+Components are loaded via the standard
+`ComposedPipelineBase.load_modules` against the FastVideo-curated
+Diffusers-format repo `FastVideo/stable-audio-open-1.0-Diffusers`
+(produced by
+`scripts/checkpoint_conversion/stable_audio_to_diffusers.py`). The DiT
+is a `BaseDiT` subclass loaded by `TransformerLoader`; the VAE is
+loaded by `VAELoader`; the multi-conditioner (T5 + NumberConditioners)
+is loaded by `ConditionerLoader` (a Stable Audio-specific addition).
+
+Stages:
+
+    InputValidationStage
+      → StableAudioConditioningStage      (T5 + NumberConditioner -> cross-attn + global cond, with CFG)
+      → StableAudioLatentPreparationStage (initial Gaussian noise; encodes A2A / inpaint refs)
+      → StableAudioDenoisingStage         (k-diffusion `dpmpp-3m-sde` over the DiT)
+      → StableAudioDecodingStage          (OobleckVAE -> waveform)
+"""
+from __future__ import annotations
+
+import functools
+
+import torch
+
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.logger import init_logger
+from fastvideo.pipelines.basic.stable_audio.stages import (
+    StableAudioConditioningStage,
+    StableAudioDecodingStage,
+    StableAudioDenoisingStage,
+    StableAudioLatentPreparationStage,
+)
+from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase
+from fastvideo.pipelines.stages import InputValidationStage
+
+logger = init_logger(__name__)
+
+
+@functools.lru_cache(maxsize=1)
+def _warn_tf32_disabled_for_stable_audio() -> None:
+    logger.warning("Stable Audio pipeline is disabling process-global "
+                   "torch.backends.{cuda.matmul.allow_tf32, cudnn.allow_tf32, "
+                   "cuda.matmul.allow_fp16_reduced_precision_reduction, "
+                   "cudnn.benchmark} for A2A renoise determinism. Other models "
+                   "loaded into this process will inherit these settings.")
+
+
+def _disable_tf32_for_stable_audio() -> None:
+    """Disable TF32 / cuDNN nondeterminism — A2A renoise-then-denoise SDE
+    amplifies per-element drift, and the published parity bounds were
+    set with these off. Process-global; the first call logs a warning.
+    """
+    _warn_tf32_disabled_for_stable_audio()
+    torch.backends.cuda.matmul.allow_tf32 = False
+    torch.backends.cudnn.allow_tf32 = False
+    torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = False
+    torch.backends.cudnn.benchmark = False
+
+
+class StableAudioPipeline(ComposedPipelineBase):
+    """Stable Audio Open 1.0 pipeline.
+
+    Mode is kwargs-driven on `generate_video()`:
+
+      * Text-to-audio (default) -- `prompt=...`, `audio_end_in_s=...`
+      * Audio-to-audio variation -- add `init_audio=ref` (and optionally
+        `init_noise_level`, lower = closer to reference)
+      * RePaint inpainting / outpainting -- add `inpaint_audio=ref` and
+        `inpaint_mask` (1-D, 1 = keep / 0 = regenerate)
+
+    See `examples/inference/basic/basic_stable_audio*.py` for runnable
+    examples of each mode.
+    """
+
+    _required_config_modules = [
+        "vae",
+        "transformer",
+        "conditioner",
+    ]
+
+    def initialize_pipeline(self, fastvideo_args: FastVideoArgs) -> None:
+        """Apply Stable Audio's process-global numerics overrides BEFORE
+        the standard component loaders run (TF32 off for A2A renoise
+        determinism)."""
+        _disable_tf32_for_stable_audio()
+
+    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None:
+        pc = fastvideo_args.pipeline_config
+
+        self.add_stage(stage_name="input_validation_stage", stage=InputValidationStage())
+
+        self.add_stage(
+            stage_name="conditioning_stage",
+            stage=StableAudioConditioningStage(conditioner=self.get_module("conditioner")),
+        )
+
+        self.add_stage(
+            stage_name="latent_preparation_stage",
+            stage=StableAudioLatentPreparationStage(
+                io_channels=64,
+                # Per-variant training window: 2,097,152 (~47.5s) for
+                # SA-1.0; 524,288 (~11.9s) for SA-small. Pulled from the
+                # pipeline config so each variant gets its own latent
+                # length.
+                sample_size=pc.sample_size,
+                vae=self.get_module("vae"),
+                sample_rate=pc.sampling_rate,
+                audio_channels=pc.audio_channels,
+            ),
+        )
+
+        self.add_stage(
+            stage_name="denoising_stage",
+            stage=StableAudioDenoisingStage(transformer=self.get_module("transformer")),
+        )
+
+        self.add_stage(
+            stage_name="decoding_stage",
+            stage=StableAudioDecodingStage(vae=self.get_module("vae")),
+        )
+
+
+EntryClass = StableAudioPipeline

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

@@ -0,0 +1,12 @@
+# SPDX-License-Identifier: Apache-2.0
+from fastvideo.pipelines.basic.stable_audio.stages.conditioning import StableAudioConditioningStage
+from fastvideo.pipelines.basic.stable_audio.stages.decoding import StableAudioDecodingStage
+from fastvideo.pipelines.basic.stable_audio.stages.denoising import StableAudioDenoisingStage
+from fastvideo.pipelines.basic.stable_audio.stages.latent_preparation import StableAudioLatentPreparationStage
+
+__all__ = [
+    "StableAudioConditioningStage",
+    "StableAudioDecodingStage",
+    "StableAudioDenoisingStage",
+    "StableAudioLatentPreparationStage",
+]

fastvideo/pipelines/basic/stable_audio/stages/conditioning.py

@@ -0,0 +1,98 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Stable Audio conditioning stage."""
+from __future__ import annotations
+
+import torch
+
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+from fastvideo.pipelines.stages.base import PipelineStage
+from fastvideo.pipelines.stages.validators import VerificationResult
+
+
+class StableAudioConditioningStage(PipelineStage):
+    """Run the conditioner over the prompt + duration and stash the
+    DiT-ready (cross_attn_cond, cross_attn_mask, global_embed) triple
+    on `batch.extra` (plus the negative-prompt triple when CFG is on).
+    """
+
+    def __init__(self, conditioner) -> None:
+        super().__init__()
+        self.conditioner = conditioner
+
+    def verify_input(self, batch, fastvideo_args):
+        return VerificationResult()
+
+    def verify_output(self, batch, fastvideo_args):
+        return VerificationResult()
+
+    @torch.inference_mode()
+    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
+        pc = fastvideo_args.pipeline_config
+        device = next(self.conditioner.parameters()).device
+
+        start_attr = getattr(batch, "audio_start_in_s", None)
+        end_attr = getattr(batch, "audio_end_in_s", None)
+        audio_start_in_s = float(start_attr if start_attr is not None else pc.audio_start_in_s)
+        audio_end_in_s = float(end_attr if end_attr is not None else pc.audio_end_in_s)
+        max_duration = float(getattr(pc, "max_audio_duration_s", 2097152 / 44100))
+        if audio_start_in_s < 0:
+            raise ValueError(f"audio_start_in_s must be >= 0, got {audio_start_in_s}.")
+        if audio_end_in_s <= audio_start_in_s:
+            raise ValueError(f"audio_end_in_s ({audio_end_in_s}) must be > audio_start_in_s "
+                             f"({audio_start_in_s}).")
+        if audio_end_in_s > max_duration:
+            raise ValueError(f"audio_end_in_s ({audio_end_in_s}s) exceeds the model's fixed "
+                             f"window of {max_duration:.4f}s. Stable Audio Open 1.0 always "
+                             f"samples a 2,097,152-frame latent and slices to "
+                             f"[start, end] after decode; values past the window are silently "
+                             f"truncated. Lower audio_end_in_s or split the request.")
+        guidance_scale = float(batch.guidance_scale or pc.guidance_scale)
+        do_cfg = guidance_scale > 1.0
+
+        if isinstance(batch.prompt, str):
+            prompt = batch.prompt
+        elif isinstance(batch.prompt, list):
+            if len(batch.prompt) > 1:
+                raise ValueError(f"Stable Audio does not support batched prompts; got "
+                                 f"{len(batch.prompt)} entries. Pass a single string or a "
+                                 f"single-element list.")
+            prompt = batch.prompt[0] if batch.prompt else ""
+        else:
+            raise TypeError(f"`prompt` must be a string or a list of strings, got "
+                            f"{type(batch.prompt).__name__}.")
+        # Send only the keys the conditioner declares (per-variant).
+        all_cond_values = {
+            "prompt": prompt,
+            "seconds_start": audio_start_in_s,
+            "seconds_total": audio_end_in_s,
+        }
+        active_ids = self.conditioner.cross_attention_cond_ids
+        cond_meta = [{k: all_cond_values[k] for k in active_ids if k in all_cond_values}]
+        cond = self.conditioner(cond_meta, device)
+        cross_attn_cond, cross_attn_mask, global_embed = self.conditioner.get_conditioning_inputs(cond)
+
+        neg_cross_attn_cond = None
+        neg_cross_attn_mask = None
+        neg_global_embed = None
+        if do_cfg:
+            neg_prompt = batch.negative_prompt or ""
+            if isinstance(neg_prompt, list):
+                neg_prompt = neg_prompt[0] if neg_prompt else ""
+            neg_values = dict(all_cond_values, prompt=neg_prompt)
+            neg_meta = [{k: neg_values[k] for k in active_ids if k in neg_values}]
+            neg = self.conditioner(neg_meta, device)
+            neg_cross_attn_cond, neg_cross_attn_mask, neg_global_embed = (self.conditioner.get_conditioning_inputs(neg))
+
+        if batch.extra is None:
+            batch.extra = {}
+        batch.extra["cross_attn_cond"] = cross_attn_cond
+        batch.extra["cross_attn_mask"] = cross_attn_mask
+        batch.extra["global_embed"] = global_embed
+        batch.extra["negative_cross_attn_cond"] = neg_cross_attn_cond
+        batch.extra["negative_cross_attn_mask"] = neg_cross_attn_mask
+        batch.extra["negative_global_embed"] = neg_global_embed
+        batch.extra["do_cfg"] = do_cfg
+        batch.extra["audio_start_in_s"] = audio_start_in_s
+        batch.extra["audio_end_in_s"] = audio_end_in_s
+        return batch

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

@@ -0,0 +1,67 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Stable Audio decoding: latent -> waveform via OobleckVAE.
+
+Slices the output to `[audio_start_in_s, audio_end_in_s]` and stashes
+the result on `batch.extra["audio"]` + `["audio_sample_rate"]` for
+`VideoGenerator._mux_audio` to pick up.
+"""
+from __future__ import annotations
+
+import torch
+
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+from fastvideo.pipelines.stages.base import PipelineStage
+from fastvideo.pipelines.stages.validators import VerificationResult
+
+
+class StableAudioDecodingStage(PipelineStage):
+    """Decode latent → audio waveform + slice to [start, end]."""
+
+    def __init__(self, vae) -> None:
+        super().__init__()
+        self.vae = vae
+
+    def verify_input(self, batch, fastvideo_args):
+        return VerificationResult()
+
+    def verify_output(self, batch, fastvideo_args):
+        return VerificationResult()
+
+    @torch.inference_mode()
+    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
+        pc = fastvideo_args.pipeline_config
+        latents = batch.latents
+
+        # VAE may be CPU-parked under `vae_cpu_offload=True`.
+        from fastvideo.distributed.parallel_state import get_local_torch_device
+        self.vae = self.vae.to(get_local_torch_device())
+        decoded = self.vae.decode(latents)
+        if hasattr(decoded, "sample"):  # tolerate tensor or dataclass
+            decoded = decoded.sample
+
+        sr = int(getattr(self.vae, "sampling_rate", pc.sampling_rate))
+        start_in_s = float(batch.extra.get("audio_start_in_s", pc.audio_start_in_s))
+        end_in_s = float(batch.extra.get("audio_end_in_s", pc.audio_end_in_s))
+        decoded = decoded[:, :, int(start_in_s * sr):int(end_in_s * sr)]
+
+        if batch.extra is None:
+            batch.extra = {}
+        # `_mux_audio` / `_write_pcm_wav` want `[samples, channels]`.
+        batch.extra["audio"] = decoded.squeeze(0).T.detach().float().cpu().numpy()
+        batch.extra["audio_sample_rate"] = sr
+        batch.extra["audio_only"] = True
+        # Raw tensor for parity tests.
+        batch.extra["decoded_audio"] = decoded.detach().cpu()
+
+        # `VideoGenerator.generate_video` is video-shaped (asserts
+        # `output_batch.output is not None`); fill with a placeholder of
+        # the expected `[B, 3, num_frames, H, W]` shape — the real audio
+        # is on `batch.extra` above. Pure-audio workload support tracked
+        # in REVIEW item 28.
+        b = decoded.shape[0]
+        n_frames = int(getattr(batch, "num_frames", 1) or 1)
+        h = int(getattr(batch, "height", 1) or 1)
+        w = int(getattr(batch, "width", 1) or 1)
+        batch.output = torch.zeros((b, 3, n_frames, h, w), dtype=torch.uint8)
+        return batch

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

@@ -0,0 +1,207 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Stable Audio denoising — k-diffusion `dpmpp-3m-sde` over the DiT.
+
+CFG-batched conditioning is built once outside the sampler loop so the
+adapter only does `cat([x, x])` + DiT call per step.
+"""
+from __future__ import annotations
+
+import math
+
+import torch
+import torch.nn as nn
+
+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.base import PipelineStage
+from fastvideo.pipelines.stages.validators import VerificationResult
+
+
+class _DiTAdapter(nn.Module):
+    """`StableAudioDiT` -> `K.external.VDenoiser` adapter.
+
+    `batch_cond` / `batch_global` are precomputed CFG-batched tensors
+    (`[2, ...]` for CFG, `[1, ...]` otherwise); building them once
+    outside the sampler loop saves ~3 cats × 100 steps per call.
+    """
+
+    def __init__(self, dit, *, batch_cond: torch.Tensor, batch_global: torch.Tensor, cfg_scale: float) -> None:
+        super().__init__()
+        self.dit = dit
+        self.batch_cond = batch_cond
+        self.batch_global = batch_global
+        self.cfg_scale = cfg_scale
+        self.do_cfg = cfg_scale != 1.0
+
+    def forward(self, x: torch.Tensor, t: torch.Tensor, **_unused) -> torch.Tensor:
+        if not self.do_cfg:
+            return self.dit(x, t, cross_attn_cond=self.batch_cond, global_embed=self.batch_global)
+        batch_x = torch.cat([x, x], dim=0)
+        batch_t = torch.cat([t, t], dim=0)
+        out = self.dit(batch_x, batch_t, cross_attn_cond=self.batch_cond, global_embed=self.batch_global)
+        cond_out, uncond_out = torch.chunk(out, 2, dim=0)
+        return uncond_out + (cond_out - uncond_out) * self.cfg_scale
+
+
+class StableAudioDenoisingStage(PipelineStage):
+    """k-diffusion `dpmpp-3m-sde` sampling loop."""
+
+    # Sampler defaults from the published model card.
+    _SIGMA_MIN = 0.3
+    _SIGMA_MAX = 500.0
+    _RHO = 1.0
+    _LOG_SIGMA_MIN = math.log(_SIGMA_MIN)
+    _LOG_SIGMA_MAX = math.log(_SIGMA_MAX)
+
+    def __init__(self, transformer) -> None:
+        super().__init__()
+        self.transformer = transformer
+
+    def _resolve_sigma_max(self, batch) -> float:
+        """Map A2A intent to `sigma_max`.
+
+        Public knob is `init_audio_strength` (0..1, higher = closer to
+        source), log-interpolated between SIGMA_MIN (= preservation) and
+        SIGMA_MAX (= full T2A). Raw `init_noise_level` is the legacy
+        sigma_max override; passing both is an error.
+        """
+        raw = getattr(batch, "init_noise_level", None)
+        strength = getattr(batch, "init_audio_strength", None)
+        if raw is not None and strength is not None:
+            raise ValueError("Pass `init_audio_strength` (0..1) OR `init_noise_level` "
+                             "(raw sigma_max), not both.")
+        if raw is not None:
+            return float(raw)
+        s = max(0.0, min(1.0, float(strength) if strength is not None else 0.6))
+        return float(math.exp(self._LOG_SIGMA_MAX - s * (self._LOG_SIGMA_MAX - self._LOG_SIGMA_MIN)))
+
+    def verify_input(self, batch, fastvideo_args):
+        return VerificationResult()
+
+    def verify_output(self, batch, fastvideo_args):
+        return VerificationResult()
+
+    @torch.inference_mode()
+    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
+        pc = fastvideo_args.pipeline_config
+        ext = batch.extra
+        device = batch.latents.device
+        guidance_scale = float(batch.guidance_scale or pc.guidance_scale)
+        steps = int(batch.num_inference_steps)
+
+        import k_diffusion as K
+
+        init_latent = ext.get("init_latent")
+        sigma_max = self._resolve_sigma_max(batch) if init_latent is not None else self._SIGMA_MAX
+
+        sigmas = K.sampling.get_sigmas_polyexponential(steps, self._SIGMA_MIN, sigma_max, self._RHO, device=device)
+        # Cast noise + conditioning to the DiT's dtype before sampling
+        # (matches `stable_audio_tools/inference/generation.py:185-187`).
+        model_dtype = next(self.transformer.parameters()).dtype
+
+        def _cast(t: torch.Tensor | None) -> torch.Tensor | None:
+            return t.to(model_dtype) if t is not None else None
+
+        x = (batch.latents * sigmas[0]).to(model_dtype)
+        if init_latent is not None:
+            x = x + init_latent.to(model_dtype)
+
+        batch_cond, batch_global = _build_cfg_conditioning(
+            cross_attn_cond=ext["cross_attn_cond"].to(model_dtype),
+            global_embed=ext["global_embed"].to(model_dtype),
+            negative_cross_attn_cond=_cast(ext.get("negative_cross_attn_cond")),
+            negative_cross_attn_mask=ext.get("negative_cross_attn_mask"),
+            negative_global_embed=_cast(ext.get("negative_global_embed")),
+            do_cfg=guidance_scale != 1.0,
+        )
+        adapter = _DiTAdapter(self.transformer,
+                              batch_cond=batch_cond,
+                              batch_global=batch_global,
+                              cfg_scale=guidance_scale)
+        denoiser = K.external.VDenoiser(adapter)
+
+        # RePaint blending hook — works on any v-prediction model, no
+        # inpaint-trained checkpoint needed.
+        inpaint_mask = ext.get("inpaint_mask_latent")
+        inpaint_ref = ext.get("inpaint_reference_latent")
+        if inpaint_mask is not None and inpaint_ref is not None:
+            inpaint_mask = inpaint_mask.to(model_dtype)
+            inpaint_ref = inpaint_ref.to(model_dtype)
+            callback = _make_inpaint_callback(inpaint_ref, inpaint_mask, sigmas)
+        else:
+            callback = None
+
+        # `LocalAttention` (in `StableAudioDiT`) reads `get_forward_context()`
+        # for `attn_metadata`; wrap the whole loop.
+        with set_forward_context(current_timestep=0, attn_metadata=None):
+            sampled = K.sampling.sample_dpmpp_3m_sde(denoiser,
+                                                     x,
+                                                     sigmas,
+                                                     disable=False,
+                                                     extra_args={},
+                                                     callback=callback)
+
+        # Final blend so the kept region of the inpaint reference is exact.
+        if inpaint_mask is not None and inpaint_ref is not None:
+            sampled = inpaint_ref * inpaint_mask + sampled * (1 - inpaint_mask)
+        batch.latents = sampled
+        return batch
+
+
+def _build_cfg_conditioning(
+    *,
+    cross_attn_cond: torch.Tensor,
+    global_embed: torch.Tensor,
+    negative_cross_attn_cond: torch.Tensor | None,
+    negative_cross_attn_mask: torch.Tensor | None,
+    negative_global_embed: torch.Tensor | None,
+    do_cfg: bool,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Build the CFG-batched `(cond, global)` tensors once.
+
+    Cond ordering is `[conditioned, unconditioned]` (the adapter splits
+    with the same convention). Masked negative cond is zero-filled where
+    `mask == 0`.
+    """
+    if not do_cfg:
+        return cross_attn_cond, global_embed
+    if negative_cross_attn_cond is not None:
+        if negative_cross_attn_mask is not None:
+            neg_mask = negative_cross_attn_mask.to(torch.bool).unsqueeze(2)
+            null_embed = torch.zeros_like(cross_attn_cond)
+            negative_cross_attn_cond = torch.where(neg_mask, negative_cross_attn_cond, null_embed)
+        batch_cond = torch.cat([cross_attn_cond, negative_cross_attn_cond], dim=0)
+    else:
+        batch_cond = torch.cat([cross_attn_cond, torch.zeros_like(cross_attn_cond)], dim=0)
+    other_global = global_embed if negative_global_embed is None else negative_global_embed
+    batch_global = torch.cat([global_embed, other_global], dim=0)
+    return batch_cond, batch_global
+
+
+def _make_inpaint_callback(reference_latent: torch.Tensor, mask: torch.Tensor, sigmas: torch.Tensor):
+    """RePaint blending callback for the k-diffusion sampler.
+
+    At every step, replaces the kept region (`mask == 1`) of the in-
+    flight latent with the reference re-noised to the next sigma —
+    pulls the kept region back onto the trajectory the model expects,
+    so RePaint-style inpainting converges on non-inpaint-trained models.
+
+    Pre-allocates the noise buffer so the ~100 sampler steps don't churn
+    ~25 MB of fresh allocations per call.
+    """
+    noise_buf = torch.empty_like(reference_latent)
+    inv_mask = 1 - mask
+
+    def cb(info: dict) -> None:
+        i = int(info["i"])
+        next_i = min(i + 1, len(sigmas) - 1)
+        sigma_next = float(sigmas[next_i])
+        noise_buf.normal_()
+        # `state["x"]` is the live latent; the dpmpp-3m-sde sampler picks
+        # up our in-place mutation between steps (verified against
+        # k_diffusion 0.1.1.post1).
+        x = info["x"]
+        x.copy_((reference_latent + noise_buf * sigma_next) * mask + x * inv_mask)
+
+    return cb

fastvideo/pipelines/basic/stable_audio/stages/latent_preparation.py

@@ -0,0 +1,174 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Stable Audio latent preparation.
+
+Seeds + samples the initial Gaussian noise; encodes `init_audio` (A2A
+variation) or `inpaint_audio` + `inpaint_mask` (RePaint inpainting) into
+latent-space tensors on `batch.extra` for the denoising stage.
+"""
+from __future__ import annotations
+
+import os
+
+import torch
+import torch.nn.functional as F
+
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+from fastvideo.pipelines.stages.base import PipelineStage
+from fastvideo.pipelines.stages.validators import VerificationResult
+
+
+class StableAudioLatentPreparationStage(PipelineStage):
+
+    def __init__(self,
+                 io_channels: int = 64,
+                 sample_size: int = 2097152,
+                 vae=None,
+                 sample_rate: int = 44100,
+                 audio_channels: int = 2) -> None:
+        super().__init__()
+        self.io_channels = io_channels
+        # Audio-domain length the model was trained for; latent length
+        # = sample_size // vae.hop_length (= 2097152 / 2048 = 1024).
+        self.sample_size = sample_size
+        self.vae = vae  # used to encode init_audio / inpaint_audio
+        self.sample_rate = sample_rate
+        self.audio_channels = audio_channels
+
+    def verify_input(self, batch, fastvideo_args):
+        return VerificationResult()
+
+    def verify_output(self, batch, fastvideo_args):
+        return VerificationResult()
+
+    @torch.inference_mode()
+    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
+        ext = batch.extra or {}
+        device = ext["cross_attn_cond"].device
+        latent_sample_size = self.sample_size // self._hop_length()
+
+        seed = int(batch.seed) if batch.seed is not None else 0
+        torch.manual_seed(seed)
+        latents = torch.randn((1, self.io_channels, latent_sample_size), device=device)
+
+        batch.latents = latents
+        if batch.extra is None:
+            batch.extra = {}
+
+        init_audio = getattr(batch, "init_audio", None)
+        inpaint_audio = getattr(batch, "inpaint_audio", None)
+        inpaint_mask = getattr(batch, "inpaint_mask", None)
+
+        # Loud-fail rather than silently falling through to T2A.
+        if inpaint_audio is not None and inpaint_mask is None:
+            raise ValueError("Stable Audio inpainting requires both `inpaint_audio` and "
+                             "`inpaint_mask` (1-D tensor in {0, 1} at the model sample rate, "
+                             "1 = keep, 0 = regenerate). Got `inpaint_audio` without `inpaint_mask`.")
+        if inpaint_mask is not None and inpaint_audio is None:
+            raise ValueError("Stable Audio inpainting requires both `inpaint_audio` and "
+                             "`inpaint_mask`. Got `inpaint_mask` without `inpaint_audio` — "
+                             "did you mean to pass `init_audio` (audio-to-audio variation)?")
+        if init_audio is not None and inpaint_audio is not None:
+            raise ValueError("Stable Audio cannot do A2A variation and inpainting in the "
+                             "same call. Pass either `init_audio` (variation) or "
+                             "`inpaint_audio` + `inpaint_mask` (inpainting), not both.")
+
+        if init_audio is not None:
+            batch.extra["init_latent"] = self._encode_audio_reference(init_audio, device)
+
+        if inpaint_audio is not None and inpaint_mask is not None:
+            batch.extra["inpaint_reference_latent"] = self._encode_audio_reference(inpaint_audio, device)
+            batch.extra["inpaint_mask_latent"] = self._prepare_mask(inpaint_mask, latent_sample_size, device)
+        return batch
+
+    def _hop_length(self) -> int:
+        return int(self.vae.hop_length)
+
+    def _encode_audio_reference(self, audio, device: torch.device) -> torch.Tensor:
+        """Pad/truncate to `sample_size` and encode via the VAE.
+
+        `audio` may be a tensor (`[samples]`, `[C, samples]`, or
+        `[B, C, samples]`) at the model's sample rate, or a path to any
+        audio-bearing file (`.wav` / `.mp3` / `.mp4` / `.m4a` / `.flac`,
+        ...) that PyAV can decode — we resample on load so callers don't
+        have to.
+        """
+        assert self.vae is not None, "VAE required for init_audio / inpaint_audio encoding"
+        # VAE may be CPU-parked under `vae_cpu_offload=True`.
+        self.vae = self.vae.to(device)
+        if isinstance(audio, str | os.PathLike):
+            audio = _decode_audio_file(audio, target_sr=self.sample_rate)
+        audio = audio.to(device=device, dtype=torch.float32)
+        if audio.dim() == 1:
+            audio = audio.unsqueeze(0).unsqueeze(0)
+        elif audio.dim() == 2:
+            audio = audio.unsqueeze(0)
+        # Match expected channel count (mono → repeat to stereo).
+        if audio.shape[1] == 1 and self.audio_channels == 2:
+            audio = audio.repeat(1, 2, 1)
+        elif audio.shape[1] == 2 and self.audio_channels == 1:
+            audio = audio.mean(dim=1, keepdim=True)
+        # Pad/truncate to model sample_size.
+        cur_len = audio.shape[-1]
+        if cur_len < self.sample_size:
+            audio = F.pad(audio, (0, self.sample_size - cur_len))
+        elif cur_len > self.sample_size:
+            audio = audio[..., :self.sample_size]
+        # Stochastic sample (the next random draw after the latent
+        # `randn` above), so encode-noise stays on the seeded sequence.
+        return self.vae.encode(audio.to(next(self.vae.parameters()).dtype)).sample()
+
+    def _prepare_mask(self, mask, latent_len: int, device: torch.device) -> torch.Tensor:
+        """Pad/truncate a binary mask to `sample_size`, then
+        nearest-resample to `[1, 1, latent_len]`. Convention: 1 = keep
+        the reference, 0 = regenerate.
+
+        `mask` may be a `[samples]` tensor at the model sample rate or a
+        `(keep_seconds, total_seconds)` tuple — the tuple form builds
+        "keep first K seconds, regenerate the rest" automatically.
+        """
+        if isinstance(mask, tuple) and len(mask) == 2:
+            keep_s, total_s = (float(x) for x in mask)
+            keep_n = int(keep_s * self.sample_rate)
+            total_n = int(total_s * self.sample_rate)
+            mask = torch.zeros(total_n, dtype=torch.float32)
+            mask[:keep_n] = 1.0
+        m = mask.to(device=device, dtype=torch.float32)
+        if m.dim() == 1:
+            m = m.unsqueeze(0)
+        cur_len = m.shape[-1]
+        if cur_len < self.sample_size:
+            m = F.pad(m, (0, self.sample_size - cur_len))
+        elif cur_len > self.sample_size:
+            m = m[..., :self.sample_size]
+        return F.interpolate(m.unsqueeze(1), size=latent_len, mode="nearest")
+
+
+def _decode_audio_file(path, target_sr: int) -> torch.Tensor:
+    """Decode any audio-bearing file (wav, mp3, mp4, m4a, flac, ...) via
+    PyAV and resample to `target_sr`. Returns `[channels, samples]`
+    float32 in roughly [-1, 1].
+
+    PyAV is already a FastVideo dep (used for muxing in
+    `VideoGenerator._mux_audio`). `torchaudio.load` on container
+    formats (mp4 / m4a) routes through `torchcodec`, which pulls in a
+    full CUDA NVRTC stack we don't otherwise need.
+    """
+    import av
+    import numpy as np
+    container = av.open(str(path))
+    audio_stream = next(s for s in container.streams if s.type == "audio")
+    resampler = av.AudioResampler(format="fltp", layout="stereo", rate=target_sr)
+    chunks: list = []
+    for frame in container.decode(audio_stream):
+        for resampled in resampler.resample(frame):
+            chunks.append(resampled.to_ndarray())
+    for resampled in resampler.resample(None):
+        chunks.append(resampled.to_ndarray())
+    container.close()
+    if not chunks:
+        raise RuntimeError(f"No audio frames decoded from {path}")
+    waveform = np.concatenate(chunks, axis=-1)
+    if waveform.ndim == 1:
+        waveform = waveform[None, :]
+    return torch.from_numpy(waveform).float()

fastvideo/pipelines/pipeline_batch_info.py

@@ -192,6 +192,21 @@ class ForwardBatch:
     ltx2_stg_blocks_video: list[int] = field(default_factory=list)
     ltx2_stg_blocks_audio: list[int] = field(default_factory=list)
 
+    # Stable Audio (T2A): clip start/end in seconds. Parallels the
+    # `SamplingParam` fields of the same name; the
+    # `StableAudioConditioningStage` / `DecodingStage` read them.
+    audio_start_in_s: float | None = None
+    audio_end_in_s: float | None = None
+
+    # Stable Audio A2A variation + inpainting payloads (parallel to
+    # `SamplingParam`). `Any` because we accept torch tensors or numpy
+    # arrays the user supplies; the latent-prep stage normalises shapes.
+    init_audio: Any = None
+    init_audio_strength: float | None = None
+    init_noise_level: float | None = None
+    inpaint_audio: Any = None
+    inpaint_mask: Any = None
+
     n_tokens: int | None = None
 
     # Other parameters that may be needed by specific schedulers

fastvideo/pipelines/preprocess/preprocess_cosmos25_overfit.py

@@ -0,0 +1,193 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Preprocess Cosmos 2.5 overfit data into parquet format.
+
+Encodes videos with the Cosmos (Wan-style) VAE and captions with the
+Reason1 (Qwen2.5-VL) text encoder into the t2v parquet schema.
+
+Usage:
+    CUDA_VISIBLE_DEVICES=0 python fastvideo/pipelines/preprocess/preprocess_cosmos25_overfit.py
+"""
+
+import json
+import os
+
+import cv2
+import numpy as np
+import pyarrow as pa
+import pyarrow.parquet as pq
+import torch
+
+from fastvideo.dataset.dataloader.schema import pyarrow_schema_t2v
+from fastvideo.utils import maybe_download_model
+
+# --- Config ---
+NUM_FRAMES = 93  # 4*23+1 → 24 latent frames
+MAX_HEIGHT = 480
+MAX_WIDTH = 832
+TRAIN_FPS = 16.0
+
+DATA_DIR = "data/cosmos_overfit"
+OUTPUT_DIR = "data/cosmos25_overfit_preprocessed"
+MODEL_REPO = "KyleShao/Cosmos-Predict2.5-2B-Diffusers"
+# The VAE is architecturally identical to Cosmos Predict2;
+# use the Predict2 model for VAE since its weights are in
+# standard diffusers format.
+VAE_REPO = "nvidia/Cosmos-Predict2-2B-Video2World"
+
+
+def load_video(path: str, num_frames: int) -> torch.Tensor:
+    """Load video as [1, C, T, H, W] in [-1, 1]."""
+    cap = cv2.VideoCapture(path)
+    frames: list[np.ndarray] = []
+    while len(frames) < num_frames:
+        ret, frame = cap.read()
+        if not ret:
+            break
+        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+        frames.append(frame)
+    cap.release()
+
+    if len(frames) < num_frames:
+        while len(frames) < num_frames:
+            frames.append(frames[-1])
+
+    frames = frames[:num_frames]
+    video = np.stack(frames, axis=0)
+    video = torch.from_numpy(video).float()
+    video = video / 127.5 - 1.0  # [0,255] -> [-1,1]
+    video = video.permute(3, 0, 1, 2).unsqueeze(0)  # [1,C,T,H,W]
+    return video
+
+
+def main() -> None:
+    device = torch.device("cuda:0")
+    model_path = maybe_download_model(MODEL_REPO)
+
+    os.makedirs(OUTPUT_DIR, exist_ok=True)
+
+    # Load captions
+    with open(os.path.join(DATA_DIR, "videos2caption.json")) as f:
+        caption_data = json.load(f)
+
+    # --- Load VAE (Wan-style, same arch for Cosmos 2 and 2.5) ---
+    print("Loading Cosmos VAE (AutoencoderKLWan)...")
+    vae_path = maybe_download_model(VAE_REPO)
+    from diffusers import AutoencoderKLWan
+    vae = AutoencoderKLWan.from_pretrained(
+        vae_path,
+        subfolder="vae",
+        torch_dtype=torch.float16,
+    ).to(device).eval()
+    print(f"VAE loaded "
+          f"({sum(p.numel() for p in vae.parameters())/1e6:.0f}M)")
+
+    # --- Load Reason1 (Qwen2.5-VL) text encoder ---
+    print("Loading Reason1 text encoder...")
+    from fastvideo.configs.pipelines.cosmos2_5 import (
+        Cosmos25Config, )
+    from fastvideo.models.encoders.reason1 import (
+        Reason1TextEncoder, )
+    pipeline_cfg = Cosmos25Config()
+    text_enc_cfg = pipeline_cfg.text_encoder_configs[0]
+    text_enc_path = os.path.join(model_path, "text_encoder")
+
+    # Instantiate Reason1TextEncoder with config and checkpoint
+    text_encoder = Reason1TextEncoder(
+        text_enc_cfg,
+        checkpoint_path=text_enc_path,
+    )
+    # Load weights from safetensors into the meta-device model.
+    # Materialize empty tensors in bf16 on the target device,
+    # then overwrite with checkpoint weights.
+    text_encoder = text_encoder.to_empty(device=device)
+    text_encoder = text_encoder.to(torch.bfloat16)
+    import glob
+    from safetensors.torch import load_file
+    sd: dict[str, torch.Tensor] = {}
+    for sf in sorted(glob.glob(os.path.join(text_enc_path, "*.safetensors"))):
+        sd.update(load_file(sf, device=str(device)))
+    sd = {k: v.to(torch.bfloat16) for k, v in sd.items()}
+    text_encoder.load_state_dict(sd, strict=False, assign=True)
+    del sd
+    torch.cuda.empty_cache()
+    text_encoder = text_encoder.eval()
+    print("Reason1 text encoder loaded")
+
+    # --- Process each video ---
+    records = []
+    for idx, item in enumerate(caption_data):
+        video_name = item["path"]
+        record_id = f"{idx:04d}_{video_name}"
+        caption = item["cap"][0]
+        video_path = os.path.join(DATA_DIR, "videos", video_name)
+
+        print(f"\nProcessing: {video_name}")
+        print(f"  Caption: {caption[:80]}...")
+
+        # Encode video
+        video = load_video(video_path, NUM_FRAMES).to(device=device, dtype=torch.float16)
+        print(f"  Video shape: {video.shape}")
+
+        with torch.no_grad():
+            latent_dist = vae.encode(video).latent_dist
+            latent = latent_dist.mean.squeeze(0).float().cpu()
+        print(f"  Latent shape: {latent.shape}")
+
+        # Encode text with Reason1 (Qwen2.5-VL)
+        with torch.no_grad():
+            text_embedding = text_encoder.compute_text_embeddings(
+                [caption],
+                device=device,
+            )
+            text_embedding = text_embedding.squeeze(0).float().cpu()
+        print(f"  Text embedding shape: {text_embedding.shape}")
+
+        record = {
+            "id": record_id,
+            "vae_latent_bytes": latent.numpy().tobytes(),
+            "vae_latent_shape": list(latent.shape),
+            "vae_latent_dtype": str(latent.dtype).replace("torch.", ""),
+            "text_embedding_bytes": (text_embedding.numpy().tobytes()),
+            "text_embedding_shape": list(text_embedding.shape),
+            "text_embedding_dtype": str(text_embedding.dtype).replace("torch.", ""),
+            "file_name": video_name,
+            "caption": caption,
+            "media_type": "video",
+            "width": MAX_WIDTH,
+            "height": MAX_HEIGHT,
+            "num_frames": NUM_FRAMES,
+            "duration_sec": NUM_FRAMES / TRAIN_FPS,
+            "fps": TRAIN_FPS,
+        }
+        records.append(record)
+
+    # Clean up
+    del text_encoder, vae
+    torch.cuda.empty_cache()
+
+    # Write parquet
+    table = pa.table(
+        {k: [r[k] for r in records]
+         for k in records[0]},
+        schema=pyarrow_schema_t2v,
+    )
+    output_path = os.path.join(OUTPUT_DIR, "data_00000.parquet")
+    pq.write_table(table, output_path)
+    print(f"\nWrote {len(records)} records to {output_path}")
+
+    # Write T2W validation prompts (no image_path for T2W)
+    val_prompts = {
+        "data": [{
+            "caption": item["cap"][0],
+        } for item in caption_data],
+    }
+    val_path = os.path.join(OUTPUT_DIR, "validation_prompts.json")
+    with open(val_path, "w") as f:
+        json.dump(val_prompts, f, indent=2)
+    print(f"Wrote validation prompts to {val_path}")
+
+    print("\nDone! Use data_path: " + OUTPUT_DIR + " in training config.")
+
+
+if __name__ == "__main__":
+    main()

fastvideo/pipelines/preprocess/preprocess_cosmos_overfit.py

@@ -0,0 +1,196 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Preprocess Cosmos-Predict2 overfit data into parquet format.
+
+Encodes videos with the Cosmos (Wan-style) VAE and captions with the
+single T5 Large text encoder into the t2v parquet schema expected by
+the training framework.
+
+Usage:
+    CUDA_VISIBLE_DEVICES=0 python fastvideo/pipelines/preprocess/preprocess_cosmos_overfit.py
+"""
+
+import json
+import os
+
+import cv2
+import numpy as np
+import pyarrow as pa
+import pyarrow.parquet as pq
+import torch
+
+from fastvideo.configs.models.encoders import T5LargeConfig
+from fastvideo.configs.models.encoders.base import BaseEncoderOutput
+from fastvideo.configs.pipelines.cosmos import t5_large_postprocess_text
+from fastvideo.dataset.dataloader.schema import pyarrow_schema_t2v
+from fastvideo.utils import maybe_download_model
+
+# --- Config ---
+NUM_FRAMES = 93  # 4*23+1 for temporal compression ratio 4 -> 24 latent frames
+MAX_HEIGHT = 480
+MAX_WIDTH = 832
+TRAIN_FPS = 16.0
+
+DATA_DIR = "data/cosmos_overfit"
+OUTPUT_DIR = "data/cosmos_overfit_preprocessed"
+MODEL_REPO = "nvidia/Cosmos-Predict2-2B-Video2World"
+
+
+def load_video(path: str, num_frames: int) -> torch.Tensor:
+    """Load video as [1, C, T, H, W] in [-1, 1]."""
+    cap = cv2.VideoCapture(path)
+    frames: list[np.ndarray] = []
+    while len(frames) < num_frames:
+        ret, frame = cap.read()
+        if not ret:
+            break
+        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+        frames.append(frame)
+    cap.release()
+
+    if len(frames) < num_frames:
+        # Repeat last frame to fill
+        while len(frames) < num_frames:
+            frames.append(frames[-1])
+
+    frames = frames[:num_frames]
+    video = np.stack(frames, axis=0)
+    video = torch.from_numpy(video).float()
+    video = video / 127.5 - 1.0  # [0,255] -> [-1,1]
+    video = video.permute(3, 0, 1, 2).unsqueeze(0)  # [1,C,T,H,W]
+    return video
+
+
+def main() -> None:
+    device = torch.device("cuda:0")
+    model_path = maybe_download_model(MODEL_REPO)
+
+    os.makedirs(OUTPUT_DIR, exist_ok=True)
+
+    # Load captions
+    with open(os.path.join(DATA_DIR, "videos2caption.json")) as f:
+        caption_data = json.load(f)
+
+    # --- Load VAE ---
+    # Cosmos-Predict2-2B ships a Wan-style VAE; vae/config.json declares
+    # `_class_name: AutoencoderKLWan`, so diffusers can load it directly.
+    print("Loading Cosmos VAE (AutoencoderKLWan)...")
+    from diffusers import AutoencoderKLWan
+    vae = AutoencoderKLWan.from_pretrained(
+        model_path,
+        subfolder="vae",
+        torch_dtype=torch.float16,
+    ).to(device).eval()
+    print(f"VAE loaded ({sum(p.numel() for p in vae.parameters())/1e6:.0f}M)")
+
+    # --- Load T5 Large text encoder ---
+    print("Loading T5 Large text encoder...")
+    from transformers import AutoTokenizer, T5EncoderModel
+
+    t5_cfg = T5LargeConfig()
+    tok_kwargs = dict(t5_cfg.tokenizer_kwargs)
+    tokenizer = AutoTokenizer.from_pretrained(os.path.join(model_path, "tokenizer"))
+    text_encoder = T5EncoderModel.from_pretrained(
+        os.path.join(model_path, "text_encoder"),
+        torch_dtype=torch.bfloat16,
+    ).to(device).eval()
+
+    # --- Process each video ---
+    records = []
+    for idx, item in enumerate(caption_data):
+        video_name = item["path"]
+        record_id = f"{idx:04d}_{video_name}"
+        caption = item["cap"][0]
+        video_path = os.path.join(DATA_DIR, "videos", video_name)
+
+        print(f"\nProcessing: {video_name}")
+        print(f"  Caption: {caption[:80]}...")
+
+        # Encode video
+        video = load_video(video_path, NUM_FRAMES).to(device=device, dtype=torch.float16)
+        print(f"  Video shape: {video.shape}")
+
+        with torch.no_grad():
+            latent_dist = vae.encode(video).latent_dist
+            # Cast to fp32 — dataloader hardcodes np.float32
+            latent = latent_dist.mean.squeeze(0).float().cpu()
+        print(f"  Latent shape: {latent.shape}")
+
+        # Encode text with T5
+        with torch.no_grad():
+            inputs = tokenizer(caption, **tok_kwargs).to(device)
+            outputs = text_encoder(**inputs)
+            enc_out = BaseEncoderOutput(
+                last_hidden_state=outputs.last_hidden_state,
+                attention_mask=inputs["attention_mask"],
+            )
+            # [1, max_len, 1024], zeros beyond real length
+            t5_embed = t5_large_postprocess_text(enc_out).squeeze(0)
+            # Trim to real sequence length so dataloader's pad() builds
+            # the correct attention mask.
+            real_len = int(inputs["attention_mask"].sum().item())
+            text_embedding = t5_embed[:real_len].float().cpu()  # [seq, 1024]
+        print(f"  Text embedding shape: {text_embedding.shape}")
+
+        record = {
+            "id": record_id,
+            "vae_latent_bytes": latent.numpy().tobytes(),
+            "vae_latent_shape": list(latent.shape),
+            "vae_latent_dtype": str(latent.dtype).replace("torch.", ""),
+            "text_embedding_bytes": text_embedding.numpy().tobytes(),
+            "text_embedding_shape": list(text_embedding.shape),
+            "text_embedding_dtype": str(text_embedding.dtype).replace("torch.", ""),
+            "file_name": video_name,
+            "caption": caption,
+            "media_type": "video",
+            "width": MAX_WIDTH,
+            "height": MAX_HEIGHT,
+            "num_frames": NUM_FRAMES,
+            "duration_sec": NUM_FRAMES / TRAIN_FPS,
+            "fps": TRAIN_FPS,
+        }
+        records.append(record)
+
+    # Clean up encoders
+    del text_encoder, tokenizer, vae
+
+    # Write parquet
+    table = pa.table(
+        {k: [r[k] for r in records]
+         for k in records[0]},
+        schema=pyarrow_schema_t2v,
+    )
+    output_path = os.path.join(OUTPUT_DIR, "data_00000.parquet")
+    pq.write_table(table, output_path)
+    print(f"\nWrote {len(records)} records to {output_path}")
+
+    # Extract first frame from first video as V2W conditioning image
+    import cv2
+    first_video = os.path.join(DATA_DIR, "videos", caption_data[0]["path"])
+    cap = cv2.VideoCapture(first_video)
+    ret, frame = cap.read()
+    cap.release()
+    cond_frame_path = os.path.join(OUTPUT_DIR, "cond_frame.png")
+    if ret:
+        cv2.imwrite(cond_frame_path, frame)
+        print(f"Saved conditioning frame to {cond_frame_path}")
+
+    # Write validation prompts for callback
+    # Wrap in "data" key — ValidationDataset expects field="data"
+    # Use "caption" field — ValidationDataset aliases it to "prompt"
+    # Include image_path for V2W conditioning during validation
+    val_prompts = {
+        "data": [{
+            "caption": item["cap"][0],
+            "image_path": "cond_frame.png",
+        } for item in caption_data]
+    }
+    val_path = os.path.join(OUTPUT_DIR, "validation_prompts.json")
+    with open(val_path, "w") as f:
+        json.dump(val_prompts, f, indent=2)
+    print(f"Wrote validation prompts to {val_path}")
+
+    print("\nDone! Use data_path: " + OUTPUT_DIR + " in training config.")
+
+
+if __name__ == "__main__":
+    main()

fastvideo/pipelines/stages/denoising.py

@@ -518,13 +518,42 @@ class DenoisingStage(PipelineStage):
 
 
 class CosmosDenoisingStage(DenoisingStage):
-    """
-    Denoising stage for Cosmos models using FlowMatchEulerDiscreteScheduler.
+    """Denoising stage for Cosmos models.
+
+    Uses FlowMatchEulerDiscreteScheduler with manual EDM
+    preconditioning (c_in, c_skip, c_out) to match the
+    pretrained Cosmos model's training convention.
     """
 
     def __init__(self, transformer, scheduler, pipeline=None) -> None:
         super().__init__(transformer, scheduler, pipeline)
 
+    def _run_transformer(
+        self,
+        hidden_states: torch.Tensor,
+        timestep: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        condition_mask: torch.Tensor,
+        padding_mask: torch.Tensor,
+        target_dtype: torch.dtype,
+        step_index: int,
+        batch: ForwardBatch,
+    ) -> torch.Tensor:
+        with set_forward_context(
+                current_timestep=step_index,
+                attn_metadata=None,
+                forward_batch=batch,
+        ):
+            return self.transformer(
+                hidden_states=hidden_states.to(target_dtype),
+                timestep=timestep.to(target_dtype),
+                encoder_hidden_states=encoder_hidden_states.to(target_dtype),
+                fps=24,
+                condition_mask=condition_mask,
+                padding_mask=padding_mask,
+                return_dict=False,
+            )[0]
+
     def forward(
         self,
         batch: ForwardBatch,
@@ -533,199 +562,188 @@ class CosmosDenoisingStage(DenoisingStage):
         pipeline = self.pipeline() if self.pipeline else None
         if not fastvideo_args.model_loaded["transformer"]:
             loader = TransformerLoader()
-            self.transformer = loader.load(fastvideo_args.model_paths["transformer"], fastvideo_args)
+            self.transformer = loader.load(
+                fastvideo_args.model_paths["transformer"],
+                fastvideo_args,
+            )
             if pipeline:
                 pipeline.add_module("transformer", self.transformer)
             fastvideo_args.model_loaded["transformer"] = True
 
-        extra_step_kwargs = self.prepare_extra_func_kwargs(
-            self.scheduler.step,
-            {
-                "generator": batch.generator,
-                "eta": batch.eta
-            },
-        )
-
-        if hasattr(self.transformer, 'module'):
+        if hasattr(self.transformer, "module"):
             transformer_dtype = next(self.transformer.module.parameters()).dtype
         else:
             transformer_dtype = next(self.transformer.parameters()).dtype
         target_dtype = transformer_dtype
-        autocast_enabled = (target_dtype != torch.float32) and not fastvideo_args.disable_autocast
+        autocast_enabled = (target_dtype != torch.float32 and not fastvideo_args.disable_autocast)
 
         latents = batch.latents
         num_inference_steps = batch.num_inference_steps
         guidance_scale = batch.guidance_scale
+        do_cfg = (batch.do_classifier_free_guidance and batch.negative_prompt_embeds is not None)
 
-        sigma_max = 80.0
-        sigma_min = 0.002
-        sigma_data = 1.0
-        final_sigmas_type = "sigma_min"
+        sigma_data = float(getattr(self.scheduler.config, "sigma_data", 1.0))
 
-        if self.scheduler is not None:
-            self.scheduler.register_to_config(
-                sigma_max=sigma_max,
-                sigma_min=sigma_min,
-                sigma_data=sigma_data,
-                final_sigmas_type=final_sigmas_type,
-            )
-
-        self.scheduler.set_timesteps(num_inference_steps, device=latents.device)
+        self.scheduler.set_timesteps(
+            num_inference_steps,
+            device=latents.device,
+        )
         timesteps = self.scheduler.timesteps
 
-        if (hasattr(self.scheduler.config, 'final_sigmas_type')
+        # Clamp terminal sigma to sigma_min (avoid zero).
+        if (hasattr(self.scheduler.config, "final_sigmas_type")
                 and self.scheduler.config.final_sigmas_type == "sigma_min" and len(self.scheduler.sigmas) > 1):
             self.scheduler.sigmas[-1] = self.scheduler.sigmas[-2]
 
-        conditioning_latents = getattr(batch, 'conditioning_latents', None)
-        unconditioning_latents = conditioning_latents
+        conditioning_latents = getattr(
+            batch,
+            "conditioning_latents",
+            None,
+        )
+        cond_indicator = getattr(batch, "cond_indicator", None)
+        uncond_indicator = getattr(
+            batch,
+            "uncond_indicator",
+            None,
+        )
 
-        with self.progress_bar(total=num_inference_steps) as progress_bar:
+        augment_sigma = torch.tensor(
+            [0.001],
+            device=latents.device,
+            dtype=torch.float32,
+        )
+
+        padding_mask = torch.zeros(
+            1,
+            1,
+            batch.height,
+            batch.width,
+            device=latents.device,
+            dtype=target_dtype,
+        )
+
+        condition_mask = (batch.cond_mask.to(target_dtype)
+                          if hasattr(batch, "cond_mask") and batch.cond_mask is not None else None)
+        uncond_condition_mask = (batch.uncond_mask.to(target_dtype)
+                                 if hasattr(batch, "uncond_mask") and batch.uncond_mask is not None else condition_mask)
+        if condition_mask is None:
+            b, c, tf, h, w = latents.shape
+            condition_mask = torch.zeros(
+                b,
+                1,
+                tf,
+                h,
+                w,
+                device=latents.device,
+                dtype=target_dtype,
+            )
+            uncond_condition_mask = condition_mask
+
+        with self.progress_bar(total=num_inference_steps, ) as progress_bar:
             for i, t in enumerate(timesteps):
-                if hasattr(self, 'interrupt') and self.interrupt:
+                if hasattr(self, "interrupt") and self.interrupt:
                     continue
 
-                current_sigma = self.scheduler.sigmas[i]
-                current_t = current_sigma / (current_sigma + 1)
-                c_in = 1 - current_t
-                c_skip = 1 - current_t
-                c_out = -current_t
+                sigma = self.scheduler.sigmas[i]
+                is_aug_greater = bool(augment_sigma >= sigma)
 
-                timestep = current_t.view(1, 1, 1, 1, 1).expand(latents.size(0), -1, latents.size(2), -1,
-                                                                -1)  # [B, 1, T, 1, 1]
+                # EDM preconditioning coefficients.
+                c_in = 1.0 / (sigma**2 + sigma_data**2)**0.5
+                c_in_aug = 1.0 / (augment_sigma**2 + sigma_data**2)**0.5
+                c_skip = sigma_data**2 / (sigma**2 + sigma_data**2)
+                c_out = (sigma * sigma_data / (sigma**2 + sigma_data**2)**0.5)
 
-                with torch.autocast(device_type="cuda", dtype=target_dtype, enabled=autocast_enabled):
+                # The model expects timestep = sigma * 1000
+                # (FlowMatchEulerDiscreteScheduler convention).
+                timestep_expanded = t.expand(latents.shape[0], ).to(target_dtype)
 
-                    cond_latent = latents * c_in
+                with torch.autocast(
+                        device_type="cuda",
+                        dtype=target_dtype,
+                        enabled=autocast_enabled,
+                ):
+                    # --- Conditioning frame injection ---
+                    cur_ci = (cond_indicator * 0 if cond_indicator is not None and is_aug_greater else cond_indicator)
 
-                    if hasattr(
+                    cond_latent = latents.clone()
+                    if (cur_ci is not None and conditioning_latents is not None):
+                        cn = torch.randn_like(
+                            latents,
+                            dtype=torch.float32,
+                        )
+                        cf = (conditioning_latents + cn * augment_sigma[:, None, None, None, None])
+                        cf = cf * c_in_aug / c_in
+                        cond_latent = (cur_ci * cf + (1 - cur_ci) * cond_latent)
+
+                    # Manual EDM input scaling.
+                    model_input = cond_latent * c_in
+
+                    noise_pred_cond = self._run_transformer(
+                        model_input,
+                        timestep_expanded,
+                        batch.prompt_embeds[0],
+                        condition_mask,
+                        padding_mask,
+                        target_dtype,
+                        i,
+                        batch,
+                    )
+
+                    # EDM output → x0 prediction.
+                    cond_x0 = (c_skip * latents + c_out * noise_pred_cond.float())
+                    if (cur_ci is not None and conditioning_latents is not None):
+                        cond_x0 = (cur_ci * conditioning_latents + (1 - cur_ci) * cond_x0)
+
+                    # --- CFG: unconditional pass ---
+                    if do_cfg:
+                        cur_ui = (uncond_indicator *
+                                  0 if uncond_indicator is not None and is_aug_greater else uncond_indicator)
+
+                        uncond_latent = latents.clone()
+                        if (cur_ui is not None and conditioning_latents is not None):
+                            un = torch.randn_like(
+                                latents,
+                                dtype=torch.float32,
+                            )
+                            uf = (conditioning_latents + un * augment_sigma[:, None, None, None, None])
+                            uf = uf * c_in_aug / c_in
+                            uncond_latent = (cur_ui * uf + (1 - cur_ui) * uncond_latent)
+
+                        uncond_input = uncond_latent * c_in
+
+                        noise_pred_uncond = (self._run_transformer(
+                            uncond_input,
+                            timestep_expanded,
+                            batch.negative_prompt_embeds[0],
+                            uncond_condition_mask,
+                            padding_mask,
+                            target_dtype,
+                            i,
                             batch,
-                            'cond_indicator') and batch.cond_indicator is not None and conditioning_latents is not None:
-                        cond_latent = batch.cond_indicator * conditioning_latents + (1 -
-                                                                                     batch.cond_indicator) * cond_latent
+                        ))
+
+                        uncond_x0 = (c_skip * latents + c_out * noise_pred_uncond.float())
+                        if (cur_ui is not None and conditioning_latents is not None):
+                            uncond_x0 = (cur_ui * conditioning_latents + (1 - cur_ui) * uncond_x0)
+
+                        final_x0 = (cond_x0 + guidance_scale * (cond_x0 - uncond_x0))
                     else:
-                        logger.warning(
-                            "Step %s: Missing conditioning data - cond_indicator: %s, conditioning_latents: %s", i,
-                            hasattr(batch, 'cond_indicator'), conditioning_latents is not None)
+                        final_x0 = cond_x0
 
-                    cond_latent = cond_latent.to(target_dtype)
+                # Convert x0 to velocity for
+                # FlowMatchEulerDiscreteScheduler.
+                velocity = (latents - final_x0) / sigma.clamp(min=1e-6)
 
-                    cond_timestep = timestep
-                    if hasattr(batch, 'cond_indicator') and batch.cond_indicator is not None:
-                        sigma_conditioning = 0.0001
-                        t_conditioning = sigma_conditioning / (sigma_conditioning + 1)
-                        cond_timestep = batch.cond_indicator * t_conditioning + (1 - batch.cond_indicator) * timestep
-                        cond_timestep = cond_timestep.to(target_dtype)
-
-                    with set_forward_context(
-                            current_timestep=i,
-                            attn_metadata=None,
-                            forward_batch=batch,
-                    ):
-                        # Use conditioning masks from CosmosLatentPreparationStage
-                        condition_mask = batch.cond_mask.to(target_dtype) if hasattr(batch, 'cond_mask') else None
-                        padding_mask = torch.zeros(1,
-                                                   1,
-                                                   batch.height,
-                                                   batch.width,
-                                                   device=cond_latent.device,
-                                                   dtype=target_dtype)
-
-                        # Fallback if masks not available
-                        if condition_mask is None:
-                            batch_size, num_channels, num_frames, height, width = cond_latent.shape
-                            condition_mask = torch.zeros(batch_size,
-                                                         1,
-                                                         num_frames,
-                                                         height,
-                                                         width,
-                                                         device=cond_latent.device,
-                                                         dtype=target_dtype)
-
-                        noise_pred = self.transformer(
-                            hidden_states=cond_latent,
-                            timestep=cond_timestep.to(target_dtype),
-                            encoder_hidden_states=batch.prompt_embeds[0].to(target_dtype),
-                            fps=24,  # TODO: get fps from batch or config
-                            condition_mask=condition_mask,
-                            padding_mask=padding_mask,
-                            return_dict=False,
-                        )[0]
-
-                    cond_pred = (c_skip * latents + c_out * noise_pred.float()).to(target_dtype)
-
-                    if hasattr(
-                            batch,
-                            'cond_indicator') and batch.cond_indicator is not None and conditioning_latents is not None:
-                        cond_pred = batch.cond_indicator * conditioning_latents + (1 - batch.cond_indicator) * cond_pred
-
-                    if batch.do_classifier_free_guidance and batch.negative_prompt_embeds is not None:
-                        uncond_latent = latents * c_in
-
-                        if hasattr(batch, 'uncond_indicator'
-                                   ) and batch.uncond_indicator is not None and unconditioning_latents is not None:
-                            uncond_latent = batch.uncond_indicator * unconditioning_latents + (
-                                1 - batch.uncond_indicator) * uncond_latent
-
-                        with set_forward_context(
-                                current_timestep=i,
-                                attn_metadata=None,
-                                forward_batch=batch,
-                        ):
-                            uncond_condition_mask = batch.uncond_mask.to(target_dtype) if hasattr(
-                                batch, 'uncond_mask') and batch.uncond_mask is not None else condition_mask
-
-                            uncond_timestep = timestep
-                            if hasattr(batch, 'uncond_indicator') and batch.uncond_indicator is not None:
-                                sigma_conditioning = 0.0001
-                                t_conditioning = sigma_conditioning / (sigma_conditioning + 1)
-                                uncond_timestep = batch.uncond_indicator * t_conditioning + (
-                                    1 - batch.uncond_indicator) * timestep
-                                uncond_timestep = uncond_timestep.to(target_dtype)
-
-                            noise_pred_uncond = self.transformer(
-                                hidden_states=uncond_latent.to(target_dtype),
-                                timestep=uncond_timestep.to(target_dtype),
-                                encoder_hidden_states=batch.negative_prompt_embeds[0].to(target_dtype),
-                                fps=24,  # TODO: get fps from batch or config
-                                condition_mask=uncond_condition_mask,
-                                padding_mask=padding_mask,
-                                return_dict=False,
-                            )[0]
-
-                        uncond_pred = (c_skip * latents + c_out * noise_pred_uncond.float()).to(target_dtype)
-
-                        if hasattr(batch, 'uncond_indicator'
-                                   ) and batch.uncond_indicator is not None and unconditioning_latents is not None:
-                            uncond_pred = batch.uncond_indicator * unconditioning_latents + (
-                                1 - batch.uncond_indicator) * uncond_pred
-
-                        guidance_diff = cond_pred - uncond_pred
-                        final_pred = cond_pred + guidance_scale * guidance_diff
-                    else:
-                        final_pred = cond_pred
-
-                # Convert to noise for scheduler step
-                if current_sigma > 1e-8:
-                    noise_for_scheduler = (latents - final_pred) / current_sigma
-                else:
-                    logger.warning("Step %s: current_sigma too small (%s), using final_pred directly", i, current_sigma)
-                    noise_for_scheduler = final_pred
-
-                if torch.isnan(noise_for_scheduler).sum() > 0:
-                    logger.error("Step %s: NaN detected in noise_for_scheduler, sum: %s", i,
-                                 noise_for_scheduler.float().sum().item())
-                    logger.error("Step %s: latents sum: %s, final_pred sum: %s, current_sigma: %s", i,
-                                 latents.float().sum().item(),
-                                 final_pred.float().sum().item(), current_sigma)
-
-                latents = self.scheduler.step(noise_for_scheduler, t, latents, **extra_step_kwargs,
-                                              return_dict=False)[0]
+                latents = self.scheduler.step(
+                    velocity,
+                    t,
+                    latents,
+                    return_dict=False,
+                )[0]
 
                 progress_bar.update()
 
         batch.latents = latents
-
         return batch
 
     def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult:
@@ -771,16 +789,21 @@ class Cosmos25DenoisingStage(CosmosDenoisingStage):
             },
         )
 
-        if hasattr(self.transformer, 'module'):
-            transformer_dtype = next(self.transformer.module.parameters()).dtype
-        else:
-            transformer_dtype = next(self.transformer.parameters()).dtype
-        target_dtype = transformer_dtype
+        # Detect the actual weight dtype.  FSDP-wrapped models may
+        # report fp32 via next(parameters()) even when the physical
+        # weights are bf16.  Walk through parameters to find one
+        # that is NOT fp32 (the real checkpoint dtype).
+        target_dtype = torch.bfloat16  # safe default for Cosmos 2.5
+        for p in self.transformer.parameters():
+            if p.dtype != torch.float32:
+                target_dtype = p.dtype
+                break
         autocast_enabled = (target_dtype != torch.float32) and not fastvideo_args.disable_autocast
 
         latents = batch.latents
         if latents is None:
-            raise ValueError("latents must be provided for Cosmos25DenoisingStage")
+            raise ValueError("latents must be provided for "
+                             "Cosmos25DenoisingStage")
         guidance_scale = batch.guidance_scale
 
         if batch.timesteps is None:

fastvideo/registry.py

@@ -48,6 +48,7 @@ from fastvideo.configs.pipelines.wan import (
     WanT2V720PConfig,
 )
 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.fastvideo_args import WorkloadType
@@ -242,6 +243,47 @@ def _register_configs() -> None:
         default_preset="ltx2_distilled",
     )
 
+    # Stable Audio Open (text-to-audio). Both variants must be loaded
+    # from the FastVideo-curated converted Diffusers-format repos —
+    # the upstream `stabilityai/stable-audio-open-{1.0,small}` repos
+    # ship `model.safetensors` as a single monolithic checkpoint with
+    # no per-component subfolders our standard loader can consume. See
+    # `scripts/checkpoint_conversion/stable_audio_to_diffusers.py`.
+    # NOTE: WorkloadType has no T2A variant yet (REVIEW item 28); using
+    # T2V as the placeholder until the enum is extended.
+    register_configs(
+        sampling_param_cls=None,
+        pipeline_config_cls=StableAudioT2AConfig,
+        workload_types=(WorkloadType.T2V, ),
+        hf_model_paths=[
+            "FastVideo/stable-audio-open-1.0-Diffusers",
+        ],
+        # Substring match against HF cache snapshot paths (the lookup
+        # runs on the resolved local directory, which uses `--` between
+        # org and repo: `models--FastVideo--stable-audio-open-1.0-Diffusers`).
+        model_detectors=[
+            lambda path: "stable-audio-open-1" in path.lower(),
+        ],
+        model_family="stable_audio",
+        default_preset="stable_audio_open_1_0_base",
+    )
+    # Small variant uses its own `pipeline_config_cls` so it picks up
+    # the smaller (524288-sample) training window in `sample_size` /
+    # `max_audio_duration_s`.
+    register_configs(
+        sampling_param_cls=None,
+        pipeline_config_cls=StableAudioOpenSmallConfig,
+        workload_types=(WorkloadType.T2V, ),
+        hf_model_paths=[
+            "FastVideo/stable-audio-open-small-Diffusers",
+        ],
+        model_detectors=[
+            lambda path: "stable-audio-open-small" in path.lower(),
+        ],
+        model_family="stable_audio",
+        default_preset="stable_audio_open_small",
+    )
+
     # Hunyuan 1.5 (specific)
     register_configs(
         sampling_param_cls=None,
@@ -786,6 +828,8 @@ def _register_presets() -> None:
         ALL_PRESETS as MATRIXGAME_PRESETS, )
     from fastvideo.pipelines.basic.sd35.presets import (
         ALL_PRESETS as SD35_PRESETS, )
+    from fastvideo.pipelines.basic.stable_audio.presets import (
+        ALL_PRESETS as STABLE_AUDIO_PRESETS, )
     from fastvideo.pipelines.basic.turbodiffusion.presets import (
         ALL_PRESETS as TURBODIFFUSION_PRESETS, )
     from fastvideo.pipelines.basic.wan.presets import (
@@ -803,6 +847,7 @@ def _register_presets() -> None:
         LTX2_PRESETS,
         MATRIXGAME_PRESETS,
         SD35_PRESETS,
+        STABLE_AUDIO_PRESETS,
         TURBODIFFUSION_PRESETS,
         WAN_PRESETS,
     )

fastvideo/tests/modal/pr_test.py

@@ -24,6 +24,10 @@ image = (modal.Image.from_registry(
     os.environ.get("BUILDKITE_COMMIT", ""),
     "BUILDKITE_PULL_REQUEST":
     os.environ.get("BUILDKITE_PULL_REQUEST", ""),
+    "BUILDKITE_BRANCH":
+    os.environ.get("BUILDKITE_BRANCH", ""),
+    "TEST_SCOPE":
+    os.environ.get("TEST_SCOPE", ""),
     "IMAGE_VERSION":
     os.environ.get("IMAGE_VERSION", ""),
 }))
@@ -66,13 +70,21 @@ def run_test(pytest_command: str):
     {pytest_command}
     """
 
+    # result = subprocess.run(["/bin/bash", "-c", command],
+    #                         stdout=sys.stdout,
+    #                         stderr=sys.stderr,
+    #                         check=False)
+
+    # sys.exit(result.returncode)
+
     result = subprocess.run(["/bin/bash", "-c", command],
                             stdout=sys.stdout,
                             stderr=sys.stderr,
                             check=False)
 
-    sys.exit(result.returncode)
-
+    if result.returncode != 0:
+        raise RuntimeError(f"Test command failed with exit code {result.returncode}")
+    # On success, just return — don't call sys.exit()
 
 @app.function(gpu="H100:1",
               image=image,
@@ -206,7 +218,7 @@ def run_self_forcing_tests():
 @app.function(gpu="L40S:1", image=image, timeout=900)
 def run_unit_test():
     run_test(
-        "pytest ./fastvideo/tests/api/ ./fastvideo/tests/dataset/ ./fastvideo/tests/workflow/ ./fastvideo/tests/entrypoints/ --ignore=./fastvideo/tests/entrypoints/test_openai_api_integration.py -vs"
+        "pytest ./fastvideo/tests/api/ ./fastvideo/tests/dataset/ ./fastvideo/tests/workflow/ ./fastvideo/tests/entrypoints/ ./fastvideo/tests/train/ --ignore=./fastvideo/tests/entrypoints/test_openai_api_integration.py -vs"
     )
 
 
@@ -228,12 +240,20 @@ def run_lora_extraction_tests():
               timeout=1800,
               secrets=[
                   modal.Secret.from_dict(
-                      {"HF_API_KEY": os.environ.get("HF_API_KEY", "")})
+                      {"HF_API_KEY": os.environ.get("HF_API_KEY", ""),
+                       "HF_REPO_ID": "FastVideo/performance-tracking"})
               ],
-              volumes={"/root/data": model_vol})
+              volumes={
+                  "/root/data": model_vol,
+              })
 def run_performance_tests():
     run_test(
-        "export HF_HOME='/root/data/.cache' && hf auth login --token $HF_API_KEY && pytest ./fastvideo/tests/performance -vs"
+        "export HF_HOME='/root/data/.cache' && "
+        "export PERFORMANCE_TRACKING_ROOT='/tmp/perf-tracking' && "
+        "hf auth login --token $HF_API_KEY && "
+        "pytest ./fastvideo/tests/performance -vs && "
+        "python ./fastvideo/tests/performance/compare_baseline.py && "
+        "python ./fastvideo/tests/performance/dashboard.py"
     )
 
 

fastvideo/tests/performance/compare_baseline.py

@@ -0,0 +1,320 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Track performance results and compare against historical baseline.
+
+This script:
+1) reads current benchmark results from fastvideo/tests/performance/results,
+2) writes normalized tracking records to the Modal volume path,
+3) compares each current record against the mean of up to 5 prior records,
+4) exits non-zero if any metric regresses by more than 15%.
+"""
+
+import glob
+import json
+import os
+import re
+import statistics
+import sys
+from huggingface_hub import HfApi, snapshot_download
+from datetime import datetime, timezone
+from typing import Any
+
+from hf_store import sync_from_hf, upload_record, load_records_for_model, sanitize, safe_float
+
+# Use the env var passed by Modal, fallback to a default if needed
+HF_REPO_ID = os.environ.get("HF_REPO_ID", "FastVideo/performance-tracking")
+HF_TOKEN = os.environ.get("HF_API_KEY")
+
+RESULTS_DIR = os.path.join(
+    os.path.dirname(os.path.abspath(__file__)),
+    "results",
+)
+TRACKING_ROOT = os.environ.get(
+    "PERFORMANCE_TRACKING_ROOT",
+    "/tmp/perf-tracking",
+)
+MAX_REGRESSION = float(os.environ.get("PERF_MAX_REGRESSION", "0.05"))
+
+def _should_persist_tracking() -> bool:
+    # test_scope = os.environ.get("TEST_SCOPE", "")
+    # branch = os.environ.get("BUILDKITE_BRANCH", "")
+    # return test_scope == "full" and branch == "main"
+    return True # only for testing purpose.
+
+def _sanitize(value: str) -> str:
+    return re.sub(r"[^A-Za-z0-9._-]", "_", value)
+
+
+def _safe_float(value: Any) -> float | None:
+    if value is None:
+        return None
+    try:
+        return float(value)
+    except (TypeError, ValueError):
+        return None
+
+
+def _load_current_results() -> list[dict[str, Any]]:
+    pattern = os.path.join(RESULTS_DIR, "perf_*.json")
+    records: list[dict[str, Any]] = []
+    for path in sorted(glob.glob(pattern)):
+        with open(path, encoding="utf-8") as f:
+            records.append(json.load(f))
+    return records
+
+
+def _normalize_record(result: dict[str, Any]) -> dict[str, Any]:
+    benchmark_id = result.get("benchmark_id", "unknown")
+    model_id = benchmark_id
+
+    timestamp = result.get("timestamp")
+    if not timestamp:
+        timestamp = datetime.now(timezone.utc).isoformat()
+
+    commit_sha = result.get("commit") or os.environ.get("BUILDKITE_COMMIT", "")
+    latency = _safe_float(result.get("avg_generation_time_s"))
+    throughput = _safe_float(result.get("throughput_fps"))
+    memory = _safe_float(result.get("max_peak_memory_mb"))
+
+    return {
+        "model_id": model_id,
+        "timestamp": timestamp,
+        "commit_sha": commit_sha,
+        "gpu_type": result.get("device", "unknown"),
+        "latency": latency,
+        "throughput": throughput,
+        "memory": memory,
+        "success": True,
+    }
+
+
+def _write_tracking_record(record: dict[str, Any]) -> str:
+    model_dir = os.path.join(TRACKING_ROOT, _sanitize(record["model_id"]))
+    os.makedirs(model_dir, exist_ok=True)
+
+    timestamp = _sanitize(record["timestamp"])
+    commit = _sanitize(record["commit_sha"] or "unknown")
+    out_path = os.path.join(model_dir, f"{timestamp}_{commit}.json")
+
+    with open(out_path, "w", encoding="utf-8") as f:
+        json.dump(record, f, indent=2)
+
+    return out_path
+
+def _baseline_metric(records: list[dict[str, Any]], key: str) -> float | None:
+    values = [
+        _safe_float(r.get(key))
+        for r in records
+    ]
+    values = [v for v in values if v is not None]
+    if not values:
+        return None
+    return statistics.median(values)
+
+
+def _check_regressions(
+    current: dict[str, Any],
+    baseline_records: list[dict[str, Any]],
+    max_regression: float,
+) -> list[str]:
+    failures: list[str] = []
+
+    for metric in ("latency", "memory"):
+        baseline = _baseline_metric(baseline_records, metric)
+        curr = _safe_float(current.get(metric))
+        if baseline is None or curr is None or baseline <= 0:
+            continue
+        regression = (curr - baseline) / baseline
+        if regression > max_regression:
+            failures.append(
+                f"{current['model_id']} {metric} regressed by {regression * 100:.1f}% "
+                f"(current={curr:.3f}, baseline_median={baseline:.3f})"
+            )
+
+    baseline_tp = _baseline_metric(baseline_records, "throughput")
+    curr_tp = _safe_float(current.get("throughput"))
+    if baseline_tp is not None and curr_tp is not None and baseline_tp > 0:
+        regression = (baseline_tp - curr_tp) / baseline_tp
+        if regression > max_regression:
+            failures.append(
+                f"{current['model_id']} throughput regressed by {regression * 100:.1f}% "
+                f"(current={curr_tp:.3f}, baseline_median={baseline_tp:.3f})"
+            )
+
+    return failures
+
+
+def _metric_delta_percent(
+    metric: str,
+    current: dict[str, Any],
+    baseline_records: list[dict[str, Any]],
+) -> float | None:
+    curr = _safe_float(current.get(metric))
+    baseline = _baseline_metric(baseline_records, metric)
+    if curr is None or baseline is None or baseline <= 0:
+        return None
+
+    if metric in ("latency", "memory"):
+        return (curr - baseline) / baseline * 100.0
+    if metric == "throughput":
+        return (baseline - curr) / baseline * 100.0
+    return None
+
+
+def _compact_value(value: float | None, precision: int = 3) -> str:
+    if value is None:
+        return "n/a"
+    return f"{value:.{precision}f}"
+
+def _build_summary_row(
+    record: dict[str, Any], 
+    baseline_records: list[dict[str, Any]], 
+    has_failed: bool
+) -> dict[str, Any]:
+    """Formats a single benchmark result into a row for the Markdown summary table."""
+    
+    latency_base = _safe_float(_baseline_metric(baseline_records, "latency"))
+    throughput_base = _safe_float(_baseline_metric(baseline_records, "throughput"))
+    memory_base = _safe_float(_baseline_metric(baseline_records, "memory"))
+    
+    # Calculate percentages for the 'Worst Regression' column
+    latency_reg = _metric_delta_percent("latency", record, baseline_records)
+    throughput_reg = _metric_delta_percent("throughput", record, baseline_records)
+    memory_reg = _metric_delta_percent("memory", record, baseline_records)
+    
+    regressions = [v for v in (latency_reg, throughput_reg, memory_reg) if v is not None]
+    worst_regression_pct = max(regressions) if regressions else None
+
+    return {
+        "model_id": record["model_id"],
+        "gpu_type": record["gpu_type"],
+        "baseline_n": len(baseline_records),
+        "latency_curr": _safe_float(record.get("latency")),
+        "latency_base": latency_base,
+        "throughput_curr": _safe_float(record.get("throughput")),
+        "throughput_base": throughput_base,
+        "memory_curr": _safe_float(record.get("memory")),
+        "memory_base": memory_base,
+        "worst_regression_pct": worst_regression_pct,
+        "failed": has_failed,
+    }
+
+def _build_markdown_summary(
+    summary_rows: list[dict[str, Any]],
+    max_regression: float,
+) -> str:
+    lines = [
+        "## Performance Baseline Comparison",
+        "",
+        f"Threshold: regressions greater than {max_regression * 100:.1f}% fail",
+        "",
+        "| Model | GPU | Baseline N | Latency (curr/base) | Throughput (curr/base) | Memory (curr/base) | Worst Regression | Status |",
+        "|---|---|---:|---|---|---|---:|---|",
+    ]
+
+    for row in summary_rows:
+        latency = f"{_compact_value(row['latency_curr'])} / {_compact_value(row['latency_base'])}"
+        throughput = f"{_compact_value(row['throughput_curr'])} / {_compact_value(row['throughput_base'])}"
+        memory = f"{_compact_value(row['memory_curr'], 1)} / {_compact_value(row['memory_base'], 1)}"
+
+        worst_reg = "n/a" if row["worst_regression_pct"] is None else f"{row['worst_regression_pct']:.1f}%"
+        status = "FAIL" if row["failed"] else "PASS"
+
+        lines.append(
+            f"| {row['model_id']} | {row['gpu_type']} | {row['baseline_n']} | "
+            f"{latency} | {throughput} | {memory} | {worst_reg} | {status} |"
+        )
+
+    return "\n".join(lines) + "\n"
+
+def _emit_markdown_summary(markdown: str, commit_sha: str) -> None:
+    print("\n" + markdown)
+    
+    # 1. Existing GitHub logic (safe to keep)
+    summary_path = os.environ.get("GITHUB_STEP_SUMMARY")
+    if summary_path:
+        with open(summary_path, "a", encoding="utf-8") as f:
+            f.write(markdown + "\n")
+
+    # 2. Write to Modal volume for Buildkite to pick up in post-run hook
+    try:
+        perf_reports_dir = "/root/data/perf_reports"
+        os.makedirs(perf_reports_dir, exist_ok=True)
+        short_sha = commit_sha[:7] if commit_sha else "unknown"
+        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+        report_path = os.path.join(perf_reports_dir, f"perf_{short_sha}_{timestamp}.md")
+        with open(report_path, "w", encoding="utf-8") as f:
+            f.write(markdown + "\n")
+        print(f"Performance report written to {report_path}")
+    except Exception as e:
+        print(f"Failed to write performance report to Modal volume: {e}")
+
+def main() -> int:
+    # Pull the current state of the world from HF
+    sync_from_hf(TRACKING_ROOT)
+
+    current_results = _load_current_results()
+    if not current_results:
+        print(f"No performance result files found in {RESULTS_DIR}")
+        return 0
+
+    all_failures = []
+    summary_rows = []
+    persist_tracking = _should_persist_tracking()
+
+    if persist_tracking:
+        print("Tracking persistence enabled: full-suite run on main branch")
+    else:
+        print("Tracking persistence disabled: only full-suite runs on main branch are persisted")
+
+    for raw in current_results:
+        record = _normalize_record(raw)
+
+        baseline_records = load_records_for_model(
+            TRACKING_ROOT, record["model_id"], record["gpu_type"],
+            last_n=5, successful_only=True
+        )
+
+        failures = _check_regressions(record, baseline_records, MAX_REGRESSION)
+
+        # Tag the current record based on the failure.
+        if not baseline_records:
+            # INITIALIZATION CASE: First run for this model/GPU
+            print(f"No baseline for {record['model_id']} on {record['gpu_type']}. Initializing...")
+            failures = []
+            record["success"] = True  # The first run is always "successful"
+        else:
+            # COMPARISON CASE: Compare against the mean of the last 5 good runs
+            failures = _check_regressions(record, baseline_records, MAX_REGRESSION)
+            if failures:
+                record["success"] = False
+                all_failures.extend(failures)
+            else:
+                record["success"] = True
+    
+        # 5. Persist to HF if we are on main
+        if persist_tracking:
+            # This writes the JSON with the "success" field to /tmp
+            current_path = _write_tracking_record(record)
+            # This pushes it to the FastVideo/performance-tracking repo
+            upload_record(current_path, record)
+
+        
+        summary_row = _build_summary_row(record, baseline_records, bool(failures))
+        summary_rows.append(summary_row)
+
+    commit_sha = os.environ.get("BUILDKITE_COMMIT", "unknown")[:7]
+
+    markdown = _build_markdown_summary(summary_rows, MAX_REGRESSION)
+    _emit_markdown_summary(markdown, commit_sha)
+
+    if all_failures:
+        print("Performance regression check failed:")
+        for item in all_failures:
+            print(f"  - {item}")
+        return 1
+
+    print("Performance baseline comparison passed")
+    return 0
+
+if __name__ == "__main__":
+    sys.exit(main())

fastvideo/tests/performance/dashboard.py

@@ -0,0 +1,102 @@
+# SPDX-License-Identifier: Apache-2.0
+import os
+import shutil
+import subprocess
+from datetime import datetime
+
+import plotly.express as px
+import pandas as pd
+
+from hf_store import sync_from_hf, load_as_dataframe
+
+# -----------------------------
+# 1. Grouping
+# -----------------------------
+def group_data(df: pd.DataFrame):
+    # Group only by model+GPU so each group produces a time-series line.
+    # config_id (commit SHA) is carried as a column for hover/color use.
+    keys = ["model_id", "gpu_type"]
+    return df.groupby(keys, dropna=False)
+
+# -----------------------------
+# 2. Plot builder
+# -----------------------------
+def build_plots(df: pd.DataFrame) -> list:
+    figs = []
+
+    for (model_id, gpu_type), g in group_data(df):
+        g = g.sort_values("timestamp")
+
+        # One chart per metric so the y-axes aren't on wildly different scales
+        for metric in ("latency", "throughput", "memory"):
+            if g[metric].isna().all():
+                continue
+
+            fig = px.line(
+                g,
+                x="timestamp",
+                y=metric,
+                markers=True,
+                hover_data=["config_id", "commit_sha"],
+                title=f"{model_id} | {gpu_type} | {metric}",
+                labels={"timestamp": "Time", metric: metric},
+            )
+            figs.append(fig)
+
+    return figs
+
+# -----------------------------
+# 3. Render HTML dashboard
+# -----------------------------
+def render_html(figs: list, days: int) -> str:
+    html_parts = [
+        "<html>",
+        "<head><meta charset='utf-8'>",
+        "<style>body { font-family: sans-serif; margin: 2rem; }</style>",
+        "</head><body>",
+        f"<h2>Performance Dashboard (last {days} days)</h2>",
+    ]
+
+    for fig in figs:
+        html_parts.append(fig.to_html(full_html=False, include_plotlyjs="cdn"))
+
+    html_parts.append("</body></html>")
+    return "\n".join(html_parts)
+
+# -----------------------------
+# 5. Main
+# -----------------------------
+def main() -> None:
+    days = int(os.environ.get("DASHBOARD_DAYS", "30"))
+
+    local_dir = sync_from_hf("/tmp/perf-tracking")
+    df = load_as_dataframe(local_dir, days=days)
+
+    if df.empty:
+        print("No data found")
+        return
+
+    # Sanity-check: log what we actually loaded
+    print(f"Loaded {len(df)} records across {df['model_id'].nunique()} model(s), "
+          f"{df['gpu_type'].nunique()} GPU type(s), "
+          f"date range: {df['timestamp'].min()} → {df['timestamp'].max()}")
+
+    figs = build_plots(df)
+    html = render_html(figs, days)
+
+    commit_sha = os.environ.get("BUILDKITE_COMMIT", "unknown")[:7]
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+
+    report_dir = "/root/data/perf_reports"
+    os.makedirs(report_dir, exist_ok=True)
+
+    filename = f"dashboard_{commit_sha}_{timestamp}.html"
+    output_file = os.path.join(report_dir, filename)
+
+    with open(output_file, "w", encoding="utf-8") as f:
+        f.write(html)
+
+    print(f"Dashboard generated: {output_file}")
+
+if __name__ == "__main__":
+    main()

fastvideo/tests/performance/hf_store.py

@@ -0,0 +1,254 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Shared HuggingFace storage utilities for performance tracking.
+
+Provides a single place for:
+- Syncing the HF dataset repo to a local directory
+- Loading raw JSON records (with optional recency filter)
+- Loading records as a normalized pandas DataFrame
+- Uploading individual result files back to HF
+- Common helpers: sanitize, safe_float
+"""
+
+import glob
+import json
+import os
+import re
+from datetime import datetime, timedelta, timezone
+from typing import Any
+
+import pandas as pd
+from huggingface_hub import HfApi, snapshot_download
+
+# ---------------------------------------------------------------------------
+# Configuration — read once at import time, shared across both consumers
+# ---------------------------------------------------------------------------
+
+HF_REPO_ID: str = os.environ.get("HF_REPO_ID", "FastVideo/performance-tracking")
+HF_TOKEN: str | None = os.environ.get("HF_API_KEY")
+
+
+# ---------------------------------------------------------------------------
+# Low-level helpers
+# ---------------------------------------------------------------------------
+
+def sanitize(value: str) -> str:
+    """Return a filesystem- and HF-path-safe version of *value*."""
+    return re.sub(r"[^A-Za-z0-9._-]", "_", value)
+
+
+def safe_float(value: Any) -> float | None:
+    """Coerce *value* to float, returning None on failure."""
+    if value is None:
+        return None
+    try:
+        return float(value)
+    except (TypeError, ValueError):
+        return None
+
+
+# ---------------------------------------------------------------------------
+# HF I/O
+# ---------------------------------------------------------------------------
+
+def sync_from_hf(local_dir: str) -> str:
+    """Download the HF dataset repo snapshot to *local_dir*.
+
+    Returns *local_dir* so callers can chain: ``load_records(sync_from_hf(...))``.
+    On failure (empty repo, no credentials, network error) the function logs a
+    warning and returns *local_dir* unchanged so the caller can still work with
+    whatever is already on disk.
+    """
+    if not HF_REPO_ID:
+        print("hf_store: HF_REPO_ID not set, skipping sync.")
+        return local_dir
+
+    print(f"hf_store: syncing from {HF_REPO_ID} → {local_dir}")
+    try:
+        snapshot_download(
+            repo_id=HF_REPO_ID,
+            repo_type="dataset",
+            local_dir=local_dir,
+            token=HF_TOKEN,
+            allow_patterns="*.json",
+        )
+    except Exception as exc:
+        print(f"hf_store: sync skipped — {exc}")
+
+    return local_dir
+
+
+def upload_record(local_path: str, record: dict[str, Any]) -> None:
+    """Upload *local_path* to the HF repo under ``<model_id>/<filename>``.
+
+    Silently skips if HF_TOKEN is absent so local/CI runs without credentials
+    don't crash.
+    """
+    if not HF_TOKEN:
+        print("hf_store: HF_API_KEY not set, skipping upload.")
+        return
+
+    model_id = record.get("model_id", "unknown")
+    path_in_repo = f"{sanitize(model_id)}/{os.path.basename(local_path)}"
+    commit_sha = (record.get("commit_sha") or "unknown")[:7]
+
+    api = HfApi(token=HF_TOKEN)
+    try:
+        api.upload_file(
+            path_or_fileobj=local_path,
+            path_in_repo=path_in_repo,
+            repo_id=HF_REPO_ID,
+            repo_type="dataset",
+            commit_message=f"Perf: {model_id} at {commit_sha}",
+        )
+        print(f"hf_store: uploaded → {HF_REPO_ID}/{path_in_repo}")
+    except Exception as exc:
+        print(f"hf_store: upload failed — {exc}")
+
+
+# ---------------------------------------------------------------------------
+# Record loading
+# ---------------------------------------------------------------------------
+
+def load_records(
+    local_dir: str,
+    *,
+    days: int | None = None,
+    successful_only: bool = False,
+) -> list[dict[str, Any]]:
+    """Return raw JSON dicts from *local_dir*.
+
+    Args:
+        local_dir: Root directory previously populated by :func:`sync_from_hf`.
+        days: When set, discard records whose ``timestamp`` is older than this
+            many days. Records with a missing/unparseable timestamp are kept.
+        successful_only: When True, only records with ``success=True`` are
+            returned. Useful when building a regression baseline.
+
+    Returns:
+        List of raw dicts sorted by ``timestamp`` ascending (records that could
+        not be parsed are silently skipped).
+    """
+    cutoff: datetime | None = None
+    if days is not None:
+        cutoff = datetime.now(timezone.utc) - timedelta(days=days)
+
+    records: list[dict[str, Any]] = []
+
+    for path in sorted(glob.glob(os.path.join(local_dir, "**", "*.json"), recursive=True)):
+        try:
+            with open(path, encoding="utf-8") as fh:
+                data: dict[str, Any] = json.load(fh)
+        except (OSError, json.JSONDecodeError):
+            continue
+
+        if successful_only and not data.get("success", True):
+            continue
+
+        if cutoff is not None:
+            raw_ts = data.get("timestamp")
+            if raw_ts:
+                try:
+                    ts = datetime.fromisoformat(raw_ts)
+                    if ts.tzinfo is None:
+                        ts = ts.replace(tzinfo=timezone.utc)
+                    if ts < cutoff:
+                        continue
+                except ValueError:
+                    pass  # keep records with unparseable timestamps
+
+        records.append(data)
+
+    return records
+
+
+def load_records_for_model(
+    local_dir: str,
+    model_id: str,
+    gpu_type: str | None = None,
+    *,
+    last_n: int | None = None,
+    successful_only: bool = True,
+) -> list[dict[str, Any]]:
+    """Return records for a specific *model_id*, optionally filtered by GPU.
+
+    Args:
+        local_dir: Root directory previously populated by :func:`sync_from_hf`.
+        model_id: Matches the ``model_id`` field inside each JSON record.
+        gpu_type: When set, only records whose ``gpu_type`` matches are returned.
+        last_n: When set, return only the most recent *n* records (after all
+            other filters). Useful for sliding-window baseline calculations.
+        successful_only: Passed through to :func:`load_records`.
+
+    Returns:
+        List of matching dicts sorted by timestamp ascending.
+    """
+    model_dir = os.path.join(local_dir, sanitize(model_id))
+    if not os.path.isdir(model_dir):
+        return []
+
+    records = load_records(model_dir, successful_only=successful_only)
+
+    if gpu_type is not None:
+        records = [r for r in records if r.get("gpu_type") == gpu_type]
+
+    if last_n is not None:
+        records = records[-last_n:]
+
+    return records
+
+
+# ---------------------------------------------------------------------------
+# DataFrame helpers (dashboard / analytics consumers)
+# ---------------------------------------------------------------------------
+
+_NUMERIC_COLS = ("latency", "throughput", "memory")
+
+
+def normalize_dataframe(df: pd.DataFrame) -> pd.DataFrame:
+    """Apply standard type coercions to a raw records DataFrame.
+
+    - Parses ``timestamp`` to UTC-aware datetime.
+    - Coerces ``latency``, ``throughput``, ``memory`` to float.
+    - Adds a ``config_id`` column (first 7 chars of ``commit_sha``).
+
+    Returns the mutated DataFrame (also modifies in place for efficiency).
+    """
+    if df.empty:
+        return df
+
+    df["timestamp"] = pd.to_datetime(df["timestamp"], utc=True, errors="coerce")
+    df["config_id"] = df.get("commit_sha", pd.Series(dtype=str)).fillna("unknown").str[:7]
+
+    for col in _NUMERIC_COLS:
+        if col in df.columns:
+            df[col] = pd.to_numeric(df[col], errors="coerce")
+
+    return df
+
+
+def load_as_dataframe(
+    local_dir: str,
+    *,
+    days: int | None = None,
+    successful_only: bool = False,
+) -> pd.DataFrame:
+    """Load and normalize records from *local_dir* into a pandas DataFrame.
+
+    Combines :func:`load_records` + :func:`normalize_dataframe` into a single
+    call for consumers (e.g. the dashboard) that work exclusively with
+    DataFrames.
+
+    Args:
+        local_dir: Root directory previously populated by :func:`sync_from_hf`.
+        days: Passed through to :func:`load_records`.
+        successful_only: Passed through to :func:`load_records`.
+
+    Returns:
+        Normalized DataFrame, or an empty DataFrame if no records were found.
+    """
+    records = load_records(local_dir, days=days, successful_only=successful_only)
+    if not records:
+        return pd.DataFrame()
+
+    df = pd.DataFrame(records)
+    return normalize_dataframe(df)

fastvideo/tests/performance/test_inference_performance.py

@@ -164,6 +164,10 @@ def test_inference_performance(cfg):
     avg_time = sum(times) / len(times)
     max_peak_memory = max(peak_memories)
     device_name = torch.cuda.get_device_name()
+    num_frames = gen_kwargs.get("num_frames")
+    throughput_fps = (1.0 / avg_time) if avg_time > 0 else None
+    if isinstance(num_frames, (int, float)) and avg_time > 0:
+        throughput_fps = num_frames / avg_time
 
     results = {
         "benchmark_id": cfg["benchmark_id"],
@@ -174,6 +178,8 @@ def test_inference_performance(cfg):
         "num_measurement_runs": num_measure,
         "avg_generation_time_s": round(avg_time, 3),
         "individual_times_s": [round(t, 3) for t in times],
+        "throughput_fps": round(throughput_fps, 3)
+        if throughput_fps is not None else None,
         "max_peak_memory_mb": round(max_peak_memory, 1),
         "individual_peak_memories_mb": [round(m, 1) for m in peak_memories],
         "thresholds": thresholds,

fastvideo/tests/train/callbacks/__init__.py

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

fastvideo/tests/train/callbacks/test_callback.py

@@ -0,0 +1,290 @@
+# SPDX-License-Identifier: Apache-2.0
+"""CPU-only unit tests for :mod:`fastvideo.train.callbacks.callback`.
+
+Covers the ``Callback`` base class no-op contract and the
+``CallbackDict`` instantiation / dispatch / state-dict logic.
+
+The concrete callback subclasses (``GradNormClipCallback``,
+``EMACallback``, ``ValidationCallback``) have their own test files.
+"""
+from __future__ import annotations
+
+import logging
+from collections.abc import Iterator
+from contextlib import contextmanager
+from typing import Any
+
+import pytest
+
+from fastvideo.train.callbacks.callback import (
+    Callback,
+    CallbackDict,
+    _BUILTIN_CALLBACKS,
+)
+
+
+# ``fastvideo.logger.init_logger`` sets ``propagate=False`` on its
+# loggers, so the standard ``caplog`` fixture cannot observe them.
+# This helper attaches a temporary handler directly to the target
+# logger and yields the captured records.
+@contextmanager
+def _capture_logger(
+    name: str, level: int = logging.WARNING
+) -> Iterator[list[logging.LogRecord]]:
+    logger = logging.getLogger(name)
+    records: list[logging.LogRecord] = []
+
+    class _Handler(logging.Handler):
+
+        def emit(self, record: logging.LogRecord) -> None:
+            records.append(record)
+
+    handler = _Handler(level=level)
+    prev_level = logger.level
+    logger.addHandler(handler)
+    logger.setLevel(level)
+    try:
+        yield records
+    finally:
+        logger.removeHandler(handler)
+        logger.setLevel(prev_level)
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+class _RecordingCallback(Callback):
+    """Callback that records every hook call into a shared list."""
+
+    def __init__(self, *, tag: str, sink: list[str]) -> None:
+        self._tag = tag
+        self._sink = sink
+
+    def on_train_start(self, method, iteration: int = 0) -> None:
+        self._sink.append(f"{self._tag}:on_train_start:{iteration}")
+
+    def on_training_step_end(
+        self, method, loss_dict, iteration: int = 0
+    ) -> None:
+        self._sink.append(f"{self._tag}:on_training_step_end:{iteration}")
+
+    def on_validation_begin(self, method, iteration: int = 0) -> None:
+        self._sink.append(f"{self._tag}:on_validation_begin:{iteration}")
+
+    def state_dict(self) -> dict[str, Any]:
+        return {"tag": self._tag, "marker": 7}
+
+    def load_state_dict(self, sd: dict[str, Any]) -> None:
+        self._sink.append(f"{self._tag}:load:{sd.get('marker')}")
+
+
+class _NotACallback:
+    """Plain class used to exercise the non-Callback type guard."""
+
+    def __init__(self, **_: Any) -> None:
+        pass
+
+
+# ---------------------------------------------------------------------------
+# A. Callback base class
+# ---------------------------------------------------------------------------
+
+
+class TestCallbackBase:
+
+    def test_default_hooks_return_none(self) -> None:
+        cb = Callback()
+        assert cb.on_train_start(method=None) is None
+        assert (
+            cb.on_training_step_end(method=None, loss_dict={}) is None
+        )
+        assert cb.on_before_optimizer_step(method=None) is None
+        assert cb.on_validation_begin(method=None) is None
+        assert cb.on_validation_end(method=None) is None
+        assert cb.on_train_end(method=None) is None
+
+    def test_default_state_dict_round_trip(self) -> None:
+        cb = Callback()
+        assert cb.state_dict() == {}
+        # Default load_state_dict accepts arbitrary state without raising.
+        assert cb.load_state_dict({"unrelated": 1}) is None
+
+
+# ---------------------------------------------------------------------------
+# B. CallbackDict construction
+# ---------------------------------------------------------------------------
+
+
+class TestCallbackDictInit:
+
+    def test_empty_config(self) -> None:
+        cb_dict = CallbackDict({}, training_config=object())
+        assert cb_dict._callbacks == {}
+
+    def test_builtin_name_resolves_without_target(self) -> None:
+        # ``grad_clip`` is a registered builtin.
+        cfg = {"grad_clip": {"max_grad_norm": 0.5}}
+        tc = object()
+        cb_dict = CallbackDict(cfg, training_config=tc)
+
+        assert "grad_clip" in cb_dict._callbacks
+        from fastvideo.train.callbacks.grad_clip import (
+            GradNormClipCallback,
+        )
+        cb = cb_dict._callbacks["grad_clip"]
+        assert isinstance(cb, GradNormClipCallback)
+        # CallbackDict wires up training_config + back-pointer.
+        assert cb.training_config is tc
+        assert cb._callback_dict is cb_dict
+
+    def test_explicit_target_overrides_name_lookup(self) -> None:
+        cfg = {
+            "anything_goes": {
+                "_target_": (
+                    "fastvideo.train.callbacks.grad_clip."
+                    "GradNormClipCallback"
+                ),
+                "max_grad_norm": 1.0,
+            }
+        }
+        cb_dict = CallbackDict(cfg, training_config=object())
+        assert "anything_goes" in cb_dict._callbacks
+
+    def test_unknown_name_without_target_is_skipped(self) -> None:
+        cfg = {"mystery": {"some_arg": 1}}
+        with _capture_logger(
+            "fastvideo.train.callbacks.callback"
+        ) as records:
+            cb_dict = CallbackDict(cfg, training_config=object())
+        assert cb_dict._callbacks == {}
+        assert any(
+            "missing" in r.getMessage() and "mystery" in r.getMessage()
+            for r in records
+        )
+
+    def test_non_callback_target_raises(self) -> None:
+        cfg = {
+            "bad": {
+                "_target_": (
+                    "fastvideo.tests.train.callbacks.test_callback."
+                    "_NotACallback"
+                )
+            }
+        }
+        with pytest.raises(TypeError, match="expected a Callback"):
+            CallbackDict(cfg, training_config=object())
+
+    def test_builtin_registry_has_expected_entries(self) -> None:
+        # Sanity: protect the builtin registry from silent shrinkage.
+        assert set(_BUILTIN_CALLBACKS) >= {
+            "grad_clip",
+            "validation",
+            "ema",
+        }
+
+
+# ---------------------------------------------------------------------------
+# C. Dispatch via __getattr__
+# ---------------------------------------------------------------------------
+
+
+class TestCallbackDictDispatch:
+
+    def _build(
+        self, sink: list[str]
+    ) -> CallbackDict:
+        cb_dict = CallbackDict({}, training_config=object())
+        cb_dict._callbacks["first"] = _RecordingCallback(
+            tag="first", sink=sink
+        )
+        cb_dict._callbacks["second"] = _RecordingCallback(
+            tag="second", sink=sink
+        )
+        return cb_dict
+
+    def test_dispatch_calls_all_in_insertion_order(self) -> None:
+        sink: list[str] = []
+        cb_dict = self._build(sink)
+
+        cb_dict.on_train_start(method=None, iteration=3)
+        assert sink == [
+            "first:on_train_start:3",
+            "second:on_train_start:3",
+        ]
+
+    def test_dispatch_to_hook_some_callbacks_skip(self) -> None:
+        sink: list[str] = []
+        cb_dict = self._build(sink)
+        # The base Callback subclass below only implements one hook;
+        # dispatch should still fan out without raising.
+
+        class _OnlyValidation(Callback):
+
+            def on_validation_end(
+                self, method, iteration: int = 0
+            ) -> None:
+                sink.append(f"vend:{iteration}")
+
+        cb_dict._callbacks["only_v"] = _OnlyValidation()
+        cb_dict.on_validation_end(method=None, iteration=11)
+        assert "vend:11" in sink
+
+    def test_dispatch_unknown_hook_is_noop(self) -> None:
+        # Methods that don't exist on any callback should not raise.
+        cb_dict = self._build([])
+        cb_dict.totally_made_up_hook(method=None, iteration=0)
+
+    def test_underscore_attribute_raises(self) -> None:
+        cb_dict = self._build([])
+        with pytest.raises(AttributeError):
+            getattr(cb_dict, "_does_not_exist")
+
+
+# ---------------------------------------------------------------------------
+# D. state_dict / load_state_dict
+# ---------------------------------------------------------------------------
+
+
+class TestCallbackDictStateDict:
+
+    def _build(self) -> tuple[CallbackDict, list[str]]:
+        sink: list[str] = []
+        cb_dict = CallbackDict({}, training_config=object())
+        cb_dict._callbacks["first"] = _RecordingCallback(
+            tag="first", sink=sink
+        )
+        cb_dict._callbacks["second"] = _RecordingCallback(
+            tag="second", sink=sink
+        )
+        return cb_dict, sink
+
+    def test_state_dict_returns_per_callback_dict(self) -> None:
+        cb_dict, _ = self._build()
+        state = cb_dict.state_dict()
+        assert set(state) == {"first", "second"}
+        assert state["first"] == {"tag": "first", "marker": 7}
+        assert state["second"] == {"tag": "second", "marker": 7}
+
+    def test_load_state_dict_dispatches_to_each(self) -> None:
+        cb_dict, sink = self._build()
+        cb_dict.load_state_dict(
+            {
+                "first": {"marker": 1},
+                "second": {"marker": 2},
+            }
+        )
+        assert sink == ["first:load:1", "second:load:2"]
+
+    def test_load_state_dict_missing_key_warns_no_raise(self) -> None:
+        cb_dict, sink = self._build()
+        with _capture_logger(
+            "fastvideo.train.callbacks.callback"
+        ) as records:
+            cb_dict.load_state_dict({"first": {"marker": 99}})
+        assert sink == ["first:load:99"]
+        assert any(
+            "second" in r.getMessage() and "not found" in r.getMessage()
+            for r in records
+        )

fastvideo/tests/train/callbacks/test_ema.py

@@ -0,0 +1,254 @@
+# SPDX-License-Identifier: Apache-2.0
+"""CPU-only unit tests for :mod:`fastvideo.train.callbacks.ema`.
+
+Exercises the EMA lifecycle (lazy init, ``start_iter`` gating, decay
+math, ``ema_context`` swap, state-dict round-trip) on a tiny CPU
+``nn.Linear``. ``EMA_FSDP`` works without ``dist.init_process_group``
+because ``dist.is_initialized()`` returns False and ``_to_local_tensor``
+falls through to raw tensors for non-DTensor inputs.
+"""
+from __future__ import annotations
+
+from typing import Any
+
+import pytest
+import torch
+
+from fastvideo.train.callbacks.ema import EMACallback
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+class _Student:
+
+    def __init__(self, transformer: torch.nn.Module | None) -> None:
+        self.transformer = transformer
+
+
+class _RecordingTracker:
+
+    def __init__(self) -> None:
+        self.entries: list[tuple[dict[str, Any], int]] = []
+
+    def log(self, payload: dict[str, Any], step: int) -> None:
+        self.entries.append((payload, step))
+
+
+class _Method:
+
+    def __init__(
+        self,
+        transformer: torch.nn.Module | None,
+        tracker: Any | None = None,
+    ) -> None:
+        self.student = _Student(transformer)
+        self.tracker = tracker
+
+
+def _tiny_transformer(*, fill: float = 0.0) -> torch.nn.Module:
+    m = torch.nn.Linear(4, 2, bias=False)
+    with torch.no_grad():
+        m.weight.fill_(fill)
+    return m
+
+
+# ---------------------------------------------------------------------------
+# A. on_train_start
+# ---------------------------------------------------------------------------
+
+
+class TestOnTrainStart:
+
+    def test_initializes_ema_from_student(self) -> None:
+        transformer = _tiny_transformer(fill=0.5)
+        cb = EMACallback(decay=0.9, start_iter=0)
+        cb.on_train_start(_Method(transformer), iteration=0)
+
+        assert cb.student_ema is not None
+        # Shadow shape matches transformer parameter.
+        shadow = cb.student_ema.shadow["weight"]
+        assert shadow.shape == transformer.weight.shape
+        assert torch.allclose(shadow, transformer.weight.detach().cpu())
+
+    def test_missing_transformer_raises(self) -> None:
+        cb = EMACallback()
+        with pytest.raises(ValueError, match="No student transformer"):
+            cb.on_train_start(_Method(transformer=None), iteration=0)
+
+
+# ---------------------------------------------------------------------------
+# B. on_training_step_end (decay math + start_iter gating)
+# ---------------------------------------------------------------------------
+
+
+class TestOnTrainingStepEnd:
+
+    def test_no_op_before_train_start(self) -> None:
+        cb = EMACallback()
+        # student_ema is None until on_train_start.
+        cb.on_training_step_end(
+            _Method(transformer=None), loss_dict={}, iteration=0
+        )
+        assert not cb._ema_started
+
+    def test_skipped_until_start_iter(self) -> None:
+        transformer = _tiny_transformer(fill=1.0)
+        cb = EMACallback(decay=0.5, start_iter=10)
+        cb.on_train_start(_Method(transformer), iteration=0)
+
+        # Mutate transformer to drift it away from initial shadow.
+        with torch.no_grad():
+            transformer.weight.fill_(7.0)
+
+        cb.on_training_step_end(
+            _Method(transformer), loss_dict={}, iteration=5
+        )
+        # Below start_iter: shadow is untouched, _ema_started False.
+        assert not cb._ema_started
+        assert torch.allclose(
+            cb.student_ema.shadow["weight"],
+            torch.full((2, 4), 1.0),
+        )
+
+    def test_first_active_step_reinits_then_updates(self) -> None:
+        transformer = _tiny_transformer(fill=1.0)
+        cb = EMACallback(decay=0.9, start_iter=10)
+        cb.on_train_start(_Method(transformer), iteration=0)
+
+        # Drift transformer so that re-init has a visible effect.
+        with torch.no_grad():
+            transformer.weight.fill_(5.0)
+
+        cb.on_training_step_end(
+            _Method(transformer), loss_dict={}, iteration=10
+        )
+        # First active step: shadow is re-initialized from the
+        # current transformer (5.0) and *then* update() applies decay
+        # against the same value, so shadow stays at 5.0.
+        assert cb._ema_started
+        assert torch.allclose(
+            cb.student_ema.shadow["weight"],
+            torch.full((2, 4), 5.0),
+        )
+
+    def test_subsequent_step_applies_decay(self) -> None:
+        transformer = _tiny_transformer(fill=2.0)
+        cb = EMACallback(decay=0.9, start_iter=0)
+        cb.on_train_start(_Method(transformer), iteration=0)
+
+        # Step 0: re-init at 2.0, then update against 2.0 → still 2.0.
+        cb.on_training_step_end(
+            _Method(transformer), loss_dict={}, iteration=0
+        )
+        # Step 1: drift transformer to 12.0, expect
+        # shadow = 0.9 * 2.0 + 0.1 * 12.0 = 3.0.
+        with torch.no_grad():
+            transformer.weight.fill_(12.0)
+        cb.on_training_step_end(
+            _Method(transformer), loss_dict={}, iteration=1
+        )
+        assert torch.allclose(
+            cb.student_ema.shadow["weight"],
+            torch.full((2, 4), 3.0),
+            atol=1e-6,
+        )
+
+    def test_tracker_logs_decay(self) -> None:
+        transformer = _tiny_transformer()
+        tracker = _RecordingTracker()
+        cb = EMACallback(decay=0.99, start_iter=0)
+        method = _Method(transformer, tracker=tracker)
+        cb.on_train_start(method, iteration=0)
+        cb.on_training_step_end(method, loss_dict={}, iteration=0)
+
+        assert any(
+            payload.get("ema/decay") == 0.99 and step == 0
+            for payload, step in tracker.entries
+        )
+
+
+# ---------------------------------------------------------------------------
+# C. ema_context
+# ---------------------------------------------------------------------------
+
+
+class TestEmaContext:
+
+    def test_passthrough_when_inactive(self) -> None:
+        transformer = _tiny_transformer(fill=3.0)
+        cb = EMACallback()
+        # No on_train_start → student_ema is None.
+        with cb.ema_context(transformer) as t:
+            assert t is transformer
+            assert torch.allclose(
+                t.weight, torch.full((2, 4), 3.0)
+            )
+
+    def test_swaps_weights_then_restores(self) -> None:
+        transformer = _tiny_transformer(fill=1.0)
+        cb = EMACallback(decay=0.0, start_iter=0)
+        method = _Method(transformer)
+        cb.on_train_start(method, iteration=0)
+
+        # decay=0 → after one step the shadow == current weights == 1.0.
+        cb.on_training_step_end(method, loss_dict={}, iteration=0)
+        # Drift transformer; ema_context should swap shadow (1.0) in
+        # for the duration and restore the post-drift value (9.0).
+        with torch.no_grad():
+            transformer.weight.fill_(9.0)
+
+        with cb.ema_context(transformer) as t:
+            assert torch.allclose(t.weight, torch.full((2, 4), 1.0))
+
+        assert torch.allclose(
+            transformer.weight, torch.full((2, 4), 9.0)
+        )
+
+
+# ---------------------------------------------------------------------------
+# D. State dict round-trip
+# ---------------------------------------------------------------------------
+
+
+class TestStateDict:
+
+    def test_state_dict_empty_before_train_start(self) -> None:
+        cb = EMACallback()
+        assert cb.state_dict() == {}
+
+    def test_round_trip_preserves_shadow_and_started_flag(self) -> None:
+        transformer = _tiny_transformer(fill=4.0)
+        cb = EMACallback(decay=0.5, start_iter=0)
+        method = _Method(transformer)
+        cb.on_train_start(method, iteration=0)
+        cb.on_training_step_end(method, loss_dict={}, iteration=0)
+
+        state = cb.state_dict()
+        assert "student_ema" in state
+        assert state["ema_started"] is True
+
+        # Build a fresh callback and load.
+        fresh = EMACallback(decay=0.5, start_iter=0)
+        fresh.on_train_start(_Method(_tiny_transformer(fill=0.0)),
+                             iteration=0)
+        # Sanity: fresh shadow != saved shadow before load.
+        assert not torch.allclose(
+            fresh.student_ema.shadow["weight"],
+            cb.student_ema.shadow["weight"],
+        )
+        fresh.load_state_dict(state)
+        assert fresh._ema_started is True
+        assert torch.allclose(
+            fresh.student_ema.shadow["weight"],
+            cb.student_ema.shadow["weight"],
+        )
+
+    def test_load_without_student_ema_only_sets_flag(self) -> None:
+        cb = EMACallback()
+        # student_ema is None — load must not attempt to assign shadow.
+        cb.load_state_dict({"ema_started": True})
+        assert cb._ema_started is True
+        assert cb.student_ema is None

fastvideo/tests/train/callbacks/test_grad_clip.py

@@ -0,0 +1,170 @@
+# SPDX-License-Identifier: Apache-2.0
+"""CPU-only unit tests for :mod:`fastvideo.train.callbacks.grad_clip`.
+
+Exercises ``GradNormClipCallback.on_before_optimizer_step`` against
+synthetic ``nn.Module`` targets with manually populated gradients.
+"""
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+
+from fastvideo.train.callbacks.grad_clip import GradNormClipCallback
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+class _RecordingTracker:
+    """Tracker stub that records every ``log`` call."""
+
+    def __init__(self) -> None:
+        self.entries: list[tuple[dict[str, Any], int]] = []
+
+    def log(self, payload: dict[str, Any], step: int) -> None:
+        self.entries.append((payload, step))
+
+
+class _Method:
+    """Minimal stand-in for ``TrainingMethod``."""
+
+    def __init__(
+        self,
+        targets: dict[str, torch.nn.Module],
+        tracker: Any | None = None,
+    ) -> None:
+        self._targets = targets
+        self.tracker = tracker
+        self.iter_seen: int | None = None
+
+    def get_grad_clip_targets(
+        self, iteration: int
+    ) -> dict[str, torch.nn.Module]:
+        self.iter_seen = iteration
+        return self._targets
+
+
+def _make_module(*, grad_value: float, n: int = 4) -> torch.nn.Module:
+    """Return an ``nn.Linear`` whose grads are filled with ``grad_value``."""
+    m = torch.nn.Linear(n, n, bias=False)
+    m.weight.grad = torch.full_like(m.weight, fill_value=grad_value)
+    return m
+
+
+def _grad_norm(module: torch.nn.Module) -> float:
+    flat = torch.cat([p.grad.flatten() for p in module.parameters()])
+    return float(flat.norm(2).item())
+
+
+# ---------------------------------------------------------------------------
+# Tests
+# ---------------------------------------------------------------------------
+
+
+class TestGradNormClipCallback:
+
+    def test_disabled_when_max_norm_non_positive(self) -> None:
+        m = _make_module(grad_value=10.0)
+        before = _grad_norm(m)
+
+        cb = GradNormClipCallback(max_grad_norm=0.0)
+        method = _Method(targets={"m": m})
+        cb.on_before_optimizer_step(method=method, iteration=0)
+
+        # No clipping applied; ``get_grad_clip_targets`` not consulted.
+        assert _grad_norm(m) == before
+        assert method.iter_seen is None
+
+    def test_large_grads_get_clipped(self) -> None:
+        m = _make_module(grad_value=10.0)
+        assert _grad_norm(m) > 1.0
+
+        cb = GradNormClipCallback(max_grad_norm=1.0)
+        cb.on_before_optimizer_step(
+            method=_Method(targets={"m": m}), iteration=0
+        )
+
+        # After clipping the L2 norm should not exceed max_grad_norm
+        # (allow a tiny epsilon for the +1e-6 in the clip helper).
+        assert _grad_norm(m) <= 1.0 + 1e-4
+
+    def test_small_grads_unchanged(self) -> None:
+        m = _make_module(grad_value=0.01)
+        before = _grad_norm(m)
+        assert before < 1.0
+
+        cb = GradNormClipCallback(max_grad_norm=1.0)
+        cb.on_before_optimizer_step(
+            method=_Method(targets={"m": m}), iteration=0
+        )
+
+        # Clip coef >1 is clamped to 1, so values are preserved
+        # (modulo the *1.0 multiply, which is exact for floats).
+        assert abs(_grad_norm(m) - before) < 1e-6
+
+    def test_iteration_forwarded_to_targets(self) -> None:
+        m = _make_module(grad_value=0.5)
+        cb = GradNormClipCallback(max_grad_norm=1.0)
+        method = _Method(targets={"m": m})
+        cb.on_before_optimizer_step(method=method, iteration=42)
+        assert method.iter_seen == 42
+
+    def test_tracker_logged_when_enabled(self) -> None:
+        m = _make_module(grad_value=5.0)
+        tracker = _RecordingTracker()
+        cb = GradNormClipCallback(
+            max_grad_norm=1.0, log_grad_norms=True
+        )
+        cb.on_before_optimizer_step(
+            method=_Method(targets={"layer": m}, tracker=tracker),
+            iteration=7,
+        )
+        assert len(tracker.entries) == 1
+        payload, step = tracker.entries[0]
+        assert step == 7
+        assert "grad_norm/layer" in payload
+        assert payload["grad_norm/layer"] > 0.0
+
+    def test_tracker_not_logged_when_disabled(self) -> None:
+        m = _make_module(grad_value=5.0)
+        tracker = _RecordingTracker()
+        cb = GradNormClipCallback(
+            max_grad_norm=1.0, log_grad_norms=False
+        )
+        cb.on_before_optimizer_step(
+            method=_Method(targets={"m": m}, tracker=tracker),
+            iteration=0,
+        )
+        assert tracker.entries == []
+
+    def test_no_tracker_does_not_raise(self) -> None:
+        m = _make_module(grad_value=5.0)
+        cb = GradNormClipCallback(max_grad_norm=1.0, log_grad_norms=True)
+        # Method without a tracker attribute at all.
+
+        class _BareMethod:
+
+            def get_grad_clip_targets(
+                self, iteration: int
+            ) -> dict[str, torch.nn.Module]:
+                return {"m": m}
+
+        cb.on_before_optimizer_step(method=_BareMethod(), iteration=0)
+        # No assertion — must simply not raise.
+
+    def test_multiple_targets_each_logged(self) -> None:
+        targets = {
+            "head": _make_module(grad_value=4.0),
+            "tail": _make_module(grad_value=8.0),
+        }
+        tracker = _RecordingTracker()
+        cb = GradNormClipCallback(max_grad_norm=1.0)
+        cb.on_before_optimizer_step(
+            method=_Method(targets=targets, tracker=tracker),
+            iteration=1,
+        )
+        keys = {next(iter(p)) for p, _ in tracker.entries}
+        assert keys == {"grad_norm/head", "grad_norm/tail"}

fastvideo/tests/train/callbacks/test_validation.py

@@ -0,0 +1,234 @@
+# SPDX-License-Identifier: Apache-2.0
+"""CPU-only unit tests for :mod:`fastvideo.train.callbacks.validation`.
+
+Covers the parts of ``ValidationCallback`` that don't need a real
+pipeline or distributed init:
+
+* constructor type coercions and defaults,
+* ``on_validation_begin`` gating logic (every_steps + modulo),
+* ``_find_ema_callback`` lookup via ``_callback_dict``,
+* ``state_dict`` / ``load_state_dict`` rng round-trip.
+
+The heavy ``_run_validation`` path needs a real diffusion pipeline plus
+distributed init and is exercised by Phase 2/3 tests.
+"""
+from __future__ import annotations
+
+import torch
+
+from fastvideo.train.callbacks.callback import CallbackDict
+from fastvideo.train.callbacks.ema import EMACallback
+from fastvideo.train.callbacks.validation import ValidationCallback
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+_PIPE_TARGET = "fastvideo.pipelines.basic.wan.wan_pipeline.WanPipeline"
+
+
+def _make_callback(
+    *,
+    every_steps: int = 100,
+    sampling_steps: list[int] | None = None,
+    guidance_scale: float | None = None,
+    num_frames: int | None = None,
+    sampling_timesteps: list[int] | None = None,
+    output_dir: str | None = None,
+) -> ValidationCallback:
+    return ValidationCallback(
+        pipeline_target=_PIPE_TARGET,
+        dataset_file="/tmp/does_not_exist.json",
+        every_steps=every_steps,
+        sampling_steps=sampling_steps,
+        guidance_scale=guidance_scale,
+        num_frames=num_frames,
+        sampling_timesteps=sampling_timesteps,
+        output_dir=output_dir,
+    )
+
+
+# ---------------------------------------------------------------------------
+# A. Constructor coercions / defaults
+# ---------------------------------------------------------------------------
+
+
+class TestConstructor:
+
+    def test_defaults(self) -> None:
+        cb = _make_callback()
+        assert cb.pipeline_target == _PIPE_TARGET
+        assert cb.dataset_file == "/tmp/does_not_exist.json"
+        assert cb.every_steps == 100
+        assert cb.sampling_steps == [40]
+        assert cb.guidance_scale is None
+        assert cb.num_frames is None
+        assert cb.sampling_timesteps is None
+        assert cb.output_dir is None
+        # Lazy fields not yet populated.
+        assert cb._pipeline is None
+        assert cb._sampling_param is None
+        assert cb.validation_random_generator is None
+
+    def test_string_inputs_are_coerced(self) -> None:
+        # YAML often produces strings for numeric fields; the
+        # constructor must coerce them.
+        cb = ValidationCallback(
+            pipeline_target=_PIPE_TARGET,
+            dataset_file="x.json",
+            every_steps="50",  # type: ignore[arg-type]
+            sampling_steps=["20", "40"],  # type: ignore[arg-type]
+            guidance_scale="4.5",  # type: ignore[arg-type]
+            num_frames="77",  # type: ignore[arg-type]
+            sampling_timesteps=["1000", "500"],
+        )
+        assert cb.every_steps == 50
+        assert cb.sampling_steps == [20, 40]
+        assert cb.guidance_scale == 4.5
+        assert cb.num_frames == 77
+        assert cb.sampling_timesteps == [1000, 500]
+
+    def test_pipeline_kwargs_collected(self) -> None:
+        cb = ValidationCallback(
+            pipeline_target=_PIPE_TARGET,
+            dataset_file="x.json",
+            extra_arg=123,
+            another="value",
+        )
+        # Unknown kwargs are stashed for the pipeline factory.
+        assert cb.pipeline_kwargs == {
+            "extra_arg": 123,
+            "another": "value",
+        }
+
+
+# ---------------------------------------------------------------------------
+# B. on_validation_begin gating
+# ---------------------------------------------------------------------------
+
+
+class _NoRunValidation(ValidationCallback):
+    """Subclass that records ``_run_validation`` calls instead of
+    actually running them — lets us assert the gating logic without a
+    real pipeline."""
+
+    def __init__(self, **kwargs) -> None:
+        super().__init__(**kwargs)
+        self.run_calls: list[int] = []
+
+    def _run_validation(self, method, step: int) -> None:  # type: ignore[override]
+        self.run_calls.append(step)
+
+
+def _make_recording(**kwargs) -> _NoRunValidation:
+    return _NoRunValidation(
+        pipeline_target=_PIPE_TARGET,
+        dataset_file="x.json",
+        **kwargs,
+    )
+
+
+class TestOnValidationBegin:
+
+    def test_skipped_when_every_steps_zero(self) -> None:
+        cb = _make_recording(every_steps=0)
+        cb.on_validation_begin(method=None, iteration=0)
+        cb.on_validation_begin(method=None, iteration=1000)
+        assert cb.run_calls == []
+
+    def test_skipped_on_off_iter(self) -> None:
+        cb = _make_recording(every_steps=50)
+        cb.on_validation_begin(method=None, iteration=49)
+        cb.on_validation_begin(method=None, iteration=51)
+        assert cb.run_calls == []
+
+    def test_runs_on_match(self) -> None:
+        cb = _make_recording(every_steps=50)
+        cb.on_validation_begin(method=None, iteration=50)
+        cb.on_validation_begin(method=None, iteration=100)
+        assert cb.run_calls == [50, 100]
+
+    def test_iter_zero_runs(self) -> None:
+        # 0 % anything == 0 → step 0 fires (matches existing
+        # validation behavior used by ValidationCallback consumers).
+        cb = _make_recording(every_steps=50)
+        cb.on_validation_begin(method=None, iteration=0)
+        assert cb.run_calls == [0]
+
+
+# ---------------------------------------------------------------------------
+# C. _find_ema_callback
+# ---------------------------------------------------------------------------
+
+
+class TestFindEmaCallback:
+
+    def test_returns_none_without_callback_dict(self) -> None:
+        cb = _make_callback()
+        # _callback_dict is not set on bare instances.
+        assert cb._find_ema_callback() is None
+
+    def test_returns_none_when_no_ema_registered(self) -> None:
+        cb = _make_callback()
+        cb_dict = CallbackDict({}, training_config=object())
+        cb._callback_dict = cb_dict
+        assert cb._find_ema_callback() is None
+
+    def test_finds_ema_callback(self) -> None:
+        cb = _make_callback()
+        cb_dict = CallbackDict({}, training_config=object())
+        ema = EMACallback(decay=0.99)
+        cb_dict._callbacks["ema"] = ema
+        cb_dict._callbacks["validation"] = cb
+        cb._callback_dict = cb_dict
+
+        found = cb._find_ema_callback()
+        assert found is ema
+
+
+# ---------------------------------------------------------------------------
+# D. state_dict / load_state_dict (rng round-trip)
+# ---------------------------------------------------------------------------
+
+
+class TestStateDict:
+
+    def test_state_dict_empty_without_generator(self) -> None:
+        cb = _make_callback()
+        # validation_random_generator is None until on_train_start.
+        assert cb.state_dict() == {}
+
+    def test_round_trip_preserves_rng_state(self) -> None:
+        cb = _make_callback()
+        gen = torch.Generator(device="cpu").manual_seed(123)
+        # Advance RNG so a default-init generator on the receiving
+        # side is observably different.
+        for _ in range(5):
+            torch.randn(4, generator=gen)
+        cb.validation_random_generator = gen
+
+        state = cb.state_dict()
+        assert "validation_rng" in state
+
+        # Receiver: fresh generator with a different seed.
+        fresh = _make_callback()
+        fresh.validation_random_generator = (
+            torch.Generator(device="cpu").manual_seed(999)
+        )
+        fresh.load_state_dict(state)
+
+        # After load, both generators draw the same next sample.
+        a = torch.randn(8, generator=cb.validation_random_generator)
+        b = torch.randn(8, generator=fresh.validation_random_generator)
+        assert torch.equal(a, b)
+
+    def test_load_without_generator_is_noop(self) -> None:
+        cb = _make_callback()
+        # Generator is None: load must not raise even when state has
+        # an rng entry.
+        cb.load_state_dict(
+            {"validation_rng": torch.tensor([1, 2, 3], dtype=torch.uint8)}
+        )
+        assert cb.validation_random_generator is None

fastvideo/tests/train/utils/test_checkpoint.py

@@ -0,0 +1,360 @@
+# SPDX-License-Identifier: Apache-2.0
+"""CPU-only unit tests for :mod:`fastvideo.train.utils.checkpoint`.
+
+Covers the pure-Python portions of the checkpoint manager: name
+parsing, resume-path resolution, metadata round-trip, rolling-delete
+cleanup, the ``_is_stateful`` predicate, and the ``maybe_save`` gating
+logic. Code paths that touch DCP (``dcp.save`` / ``dcp.load``) and
+CUDA RNG snapshots are intentionally not covered here — those need a
+GPU runner and will be tested in later phases.
+"""
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any
+
+import pytest
+
+from fastvideo.train.utils.checkpoint import (
+    CheckpointConfig,
+    CheckpointManager,
+    _find_latest_checkpoint,
+    _is_stateful,
+    _parse_step_from_dir,
+    _resolve_resume_checkpoint,
+)
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def _make_checkpoint_dir(
+    output_dir: Path,
+    step: int,
+    *,
+    with_dcp: bool = True,
+) -> Path:
+    """Create a fake ``checkpoint-<step>/dcp`` directory tree."""
+    ckpt_dir = output_dir / f"checkpoint-{step}"
+    ckpt_dir.mkdir(parents=True, exist_ok=True)
+    if with_dcp:
+        (ckpt_dir / "dcp").mkdir(exist_ok=True)
+    return ckpt_dir
+
+
+def _make_manager(
+    tmp_path: Path,
+    *,
+    save_steps: int = 0,
+    keep_last: int = 0,
+    raw_config: dict[str, Any] | None = None,
+) -> CheckpointManager:
+    """Build a minimal ``CheckpointManager`` for tests that don't touch DCP."""
+    return CheckpointManager(
+        method=None,
+        dataloader=None,
+        output_dir=str(tmp_path),
+        config=CheckpointConfig(save_steps=save_steps, keep_last=keep_last),
+        raw_config=raw_config,
+    )
+
+
+# ---------------------------------------------------------------------------
+# A. _is_stateful predicate
+# ---------------------------------------------------------------------------
+
+
+class _Full:
+
+    def state_dict(self) -> dict[str, Any]:
+        return {}
+
+    def load_state_dict(self, sd: dict[str, Any]) -> None:
+        pass
+
+
+class _MissingStateDict:
+
+    def load_state_dict(self, sd: dict[str, Any]) -> None:
+        pass
+
+
+class _MissingLoad:
+
+    def state_dict(self) -> dict[str, Any]:
+        return {}
+
+
+def test_is_stateful_true_for_full_object() -> None:
+    assert _is_stateful(_Full()) is True
+
+
+def test_is_stateful_false_when_missing_state_dict() -> None:
+    assert _is_stateful(_MissingStateDict()) is False
+
+
+def test_is_stateful_false_when_missing_load_state_dict() -> None:
+    assert _is_stateful(_MissingLoad()) is False
+
+
+# ---------------------------------------------------------------------------
+# B. _parse_step_from_dir
+# ---------------------------------------------------------------------------
+
+
+def test_parse_step_valid(tmp_path: Path) -> None:
+    assert _parse_step_from_dir(tmp_path / "checkpoint-100") == 100
+
+
+def test_parse_step_zero(tmp_path: Path) -> None:
+    assert _parse_step_from_dir(tmp_path / "checkpoint-0") == 0
+
+
+def test_parse_step_invalid_raises(tmp_path: Path) -> None:
+    with pytest.raises(ValueError, match="Invalid checkpoint directory"):
+        _parse_step_from_dir(tmp_path / "not-a-checkpoint")
+
+
+# ---------------------------------------------------------------------------
+# C. _find_latest_checkpoint
+# ---------------------------------------------------------------------------
+
+
+def test_find_latest_returns_none_on_nonexistent_dir(tmp_path: Path) -> None:
+    assert _find_latest_checkpoint(tmp_path / "missing") is None
+
+
+def test_find_latest_returns_none_on_empty_dir(tmp_path: Path) -> None:
+    assert _find_latest_checkpoint(tmp_path) is None
+
+
+def test_find_latest_returns_largest_step(tmp_path: Path) -> None:
+    _make_checkpoint_dir(tmp_path, 10)
+    _make_checkpoint_dir(tmp_path, 200)
+    _make_checkpoint_dir(tmp_path, 50)
+    latest = _find_latest_checkpoint(tmp_path)
+    assert latest is not None
+    assert latest.name == "checkpoint-200"
+
+
+def test_find_latest_skips_dirs_without_dcp_subdir(tmp_path: Path) -> None:
+    # checkpoint-10 is "corrupted" — has no dcp/ subdir, must be skipped.
+    _make_checkpoint_dir(tmp_path, 10, with_dcp=False)
+    _make_checkpoint_dir(tmp_path, 5, with_dcp=True)
+    latest = _find_latest_checkpoint(tmp_path)
+    assert latest is not None
+    assert latest.name == "checkpoint-5"
+
+
+def test_find_latest_skips_non_checkpoint_dirs(tmp_path: Path) -> None:
+    (tmp_path / "logs").mkdir()
+    (tmp_path / "wandb").mkdir()
+    (tmp_path / "some_file.txt").write_text("noise")
+    _make_checkpoint_dir(tmp_path, 7)
+    latest = _find_latest_checkpoint(tmp_path)
+    assert latest is not None
+    assert latest.name == "checkpoint-7"
+
+
+# ---------------------------------------------------------------------------
+# D. _resolve_resume_checkpoint
+# ---------------------------------------------------------------------------
+
+
+def test_resolve_latest_with_no_checkpoints_returns_none(
+        tmp_path: Path) -> None:
+    out = tmp_path / "outputs"
+    out.mkdir()
+    assert _resolve_resume_checkpoint("latest", output_dir=str(out)) is None
+
+
+def test_resolve_latest_returns_latest_checkpoint(tmp_path: Path) -> None:
+    _make_checkpoint_dir(tmp_path, 30)
+    _make_checkpoint_dir(tmp_path, 10)
+    resolved = _resolve_resume_checkpoint("latest", output_dir=str(tmp_path))
+    assert resolved is not None
+    assert resolved.name == "checkpoint-30"
+
+
+def test_resolve_explicit_checkpoint_dir(tmp_path: Path) -> None:
+    ckpt = _make_checkpoint_dir(tmp_path, 42)
+    resolved = _resolve_resume_checkpoint(str(ckpt),
+                                          output_dir=str(tmp_path))
+    assert resolved is not None
+    assert resolved.name == "checkpoint-42"
+
+
+def test_resolve_dcp_subdir_returns_parent_checkpoint(tmp_path: Path) -> None:
+    ckpt = _make_checkpoint_dir(tmp_path, 42)
+    dcp_path = ckpt / "dcp"
+    resolved = _resolve_resume_checkpoint(str(dcp_path),
+                                          output_dir=str(tmp_path))
+    assert resolved is not None
+    assert resolved.name == "checkpoint-42"
+
+
+def test_resolve_output_dir_returns_latest(tmp_path: Path) -> None:
+    out = tmp_path / "outputs"
+    out.mkdir()
+    _make_checkpoint_dir(out, 100)
+    _make_checkpoint_dir(out, 50)
+    resolved = _resolve_resume_checkpoint(str(out), output_dir=str(tmp_path))
+    assert resolved is not None
+    assert resolved.name == "checkpoint-100"
+
+
+def test_resolve_nonexistent_path_raises(tmp_path: Path) -> None:
+    with pytest.raises(FileNotFoundError):
+        _resolve_resume_checkpoint(str(tmp_path / "missing"),
+                                   output_dir=str(tmp_path))
+
+
+def test_resolve_checkpoint_without_dcp_raises(tmp_path: Path) -> None:
+    ckpt = _make_checkpoint_dir(tmp_path, 5, with_dcp=False)
+    with pytest.raises(FileNotFoundError, match="dcp"):
+        _resolve_resume_checkpoint(str(ckpt), output_dir=str(tmp_path))
+
+
+def test_resolve_unknown_dir_raises(tmp_path: Path) -> None:
+    """A dir that is neither a checkpoint nor an output_dir-with-checkpoints."""
+    bogus = tmp_path / "bogus"
+    bogus.mkdir()
+    with pytest.raises(ValueError, match="Could not resolve"):
+        _resolve_resume_checkpoint(str(bogus), output_dir=str(tmp_path))
+
+
+# ---------------------------------------------------------------------------
+# E. metadata read/write
+# ---------------------------------------------------------------------------
+
+
+def test_write_metadata_roundtrip_with_step(tmp_path: Path) -> None:
+    mgr = _make_manager(tmp_path)
+    ckpt_dir = _make_checkpoint_dir(tmp_path, 7)
+    mgr._write_metadata(ckpt_dir, step=7)
+    loaded = CheckpointManager.load_metadata(ckpt_dir)
+    assert loaded == {"step": 7}
+
+
+def test_write_metadata_includes_raw_config(tmp_path: Path) -> None:
+    raw = {
+        "models": {
+            "student": {
+                "_target_": "X"
+            }
+        },
+        "training": {
+            "distributed": {
+                "num_gpus": 4
+            }
+        },
+    }
+    mgr = _make_manager(tmp_path, raw_config=raw)
+    ckpt_dir = _make_checkpoint_dir(tmp_path, 7)
+    mgr._write_metadata(ckpt_dir, step=7)
+    loaded = CheckpointManager.load_metadata(ckpt_dir)
+    assert loaded["step"] == 7
+    assert loaded["config"] == raw
+
+
+def test_load_metadata_raises_on_missing_file(tmp_path: Path) -> None:
+    ckpt_dir = _make_checkpoint_dir(tmp_path, 7)
+    # No metadata.json written.
+    with pytest.raises(FileNotFoundError, match="metadata"):
+        CheckpointManager.load_metadata(ckpt_dir)
+
+
+# ---------------------------------------------------------------------------
+# F. _cleanup_old_checkpoints (rolling delete)
+# ---------------------------------------------------------------------------
+
+
+def test_cleanup_keep_last_zero_is_noop(tmp_path: Path) -> None:
+    mgr = _make_manager(tmp_path, keep_last=0)
+    for step in (1, 2, 3):
+        _make_checkpoint_dir(tmp_path, step)
+    mgr._cleanup_old_checkpoints()
+    remaining = sorted(p.name for p in tmp_path.iterdir())
+    assert remaining == ["checkpoint-1", "checkpoint-2", "checkpoint-3"]
+
+
+def test_cleanup_keeps_newest_when_over_limit(tmp_path: Path) -> None:
+    mgr = _make_manager(tmp_path, keep_last=2)
+    for step in (1, 5, 10, 50, 100):
+        _make_checkpoint_dir(tmp_path, step)
+    mgr._cleanup_old_checkpoints()
+    remaining = sorted(p.name for p in tmp_path.iterdir())
+    assert remaining == ["checkpoint-100", "checkpoint-50"]
+
+
+def test_cleanup_no_op_when_under_limit(tmp_path: Path) -> None:
+    mgr = _make_manager(tmp_path, keep_last=3)
+    for step in (1, 2):
+        _make_checkpoint_dir(tmp_path, step)
+    mgr._cleanup_old_checkpoints()
+    remaining = sorted(p.name for p in tmp_path.iterdir())
+    assert remaining == ["checkpoint-1", "checkpoint-2"]
+
+
+def test_cleanup_skips_non_checkpoint_dirs(tmp_path: Path) -> None:
+    mgr = _make_manager(tmp_path, keep_last=1)
+    for step in (1, 2, 3):
+        _make_checkpoint_dir(tmp_path, step)
+    (tmp_path / "logs").mkdir()
+    (tmp_path / "wandb").mkdir()
+    mgr._cleanup_old_checkpoints()
+    remaining = sorted(p.name for p in tmp_path.iterdir())
+    assert remaining == ["checkpoint-3", "logs", "wandb"]
+
+
+# ---------------------------------------------------------------------------
+# G. maybe_save gating logic
+# ---------------------------------------------------------------------------
+
+
+def _record_save_calls(mgr: CheckpointManager) -> list[int]:
+    """Replace ``mgr.save`` with a recorder that mimics the side effect
+    of advancing ``_last_saved_step`` so dedup logic still works."""
+    calls: list[int] = []
+
+    def fake_save(step: int) -> None:
+        calls.append(step)
+        mgr._last_saved_step = step
+
+    mgr.save = fake_save  # type: ignore[method-assign]
+    return calls
+
+
+def test_maybe_save_skipped_when_save_steps_is_zero(tmp_path: Path) -> None:
+    mgr = _make_manager(tmp_path, save_steps=0)
+    calls = _record_save_calls(mgr)
+    mgr.maybe_save(step=10)
+    mgr.maybe_save(step=100)
+    assert calls == []
+
+
+def test_maybe_save_skipped_when_step_not_on_interval(tmp_path: Path) -> None:
+    mgr = _make_manager(tmp_path, save_steps=10)
+    calls = _record_save_calls(mgr)
+    for step in (1, 5, 9, 11, 15):
+        mgr.maybe_save(step=step)
+    assert calls == []
+
+
+def test_maybe_save_dedupes_on_repeated_call(tmp_path: Path) -> None:
+    mgr = _make_manager(tmp_path, save_steps=10)
+    calls = _record_save_calls(mgr)
+    mgr.maybe_save(step=20)
+    mgr.maybe_save(step=20)
+    mgr.maybe_save(step=20)
+    assert calls == [20]
+
+
+def test_maybe_save_triggers_on_each_interval(tmp_path: Path) -> None:
+    mgr = _make_manager(tmp_path, save_steps=10)
+    calls = _record_save_calls(mgr)
+    for step in range(1, 41):
+        mgr.maybe_save(step=step)
+    assert calls == [10, 20, 30, 40]

fastvideo/tests/train/utils/test_config.py

@@ -0,0 +1,314 @@
+# SPDX-License-Identifier: Apache-2.0
+"""CPU-only unit tests for :func:`load_run_config`."""
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any
+
+import pytest
+import yaml
+
+from fastvideo.train.utils.config import RunConfig, load_run_config
+from fastvideo.train.utils.training_config import TrainingConfig
+
+
+def _write_yaml(tmp_path: Path, data: dict[str, Any]) -> str:
+    path = tmp_path / "run.yaml"
+    path.write_text(yaml.safe_dump(data), encoding="utf-8")
+    return str(path)
+
+
+def _minimal_yaml() -> dict[str, Any]:
+    return {
+        "models": {
+            "student": {
+                "_target_": "fastvideo.train.models.wan.WanModel",
+                "init_from": "fake/model",
+            },
+        },
+        "method": {
+            "_target_":
+            "fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod",
+        },
+        "training": {},
+    }
+
+
+# ---------------------------------------------------------------------------
+# Happy path
+# ---------------------------------------------------------------------------
+
+
+def test_minimal_yaml_loads_happy_path(tmp_path: Path) -> None:
+    cfg = load_run_config(_write_yaml(tmp_path, _minimal_yaml()))
+
+    assert isinstance(cfg, RunConfig)
+    assert isinstance(cfg.training, TrainingConfig)
+    assert cfg.method["_target_"] == (
+        "fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod")
+    assert "student" in cfg.models
+    assert cfg.callbacks == {}
+    # raw retains the original YAML dict for downstream logging.
+    assert "models" in cfg.raw and "method" in cfg.raw
+
+
+def test_minimal_yaml_applies_all_defaults(tmp_path: Path) -> None:
+    cfg = load_run_config(_write_yaml(tmp_path, _minimal_yaml()))
+    t = cfg.training
+
+    assert t.distributed.num_gpus == 1
+    assert t.distributed.tp_size == 1
+    assert t.distributed.sp_size == 1
+    assert t.distributed.hsdp_replicate_dim == 1
+    assert t.distributed.pin_cpu_memory is False
+
+    assert t.data.train_batch_size == 1
+    assert t.data.dataloader_num_workers == 0
+    assert t.data.training_cfg_rate == 0.0
+    assert t.data.seed == 0
+
+    assert t.optimizer.learning_rate == 0.0
+    assert t.optimizer.betas == (0.9, 0.999)
+    assert t.optimizer.weight_decay == 0.0
+    assert t.optimizer.lr_scheduler == "constant"
+    assert t.optimizer.min_lr_ratio == 0.5
+
+    assert t.loop.max_train_steps == 0
+    assert t.loop.gradient_accumulation_steps == 1
+
+    assert t.checkpoint.output_dir == ""
+    assert t.checkpoint.checkpoints_total_limit == 0
+
+    assert t.tracker.trackers == []
+    assert t.tracker.project_name == "fastvideo"
+    assert t.tracker.run_name == ""
+
+    assert t.model.weighting_scheme == "uniform"
+    assert t.model.precondition_outputs is False
+    assert t.model.moba_config == {}
+
+    assert t.dit_precision == "fp32"
+    assert t.vsa_sparsity == 0.0
+    assert t.pipeline_config is None
+
+
+def test_full_yaml_populates_all_training_fields(tmp_path: Path) -> None:
+    data = _minimal_yaml()
+    data["training"] = {
+        "distributed": {
+            "num_gpus": 4,
+            "tp_size": 2,
+            "sp_size": 2,
+            "hsdp_replicate_dim": 2,
+            "hsdp_shard_dim": 2,
+            "pin_cpu_memory": True,
+        },
+        "data": {
+            "data_path": "/some/path",
+            "train_batch_size": 2,
+            "dataloader_num_workers": 4,
+            "training_cfg_rate": 0.1,
+            "seed": 42,
+            "num_height": 256,
+            "num_width": 512,
+            "num_latent_t": 8,
+            "num_frames": 33,
+        },
+        "optimizer": {
+            "learning_rate": 1e-4,
+            "betas": [0.9, 0.95],
+            "weight_decay": 0.01,
+            "lr_scheduler": "cosine",
+            "lr_warmup_steps": 100,
+            "min_lr_ratio": 0.1,
+        },
+        "loop": {
+            "max_train_steps": 1000,
+            "gradient_accumulation_steps": 4,
+        },
+        "checkpoint": {
+            "output_dir": "/out",
+            "training_state_checkpointing_steps": 50,
+            "checkpoints_total_limit": 3,
+        },
+        "tracker": {
+            "trackers": ["wandb"],
+            "project_name": "myproj",
+            "run_name": "myrun",
+        },
+        "vsa": {
+            "sparsity": 0.5
+        },
+        "model": {
+            "weighting_scheme": "logit_normal",
+            "logit_mean": 0.5,
+            "logit_std": 1.5,
+            "precondition_outputs": True,
+        },
+        "dit_precision": "bf16",
+    }
+    cfg = load_run_config(_write_yaml(tmp_path, data))
+    t = cfg.training
+
+    assert t.distributed.num_gpus == 4
+    assert t.distributed.tp_size == 2
+    assert t.distributed.pin_cpu_memory is True
+
+    assert t.data.train_batch_size == 2
+    assert t.data.num_frames == 33
+    assert t.data.seed == 42
+
+    assert t.optimizer.learning_rate == pytest.approx(1e-4)
+    assert t.optimizer.betas == (0.9, 0.95)
+    assert t.optimizer.lr_scheduler == "cosine"
+    assert t.optimizer.lr_warmup_steps == 100
+    assert t.optimizer.min_lr_ratio == pytest.approx(0.1)
+
+    assert t.loop.max_train_steps == 1000
+    assert t.loop.gradient_accumulation_steps == 4
+
+    assert t.checkpoint.output_dir == "/out"
+    assert t.checkpoint.checkpoints_total_limit == 3
+
+    assert t.tracker.trackers == ["wandb"]
+    assert t.tracker.project_name == "myproj"
+
+    assert t.vsa_sparsity == pytest.approx(0.5)
+    assert t.model.weighting_scheme == "logit_normal"
+    assert t.model.precondition_outputs is True
+    assert t.dit_precision == "bf16"
+
+
+# ---------------------------------------------------------------------------
+# Schema validation
+# ---------------------------------------------------------------------------
+
+
+def test_missing_models_raises(tmp_path: Path) -> None:
+    data = _minimal_yaml()
+    del data["models"]
+    with pytest.raises(ValueError, match="models"):
+        load_run_config(_write_yaml(tmp_path, data))
+
+
+def test_missing_method_raises(tmp_path: Path) -> None:
+    data = _minimal_yaml()
+    del data["method"]
+    with pytest.raises(ValueError, match="method"):
+        load_run_config(_write_yaml(tmp_path, data))
+
+
+def test_missing_training_raises(tmp_path: Path) -> None:
+    data = _minimal_yaml()
+    del data["training"]
+    with pytest.raises(ValueError, match="training"):
+        load_run_config(_write_yaml(tmp_path, data))
+
+
+def test_models_role_without_target_raises(tmp_path: Path) -> None:
+    data = _minimal_yaml()
+    data["models"]["student"] = {"init_from": "fake/model"}
+    with pytest.raises(ValueError, match="_target_"):
+        load_run_config(_write_yaml(tmp_path, data))
+
+
+def test_method_without_target_raises(tmp_path: Path) -> None:
+    data = _minimal_yaml()
+    data["method"] = {"some_param": 5}
+    with pytest.raises(ValueError, match="_target_"):
+        load_run_config(_write_yaml(tmp_path, data))
+
+
+def test_missing_config_file_raises(tmp_path: Path) -> None:
+    with pytest.raises(FileNotFoundError):
+        load_run_config(str(tmp_path / "does_not_exist.yaml"))
+
+
+# ---------------------------------------------------------------------------
+# Special parsing
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("betas_value, expected", [
+    ([0.8, 0.9], (0.8, 0.9)),
+    ("0.9,0.999", (0.9, 0.999)),
+])
+def test_betas_parses_list_and_string_forms(
+    tmp_path: Path,
+    betas_value: Any,
+    expected: tuple[float, float],
+) -> None:
+    data = _minimal_yaml()
+    data["training"] = {"optimizer": {"betas": betas_value}}
+    cfg = load_run_config(_write_yaml(tmp_path, data))
+    assert cfg.training.optimizer.betas == expected
+
+
+def test_dotted_overrides_apply_with_type_coercion(tmp_path: Path) -> None:
+    path = _write_yaml(tmp_path, _minimal_yaml())
+    overrides = [
+        "--training.distributed.num_gpus=4",
+        "--training.optimizer.learning_rate=1e-3",
+        "--training.distributed.pin_cpu_memory=true",
+        "--training.tracker.project_name=overridden",
+    ]
+    cfg = load_run_config(path, overrides=overrides)
+
+    assert cfg.training.distributed.num_gpus == 4
+    assert cfg.training.optimizer.learning_rate == pytest.approx(1e-3)
+    assert cfg.training.distributed.pin_cpu_memory is True
+    assert cfg.training.tracker.project_name == "overridden"
+
+
+def test_dotted_overrides_accept_separate_value_token(tmp_path: Path) -> None:
+    path = _write_yaml(tmp_path, _minimal_yaml())
+    cfg = load_run_config(
+        path,
+        overrides=["--training.distributed.num_gpus", "8"],
+    )
+    assert cfg.training.distributed.num_gpus == 8
+
+
+def test_overrides_create_intermediate_keys(tmp_path: Path) -> None:
+    """Overrides into a nested key absent from YAML should still apply."""
+    data = _minimal_yaml()
+    # No `training.checkpoint` block in the minimal YAML.
+    path = _write_yaml(tmp_path, data)
+    cfg = load_run_config(
+        path,
+        overrides=["--training.checkpoint.checkpoints_total_limit=5"],
+    )
+    assert cfg.training.checkpoint.checkpoints_total_limit == 5
+
+
+def test_hsdp_shard_dim_defaults_to_num_gpus(tmp_path: Path) -> None:
+    """When unset, hsdp_shard_dim and sp_size fall back to num_gpus."""
+    data = _minimal_yaml()
+    data["training"] = {"distributed": {"num_gpus": 4}}
+    cfg = load_run_config(_write_yaml(tmp_path, data))
+    assert cfg.training.distributed.hsdp_shard_dim == 4
+    assert cfg.training.distributed.sp_size == 4
+
+
+def test_model_path_derived_from_student_init_from(tmp_path: Path) -> None:
+    cfg = load_run_config(_write_yaml(tmp_path, _minimal_yaml()))
+    assert cfg.training.model_path == "fake/model"
+
+
+def test_callbacks_default_to_empty_when_absent(tmp_path: Path) -> None:
+    cfg = load_run_config(_write_yaml(tmp_path, _minimal_yaml()))
+    assert cfg.callbacks == {}
+
+
+def test_callbacks_passed_through_when_present(tmp_path: Path) -> None:
+    data = _minimal_yaml()
+    data["callbacks"] = {
+        "grad_clip": {
+            "_target_":
+            "fastvideo.train.callbacks.grad_clip.GradNormClipCallback",
+            "max_grad_norm": 1.0,
+        },
+    }
+    cfg = load_run_config(_write_yaml(tmp_path, data))
+    assert "grad_clip" in cfg.callbacks
+    assert cfg.callbacks["grad_clip"]["max_grad_norm"] == 1.0

fastvideo/train/callbacks/callback.py

@@ -38,6 +38,11 @@ class Callback:
     training_config: TrainingConfig
     method: TrainingMethod
     _callback_dict: CallbackDict | None
+    # Yaml dict key under which this callback was declared (e.g.
+    # "validation_short").  Set by ``CallbackDict`` after instantiation.
+    # Useful for callbacks that want to disambiguate themselves from
+    # sibling instances in tracker keys, log paths, etc.
+    name: str = ""
 
     def on_train_start(
         self,
@@ -132,6 +137,7 @@ class CallbackDict:
                                 f"Callback subclass.")
             cb.training_config = training_config
             cb._callback_dict = self
+            cb.name = name
             self._callbacks[name] = cb
 
     def __getattr__(

fastvideo/train/callbacks/validation.py

@@ -288,7 +288,8 @@ class ValidationCallback(Callback):
             "sp_size": tc.distributed.sp_size,
             "num_gpus": tc.distributed.num_gpus,
             "pin_cpu_memory": (tc.distributed.pin_cpu_memory),
-            "dit_cpu_offload": True,
+            "dit_cpu_offload": False,
+            "dit_layerwise_offload": False,
         }
         if flow_shift is not None:
             kwargs["flow_shift"] = float(flow_shift)

fastvideo/train/models/cosmos/__init__.py

@@ -0,0 +1,5 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Cosmos model plugin package."""
+
+from fastvideo.train.models.cosmos.cosmos import (
+    CosmosModel as CosmosModel, )

fastvideo/train/models/cosmos/cosmos.py

@@ -0,0 +1,364 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Cosmos model plugin (per-role instance).
+
+Subclasses WanModel since Cosmos uses the same
+FlowMatchEulerDiscreteScheduler.  Differences:
+  - transformer class name: CosmosTransformer3DModel
+  - normalize_dit_input("cosmos", ...) instead of ("wan", ...)
+  - forward kwargs: no encoder_attention_mask, needs
+    condition_mask + padding_mask + fps
+  - hidden_states in (B,C,T,H,W) — no permute needed
+  - default flow_shift = 1.0
+  - single T5 text encoder (not dual like Hunyuan)
+"""
+
+from __future__ import annotations
+
+import copy
+from typing import Any, Literal, TYPE_CHECKING
+
+import torch
+
+from fastvideo.forward_context import set_forward_context
+from fastvideo.pipelines import TrainingBatch
+from fastvideo.training.training_utils import (
+    normalize_dit_input, )
+
+from fastvideo.train.models.wan.wan import WanModel
+
+if TYPE_CHECKING:
+    from fastvideo.train.utils.training_config import (
+        TrainingConfig, )
+
+
+class CosmosModel(WanModel):
+    """Cosmos 2.5 per-role model.
+
+    Inherits most behaviour from WanModel (noise scheduler,
+    timestep sampling, attention metadata, backward).  Overrides
+    only the pieces that differ for Cosmos 2.5.
+
+    Cosmos 2.5 uses:
+    - Cosmos25Transformer3DModel (velocity prediction)
+    - EDM noise schedule: x_t = x_0 + sigma * eps
+    - No input/output preconditioning (raw latents)
+    - Timestep = raw sigma value
+    - Model output = velocity ≈ noise
+    """
+
+    _transformer_cls_name: str = "Cosmos25Transformer3DModel"
+
+    def __init__(
+        self,
+        *,
+        init_from: str,
+        training_config: TrainingConfig,
+        trainable: bool = True,
+        disable_custom_init_weights: bool = False,
+        flow_shift: float = 1.0,
+        enable_gradient_checkpointing_type: str
+        | None = None,
+        transformer_override_safetensor: str
+        | None = None,
+    ) -> None:
+        super().__init__(
+            init_from=init_from,
+            training_config=training_config,
+            trainable=trainable,
+            disable_custom_init_weights=(disable_custom_init_weights),
+            flow_shift=flow_shift,
+            enable_gradient_checkpointing_type=(enable_gradient_checkpointing_type),
+            transformer_override_safetensor=(transformer_override_safetensor),
+        )
+
+    # ------------------------------------------------------------------
+    # Overrides
+    # ------------------------------------------------------------------
+
+    def prepare_batch(
+        self,
+        raw_batch: dict[str, Any],
+        *,
+        generator: torch.Generator,
+        latents_source: Literal["data", "zeros"] = "data",
+    ) -> TrainingBatch:
+        """Same flow as Wan, but uses Cosmos VAE
+        normalisation."""
+        self.ensure_negative_conditioning()
+        assert self.training_config is not None
+        tc = self.training_config
+
+        dtype = self._get_training_dtype()
+        device = self.device
+
+        training_batch = TrainingBatch()
+        encoder_hidden_states = raw_batch["text_embedding"]
+        encoder_attention_mask = raw_batch["text_attention_mask"]
+        infos = raw_batch.get("info_list")
+
+        if latents_source == "zeros":
+            batch_size = encoder_hidden_states.shape[0]
+            vae_config = (
+                tc.pipeline_config.vae_config  # type: ignore[union-attr]
+                .arch_config)
+            num_channels = getattr(
+                vae_config,
+                "z_dim",
+                getattr(vae_config, "latent_channels", 16),
+            )
+            spatial_compression_ratio = (vae_config.spatial_compression_ratio)
+            latent_height = (tc.data.num_height // spatial_compression_ratio)
+            latent_width = (tc.data.num_width // spatial_compression_ratio)
+            latents = torch.zeros(
+                batch_size,
+                num_channels,
+                tc.data.num_latent_t,
+                latent_height,
+                latent_width,
+                device=device,
+                dtype=dtype,
+            )
+        elif latents_source == "data":
+            if "vae_latent" not in raw_batch:
+                raise ValueError("vae_latent not found in batch "
+                                 "and latents_source='data'")
+            latents = raw_batch["vae_latent"]
+            latents = latents[:, :, :tc.data.num_latent_t]
+            latents = latents.to(device, dtype=dtype)
+        else:
+            raise ValueError(f"Unknown latents_source: "
+                             f"{latents_source!r}")
+
+        training_batch.latents = latents
+        training_batch.encoder_hidden_states = (encoder_hidden_states.to(device, dtype=dtype))
+        training_batch.encoder_attention_mask = (encoder_attention_mask.to(device, dtype=dtype))
+        training_batch.infos = infos
+
+        # KEY DIFFERENCE: "cosmos" normalisation
+        training_batch.latents = normalize_dit_input(
+            "cosmos",
+            training_batch.latents,
+            self.vae,
+        )
+        training_batch = self._prepare_dit_inputs(training_batch, generator)
+        training_batch = self._build_attention_metadata(training_batch)
+
+        training_batch.attn_metadata_vsa = copy.deepcopy(training_batch.attn_metadata)
+        if training_batch.attn_metadata is not None:
+            training_batch.attn_metadata.VSA_sparsity = 0.0  # type: ignore[attr-defined]
+
+        return training_batch
+
+    def predict_noise(
+        self,
+        noisy_latents: torch.Tensor,
+        timestep: torch.Tensor,
+        batch: TrainingBatch,
+        *,
+        conditional: bool,
+        cfg_uncond: dict[str, Any] | None = None,
+        attn_kind: Literal["dense", "vsa"] = "dense",
+    ) -> torch.Tensor:
+        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}")
+
+        # finetune.py hands us `noisy_latents` in (B, T, C, H, W)
+        # (Wan canonical). Un-permute back to the (B, C, T, H, W)
+        # layout Cosmos's transformer expects.
+        noisy_latents = noisy_latents.permute(0, 2, 1, 3, 4)
+
+        # ----------------------------------------------------------
+        # Cosmos 2.5: velocity prediction, no preconditioning.
+        #
+        # Noise schedule (EDM):  x_t = x_0 + sigma * eps
+        # Model input:   raw x_t (no c_in scaling)
+        # Timestep:      raw sigma value
+        # Model output:  velocity v ≈ noise
+        #
+        # finetune.py (precondition_outputs=True) computes:
+        #   pred_x0 = x_t - pred * sigma
+        # Since v ≈ noise, this gives:
+        #   pred_x0 = (x_0 + sigma*eps) - eps*sigma = x_0  ✓
+        # ----------------------------------------------------------
+        assert batch.sigmas is not None
+        sigma_1d = batch.sigmas.flatten()[:noisy_latents.shape[0]]
+
+        # Cosmos 2.5 timestep = raw sigma (matching inference
+        # convention where scheduler timestep * 0.001 = sigma).
+        cosmos_timestep = sigma_1d
+
+        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,
+                cosmos_timestep,
+                text_dict,
+            )
+            transformer = self._get_transformer(timestep)
+            model_output = transformer(**input_kwargs)
+
+        # Raw velocity output (≈ noise).
+        # Re-permute (B, C, T, H, W) → (B, T, C, H, W)
+        pred = model_output.permute(0, 2, 1, 3, 4)
+        return pred
+
+    def _build_distill_input_kwargs(
+        self,
+        noise_input: torch.Tensor,
+        timestep: torch.Tensor,
+        text_dict: dict[str, torch.Tensor] | None,
+    ) -> dict[str, Any]:
+        """Build transformer forward kwargs for Cosmos 2.5.
+
+        Cosmos 2.5 transformer expects:
+        - hidden_states in (B, C, T, H, W)
+        - timestep as (B, 1) for T2W (scalar per sample)
+        - condition_mask and padding_mask
+        - fps tensor
+        """
+        if text_dict is None:
+            raise ValueError("text_dict cannot be None for Cosmos forward")
+
+        b, c, t, h, w = noise_input.shape
+
+        # For T2W fine-tuning: no conditioning frames.
+        condition_mask = torch.zeros(
+            b,
+            1,
+            t,
+            h,
+            w,
+            device=noise_input.device,
+            dtype=noise_input.dtype,
+        )
+
+        # Padding mask: zeros (no padding).
+        padding_mask = torch.zeros(
+            1,
+            1,
+            h,
+            w,
+            device=noise_input.device,
+            dtype=noise_input.dtype,
+        )
+
+        # Cosmos 2.5 expects timestep as (B, 1) for T2W.
+        if timestep.ndim == 1:
+            timestep = timestep.unsqueeze(1)
+
+        return {
+            "hidden_states": noise_input,
+            "encoder_hidden_states": (text_dict["encoder_hidden_states"]),
+            "timestep": timestep,
+            "condition_mask": condition_mask,
+            "padding_mask": padding_mask,
+            "fps": 16,
+        }
+
+    def _prepare_dit_inputs(
+        self,
+        training_batch: TrainingBatch,
+        generator: torch.Generator,
+    ) -> TrainingBatch:
+        """Prepare DiT inputs for Cosmos 2.5.
+
+        Uses the same flow-matching noise schedule as the parent
+        WanModel but with the Karras sigma schedule from the
+        Cosmos scheduler.  The noisy input is:
+
+            x_t = (1 - sigma) * x_0 + sigma * noise
+
+        where sigma comes from the scheduler's Karras schedule.
+        The model predicts the flow-matching velocity
+        ``v = noise - x_0``.
+        """
+        # Delegate entirely to the parent WanModel implementation.
+        # This uses the noise scheduler's sigmas (Karras, [0..80])
+        # with standard flow-matching interpolation:
+        #   x_t = (1 - sigma) * x_0 + sigma * noise
+        training_batch = super()._prepare_dit_inputs(
+            training_batch,
+            generator,
+        )
+        return training_batch
+
+    def ensure_negative_conditioning(self) -> None:
+        """Create negative (unconditional) prompt embeddings.
+
+        Cosmos 2.5 uses Reason1 (Qwen2.5-VL) which is expensive
+        to load.  This method only supports ``training_cfg_rate=0``
+        (no classifier-free guidance dropout), in which case the
+        negative embedding is never used and a zero placeholder
+        sized to match the text embedding dimension is sufficient.
+        ``training_cfg_rate>0`` would require real Reason1 negative
+        embeddings and is rejected here to avoid silently training
+        with zero-vector "unconditional" inputs.
+        """
+        if self.negative_prompt_embeds is not None:  # type: ignore[has-type]
+            return
+
+        assert self.training_config is not None
+        tc = self.training_config
+
+        cfg_rate = float(tc.data.training_cfg_rate or 0.0)
+        if cfg_rate > 0.0:
+            raise NotImplementedError("Cosmos 2.5 currently only supports training_cfg_rate=0; "
+                                      f"got training_cfg_rate={cfg_rate}. Real negative-prompt "
+                                      "embeddings via Reason1 (Qwen2.5-VL) are not implemented "
+                                      "yet — using the zero placeholder with CFG dropout would "
+                                      "train against zero-vector \"unconditional\" inputs and "
+                                      "produce wrong gradients. Set "
+                                      "training.data.training_cfg_rate=0.")
+
+        device = self.device
+        dtype = self._get_training_dtype()
+
+        # Infer embedding dimension from the pipeline config's
+        # text encoder settings, or fall back to a reasonable
+        # default for Cosmos 2.5 (Reason1 full_concat: 100352).
+        text_enc_cfgs = tc.pipeline_config.text_encoder_configs
+        if text_enc_cfgs:
+            arch = text_enc_cfgs[0].arch_config
+            embed_dim = getattr(arch, "hidden_size", 100352)
+        else:
+            embed_dim = 100352
+
+        num_tokens = 512  # Reason1 default padding length
+
+        neg_embeds = torch.zeros(
+            1,
+            num_tokens,
+            embed_dim,
+            device=device,
+            dtype=dtype,
+        )
+        neg_mask = torch.ones(
+            1,
+            num_tokens,
+            device=device,
+            dtype=dtype,
+        )
+
+        self.negative_prompt_embeds = neg_embeds
+        self.negative_prompt_attention_mask = neg_mask

fastvideo/train/models/wan/wan.py

@@ -4,7 +4,6 @@
 from __future__ import annotations
 
 import copy
-import gc
 from typing import Any, Literal, TYPE_CHECKING
 
 import torch
@@ -19,10 +18,6 @@ from fastvideo.forward_context import set_forward_context
 from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import (
     FlowMatchEulerDiscreteScheduler, )
 from fastvideo.pipelines import TrainingBatch
-from fastvideo.pipelines.basic.wan.wan_pipeline import (
-    WanPipeline, )
-from fastvideo.pipelines.pipeline_batch_info import (
-    ForwardBatch, )
 from fastvideo.training.activation_checkpoint import (
     apply_activation_checkpointing, )
 from fastvideo.training.training_utils import (
@@ -41,6 +36,7 @@ from fastvideo.train.utils.module_state import (
     apply_trainable, )
 from fastvideo.train.utils.moduleloader import (
     load_module_from_path, )
+from fastvideo.train.utils.negative_prompt import encode_negative_prompt
 
 if TYPE_CHECKING:
     from fastvideo.train.utils.training_config import (
@@ -341,98 +337,15 @@ class WanModel(ModelBase):
 
         assert self.training_config is not None
         tc = self.training_config
-        world_group = self.world_group
-        device = self.device
-        dtype = self._get_training_dtype()
-
-        from fastvideo.train.utils.moduleloader import (
-            make_inference_args, )
-
-        neg_embeds: torch.Tensor | None = None
-        neg_mask: torch.Tensor | None = None
-
-        if world_group.rank_in_group == 0:
-            sampling_param = SamplingParam.from_pretrained(tc.model_path)
-            negative_prompt = sampling_param.negative_prompt
-
-            inference_args = make_inference_args(tc, model_path=tc.model_path)
-
-            prompt_pipeline = WanPipeline.from_pretrained(
-                tc.model_path,
-                args=inference_args,
-                inference_mode=True,
-                loaded_modules={"transformer": self.transformer},
-                tp_size=tc.distributed.tp_size,
-                sp_size=tc.distributed.sp_size,
-                num_gpus=tc.distributed.num_gpus,
-                pin_cpu_memory=(tc.distributed.pin_cpu_memory),
-                dit_cpu_offload=True,
-            )
-
-            batch_negative = ForwardBatch(
-                data_type="video",
-                prompt=negative_prompt,
-                prompt_embeds=[],
-                prompt_attention_mask=[],
-            )
-            result_batch = prompt_pipeline.prompt_encoding_stage(  # type: ignore[attr-defined]
-                batch_negative,
-                inference_args,
-            )
-
-            neg_embeds = result_batch.prompt_embeds[0].to(device=device, dtype=dtype)
-            neg_mask = (result_batch.prompt_attention_mask[0].to(device=device, dtype=dtype))
-
-            del prompt_pipeline
-            gc.collect()
-            if torch.cuda.is_available():
-                torch.cuda.empty_cache()
-
-        meta = torch.zeros((2, ), device=device, dtype=torch.int64)
-        if world_group.rank_in_group == 0:
-            assert neg_embeds is not None
-            assert neg_mask is not None
-            meta[0] = neg_embeds.ndim
-            meta[1] = neg_mask.ndim
-        world_group.broadcast(meta, src=0)
-        embed_ndim, mask_ndim = (
-            int(meta[0].item()),
-            int(meta[1].item()),
+        sampling_param = SamplingParam.from_pretrained(tc.model_path)
+        embeds, mask = encode_negative_prompt(
+            tc,
+            prompt=sampling_param.negative_prompt,
+            device=self.device,
+            dtype=self._get_training_dtype(),
         )
-
-        max_ndim = 8
-        embed_shape = torch.full((max_ndim, ), -1, device=device, dtype=torch.int64)
-        mask_shape = torch.full((max_ndim, ), -1, device=device, dtype=torch.int64)
-        if world_group.rank_in_group == 0:
-            assert neg_embeds is not None
-            assert neg_mask is not None
-            embed_shape[:embed_ndim] = torch.tensor(
-                list(neg_embeds.shape),
-                device=device,
-                dtype=torch.int64,
-            )
-            mask_shape[:mask_ndim] = torch.tensor(
-                list(neg_mask.shape),
-                device=device,
-                dtype=torch.int64,
-            )
-        world_group.broadcast(embed_shape, src=0)
-        world_group.broadcast(mask_shape, src=0)
-
-        embed_sizes = tuple(int(x) for x in embed_shape[:embed_ndim].tolist())
-        mask_sizes = tuple(int(x) for x in mask_shape[:mask_ndim].tolist())
-
-        if world_group.rank_in_group != 0:
-            neg_embeds = torch.empty(embed_sizes, device=device, dtype=dtype)
-            neg_mask = torch.empty(mask_sizes, device=device, dtype=dtype)
-        assert neg_embeds is not None
-        assert neg_mask is not None
-
-        world_group.broadcast(neg_embeds, src=0)
-        world_group.broadcast(neg_mask, src=0)
-
-        self.negative_prompt_embeds = neg_embeds
-        self.negative_prompt_attention_mask = neg_mask
+        self.negative_prompt_embeds = embeds
+        self.negative_prompt_attention_mask = mask
 
     def _sample_timesteps(
         self,

fastvideo/train/utils/negative_prompt.py

@@ -0,0 +1,99 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Per-rank negative-prompt encoding shared by training model plugins.
+
+Encoding the negative prompt only on rank 0 and broadcasting (the
+previous Wan path) ran ``Pipeline.from_pretrained`` asymmetrically across
+ranks, which deadlocked on any collective fired during text-encoder load
+(FSDP device-mesh init, weight broadcast, etc.). The text encoder is
+small and only loaded once at startup, so loading it on every rank
+sidesteps the deadlock entirely.
+"""
+
+from __future__ import annotations
+
+import os
+from typing import TYPE_CHECKING
+
+import torch
+from transformers import AutoTokenizer
+
+from fastvideo.forward_context import set_forward_context
+from fastvideo.models.loader.component_loader import TextEncoderLoader
+from fastvideo.train.utils.moduleloader import make_inference_args
+from fastvideo.utils import maybe_download_model
+
+if TYPE_CHECKING:
+    from fastvideo.train.utils.training_config import TrainingConfig
+
+
+def encode_negative_prompt(
+    training_config: TrainingConfig,
+    *,
+    prompt: str,
+    device: torch.device,
+    dtype: torch.dtype,
+    encoder_index: int = 0,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Per-rank encode of ``prompt`` using encoder ``encoder_index``.
+
+    Reads ``pipeline_config.text_encoder_configs[encoder_index]`` so the
+    encoder class (e.g. UMT5 for Wan) and tokenizer kwargs match the
+    inference path, and applies the matching ``postprocess_text_funcs``
+    entry. Returns ``(embeds, mask)`` on ``device`` cast to ``dtype``.
+    """
+    tc = training_config
+    pipeline_config = tc.pipeline_config
+    if pipeline_config is None:
+        raise ValueError("training_config.pipeline_config is required for negative "
+                         "prompt encoding")
+
+    encoder_configs = pipeline_config.text_encoder_configs
+    postprocess_funcs = pipeline_config.postprocess_text_funcs
+    preprocess_funcs = getattr(pipeline_config, "preprocess_text_funcs", None)
+
+    if encoder_index < 0 or encoder_index >= len(encoder_configs):
+        raise IndexError(f"encoder_index {encoder_index} out of range for "
+                         f"text_encoder_configs (len={len(encoder_configs)})")
+    encoder_config = encoder_configs[encoder_index]
+    postprocess_text = postprocess_funcs[encoder_index]
+    preprocess_text = (preprocess_funcs[encoder_index] if preprocess_funcs is not None else None)
+
+    # HF convention: text_encoder / tokenizer for index 0,
+    # text_encoder_2 / tokenizer_2 for index 1, etc.
+    suffix = "" if encoder_index == 0 else f"_{encoder_index + 1}"
+    encoder_subdir = f"text_encoder{suffix}"
+    tokenizer_subdir = f"tokenizer{suffix}"
+
+    model_path = maybe_download_model(tc.model_path)
+    inference_args = make_inference_args(tc, model_path=model_path)
+    # Keep the encoder on-device; CPU offload would init an FSDP device
+    # mesh and reintroduce the collective at load time.
+    inference_args.text_encoder_cpu_offload = False
+
+    loader = TextEncoderLoader()
+    text_encoder = loader.load(
+        os.path.join(model_path, encoder_subdir),
+        inference_args,
+    ).to(device).eval()
+    tokenizer = AutoTokenizer.from_pretrained(os.path.join(model_path, tokenizer_subdir))
+
+    tok_kwargs = dict(encoder_config.tokenizer_kwargs)
+    text = preprocess_text(prompt) if preprocess_text is not None else prompt
+
+    with torch.no_grad(), set_forward_context(
+            current_timestep=0,
+            attn_metadata=None,
+    ):
+        text_inputs = tokenizer(text, **tok_kwargs).to(device)
+        outputs = text_encoder(
+            input_ids=text_inputs.input_ids,
+            attention_mask=text_inputs.attention_mask,
+        )
+        # Mirror TextEncodingStage: postprocess reads outputs.attention_mask.
+        outputs.attention_mask = text_inputs["attention_mask"]
+        embeds = postprocess_text(outputs).to(device=device, dtype=dtype)
+        mask = text_inputs["attention_mask"].to(device=device, dtype=dtype)
+
+    del text_encoder, tokenizer
+
+    return embeds, mask

fastvideo/training/distillation_pipeline.py

@@ -742,8 +742,11 @@ class DistillationPipeline(TrainingPipeline):
                 max_grad_norm,
                 foreach=None,
             )
-            assert grad_norm is not float('nan') or grad_norm is not float('inf')
-            grad_norm = grad_norm.item() if grad_norm is not None else 0.0
+            if grad_norm is not None:
+                assert torch.isfinite(grad_norm), (f"grad_norm is not finite: {grad_norm}")
+                grad_norm = grad_norm.item()
+            else:
+                grad_norm = 0.0
         else:
             grad_norm = 0.0
         training_batch.grad_norm = grad_norm

fastvideo/training/training_pipeline.py

@@ -448,8 +448,11 @@ class TrainingPipeline(LoRAPipeline, ABC):
                     max_grad_norm,
                     foreach=None,
                 )
-                assert grad_norm is not float('nan') or grad_norm is not float('inf')
-                grad_norm = grad_norm.item() if grad_norm is not None else 0.0
+                if grad_norm is not None:
+                    assert torch.isfinite(grad_norm), (f"grad_norm is not finite: {grad_norm}")
+                    grad_norm = grad_norm.item()
+                else:
+                    grad_norm = 0.0
         else:
             grad_norm = 0.0
         training_batch.grad_norm = grad_norm

fastvideo/training/training_utils.py

@@ -806,6 +806,25 @@ def normalize_dit_input(model_type, latents, vae) -> torch.Tensor:
         latents_std = latents_std.view(1, -1, 1, 1, 1).to(device=latents.device)
         latents = ((latents.float() - latents_mean) * latents_std).to(latents)
         return latents
+    elif model_type == "cosmos":
+        # Support both diffusers VAE (vae.config.latents_mean)
+        # and Cosmos25WanVAE/Adapter (vae._latents_mean buffer).
+        cfg = getattr(vae, "config", None)
+        if cfg is not None and hasattr(cfg, "latents_mean"):
+            lm = torch.tensor(cfg.latents_mean)
+            ls = torch.tensor(cfg.latents_std)
+        elif hasattr(vae, "_latents_mean"):
+            lm = vae._latents_mean.flatten()
+            ls = vae._latents_std.flatten()
+        elif hasattr(vae, "latents_mean"):
+            lm = torch.tensor(vae.latents_mean)
+            ls = torch.tensor(vae.latents_std)
+        else:
+            raise ValueError("Cannot find latents_mean/latents_std on VAE")
+        latents_mean = lm.view(1, -1, 1, 1, 1).to(device=latents.device)
+        latents_std = ls.view(1, -1, 1, 1, 1).to(device=latents.device)
+        latents = ((latents.float() - latents_mean) / latents_std).to(latents)
+        return latents
     else:
         raise NotImplementedError(f"model_type {model_type} not supported")
 

fastvideo/utils.py

@@ -636,6 +636,23 @@ def maybe_download_model_index(model_name_or_path: str) -> dict[str, Any]:
         raise ValueError(f"Failed to download or parse model_index.json for {model_name_or_path}: {e}") from e
 
 
+_HF_TOKEN_ENV_VARS = ("HF_TOKEN", "HUGGINGFACE_HUB_TOKEN", "HF_API_KEY")
+
+
+def resolve_hf_token() -> str | None:
+    """Return the first non-empty HF token from the standard env vars.
+
+    Order: `HF_TOKEN`, `HUGGINGFACE_HUB_TOKEN`, `HF_API_KEY` (the last is
+    a FastVideo convention; `huggingface_hub` itself doesn't read it).
+    Does not mutate `os.environ`.
+    """
+    for src in _HF_TOKEN_ENV_VARS:
+        v = os.environ.get(src)
+        if v:
+            return v
+    return None
+
+
 def update_environment_variables(envs: dict[str, str]):
     for k, v in envs.items():
         if k in os.environ and os.environ[k] != v:

pyproject.toml

@@ -124,6 +124,8 @@ test = [
     "av",
     "pytorch-msssim==1.0.0",
     "pytest",
+    "pandas",
+    "plotly",
 ]
 
 dev = [ "fastvideo[lint]", "fastvideo[test]", ]

scripts/checkpoint_conversion/stable_audio_to_diffusers.py

@@ -0,0 +1,295 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Convert a `stable_audio_tools` raw checkpoint into a Diffusers-style
+per-component layout consumable by FastVideo's standard
+`ComposedPipelineBase` loader.
+
+Why: the published Stability AI Stable Audio Open repos ship a single
+monolithic `model.safetensors` containing DiT (`model.model.*`),
+VAE / pretransform (`pretransform.model.*`), and NumberConditioner
+(`conditioner.conditioners.*`) state dicts under shared prefixes.
+FastVideo's standard loader expects the Diffusers per-subfolder layout
+(`model_index.json` + per-component subfolders). This script splits the
+monolithic checkpoint, drops the host-pipeline prefixes, normalizes
+LayerNorm `gamma`/`beta` keys to `weight`/`bias`, and writes a
+Diffusers-style tree.
+
+Usage::
+
+    python scripts/checkpoint_conversion/stable_audio_to_diffusers.py \
+        --src stabilityai/stable-audio-open-1.0 \
+        --dst converted_weights/stable_audio_open_1_0_diffusers
+
+    # Or with a local snapshot directory:
+    python scripts/checkpoint_conversion/stable_audio_to_diffusers.py \
+        --src /path/to/stable-audio-open-1.0 \
+        --dst converted_weights/stable_audio_open_small_diffusers
+
+The output tree is::
+
+    <dst>/
+    ├── model_index.json
+    ├── transformer/{config.json, diffusion_pytorch_model.safetensors}
+    ├── vae/{config.json, diffusion_pytorch_model.safetensors}
+    ├── conditioner/{config.json, diffusion_pytorch_model.safetensors}
+    ├── text_encoder/  (copied from src if present)
+    ├── tokenizer/     (copied from src if present)
+    └── scheduler/     (copied from src if present)
+
+Push to the FastVideo HF org with `huggingface-cli upload`. Confirm
+visibility (`--private` for staging, omit for public release)::
+
+    huggingface-cli upload --private FastVideo/stable-audio-open-1.0-Diffusers \
+        converted_weights/stable-audio-open-1.0-Diffusers
+    huggingface-cli upload --private FastVideo/stable-audio-open-small-Diffusers \
+        converted_weights/stable-audio-open-small-Diffusers
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import shutil
+from pathlib import Path
+from typing import Any
+
+from huggingface_hub import snapshot_download
+from safetensors import safe_open
+from safetensors.torch import save_file
+
+# Map host-pipeline prefix → output subfolder. Order matters only for
+# logging.
+_COMPONENT_PREFIXES: dict[str, str] = {
+    "model.model.": "transformer",
+    "pretransform.model.": "vae",
+    "conditioner.": "conditioner",
+}
+
+# Subfolders we copy verbatim from the source repo. The `vae/` subfolder
+# in `stabilityai/stable-audio-open-1.0` already ships in Diffusers
+# format with the key naming our `OobleckVAE` expects (`decoder.block.X.
+# conv_t1.weight_g` etc.) — preferred over the raw `pretransform.model.*`
+# extraction below, which uses a different module nesting (`decoder.
+# layers.X.layers.Y.*`) that would need a separate remap.
+_PASSTHROUGH_SUBFOLDERS = ("text_encoder", "tokenizer", "scheduler", "vae")
+
+_MODEL_INDEX = {
+    "_class_name": "StableAudioPipeline",
+    # `ComposedPipelineBase._load_config` requires this field; the value
+    # is informational (we don't pin a Diffusers version).
+    "_diffusers_version": "0.30.0",
+    "_fastvideo_converted_from": None,  # filled in at write time
+    "transformer": ["fastvideo.models.dits.stable_audio", "StableAudioDiT"],
+    "vae": ["fastvideo.models.vaes.oobleck", "OobleckVAE"],
+    "conditioner": [
+        "fastvideo.models.encoders.stable_audio_conditioner",
+        "StableAudioMultiConditioner",
+    ],
+    "text_encoder": ["transformers", "T5EncoderModel"],
+    "tokenizer": ["transformers", "T5TokenizerFast"],
+    "scheduler": ["diffusers", "CosineDPMSolverMultistepScheduler"],
+}
+
+
+def _resolve_src(src: str) -> Path:
+    """Accept either a HF repo id or a local directory."""
+    if os.path.isdir(src):
+        return Path(src)
+    from fastvideo.utils import resolve_hf_token
+    return Path(snapshot_download(repo_id=src, token=resolve_hf_token()))
+
+
+def _rename_layernorm_keys(state: dict[str, Any]) -> dict[str, Any]:
+    """`nn.LayerNorm` ships `gamma`/`beta` in the official checkpoint;
+    torch's `nn.LayerNorm` uses `weight`/`bias`."""
+    out: dict[str, Any] = {}
+    for k, v in state.items():
+        if k.endswith(".gamma"):
+            out[k[: -len(".gamma")] + ".weight"] = v
+        elif k.endswith(".beta"):
+            out[k[: -len(".beta")] + ".bias"] = v
+        else:
+            out[k] = v
+    return out
+
+
+def _split_state_dict(src_safetensors: Path) -> dict[str, dict[str, Any]]:
+    """Read `model.safetensors` and bucket its keys by the host-pipeline
+    prefix listed in `_COMPONENT_PREFIXES`."""
+    buckets: dict[str, dict[str, Any]] = {name: {} for name in _COMPONENT_PREFIXES.values()}
+    skipped = 0
+    with safe_open(str(src_safetensors), framework="pt") as f:
+        for k in f.keys():
+            for prefix, subfolder in _COMPONENT_PREFIXES.items():
+                if k.startswith(prefix):
+                    buckets[subfolder][k[len(prefix):]] = f.get_tensor(k)
+                    break
+            else:
+                skipped += 1
+    if skipped:
+        print(f"  skipped {skipped} keys (no matching prefix)")
+    for subfolder, state in buckets.items():
+        if subfolder == "transformer":
+            buckets[subfolder] = _rename_layernorm_keys(state)
+        print(f"  {subfolder}: {len(buckets[subfolder])} keys")
+    return buckets
+
+
+def _detect_projection_flags(
+        diff_cfg: dict[str, Any],
+        dit_state: dict[str, Any]) -> tuple[bool, bool]:
+    """Decide `project_cond_tokens` / `project_global_cond` by comparing
+    the actual `to_cond_embed.0` / `to_global_embed.0` weight shapes
+    in the DiT state dict against `cond_token_dim` / `global_cond_dim`.
+    Upstream's factory toggles these per-variant in ways that aren't
+    derivable from `embed_dim` alone (SA-1.0 keeps `cond_embed_dim` =
+    `cond_token_dim` and uses GQA in cross-attn; SA-small projects to
+    `embed_dim` and uses MHA).
+    """
+    cond_dim = diff_cfg.get("cond_token_dim")
+    glob_dim = diff_cfg.get("global_cond_dim")
+    cond_w = dit_state.get("to_cond_embed.0.weight")
+    glob_w = dit_state.get("to_global_embed.0.weight")
+    project_cond = (cond_w is not None and cond_w.shape[0] != cond_dim)
+    project_glob = (glob_w is not None and glob_w.shape[0] != glob_dim)
+    return project_cond, project_glob
+
+
+def _component_config(model_config: dict[str, Any]) -> dict[str, dict[str, Any]]:
+    """Pull per-component sizing out of the official `model_config.json`
+    so the converted repo records the architecture authoritatively."""
+    diff = model_config["model"]["diffusion"]
+    pre = model_config["model"]["pretransform"]["config"]
+    cond = model_config["model"]["conditioning"]
+    # Translate upstream `model_config.json` field names to FastVideo's
+    # `StableAudioArchConfig` fields. The standard `TransformerLoader`
+    # calls `dit_config.update_model_arch(json_dict)` which raises on
+    # unknown keys.
+    diff_cfg = dict(diff["config"])
+    if "num_heads" in diff_cfg:
+        diff_cfg["num_attention_heads"] = diff_cfg.pop("num_heads")
+    diff_cfg.pop("transformer_type", None)  # always "continuous_transformer"
+    # `attn_kwargs.qk_norm` (small variant) → top-level `qk_norm`.
+    attn_kwargs = diff_cfg.pop("attn_kwargs", None) or {}
+    if "qk_norm" in attn_kwargs:
+        diff_cfg["qk_norm"] = attn_kwargs["qk_norm"]
+    # `cross_attention_cond_ids` / `global_cond_ids` belong to the
+    # conditioner, not the DiT — keep them out of transformer/config.json.
+    transformer_cfg = {
+        "_class_name": "StableAudioDiT",
+        **diff_cfg,
+    }
+    # Mirror Diffusers' `AutoencoderOobleck` config field naming so
+    # `OobleckVAEArchConfig.update_model_arch` (called by `VAELoader`)
+    # accepts the keys.
+    enc_cfg = pre["encoder"]["config"]
+    dec_cfg = pre["decoder"]["config"]
+    vae_cfg = {
+        "_class_name": "AutoencoderOobleck",
+        "encoder_hidden_size": enc_cfg["channels"],
+        "downsampling_ratios": enc_cfg["strides"],
+        "channel_multiples": enc_cfg["c_mults"],
+        "decoder_channels": dec_cfg["channels"],
+        "decoder_input_channels": dec_cfg["latent_dim"],
+        "audio_channels": pre.get("io_channels", 2),
+        "sampling_rate": model_config.get("sample_rate", 44100),
+    }
+    conditioner_cfg = {
+        "_class_name": "StableAudioMultiConditioner",
+        "cond_dim": cond["cond_dim"],
+        "configs": cond["configs"],
+        "cross_attention_cond_ids": diff["cross_attention_cond_ids"],
+        "global_cond_ids": diff["global_cond_ids"],
+    }
+    return {
+        "transformer": transformer_cfg,
+        "vae": vae_cfg,
+        "conditioner": conditioner_cfg,
+    }
+
+
+def _write_component(dst_dir: Path, name: str, state: dict[str, Any], config: dict[str, Any]) -> None:
+    sub = dst_dir / name
+    sub.mkdir(parents=True, exist_ok=True)
+    save_file(state, str(sub / "diffusion_pytorch_model.safetensors"))
+    with open(sub / "config.json", "w") as f:
+        json.dump(config, f, indent=2)
+    print(f"  wrote {sub}/")
+
+
+def _copy_passthrough(src_dir: Path, dst_dir: Path) -> list[str]:
+    copied: list[str] = []
+    for sub in _PASSTHROUGH_SUBFOLDERS:
+        s = src_dir / sub
+        if not s.is_dir():
+            continue
+        d = dst_dir / sub
+        if d.exists():
+            shutil.rmtree(d)
+        shutil.copytree(s, d, ignore=shutil.ignore_patterns("*.bin", "*.ckpt"))
+        copied.append(sub)
+        print(f"  copied {sub}/")
+    return copied
+
+
+def convert(src: str, dst: str) -> None:
+    src_dir = _resolve_src(src)
+    dst_dir = Path(dst)
+    dst_dir.mkdir(parents=True, exist_ok=True)
+    print(f"src: {src_dir}")
+    print(f"dst: {dst_dir}")
+
+    monolithic = src_dir / "model.safetensors"
+    if not monolithic.is_file():
+        raise FileNotFoundError(f"Expected monolithic safetensors at {monolithic}")
+    cfg_path = src_dir / "model_config.json"
+    if not cfg_path.is_file():
+        raise FileNotFoundError(
+            f"Expected model_config.json at {cfg_path} (Stable Audio's authoritative "
+            f"per-component arch config).")
+    with open(cfg_path) as f:
+        model_config = json.load(f)
+
+    print("\n[1/4] Splitting monolithic state dict by component prefix:")
+    buckets = _split_state_dict(monolithic)
+
+    print("\n[2/4] Copying passthrough subfolders (text_encoder/tokenizer/scheduler/vae):")
+    copied = _copy_passthrough(src_dir, dst_dir)
+
+    print("\n[3/4] Writing per-component configs + safetensors:")
+    component_cfgs = _component_config(model_config)
+    project_cond, project_glob = _detect_projection_flags(component_cfgs["transformer"],
+                                                          buckets["transformer"])
+    component_cfgs["transformer"]["project_cond_tokens"] = project_cond
+    component_cfgs["transformer"]["project_global_cond"] = project_glob
+    for name in ("transformer", "vae", "conditioner"):
+        if name in copied:
+            print(f"  skipped {name}/ (already copied from source in Diffusers format)")
+            continue
+        _write_component(dst_dir, name, buckets[name], component_cfgs[name])
+
+    print("\n[4/4] Writing model_index.json:")
+    index = dict(_MODEL_INDEX)
+    # Prefer the HF repo id over a snapshot directory (which is
+    # machine-local and full of hash gunk).
+    index["_fastvideo_converted_from"] = src if "/" in src and not os.path.isdir(src) else src_dir.name
+    # Drop entries for subfolders we didn't actually populate.
+    available = {"transformer", "vae", "conditioner", *copied}
+    index = {k: v for k, v in index.items() if not k or k.startswith("_") or k in available}
+    with open(dst_dir / "model_index.json", "w") as f:
+        json.dump(index, f, indent=2)
+    print(f"  wrote {dst_dir}/model_index.json")
+    print(f"\nDone. Push to HF with:")
+    print(f"  huggingface-cli upload <FastVideo/repo-name> {dst_dir}")
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
+    parser.add_argument("--src", required=True,
+                        help="HF repo id or local directory containing model.safetensors + model_config.json")
+    parser.add_argument("--dst", required=True, help="Output directory for the Diffusers-format repo")
+    args = parser.parse_args()
+    convert(args.src, args.dst)
+
+
+if __name__ == "__main__":
+    main()

tests/README.md

@@ -1,7 +1,10 @@
 # Tests
 
-- `tests/local_tests/` are local-only tests that require a checked-out `LTX-2/`
-  directory under the repo root (`FastVideo/LTX-2`). Without that repo, they
-  will skip or fail.
+- `tests/local_tests/` holds local-only parity, smoke, and component tests
+  organized **by model family** (e.g. `tests/local_tests/ltx2/`,
+  `tests/local_tests/stable_audio/`). They typically require a checked-out
+  reference repo and/or model weights and are skipped in CI; see each family's
+  `README.md` for setup. The top-level
+  [`tests/local_tests/README.md`](./local_tests/README.md) is the family index.
 - The CI-backed test suite still lives in `fastvideo/tests/`.
 - Eventually, all tests will move under `tests/`.

tests/local_tests/README.md

@@ -1,7 +1,49 @@
-# Local LTX-2 Tests
+# Local Tests
 
-These tests depend on a checked-out `LTX-2/` directory under the repo root
-(`FastVideo/LTX-2`). Without that local repo, the tests will skip or fail.
+Local-only parity, smoke, and component tests for FastVideo model ports. They
+compare FastVideo against the official reference implementations and are
+**skipped in CI**; run them locally on a single GPU (or CPU where noted).
 
-For the CI-backed test suite, see `fastvideo/tests/`. Those are the tests
-currently exercised in CI. (Eventually all tests will move to `tests/`.)
+For the CI-backed test suite, see [`fastvideo/tests/`](../../fastvideo/tests/).
+
+## Layout
+
+Tests are organized by **model family**, one directory per port. Each family
+directory follows the layout produced by the
+[`add-model-prep`](https://github.com/anthropic-skills/add-model-prep) skill:
+
+```
+tests/local_tests/<family>/
+├── README.md         # reference assets, setup, test inventory
+├── __init__.py
+└── test_<family>_<component>_*.py
+```
+
+## Families
+
+| Family | Workload | Dir |
+|---|---|---|
+| Hunyuan GameCraft | T2V / I2V | [`gamecraft/`](./gamecraft/) |
+| GEN3C | T2V | [`gen3c/`](./gen3c/) |
+| Kandinsky-5 | T2V | [`kandinsky5/`](./kandinsky5/) |
+| LTX-2 | T2V (+ audio) | [`ltx2/`](./ltx2/) |
+| Stable Diffusion 3.5 | T2I | [`sd35/`](./sd35/) |
+| Stable Audio Open 1.0 | T2A | [`stable_audio/`](./stable_audio/) |
+| Wan2.2 | I2V | [`wan22/`](./wan22/) |
+
+## Running a family
+
+```bash
+pytest tests/local_tests/<family>/ -v -s
+```
+
+See each family's `README.md` for setup (clones, weights, env vars, gated
+tokens) and the per-component parity coverage.
+
+## Adding a new family
+
+Run the [`add-model-prep`](https://github.com/anthropic-skills/add-model-prep)
+skill from the FastVideo repo root. It scaffolds the family directory and
+`README.md` and stages reference clones + weights. (In-progress ports may also
+add a `PORT_STATUS.md` from the same skill template; remove it once the port
+lands.)

tests/local_tests/gamecraft/README.md

@@ -0,0 +1,67 @@
+# Hunyuan GameCraft Local Tests
+
+Local-only parity tests for the `gamecraft` FastVideo port. They compare
+FastVideo against the official `Hunyuan-GameCraft-1.0` implementation and are
+skipped in CI; run locally on a single GPU (~40 GB+ memory).
+
+## Reference Assets
+
+| Field | Value |
+|---|---|
+| Model family | `gamecraft` |
+| Workload types | `T2V`, `I2V` |
+| Official reference | `Hunyuan-GameCraft-1.0` (local clone, see test files) |
+| Local reference dir | `Hunyuan-GameCraft-1.0/` (env: `GAMECRAFT_OFFICIAL_PATH`) |
+| Official commit/version | `<TODO>` |
+| HF weights | `<TODO>` |
+| HF revision | `<TODO>` |
+| Local weights dir | `Hunyuan-GameCraft-1.0/weights/` (env: `GAMECRAFT_WEIGHTS_PATH`) |
+| Source layout | `<TODO>` |
+| Needs conversion | `<TODO>` |
+
+> Use only the env-var **name** for tokens (e.g., `HF_TOKEN`). Never paste a token value.
+
+## Shared Environment Setup
+
+Run from the FastVideo repo root in the same env used for FastVideo. See the
+`add-model-prep` skill for canonical clone + install commands.
+
+```bash
+# Clone the official reference under the FastVideo repo root.
+git clone <official_repo_url> Hunyuan-GameCraft-1.0
+uv pip install --no-deps -e ./Hunyuan-GameCraft-1.0
+```
+
+Do not change core dependency versions (`torch`, `diffusers`, `transformers`,
+`flash-attn`, `triton`, CUDA packages) without explicit approval.
+
+## Official Environment Status
+
+```text
+dependency_changes: <TODO>
+official_env_status: <TODO: imports_ok | private_deps_need_stubs | blocked>
+private_dep_stubs: <TODO>
+blocked_on: <TODO>
+```
+
+## Tests in this directory
+
+Run the whole family:
+
+```bash
+DISABLE_SP=1 pytest tests/local_tests/gamecraft/ -v -s
+```
+
+| Component | Test | Concerns | Status |
+|---|---|---|---|
+| `text encoders` (LLaVA-LLaMA-3-8B + CLIP ViT-L/14) | [`test_gamecraft_encoders_parity.py`](./test_gamecraft_encoders_parity.py) | `<TODO>` | `<TODO>` |
+| `transformer/DiT` | [`test_gamecraft_parity.py`](./test_gamecraft_parity.py) | `<TODO>` | `<TODO>` |
+| `vae` (`AutoencoderKLCausal3D`) | [`test_gamecraft_vae_parity.py`](./test_gamecraft_vae_parity.py) | `<TODO>` | `<TODO>` |
+| `pipeline` | [`test_gamecraft_pipeline_parity.py`](./test_gamecraft_pipeline_parity.py) | text encoding -> denoising -> VAE decode | `<TODO>` |
+
+## Review Notes
+
+- Required before handoff: non-skip PASS for each component parity test,
+  including reused components that own weights or numerical behavior.
+- Pipeline parity may start as a scaffold; final handoff requires non-skip
+  PASS or an explicit blocker accepted via the escape-hatch process.

tests/local_tests/gamecraft/__init__.py

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

tests/local_tests/gamecraft/test_gamecraft_encoders_parity.py

@@ -7,7 +7,7 @@ GameCraft uses:
 2. CLIP ViT-L/14 for secondary text encoding (768-dim pooled)
 
 Usage:
-    DISABLE_SP=1 pytest tests/local_tests/encoders/test_gamecraft_encoders_parity.py -v
+    DISABLE_SP=1 pytest tests/local_tests/gamecraft/test_gamecraft_encoders_parity.py -v
 """
 import os
 import sys

tests/local_tests/gamecraft/test_gamecraft_parity.py

@@ -7,15 +7,15 @@ Hunyuan-GameCraft-1.0 implementation to ensure numerical alignment.
 
 Usage:
     # Run with default paths
-    pytest tests/local_tests/transformers/test_gamecraft_parity.py -v
+    pytest tests/local_tests/gamecraft/test_gamecraft_parity.py -v
     
     # Run with custom paths
     GAMECRAFT_OFFICIAL_PATH=path/to/Hunyuan-GameCraft-1.0 \
     GAMECRAFT_WEIGHTS_PATH=path/to/weights/gamecraft_models/mp_rank_00_model_states.pt \
-    pytest tests/local_tests/transformers/test_gamecraft_parity.py -v
+    pytest tests/local_tests/gamecraft/test_gamecraft_parity.py -v
     
     # Enable debug logging
-    GAMECRAFT_DEBUG_LOGS=1 pytest tests/local_tests/transformers/test_gamecraft_parity.py -v
+    GAMECRAFT_DEBUG_LOGS=1 pytest tests/local_tests/gamecraft/test_gamecraft_parity.py -v
 """
 import os
 import sys

tests/local_tests/gamecraft/test_gamecraft_pipeline_parity.py

@@ -8,12 +8,12 @@ across the complete pipeline (text encoding -> denoising -> VAE decode).
 
 Usage:
     # Run with default paths
-    DISABLE_SP=1 pytest tests/local_tests/pipelines/test_gamecraft_pipeline_parity.py -v -s
+    DISABLE_SP=1 pytest tests/local_tests/gamecraft/test_gamecraft_pipeline_parity.py -v -s
     
     # Run with custom paths
     GAMECRAFT_OFFICIAL_PATH=path/to/Hunyuan-GameCraft-1.0 \
     GAMECRAFT_WEIGHTS_PATH=path/to/weights \
-    DISABLE_SP=1 pytest tests/local_tests/pipelines/test_gamecraft_pipeline_parity.py -v -s
+    DISABLE_SP=1 pytest tests/local_tests/gamecraft/test_gamecraft_pipeline_parity.py -v -s
 
 Notes:
     - This test requires the official Hunyuan-GameCraft-1.0 repo

tests/local_tests/gamecraft/test_gamecraft_vae_parity.py

@@ -12,7 +12,7 @@ Follows the same pattern as test_ltx2_vae.py:
 - Checks for non-finite values
 
 Usage:
-    DISABLE_SP=1 pytest tests/local_tests/vaes/test_gamecraft_vae_parity.py -v
+    DISABLE_SP=1 pytest tests/local_tests/gamecraft/test_gamecraft_vae_parity.py -v
 """
 import json
 import os

tests/local_tests/gen3c/README.md

@@ -0,0 +1,65 @@
+# GEN3C Local Tests
+
+Local-only parity and smoke tests for the `gen3c` FastVideo port. The pipeline
+smoke test runs without weights; the transformer parity test compares against
+the official `GEN3C` repo and converted weights, and is skipped in CI.
+
+## Reference Assets
+
+| Field | Value |
+|---|---|
+| Model family | `gen3c` |
+| Workload types | `T2V` |
+| Official reference | `GEN3C` (local clone, see test files) |
+| Local reference dir | `GEN3C/` |
+| Official commit/version | `<TODO>` |
+| HF weights | `<TODO>` |
+| HF revision | `<TODO>` |
+| Local weights dir | `<TODO>` (env: `GEN3C_FASTVIDEO_PATH`, `GEN3C_DIFFUSERS_PATH`) |
+| Source layout | `<TODO>` |
+| Needs conversion | `<TODO>` |
+
+> Use only the env-var **name** for tokens (e.g., `HF_TOKEN`). Never paste a token value.
+
+## Shared Environment Setup
+
+Run from the FastVideo repo root in the same env used for FastVideo. See the
+`add-model-prep` skill for canonical clone + install commands.
+
+```bash
+# Clone the official reference under the FastVideo repo root.
+git clone <official_repo_url> GEN3C
+uv pip install --no-deps -e ./GEN3C
+```
+
+Do not change core dependency versions (`torch`, `diffusers`, `transformers`,
+`flash-attn`, `triton`, CUDA packages) without explicit approval.
+
+## Official Environment Status
+
+```text
+dependency_changes: <TODO>
+official_env_status: <TODO: imports_ok | private_deps_need_stubs | blocked>
+private_dep_stubs: <TODO>
+blocked_on: <TODO>
+```
+
+## Tests in this directory
+
+Run the whole family:
+
+```bash
+pytest tests/local_tests/gen3c/ -v
+```
+
+| Component | Test | Concerns | Status |
+|---|---|---|---|
+| `transformer` (Gen3CTransformer3DModel) | [`test_gen3c.py`](./test_gen3c.py) | `<TODO>` | `<TODO>` |
+| `pipeline` (smoke) | [`test_gen3c_pipeline_smoke.py`](./test_gen3c_pipeline_smoke.py) | random-weight smoke; full mode requires `GEN3C_DIFFUSERS_PATH` | `<TODO>` |
+
+## Review Notes
+
+- Required before handoff: non-skip PASS for each component parity test,
+  including reused components that own weights or numerical behavior.
+- Pipeline parity may start as a scaffold; final handoff requires non-skip
+  PASS or an explicit blocker accepted via the escape-hatch process.

tests/local_tests/gen3c/__init__.py

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

tests/local_tests/gen3c/test_gen3c.py

@@ -2,10 +2,10 @@
 """
 Usage:
     # Basic tests (no weights required)
-    pytest tests/local_tests/transformers/test_gen3c.py -v
+    pytest tests/local_tests/gen3c/test_gen3c.py -v
     
     # Full parity tests (requires converted weights)
-    GEN3C_FASTVIDEO_PATH=./gen3c_converted/transformer pytest tests/local_tests/transformers/test_gen3c.py -v
+    GEN3C_FASTVIDEO_PATH=./gen3c_converted/transformer pytest tests/local_tests/gen3c/test_gen3c.py -v
 """
 
 import os

tests/local_tests/gen3c/test_gen3c_pipeline_smoke.py

@@ -9,11 +9,11 @@ required.
 
 Usage:
     # Quick smoke test (no weights required, CPU or CUDA)
-    pytest tests/local_tests/pipelines/test_gen3c_pipeline_smoke.py -v
+    pytest tests/local_tests/gen3c/test_gen3c_pipeline_smoke.py -v
 
     # Full pipeline smoke test with converted weights
     GEN3C_DIFFUSERS_PATH=./gen3c_fastvideo \
-        pytest tests/local_tests/pipelines/test_gen3c_pipeline_smoke.py -v -k full
+        pytest tests/local_tests/gen3c/test_gen3c_pipeline_smoke.py -v -k full
 """
 
 import os

tests/local_tests/kandinsky5/README.md

@@ -0,0 +1,65 @@
+# Kandinsky-5 Local Tests
+
+Local-only parity tests for the `kandinsky5` FastVideo port. Compares the
+FastVideo Kandinsky-5 transformer against the Diffusers reference
+(`diffusers.Kandinsky5Transformer3DModel`); skipped in CI.
+
+## Reference Assets
+
+| Field | Value |
+|---|---|
+| Model family | `kandinsky5` |
+| Workload types | `T2V` |
+| Official reference | `diffusers.Kandinsky5Transformer3DModel` (Diffusers package) |
+| Local reference dir | `none` (Diffusers reference, no clone) |
+| Official commit/version | `<TODO: diffusers version>` |
+| HF weights | `kandinskylab/Kandinsky-5.0-T2V-Lite-sft-5s-Diffusers` |
+| HF revision | `<TODO>` |
+| Local weights dir | `official_weights/kandinskylab/Kandinsky-5.0-T2V-Lite-sft-5s-Diffusers` (env: `KANDINSKY5_DIFFUSERS_PATH`, `KANDINSKY5_TRANSFORMER_PATH`) |
+| Source layout | `diffusers` |
+| Needs conversion | `<TODO>` |
+
+> Use only the env-var **name** for tokens (e.g., `HF_TOKEN`). Never paste a token value.
+
+## Shared Environment Setup
+
+Run from the FastVideo repo root in the same env used for FastVideo. The
+reference is the published Diffusers class — no clone or upstream install is
+required beyond the FastVideo `diffusers` pin.
+
+Do not change core dependency versions (`torch`, `diffusers`, `transformers`,
+`flash-attn`, `triton`, CUDA packages) without explicit approval.
+
+## Official Environment Status
+
+```text
+dependency_changes: none
+official_env_status: <TODO: imports_ok | private_deps_need_stubs | blocked>
+private_dep_stubs: none
+blocked_on: <TODO>
+```
+
+## Weight Setup
+
+```bash
+python "$HOME/.config/opencode/skill/add-model-prep/scripts/download_hf_weights.py" \
+    "kandinskylab/Kandinsky-5.0-T2V-Lite-sft-5s-Diffusers" \
+    "official_weights/kandinskylab/Kandinsky-5.0-T2V-Lite-sft-5s-Diffusers"
+```
+
+## Tests in this directory
+
+```bash
+pytest tests/local_tests/kandinsky5/ -v
+```
+
+| Component | Test | Concerns | Status |
+|---|---|---|---|
+| `transformer` (Kandinsky5Transformer3DModel) | [`test_kandinsky5_lite_transformer_parity.py`](./test_kandinsky5_lite_transformer_parity.py) | `<TODO>` | `<TODO>` |
+
+## Review Notes
+
+- Required before handoff: non-skip PASS for each component parity test,
+  including reused components that own weights or numerical behavior.
+- Pipeline parity may start as a scaffold; final handoff requires non-skip
+  PASS or an explicit blocker accepted via the escape-hatch process.

tests/local_tests/kandinsky5/__init__.py

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

tests/local_tests/ltx2/README.md

@@ -0,0 +1,73 @@
+# LTX-2 Local Tests
+
+Local-only parity, registry, and pipeline-smoke tests for the `ltx2` FastVideo
+port. Most tests depend on a checked-out `LTX-2/` directory at the repo root
+(`FastVideo/LTX-2`) and converted Diffusers weights; without those they will
+skip or fail. Skipped in CI; run locally on a single GPU.
+
+## Reference Assets
+
+| Field | Value |
+|---|---|
+| Model family | `ltx2` |
+| Workload types | `T2V` (video + optional audio) |
+| Official reference | `Lightricks/LTX-2` (local clone) |
+| Local reference dir | `LTX-2/` (with `LTX-2/packages/ltx-core/src` on `sys.path`) |
+| Official commit/version | `<TODO>` |
+| HF weights | `Lightricks/LTX-2`, `FastVideo/LTX2-base`, `FastVideo/LTX2-Distilled-Diffusers` |
+| HF revision | `<TODO>` |
+| Local weights dir | `<TODO>` (env: `LTX2_DIFFUSERS_PATH`, `LTX2_FASTVIDEO_GEMMA_LOG`, etc.) |
+| Source layout | `<TODO: diffusers / raw_official / mixed>` |
+| Needs conversion | `<TODO>` |
+
+> Use only the env-var **name** for tokens (e.g., `HF_TOKEN`). Never paste a token value.
+
+## Shared Environment Setup
+
+Run from the FastVideo repo root in the same env used for FastVideo. See the
+`add-model-prep` skill for canonical clone + install commands.
+
+```bash
+# Clone the official reference under the FastVideo repo root.
+git clone <official_repo_url> LTX-2
+uv pip install --no-deps -e ./LTX-2
+```
+
+Do not change core dependency versions (`torch`, `diffusers`, `transformers`,
+`flash-attn`, `triton`, CUDA packages) without explicit approval.
+
+## Official Environment Status
+
+```text
+dependency_changes: <TODO>
+official_env_status: <TODO: imports_ok | private_deps_need_stubs | blocked>
+private_dep_stubs: <TODO>
+blocked_on: <TODO>
+```
+
+## Tests in this directory
+
+Run the whole family:
+
+```bash
+pytest tests/local_tests/ltx2/ -v -s
+```
+
+| Component | Test | Concerns | Status |
+|---|---|---|---|
+| `text encoder` (Gemma) | [`test_ltx2_gemma_encoder.py`](./test_ltx2_gemma_encoder.py) | encoder + connector | `<TODO>` |
+| `text encoder` (Gemma, full parity) | [`test_ltx2_gemma_parity.py`](./test_ltx2_gemma_parity.py) | full Gemma parity vs official | `<TODO>` |
+| `transformer/DiT` | [`test_ltx2.py`](./test_ltx2.py) | LTX2VideoConfig DiT parity | `<TODO>` |
+| `transformer/DiT` (audio modality) | [`test_ltx2_audio.py`](./test_ltx2_audio.py) | audio-conditioning DiT parity (imports helpers from `test_ltx2.py`) | `<TODO>` |
+| `vae` (video) | [`test_ltx2_vae.py`](./test_ltx2_vae.py) | LTX2VideoEncoder/Decoder parity | `<TODO>` |
+| `vae` (video, official path) | [`test_ltx2_vae_official.py`](./test_ltx2_vae_official.py) | LTX2VAEConfig + VAELoader path | `<TODO>` |
+| `vae` (audio) | [`test_ltx2_audio_vae.py`](./test_ltx2_audio_vae.py) | audio encoder/decoder/vocoder parity | `<TODO>` |
+| `pipeline` (smoke) | [`test_ltx2_pipeline_smoke.py`](./test_ltx2_pipeline_smoke.py) | VideoGenerator end-to-end smoke | `<TODO>` |
+| `registry` | [`test_ltx2_registry.py`](./test_ltx2_registry.py) | sampling/pipeline registry resolution | `<TODO>` |
+
+## Review Notes
+
+- Required before handoff: non-skip PASS for each component parity test,
+  including reused components that own weights or numerical behavior.
+- Pipeline parity may start as a scaffold; final handoff requires non-skip
+  PASS or an explicit blocker accepted via the escape-hatch process.

tests/local_tests/ltx2/__init__.py

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

tests/local_tests/sd35/README.md

@@ -0,0 +1,66 @@
+# Stable Diffusion 3.5 Local Tests
+
+Local-only component parity tests for the `sd35` FastVideo port. Compares
+FastVideo's SD3.5 components (text encoders, transformer, VAE, scheduler)
+against the Diffusers reference and the published `stabilityai/stable-diffusion-3.5-medium`
+checkpoint. Skipped in CI; CUDA required.
+
+## Reference Assets
+
+| Field | Value |
+|---|---|
+| Model family | `sd35` |
+| Workload types | `T2I` |
+| Official reference | `diffusers.SD3Transformer2DModel`, `diffusers.AutoencoderKL`, `transformers` text encoders |
+| Local reference dir | `none` (Diffusers + transformers reference, no clone) |
+| Official commit/version | `<TODO: diffusers + transformers versions>` |
+| HF weights | `stabilityai/stable-diffusion-3.5-medium` |
+| HF revision | `<TODO>` |
+| Local weights dir | `official_weights/stabilityai__stable-diffusion-3.5-medium` (env: `SD35_MODEL_DIR`) |
+| Source layout | `diffusers` |
+| Needs conversion | `<TODO>` |
+
+> Use only the env-var **name** for tokens (e.g., `HF_TOKEN`). Never paste a token value.
+
+## Shared Environment Setup
+
+Run from the FastVideo repo root in the same env used for FastVideo. The
+reference is the published Diffusers + transformers classes — no clone or
+upstream install is required beyond the FastVideo pins.
+
+Do not change core dependency versions (`torch`, `diffusers`, `transformers`,
+`flash-attn`, `triton`, CUDA packages) without explicit approval.
+
+## Official Environment Status
+
+```text
+dependency_changes: none
+official_env_status: <TODO: imports_ok | private_deps_need_stubs | blocked>
+private_dep_stubs: none
+blocked_on: <TODO>
+```
+
+## Weight Setup
+
+```bash
+python "$HOME/.config/opencode/skill/add-model-prep/scripts/download_hf_weights.py" \
+    "stabilityai/stable-diffusion-3.5-medium" \
+    "official_weights/stabilityai__stable-diffusion-3.5-medium"
+```
+
+## Tests in this directory
+
+```bash
+pytest tests/local_tests/sd35/ -v -s
+```
+
+| Component | Test | Concerns | Status |
+|---|---|---|---|
+| all components (text encoders, transformer, VAE, scheduler) | [`test_sd35_component_parity.py`](./test_sd35_component_parity.py) | single-file omnibus parity test | `<TODO>` |
+
+## Review Notes
+
+- Required before handoff: non-skip PASS for each component parity test,
+  including reused components that own weights or numerical behavior.
+- Pipeline parity may start as a scaffold; final handoff requires non-skip
+  PASS or an explicit blocker accepted via the escape-hatch process.

tests/local_tests/stable_audio/README.md

@@ -0,0 +1,137 @@
+# Local Stable Audio Tests
+
+End-to-end parity tests for the Stable Audio Open 1.0 pipeline + the
+Oobleck VAE component. They compare FastVideo against the published
+weights and the official `Stability-AI/stable-audio-tools` reference,
+so they're skipped in CI and run locally on a single GPU.
+
+## Setup
+
+### 1. Hugging Face access
+
+Stable Audio Open 1.0 is a gated repo. Accept the terms once at
+https://huggingface.co/stabilityai/stable-audio-open-1.0 and export
+your token in the shell:
+
+```bash
+export HF_TOKEN=hf_...
+# any of HF_TOKEN / HUGGINGFACE_HUB_TOKEN / HF_API_KEY works
+```
+
+The tests skip cleanly with a helpful message if no token is set.
+
+### 2. Optional inference dependencies
+
+`k_diffusion` and friends are not part of FastVideo's base install.
+One-shot:
+
+```bash
+pip install k_diffusion einops_exts alias_free_torch torchsde
+```
+
+### 3. Clone the upstream reference repo
+
+The pipeline-level parity tests (`test_stable_audio_pipeline_parity.py`,
+`test_stable_audio_a2a_parity.py`) and the VAE-component parity tests
+(`test_oobleck_vae_parity.py`, `test_oobleck_vae_official_parity.py`)
+import directly from `stable_audio_tools`. Clone it under the repo
+root and `pip install` editable, then add it to `.gitignore`:
+
+```bash
+cd <FastVideo repo root>
+git clone --depth 1 https://github.com/Stability-AI/stable-audio-tools.git
+pip install --no-deps -e ./stable-audio-tools
+echo "/stable-audio-tools/" >> .git/info/exclude   # personal ignore
+```
+
+The smoke test and the inpainting self-consistency test do not need the
+upstream clone.
+
+### 4. (Optional) Pre-warm the model cache
+
+The first parity-test run downloads the full Stable Audio Open 1.0
+checkpoint (~3 GB). To avoid the download blocking your first test
+run, fetch it ahead of time:
+
+```bash
+python -c "from huggingface_hub import snapshot_download; \
+           snapshot_download('stabilityai/stable-audio-open-1.0')"
+```
+
+## Running the tests
+
+All Stable Audio tests in one shot (~80 s on a single B200):
+
+```bash
+pytest tests/local_tests/stable_audio/test_stable_audio_pipeline_smoke.py \
+       tests/local_tests/stable_audio/test_stable_audio_pipeline_parity.py \
+       tests/local_tests/stable_audio/test_stable_audio_a2a_parity.py \
+       tests/local_tests/stable_audio/test_stable_audio_inpaint_parity.py \
+       tests/local_tests/stable_audio/test_oobleck_vae_parity.py \
+       tests/local_tests/stable_audio/test_oobleck_vae_official_parity.py \
+       -v -s
+```
+
+Or, to run the entire family:
+
+```bash
+pytest tests/local_tests/stable_audio/ -v -s
+```
+
+Add `-s` to print the per-test diff numbers (shape / abs_mean / max
+diff / drift / RMS ratios).
+
+### What each test covers
+
+| Test | Compares against | Notes |
+|---|---|---|
+| `test_stable_audio_pipeline_smoke.py` | — (no GPU) | Imports + registry + preset wiring. CPU-only. |
+| `test_stable_audio_pipeline_parity.py` | `stable_audio_tools.inference.generate_diffusion_cond` | T2A end-to-end. Drift bound: < 1% / 0.05 element-wise. |
+| `test_stable_audio_a2a_parity.py` | `generate_diffusion_cond(init_audio=...)` | Audio-to-audio variation. Drift bound: < 2% / 0.05 mean. |
+| `test_stable_audio_inpaint_parity.py` | self-consistency | RePaint blending — kept-region RMS ratio in [0.5, 2.0], unkept region non-silent + bounded. SA Open 1.0 isn't inpaint-trained so there's no upstream reference to compare against. |
+| `test_oobleck_vae_parity.py` | `stable_audio_tools.models.AudioAutoencoder` (real weights) | Decode + encode + round-trip; should be bit-identical (fp32, diff = 0). |
+| `test_oobleck_vae_official_parity.py` | `OobleckEncoder` / `OobleckDecoder` (random init) | Architectural parity — verifies the FastVideo port matches upstream's structure. Bit-identical (diff = 0). |
+
+### Reproducing a single test
+
+Each test file is independent. Run one:
+
+```bash
+pytest tests/local_tests/stable_audio/test_stable_audio_a2a_parity.py -v -s
+```
+
+Expected output (approximate, on a B200, fp32, default seed = 0):
+
+```
+off shape=(1, 2, 66150) abs_mean=0.181810
+fv  shape=(1, 2, 66150) abs_mean=0.180913
+diff max=0.030 mean=0.003 median=0.002  abs_mean rel drift=0.49%
+PASSED
+```
+
+## Troubleshooting
+
+- **All tests skip with "gated repo" reason.** Check `echo $HF_TOKEN`
+  and that you've accepted the model terms at
+  https://huggingface.co/stabilityai/stable-audio-open-1.0.
+- **`ModuleNotFoundError: stable_audio_tools` / `k_diffusion`.** Re-run
+  the setup steps; the parity tests skip cleanly if these aren't
+  importable but only after detecting both, so a partial install can
+  produce an unhelpful error mid-test.
+- **A2A or inpaint test passes individually but fails in a batch
+  run.** The deterministic-math toggles (TF32 / cuDNN benchmark off)
+  are set in `StableAudioPipeline.load_modules`. If a test earlier in
+  the batch flipped them on, drift can spike — re-run the failing
+  test alone to confirm.
+- **Wall-clock blew up past 5 min.** First call downloads the
+  checkpoint and the T5-base weights (~4 GB total) — see the
+  pre-warm step above.
+
+## Adding new parity tests for this family
+
+Use any of the existing `test_stable_audio_*` files as a template — the
+HF-token detection (`_hf_token` / `_can_access` / `_setup_hf_env`) and
+the loader for the official model (`_load_official`) are duplicated
+across the parity-test files; copy them as a starting point and call
+out the comparison reference (upstream class + invocation) in the file
+docstring.

tests/local_tests/stable_audio/__init__.py

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