bundle variant -> preset resolution, pipeline pin, and LTX-2.5 examples

A bundle's training variant (header-declared `variant`, filename token as
fallback) now selects its sampling preset: distilled -> ltx2_distilled
(8 steps, cfg 1.0, distilled sigmas engage at <=8 steps), sft/base ->
ltx2_base. The bundle table also pins the pipeline class, so
VideoGenerator.from_pretrained(<bundle>) resolves without an override.
Adds basic_ltx2_5{,_distilled}.py examples (--model-path takes the FILE,
--gemma-root declares the paired encoder root), a README section, and CPU
tests for variant/preset resolution, table integrity, and the examples.

.gitignore

@@ -23,6 +23,7 @@ Miniconda3-latest-Linux-x86_64.sh
 *validation/
 data/
 outputs/
+outputs_audio/
 outputs_video
 checkpoints/
 sbatch.sh

docs/inference/support_matrix.md

@@ -64,6 +64,7 @@ column links a runnable script in `examples/inference/basic/` where one exists.
 | ltx2 | `FastVideo/LTX2-Distilled-Diffusers`<br>`FastVideo/LTX2.3-Distilled-Diffusers`<br>`FastVideo/LTX-2.3-Distilled-Diffusers` | T2V | [basic_ltx2_distilled.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_ltx2_distilled.py) |
 | ltx2 | `Lightricks/LTX-2.3`<br>`FastVideo/LTX2.3-base`<br>`FastVideo/LTX2.3-Diffusers` | T2V | [basic_ltx2.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_ltx2.py) |
 | ltx2 | `Lightricks/LTX-2`<br>`FastVideo/LTX2-base`<br>`FastVideo/LTX2-Diffusers` | T2V | [basic_ltx2.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_ltx2.py) |
+| mmaudio | `FastVideo/MMAudio-large-44k-v2-Diffusers` | V2A, T2A | [basic_mmaudio.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_mmaudio.py) |
 | matrixgame | `FastVideo/Matrix-Game-2.0-Base-Distilled-Diffusers`<br>`FastVideo/Matrix-Game-2.0-GTA-Distilled-Diffusers`<br>`FastVideo/Matrix-Game-2.0-TempleRun-Distilled-Diffusers`<br>`FastVideo/Matrix-Game-2.0-Base-Diffusers`<br>`FastVideo/Matrix-Game-2.0-GTA-Diffusers`<br>`FastVideo/Matrix-Game-2.0-TempleRun-Diffusers`<br>`mignonjia/mg_longtuning_distilled_zelda`<br>`mignonjia/mg_sf_distilled_zelda_1k_steps`<br>`mignonjia/mg_sf_distilled_zelda`<br>`mignonjia/mg_causal_zelda`<br>`mignonjia/mg_bidirectional_zelda` | I2V | [basic_matrixgame2.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_matrixgame2.py) |
 | matrixgame | `FastVideo/Matrix-Game-3.0-Base-Distilled-Diffusers` | I2V | [basic_matrixgame3.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_matrixgame3.py) |
 | minimax_h3 | `MiniMaxAI/MiniMax-H3` | T2V, I2V | [T2VA](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_minimax_h3_t2v.py)<br>[FL2VA](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_minimax_h3_fl2va.py)<br>[Ref2VA](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_minimax_h3_ref2va.py) |
@@ -94,6 +95,10 @@ column links a runnable script in `examples/inference/basic/` where one exists.
 (`StableAudioT2AConfig` / `StableAudioOpenSmallConfig`); they are registered
 under the generic T2V workload option in the registry.
 
+**Note (MMAudio)**: the registered Hugging Face model ID is reserved but not
+yet public. Follow the [MMAudio inference guide](https://github.com/hao-ai-lab/FastVideo/blob/main/fastvideo/pipelines/basic/mmaudio/README.md)
+to convert the official weights locally and set `MMAUDIO_MODEL_PATH`.
+
 **Note (MiniMax H3)**: T2VA, FL2VA, and Ref2VA all generate video with stereo
 audio. Use the Ref2VA example when passing ordered image, video, or audio
 references.

examples/inference/basic/README.md

@@ -33,6 +33,35 @@ For the typed config/request path added during the inference API refactor:
 python examples/inference/basic/basic_dmd_new_api.py
 ```
 
+## LTX-2.5 single-file bundles
+
+These examples load LTX-2.5 from a single-file bundle: one `.safetensors`
+file carrying every component (transformer, video VAE, audio VAE, vocoder,
+text projection). (The Hugging Face release ships as a split pack — one
+file per component — which loads through the standard repo path instead.)
+Pass the bundle FILE path directly:
+
+```bash
+python examples/inference/basic/basic_ltx2_5_distilled.py \
+    --model-path /path/to/bundle.safetensors \
+    --gemma-root /path/to/gemma_root
+```
+
+| bundle variant | preset | sampling defaults |
+|---|---|---|
+| distilled | `ltx2_distilled` | one stage, 8 steps, CFG 1.0 (no CFG/STG, no spatial upscaler) |
+| sft / base | `ltx2_base` | 40 steps, CFG 3.0 with STG |
+
+- The variant is read from the bundle header when declared; otherwise the
+  `distilled` token in the file name decides.
+- The Gemma text encoder is NOT in the bundle. Declare its root with
+  `--gemma-root` (or `FASTVIDEO_LTX_ENCODER_ROOT`), and use the root shipped
+  with the same transformer variant — the roots differ in prompt templating.
+- Decoder caveat: bundles that declare a diffusion VAE decoder
+  (`CausalDiffusionVAE`) are not supported yet; loading them currently
+  fails at VAE build time. Bundles with the classic `CausalVideoAutoencoder`
+  run end to end.
+
 ## Basic Walkthrough
 
 All you need to generate videos using multi-gpus from state-of-the-art diffusion pipelines is the following few lines!

examples/inference/basic/basic_ltx2_5.py

@@ -0,0 +1,103 @@
+# SPDX-License-Identifier: Apache-2.0
+"""LTX-2.5 (sft/base) text-to-video from a single-file bundle.
+
+Loads LTX-2.5 from a single-file bundle: ONE ``.safetensors`` file carrying
+every component (transformer, video VAE, audio VAE, vocoder, text
+projection), so ``--model-path`` takes the bundle FILE, not a repo
+directory. (The split-pack Hugging Face release loads through the standard
+repo path instead.)
+
+The Gemma text encoder is NOT in the bundle: pass its root directory with
+``--gemma-root`` (or set ``FASTVIDEO_LTX_ENCODER_ROOT``). Use the encoder
+root shipped WITH this transformer variant -- the published roots share
+weights but differ in prompt templating, so a mismatched root silently
+changes prompting.
+
+A non-distilled bundle resolves to the standard ``ltx2_base`` preset
+(40 steps, CFG 3.0 with STG); run without sampling flags to use it as-is.
+For the 8-step distilled recipe, see ``basic_ltx2_5_distilled.py``.
+
+Audio: the bundle carries an audio VAE + vocoder, so generated videos get
+an audio track. A bundle that declares no audio decoder skips audio
+decoding and still produces video.
+
+Decoder caveat: bundles that declare a *diffusion* VAE decoder
+(``CausalDiffusionVAE``) are not implemented yet -- loading such a bundle
+currently fails at VAE build time with an unsupported-architecture error
+(no classic-decoder fallback is wired). Bundles with the classic
+``CausalVideoAutoencoder`` decode end to end.
+"""
+import argparse
+import os
+
+from fastvideo import VideoGenerator
+
+PROMPT = ("A warm sunny backyard. The camera starts in a tight cinematic close-up "
+          "of a woman and a man in their 30s, facing each other with serious "
+          "expressions. The woman, emotional and dramatic, says softly, \"That's "
+          "it... Dad's lost it. And we've lost Dad.\" The man exhales, slightly "
+          "annoyed: \"Stop being so dramatic, Jess.\" A beat. He glances aside, "
+          "then mutters defensively, \"He's just having fun.\" The camera slowly "
+          "pans right, revealing the grandfather in the garden wearing enormous "
+          "butterfly wings, waving his arms in the air like he's trying to take "
+          "off. He shouts, \"Wheeeew!\" as he flaps his wings with full commitment. "
+          "The woman covers her face, on the verge of tears. The tone is deadpan, "
+          "absurd, and quietly tragic.")
+
+# Sampling flags default to None: unset flags are NOT passed to
+# `generate_video`, so the bundle's preset supplies them via
+# `SamplingParam.from_pretrained` (ltx2_base: 40 steps, cfg 3.0, 512x768,
+# 121 frames).
+_SAMPLING_FLAGS = ("height", "width", "num_frames", "num_inference_steps", "guidance_scale", "seed")
+
+
+def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="LTX-2.5 (sft/base) inference from a single-file bundle.")
+    parser.add_argument(
+        "--model-path",
+        required=True,
+        help="Path to the LTX-2.5 bundle (.safetensors FILE).",
+    )
+    parser.add_argument(
+        "--gemma-root",
+        default=None,
+        help="Gemma text-encoder root directory. Must be the root paired with "
+        "this transformer variant; the roots differ in prompt templating.",
+    )
+    parser.add_argument("--prompt", default=PROMPT)
+    parser.add_argument("--output-path", default="outputs_video/ltx2_5/output_ltx2_5_t2v.mp4")
+    parser.add_argument("--num-gpus", type=int, default=1)
+    parser.add_argument("--height", type=int, default=None)
+    parser.add_argument("--width", type=int, default=None)
+    parser.add_argument("--num-frames", type=int, default=None)
+    parser.add_argument("--num-inference-steps", type=int, default=None)
+    parser.add_argument("--guidance-scale", type=float, default=None)
+    parser.add_argument("--seed", type=int, default=None)
+    return parser.parse_args(argv)
+
+
+def sampling_overrides(args: argparse.Namespace) -> dict:
+    """Only sampling flags the user actually passed; the preset supplies the rest."""
+    return {name: getattr(args, name) for name in _SAMPLING_FLAGS if getattr(args, name) is not None}
+
+
+def main(argv: list[str] | None = None) -> None:
+    args = parse_args(argv)
+    if args.gemma_root:
+        os.environ["FASTVIDEO_LTX_ENCODER_ROOT"] = args.gemma_root
+
+    generator = VideoGenerator.from_pretrained(
+        args.model_path,
+        num_gpus=args.num_gpus,
+    )
+    generator.generate_video(
+        prompt=args.prompt,
+        output_path=args.output_path,
+        save_video=True,
+        **sampling_overrides(args),
+    )
+    generator.shutdown()
+
+
+if __name__ == "__main__":
+    main()

examples/inference/basic/basic_ltx2_5_distilled.py

@@ -0,0 +1,104 @@
+# SPDX-License-Identifier: Apache-2.0
+"""LTX-2.5 distilled text-to-video from a single-file bundle.
+
+Loads LTX-2.5 from a single-file bundle: ONE ``.safetensors`` file carrying
+every component (transformer, video VAE, audio VAE, vocoder, text
+projection), so ``--model-path`` takes the bundle FILE, not a repo
+directory. (The split-pack Hugging Face release loads through the standard
+repo path instead.)
+
+The Gemma text encoder is NOT in the bundle: pass its root directory with
+``--gemma-root`` (or set ``FASTVIDEO_LTX_ENCODER_ROOT``). Use the encoder
+root shipped WITH this transformer variant -- the published roots share
+weights but differ in prompt templating, so a mismatched root silently
+changes prompting.
+
+Distilled recipe: ONE stage -- 8 steps on the distilled sigma schedule,
+CFG 1.0, no CFG/STG, and no spatial upscaler needed. All of it comes from
+the ``ltx2_distilled`` preset the bundle resolves to; run without sampling
+flags to use it as-is.
+
+Audio: the bundle carries an audio VAE + vocoder, so generated videos get
+an audio track. A bundle that declares no audio decoder skips audio
+decoding and still produces video.
+
+Decoder caveat: bundles that declare a *diffusion* VAE decoder
+(``CausalDiffusionVAE``) are not implemented yet -- loading such a bundle
+currently fails at VAE build time with an unsupported-architecture error
+(no classic-decoder fallback is wired). Bundles with the classic
+``CausalVideoAutoencoder`` decode end to end.
+"""
+import argparse
+import os
+
+from fastvideo import VideoGenerator
+
+PROMPT = ("A warm sunny backyard. The camera starts in a tight cinematic close-up "
+          "of a woman and a man in their 30s, facing each other with serious "
+          "expressions. The woman, emotional and dramatic, says softly, \"That's "
+          "it... Dad's lost it. And we've lost Dad.\" The man exhales, slightly "
+          "annoyed: \"Stop being so dramatic, Jess.\" A beat. He glances aside, "
+          "then mutters defensively, \"He's just having fun.\" The camera slowly "
+          "pans right, revealing the grandfather in the garden wearing enormous "
+          "butterfly wings, waving his arms in the air like he's trying to take "
+          "off. He shouts, \"Wheeeew!\" as he flaps his wings with full commitment. "
+          "The woman covers her face, on the verge of tears. The tone is deadpan, "
+          "absurd, and quietly tragic.")
+
+# Sampling flags default to None: unset flags are NOT passed to
+# `generate_video`, so the bundle's preset supplies them via
+# `SamplingParam.from_pretrained` (ltx2_distilled: 8 steps, cfg 1.0,
+# 1024x1536, 121 frames).
+_SAMPLING_FLAGS = ("height", "width", "num_frames", "num_inference_steps", "guidance_scale", "seed")
+
+
+def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="LTX-2.5 distilled inference from a single-file bundle.")
+    parser.add_argument(
+        "--model-path",
+        required=True,
+        help="Path to the LTX-2.5 distilled bundle (.safetensors FILE).",
+    )
+    parser.add_argument(
+        "--gemma-root",
+        default=None,
+        help="Gemma text-encoder root directory. Must be the root paired with "
+        "this transformer variant; the roots differ in prompt templating.",
+    )
+    parser.add_argument("--prompt", default=PROMPT)
+    parser.add_argument("--output-path", default="outputs_video/ltx2_5/output_ltx2_5_distilled_t2v.mp4")
+    parser.add_argument("--num-gpus", type=int, default=1)
+    parser.add_argument("--height", type=int, default=None)
+    parser.add_argument("--width", type=int, default=None)
+    parser.add_argument("--num-frames", type=int, default=None)
+    parser.add_argument("--num-inference-steps", type=int, default=None)
+    parser.add_argument("--guidance-scale", type=float, default=None)
+    parser.add_argument("--seed", type=int, default=None)
+    return parser.parse_args(argv)
+
+
+def sampling_overrides(args: argparse.Namespace) -> dict:
+    """Only sampling flags the user actually passed; the preset supplies the rest."""
+    return {name: getattr(args, name) for name in _SAMPLING_FLAGS if getattr(args, name) is not None}
+
+
+def main(argv: list[str] | None = None) -> None:
+    args = parse_args(argv)
+    if args.gemma_root:
+        os.environ["FASTVIDEO_LTX_ENCODER_ROOT"] = args.gemma_root
+
+    generator = VideoGenerator.from_pretrained(
+        args.model_path,
+        num_gpus=args.num_gpus,
+    )
+    generator.generate_video(
+        prompt=args.prompt,
+        output_path=args.output_path,
+        save_video=True,
+        **sampling_overrides(args),
+    )
+    generator.shutdown()
+
+
+if __name__ == "__main__":
+    main()

examples/inference/basic/basic_mmaudio.py

@@ -0,0 +1,44 @@
+# SPDX-License-Identifier: Apache-2.0
+"""MMAudio large-44k-v2 video-to-audio example."""
+
+import argparse
+import os
+
+from fastvideo import VideoGenerator
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--video-path", required=True)
+    parser.add_argument("--output-path", default="outputs_audio/mmaudio.wav")
+    parser.add_argument("--duration-seconds", type=float, default=8.0)
+    parser.add_argument("--prompt", default="")
+    parser.add_argument("--negative-prompt", default="music")
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    generator = VideoGenerator.from_pretrained(
+        os.environ.get(
+            "MMAUDIO_MODEL_PATH",
+            "converted_weights/mmaudio/large_44k_v2",
+        ),
+        workload_type="v2a",
+        num_gpus=1,
+    )
+    result = generator.generate_video(
+        prompt=args.prompt,
+        negative_prompt=args.negative_prompt,
+        video_path=args.video_path,
+        audio_end_in_s=args.duration_seconds,
+        output_path=args.output_path,
+        save_video=True,
+        return_frames=False,
+    )
+    print(result["video_path"])
+    generator.shutdown()
+
+
+if __name__ == "__main__":
+    main()

examples/train/configs/fine_tuning/ltx2_5/nvfp4_qat_t2v.yaml

@@ -0,0 +1,130 @@
+# Dual-stream T2V NVFP4 quantization-aware fine-tune (QAT).
+#
+# Adapted from fine_tuning/ltx2_3/nvfp4_qat_t2v.yaml. The quantization wiring
+# is UNCHANGED and deliberately so: the curated linear prefix set in
+# fastvideo/layers/quantization/nvfp4_config.py targets
+# ltx2.blocks.{0..47}.* and attaches to this architecture unmodified --
+# verified against the constructed module tree, 576 linears / 12.9B params,
+# the same 12 suffix classes x 48 blocks as before. Do NOT narrow the set for
+# this architecture: the cross-modal projections it targets
+# (audio_to_video_attn, video_to_audio_attn) are deliberately enumerated in the
+# DEPLOYMENT set, so dropping them in training would break train/deploy
+# symmetry.
+#
+# What is NOT targeted, and why (all deliberate, none accidental):
+#   * head-dim-64 audio self/cross attention -- same constraint that kept it
+#     dense before;
+#   * per-head gate projections (to_gate_logits, [32, width]) -- tiny and
+#     precision-sensitive;
+#   * the embeddings connectors -- they live in the TEXT ENCODER, not the DiT,
+#     so deployment targeting must not reach across the component boundary.
+# The coverage audit prints these as skipped-by-rule so a reader can see the
+# exclusions were chosen, and FAILS if any declared class matches zero modules.
+#
+# Video attention uses ATTN_QAT_TRAIN (quantized forward, STE backward).
+#
+# Data paths, step counts and learning rate below are placeholders to adapt.
+
+models:
+  student:
+    _target_: fastvideo.train.models.ltx2.LTX2Model
+    # REQUIRED: single-file bundle path. The config parser does NOT expand
+    # environment interpolations, so override this on the command line:
+    #   --models.student.init_from /path/to/bundle.safetensors
+    init_from: SET_ME_ON_THE_COMMAND_LINE
+    trainable: true
+    enable_gradient_checkpointing_type: full
+    attention_backend: ATTN_QAT_TRAIN
+
+method:
+  _target_: fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod
+
+training:
+  distributed:
+    # ~19B params: bf16 weights 35.4 GiB + bf16 grads 35.4 GiB + Adam fp32
+    # m/v 141.5 GiB = ~212 GiB, against ~173 GiB usable per GPU. Single-GPU is
+    # arithmetically impossible; 4-way shard puts it at ~53 GiB/GPU.
+    num_gpus: 4
+    sp_size: 1
+    tp_size: 1
+    hsdp_replicate_dim: 1
+    hsdp_shard_dim: 4
+
+  data:
+    # REQUIRED: path to your preprocessed dataset.
+    data_path: data/your_dataset_preprocessed
+    dataloader_num_workers: 4
+    train_batch_size: 1
+    # LTX2Model requires 0.0: CFG dropout would zero post-connector
+    # embeddings, which is not the model's unconditional input.
+    training_cfg_rate: 0.0
+    seed: 42
+    # Must match your preprocessed resolution/length.
+    # num_latent_t = (num_frames - 1) / 8 + 1.
+    num_latent_t: 11
+    num_height: 480
+    num_width: 832
+    num_frames: 81
+
+  optimizer:
+    # Carried from the validated LTX-2.0 QAT overfit run as a starting
+    # point; tune for your dataset size and batch configuration.
+    learning_rate: 5.0e-5
+    betas: [0.9, 0.999]
+    weight_decay: 0.0
+    lr_scheduler: constant
+    lr_warmup_steps: 0
+
+  loop:
+    # Workload-dependent: set from your dataset size and target epochs.
+    max_train_steps: 2000
+    gradient_accumulation_steps: 1
+
+  checkpoint:
+    output_dir: outputs/dual_stream_t2v_nvfp4_qat_finetune
+    # A full training-state checkpoint is ~150GB for the 13B trainable
+    # video branch — size the interval and total limit to your storage.
+    training_state_checkpointing_steps: 500
+    checkpoints_total_limit: 2
+    resume_from_checkpoint: latest
+
+  tracker:
+    # No experiment tracking.
+    # It must be ["none"], NOT []: build_tracker appends wandb whenever the
+    # list is empty and project_name is non-empty, so [] silently means wandb.
+    trackers: ["none"]
+
+  model:
+    # LTX2Model.predict_noise converts the DiT's x0 output to velocity,
+    # so the default noise-minus-clean target reproduces the official
+    # unweighted masked-MSE (mask is all-ones for plain T2V).
+    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.ltx2.ltx2_pipeline.LTX2Pipeline
+    # REQUIRED: prompts to sample during training.
+    dataset_file: data/your_dataset_preprocessed/validation_prompts.json
+    every_steps: 250
+    # 8-step single-pass sampling matches the distilled checkpoint
+    # (validated in the LTX-2.3 overfit runs).
+    sampling_steps: [8]
+    guidance_scale: 1.0
+    num_frames: 81
+    # Validation reuses the live transformer and temporarily replaces its
+    # ATTN_QAT_TRAIN implementations with the sm120 inference kernel.
+    # Set false on GB200 (see header).
+    attn_qat_infer: true
+
+# quant_config selects NVFP4 QAT for the same curated linear prefixes as
+# LTX-2 NVFP4 deployment (see the LTX-2.3 applicability note in the
+# header). They fake-quantize through real FP4 GEMMs with an STE
+# backward. The string resolves to NVFP4QATTrainConfig at parse.
+pipeline:
+  dit_config:
+    quant_config: nvfp4_qat_train

examples/training/finetune/wan_t2v_1.3B/crush_smol/profile_finetune_t2v.sh

@@ -19,16 +19,13 @@ export FASTVIDEO_TORCH_PROFILER_RECORD_SHAPES=1
 export FASTVIDEO_TORCH_PROFILER_WITH_STACK=1  
 export FASTVIDEO_TORCH_PROFILER_WITH_FLOPS=1
 export FASTVIDEO_TORCH_PROFILER_WITH_PROFILE_MEMORY=1
-export FASTVIDEO_TORCH_PROFILER_WAIT_STEPS=2
-export FASTVIDEO_TORCH_PROFILER_WARMUP_STEPS=1
-export FASTVIDEO_TORCH_PROFILER_ACTIVE_STEPS=1
 export FASTVIDEO_TORCH_PROFILER_DIR="../profiler_traces/wan_t2v_finetune/"
 
 # Training arguments
 training_args=(
   --tracker_project_name "wan_t2v_finetune"
   --output_dir "checkpoints/wan_t2v_finetune"
-  --max_train_steps 5000
+  --max_train_steps 1  # Profiler captures every selected region until shutdown.
   --train_batch_size 1
   --train_sp_batch_size 1
   --gradient_accumulation_steps 8

fastvideo/api/schema.py

@@ -100,7 +100,7 @@ class ComponentConfig:
 
 @dataclass
 class PipelineSelection:
-    workload_type: Literal["t2v", "i2v", "t2i", "i2i"] | None = None
+    workload_type: Literal["t2v", "i2v", "t2i", "i2i", "v2a", "t2a"] | None = None
     preset: str | None = None
     preset_version: int | None = None
     components: ComponentConfig = field(default_factory=ComponentConfig)

fastvideo/attention/backends/flash_attn.py

@@ -168,8 +168,12 @@ class FlashAttentionBackend(AttentionBackend):
 def _key_padding_mask_from_attn_mask(attn_mask: torch.Tensor, key_len: int) -> torch.Tensor:
     # Normalize attn_mask to [B, key_len] where True means valid token.
     if attn_mask.dim() == 4:
+        if attn_mask.shape[1] != 1 or attn_mask.shape[-2] != 1:
+            raise ValueError("FLASH_ATTN only supports 4D key-padding masks with shape [B, 1, 1, K]")
         attn_mask = attn_mask[:, 0, 0, :]
     elif attn_mask.dim() == 3:
+        if attn_mask.shape[-2] != 1:
+            raise ValueError("FLASH_ATTN only supports 3D key-padding masks with shape [B, 1, K]")
         attn_mask = attn_mask[:, 0, :]
     elif attn_mask.dim() != 2:
         raise ValueError(f"Unsupported attn_mask shape for FLASH_ATTN: {attn_mask.shape}")
@@ -274,6 +278,14 @@ class FlashAttentionImpl(AttentionImpl):
             )
 
             attn_mask = attn_metadata.attn_mask
+            if getattr(attn_metadata, "is_causal", False):
+                return flash_attn_func_compilable(
+                    query,
+                    key,
+                    value,
+                    softmax_scale=self.softmax_scale,
+                    causal=True,
+                )
 
             # flash_attn_no_pad packs q/k/v as one tensor and assumes equal q/k
             # sequence lengths. Cross-attention can violate this.
@@ -302,7 +314,7 @@ class FlashAttentionImpl(AttentionImpl):
                 raise ValueError("Invalid key padding mask length for FLASH_ATTN: "
                                  f"expected at most {qkv.shape[1]}, got {key_padding_mask.shape[-1]}")
             attn_mask_padded = F.pad(key_padding_mask, (qkv.shape[1] - key_padding_mask.shape[-1], 0), value=True)
-            output = flash_attn_no_pad(qkv, attn_mask_padded, causal=False, dropout_p=0, softmax_scale=None)
+            output = flash_attn_no_pad(qkv, attn_mask_padded, causal=self.causal, dropout_p=0, softmax_scale=None)
         elif self.nvfp4_fa4:
             output = self._forward_nvfp4(query, key, value)
 

fastvideo/attention/backends/sdpa.py

@@ -41,6 +41,8 @@ class SDPABackend(AttentionBackend):
 class SDPAMetadata(AttentionMetadata):
     current_timestep: int
     attn_mask: torch.Tensor | None = None
+    # The mask is exactly native causal attention, with no additional padding.
+    is_causal: bool = False
 
 
 class SDPAMetadataBuilder(AttentionMetadataBuilder):

fastvideo/configs/models/__init__.py

@@ -2,7 +2,13 @@ from fastvideo.configs.models.base import ModelConfig
 from fastvideo.configs.models.dits.base import DiTConfig
 from fastvideo.configs.models.encoders.base import EncoderConfig
 from fastvideo.configs.models.vaes.base import VAEConfig
-from fastvideo.configs.models.audio import (LTX2AudioDecoderConfig, LTX2AudioEncoderConfig, LTX2VocoderConfig)
+from fastvideo.configs.models.audio import (
+    BigVGANV2Config,
+    LTX2AudioDecoderConfig,
+    LTX2AudioEncoderConfig,
+    LTX2VocoderConfig,
+    MMAudioVAEConfig,
+)
 from fastvideo.configs.models.upsamplers.base import UpsamplerConfig
 
 __all__ = [
@@ -13,5 +19,7 @@ __all__ = [
     "LTX2AudioEncoderConfig",
     "LTX2AudioDecoderConfig",
     "LTX2VocoderConfig",
+    "MMAudioVAEConfig",
+    "BigVGANV2Config",
     "UpsamplerConfig",
 ]

fastvideo/configs/models/audio/__init__.py

@@ -5,9 +5,21 @@ from fastvideo.configs.models.audio.ltx2_audio_vae import (
     LTX2AudioEncoderConfig,
     LTX2VocoderConfig,
 )
+from fastvideo.configs.models.audio.mmaudio_vae import (
+    MMAudioVAEArchConfig,
+    MMAudioVAEConfig,
+)
+from fastvideo.configs.models.audio.bigvgan import (
+    BigVGANV2ArchConfig,
+    BigVGANV2Config,
+)
 
 __all__ = [
     "LTX2AudioEncoderConfig",
     "LTX2AudioDecoderConfig",
     "LTX2VocoderConfig",
+    "MMAudioVAEArchConfig",
+    "MMAudioVAEConfig",
+    "BigVGANV2ArchConfig",
+    "BigVGANV2Config",
 ]

fastvideo/configs/models/audio/bigvgan.py

@@ -0,0 +1,18 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Reusable BigVGAN-v2 vocoder configuration."""
+
+from dataclasses import dataclass, field
+
+from fastvideo.configs.models.base import ArchConfig, ModelConfig
+
+
+@dataclass
+class BigVGANV2ArchConfig(ArchConfig):
+    architectures: list[str] = field(default_factory=lambda: ["BigVGANV2"])
+    sample_rate: int = 44100
+    num_mels: int = 128
+
+
+@dataclass
+class BigVGANV2Config(ModelConfig):
+    arch_config: ArchConfig = field(default_factory=BigVGANV2ArchConfig)

fastvideo/configs/models/audio/mmaudio_vae.py

@@ -0,0 +1,21 @@
+# SPDX-License-Identifier: Apache-2.0
+"""MMAudio audio VAE configuration."""
+
+from dataclasses import dataclass, field
+
+from fastvideo.configs.models.base import ArchConfig, ModelConfig
+
+
+@dataclass
+class MMAudioVAEArchConfig(ArchConfig):
+    architectures: list[str] = field(default_factory=lambda: ["MMAudioVAE"])
+    mode: str = "44k"
+    data_dim: int = 128
+    embed_dim: int = 40
+    hidden_dim: int = 512
+    need_encoder: bool = False
+
+
+@dataclass
+class MMAudioVAEConfig(ModelConfig):
+    arch_config: ArchConfig = field(default_factory=MMAudioVAEArchConfig)

fastvideo/configs/models/dits/__init__.py

@@ -11,6 +11,7 @@ from fastvideo.configs.models.dits.longcat import LongCatVideoConfig
 from fastvideo.configs.models.dits.ltx2 import LTX2VideoConfig
 from fastvideo.configs.models.dits.magi_human import MagiHumanVideoConfig
 from fastvideo.configs.models.dits.minimax_h3 import MiniMaxH3Config
+from fastvideo.configs.models.dits.mmaudio import MMAudioArchConfig, MMAudioTransformerConfig
 from fastvideo.configs.models.dits.stable_audio import StableAudioConfig
 from fastvideo.configs.models.dits.wanvideo import WanVideoConfig
 from fastvideo.configs.models.dits.zimage import ZImageDiTConfig
@@ -24,5 +25,5 @@ __all__ = [
     "DreamXWorldARConfig", "CosmosVideoConfig", "Cosmos25VideoConfig", "FluxDiTConfig", "Flux2Config",
     "LongCatVideoConfig", "LTX2VideoConfig", "HYWorldConfig", "Kandinsky5VideoConfig", "MagiHumanVideoConfig",
     "StableAudioConfig", "GlmImageDiTConfig", "LingBotWorld2CausalFastVideoConfig", "LingBotVideoConfig",
-    "MiniMaxH3Config", "ZImageDiTConfig"
+    "MiniMaxH3Config", "ZImageDiTConfig", "MMAudioArchConfig", "MMAudioTransformerConfig"
 ]

fastvideo/configs/models/dits/ltx2.py

@@ -55,6 +55,13 @@ class LTX2VideoArchConfig(DiTArchConfig):
     # LTX-2.3 gated extensions. All default OFF == LTX-2.0 behavior.
     cross_attention_adaln: bool = False
     caption_proj_before_connector: bool = False
+    # FFN bias, per stream. Some checkpoints ship the video FFN without bias
+    # (``ff_bias: false`` in their metadata), which drops the 96
+    # ``transformer_blocks.*.ff.net.{0.proj,2}.bias`` tensors; the audio FFN is
+    # configured independently and commonly keeps its biases. Both default True,
+    # which is the existing behavior.
+    ff_bias: bool = True
+    audio_ff_bias: bool = True
 
     positional_embedding_theta: float = 10000.0
     positional_embedding_max_pos: list[int] = field(default_factory=lambda: [20, 2048, 2048])

fastvideo/configs/models/dits/mmaudio.py

@@ -0,0 +1,49 @@
+# SPDX-License-Identifier: Apache-2.0
+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_mmaudio_transformer_block(name: str, module) -> bool:
+    del module
+    parts = name.split(".")
+    return len(parts) >= 2 and parts[-1].isdigit() and parts[-2] in {"joint_blocks", "fused_blocks"}
+
+
+@dataclass
+class MMAudioArchConfig(DiTArchConfig):
+    _fsdp_shard_conditions: list = field(default_factory=lambda: [_is_mmaudio_transformer_block])
+    param_names_mapping: dict = field(default_factory=lambda: {r"^(.*)$": r"\1"})
+    _supported_attention_backends: tuple[AttentionBackendEnum, ...] = (AttentionBackendEnum.TORCH_SDPA, )
+
+    latent_dim: int = 40
+    clip_dim: int = 1024
+    sync_dim: int = 768
+    text_dim: int = 1024
+    hidden_dim: int = 896
+    depth: int = 21
+    fused_depth: int = 14
+    num_heads: int = 14
+    mlp_ratio: float = 4.0
+    latent_seq_len: int = 345
+    clip_seq_len: int = 64
+    sync_seq_len: int = 192
+    text_seq_len: int = 77
+    v2: bool = True
+
+    def __post_init__(self) -> None:
+        super().__post_init__()
+        self.hidden_size = self.hidden_dim
+        self.num_attention_heads = self.num_heads
+        self.num_channels_latents = self.latent_dim
+        self.in_channels = self.latent_dim
+        self.out_channels = self.latent_dim
+
+
+@dataclass
+class MMAudioTransformerConfig(DiTConfig):
+    arch_config: DiTArchConfig = field(default_factory=MMAudioArchConfig)
+    prefix: str = "MMAudio"

fastvideo/configs/models/encoders/__init__.py

@@ -1,6 +1,10 @@
-from fastvideo.configs.models.encoders.base import (BaseEncoderOutput, EncoderConfig, ImageEncoderConfig,
-                                                    TextEncoderConfig)
-from fastvideo.configs.models.encoders.clip import (CLIPTextConfig, CLIPVisionConfig, WAN2_1ControlCLIPVisionConfig)
+from fastvideo.configs.models.encoders.base import (
+    BaseEncoderOutput,
+    EncoderConfig,
+    ImageEncoderConfig,
+    TextEncoderConfig,
+)
+from fastvideo.configs.models.encoders.clip import CLIPTextConfig, CLIPVisionConfig, WAN2_1ControlCLIPVisionConfig
 from fastvideo.configs.models.encoders.llama import LlamaConfig
 from fastvideo.configs.models.encoders.lingbotworld2_t5 import LingBotWorld2UMT5ArchConfig, LingBotWorld2UMT5Config
 from fastvideo.configs.models.encoders.t5 import T5Config, T5LargeConfig
@@ -12,15 +16,49 @@ from fastvideo.configs.models.encoders.mistral3 import Mistral3TextConfig
 from fastvideo.configs.models.encoders.minimax_h3_qwen3_vl import (MiniMaxH3Qwen3VLArchConfig, MiniMaxH3Qwen3VLConfig)
 from fastvideo.configs.models.encoders.qwen3 import Qwen3TextConfig
 from fastvideo.configs.models.encoders.lingbot_video import LingBotVideoQwen3VLTextConfig
-from fastvideo.configs.models.encoders.stable_audio_conditioner import (StableAudioConditionerArchConfig,
-                                                                        StableAudioConditionerConfig)
+from fastvideo.configs.models.encoders.stable_audio_conditioner import (
+    StableAudioConditionerArchConfig,
+    StableAudioConditionerConfig,
+)
 from fastvideo.configs.models.encoders.t5gemma import T5GemmaEncoderConfig
+from fastvideo.configs.models.encoders.mmaudio_synchformer import MMAudioSynchformerArchConfig, MMAudioSynchformerConfig
+from fastvideo.configs.models.encoders.mmaudio_clip import (
+    MMAudioDFNCLIPTextArchConfig,
+    MMAudioDFNCLIPTextConfig,
+    MMAudioDFNCLIPVisionArchConfig,
+    MMAudioDFNCLIPVisionConfig,
+)
 
 __all__ = [
-    "EncoderConfig", "TextEncoderConfig", "ImageEncoderConfig", "BaseEncoderOutput", "CLIPTextConfig",
-    "CLIPVisionConfig", "WAN2_1ControlCLIPVisionConfig", "LlamaConfig", "T5Config", "T5LargeConfig", "Qwen2_5_VLConfig",
-    "Reason1ArchConfig", "Reason1Config", "LTX2GemmaConfig", "SiglipVisionConfig", "StableAudioConditionerArchConfig",
-    "StableAudioConditionerConfig", "T5GemmaEncoderConfig", "Qwen3TextConfig", "Mistral3TextConfig",
-    "LingBotWorld2UMT5ArchConfig", "LingBotWorld2UMT5Config", "LingBotVideoQwen3VLTextConfig",
-    "MiniMaxH3Qwen3VLArchConfig", "MiniMaxH3Qwen3VLConfig"
+    "EncoderConfig",
+    "TextEncoderConfig",
+    "ImageEncoderConfig",
+    "BaseEncoderOutput",
+    "CLIPTextConfig",
+    "CLIPVisionConfig",
+    "WAN2_1ControlCLIPVisionConfig",
+    "LlamaConfig",
+    "T5Config",
+    "T5LargeConfig",
+    "Qwen2_5_VLConfig",
+    "Reason1ArchConfig",
+    "Reason1Config",
+    "LTX2GemmaConfig",
+    "SiglipVisionConfig",
+    "StableAudioConditionerArchConfig",
+    "StableAudioConditionerConfig",
+    "T5GemmaEncoderConfig",
+    "Qwen3TextConfig",
+    "Mistral3TextConfig",
+    "LingBotWorld2UMT5ArchConfig",
+    "LingBotWorld2UMT5Config",
+    "LingBotVideoQwen3VLTextConfig",
+    "MiniMaxH3Qwen3VLArchConfig",
+    "MiniMaxH3Qwen3VLConfig",
+    "MMAudioSynchformerArchConfig",
+    "MMAudioSynchformerConfig",
+    "MMAudioDFNCLIPTextArchConfig",
+    "MMAudioDFNCLIPTextConfig",
+    "MMAudioDFNCLIPVisionArchConfig",
+    "MMAudioDFNCLIPVisionConfig",
 ]

fastvideo/configs/models/encoders/mmaudio_clip.py

@@ -0,0 +1,67 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Native DFN5B CLIP conditioner configurations for MMAudio."""
+
+from dataclasses import dataclass, field
+
+from fastvideo.configs.models.encoders.clip import (
+    CLIPTextArchConfig,
+    CLIPTextConfig,
+    CLIPVisionArchConfig,
+    CLIPVisionConfig,
+)
+
+
+@dataclass
+class MMAudioDFNCLIPTextArchConfig(CLIPTextArchConfig):
+    architectures: list[str] = field(default_factory=lambda: ["MMAudioDFNCLIPTextEncoder"])
+    vocab_size: int = 49408
+    hidden_size: int = 1024
+    intermediate_size: int = 4096
+    projection_dim: int = 1024
+    num_hidden_layers: int = 24
+    num_attention_heads: int = 16
+    max_position_embeddings: int = 77
+    text_len: int = 77
+    hidden_act: str = "quick_gelu"
+    layer_norm_eps: float = 1e-5
+    pad_token_id: int = 0
+    bos_token_id: int = 49406
+    eos_token_id: int = 49407
+    # MMAudio moves this encoder as one unit between CPU and GPU. Keeping it a
+    # plain module avoids nesting FSDP CPU-offload semantics inside the custom
+    # OpenCLIP causal-mask forward used by this pipeline.
+    _fsdp_shard_conditions: list = field(default_factory=list)
+
+
+@dataclass
+class MMAudioDFNCLIPTextConfig(CLIPTextConfig):
+    arch_config: CLIPTextArchConfig = field(default_factory=MMAudioDFNCLIPTextArchConfig)
+    # OpenCLIP supplies an explicit additive triangular mask to
+    # nn.MultiheadAttention. The MMAudio adapter reproduces that path instead
+    # of using SDPA's is_causal shortcut, which rounds differently in bf16.
+    is_causal: bool = False
+    prefix: str = "mmaudio_dfn_clip_text"
+
+
+@dataclass
+class MMAudioDFNCLIPVisionArchConfig(CLIPVisionArchConfig):
+    architectures: list[str] = field(default_factory=lambda: ["MMAudioDFNCLIPVisionEncoder"])
+    hidden_size: int = 1280
+    intermediate_size: int = 5120
+    projection_dim: int = 1024
+    num_hidden_layers: int = 32
+    num_attention_heads: int = 16
+    image_size: int = 378
+    patch_size: int = 14
+    hidden_act: str = "quick_gelu"
+    layer_norm_eps: float = 1e-5
+
+
+@dataclass
+class MMAudioDFNCLIPVisionConfig(CLIPVisionConfig):
+    arch_config: CLIPVisionArchConfig = field(default_factory=MMAudioDFNCLIPVisionArchConfig)
+    num_hidden_layers_override: int | None = None
+    require_post_norm: bool | None = True
+    enable_scale: bool = True
+    is_causal: bool = False
+    prefix: str = "mmaudio_dfn_clip_vision"

fastvideo/configs/models/encoders/mmaudio_synchformer.py

@@ -0,0 +1,26 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Configuration for MMAudio's Synchformer visual conditioner."""
+
+from dataclasses import dataclass, field
+
+from fastvideo.configs.models.encoders.base import (
+    ImageEncoderArchConfig,
+    ImageEncoderConfig,
+)
+
+
+@dataclass
+class MMAudioSynchformerArchConfig(ImageEncoderArchConfig):
+    architectures: list[str] = field(default_factory=lambda: ["MMAudioSynchformerVisualEncoder"])
+    image_size: int = 224
+    num_channels: int = 3
+    segment_size: int = 16
+    segment_stride: int = 8
+    hidden_size: int = 768
+    tokens_per_segment: int = 8
+
+
+@dataclass
+class MMAudioSynchformerConfig(ImageEncoderConfig):
+    arch_config: ImageEncoderArchConfig = field(default_factory=MMAudioSynchformerArchConfig)
+    prefix: str = "synchformer"

fastvideo/configs/pipelines/__init__.py

@@ -11,6 +11,7 @@ from fastvideo.configs.pipelines.lingbotworld2 import LingBotWorld2CausalFastI2V
 from fastvideo.configs.pipelines.lingbot_video import LingBotVideoT2VConfig
 from fastvideo.configs.pipelines.matrixgame2 import MatrixGame2I2V480PConfig
 from fastvideo.configs.pipelines.matrixgame3 import MatrixGame3I2V720PConfig
+from fastvideo.configs.pipelines.mmaudio import MMAudioV2AConfig
 from fastvideo.pipelines.basic.ltx2.pipeline_configs import LTX2T2VConfig
 from fastvideo.registry import get_pipeline_config_cls_from_name
 from fastvideo.configs.pipelines.wan import (LucyEditDevConfig, SelfForcingWanT2V480PConfig, WanI2V480PConfig,
@@ -22,5 +23,5 @@ __all__ = [
     "SelfForcingWanT2V480PConfig", "LucyEditDevConfig", "CosmosConfig", "Cosmos25Config", "LTX2T2VConfig",
     "DreamXWorld5BCamPipelineConfig", "DreamXWorld5BARPipelineConfig", "HYWorldConfig", "Kandinsky5T2VConfig",
     "Kandinsky5I2VConfig", "Kandinsky5DMDConfig", "LingBotWorld2CausalFastI2V480PConfig", "LingBotVideoT2VConfig",
-    "MatrixGame2I2V480PConfig", "MatrixGame3I2V720PConfig", "get_pipeline_config_cls_from_name"
+    "MatrixGame2I2V480PConfig", "MatrixGame3I2V720PConfig", "MMAudioV2AConfig", "get_pipeline_config_cls_from_name"
 ]

fastvideo/configs/pipelines/base.py

@@ -63,6 +63,11 @@ class PipelineConfig:
     # Image encoder configuration
     image_encoder_config: EncoderConfig = field(default_factory=EncoderConfig)
     image_encoder_precision: str = "fp32"
+    # Optional multi-encoder contract. Existing pipelines continue to use the
+    # singular fields above; V2A and other multimodal pipelines can opt into
+    # indexed ``image_encoder``, ``image_encoder_2``, ... components.
+    image_encoder_configs: tuple[EncoderConfig, ...] | None = None
+    image_encoder_precisions: tuple[str, ...] | None = None
 
     # Text encoder configuration
     DEFAULT_TEXT_ENCODER_PRECISIONS = ("fp32", )

fastvideo/configs/pipelines/mmaudio.py

@@ -0,0 +1,66 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Pipeline configuration for the native MMAudio video-to-audio port."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+from fastvideo.configs.models import DiTConfig, EncoderConfig, ModelConfig
+from fastvideo.configs.models.audio import BigVGANV2Config, MMAudioVAEConfig
+from fastvideo.configs.models.dits import MMAudioTransformerConfig
+from fastvideo.configs.models.encoders import (
+    MMAudioDFNCLIPTextConfig,
+    MMAudioDFNCLIPVisionConfig,
+    MMAudioSynchformerConfig,
+)
+from fastvideo.configs.pipelines.base import PipelineConfig
+
+
+@dataclass
+class MMAudioV2AConfig(PipelineConfig):
+    """MMAudio large-44k-v2 inference defaults.
+
+    The published demo moves every module to bfloat16. Keeping the same
+    per-component precision here is important: condition features seed the
+    complete flow trajectory, so silently encoding them in fp32 changes the
+    generated waveform even when the transformer weights are identical.
+    """
+
+    dit_config: DiTConfig = field(default_factory=MMAudioTransformerConfig)
+    dit_precision: str = "bf16"
+
+    text_encoder_configs: tuple[EncoderConfig, ...] = field(default_factory=lambda: (MMAudioDFNCLIPTextConfig(), ))
+    text_encoder_precisions: tuple[str, ...] = field(default_factory=lambda: ("bf16", ))
+    image_encoder_configs: tuple[EncoderConfig, ...] = field(default_factory=lambda: (
+        MMAudioDFNCLIPVisionConfig(),
+        MMAudioSynchformerConfig(),
+    ))
+    image_encoder_precisions: tuple[str, ...] = field(default_factory=lambda: ("bf16", "bf16"))
+
+    audio_decoder_config: ModelConfig = field(default_factory=MMAudioVAEConfig)
+    audio_decoder_precision: str = "bf16"
+    vocoder_config: ModelConfig = field(default_factory=BigVGANV2Config)
+    vocoder_precision: str = "bf16"
+
+    # Published large_44k_v2 default sequence contract. The official demo
+    # supports other durations, although quality can drop far away from the
+    # eight-second training duration.
+    duration_s: float = 8.0
+    max_audio_duration_s: float | None = None
+    sampling_rate: int = 44100
+    spectrogram_frame_rate: int = 512
+    latent_downsample_rate: int = 2
+    clip_frame_rate: int = 8
+    sync_frame_rate: int = 25
+    sync_segment_size: int = 16
+    sync_segment_stride: int = 8
+    sync_downsample_rate: int = 2
+    clip_image_size: int = 384
+    sync_image_size: int = 224
+    clip_batch_size_multiplier: int = 40
+    sync_batch_size_multiplier: int = 40
+
+    num_inference_steps: int = 25
+    guidance_scale: float = 4.5
+    vae_tiling: bool = False
+    vae_sp: bool = False

fastvideo/entrypoints/video_generator.py

@@ -50,7 +50,7 @@ from fastvideo.api.schema import (
     SamplingConfig,
 )
 from fastvideo.api.sampling_param import SamplingParam
-from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.fastvideo_args import FastVideoArgs, WorkloadType
 from fastvideo.logger import init_logger
 from fastvideo.pipelines import ForwardBatch
 from fastvideo.utils import align_to, shallow_asdict
@@ -605,7 +605,13 @@ class VideoGenerator:
 
         # Single prompt generation (original behavior)
         if prompt is None:
-            raise ValueError("Either prompt or prompt_txt must be provided")
+            if fastvideo_args.workload_type is WorkloadType.V2A:
+                # Video semantics are sufficient conditioning for V2A models;
+                # model-specific text stages interpret the empty string using
+                # their native tokenizer/empty-prompt contract.
+                prompt = ""
+            else:
+                raise ValueError("Either prompt or prompt_txt must be provided")
         output_path = self._prepare_output_path(sampling_param.output_path, prompt)
         kwargs["output_path"] = output_path
         if prompt_embeds is not None:
@@ -624,6 +630,13 @@ class VideoGenerator:
             return False
         return args.workload_type.value.endswith("2i")
 
+    def _is_audio_workload(self) -> bool:
+        """Return True when the workload produces standalone audio."""
+        args = getattr(self, "fastvideo_args", None)
+        if args is None:
+            return False
+        return args.workload_type.value.endswith("2a")
+
     def _prepare_output_path(
         self,
         output_path: str,
@@ -643,7 +656,12 @@ class VideoGenerator:
           warning is logged.
         - If the target path already exists, a numeric suffix is appended.
         """
-        target_ext = ".png" if self._is_image_workload() else ".mp4"
+        if self._is_image_workload():
+            target_ext = ".png"
+        elif self._is_audio_workload():
+            target_ext = ".wav"
+        else:
+            target_ext = ".mp4"
 
         def _sanitize_filename_component(name: str) -> str:
             # Remove characters invalid on common filesystems, strip spaces/dots

fastvideo/envs.py

@@ -30,11 +30,8 @@ if TYPE_CHECKING:
     FASTVIDEO_TORCH_PROFILER_DIR: str | None = None
     FASTVIDEO_TORCH_PROFILER_RECORD_SHAPES: bool = False
     FASTVIDEO_TORCH_PROFILER_WITH_PROFILE_MEMORY: bool = False
-    FASTVIDEO_TORCH_PROFILER_WITH_STACK: bool = True
+    FASTVIDEO_TORCH_PROFILER_WITH_STACK: bool = False
     FASTVIDEO_TORCH_PROFILER_WITH_FLOPS: bool = False
-    FASTVIDEO_TORCH_PROFILER_WAIT_STEPS: int = 2
-    FASTVIDEO_TORCH_PROFILER_WARMUP_STEPS: int = 1
-    FASTVIDEO_TORCH_PROFILER_ACTIVE_STEPS: int = 2
     FASTVIDEO_TORCH_PROFILE_REGIONS: str = ""
     FASTVIDEO_TRACE_ACTIVATIONS: bool = False
     FASTVIDEO_TRACE_LAYERS: str = ""
@@ -242,11 +239,11 @@ environment_variables: dict[str, Callable[[], Any]] = {
     "FASTVIDEO_TORCH_PROFILER_WITH_PROFILE_MEMORY":
     lambda: bool(os.getenv("FASTVIDEO_TORCH_PROFILER_WITH_PROFILE_MEMORY", "0") != "0"),
 
-    # Enable torch profiler to profile stack if set
-    # FASTVIDEO_TORCH_PROFILER_WITH_STACK=1. If not set, torch profiler WILL
-    # profile stack by default.
+    # Enable torch profiler stack capture with
+    # FASTVIDEO_TORCH_PROFILER_WITH_STACK=1. Off by default: stack capture
+    # costs ~1.5x runtime overhead and ~1.4x trace size.
     "FASTVIDEO_TORCH_PROFILER_WITH_STACK":
-    lambda: bool(os.getenv("FASTVIDEO_TORCH_PROFILER_WITH_STACK", "1") != "0"),
+    lambda: bool(os.getenv("FASTVIDEO_TORCH_PROFILER_WITH_STACK", "0") != "0"),
 
     # Enable torch profiler to profile flops if set
     # FASTVIDEO_TORCH_PROFILER_WITH_FLOPS=1. If not set, torch profiler will
@@ -255,16 +252,10 @@ environment_variables: dict[str, Callable[[], Any]] = {
     lambda: bool(os.getenv("FASTVIDEO_TORCH_PROFILER_WITH_FLOPS", "0") != "0"),
     # Wait steps per profiling cycle (torch.profiler.schedule wait parameter)
     # Defaults to 2 if not set.
-    "FASTVIDEO_TORCH_PROFILER_WAIT_STEPS":
-    lambda: int(os.getenv("FASTVIDEO_TORCH_PROFILER_WAIT_STEPS", "2")),
     # Warmup steps per profiling cycle (torch.profiler.schedule warmup parameter)
     # Defaults to 1 if not set.
-    "FASTVIDEO_TORCH_PROFILER_WARMUP_STEPS":
-    lambda: int(os.getenv("FASTVIDEO_TORCH_PROFILER_WARMUP_STEPS", "1")),
     # Active steps per profiling cycle (torch.profiler.schedule active parameter)
     # Defaults to 2 if not set.
-    "FASTVIDEO_TORCH_PROFILER_ACTIVE_STEPS":
-    lambda: int(os.getenv("FASTVIDEO_TORCH_PROFILER_ACTIVE_STEPS", "2")),
     "FASTVIDEO_TORCH_PROFILE_REGIONS":
     lambda: os.getenv("FASTVIDEO_TORCH_PROFILE_REGIONS", ""),
 

fastvideo/fastvideo_args.py

@@ -61,6 +61,8 @@ class WorkloadType(str, Enum):
     T2V = "t2v"  # Text to Video
     T2I = "t2i"  # Text to Image
     I2I = "i2i"  # Image to Image
+    V2A = "v2a"  # Video to Audio
+    T2A = "t2a"  # Text to Audio
 
     @classmethod
     def from_string(cls, value: str) -> "WorkloadType":

fastvideo/layers/quantization/nvfp4_config.py

@@ -82,6 +82,79 @@ def is_ltx2_nvfp4_linear_prefix(prefix: str) -> bool:
     return prefix in _LTX2_NVFP4_LINEAR_PREFIXES
 
 
+def audit_nvfp4_coverage(
+    model,
+    predicate=None,
+    target_classes=None,
+    skip_classes=(),
+) -> dict[str, object]:
+    """Report which declared target classes actually attached, and raise if any did not.
+
+    The target set is written as module-prefix suffixes, but a tensor carries
+    three different names on its way here -- the checkpoint key, the
+    ``named_modules()`` path, and the ``prefix=`` string a layer is constructed
+    with -- and only the last one reaches the predicate. A rename on any of the
+    others leaves this set syntactically valid and silently matching nothing,
+    and the attached COUNT alone cannot show it: a set that quietly stops
+    covering a whole class still reports a large, healthy-looking number.
+
+    So: a declared class matching zero modules is an error, not a warning. This
+    guard did NOT fire when it was written -- the shipped set is correct today.
+    It exists for the next rename, which is the only kind of failure that can
+    reach production looking like success.
+
+    Returns the receipt: attached count and params, per-class attached counts,
+    and the classes deliberately NOT targeted, so a reader can see the
+    exclusions were chosen rather than lost.
+    """
+    from fastvideo.layers.linear import LinearBase
+
+    if predicate is None:
+        predicate = is_ltx2_nvfp4_linear_prefix
+    if target_classes is None:
+        target_classes = _LTX2_NVFP4_BLOCK_LINEAR_SUFFIXES
+
+    linears = [(getattr(m, "prefix", None) or n, n, m)
+               for n, m in model.named_modules() if isinstance(m, LinearBase)]
+
+    def _params(mod):
+        w = getattr(mod, "weight", None)
+        return int(w.shape[0] * w.shape[1]) if w is not None and w.dim() == 2 else 0
+
+    per_class: dict[str, int] = {}
+    for suffix in target_classes:
+        per_class[suffix] = sum(1 for pfx, _, _ in linears
+                                if pfx.endswith("." + suffix) and predicate(pfx))
+    empty = sorted(k for k, v in per_class.items() if v == 0)
+    if empty:
+        raise ValueError(
+            "NVFP4 target classes matched ZERO modules: " + ", ".join(empty) +
+            ". The model's actual Linear prefixes look like: " +
+            ", ".join(sorted({p for p, _, _ in linears})[:5]) +
+            ". Either the module naming changed or the target set is stale -- "
+            "refusing to train a model that silently quantizes less than declared.")
+
+    attached = [(p, n, m) for p, n, m in linears if predicate(p)]
+    skipped: dict[str, int] = {}
+    for pfx, name, _ in linears:
+        if predicate(pfx):
+            continue
+        tail = name.rsplit(".", 1)[-1]
+        for known in skip_classes:
+            if known in name:
+                skipped[known] = skipped.get(known, 0) + 1
+                break
+        else:
+            skipped["other:" + tail] = skipped.get("other:" + tail, 0) + 1
+    return {
+        "linears_total": len(linears),
+        "attached": len(attached),
+        "attached_params_M": round(sum(_params(m) for _, _, m in attached) / 1e6, 1),
+        "attached_per_class": per_class,
+        "skipped_by_rule": dict(sorted(skipped.items())),
+    }
+
+
 def _is_ltx2_refine_only_prefix(prefix: str) -> bool:
     return any(prefix.endswith(suffix) for suffix in _LTX2_REFINE_ONLY_SUFFIXES)
 
@@ -552,4 +625,5 @@ __all__ = [
     "NVFP4QuantizeMethod",
     "convert_model_to_nvfp4",
     "is_ltx2_nvfp4_linear_prefix",
+    "audit_nvfp4_coverage",
 ]

fastvideo/models/audio/bigvgan.py

@@ -0,0 +1,376 @@
+# SPDX-License-Identifier: MIT
+"""Reusable native BigVGAN-v2 vocoder.
+
+Adapted from NVIDIA BigVGAN-v2 and its alias-free activation implementation.
+The CUDA activation kernel is intentionally excluded; FastVideo uses the
+portable PyTorch path for deterministic loading and parity.
+"""
+
+from __future__ import annotations
+
+import math
+from typing import Any
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn import Conv1d, ConvTranspose1d
+from torch.nn.utils.parametrizations import weight_norm
+from torch.nn.utils.parametrize import remove_parametrizations
+
+
+class AttrDict(dict):
+    def __init__(self, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+        self.__dict__ = self
+
+
+def get_padding(kernel_size: int, dilation: int = 1) -> int:
+    return int((kernel_size * dilation - dilation) / 2)
+
+
+def init_weights(module: nn.Module, mean: float = 0.0, std: float = 0.01) -> None:
+    if "Conv" in module.__class__.__name__:
+        module.weight.data.normal_(mean, std)
+
+
+def kaiser_sinc_filter1d(cutoff: float, half_width: float, kernel_size: int) -> torch.Tensor:
+    even = kernel_size % 2 == 0
+    half_size = kernel_size // 2
+    delta_f = 4 * half_width
+    amplitude = 2.285 * (half_size - 1) * math.pi * delta_f + 7.95
+    if amplitude > 50.0:
+        beta = 0.1102 * (amplitude - 8.7)
+    elif amplitude >= 21.0:
+        beta = 0.5842 * (amplitude - 21) ** 0.4 + 0.07886 * (amplitude - 21.0)
+    else:
+        beta = 0.0
+    window = torch.kaiser_window(kernel_size, beta=beta, periodic=False)
+    time = torch.arange(-half_size, half_size) + 0.5 if even else torch.arange(kernel_size) - half_size
+    if cutoff == 0:
+        kernel = torch.zeros_like(time)
+    else:
+        kernel = 2 * cutoff * window * torch.sinc(2 * cutoff * time)
+        kernel /= kernel.sum()
+    return kernel.view(1, 1, kernel_size)
+
+
+class LowPassFilter1d(nn.Module):
+    def __init__(
+        self,
+        cutoff: float = 0.5,
+        half_width: float = 0.6,
+        stride: int = 1,
+        padding: bool = True,
+        padding_mode: str = "replicate",
+        kernel_size: int = 12,
+    ) -> None:
+        super().__init__()
+        self.kernel_size = kernel_size
+        self.even = kernel_size % 2 == 0
+        self.pad_left = kernel_size // 2 - int(self.even)
+        self.pad_right = kernel_size // 2
+        self.stride = stride
+        self.padding = padding
+        self.padding_mode = padding_mode
+        self.register_buffer("filter", kaiser_sinc_filter1d(cutoff, half_width, kernel_size))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        _, channels, _ = x.shape
+        if self.padding:
+            x = F.pad(x, (self.pad_left, self.pad_right), mode=self.padding_mode)
+        return F.conv1d(x, self.filter.expand(channels, -1, -1), stride=self.stride, groups=channels)
+
+
+class UpSample1d(nn.Module):
+    def __init__(self, ratio: int = 2, kernel_size: int | None = None) -> None:
+        super().__init__()
+        self.ratio = ratio
+        self.kernel_size = int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size
+        self.stride = ratio
+        self.pad = self.kernel_size // ratio - 1
+        self.pad_left = self.pad * self.stride + (self.kernel_size - self.stride) // 2
+        self.pad_right = self.pad * self.stride + (self.kernel_size - self.stride + 1) // 2
+        self.register_buffer(
+            "filter", kaiser_sinc_filter1d(cutoff=0.5 / ratio, half_width=0.6 / ratio, kernel_size=self.kernel_size)
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        _, channels, _ = x.shape
+        x = F.pad(x, (self.pad, self.pad), mode="replicate")
+        x = self.ratio * F.conv_transpose1d(
+            x, self.filter.expand(channels, -1, -1), stride=self.stride, groups=channels
+        )
+        return x[..., self.pad_left : -self.pad_right]
+
+
+class DownSample1d(nn.Module):
+    def __init__(self, ratio: int = 2, kernel_size: int | None = None) -> None:
+        super().__init__()
+        self.ratio = ratio
+        self.kernel_size = int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size
+        self.lowpass = LowPassFilter1d(
+            cutoff=0.5 / ratio, half_width=0.6 / ratio, stride=ratio, kernel_size=self.kernel_size
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.lowpass(x)
+
+
+class Activation1d(nn.Module):
+    def __init__(
+        self,
+        activation: nn.Module,
+        up_ratio: int = 2,
+        down_ratio: int = 2,
+        up_kernel_size: int = 12,
+        down_kernel_size: int = 12,
+    ) -> None:
+        super().__init__()
+        self.up_ratio = up_ratio
+        self.down_ratio = down_ratio
+        self.act = activation
+        self.upsample = UpSample1d(up_ratio, up_kernel_size)
+        self.downsample = DownSample1d(down_ratio, down_kernel_size)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.downsample(self.act(self.upsample(x)))
+
+
+class Snake(nn.Module):
+    def __init__(
+        self, in_features: int, alpha: float = 1.0, alpha_trainable: bool = True, alpha_logscale: bool = False
+    ) -> None:
+        super().__init__()
+        self.in_features = in_features
+        self.alpha_logscale = alpha_logscale
+        initial = torch.zeros(in_features) * alpha if alpha_logscale else torch.ones(in_features) * alpha
+        self.alpha = nn.Parameter(initial, requires_grad=alpha_trainable)
+        self.no_div_by_zero = 1e-9
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        alpha = self.alpha.unsqueeze(0).unsqueeze(-1)
+        if self.alpha_logscale:
+            alpha = torch.exp(alpha)
+        return x + (1.0 / (alpha + self.no_div_by_zero)) * torch.pow(
+            torch.sin(x * alpha), 2)
+
+
+class SnakeBeta(nn.Module):
+    def __init__(
+        self, in_features: int, alpha: float = 1.0, alpha_trainable: bool = True, alpha_logscale: bool = False
+    ) -> None:
+        super().__init__()
+        self.in_features = in_features
+        self.alpha_logscale = alpha_logscale
+        initial = torch.zeros(in_features) * alpha if alpha_logscale else torch.ones(in_features) * alpha
+        self.alpha = nn.Parameter(initial.clone(), requires_grad=alpha_trainable)
+        self.beta = nn.Parameter(initial.clone(), requires_grad=alpha_trainable)
+        self.no_div_by_zero = 1e-9
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        alpha = self.alpha.unsqueeze(0).unsqueeze(-1)
+        beta = self.beta.unsqueeze(0).unsqueeze(-1)
+        if self.alpha_logscale:
+            alpha = torch.exp(alpha)
+            beta = torch.exp(beta)
+        return x + (1.0 / (beta + self.no_div_by_zero)) * torch.pow(
+            torch.sin(x * alpha), 2)
+
+
+def _activation(name: str, channels: int, logscale: bool) -> Activation1d:
+    if name == "snake":
+        activation = Snake(channels, alpha_logscale=logscale)
+    elif name == "snakebeta":
+        activation = SnakeBeta(channels, alpha_logscale=logscale)
+    else:
+        raise ValueError(f"Unsupported BigVGAN activation: {name}")
+    return Activation1d(activation)
+
+
+class AMPBlock1(nn.Module):
+    def __init__(
+        self,
+        config: AttrDict,
+        channels: int,
+        kernel_size: int = 3,
+        dilation: tuple[int, ...] = (1, 3, 5),
+        activation: str = "snake",
+    ) -> None:
+        super().__init__()
+        self.convs1 = nn.ModuleList(
+            [
+                weight_norm(
+                    Conv1d(
+                        channels, channels, kernel_size, stride=1, dilation=rate, padding=get_padding(kernel_size, rate)
+                    )
+                )
+                for rate in dilation
+            ]
+        )
+        self.convs1.apply(init_weights)
+        self.convs2 = nn.ModuleList(
+            [
+                weight_norm(
+                    Conv1d(channels, channels, kernel_size, stride=1, dilation=1, padding=get_padding(kernel_size, 1))
+                )
+                for _ in dilation
+            ]
+        )
+        self.convs2.apply(init_weights)
+        self.activations = nn.ModuleList(
+            [
+                _activation(activation, channels, config.snake_logscale)
+                for _ in range(len(self.convs1) + len(self.convs2))
+            ]
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        acts1, acts2 = self.activations[::2], self.activations[1::2]
+        for conv1, conv2, act1, act2 in zip(self.convs1, self.convs2, acts1, acts2, strict=True):
+            residual = conv2(act2(conv1(act1(x))))
+            x = residual + x
+        return x
+
+    def remove_weight_norm(self) -> None:
+        for layer in self.convs1:
+            remove_parametrizations(layer, "weight")
+        for layer in self.convs2:
+            remove_parametrizations(layer, "weight")
+
+
+class AMPBlock2(nn.Module):
+    def __init__(
+        self,
+        config: AttrDict,
+        channels: int,
+        kernel_size: int = 3,
+        dilation: tuple[int, ...] = (1, 3, 5),
+        activation: str = "snake",
+    ) -> None:
+        super().__init__()
+        self.convs = nn.ModuleList(
+            [
+                weight_norm(
+                    Conv1d(
+                        channels, channels, kernel_size, stride=1, dilation=rate, padding=get_padding(kernel_size, rate)
+                    )
+                )
+                for rate in dilation
+            ]
+        )
+        self.convs.apply(init_weights)
+        self.activations = nn.ModuleList([_activation(activation, channels, config.snake_logscale) for _ in self.convs])
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        for conv, activation in zip(self.convs, self.activations, strict=True):
+            x = conv(activation(x)) + x
+        return x
+
+    def remove_weight_norm(self) -> None:
+        for layer in self.convs:
+            remove_parametrizations(layer, "weight")
+
+
+class BigVGANV2(nn.Module):
+    """BigVGAN-v2 generator compatible with NVIDIA checkpoint keys."""
+
+    def __init__(self, config: dict[str, Any]) -> None:
+        super().__init__()
+        config = dict(config)
+        config.pop("_class_name", None)
+        config.pop("architectures", None)
+        weight_norm_removed = bool(config.pop("weight_norm_removed", False))
+        self.config = AttrDict(config)
+        if self.config.get("use_cuda_kernel", False):
+            raise ValueError("FastVideo BigVGANV2 supports only the portable PyTorch path")
+        self.config["use_cuda_kernel"] = False
+        self.num_kernels = len(self.config.resblock_kernel_sizes)
+        self.num_upsamples = len(self.config.upsample_rates)
+        self.conv_pre = weight_norm(Conv1d(self.config.num_mels, self.config.upsample_initial_channel, 7, 1, padding=3))
+        if self.config.resblock == "1":
+            block_class = AMPBlock1
+        elif self.config.resblock == "2":
+            block_class = AMPBlock2
+        else:
+            raise ValueError(f"Unsupported BigVGAN resblock: {self.config.resblock}")
+
+        self.ups = nn.ModuleList()
+        for index, (rate, kernel) in enumerate(
+            zip(self.config.upsample_rates, self.config.upsample_kernel_sizes, strict=True)
+        ):
+            self.ups.append(
+                nn.ModuleList(
+                    [
+                        weight_norm(
+                            ConvTranspose1d(
+                                self.config.upsample_initial_channel // (2**index),
+                                self.config.upsample_initial_channel // (2 ** (index + 1)),
+                                kernel,
+                                rate,
+                                padding=(kernel - rate) // 2,
+                            )
+                        )
+                    ]
+                )
+            )
+
+        self.resblocks = nn.ModuleList()
+        for index in range(len(self.ups)):
+            channels = self.config.upsample_initial_channel // (2 ** (index + 1))
+            for kernel, dilation in zip(
+                self.config.resblock_kernel_sizes, self.config.resblock_dilation_sizes, strict=True
+            ):
+                self.resblocks.append(
+                    block_class(self.config, channels, kernel, tuple(dilation), activation=self.config.activation)
+                )
+
+        channels = self.config.upsample_initial_channel // (2 ** len(self.ups))
+        self.activation_post = _activation(self.config.activation, channels, self.config.snake_logscale)
+        self.use_bias_at_final = self.config.get("use_bias_at_final", True)
+        self.conv_post = weight_norm(Conv1d(channels, 1, 7, 1, padding=3, bias=self.use_bias_at_final))
+        for upsampler in self.ups:
+            upsampler.apply(init_weights)
+        self.conv_post.apply(init_weights)
+        self.use_tanh_at_final = self.config.get("use_tanh_at_final", True)
+        if weight_norm_removed:
+            self.remove_weight_norm()
+
+    def forward(self, mel: torch.Tensor) -> torch.Tensor:
+        hidden = self.conv_pre(mel)
+        for index in range(self.num_upsamples):
+            for upsampler in self.ups[index]:
+                hidden = upsampler(hidden)
+            accumulated = None
+            for kernel in range(self.num_kernels):
+                block_output = self.resblocks[
+                    index * self.num_kernels + kernel](hidden)
+                if accumulated is None:
+                    accumulated = block_output
+                else:
+                    # Preserve BigVGAN's published sequential accumulation
+                    # order. Tree reduction drifts through later nonlinear
+                    # upsampling stages with the full checkpoint.
+                    accumulated += block_output
+            assert accumulated is not None
+            hidden = accumulated / self.num_kernels
+        hidden = self.conv_post(self.activation_post(hidden))
+        if self.use_tanh_at_final:
+            return torch.tanh(hidden)
+        return torch.clamp(hidden, min=-1.0, max=1.0)
+
+    def remove_weight_norm(self) -> None:
+        try:
+            for upsamplers in self.ups:
+                for upsampler in upsamplers:
+                    remove_parametrizations(upsampler, "weight")
+            for block in self.resblocks:
+                block.remove_weight_norm()
+            remove_parametrizations(self.conv_pre, "weight")
+            remove_parametrizations(self.conv_post, "weight")
+        except ValueError:
+            # Idempotent for pipeline setup and converted checkpoints.
+            return
+
+
+EntryClass = BigVGANV2

fastvideo/models/audio/mmaudio_vae.py

@@ -0,0 +1,827 @@
+# SPDX-License-Identifier: MIT
+#
+# MIT License
+#
+# Copyright (c) 2024 Sony Research Inc.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+"""Native 1D audio VAE used by MMAudio."""
+
+from __future__ import annotations
+
+import logging
+import math
+from collections.abc import Iterable
+from typing import Any
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+
+from fastvideo.models.loader.weight_utils import default_weight_loader
+
+logger = logging.getLogger(__name__)
+
+
+DATA_MEAN_80D = [
+    -1.6058,
+    -1.3676,
+    -1.2520,
+    -1.2453,
+    -1.2078,
+    -1.2224,
+    -1.2419,
+    -1.2439,
+    -1.2922,
+    -1.2927,
+    -1.3170,
+    -1.3543,
+    -1.3401,
+    -1.3836,
+    -1.3907,
+    -1.3912,
+    -1.4313,
+    -1.4152,
+    -1.4527,
+    -1.4728,
+    -1.4568,
+    -1.5101,
+    -1.5051,
+    -1.5172,
+    -1.5623,
+    -1.5373,
+    -1.5746,
+    -1.5687,
+    -1.6032,
+    -1.6131,
+    -1.6081,
+    -1.6331,
+    -1.6489,
+    -1.6489,
+    -1.6700,
+    -1.6738,
+    -1.6953,
+    -1.6969,
+    -1.7048,
+    -1.7280,
+    -1.7361,
+    -1.7495,
+    -1.7658,
+    -1.7814,
+    -1.7889,
+    -1.8064,
+    -1.8221,
+    -1.8377,
+    -1.8417,
+    -1.8643,
+    -1.8857,
+    -1.8929,
+    -1.9173,
+    -1.9379,
+    -1.9531,
+    -1.9673,
+    -1.9824,
+    -2.0042,
+    -2.0215,
+    -2.0436,
+    -2.0766,
+    -2.1064,
+    -2.1418,
+    -2.1855,
+    -2.2319,
+    -2.2767,
+    -2.3161,
+    -2.3572,
+    -2.3954,
+    -2.4282,
+    -2.4659,
+    -2.5072,
+    -2.5552,
+    -2.6074,
+    -2.6584,
+    -2.7107,
+    -2.7634,
+    -2.8266,
+    -2.8981,
+    -2.9673,
+]
+
+DATA_STD_80D = [
+    1.0291,
+    1.0411,
+    1.0043,
+    0.9820,
+    0.9677,
+    0.9543,
+    0.9450,
+    0.9392,
+    0.9343,
+    0.9297,
+    0.9276,
+    0.9263,
+    0.9242,
+    0.9254,
+    0.9232,
+    0.9281,
+    0.9263,
+    0.9315,
+    0.9274,
+    0.9247,
+    0.9277,
+    0.9199,
+    0.9188,
+    0.9194,
+    0.9160,
+    0.9161,
+    0.9146,
+    0.9161,
+    0.9100,
+    0.9095,
+    0.9145,
+    0.9076,
+    0.9066,
+    0.9095,
+    0.9032,
+    0.9043,
+    0.9038,
+    0.9011,
+    0.9019,
+    0.9010,
+    0.8984,
+    0.8983,
+    0.8986,
+    0.8961,
+    0.8962,
+    0.8978,
+    0.8962,
+    0.8973,
+    0.8993,
+    0.8976,
+    0.8995,
+    0.9016,
+    0.8982,
+    0.8972,
+    0.8974,
+    0.8949,
+    0.8940,
+    0.8947,
+    0.8936,
+    0.8939,
+    0.8951,
+    0.8956,
+    0.9017,
+    0.9167,
+    0.9436,
+    0.9690,
+    1.0003,
+    1.0225,
+    1.0381,
+    1.0491,
+    1.0545,
+    1.0604,
+    1.0761,
+    1.0929,
+    1.1089,
+    1.1196,
+    1.1176,
+    1.1156,
+    1.1117,
+    1.1070,
+]
+
+DATA_MEAN_128D = [
+    -3.3462,
+    -2.6723,
+    -2.4893,
+    -2.3143,
+    -2.2664,
+    -2.3317,
+    -2.1802,
+    -2.4006,
+    -2.2357,
+    -2.4597,
+    -2.3717,
+    -2.4690,
+    -2.5142,
+    -2.4919,
+    -2.6610,
+    -2.5047,
+    -2.7483,
+    -2.5926,
+    -2.7462,
+    -2.7033,
+    -2.7386,
+    -2.8112,
+    -2.7502,
+    -2.9594,
+    -2.7473,
+    -3.0035,
+    -2.8891,
+    -2.9922,
+    -2.9856,
+    -3.0157,
+    -3.1191,
+    -2.9893,
+    -3.1718,
+    -3.0745,
+    -3.1879,
+    -3.2310,
+    -3.1424,
+    -3.2296,
+    -3.2791,
+    -3.2782,
+    -3.2756,
+    -3.3134,
+    -3.3509,
+    -3.3750,
+    -3.3951,
+    -3.3698,
+    -3.4505,
+    -3.4509,
+    -3.5089,
+    -3.4647,
+    -3.5536,
+    -3.5788,
+    -3.5867,
+    -3.6036,
+    -3.6400,
+    -3.6747,
+    -3.7072,
+    -3.7279,
+    -3.7283,
+    -3.7795,
+    -3.8259,
+    -3.8447,
+    -3.8663,
+    -3.9182,
+    -3.9605,
+    -3.9861,
+    -4.0105,
+    -4.0373,
+    -4.0762,
+    -4.1121,
+    -4.1488,
+    -4.1874,
+    -4.2461,
+    -4.3170,
+    -4.3639,
+    -4.4452,
+    -4.5282,
+    -4.6297,
+    -4.7019,
+    -4.7960,
+    -4.8700,
+    -4.9507,
+    -5.0303,
+    -5.0866,
+    -5.1634,
+    -5.2342,
+    -5.3242,
+    -5.4053,
+    -5.4927,
+    -5.5712,
+    -5.6464,
+    -5.7052,
+    -5.7619,
+    -5.8410,
+    -5.9188,
+    -6.0103,
+    -6.0955,
+    -6.1673,
+    -6.2362,
+    -6.3120,
+    -6.3926,
+    -6.4797,
+    -6.5565,
+    -6.6511,
+    -6.8130,
+    -6.9961,
+    -7.1275,
+    -7.2457,
+    -7.3576,
+    -7.4663,
+    -7.6136,
+    -7.7469,
+    -7.8815,
+    -8.0132,
+    -8.1515,
+    -8.3071,
+    -8.4722,
+    -8.7418,
+    -9.3975,
+    -9.6628,
+    -9.7671,
+    -9.8863,
+    -9.9992,
+    -10.0860,
+    -10.1709,
+    -10.5418,
+    -11.2795,
+    -11.3861,
+]
+
+DATA_STD_128D = [
+    2.3804,
+    2.4368,
+    2.3772,
+    2.3145,
+    2.2803,
+    2.2510,
+    2.2316,
+    2.2083,
+    2.1996,
+    2.1835,
+    2.1769,
+    2.1659,
+    2.1631,
+    2.1618,
+    2.1540,
+    2.1606,
+    2.1571,
+    2.1567,
+    2.1612,
+    2.1579,
+    2.1679,
+    2.1683,
+    2.1634,
+    2.1557,
+    2.1668,
+    2.1518,
+    2.1415,
+    2.1449,
+    2.1406,
+    2.1350,
+    2.1313,
+    2.1415,
+    2.1281,
+    2.1352,
+    2.1219,
+    2.1182,
+    2.1327,
+    2.1195,
+    2.1137,
+    2.1080,
+    2.1179,
+    2.1036,
+    2.1087,
+    2.1036,
+    2.1015,
+    2.1068,
+    2.0975,
+    2.0991,
+    2.0902,
+    2.1015,
+    2.0857,
+    2.0920,
+    2.0893,
+    2.0897,
+    2.0910,
+    2.0881,
+    2.0925,
+    2.0873,
+    2.0960,
+    2.0900,
+    2.0957,
+    2.0958,
+    2.0978,
+    2.0936,
+    2.0886,
+    2.0905,
+    2.0845,
+    2.0855,
+    2.0796,
+    2.0840,
+    2.0813,
+    2.0817,
+    2.0838,
+    2.0840,
+    2.0917,
+    2.1061,
+    2.1431,
+    2.1976,
+    2.2482,
+    2.3055,
+    2.3700,
+    2.4088,
+    2.4372,
+    2.4609,
+    2.4731,
+    2.4847,
+    2.5072,
+    2.5451,
+    2.5772,
+    2.6147,
+    2.6529,
+    2.6596,
+    2.6645,
+    2.6726,
+    2.6803,
+    2.6812,
+    2.6899,
+    2.6916,
+    2.6931,
+    2.6998,
+    2.7062,
+    2.7262,
+    2.7222,
+    2.7158,
+    2.7041,
+    2.7485,
+    2.7491,
+    2.7451,
+    2.7485,
+    2.7233,
+    2.7297,
+    2.7233,
+    2.7145,
+    2.6958,
+    2.6788,
+    2.6439,
+    2.6007,
+    2.4786,
+    2.2469,
+    2.1877,
+    2.1392,
+    2.0717,
+    2.0107,
+    1.9676,
+    1.9140,
+    1.7102,
+    0.9101,
+    0.7164,
+]
+
+
+_NORMALIZATION_EPSILON = 1e-4
+_SILU_DIVISOR = 0.596
+_RESIDUAL_BLEND = 0.3
+_RESIDUAL_DIVISOR = math.hypot(1.0 - _RESIDUAL_BLEND, _RESIDUAL_BLEND)
+
+
+def _rms_normalize(x: torch.Tensor, dim: int | tuple[int, ...]) -> torch.Tensor:
+    dims = (dim,) if isinstance(dim, int) else dim
+    element_count = math.prod(x.shape[axis] for axis in dims)
+    l2_norm = torch.linalg.vector_norm(x, dim=dims, keepdim=True, dtype=torch.float32)
+    rms = torch.add(_NORMALIZATION_EPSILON, l2_norm, alpha=element_count**-0.5)
+    return x / rms.to(x.dtype)
+
+
+def _conv1d(in_channels: int, out_channels: int, kernel_size: int) -> nn.Conv1d:
+    return nn.Conv1d(in_channels, out_channels, kernel_size, padding=kernel_size // 2, bias=False)
+
+
+def _conv1d_with_gain(conv: nn.Conv1d, x: torch.Tensor, gain: torch.Tensor | float) -> torch.Tensor:
+    return F.conv1d(x, conv.weight * gain, stride=conv.stride, padding=conv.padding, dilation=conv.dilation,
+                    groups=conv.groups)
+
+
+class DiagonalGaussianDistribution:
+    def __init__(self, parameters: torch.Tensor, deterministic: bool = False) -> None:
+        self.parameters = parameters
+        self.mean, self.logvar = torch.chunk(parameters, 2, dim=1)
+        self.logvar = torch.clamp(self.logvar, -30.0, 20.0)
+        self.deterministic = deterministic
+        self.std = torch.exp(0.5 * self.logvar)
+        self.var = torch.exp(self.logvar)
+        if deterministic:
+            self.var = self.std = torch.zeros_like(self.mean)
+
+    def sample(self, generator: torch.Generator | None = None) -> torch.Tensor:
+        noise = torch.empty_like(self.mean).normal_(generator=generator)
+        return self.mean + self.std * noise
+
+    def mode(self) -> torch.Tensor:
+        return self.mean
+
+
+class ResnetBlock1D(nn.Module):
+    def __init__(
+        self,
+        *,
+        in_dim: int,
+        out_dim: int | None = None,
+        conv_shortcut: bool = False,
+        kernel_size: int = 3,
+        use_norm: bool = True,
+    ) -> None:
+        super().__init__()
+        self.in_dim = in_dim
+        self.out_dim = in_dim if out_dim is None else out_dim
+        self.use_conv_shortcut = conv_shortcut
+        self.use_norm = use_norm
+        self.conv1 = _conv1d(in_dim, self.out_dim, kernel_size)
+        self.conv2 = _conv1d(self.out_dim, self.out_dim, kernel_size)
+        if self.in_dim != self.out_dim:
+            if conv_shortcut:
+                self.conv_shortcut = _conv1d(in_dim, self.out_dim, kernel_size)
+            else:
+                self.nin_shortcut = _conv1d(in_dim, self.out_dim, 1)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.use_norm:
+            x = _rms_normalize(x, dim=1)
+        hidden = self.conv1(F.silu(x) / _SILU_DIVISOR)
+        hidden = self.conv2(F.silu(hidden) / _SILU_DIVISOR)
+        if self.in_dim != self.out_dim:
+            shortcut = self.conv_shortcut if self.use_conv_shortcut else self.nin_shortcut
+            x = shortcut(x)
+        return torch.lerp(x, hidden, _RESIDUAL_BLEND) / _RESIDUAL_DIVISOR
+
+
+class AttnBlock1D(nn.Module):
+    def __init__(self, in_channels: int, num_heads: int = 1) -> None:
+        super().__init__()
+        self.in_channels = in_channels
+        self.num_heads = num_heads
+        self.qkv = _conv1d(in_channels, in_channels * 3, 1)
+        self.proj_out = _conv1d(in_channels, in_channels, 1)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        qkv = self.qkv(x).reshape(x.shape[0], self.num_heads, -1, 3, x.shape[-1])
+        query, key, value = _rms_normalize(qkv, dim=2).unbind(3)
+        query = rearrange(query, "b h c l -> b h l c")
+        key = rearrange(key, "b h c l -> b h l c")
+        value = rearrange(value, "b h c l -> b h l c")
+        hidden = F.scaled_dot_product_attention(query, key, value)
+        hidden = rearrange(hidden, "b h l c -> b (h c) l")
+        return torch.lerp(x, self.proj_out(hidden), _RESIDUAL_BLEND) / _RESIDUAL_DIVISOR
+
+
+class Upsample1D(nn.Module):
+    def __init__(self, in_channels: int, with_conv: bool) -> None:
+        super().__init__()
+        self.with_conv = with_conv
+        if with_conv:
+            self.conv = _conv1d(in_channels, in_channels, 3)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = F.interpolate(x, scale_factor=2.0, mode="nearest-exact")
+        return self.conv(x) if self.with_conv else x
+
+
+class Downsample1D(nn.Module):
+    def __init__(self, in_channels: int, with_conv: bool) -> None:
+        super().__init__()
+        self.with_conv = with_conv
+        if with_conv:
+            self.conv1 = _conv1d(in_channels, in_channels, 1)
+            self.conv2 = _conv1d(in_channels, in_channels, 1)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.with_conv:
+            x = self.conv1(x)
+        x = F.avg_pool1d(x, kernel_size=2, stride=2)
+        return self.conv2(x) if self.with_conv else x
+
+
+class Encoder1D(nn.Module):
+    def __init__(
+        self,
+        *,
+        dim: int,
+        ch_mult: tuple[int, ...],
+        num_res_blocks: int,
+        attn_layers: list[int],
+        down_layers: list[int],
+        in_dim: int,
+        embed_dim: int,
+        resamp_with_conv: bool = True,
+        double_z: bool = True,
+        kernel_size: int = 3,
+        clip_act: float = 256.0,
+    ) -> None:
+        super().__init__()
+        self.dim = dim
+        self.num_layers = len(ch_mult)
+        self.num_res_blocks = num_res_blocks
+        self.in_channels = in_dim
+        self.clip_act = clip_act
+        self.down_layers = down_layers
+        self.attn_layers = attn_layers
+        self.conv_in = _conv1d(in_dim, dim, kernel_size)
+
+        in_ch_mult = (1,) + ch_mult
+        self.in_ch_mult = in_ch_mult
+        self.down = nn.ModuleList()
+        for level in range(self.num_layers):
+            block = nn.ModuleList()
+            attn = nn.ModuleList()
+            block_in = dim * in_ch_mult[level]
+            block_out = dim * ch_mult[level]
+            for _ in range(num_res_blocks):
+                block.append(ResnetBlock1D(in_dim=block_in, out_dim=block_out, kernel_size=kernel_size, use_norm=True))
+                block_in = block_out
+                if level in attn_layers:
+                    attn.append(AttnBlock1D(block_in))
+            down = nn.Module()
+            down.block = block
+            down.attn = attn
+            if level in down_layers:
+                down.downsample = Downsample1D(block_in, resamp_with_conv)
+            self.down.append(down)
+
+        self.mid = nn.Module()
+        self.mid.block_1 = ResnetBlock1D(in_dim=block_in, out_dim=block_in, kernel_size=kernel_size, use_norm=True)
+        self.mid.attn_1 = AttnBlock1D(block_in)
+        self.mid.block_2 = ResnetBlock1D(in_dim=block_in, out_dim=block_in, kernel_size=kernel_size, use_norm=True)
+        output_dim = 2 * embed_dim if double_z else embed_dim
+        self.conv_out = _conv1d(block_in, output_dim, kernel_size)
+        self.learnable_gain = nn.Parameter(torch.zeros([]))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        states = [self.conv_in(x)]
+        for level in range(self.num_layers):
+            for block_index in range(self.num_res_blocks):
+                hidden = self.down[level].block[block_index](states[-1])
+                if len(self.down[level].attn) > 0:
+                    hidden = self.down[level].attn[block_index](hidden)
+                states.append(hidden.clamp(-self.clip_act, self.clip_act))
+            if level in self.down_layers:
+                states.append(self.down[level].downsample(states[-1]))
+        hidden = self.mid.block_1(states[-1])
+        hidden = self.mid.attn_1(hidden)
+        hidden = self.mid.block_2(hidden).clamp(-self.clip_act, self.clip_act)
+        return _conv1d_with_gain(self.conv_out, F.silu(hidden) / _SILU_DIVISOR, self.learnable_gain + 1)
+
+
+class Decoder1D(nn.Module):
+    def __init__(
+        self,
+        *,
+        dim: int,
+        out_dim: int,
+        ch_mult: tuple[int, ...],
+        num_res_blocks: int,
+        attn_layers: list[int],
+        down_layers: list[int],
+        in_dim: int,
+        embed_dim: int,
+        kernel_size: int = 3,
+        resamp_with_conv: bool = True,
+        clip_act: float = 256.0,
+    ) -> None:
+        super().__init__()
+        self.ch = dim
+        self.num_layers = len(ch_mult)
+        self.num_res_blocks = num_res_blocks
+        self.in_channels = in_dim
+        self.clip_act = clip_act
+        self.down_layers = [level + 1 for level in down_layers]
+        block_in = dim * ch_mult[-1]
+        self.conv_in = _conv1d(embed_dim, block_in, kernel_size)
+        self.mid = nn.Module()
+        self.mid.block_1 = ResnetBlock1D(in_dim=block_in, out_dim=block_in, use_norm=True)
+        self.mid.attn_1 = AttnBlock1D(block_in)
+        self.mid.block_2 = ResnetBlock1D(in_dim=block_in, out_dim=block_in, use_norm=True)
+
+        self.up = nn.ModuleList()
+        for level in reversed(range(self.num_layers)):
+            block = nn.ModuleList()
+            attn = nn.ModuleList()
+            block_out = dim * ch_mult[level]
+            for _ in range(num_res_blocks + 1):
+                block.append(ResnetBlock1D(in_dim=block_in, out_dim=block_out, use_norm=True))
+                block_in = block_out
+                if level in attn_layers:
+                    attn.append(AttnBlock1D(block_in))
+            up = nn.Module()
+            up.block = block
+            up.attn = attn
+            if level in self.down_layers:
+                up.upsample = Upsample1D(block_in, resamp_with_conv)
+            self.up.insert(0, up)
+
+        self.conv_out = _conv1d(block_in, out_dim, kernel_size)
+        self.learnable_gain = nn.Parameter(torch.zeros([]))
+
+    def forward(self, z: torch.Tensor) -> torch.Tensor:
+        hidden = self.conv_in(z)
+        hidden = self.mid.block_1(hidden)
+        hidden = self.mid.attn_1(hidden)
+        hidden = self.mid.block_2(hidden).clamp(-self.clip_act, self.clip_act)
+        for level in reversed(range(self.num_layers)):
+            for block_index in range(self.num_res_blocks + 1):
+                hidden = self.up[level].block[block_index](hidden)
+                if len(self.up[level].attn) > 0:
+                    hidden = self.up[level].attn[block_index](hidden)
+                hidden = hidden.clamp(-self.clip_act, self.clip_act)
+            if level in self.down_layers:
+                hidden = self.up[level].upsample(hidden)
+        return _conv1d_with_gain(self.conv_out, F.silu(hidden) / _SILU_DIVISOR, self.learnable_gain + 1)
+
+
+class MMAudioVAE(nn.Module):
+    """MMAudio mel-spectrogram VAE for 16 kHz or 44.1 kHz audio."""
+
+    def __init__(self, mode: str | dict[str, Any] = "44k", need_encoder: bool = False) -> None:
+        super().__init__()
+        if isinstance(mode, dict):
+            config = mode
+            mode = config.get("mode", "44k")
+            need_encoder = config.get("need_encoder", need_encoder)
+        if mode == "16k":
+            data_dim, embed_dim, hidden_dim = 80, 20, 384
+            data_mean, data_std = DATA_MEAN_80D, DATA_STD_80D
+        elif mode == "44k":
+            data_dim, embed_dim, hidden_dim = 128, 40, 512
+            data_mean, data_std = DATA_MEAN_128D, DATA_STD_128D
+        else:
+            raise ValueError(f"Unknown MMAudio VAE mode: {mode}")
+
+        self.mode = mode
+        self.embed_dim = embed_dim
+        self._weights_normalized = False
+        self.register_buffer("data_mean", torch.tensor(data_mean, dtype=torch.float32).view(1, -1, 1))
+        self.register_buffer("data_std", torch.tensor(data_std, dtype=torch.float32).view(1, -1, 1))
+        if need_encoder:
+            self.encoder = Encoder1D(
+                dim=hidden_dim,
+                ch_mult=(1, 2, 4),
+                num_res_blocks=2,
+                attn_layers=[3],
+                down_layers=[0],
+                in_dim=data_dim,
+                embed_dim=embed_dim,
+            )
+        self.decoder = Decoder1D(
+            dim=hidden_dim,
+            ch_mult=(1, 2, 4),
+            num_res_blocks=2,
+            attn_layers=[3],
+            down_layers=[0],
+            in_dim=data_dim,
+            out_dim=data_dim,
+            embed_dim=embed_dim,
+        )
+
+    def encode(self, mel: torch.Tensor, normalize_input: bool = True) -> DiagonalGaussianDistribution:
+        self._require_normalized_weights()
+        if not hasattr(self, "encoder"):
+            raise RuntimeError("This MMAudio VAE was loaded decoder-only")
+        if normalize_input:
+            mel = self.normalize(mel)
+        return DiagonalGaussianDistribution(self.encoder(mel))
+
+    def decode(self, latent: torch.Tensor, unnormalize_output: bool = True) -> torch.Tensor:
+        self._require_normalized_weights()
+        mel = self.decoder(latent)
+        return self.unnormalize(mel) if unnormalize_output else mel
+
+    def forward(self, latent: torch.Tensor) -> torch.Tensor:
+        return self.decode(latent)
+
+    def normalize(self, mel: torch.Tensor) -> torch.Tensor:
+        return (mel - self.data_mean) / self.data_std
+
+    def unnormalize(self, mel: torch.Tensor) -> torch.Tensor:
+        return mel * self.data_std + self.data_mean
+
+    def _require_normalized_weights(self) -> None:
+        if not self._weights_normalized:
+            raise RuntimeError("call remove_weight_norm() before inference")
+
+    @torch.no_grad()
+    def remove_weight_norm(self):
+        for name, module in self.named_modules():
+            if isinstance(module, nn.Conv1d):
+                weight = _rms_normalize(module.weight.to(torch.float32), dim=(1, 2))
+                weight = weight / math.sqrt(weight[0].numel())
+                module.weight.copy_(weight.to(module.weight.dtype))
+                logger.debug("Removed weight norm from %s", name)
+        self._weights_normalized = True
+        return self
+
+    def load_weights(
+        self,
+        weights: Iterable[tuple[str, torch.Tensor]],
+    ) -> set[str]:
+        params = dict(self.named_parameters())
+        loaded: set[str] = set()
+        for name, tensor in weights:
+            if name not in params:
+                continue
+            parameter = params[name]
+            loader = getattr(parameter, "weight_loader", default_weight_loader)
+            loader(parameter, tensor)
+            loaded.add(name)
+        return loaded
+
+
+EntryClass = MMAudioVAE

fastvideo/models/dits/ltx2.py

@@ -333,6 +333,7 @@ class GELUApprox(nn.Module):
         self,
         in_features: int,
         out_features: int,
+        bias: bool = True,
         quant_config: QuantizationConfig | None = None,
         prefix: str = "",
     ):
@@ -340,6 +341,7 @@ class GELUApprox(nn.Module):
         self.proj = ReplicatedLinear(
             in_features,
             out_features,
+            bias=bias,
             quant_config=quant_config,
             prefix=f"{prefix}.fc_in",
         )
@@ -357,6 +359,7 @@ class FeedForward(nn.Module):
         dim: int,
         dim_out: int,
         mult: int = 4,
+        bias: bool = True,
         quant_config: QuantizationConfig | None = None,
         prefix: str = "",
     ) -> None:
@@ -365,12 +368,14 @@ class FeedForward(nn.Module):
         project_in = GELUApprox(
             dim,
             inner_dim,
+            bias=bias,
             quant_config=quant_config,
             prefix=f"{prefix}.ffn",
         )
         project_out = ReplicatedLinear(
             inner_dim,
             dim_out,
+            bias=bias,
             quant_config=quant_config,
             prefix=f"{prefix}.ffn.fc_out",
         )
@@ -1206,6 +1211,9 @@ class TransformerConfig:
     # LTX-2.3 gated extensions (default OFF == LTX-2.0 behavior).
     apply_gated_attention: bool = False
     cross_attention_adaln: bool = False
+    # FFN bias is per stream: a checkpoint may drop it on video while keeping
+    # it on audio. Default True preserves existing behavior.
+    ff_bias: bool = True
 
 
 class LTXDistributedAttention(DistributedAttention):
@@ -1846,6 +1854,7 @@ class BasicAVTransformerBlock(torch.nn.Module):
             self.ff = FeedForward(
                 video.dim,
                 dim_out=video.dim,
+                bias=video.ff_bias,
                 quant_config=quant_config,
                 prefix=f"{prefix}.blocks.{idx}",
             )
@@ -1884,6 +1893,7 @@ class BasicAVTransformerBlock(torch.nn.Module):
             self.audio_ff = FeedForward(
                 audio.dim,
                 dim_out=audio.dim,
+                bias=audio.ff_bias,
                 quant_config=quant_config,
                 prefix=f"{prefix}.blocks.{idx}.audio",
             )
@@ -2353,6 +2363,8 @@ class LTXModel(torch.nn.Module):
         cross_attention_adaln: bool = False,
         caption_proj_before_connector: bool = False,
         apply_gated_attention: bool = False,
+        ff_bias: bool = True,
+        audio_ff_bias: bool = True,
         stg_block_idx: int = 29,
         use_distributed_attention: bool = False,
         quant_config: QuantizationConfig | None = None,
@@ -2364,6 +2376,9 @@ class LTXModel(torch.nn.Module):
         self.cross_attention_adaln = cross_attention_adaln
         self.caption_proj_before_connector = caption_proj_before_connector
         self.apply_gated_attention = apply_gated_attention
+        # Per-stream FFN bias; default True preserves existing behavior.
+        self.ff_bias = ff_bias
+        self.audio_ff_bias = audio_ff_bias
         self.stg_block_idx = stg_block_idx
         self.use_middle_indices_grid = use_middle_indices_grid
         self.rope_type = rope_type
@@ -2584,6 +2599,7 @@ class LTXModel(torch.nn.Module):
             context_dim=cross_attention_dim,
             apply_gated_attention=self.apply_gated_attention,
             cross_attention_adaln=self.cross_attention_adaln,
+            ff_bias=self.ff_bias,
         ) if self.model_type.is_video_enabled() else None)
         audio_config = (TransformerConfig(
             dim=self.audio_inner_dim,
@@ -2592,6 +2608,7 @@ class LTXModel(torch.nn.Module):
             context_dim=audio_cross_attention_dim,
             apply_gated_attention=self.apply_gated_attention,
             cross_attention_adaln=self.cross_attention_adaln,
+            ff_bias=self.audio_ff_bias,
         ) if self.model_type.is_audio_enabled() else None)
         self.use_distributed_attention = use_distributed_attention
         self.transformer_blocks = torch.nn.ModuleList([
@@ -2778,6 +2795,8 @@ class LTX2Transformer3DModel(BaseDiT):
             cross_attention_adaln=arch.cross_attention_adaln,
             caption_proj_before_connector=arch.caption_proj_before_connector,
             apply_gated_attention=arch.apply_gated_attention,
+            ff_bias=arch.ff_bias,
+            audio_ff_bias=arch.audio_ff_bias,
             stg_block_idx=arch.stg_block_idx,
             use_distributed_attention=use_distributed_attention,
             quant_config=config.quant_config,

fastvideo/models/dits/mmaudio.py

@@ -0,0 +1,588 @@
+# SPDX-License-Identifier: Apache-2.0
+"""MMAudio multimodal flow-prediction transformer.
+
+The model operates on one-dimensional audio latent sequences and jointly
+attends to semantic video, synchronization-video, and text conditions. The
+initial implementation intentionally uses torch SDPA so its single-GPU numeric
+contract remains explicit; sequence/tensor parallel support is a later,
+separately verified optimization.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+import torch.nn.functional as F
+from einops import rearrange
+from einops.layers.torch import Rearrange
+from torch import nn
+
+from fastvideo.configs.models.dits.mmaudio import MMAudioTransformerConfig
+from fastvideo.models.dits.base import BaseDiT
+
+
+def compute_rope_rotations(
+    length: int, dim: int, theta: int, *, freq_scaling: float = 1.0, device: torch.device | str = "cpu"
+) -> torch.Tensor:
+    if dim % 2 != 0:
+        raise ValueError(f"RoPE dimension must be even, got {dim}.")
+    with torch.amp.autocast(device_type="cuda", enabled=False):
+        pos = torch.arange(length, dtype=torch.float32, device=device)
+        freqs = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float32, device=device) / dim))
+        freqs *= freq_scaling
+        rotations = torch.einsum("..., f -> ... f", pos, freqs)
+        rotations = torch.stack(
+            [
+                torch.cos(rotations),
+                -torch.sin(rotations),
+                torch.sin(rotations),
+                torch.cos(rotations),
+            ],
+            dim=-1,
+        )
+        return rearrange(rotations, "n d (i j) -> 1 n d i j", i=2, j=2)
+
+
+def apply_rope(x: torch.Tensor, rotations: torch.Tensor) -> torch.Tensor:
+    with torch.amp.autocast(device_type="cuda", enabled=False):
+        source = x.float().view(*x.shape[:-1], -1, 1, 2)
+        output = rotations[..., 0] * source[..., 0] + rotations[..., 1] * source[..., 1]
+        return output.reshape(*x.shape).to(dtype=x.dtype)
+
+
+class ChannelLastConv1d(nn.Conv1d):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return super().forward(x.permute(0, 2, 1)).permute(0, 2, 1)
+
+
+class MMAudioMLP(nn.Module):
+    def __init__(self, dim: int, hidden_dim: int, multiple_of: int = 256) -> None:
+        super().__init__()
+        hidden_dim = int(2 * hidden_dim / 3)
+        hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of)
+        self.w1 = nn.Linear(dim, hidden_dim, bias=False)
+        self.w2 = nn.Linear(hidden_dim, dim, bias=False)
+        self.w3 = nn.Linear(dim, hidden_dim, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.w2(F.silu(self.w1(x)) * self.w3(x))
+
+
+class MMAudioConvMLP(nn.Module):
+    def __init__(
+        self, dim: int, hidden_dim: int, multiple_of: int = 256, kernel_size: int = 3, padding: int = 1
+    ) -> None:
+        super().__init__()
+        hidden_dim = int(2 * hidden_dim / 3)
+        hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of)
+        self.w1 = ChannelLastConv1d(dim, hidden_dim, bias=False, kernel_size=kernel_size, padding=padding)
+        self.w2 = ChannelLastConv1d(hidden_dim, dim, bias=False, kernel_size=kernel_size, padding=padding)
+        self.w3 = ChannelLastConv1d(dim, hidden_dim, bias=False, kernel_size=kernel_size, padding=padding)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.w2(F.silu(self.w1(x)) * self.w3(x))
+
+
+class TimestepEmbedder(nn.Module):
+    def __init__(self, dim: int, frequency_embedding_size: int, max_period: int) -> None:
+        super().__init__()
+        if dim % 2 != 0:
+            raise ValueError(f"Timestep embedding dim must be even, got {dim}.")
+        self.mlp = nn.Sequential(
+            nn.Linear(frequency_embedding_size, dim),
+            nn.SiLU(),
+            nn.Linear(dim, dim),
+        )
+        self.dim = dim
+        self.max_period = max_period
+        with torch.autocast("cuda", enabled=False):
+            freqs = 1.0 / (
+                10000 ** (torch.arange(0, frequency_embedding_size, 2, dtype=torch.float32) / frequency_embedding_size)
+            )
+            self.register_buffer("freqs", (10000 / max_period) * freqs, persistent=False)
+
+    def timestep_embedding(self, timestep: torch.Tensor) -> torch.Tensor:
+        args = timestep[:, None].float() * self.freqs[None]
+        return torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+
+    def forward(self, timestep: torch.Tensor) -> torch.Tensor:
+        embedding = self.timestep_embedding(timestep).to(self.mlp[0].weight.dtype)
+        return self.mlp(embedding)
+
+
+def modulate(x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
+    return x * (1 + scale) + shift
+
+
+def mmaudio_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
+    output = F.scaled_dot_product_attention(q.contiguous(), k.contiguous(), v.contiguous())
+    return rearrange(output, "b h n d -> b n (h d)").contiguous()
+
+
+class SelfAttention(nn.Module):
+    def __init__(self, dim: int, num_heads: int) -> None:
+        super().__init__()
+        self.dim = dim
+        self.num_heads = num_heads
+        self.qkv = nn.Linear(dim, dim * 3, bias=True)
+        self.q_norm = nn.RMSNorm(dim // num_heads)
+        self.k_norm = nn.RMSNorm(dim // num_heads)
+        self.split_into_heads = Rearrange(
+            "b n (h d j) -> b h n d j",
+            h=num_heads,
+            d=dim // num_heads,
+            j=3,
+        )
+
+    def pre_attention(
+        self, x: torch.Tensor, rotations: torch.Tensor | None
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        q, k, v = self.split_into_heads(self.qkv(x)).chunk(3, dim=-1)
+        q = self.q_norm(q.squeeze(-1))
+        k = self.k_norm(k.squeeze(-1))
+        v = v.squeeze(-1)
+        if rotations is not None:
+            q = apply_rope(q, rotations)
+            k = apply_rope(k, rotations)
+        return q, k, v
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return mmaudio_attention(*self.pre_attention(x, None))
+
+
+class MMDitSingleBlock(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int,
+        mlp_ratio: float = 4.0,
+        pre_only: bool = False,
+        kernel_size: int = 7,
+        padding: int = 3,
+    ) -> None:
+        super().__init__()
+        self.norm1 = nn.LayerNorm(dim, elementwise_affine=False)
+        self.attn = SelfAttention(dim, num_heads)
+        self.pre_only = pre_only
+        if pre_only:
+            self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(dim, 2 * dim, bias=True))
+        else:
+            self.linear1 = (
+                nn.Linear(dim, dim)
+                if kernel_size == 1
+                else ChannelLastConv1d(dim, dim, kernel_size=kernel_size, padding=padding)
+            )
+            self.norm2 = nn.LayerNorm(dim, elementwise_affine=False)
+            self.ffn = (
+                MMAudioMLP(dim, int(dim * mlp_ratio))
+                if kernel_size == 1
+                else MMAudioConvMLP(
+                    dim,
+                    int(dim * mlp_ratio),
+                    kernel_size=kernel_size,
+                    padding=padding,
+                )
+            )
+            self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(dim, 6 * dim, bias=True))
+
+    def pre_attention(self, x: torch.Tensor, condition: torch.Tensor, rotations: torch.Tensor | None):
+        modulation = self.adaLN_modulation(condition)
+        if self.pre_only:
+            shift_msa, scale_msa = modulation.chunk(2, dim=-1)
+            gate_msa = shift_mlp = scale_mlp = gate_mlp = None
+        else:
+            (shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp) = modulation.chunk(6, dim=-1)
+        normalized = modulate(self.norm1(x), shift_msa, scale_msa)
+        qkv = self.attn.pre_attention(normalized, rotations)
+        return qkv, (gate_msa, shift_mlp, scale_mlp, gate_mlp)
+
+    def post_attention(self, x: torch.Tensor, attention_output: torch.Tensor, condition):
+        if self.pre_only:
+            return x
+        gate_msa, shift_mlp, scale_mlp, gate_mlp = condition
+        x = x + self.linear1(attention_output) * gate_msa
+        residual = modulate(self.norm2(x), shift_mlp, scale_mlp)
+        return x + self.ffn(residual) * gate_mlp
+
+    def forward(self, x: torch.Tensor, condition: torch.Tensor, rotations: torch.Tensor | None) -> torch.Tensor:
+        qkv, modulation = self.pre_attention(x, condition, rotations)
+        return self.post_attention(x, mmaudio_attention(*qkv), modulation)
+
+
+class JointBlock(nn.Module):
+    def __init__(self, dim: int, num_heads: int, mlp_ratio: float = 4.0, pre_only: bool = False) -> None:
+        super().__init__()
+        self.pre_only = pre_only
+        self.latent_block = MMDitSingleBlock(dim, num_heads, mlp_ratio, pre_only=False, kernel_size=3, padding=1)
+        self.clip_block = MMDitSingleBlock(dim, num_heads, mlp_ratio, pre_only=pre_only, kernel_size=3, padding=1)
+        self.text_block = MMDitSingleBlock(dim, num_heads, mlp_ratio, pre_only=pre_only, kernel_size=1)
+
+    def forward(
+        self,
+        latent: torch.Tensor,
+        clip_features: torch.Tensor,
+        text_features: torch.Tensor,
+        global_condition: torch.Tensor,
+        extended_condition: torch.Tensor,
+        latent_rotations: torch.Tensor,
+        clip_rotations: torch.Tensor,
+    ):
+        latent_qkv, latent_mod = self.latent_block.pre_attention(latent, extended_condition, latent_rotations)
+        clip_qkv, clip_mod = self.clip_block.pre_attention(clip_features, global_condition, clip_rotations)
+        text_qkv, text_mod = self.text_block.pre_attention(text_features, global_condition, None)
+
+        latent_len = latent.shape[1]
+        clip_len = clip_features.shape[1]
+        joint_qkv = [torch.cat([latent_qkv[i], clip_qkv[i], text_qkv[i]], dim=2) for i in range(3)]
+        attention_output = mmaudio_attention(*joint_qkv)
+        latent_output = attention_output[:, :latent_len]
+        clip_output = attention_output[:, latent_len : latent_len + clip_len]
+        text_output = attention_output[:, latent_len + clip_len :]
+
+        latent = self.latent_block.post_attention(latent, latent_output, latent_mod)
+        if not self.pre_only:
+            clip_features = self.clip_block.post_attention(clip_features, clip_output, clip_mod)
+            text_features = self.text_block.post_attention(text_features, text_output, text_mod)
+        return latent, clip_features, text_features
+
+
+class FinalBlock(nn.Module):
+    def __init__(self, dim: int, out_dim: int) -> None:
+        super().__init__()
+        self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(dim, 2 * dim, bias=True))
+        self.norm = nn.LayerNorm(dim, elementwise_affine=False)
+        self.conv = ChannelLastConv1d(dim, out_dim, kernel_size=7, padding=3)
+
+    def forward(self, latent: torch.Tensor, condition: torch.Tensor) -> torch.Tensor:
+        shift, scale = self.adaLN_modulation(condition).chunk(2, dim=-1)
+        return self.conv(modulate(self.norm(latent), shift, scale))
+
+
+@dataclass
+class PreprocessedConditions:
+    clip_f: torch.Tensor
+    sync_f: torch.Tensor
+    text_f: torch.Tensor
+    clip_f_c: torch.Tensor
+    text_f_c: torch.Tensor
+
+
+_DEFAULT_CONFIG = MMAudioTransformerConfig()
+
+
+class MMAudioTransformer(BaseDiT):
+    _fsdp_shard_conditions = _DEFAULT_CONFIG.arch_config._fsdp_shard_conditions
+    _compile_conditions = _DEFAULT_CONFIG.arch_config._compile_conditions
+    _supported_attention_backends = _DEFAULT_CONFIG.arch_config._supported_attention_backends
+    # The shared FSDP loader converts this mapping dictionary into its callable
+    # form. Keeping the class attribute as a dict matches every other native
+    # FastVideo DiT and also supports identity-mapped converted checkpoints.
+    param_names_mapping = _DEFAULT_CONFIG.arch_config.param_names_mapping
+    reverse_param_names_mapping = _DEFAULT_CONFIG.arch_config.reverse_param_names_mapping
+
+    def __init__(self, config: MMAudioTransformerConfig, hf_config: dict[str, Any], **kwargs) -> None:
+        del kwargs
+        super().__init__(config=config, hf_config=hf_config)
+        arch = config.arch_config
+        self.v2 = arch.v2
+        self.latent_dim = arch.latent_dim
+        self._latent_seq_len = arch.latent_seq_len
+        self._clip_seq_len = arch.clip_seq_len
+        self._sync_seq_len = arch.sync_seq_len
+        self._text_seq_len = arch.text_seq_len
+        self.hidden_dim = arch.hidden_dim
+        self.num_heads = arch.num_heads
+        self.hidden_size = arch.hidden_size
+        self.num_attention_heads = arch.num_attention_heads
+        self.num_channels_latents = arch.num_channels_latents
+
+        activation = nn.SiLU if arch.v2 else nn.SELU
+        self.audio_input_proj = nn.Sequential(
+            ChannelLastConv1d(arch.latent_dim, arch.hidden_dim, kernel_size=7, padding=3),
+            activation(),
+            MMAudioConvMLP(arch.hidden_dim, arch.hidden_dim * 4, kernel_size=7, padding=3),
+        )
+        clip_layers: list[nn.Module] = [nn.Linear(arch.clip_dim, arch.hidden_dim)]
+        if arch.v2:
+            clip_layers.append(nn.SiLU())
+        clip_layers.append(MMAudioConvMLP(arch.hidden_dim, arch.hidden_dim * 4, kernel_size=3, padding=1))
+        self.clip_input_proj = nn.Sequential(*clip_layers)
+        self.sync_input_proj = nn.Sequential(
+            ChannelLastConv1d(arch.sync_dim, arch.hidden_dim, kernel_size=7, padding=3),
+            activation(),
+            MMAudioConvMLP(arch.hidden_dim, arch.hidden_dim * 4, kernel_size=3, padding=1),
+        )
+        text_layers: list[nn.Module] = [nn.Linear(arch.text_dim, arch.hidden_dim)]
+        if arch.v2:
+            text_layers.append(nn.SiLU())
+        text_layers.append(MMAudioMLP(arch.hidden_dim, arch.hidden_dim * 4))
+        self.text_input_proj = nn.Sequential(*text_layers)
+
+        self.clip_cond_proj = nn.Linear(arch.hidden_dim, arch.hidden_dim)
+        self.text_cond_proj = nn.Linear(arch.hidden_dim, arch.hidden_dim)
+        self.global_cond_mlp = MMAudioMLP(arch.hidden_dim, arch.hidden_dim * 4)
+        self.sync_pos_emb = nn.Parameter(torch.zeros((1, 1, 8, arch.sync_dim)))
+        self.final_layer = FinalBlock(arch.hidden_dim, arch.latent_dim)
+        self.t_embed = TimestepEmbedder(
+            arch.hidden_dim,
+            frequency_embedding_size=(arch.hidden_dim if arch.v2 else 256),
+            max_period=(1 if arch.v2 else 10000),
+        )
+        self.joint_blocks = nn.ModuleList(
+            [
+                JointBlock(
+                    arch.hidden_dim,
+                    arch.num_heads,
+                    mlp_ratio=arch.mlp_ratio,
+                    pre_only=(index == arch.depth - arch.fused_depth - 1),
+                )
+                for index in range(arch.depth - arch.fused_depth)
+            ]
+        )
+        self.fused_blocks = nn.ModuleList(
+            [
+                MMDitSingleBlock(arch.hidden_dim, arch.num_heads, mlp_ratio=arch.mlp_ratio, kernel_size=3, padding=1)
+                for _ in range(arch.fused_depth)
+            ]
+        )
+
+        self.latent_mean = nn.Parameter(torch.full((1, 1, arch.latent_dim), float("nan")), requires_grad=False)
+        self.latent_std = nn.Parameter(torch.full((1, 1, arch.latent_dim), float("nan")), requires_grad=False)
+        self.empty_string_feat = nn.Parameter(torch.zeros((arch.text_seq_len, arch.text_dim)), requires_grad=False)
+        self.empty_clip_feat = nn.Parameter(torch.zeros(1, arch.clip_dim), requires_grad=True)
+        self.empty_sync_feat = nn.Parameter(torch.zeros(1, arch.sync_dim), requires_grad=True)
+
+        self.initialize_weights()
+        self.initialize_rotations()
+        self.__post_init__()
+
+    @property
+    def device(self) -> torch.device:
+        return self.latent_mean.device
+
+    @property
+    def latent_seq_len(self) -> int:
+        return self._latent_seq_len
+
+    @property
+    def clip_seq_len(self) -> int:
+        return self._clip_seq_len
+
+    @property
+    def sync_seq_len(self) -> int:
+        return self._sync_seq_len
+
+    def initialize_rotations(self) -> None:
+        head_dim = self.hidden_dim // self.num_heads
+        latent_rotations = compute_rope_rotations(self._latent_seq_len, head_dim, 10000, device=self.device)
+        clip_rotations = compute_rope_rotations(
+            self._clip_seq_len,
+            head_dim,
+            10000,
+            freq_scaling=self._latent_seq_len / self._clip_seq_len,
+            device=self.device,
+        )
+        self.register_buffer("latent_rot", latent_rotations, persistent=False)
+        self.register_buffer("clip_rot", clip_rotations, persistent=False)
+
+    def materialize_non_persistent_buffers(
+        self,
+        device: torch.device,
+        dtype: torch.dtype | None = None,
+    ) -> None:
+        """Rebuild derived buffers after meta-device production loading."""
+        if self.t_embed.freqs.is_meta:
+            frequency_dim = self.t_embed.mlp[0].in_features
+            freqs = 1.0 / (
+                10000
+                ** (
+                    torch.arange(0, frequency_dim, 2, dtype=torch.float32, device=device)
+                    / frequency_dim
+                )
+            )
+            freqs = (10000 / self.t_embed.max_period) * freqs
+            self.t_embed._buffers["freqs"] = freqs.to(dtype=dtype or torch.float32)
+        if self.latent_rot.is_meta or self.clip_rot.is_meta:
+            head_dim = self.hidden_dim // self.num_heads
+            self._buffers["latent_rot"] = compute_rope_rotations(
+                self._latent_seq_len, head_dim, 10000, device=device
+            )
+            self._buffers["clip_rot"] = compute_rope_rotations(
+                self._clip_seq_len,
+                head_dim,
+                10000,
+                freq_scaling=self._latent_seq_len / self._clip_seq_len,
+                device=device,
+            )
+
+    def update_seq_lengths(self, latent_seq_len: int, clip_seq_len: int, sync_seq_len: int) -> None:
+        self._latent_seq_len = latent_seq_len
+        self._clip_seq_len = clip_seq_len
+        self._sync_seq_len = sync_seq_len
+        self.initialize_rotations()
+
+    def initialize_weights(self) -> None:
+
+        def basic_init(module: nn.Module) -> None:
+            if isinstance(module, nn.Linear):
+                nn.init.xavier_uniform_(module.weight)
+                if module.bias is not None:
+                    nn.init.constant_(module.bias, 0)
+
+        self.apply(basic_init)
+        nn.init.normal_(self.t_embed.mlp[0].weight, std=0.02)
+        nn.init.normal_(self.t_embed.mlp[2].weight, std=0.02)
+        for block in self.joint_blocks:
+            for stream in (block.latent_block, block.clip_block, block.text_block):
+                nn.init.constant_(stream.adaLN_modulation[-1].weight, 0)
+                nn.init.constant_(stream.adaLN_modulation[-1].bias, 0)
+        for block in self.fused_blocks:
+            nn.init.constant_(block.adaLN_modulation[-1].weight, 0)
+            nn.init.constant_(block.adaLN_modulation[-1].bias, 0)
+        nn.init.constant_(self.final_layer.adaLN_modulation[-1].weight, 0)
+        nn.init.constant_(self.final_layer.adaLN_modulation[-1].bias, 0)
+        nn.init.constant_(self.final_layer.conv.weight, 0)
+        nn.init.constant_(self.final_layer.conv.bias, 0)
+        nn.init.constant_(self.sync_pos_emb, 0)
+        nn.init.constant_(self.empty_clip_feat, 0)
+        nn.init.constant_(self.empty_sync_feat, 0)
+
+    def normalize(self, latent: torch.Tensor) -> torch.Tensor:
+        return latent.sub_(self.latent_mean).div_(self.latent_std)
+
+    def unnormalize(self, latent: torch.Tensor) -> torch.Tensor:
+        return latent.mul_(self.latent_std).add_(self.latent_mean)
+
+    def preprocess_conditions(
+        self, clip_features: torch.Tensor, sync_features: torch.Tensor, text_features: torch.Tensor
+    ) -> PreprocessedConditions:
+        if clip_features.shape[1] != self._clip_seq_len:
+            raise ValueError(f"Expected {self._clip_seq_len} CLIP tokens, got {clip_features.shape}.")
+        if sync_features.shape[1] != self._sync_seq_len:
+            raise ValueError(f"Expected {self._sync_seq_len} sync tokens, got {sync_features.shape}.")
+        if text_features.shape[1] != self._text_seq_len:
+            raise ValueError(f"Expected {self._text_seq_len} text tokens, got {text_features.shape}.")
+        batch_size = clip_features.shape[0]
+        sync_features = sync_features.view(batch_size, self._sync_seq_len // 8, 8, -1) + self.sync_pos_emb
+        sync_features = sync_features.flatten(1, 2)
+        clip_features = self.clip_input_proj(clip_features)
+        sync_features = self.sync_input_proj(sync_features)
+        text_features = self.text_input_proj(text_features)
+        sync_features = F.interpolate(
+            sync_features.transpose(1, 2),
+            size=self._latent_seq_len,
+            mode="nearest-exact",
+        ).transpose(1, 2)
+        return PreprocessedConditions(
+            clip_f=clip_features,
+            sync_f=sync_features,
+            text_f=text_features,
+            clip_f_c=self.clip_cond_proj(clip_features.mean(dim=1)),
+            text_f_c=self.text_cond_proj(text_features.mean(dim=1)),
+        )
+
+    def predict_flow(
+        self, latent: torch.Tensor, timestep: torch.Tensor, conditions: PreprocessedConditions
+    ) -> torch.Tensor:
+        if latent.shape[1] != self._latent_seq_len:
+            raise ValueError(f"Expected latent length {self._latent_seq_len}, got {latent.shape}.")
+        clip_features = conditions.clip_f
+        text_features = conditions.text_f
+        latent = self.audio_input_proj(latent)
+        global_condition = self.global_cond_mlp(conditions.clip_f_c + conditions.text_f_c)
+        global_condition = self.t_embed(timestep).unsqueeze(1) + global_condition.unsqueeze(1)
+        extended_condition = global_condition + conditions.sync_f
+        for block in self.joint_blocks:
+            latent, clip_features, text_features = block(
+                latent,
+                clip_features,
+                text_features,
+                global_condition,
+                extended_condition,
+                self.latent_rot,
+                self.clip_rot,
+            )
+        for block in self.fused_blocks:
+            latent = block(latent, extended_condition, self.latent_rot)
+        # The released checkpoint was trained with global rather than sync-
+        # extended conditioning at the final layer; preserve that contract.
+        return self.final_layer(latent, global_condition)
+
+    def get_empty_string_sequence(self, batch_size: int) -> torch.Tensor:
+        return self.empty_string_feat.unsqueeze(0).expand(batch_size, -1, -1)
+
+    def get_empty_clip_sequence(self, batch_size: int) -> torch.Tensor:
+        return self.empty_clip_feat.unsqueeze(0).expand(batch_size, self._clip_seq_len, -1)
+
+    def get_empty_sync_sequence(self, batch_size: int) -> torch.Tensor:
+        return self.empty_sync_feat.unsqueeze(0).expand(batch_size, self._sync_seq_len, -1)
+
+    def get_empty_conditions(
+        self,
+        batch_size: int,
+        *,
+        negative_text_features: torch.Tensor | None = None,
+    ) -> PreprocessedConditions:
+        empty_text = negative_text_features if negative_text_features is not None else self.get_empty_string_sequence(1)
+        conditions = self.preprocess_conditions(
+            self.get_empty_clip_sequence(1),
+            self.get_empty_sync_sequence(1),
+            empty_text,
+        )
+        conditions.clip_f = conditions.clip_f.expand(batch_size, -1, -1)
+        conditions.sync_f = conditions.sync_f.expand(batch_size, -1, -1)
+        conditions.clip_f_c = conditions.clip_f_c.expand(batch_size, -1)
+        if negative_text_features is None:
+            conditions.text_f = conditions.text_f.expand(batch_size, -1, -1)
+            conditions.text_f_c = conditions.text_f_c.expand(batch_size, -1)
+        return conditions
+
+    def guided_flow(
+        self,
+        timestep: torch.Tensor,
+        latent: torch.Tensor,
+        conditions: PreprocessedConditions,
+        empty_conditions: PreprocessedConditions,
+        guidance_scale: float,
+    ) -> torch.Tensor:
+        timestep = timestep * torch.ones(len(latent), device=latent.device, dtype=latent.dtype)
+        if guidance_scale < 1.0:
+            return self.predict_flow(latent, timestep, conditions)
+        return guidance_scale * self.predict_flow(latent, timestep, conditions) + (
+            1 - guidance_scale
+        ) * self.predict_flow(latent, timestep, empty_conditions)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: PreprocessedConditions
+        | tuple[torch.Tensor, torch.Tensor, torch.Tensor]
+        | list[torch.Tensor]
+        | dict[str, torch.Tensor],
+        timestep: torch.LongTensor,
+        encoder_hidden_states_image: torch.Tensor | list[torch.Tensor] | None = None,
+        guidance=None,
+        **kwargs,
+    ) -> torch.Tensor:
+        del encoder_hidden_states_image, guidance, kwargs
+        if isinstance(encoder_hidden_states, PreprocessedConditions):
+            conditions = encoder_hidden_states
+        elif isinstance(encoder_hidden_states, dict):
+            conditions = self.preprocess_conditions(
+                encoder_hidden_states["clip_features"],
+                encoder_hidden_states["sync_features"],
+                encoder_hidden_states["text_features"],
+            )
+        elif isinstance(encoder_hidden_states, (tuple, list)) and len(encoder_hidden_states) == 3:
+            conditions = self.preprocess_conditions(*encoder_hidden_states)
+        else:
+            raise TypeError(
+                "MMAudio encoder_hidden_states must be PreprocessedConditions, "
+                "a (clip, sync, text) triple, or a named condition dict."
+            )
+        return self.predict_flow(hidden_states, timestep, conditions)
+
+
+EntryClass = MMAudioTransformer

fastvideo/models/encoders/gemma.py

@@ -7,7 +7,7 @@ from typing import Any, Iterable
 
 import torch
 from torch import nn
-from transformers import AutoTokenizer, Gemma3ForConditionalGeneration
+from transformers import AutoModelForImageTextToText, AutoTokenizer
 
 from fastvideo.configs.models.encoders import BaseEncoderOutput, TextEncoderConfig
 from fastvideo.models.encoders.base import TextEncoder
@@ -398,7 +398,7 @@ class LTX2GemmaTextEncoderModel(TextEncoder):
         self.gemma_model_path = arch.gemma_model_path
         self.gemma_dtype = arch.gemma_dtype
         self.padding_side = arch.padding_side
-        self._gemma_model: Gemma3ForConditionalGeneration | None = None
+        self._gemma_model: nn.Module | None = None
 
     def named_parameters(self, prefix: str = "", recurse: bool = True):
         for name, param in super().named_parameters(prefix=prefix, recurse=recurse):
@@ -436,17 +436,22 @@ class LTX2GemmaTextEncoderModel(TextEncoder):
             )
 
     @property
-    def gemma_model(self) -> Gemma3ForConditionalGeneration:
+    def gemma_model(self) -> nn.Module:
         if self._gemma_model is None:
             gemma_path = self.gemma_model_path
             if not gemma_path:
                 raise ValueError("gemma_model_path must be set (expected text_encoder/gemma).")
             dtype = getattr(torch, self.gemma_dtype, torch.bfloat16)
-            self._gemma_model = Gemma3ForConditionalGeneration.from_pretrained(
+            # Resolve the class from the root config's own ``model_type``
+            # instead of hardcoding one: a checkpoint may declare its own
+            # encoder pairing. The auto class still resolves to the previously
+            # hardcoded class for the roots shipped so far.
+            self._gemma_model = AutoModelForImageTextToText.from_pretrained(
                 gemma_path,
                 local_files_only=True,
                 torch_dtype=dtype,
             )
+            self._sync_arch_from_gemma_config(self._gemma_model.config)
             # Configure model-level attention implementation when using TORCH_SDPA.
             # Note: torch.backends.cuda.enable_*_sdp() settings should be configured
             # at application/pipeline initialization level, not here, to avoid
@@ -464,6 +469,33 @@ class LTX2GemmaTextEncoderModel(TextEncoder):
             self._gemma_model.eval()
         return self._gemma_model
 
+    def _sync_arch_from_gemma_config(self, gemma_config: Any) -> None:
+        """Take language-model geometry from the loaded Gemma config.
+
+        A multimodal root nests it under ``text_config``; a text-only config
+        carries it at the root. The dataclass defaults hold for every encoder
+        root shipped so far, so this is not a behavior change -- it just stops
+        the defaults from being authoritative.
+        The feature-extractor linears are sized in ``__init__`` from
+        ``feature_extractor_in_features``, before the Gemma weights load, so a
+        geometry mismatch is raised here rather than surfacing later as an
+        opaque matmul shape error.
+        """
+        text_config = getattr(gemma_config, "text_config", gemma_config)
+        arch = self.config.arch_config
+        arch.hidden_size = text_config.hidden_size
+        arch.num_hidden_layers = text_config.num_hidden_layers
+        # +1: the stacked hidden states include the embedding output.
+        expected_in_features = arch.hidden_size * (arch.num_hidden_layers + 1)
+        if expected_in_features != arch.feature_extractor_in_features:
+            raise ValueError(
+                "Gemma geometry does not match the configured feature "
+                f"extractor: loaded config gives hidden_size={arch.hidden_size} "
+                f"x (num_hidden_layers={arch.num_hidden_layers} + 1) = "
+                f"{expected_in_features}, but feature_extractor_in_features is "
+                f"{arch.feature_extractor_in_features}."
+            )
+
     def _run_feature_extractor(
         self,
         hidden_states: tuple[torch.Tensor, ...],
@@ -670,6 +702,10 @@ class LTX2GemmaTextEncoderModel(TextEncoder):
                 name = "feature_extractor_linear.aggregate_embed.weight"
             elif name.startswith("video_connector."):
                 name = name.replace("video_connector.", "embeddings_connector.", 1)
+            elif name.startswith("video_embeddings_connector."):
+                # Some checkpoints name the video connector sub-tree after its
+                # modality; the module is just ``embeddings_connector``.
+                name = name.replace("video_embeddings_connector.", "embeddings_connector.", 1)
             elif name.startswith("audio_connector."):
                 name = name.replace("audio_connector.", "audio_embeddings_connector.", 1)
             if name not in params_dict:

fastvideo/models/encoders/mmaudio_clip.py

@@ -0,0 +1,172 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Native split DFN5B CLIP encoders used by MMAudio.
+
+MMAudio uses one OpenCLIP checkpoint in two different ways: normalized
+projected image embeddings and normalized token-wise text hidden states. The
+classes here reuse FastVideo's native CLIP transformer while exposing those
+two contracts as independently loadable pipeline components.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Iterable
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from fastvideo.attention.backends.sdpa import SDPAMetadata
+from fastvideo.configs.models.encoders import BaseEncoderOutput
+from fastvideo.configs.models.encoders.mmaudio_clip import (
+    MMAudioDFNCLIPTextConfig,
+    MMAudioDFNCLIPVisionConfig,
+)
+from fastvideo.models.encoders.clip import CLIPTextModel, CLIPVisionModel
+from fastvideo.forward_context import set_forward_context
+from fastvideo.models.loader.weight_utils import default_weight_loader
+
+
+class MMAudioDFNCLIPTextEncoder(CLIPTextModel):
+    """Return normalized CLIP hidden states for every text token."""
+
+    def __init__(self, config: MMAudioDFNCLIPTextConfig) -> None:
+        super().__init__(config)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor | None,
+        position_ids: torch.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        output_hidden_states: bool | None = None,
+        **kwargs,
+    ) -> BaseEncoderOutput:
+        if input_ids is not None:
+            sequence_length = input_ids.shape[-1]
+            device = input_ids.device
+        elif inputs_embeds is not None:
+            sequence_length = inputs_embeds.shape[-2]
+            device = inputs_embeds.device
+        else:
+            raise ValueError("MMAudio CLIP text encoding requires input_ids or inputs_embeds.")
+        model_dtype = next(self.parameters()).dtype
+        causal_mask = torch.empty(
+            (sequence_length, sequence_length),
+            device=device,
+            dtype=model_dtype,
+        )
+        causal_mask.fill_(float("-inf"))
+        causal_mask.triu_(1)
+        metadata = SDPAMetadata(
+            current_timestep=0,
+            attn_mask=causal_mask[None, None],
+            is_causal=True,
+        )
+        with set_forward_context(current_timestep=0, attn_metadata=metadata):
+            output = super().forward(
+                input_ids=input_ids,
+                position_ids=position_ids,
+                attention_mask=attention_mask,
+                inputs_embeds=inputs_embeds,
+                output_hidden_states=output_hidden_states,
+                **kwargs,
+            )
+        assert output.last_hidden_state is not None
+        normalized = F.normalize(output.last_hidden_state, dim=-1)
+        return BaseEncoderOutput(
+            last_hidden_state=normalized,
+            pooler_output=output.pooler_output,
+            hidden_states=output.hidden_states,
+            attentions=output.attentions,
+            attention_mask=output.attention_mask,
+        )
+
+    def load_weights(
+        self,
+        weights: Iterable[tuple[str, torch.Tensor]],
+    ) -> set[str]:
+        """Load either split OpenCLIP Q/K/V or converted fused QKV weights."""
+        params = dict(self.named_parameters())
+        loaded: set[str] = set()
+        for name, tensor in weights:
+            if name in params:
+                parameter = params[name]
+                loader = getattr(parameter, "weight_loader", default_weight_loader)
+                loader(parameter, tensor)
+                loaded.add(name)
+                continue
+            for param_name, weight_name, shard_id in self.config.arch_config.stacked_params_mapping:
+                if weight_name not in name:
+                    continue
+                target_name = name.replace(weight_name, param_name)
+                if target_name not in params:
+                    continue
+                parameter = params[target_name]
+                parameter.weight_loader(parameter, tensor, shard_id)
+                loaded.add(target_name)
+                break
+        return loaded
+
+
+class MMAudioDFNCLIPVisionEncoder(CLIPVisionModel):
+    """Return normalized projected CLS embeddings for individual frames."""
+
+    def __init__(self, config: MMAudioDFNCLIPVisionConfig) -> None:
+        super().__init__(config)
+        self.visual_projection = nn.Linear(config.hidden_size, config.projection_dim, bias=False)
+
+    def forward(
+        self,
+        pixel_values: torch.Tensor,
+        feature_sample_layers: list[int] | None = None,
+        **kwargs,
+    ) -> BaseEncoderOutput:
+        del kwargs
+        if feature_sample_layers is not None:
+            raise ValueError("MMAudio DFN5B vision encoding requires the final CLS token")
+        tokens = self.vision_model(pixel_values, feature_sample_layers=None)
+        image_features = F.normalize(self.visual_projection(tokens[:, 0]), dim=-1)
+        return BaseEncoderOutput(last_hidden_state=image_features, pooler_output=image_features)
+
+    def load_weights(
+        self,
+        weights: Iterable[tuple[str, torch.Tensor]],
+    ) -> set[str]:
+        params = dict(self.named_parameters())
+        loaded: set[str] = set()
+        layer_count = len(self.vision_model.encoder.layers)
+        for name, tensor in weights:
+            if name.startswith("vision_model.encoder.layers"):
+                layer_index = int(name.split(".")[3])
+                if layer_index >= layer_count:
+                    continue
+            # Converted checkpoints already contain fused qkv_proj tensors.
+            # Check exact names before matching the ``v_proj`` substring that
+            # also occurs at the end of ``qkv_proj``.
+            if name in params:
+                parameter = params[name]
+                loader = getattr(parameter, "weight_loader", default_weight_loader)
+                loader(parameter, tensor)
+                loaded.add(name)
+                continue
+            for param_name, weight_name, shard_id in self.config.arch_config.stacked_params_mapping:
+                if weight_name not in name:
+                    continue
+                target_name = name.replace(weight_name, param_name)
+                if target_name not in params:
+                    continue
+                parameter = params[target_name]
+                parameter.weight_loader(parameter, tensor, shard_id)
+                loaded.add(target_name)
+                break
+            else:
+                if name not in params:
+                    continue
+                parameter = params[name]
+                loader = getattr(parameter, "weight_loader", default_weight_loader)
+                loader(parameter, tensor)
+                loaded.add(name)
+        return loaded
+
+
+EntryClass = [MMAudioDFNCLIPTextEncoder, MMAudioDFNCLIPVisionEncoder]

fastvideo/models/encoders/mmaudio_synchformer.py

@@ -0,0 +1,104 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Synchformer visual conditioner used by MMAudio.
+
+The MotionFormer implementation is shared with FastVideo's existing
+audio/video synchronization evaluator. This production adapter deliberately
+owns only the visual feature extractor used by MMAudio, so its state-dict and
+forward numerics match the official ``Synchformer`` module without carrying
+the evaluator's unused audio and classification heads.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Iterable
+
+import torch
+
+from fastvideo.configs.models.encoders import BaseEncoderOutput
+from fastvideo.configs.models.encoders.mmaudio_synchformer import (
+    MMAudioSynchformerConfig,
+)
+from fastvideo.models.encoders.base import ImageEncoder
+from fastvideo.models.loader.weight_utils import default_weight_loader
+from fastvideo.third_party.synchformer.motionformer import MotionFormer
+
+
+class MMAudioSynchformerVisualEncoder(ImageEncoder):
+    """Extract temporal synchronization tokens from 25 FPS video frames."""
+
+    def __init__(self, config: MMAudioSynchformerConfig) -> None:
+        super().__init__(config)
+        self.vfeat_extractor = MotionFormer(
+            extract_features=True,
+            factorize_space_time=True,
+            agg_space_module="TransformerEncoderLayer",
+            agg_time_module="torch.nn.Identity",
+            add_global_repr=False,
+        )
+
+    def forward_segmented(self, segments: torch.Tensor) -> torch.Tensor:
+        """Encode ``[B, S, 16, 3, 224, 224]`` frame segments."""
+        if segments.ndim != 6:
+            raise ValueError(f"Synchformer segments must have shape [B, S, T, C, H, W], got {tuple(segments.shape)}")
+        _, _, frames, channels, height, width = segments.shape
+        expected = self.config.arch_config
+        if (
+            frames != expected.segment_size
+            or channels != expected.num_channels
+            or height != expected.image_size
+            or width != expected.image_size
+        ):
+            raise ValueError(
+                "Synchformer requires segments shaped "
+                f"[B, S, {expected.segment_size}, {expected.num_channels}, "
+                f"{expected.image_size}, {expected.image_size}], got "
+                f"{tuple(segments.shape)}"
+            )
+        visual = segments.permute(0, 1, 3, 2, 4, 5)
+        return self.vfeat_extractor(visual)
+
+    def forward(
+        self,
+        pixel_values: torch.Tensor,
+        **kwargs,
+    ) -> BaseEncoderOutput:
+        """Encode contiguous ``[B, T, 3, 224, 224]`` 25 FPS frames."""
+        del kwargs
+        if pixel_values.ndim != 5:
+            raise ValueError(f"Synchformer video must have shape [B, T, C, H, W], got {tuple(pixel_values.shape)}")
+
+        config = self.config.arch_config
+        frame_count = pixel_values.shape[1]
+        if frame_count < config.segment_size:
+            raise ValueError(f"Synchformer needs at least {config.segment_size} frames, got {frame_count}")
+
+        segments = (
+            pixel_values.unfold(
+                dimension=1,
+                size=config.segment_size,
+                step=config.segment_stride,
+            )
+            .permute(0, 1, 5, 2, 3, 4)
+            .contiguous()
+        )
+        features = self.forward_segmented(segments)
+        features = features.flatten(1, 2)
+        return BaseEncoderOutput(last_hidden_state=features)
+
+    def load_weights(
+        self,
+        weights: Iterable[tuple[str, torch.Tensor]],
+    ) -> set[str]:
+        params = dict(self.named_parameters())
+        loaded: set[str] = set()
+        for name, tensor in weights:
+            if not name.startswith("vfeat_extractor.") or name not in params:
+                continue
+            parameter = params[name]
+            weight_loader = getattr(parameter, "weight_loader", default_weight_loader)
+            weight_loader(parameter, tensor)
+            loaded.add(name)
+        return loaded
+
+
+EntryClass = MMAudioSynchformerVisualEncoder

fastvideo/models/loader/component_loader.py

@@ -9,7 +9,7 @@ from abc import ABC, abstractmethod
 from collections.abc import Generator, Iterable
 from contextlib import nullcontext
 from copy import deepcopy
-from typing import cast
+from typing import Any, cast
 
 import torch
 import torch.distributed as dist
@@ -33,6 +33,11 @@ from fastvideo.logger import init_logger
 from fastvideo.models.encoders.base import TextEncoder
 from fastvideo.models.hf_transformer_utils import get_diffusers_config
 from fastvideo.models.loader.fsdp_load import maybe_load_fsdp_model, shard_model
+from fastvideo.models.loader.ltx_single_file import (
+    component_weights,
+    is_single_file_bundle,
+    read_ltx_metadata,
+)
 from fastvideo.models.loader.utils import set_default_torch_dtype
 from fastvideo.models.loader.weight_utils import (
     filter_duplicate_safetensors_files,
@@ -102,6 +107,8 @@ class ComponentLoader(ABC):
             "image_processor": (ImageProcessorLoader, "transformers"),
             "feature_extractor": (ImageProcessorLoader, "transformers"),
             "image_encoder": (ImageEncoderLoader, "transformers"),
+            "image_encoder_2": (ImageEncoderLoader, "transformers"),
+            "image_encoder_3": (ImageEncoderLoader, "transformers"),
             "vision_language_encoder": (VisionLanguageEncoderLoader, "transformers"),
             "processor": (ProcessorLoader, "transformers"),
             "upsampler": (UpsamplerLoader, "diffusers"),
@@ -136,6 +143,41 @@ class ComponentLoader(ABC):
         return GenericComponentLoader(transformers_or_diffusers)
 
 
+def _check_connector_widths(transformer_config: dict[str, Any]) -> None:
+    """Fail loudly when a connector's declared factorization is inconsistent.
+
+    A connector's width is ``num_attention_heads * attention_head_dim``, and it
+    has to equal the cross-attention width of the stream it feeds. This is the
+    class of mistake that loads cleanly and computes wrong -- a factorization
+    that multiplies out to the wrong number produces correctly shaped tensors
+    everywhere except the one parameter it silently mis-sizes. Check it while
+    both numbers are still in hand rather than discovering it as a shape
+    assertion thousands of tensors later, or not at all.
+
+    Only checks what the checkpoint actually declares; a missing field means
+    the arch default applies and there is nothing to contradict.
+    """
+    for heads_key, dim_key, width_key, label in (
+        ("connector_num_attention_heads", "connector_attention_head_dim",
+         "cross_attention_dim", "video"),
+        ("audio_connector_num_attention_heads",
+         "audio_connector_attention_head_dim", "audio_cross_attention_dim",
+         "audio"),
+    ):
+        heads = transformer_config.get(heads_key)
+        head_dim = transformer_config.get(dim_key)
+        width = transformer_config.get(width_key)
+        if heads is None or head_dim is None or width is None:
+            continue
+        if heads * head_dim != width:
+            raise ValueError(
+                f"Checkpoint declares an inconsistent {label} connector shape: "
+                f"{heads_key}={heads} x {dim_key}={head_dim} = {heads * head_dim}, "
+                f"but {width_key}={width}. The connector feeds that stream, so "
+                "the two must agree; refusing to build a model whose weights "
+                "would load into mis-sized parameters.")
+
+
 class TextEncoderLoader(ComponentLoader):
     """Loader for text encoders."""
 
@@ -249,6 +291,13 @@ class TextEncoderLoader(ComponentLoader):
         #     revision=fastvideo_args.revision,
         #     model_override_args=None,
         # )
+        # For a bundle, `model_path` is the declared text-encoder root rather
+        # than a `text_encoder/` directory: the language model lives there, its
+        # projection and connectors live inside the bundle, and the fields the
+        # directory layout reads out of `<repo>/transformer/config.json` come
+        # from the bundle's metadata instead.
+        bundle_path = fastvideo_args.model_path
+        single_file = is_single_file_bundle(bundle_path)
         model_config = get_diffusers_config(model=model_path)
         model_config.pop("_name_or_path", None)
         model_config.pop("transformers_version", None)
@@ -275,6 +324,15 @@ class TextEncoderLoader(ComponentLoader):
         if gemma_path and not gemma_path_from_candidate:
             if not os.path.isabs(gemma_path):
                 model_config["gemma_model_path"] = os.path.normpath(os.path.join(repo_root, gemma_path))
+        if single_file:
+            # The encoder root was declared, not discovered, so it wins over
+            # anything the probing above turned up.
+            model_config["gemma_model_path"] = model_path
+            # That root's config.json describes the language model it holds,
+            # which is not the class to build here: the root supplies the
+            # encoder's dimensions, the wrapper class supplies the connectors,
+            # and only the wrapper is buildable from this loader.
+            model_config["architectures"] = ["LTX2GemmaTextEncoderModel"]
         transformer_config_path = os.path.join(repo_root, "transformer", "config.json")
         if os.path.isfile(transformer_config_path):
             try:
@@ -288,6 +346,40 @@ class TextEncoderLoader(ComponentLoader):
                     rope_type = transformer_config.get("rope_type")
                     if rope_type is not None:
                         model_config["connector_rope_type"] = rope_type
+                # Each per-modality projection feeds its connector, so its
+                # output width is that stream's cross-attention width, which
+                # only the transformer config declares. Absent from both
+                # configs -> the arch config default stands.
+                for src, dst in (
+                    ("cross_attention_dim",
+                     "video_feature_extractor_out_features"),
+                    ("audio_cross_attention_dim",
+                     "audio_feature_extractor_out_features"),
+                    # The rest of the text stack's shape is declared under these
+                    # exact names too. Leaving any of them to the arch default
+                    # builds a DIFFERENT model than the checkpoint describes:
+                    # the audio connector silently inherits the video
+                    # connector's width, and the feature extractor is built in
+                    # the wrong one of its two forms. Same names on both sides,
+                    # so the mapping is identity.
+                    *((name, name) for name in (
+                        "connector_num_attention_heads",
+                        "connector_attention_head_dim",
+                        "connector_num_layers",
+                        "connector_num_learnable_registers",
+                        "connector_positional_embedding_theta",
+                        "connector_positional_embedding_max_pos",
+                        "connector_apply_gated_attention",
+                        "audio_connector_num_attention_heads",
+                        "audio_connector_attention_head_dim",
+                        "audio_connector_num_layers",
+                        # Selects which feature-extractor form is built.
+                        "caption_proj_before_connector",
+                    )),
+                ):
+                    if dst not in model_config and src in transformer_config:
+                        model_config[dst] = transformer_config[src]
+                _check_connector_widths(transformer_config)
             except json.JSONDecodeError:
                 pass
         logger.info("HF Model config: %s", model_config)
@@ -330,6 +422,12 @@ class TextEncoderLoader(ComponentLoader):
             fastvideo_args,
             encoder_precision,
             use_text_encoder_override=True,
+            # Everything this model owns is inside the bundle: the projection
+            # and both connectors. The language model at the encoder root is
+            # loaded separately and lazily by the model itself, and is filtered
+            # out of `named_parameters`, so it must not be routed through here.
+            weight_iterator=(component_weights(bundle_path, "text_encoder")
+                             if single_file else None),
         )
 
     def load_model(
@@ -340,13 +438,16 @@ class TextEncoderLoader(ComponentLoader):
             fastvideo_args: FastVideoArgs,
             dtype: str = "fp16",
             use_text_encoder_override: bool = False,  # prevent subclasses from misusing
+            cpu_offload: bool | None = None,
+            weight_iterator: Iterable[tuple[str, torch.Tensor]] | None = None,
     ):
-        use_cpu_offload = (fastvideo_args.text_encoder_cpu_offload
-                           and len(getattr(model_config, "_fsdp_shard_conditions", [])) > 0)
+        if cpu_offload is None:
+            cpu_offload = fastvideo_args.text_encoder_cpu_offload
+        use_cpu_offload = (cpu_offload and len(getattr(model_config, "_fsdp_shard_conditions", [])) > 0)
 
         from fastvideo.platforms import current_platform
 
-        if fastvideo_args.text_encoder_cpu_offload:
+        if cpu_offload:
             target_device = (torch.device("mps") if current_platform.is_mps() else torch.device("cpu"))
 
         # Set quantization config if specified
@@ -382,12 +483,18 @@ class TextEncoderLoader(ComponentLoader):
                         [fastvideo_args.override_text_encoder_safetensors],
                         to_cpu=use_cpu_offload,
                     ))  # type: ignore
+            elif weight_iterator is not None:
+                # Same contract as `maybe_load_fsdp_model`'s `weight_iterator`:
+                # already-routed `(name, cpu_tensor)` pairs, for checkpoints
+                # that are not a directory of per-component weight files.
+                self.counter_before_loading_weights = time.perf_counter()
+                loaded_weights: set[str] = model.load_weights(weight_iterator)  # type: ignore
             else:
                 loaded_weights: set[str] = model.load_weights(
                     self._get_all_weights(
                         model,
                         model_path,
-                        to_cpu=fastvideo_args.text_encoder_cpu_offload,
+                        to_cpu=cpu_offload,
                     ))  # type: ignore
 
             self.counter_after_loading_weights = time.perf_counter()
@@ -463,7 +570,36 @@ class ImageEncoderLoader(TextEncoderLoader):
         model_config.pop("model_type", None)
         logger.info("HF Model config: %s", model_config)
 
-        encoder_config = fastvideo_args.pipeline_config.image_encoder_config
+        base = os.path.basename(os.path.normpath(model_path))
+        index = 0
+        if base.startswith("image_encoder_"):
+            try:
+                index = int(base.split("_")[-1]) - 1
+            except ValueError:
+                index = 0
+
+        encoder_configs = getattr(
+            fastvideo_args.pipeline_config, "image_encoder_configs", None)
+        encoder_precisions = getattr(
+            fastvideo_args.pipeline_config, "image_encoder_precisions", None)
+        if encoder_configs is None:
+            if index != 0:
+                raise IndexError(
+                    f"image encoder index {index} requires pipeline_config."
+                    "image_encoder_configs")
+            encoder_config = fastvideo_args.pipeline_config.image_encoder_config
+            encoder_precision = fastvideo_args.pipeline_config.image_encoder_precision
+        else:
+            if index < 0 or index >= len(encoder_configs):
+                raise IndexError(
+                    f"image encoder index {index} out of range for "
+                    f"image_encoder_configs (len={len(encoder_configs)})")
+            encoder_config = encoder_configs[index]
+            if encoder_precisions is None or index >= len(
+                    encoder_precisions):
+                raise IndexError(
+                    f"image encoder index {index} has no matching precision")
+            encoder_precision = encoder_precisions[index]
         encoder_config.update_model_arch(model_config)
         record_resolved_attention_backend(encoder_config)
 
@@ -479,7 +615,8 @@ class ImageEncoderLoader(TextEncoderLoader):
             encoder_config,
             target_device,
             fastvideo_args,
-            fastvideo_args.pipeline_config.image_encoder_precision,
+            encoder_precision,
+            cpu_offload=fastvideo_args.image_encoder_cpu_offload,
         )
 
 
@@ -664,7 +801,17 @@ class VAELoader(ComponentLoader):
 
     def load(self, model_path: str, fastvideo_args: FastVideoArgs):
         """Load the VAE based on the model path, and inference args."""
-        config = get_diffusers_config(model=model_path)
+        if is_single_file_bundle(model_path):
+            # The bundle stores the VAE config flat; the converted directory
+            # layout nests the same fields under "vae" and hoists the class out
+            # (`convert_ltx2_weights.py::_wrap_component_config`), and the build
+            # path below keys off that nesting. Match the shape rather than
+            # forking the build path.
+            section = deepcopy(read_ltx_metadata(model_path).config["vae"])
+            declared_class = section.pop("_class_name", None)
+            config = {"_class_name": declared_class, "vae": section}
+        else:
+            config = get_diffusers_config(model=model_path)
         class_name = config.pop("_class_name")
         config.pop("_name_or_path", None)
         assert class_name is not None, (
@@ -783,15 +930,21 @@ class VAELoader(ComponentLoader):
                 vae_cls, _ = ModelRegistry.resolve_model_cls(class_name)
                 vae = vae_cls(vae_config).to(target_device)
 
-        # Find all safetensors files
-        safetensors_list = glob.glob(os.path.join(str(model_path), "*.safetensors"))
-        if not safetensors_list:
-            raise ValueError(f"No safetensors files found in {model_path}")
-        # Common case: a single `.safetensors` checkpoint file.
-        # Some models may be sharded into multiple files; in that case we merge.
-        loaded = {}
-        for sf_file in safetensors_list:
-            loaded.update(safetensors_load_file(sf_file))
+        if is_single_file_bundle(model_path):
+            # `safetensors_load_file` takes a path, not an iterator, so the
+            # bundle's prefix-routed tensors are materialized here instead.
+            # The VAE is small enough for that (see `component_weights`).
+            loaded = dict(component_weights(model_path, "vae"))
+        else:
+            # Find all safetensors files
+            safetensors_list = glob.glob(os.path.join(str(model_path), "*.safetensors"))
+            if not safetensors_list:
+                raise ValueError(f"No safetensors files found in {model_path}")
+            # Common case: a single `.safetensors` checkpoint file.
+            # Some models may be sharded into multiple files; in that case we merge.
+            loaded = {}
+            for sf_file in safetensors_list:
+                loaded.update(safetensors_load_file(sf_file))
 
         # LTX-2 CausalVideoAutoencoder needs per_channel_statistics remapping
         if class_name == "CausalVideoAutoencoder" and "vae" in config:
@@ -859,14 +1012,25 @@ class AudioDecoderLoader(ComponentLoader):
             return audio_vae.eval()
 
         precision = getattr(fastvideo_args.pipeline_config, "audio_decoder_precision", "bf16")
-        with set_default_torch_dtype(PRECISION_TO_TYPE[precision]):
-            audio_decoder = model_cls(config).to(target_device)
+        # MMAudio normalizes its magnitude-preserving convolution weights in
+        # fp32 and only then casts the whole feature utility module to bf16.
+        # Constructing/loading directly in bf16 quantizes the unnormalized
+        # checkpoint first and changes the decoded mel trajectory.
+        construction_precision = "fp32" if class_name == "MMAudioVAE" else precision
+        construction_device = torch.device("cpu") if class_name == "MMAudioVAE" else target_device
+        with set_default_torch_dtype(PRECISION_TO_TYPE[construction_precision]):
+            audio_decoder = model_cls(config).to(construction_device)
 
         safetensors_list = glob.glob(os.path.join(str(model_path), "*.safetensors"))
         loaded: dict[str, torch.Tensor] = {}
         for sf_file in safetensors_list:
             loaded.update(safetensors_load_file(sf_file))
 
+        if class_name == "MMAudioVAE":
+            audio_decoder.load_state_dict(loaded, strict=True)
+            audio_decoder.remove_weight_norm()
+            return audio_decoder.to(device=target_device, dtype=PRECISION_TO_TYPE[precision]).eval()
+
         decoder_state = {}
         for name, tensor in loaded.items():
             if name.startswith("decoder."):
@@ -880,7 +1044,7 @@ class AudioDecoderLoader(ComponentLoader):
 
 
 class VocoderLoader(ComponentLoader):
-    """Loader for LTX-2 vocoder."""
+    """Loader for native vocoders."""
 
     def load(self, model_path: str, fastvideo_args: FastVideoArgs):
         config = get_diffusers_config(model=model_path)
@@ -890,14 +1054,23 @@ class VocoderLoader(ComponentLoader):
         target_device = get_local_torch_device()
 
         precision = getattr(fastvideo_args.pipeline_config, "vocoder_precision", "bf16")
-        with set_default_torch_dtype(PRECISION_TO_TYPE[precision]):
-            vocoder = model_cls(config).to(target_device)
+        # Canonical BigVGAN likewise removes parametrized weight norm in fp32
+        # before the official MMAudio feature module is cast to bf16.
+        construction_precision = "fp32" if class_name == "BigVGANV2" else precision
+        construction_device = torch.device("cpu") if class_name == "BigVGANV2" else target_device
+        with set_default_torch_dtype(PRECISION_TO_TYPE[construction_precision]):
+            vocoder = model_cls(config).to(construction_device)
 
         safetensors_list = glob.glob(os.path.join(str(model_path), "*.safetensors"))
         loaded: dict[str, torch.Tensor] = {}
         for sf_file in safetensors_list:
             loaded.update(safetensors_load_file(sf_file))
 
+        if class_name == "BigVGANV2":
+            vocoder.load_state_dict(loaded, strict=True)
+            vocoder.remove_weight_norm()
+            return vocoder.to(device=target_device, dtype=PRECISION_TO_TYPE[precision]).eval()
+
         target_module = getattr(vocoder, "model", vocoder)
         target_module.load_state_dict(loaded, strict=False)
         return vocoder.eval()
@@ -920,7 +1093,16 @@ class TransformerLoader(ComponentLoader):
 
     def load(self, model_path: str, fastvideo_args: FastVideoArgs):
         """Load the transformer based on the model path, and inference args."""
-        config = get_diffusers_config(model=model_path)
+        # LTX ships one bundle holding every component instead of a
+        # per-component directory: the config lives in the file's
+        # ``__metadata__`` and the transformer's tensors are selected by their
+        # top-level key prefix.
+        single_file = is_single_file_bundle(model_path)
+        if single_file:
+            config = deepcopy(
+                read_ltx_metadata(model_path).config["transformer"])
+        else:
+            config = get_diffusers_config(model=model_path)
         hf_config = deepcopy(config)
         cls_name = config.pop("_class_name")
         config.pop("_name_or_path", None)
@@ -954,9 +1136,12 @@ class TransformerLoader(ComponentLoader):
         model_cls, _ = ModelRegistry.resolve_model_cls(cls_name)
 
         # Find all safetensors files
-        safetensors_list = glob.glob(os.path.join(str(model_path), "*.safetensors"))
-        if not safetensors_list:
-            raise ValueError(f"No safetensors files found in {model_path}")
+        if single_file:
+            safetensors_list = [str(model_path)]
+        else:
+            safetensors_list = glob.glob(os.path.join(str(model_path), "*.safetensors"))
+            if not safetensors_list:
+                raise ValueError(f"No safetensors files found in {model_path}")
 
         # arch_config can infer architecture from weight keys (e.g. Flux2 layer counts)
         update_fn = getattr(dit_config.arch_config, "update_from_weight_keys", None)
@@ -1011,6 +1196,11 @@ class TransformerLoader(ComponentLoader):
                     "hf_config": hf_config
                 },
                 weight_dir_list=safetensors_list,
+                # Route the bundle's tensors to the transformer, prefix
+                # stripped. Skipped when custom init weights replaced the
+                # file list above.
+                weight_iterator=(component_weights(model_path, "transformer")
+                                 if single_file and not use_custom_weights else None),
                 device=get_local_torch_device(),
                 hsdp_replicate_dim=fastvideo_args.hsdp_replicate_dim,
                 hsdp_shard_dim=fastvideo_args.hsdp_shard_dim,
@@ -1055,7 +1245,12 @@ class SchedulerLoader(ComponentLoader):
 
     def load(self, model_path: str, fastvideo_args: FastVideoArgs):
         """Load the scheduler based on the model path, and inference args."""
-        config = get_diffusers_config(model=model_path)
+        # A bundle keeps the scheduler's config in its own metadata; there is
+        # no `scheduler/` directory holding a `scheduler_config.json`.
+        if is_single_file_bundle(model_path):
+            config = deepcopy(read_ltx_metadata(model_path).config["scheduler"])
+        else:
+            config = get_diffusers_config(model=model_path)
 
         class_name = config.pop("_class_name")
         assert class_name is not None, (

fastvideo/models/loader/fsdp_load.py

@@ -7,7 +7,7 @@
 from __future__ import annotations
 import os
 import contextlib
-from collections.abc import Callable, Generator
+from collections.abc import Callable, Generator, Iterable
 from itertools import chain
 from typing import Any
 
@@ -136,9 +136,15 @@ def maybe_load_fsdp_model(
     pin_cpu_memory: bool = True,
     enable_torch_compile: bool = False,
     torch_compile_kwargs: dict[str, Any] | None = None,
+    weight_iterator: Iterable[tuple[str, torch.Tensor]] | None = None,
 ) -> torch.nn.Module:
     """
     Load the model with FSDP if is training, else load the model without FSDP.
+
+    ``weight_iterator`` overrides reading ``weight_dir_list``, for checkpoints
+    whose keys need routing before they reach the model (see
+    ``ltx_single_file.component_weights``). It must yield the same
+    ``(name, cpu_tensor)`` pairs that ``safetensors_weights_iterator`` does.
     """
     # NOTE(will): cast_forward_inputs=True shouldn't be needed as we are
     # manually casting the inputs to the model
@@ -201,7 +207,8 @@ def maybe_load_fsdp_model(
                     fsdp_shard_conditions=model._fsdp_shard_conditions,
                     pin_cpu_memory=pin_cpu_memory)
 
-    weight_iterator = safetensors_weights_iterator(weight_dir_list, to_cpu=True)
+    if weight_iterator is None:
+        weight_iterator = safetensors_weights_iterator(weight_dir_list, to_cpu=True)
     param_names_mapping_fn = get_param_names_mapping(model.param_names_mapping)
     load_model_from_full_model_state_dict(
         model,

fastvideo/models/loader/ltx_single_file.py

@@ -0,0 +1,296 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Single-file LTX checkpoint support.
+
+LTX ships one ``.safetensors`` bundle holding every component, so the two
+assumptions the directory-based loaders make do not hold:
+
+* there is no per-component ``config.json`` -- every component's config lives
+  in the file's ``__metadata__`` (safetensors stores an 8-byte little-endian
+  header length, then that many bytes of JSON whose ``__metadata__`` object is
+  a flat ``str -> str`` map; ``safe_open(...).metadata()`` returns it);
+* there is no per-component directory -- components are told apart by the
+  top-level prefix on each tensor key.
+
+This module reads that metadata and routes tensors by prefix. It does not
+convert anything to a diffusers layout.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+from collections.abc import Generator
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+from safetensors import safe_open
+
+from fastvideo.configs.models.dits.ltx2 import LTX2VideoConfig
+from fastvideo.logger import init_logger
+
+logger = init_logger(__name__)
+
+# Top-level tensor key prefix per component. Verified against the LTX bundle
+# header; every tensor in the file starts with one of these except
+# ``duration_head.``, which has no FastVideo component today.
+COMPONENT_PREFIXES: dict[str, str] = {
+    "transformer": "model.diffusion_model.",
+    "vae": "vae.",
+    "audio_vae": "audio_vae.",
+    "vocoder": "vocoder.",
+    "text_encoder": "text_embedding_projection.",
+}
+
+# Both embeddings connectors are stored *under* the transformer prefix, but
+# FastVideo builds and runs them in the text encoder
+# (``LTX2GemmaTextEncoderModel``'s ``Embeddings1DConnector``), not in the DiT,
+# so they are routed to the text encoder instead. Only the transformer prefix
+# is stripped: the sub-tree name survives so the text encoder's own rename
+# table can map it onto the module name, the same way the already-converted
+# repo layout is remapped. Matches how ``convert_ltx2_weights.py`` splits them.
+TEXT_STACK_SUBPREFIXES: tuple[str, ...] = (
+    "video_embeddings_connector.",
+    "audio_embeddings_connector.",
+)
+
+
+def is_single_file_bundle(path: str) -> bool:
+    """True when a model path names a bundle rather than a component directory."""
+    return str(path).endswith(".safetensors")
+
+
+@dataclass(frozen=True)
+class LTXCheckpointMetadata:
+    """Parsed ``__metadata__`` of a single-file LTX checkpoint.
+
+    ``config`` holds the per-component config sections keyed by component
+    (``transformer``, ``vae``, ``scheduler``, ``audio_vae``, ``vocoder``).
+    ``model_version`` and ``gemma_source_checkpoint`` are carried through so
+    callers can validate the text encoder against the checkpoint that trained
+    it; ``gemma_source_checkpoint`` is absent on some variants. ``variant`` is
+    the training variant the header declares, when it declares one (see
+    :func:`bundle_variant`).
+    """
+
+    config: dict[str, Any]
+    model_version: str | None
+    gemma_source_checkpoint: dict[str, Any] | None
+    variant: str | None = None
+
+
+def read_ltx_metadata(path: str) -> LTXCheckpointMetadata:
+    """Read the ``__metadata__`` config out of a single-file LTX checkpoint.
+
+    Only the header is parsed; no tensor data is touched. The remaining
+    metadata entries are deliberately not returned or logged -- the bundle also
+    carries a full license text and an ``encrypted_wandb_properties`` blob.
+    """
+    with safe_open(path, framework="pt") as f:
+        metadata = f.metadata() or {}
+
+    config = json.loads(metadata["config"])
+
+    transformer = config.get("transformer")
+    if transformer is not None:
+        # ``frequencies_precision`` is the checkpoint's name for what the arch
+        # config calls ``double_precision_rope``; without this the field would
+        # silently fall back to its dataclass default.
+        precision = transformer.get("frequencies_precision")
+        if precision is not None:
+            transformer["double_precision_rope"] = precision == "float64"
+
+    gemma_source = metadata.get("gemma_source_checkpoint")
+    return LTXCheckpointMetadata(
+        config=config,
+        model_version=metadata.get("model_version"),
+        gemma_source_checkpoint=(json.loads(gemma_source) if gemma_source is not None else None),
+        variant=metadata.get("variant"),
+    )
+
+
+def bundle_variant(metadata: LTXCheckpointMetadata, path: str) -> str:
+    """The bundle's training variant: ``"distilled"`` or ``"base"``.
+
+    Decides which sampling preset a bundle gets (a distilled model wants its
+    short no-CFG schedule; everything else wants the standard one), so it must
+    be answerable from the header alone, without loading weights. An explicit
+    ``variant`` entry in the file metadata wins when the checkpoint declares
+    one.
+
+    ponytail: filename fallback -- the known bundles declare no variant marker
+    in their headers (a distilled file and its sft sibling differ only in
+    fields incidental to the variant), so the "distilled" token in the file
+    name is the only signal available today. Drop the fallback when
+    checkpoints start declaring ``variant``.
+    """
+    declared = metadata.variant or os.path.basename(path)
+    return "distilled" if "distilled" in declared.lower() else "base"
+
+
+def bundle_model_index(path: str) -> dict[str, Any]:
+    """Build a ``model_index.json``-shaped dict out of a bundle's own metadata.
+
+    The pipeline loader is written against a diffusers repo layout, which
+    answers two questions: which components exist, and what class is each. A
+    bundle already answers both in its ``__metadata__``, so this only reshapes
+    the answer -- it mirrors the entries
+    ``convert_ltx2_weights.py::_build_model_index`` writes for the converted
+    directory layout, including the library each is declared under.
+
+    A section that exists but declares no class is emitted as ``[None, None]``
+    rather than dropped. ``ComposedPipelineBase.load_modules`` already treats a
+    null library as "declared, but not something to build" and removes the
+    component from the required set; dropping the key instead would fail its
+    required-module check for a component the checkpoint does carry.
+
+    ``text_encoder`` and ``tokenizer`` are always declared: they live outside
+    the bundle, but the pipeline needs both.
+    """
+    model_index: dict[str, Any] = {
+        # ponytail: the pipeline class is pinned by the registry's bundle
+        # table (`registry._bundle_config_info`) or an explicit
+        # `override_pipeline_cls_name`, and `load_modules` pops both of these
+        # without reading them. Nothing here is entitled to name a pipeline,
+        # so these are placeholders -- they exist only because those pops have
+        # no default.
+        "_class_name": None,
+        "_diffusers_version": None,
+    }
+    for component, section in read_ltx_metadata(path).config.items():
+        cls_name = (section.get("_class_name") if isinstance(section, dict) else None)
+        model_index[component] = (["diffusers", cls_name] if cls_name else [None, None])
+    model_index["text_encoder"] = ["transformers", "LTX2GemmaTextEncoderModel"]
+    model_index["tokenizer"] = ["transformers", "AutoTokenizer"]
+    return model_index
+
+
+def build_dit_config(metadata: LTXCheckpointMetadata) -> LTX2VideoConfig:
+    """Build the LTX-2 DiT config from checkpoint metadata.
+
+    Reuses ``update_model_arch``, so the metadata keys that name an arch field
+    win and the rest of the section is ignored -- no hand-written constants.
+    """
+    config = LTX2VideoConfig()
+    config.update_model_arch(metadata.config["transformer"])
+    return config
+
+
+def component_weights(
+    path: str,
+    component: str,
+) -> Generator[tuple[str, torch.Tensor], None, None]:
+    """Yield ``(key_with_prefix_stripped, tensor)`` for one component.
+
+    The transformer prefix covers two owners -- see
+    ``TEXT_STACK_SUBPREFIXES`` -- so keys under it are split before the
+    component filter applies.
+
+    ``safe_open`` mmaps the file and ``get_tensor`` materializes one tensor at
+    a time, so iterating never holds more than a single tensor in RAM. Callers
+    that need a state dict for a small component can do
+    ``dict(component_weights(path, "vocoder"))``; do not do that for the
+    transformer.
+
+    ponytail: every rank reads the file itself, unlike
+    ``safetensors_weights_iterator``, which has local rank 0 read and broadcast.
+    Correct either way, but the ceiling is filesystem bandwidth -- N ranks read
+    N copies. Add the broadcast if loading a bundle over a shared filesystem
+    turns out to be the bottleneck.
+    """
+    prefix = COMPONENT_PREFIXES[component]
+    transformer_prefix = COMPONENT_PREFIXES["transformer"]
+    with safe_open(path, framework="pt") as f:
+        for key in f.keys():
+            if key.startswith(transformer_prefix):
+                name = key[len(transformer_prefix):]
+                # str.startswith takes a tuple.
+                owner = ("text_encoder" if name.startswith(TEXT_STACK_SUBPREFIXES) else "transformer")
+                if owner != component:
+                    continue
+            elif key.startswith(prefix):
+                name = key[len(prefix):]
+            else:
+                continue
+            yield name, f.get_tensor(key)
+
+
+def resolve_text_encoder_root(
+    configured_path: str | None,
+    override_path: str | None = None,
+    metadata: LTXCheckpointMetadata | None = None,
+) -> str:
+    """Locate the text-encoder root that goes with a single-file bundle.
+
+    A bundle carries every component's *weights* but no pointer to the text
+    encoder, which lives outside it. The root therefore has to be declared,
+    and there are exactly two ways to declare it:
+
+    * ``override_path`` -- an explicit per-run argument, so a one-off run needs
+      no config edit;
+    * ``configured_path`` -- the pipeline config's encoder path, which is where
+      model composition already lives and which survives files being moved.
+
+    Deliberately absent: any search of the bundle's directory. Picking up an
+    encoder nobody asked for changes what gets loaded without being requested,
+    and silently picks the wrong one when two sit side by side. When neither
+    source is set this raises and names both, rather than guessing.
+
+    ``metadata`` is used only to *validate* a declared root, never to find one:
+    a bundle may not declare an encoder pairing at all, so discovery cannot
+    depend on it.
+    """
+    root = override_path or configured_path
+    if not root:
+        raise ValueError("A single-file checkpoint does not carry its text encoder, so the "
+                         "encoder root must be declared. Set `gemma_model_path` in the "
+                         "pipeline config, or pass the text-encoder path explicitly for "
+                         "this run. It is not inferred from the checkpoint's directory: "
+                         "loading whichever encoder happens to sit beside the file would "
+                         "silently pick the wrong one when several are present.")
+
+    expected = (metadata.gemma_source_checkpoint or {}).get("gemma_version") if metadata is not None else None
+    if expected:
+        # ponytail: warn, don't raise -- the declared root is the user's
+        # explicit instruction and the pairing is advisory. Promote to a hard
+        # error if a mismatch ever turns out to produce silent garbage rather
+        # than an obvious shape failure.
+        actual = _read_encoder_version(root)
+        if actual is not None and actual != expected:
+            logger.warning(
+                "Checkpoint expects text-encoder version %r but the encoder at "
+                "%s reports %r; continuing with the declared root.", expected, root, actual)
+    return root
+
+
+def _read_encoder_version(root: str) -> str | None:
+    """The encoder's declared version, or None if it declares none."""
+    config_path = os.path.join(root, "config.json")
+    if not os.path.isfile(config_path):
+        return None
+    try:
+        with open(config_path, encoding="utf-8") as f:
+            return json.load(f).get("gemma_version")
+    except (OSError, json.JSONDecodeError):
+        return None
+
+
+def model_index_and_component_path(model_path: str, module_type: str) -> tuple[dict[str, Any], str]:
+    """``(model_index, component_path)`` for a directory repo OR a bundle.
+
+    The two differ in both halves: a repo answers "what components exist" from
+    ``model_index.json`` and puts each in its own subdirectory, while a bundle
+    declares its components in its own metadata and holds them all in one file.
+    Callers that resolve those two things together should route through here so
+    the bundle case is handled once instead of at every site.
+
+    ponytail: a bundle's text encoder and tokenizer live OUTSIDE the file, so
+    this returns the bundle path for them too, which is wrong for those two
+    module types. Training does not load them (text embeddings are
+    preprocessed), so it does not arise. Give this the resolved encoder root the
+    day a caller needs them.
+    """
+    if is_single_file_bundle(model_path):
+        return bundle_model_index(model_path), model_path
+    from fastvideo.utils import verify_model_config_and_directory
+    return verify_model_config_and_directory(model_path), os.path.join(model_path, module_type)

fastvideo/models/registry.py

@@ -23,6 +23,7 @@ logger = init_logger(__name__)
 
 # huggingface class name: (component_name, fastvideo module name, fastvideo class name)
 _TEXT_TO_VIDEO_DIT_MODELS = {
+    "MMAudioTransformer": ("dits", "mmaudio", "MMAudioTransformer"),
     "HunyuanVideoTransformer3DModel": ("dits", "hunyuanvideo", "HunyuanVideoTransformer3DModel"),
     "HunyuanGameCraftTransformer3DModel": ("dits", "hunyuangamecraft", "HunyuanGameCraftTransformer3DModel"),
     "HunyuanVideo15Transformer3DModel": ("dits", "hunyuanvideo15", "HunyuanVideo15Transformer3DModel"),
@@ -37,6 +38,8 @@ _TEXT_TO_VIDEO_DIT_MODELS = {
     ("dits", "longcat_video_dit", "LongCatVideoTransformer3DModel"),  # Wrapper (Phase 1)
     "LongCatTransformer3DModel": ("dits", "longcat", "LongCatTransformer3DModel"),  # Native (Phase 2)
     "LTX2Transformer3DModel": ("dits", "ltx2", "LTX2Transformer3DModel"),
+    # ``_class_name`` declared by single-file LTX checkpoint metadata.
+    "AVTransformer3DModel": ("dits", "ltx2", "LTX2Transformer3DModel"),
     "SD3Transformer2DModel": ("dits", "sd3", "SD3Transformer2DModel"),
     "LingBotWorldTransformer3DModel": ("dits", "lingbotworld", "LingBotWorldTransformer3DModel"),
     "LingBotWorld2CausalFastTransformer3DModel": (
@@ -75,6 +78,7 @@ _TEXT_TO_IMAGE_DIT_MODELS = {
 }
 
 _TEXT_ENCODER_MODELS = {
+    "MMAudioDFNCLIPTextEncoder": ("encoders", "mmaudio_clip", "MMAudioDFNCLIPTextEncoder"),
     "CLIPTextModel": ("encoders", "clip", "CLIPTextModel"),
     "CLIPTextModelWithProjection": ("encoders", "clip", "CLIPTextModelWithProjection"),
     "LlamaModel": ("encoders", "llama", "LlamaModel"),
@@ -94,6 +98,12 @@ _TEXT_ENCODER_MODELS = {
 }
 
 _IMAGE_ENCODER_MODELS: dict[str, tuple] = {
+    "MMAudioDFNCLIPVisionEncoder": ("encoders", "mmaudio_clip", "MMAudioDFNCLIPVisionEncoder"),
+    "MMAudioSynchformerVisualEncoder": (
+        "encoders",
+        "mmaudio_synchformer",
+        "MMAudioSynchformerVisualEncoder",
+    ),
     # "HunyuanVideoTransformer3DModel": ("image_encoder", "hunyuanvideo", "HunyuanVideoImageEncoder"),
     "CLIPVisionModelWithProjection": ("encoders", "clip", "CLIPVisionModel"),
     "CLIPVisionModel": ("encoders", "clip", "CLIPVisionModel"),
@@ -118,6 +128,8 @@ _VAE_MODELS = {
 }
 
 _AUDIO_MODELS = {
+    "MMAudioVAE": ("audio", "mmaudio_vae", "MMAudioVAE"),
+    "BigVGANV2": ("audio", "bigvgan", "BigVGANV2"),
     "LTX2AudioEncoder": ("audio", "ltx2_audio_vae", "LTX2AudioEncoder"),
     "LTX2AudioDecoder": ("audio", "ltx2_audio_vae", "LTX2AudioDecoder"),
     "LTX2Vocoder": ("audio", "ltx2_audio_vae", "LTX2Vocoder"),

fastvideo/pipelines/basic/ltx2/ltx2_pipeline.py

@@ -277,7 +277,7 @@ class LTX2Pipeline(LoRAPipeline):
                 modules[module_name] = loaded_modules[module_name]
                 continue
 
-            component_model_path = os.path.join(self.model_path, module_name)
+            component_model_path = self._component_path(module_name, fastvideo_args)
             if module_name == "tokenizer" and not os.path.isdir(component_model_path):
                 gemma_path = os.path.join(self.model_path, "text_encoder", "gemma")
                 if os.path.isdir(gemma_path):

fastvideo/pipelines/basic/ltx2/stages/ltx2_audio_decoding.py

@@ -49,6 +49,16 @@ class LTX2AudioDecodingStage(PipelineStage):
         audio_latents = batch.extra.get("ltx2_audio_latents")
         if audio_latents is None:
             return batch
+        # A checkpoint that declares no audio decoder builds these as None
+        # (``get_module`` returns its default), while the transformer is still
+        # audio-video and still produces audio latents. Having latents is
+        # therefore not evidence that anything can decode them -- guard on the
+        # modules too, and leave the video path unaffected.
+        if self.audio_decoder is None or self.vocoder is None:
+            logger.info(
+                "Skipping audio decoding: this checkpoint declares no audio "
+                "decoder/vocoder. Video output is unaffected.")
+            return batch
 
         device = get_local_torch_device()
         self.audio_decoder = self.audio_decoder.to(device)

fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_denoising.py

@@ -3,7 +3,6 @@
 
 from __future__ import annotations
 
-import contextlib
 from typing import Any
 
 import torch
@@ -12,7 +11,7 @@ from fastvideo.attention.selector import component_attention_backend, get_attn_b
 from fastvideo.distributed import get_local_torch_device
 from fastvideo.fastvideo_args import FastVideoArgs
 from fastvideo.forward_context import set_forward_context
-from fastvideo.profiler import get_global_controller
+from fastvideo.profiler import profiler_region
 from fastvideo.hooks.activation_trace import trace_step
 from fastvideo.pipelines.basic.minimax_h3.packing import (
     MINIMAX_H3_KEYFRAME_NOISE_AUG,
@@ -147,11 +146,8 @@ class MiniMaxH3DenoisingStage(PipelineStage):
             vsa_dense_layers = tuple(batch.extra.get("vsa_dense_layers", ()))
             vsa_dense_first_n = int(batch.extra.get("vsa_dense_first_n_steps", 0))
 
-        controller = get_global_controller()
-        denoise_region = (controller.region("profiler_region_inference_denoising")
-                          if controller is not None else contextlib.nullcontext())
         try:
-            with denoise_region:
+            with profiler_region("inference_denoising"):
                 for index, (video_timestep,
                             audio_timestep) in enumerate(zip(video_timesteps, audio_timesteps, strict=True)):
                     unique_timesteps, timestep_indices = row_timestep_plan[index]

fastvideo/pipelines/basic/mmaudio/README.md

@@ -0,0 +1,222 @@
+# MMAudio Video-to-Audio Inference
+
+This directory contains FastVideo's native `MMAudioPipeline` implementation for
+video-to-audio (V2A) and text-to-audio (T2A) generation. The first supported
+checkpoint is `large_44k_v2`, which produces mono audio at 44.1 kHz.
+
+The production pipeline is implemented with FastVideo components and does not
+import the upstream `mmaudio` Python package. It loads the transformer, DFN5B
+CLIP text/vision encoders, Synchformer visual encoder, audio VAE, BigVGAN-v2,
+tokenizer, and scheduler from a standard FastVideo/Diffusers-style checkpoint
+directory.
+
+## Status
+
+- `large_44k_v2` V2A inference: supported
+- T2A pipeline routing: supported; the provided example currently targets V2A
+- Output: mono WAV, 44.1 kHz
+- Default inference duration: 8 seconds
+- Variable-duration inference: supported
+- Single-GPU inference with FastVideo's default offloading: supported
+- Source-video/audio muxing: not yet part of the pipeline output
+- Training and small/medium/16 kHz variants: not included in this port
+
+The native pipeline has passed exact official-vs-FastVideo real-weight parity
+for a 25-step two-second waveform and a real ten-second V2A inference smoke
+test.
+
+## Requirements
+
+Follow FastVideo's main NVIDIA installation guide. The relevant baseline is:
+
+- Linux or Windows WSL
+- Python 3.10-3.12
+- CUDA 12.6 or CUDA 13.0
+- PyTorch 2.12.0
+- One NVIDIA GPU
+
+The port was validated with Python 3.12, PyTorch 2.12.0+cu126, and an RTX 6000
+Ada 48 GB. This is a validated configuration, not a minimum VRAM claim. Default
+layerwise/component offloading is enabled. FlashAttention is optional; the
+pipeline falls back to Torch SDPA when FlashAttention is unavailable.
+
+Install FastVideo from this source tree with `uv`:
+
+```bash
+cd FastVideo
+uv venv --python 3.12 --seed
+source .venv/bin/activate
+UV_TORCH_BACKEND=cu126 uv pip install -e .
+```
+
+Use `UV_TORCH_BACKEND=cu130` instead on CUDA 13. Conda is not required.
+
+## Checkpoint Layout
+
+Model weights are intentionally excluded from the FastVideo Git repository.
+For local development, the example expects a converted checkpoint at:
+
+```text
+converted_weights/mmaudio/large_44k_v2/
+├── model_index.json
+├── transformer/
+├── text_encoder/
+├── tokenizer/
+├── image_encoder/
+├── image_encoder_2/
+├── audio_vae/
+├── vocoder/
+└── scheduler/
+```
+
+`official_weights/`, `converted_weights/`, and inference outputs are ignored by
+Git. Cloning a code branch therefore does not clone the model.
+
+### Pre-converted checkpoint
+
+FastVideo accepts either a local directory or a Hugging Face model ID. When a
+complete converted checkpoint is published, select it with:
+
+```bash
+export MMAUDIO_MODEL_PATH=ORG/MMAudio-large-44k-v2-Diffusers
+```
+
+FastVideo will then download the complete snapshot on first use and reuse the
+Hugging Face cache on later runs. At the time of this port, the registered
+`FastVideo/MMAudio-large-44k-v2-Diffusers` name is reserved but is not yet a
+public checkpoint, so use the local conversion below.
+
+### Convert the official weights locally
+
+Only checkpoint conversion requires `open_clip_torch`; native FastVideo
+inference does not depend on the upstream MMAudio package.
+
+```bash
+uv pip install open_clip_torch
+
+mkdir -p official_weights/mmaudio/raw/weights
+mkdir -p official_weights/mmaudio/raw/ext_weights
+mkdir -p official_weights/mmaudio/DFN5B-CLIP-ViT-H-14-384
+mkdir -p official_weights/mmaudio/bigvgan_v2_44khz_128band_512x
+```
+
+Download the three MMAudio assets:
+
+```bash
+curl -L --continue-at - \
+  https://huggingface.co/hkchengrex/MMAudio/resolve/main/weights/mmaudio_large_44k_v2.pth \
+  -o official_weights/mmaudio/raw/weights/mmaudio_large_44k_v2.pth
+
+curl -L --continue-at - \
+  https://github.com/hkchengrex/MMAudio/releases/download/v0.1/v1-44.pth \
+  -o official_weights/mmaudio/raw/ext_weights/v1-44.pth
+
+curl -L --continue-at - \
+  https://github.com/hkchengrex/MMAudio/releases/download/v0.1/synchformer_state_dict.pth \
+  -o official_weights/mmaudio/raw/ext_weights/synchformer_state_dict.pth
+```
+
+Download only the DFN5B and BigVGAN files used by the converter:
+
+```bash
+python - <<'PY'
+from huggingface_hub import snapshot_download
+
+snapshot_download(
+    repo_id="apple/DFN5B-CLIP-ViT-H-14-384",
+    local_dir="official_weights/mmaudio/DFN5B-CLIP-ViT-H-14-384",
+    allow_patterns=["open_clip_config.json", "open_clip_pytorch_model.bin"],
+)
+snapshot_download(
+    repo_id="nvidia/bigvgan_v2_44khz_128band_512x",
+    local_dir="official_weights/mmaudio/bigvgan_v2_44khz_128band_512x",
+    allow_patterns=["config.json", "bigvgan_generator.pt"],
+)
+PY
+```
+
+Convert the source assets into the component tree consumed by FastVideo:
+
+```bash
+python scripts/checkpoint_conversion/convert_mmaudio_to_diffusers.py \
+  --transformer-checkpoint official_weights/mmaudio/raw/weights/mmaudio_large_44k_v2.pth \
+  --audio-vae-checkpoint official_weights/mmaudio/raw/ext_weights/v1-44.pth \
+  --synchformer-checkpoint official_weights/mmaudio/raw/ext_weights/synchformer_state_dict.pth \
+  --dfn5b-dir official_weights/mmaudio/DFN5B-CLIP-ViT-H-14-384 \
+  --bigvgan-dir official_weights/mmaudio/bigvgan_v2_44khz_128band_512x \
+  --output converted_weights/mmaudio/large_44k_v2
+```
+
+The converted checkpoint is approximately 9 GB.
+
+## Run V2A Inference
+
+The runnable example is
+[`examples/inference/basic/basic_mmaudio.py`](../../../../examples/inference/basic/basic_mmaudio.py).
+Run all commands from the FastVideo repository root.
+
+```bash
+export MMAUDIO_MODEL_PATH=converted_weights/mmaudio/large_44k_v2
+
+python examples/inference/basic/basic_mmaudio.py \
+  --video-path /path/to/input.mp4 \
+  --duration-seconds 8 \
+  --prompt "A skateboarder rolls over concrete and lands on a metal rail." \
+  --negative-prompt "music, speech" \
+  --output-path outputs_audio/mmaudio.wav
+```
+
+To select a GPU explicitly:
+
+```bash
+CUDA_VISIBLE_DEVICES=0 \
+MMAUDIO_MODEL_PATH=converted_weights/mmaudio/large_44k_v2 \
+python examples/inference/basic/basic_mmaudio.py \
+  --video-path /path/to/input.mp4 \
+  --duration-seconds 10 \
+  --output-path outputs_audio/mmaudio_10s.wav
+```
+
+The text prompt is optional, but a short description of audible events usually
+provides better control. The negative prompt can suppress unwanted categories
+such as music or speech.
+
+Eight seconds is the published training/default duration, not a hard inference
+limit. Shorter and longer clips use dynamic sequence lengths. As in the
+official MMAudio demo, quality may decrease when the requested duration is far
+from eight seconds. If the source video is shorter than the requested duration,
+the pipeline uses the available decoded duration. Synchformer requires at least
+16 frames at 25 FPS, so V2A input must cover at least 0.64 seconds.
+
+## Listen with the Source Video
+
+The pipeline currently writes a WAV file. To create a preview MP4 that replaces
+the source audio with the generated waveform:
+
+```bash
+ffmpeg -y \
+  -i /path/to/input.mp4 \
+  -i outputs_audio/mmaudio.wav \
+  -map 0:v:0 -map 1:a:0 \
+  -c:v copy -c:a aac -shortest \
+  outputs_audio/mmaudio_preview.mp4
+```
+
+## Troubleshooting
+
+- **Model download fails:** verify that `MMAUDIO_MODEL_PATH` is an existing
+  local converted directory or a public Hugging Face repository containing
+  `model_index.json` and all eight components.
+- **`FlashAttention-2 ... not found`:** this is informational. Torch SDPA is a
+  supported fallback and was used for exact parity validation.
+- **The original video already has audio:** V2A conditioning reads video frames
+  only. The source audio is not passed into MMAudio.
+- **No MP4 is returned:** the native result is audio-only by design; use the
+  `ffmpeg` command above for a preview mux.
+- **Do not commit checkpoints:** never force-add `official_weights/` or
+  `converted_weights/` to the FastVideo Git repository.
+
+## License and Attribution
+
+The upstream MMAudio checkpoint is distributed as CC-BY-NC 4.0. Preserve its
+license and attribution when redistributing raw or converted checkpoints.

fastvideo/pipelines/basic/mmaudio/__init__.py

@@ -0,0 +1,5 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from fastvideo.pipelines.basic.mmaudio.mmaudio_pipeline import MMAudioPipeline
+
+__all__ = ["MMAudioPipeline"]

fastvideo/pipelines/basic/mmaudio/mmaudio_pipeline.py

@@ -0,0 +1,63 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Native MMAudio video/text-to-audio pipeline."""
+
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.pipelines.basic.mmaudio.stages import (
+    MMAudioDecodingStage,
+    MMAudioDenoisingStage,
+    MMAudioInputValidationStage,
+    MMAudioLatentPreparationStage,
+    MMAudioTextConditioningStage,
+    MMAudioVideoConditioningStage,
+)
+from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase
+
+
+class MMAudioPipeline(ComposedPipelineBase):
+    """MMAudio large-44k-v2 V2A/T2A inference pipeline."""
+
+    _required_config_modules = [
+        "transformer",
+        "scheduler",
+        "text_encoder",
+        "tokenizer",
+        "image_encoder",
+        "image_encoder_2",
+        "audio_vae",
+        "vocoder",
+    ]
+
+    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None:
+        transformer = self.get_module("transformer")
+        self.add_stage("input_validation_stage", MMAudioInputValidationStage())
+        self.add_stage(
+            "video_conditioning_stage",
+            MMAudioVideoConditioningStage(
+                image_encoder=self.get_module("image_encoder"),
+                sync_encoder=self.get_module("image_encoder_2"),
+                transformer=transformer,
+            ),
+        )
+        self.add_stage(
+            "text_conditioning_stage",
+            MMAudioTextConditioningStage(
+                text_encoder=self.get_module("text_encoder"),
+                tokenizer=self.get_module("tokenizer"),
+                transformer=transformer,
+            ),
+        )
+        self.add_stage("latent_preparation_stage", MMAudioLatentPreparationStage(transformer))
+        self.add_stage(
+            "denoising_stage",
+            MMAudioDenoisingStage(transformer, self.get_module("scheduler")),
+        )
+        self.add_stage(
+            "decoding_stage",
+            MMAudioDecodingStage(
+                audio_vae=self.get_module("audio_vae"),
+                vocoder=self.get_module("vocoder"),
+            ),
+        )
+
+
+EntryClass = MMAudioPipeline

fastvideo/pipelines/basic/mmaudio/presets.py

@@ -0,0 +1,40 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Published MMAudio large-44k-v2 inference defaults."""
+
+from fastvideo.api.presets import InferencePreset, PresetStageSpec
+
+_DENOISE_STAGE = PresetStageSpec(
+    name="denoise",
+    kind="denoising",
+    description="MMAudio forward-time Euler flow with multimodal CFG.",
+    allowed_overrides=frozenset({"num_inference_steps", "guidance_scale"}),
+)
+
+_DEFAULTS = {
+    "seed": 42,
+    "guidance_scale": 4.5,
+    "num_inference_steps": 25,
+    "negative_prompt": "",
+    "audio_start_in_s": 0.0,
+    "audio_end_in_s": 8.0,
+    # The shared generator still validates these fields, but MMAudio returns
+    # audio metadata rather than materializing the placeholder pixels.
+    "height": 8,
+    "width": 8,
+    "num_frames": 1,
+    "fps": 25,
+    "return_frames": False,
+}
+
+MMAUDIO_LARGE_44K_V2 = InferencePreset(
+    name="mmaudio_large_44k_v2",
+    version=1,
+    model_family="mmaudio",
+    description=("MMAudio large-44k-v2 video-to-audio generation with DFN5B CLIP, "
+                 "Synchformer, a 44.1 kHz audio VAE, and BigVGAN-v2."),
+    workload_type="v2a",
+    stage_schemas=(_DENOISE_STAGE, ),
+    defaults=dict(_DEFAULTS),
+)
+
+ALL_PRESETS = (MMAUDIO_LARGE_44K_V2, )

fastvideo/pipelines/basic/mmaudio/stages.py

@@ -0,0 +1,455 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Inference stages for MMAudio V2A/T2A.
+
+The video sampler intentionally mirrors ``mmaudio.data.av_utils.read_frames``:
+timestamps are sampled independently at 8 FPS and 25 FPS, and a decoded frame
+is repeated when the source FPS is lower than a requested sampling rate.
+"""
+
+from __future__ import annotations
+
+import math
+from pathlib import Path
+
+import numpy as np
+import torch
+from torchvision.transforms import v2
+
+from fastvideo.distributed import get_local_torch_device
+from fastvideo.fastvideo_args import FastVideoArgs, WorkloadType
+from fastvideo.forward_context import set_forward_context
+from fastvideo.logger import init_logger
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+from fastvideo.pipelines.stages.base import PipelineStage
+from fastvideo.pipelines.stages.validators import VerificationResult
+
+logger = init_logger(__name__)
+
+_CLIP_MEAN = (0.48145466, 0.4578275, 0.40821073)
+_CLIP_STD = (0.26862954, 0.26130258, 0.27577711)
+
+
+def _read_frames_at_fps(
+    video_path: str | Path,
+    frame_rates: tuple[float, ...],
+    *,
+    start_s: float,
+    end_s: float,
+) -> list[np.ndarray]:
+    """Decode RGB frames with MMAudio's timestamp/duplication contract."""
+    import av
+
+    outputs: list[list[np.ndarray]] = [[] for _ in frame_rates]
+    next_times = [0.0 for _ in frame_rates]
+    deltas = [1.0 / fps for fps in frame_rates]
+
+    with av.open(str(video_path)) as container:
+        stream = container.streams.video[0]
+        stream.thread_type = "AUTO"
+        for packet in container.demux(stream):
+            for frame in packet.decode():
+                frame_time = frame.time
+                if frame_time is None or frame_time < start_s:
+                    continue
+                if frame_time > end_s:
+                    break
+
+                frame_array = None
+                for index in range(len(frame_rates)):
+                    while frame_time >= next_times[index]:
+                        if frame_array is None:
+                            frame_array = frame.to_ndarray(format="rgb24")
+                        outputs[index].append(frame_array)
+                        next_times[index] += deltas[index]
+
+    if any(not frames for frames in outputs):
+        raise ValueError(f"Could not decode enough video frames from {video_path} in "
+                         f"[{start_s}, {end_s}] seconds.")
+    return [np.stack(frames) for frames in outputs]
+
+
+def preprocess_mmaudio_video(
+    video_path: str | Path,
+    *,
+    duration_s: float,
+    clip_fps: int = 8,
+    sync_fps: int = 25,
+    clip_size: int = 384,
+    sync_size: int = 224,
+) -> tuple[torch.Tensor, torch.Tensor, float]:
+    """Return official-format CLIP frames, sync frames, and effective duration.
+
+    CLIP output is float32 ``[T,3,384,384]`` in ``[0,1]``. Synchformer
+    output is float32 ``[T,3,224,224]`` in ``[-1,1]``.
+    """
+    clip_array, sync_array = _read_frames_at_fps(
+        video_path,
+        (float(clip_fps), float(sync_fps)),
+        start_s=0.0,
+        end_s=duration_s,
+    )
+    clip_frames = torch.from_numpy(clip_array).permute(0, 3, 1, 2)
+    sync_frames = torch.from_numpy(sync_array).permute(0, 3, 1, 2)
+
+    clip_transform = v2.Compose([
+        v2.Resize((clip_size, clip_size), interpolation=v2.InterpolationMode.BICUBIC),
+        v2.ToImage(),
+        v2.ToDtype(torch.float32, scale=True),
+    ])
+    sync_transform = v2.Compose([
+        v2.Resize(sync_size, interpolation=v2.InterpolationMode.BICUBIC),
+        v2.CenterCrop(sync_size),
+        v2.ToImage(),
+        v2.ToDtype(torch.float32, scale=True),
+        v2.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
+    ])
+    clip_frames = clip_transform(clip_frames)
+    sync_frames = sync_transform(sync_frames)
+
+    effective_duration = min(
+        duration_s,
+        clip_frames.shape[0] / clip_fps,
+        sync_frames.shape[0] / sync_fps,
+    )
+    clip_frames = clip_frames[:int(clip_fps * effective_duration)]
+    sync_frames = sync_frames[:int(sync_fps * effective_duration)]
+    return clip_frames, sync_frames, effective_duration
+
+
+def mmaudio_sequence_lengths(duration_s: float, pc) -> tuple[int, int, int]:
+    latent_length = math.ceil(duration_s * pc.sampling_rate / pc.spectrogram_frame_rate / pc.latent_downsample_rate)
+    clip_length = int(duration_s * pc.clip_frame_rate)
+    sync_frame_count = duration_s * pc.sync_frame_rate
+    sync_segments = ((sync_frame_count - pc.sync_segment_size) // pc.sync_segment_stride + 1)
+    sync_length = int(sync_segments * pc.sync_segment_size / pc.sync_downsample_rate)
+    return latent_length, clip_length, sync_length
+
+
+class MMAudioInputValidationStage(PipelineStage):
+    """Validate audio-generation inputs without invoking video pipeline logic."""
+
+    def verify_input(self, batch, fastvideo_args):
+        return VerificationResult()
+
+    def verify_output(self, batch, fastvideo_args):
+        return VerificationResult()
+
+    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
+        if batch.num_inference_steps <= 0:
+            raise ValueError("MMAudio num_inference_steps must be positive.")
+        if batch.num_videos_per_prompt != 1:
+            raise ValueError("MMAudio currently supports one output per request.")
+        if isinstance(batch.prompt, list) and len(batch.prompt) != 1:
+            raise ValueError("MMAudio currently supports one prompt per request.")
+        if isinstance(batch.negative_prompt, list) and len(batch.negative_prompt) != 1:
+            raise ValueError("MMAudio currently supports one negative prompt per request.")
+
+        direct_video = (batch.extra.get("mmaudio_clip_frames") is not None
+                        and batch.extra.get("mmaudio_sync_frames") is not None)
+        if (fastvideo_args.workload_type is WorkloadType.V2A and batch.video_path is None and not direct_video):
+            raise ValueError("MMAudio V2A requires `video_path` or preprocessed MMAudio frame tensors.")
+
+        pc = fastvideo_args.pipeline_config
+        duration_s = batch.audio_end_in_s
+        duration_s = pc.duration_s if duration_s is None else float(duration_s)
+        start_s = 0.0 if batch.audio_start_in_s is None else float(batch.audio_start_in_s)
+        if start_s != 0.0:
+            raise ValueError("MMAudio currently generates from time zero; audio_start_in_s must be 0.")
+        max_duration_s = pc.max_audio_duration_s
+        if max_duration_s is not None and duration_s > max_duration_s:
+            raise ValueError(f"MMAudio duration {duration_s}s exceeds this checkpoint's "
+                             f"{max_duration_s}s maximum.")
+        minimum_duration = pc.sync_segment_size / pc.sync_frame_rate
+        if duration_s < minimum_duration:
+            raise ValueError(f"MMAudio needs at least {minimum_duration:.2f}s "
+                             f"({pc.sync_segment_size} sync frames), got {duration_s}s.")
+        batch.extra["mmaudio_duration_s"] = duration_s
+        batch.seed = 0 if batch.seed is None else int(batch.seed)
+        return batch
+
+
+class MMAudioVideoConditioningStage(PipelineStage):
+    """Decode video and run DFN5B/Synchformer with official preprocessing."""
+
+    def __init__(self, image_encoder, sync_encoder, transformer) -> None:
+        super().__init__()
+        self.image_encoder = image_encoder
+        self.sync_encoder = sync_encoder
+        self.transformer = transformer
+
+    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
+        duration_s = float(batch.extra["mmaudio_duration_s"])
+        clip_frames = batch.extra.get("mmaudio_clip_frames")
+        sync_frames = batch.extra.get("mmaudio_sync_frames")
+        use_video = batch.video_path is not None or (clip_frames is not None and sync_frames is not None)
+
+        if batch.video_path is not None:
+            clip_frames, sync_frames, duration_s = preprocess_mmaudio_video(
+                batch.video_path,
+                duration_s=duration_s,
+                clip_fps=pc.clip_frame_rate,
+                sync_fps=pc.sync_frame_rate,
+                clip_size=pc.clip_image_size,
+                sync_size=pc.sync_image_size,
+            )
+        elif use_video:
+            if clip_frames.ndim == 5:
+                if clip_frames.shape[0] != 1:
+                    raise ValueError("MMAudio preprocessed CLIP frames must have batch size 1.")
+                clip_frames = clip_frames[0]
+            if sync_frames.ndim == 5:
+                if sync_frames.shape[0] != 1:
+                    raise ValueError("MMAudio preprocessed sync frames must have batch size 1.")
+                sync_frames = sync_frames[0]
+            duration_s = min(
+                duration_s,
+                clip_frames.shape[0] / pc.clip_frame_rate,
+                sync_frames.shape[0] / pc.sync_frame_rate,
+            )
+            clip_frames = clip_frames[:int(duration_s * pc.clip_frame_rate)]
+            sync_frames = sync_frames[:int(duration_s * pc.sync_frame_rate)]
+
+        latent_length, clip_length, sync_length = mmaudio_sequence_lengths(duration_s, pc)
+        if clip_length <= 0 or sync_length <= 0:
+            raise ValueError(f"MMAudio duration {duration_s}s produces an empty condition sequence.")
+        self.transformer.update_seq_lengths(latent_length, clip_length, sync_length)
+        batch.extra["mmaudio_duration_s"] = duration_s
+        batch.extra["mmaudio_sequence_lengths"] = (latent_length, clip_length, sync_length)
+
+        if not use_video:
+            batch.extra["mmaudio_clip_features"] = self.transformer.get_empty_clip_sequence(1)
+            batch.extra["mmaudio_sync_features"] = self.transformer.get_empty_sync_sequence(1)
+            return batch
+
+        assert clip_frames is not None and sync_frames is not None
+        if clip_frames.shape != (clip_length, 3, pc.clip_image_size, pc.clip_image_size):
+            raise ValueError(
+                "MMAudio CLIP frames must have shape "
+                f"[{clip_length},3,{pc.clip_image_size},{pc.clip_image_size}], got {tuple(clip_frames.shape)}.")
+        expected_sync_frames = int(duration_s * pc.sync_frame_rate)
+        if sync_frames.shape != (expected_sync_frames, 3, pc.sync_image_size, pc.sync_image_size):
+            raise ValueError("MMAudio sync frames must have shape "
+                             f"[{expected_sync_frames},3,{pc.sync_image_size},{pc.sync_image_size}], "
+                             f"got {tuple(sync_frames.shape)}.")
+
+        device = get_local_torch_device()
+        model_dtype = next(self.transformer.parameters()).dtype
+
+        self.image_encoder = self.image_encoder.to(device)
+        clip_video = clip_frames.to(device=device, dtype=model_dtype, non_blocking=True)
+        mean = torch.tensor(_CLIP_MEAN, device=device, dtype=model_dtype).view(1, 3, 1, 1)
+        std = torch.tensor(_CLIP_STD, device=device, dtype=model_dtype).view(1, 3, 1, 1)
+        clip_video = (clip_video - mean) / std
+        clip_outputs: list[torch.Tensor] = []
+        chunk_size = pc.clip_batch_size_multiplier
+        with set_forward_context(current_timestep=0, attn_metadata=None):
+            for start in range(0, clip_length, chunk_size):
+                encoded = self.image_encoder(clip_video[start:start + chunk_size]).last_hidden_state
+                clip_outputs.append(encoded)
+        clip_features = torch.cat(clip_outputs, dim=0).unsqueeze(0)
+        if fastvideo_args.image_encoder_cpu_offload:
+            self.image_encoder = self.image_encoder.to("cpu")
+
+        self.sync_encoder = self.sync_encoder.to(device)
+        sync_video = sync_frames.to(device=device, dtype=model_dtype, non_blocking=True).unsqueeze(0)
+        with set_forward_context(current_timestep=0, attn_metadata=None):
+            sync_features = self.sync_encoder(sync_video).last_hidden_state
+        if fastvideo_args.image_encoder_cpu_offload:
+            self.sync_encoder = self.sync_encoder.to("cpu")
+
+        if sync_features.shape[1] != sync_length:
+            raise RuntimeError(f"Synchformer produced {sync_features.shape[1]} tokens; expected {sync_length}.")
+        batch.extra["mmaudio_clip_features"] = clip_features
+        batch.extra["mmaudio_sync_features"] = sync_features
+        return batch
+
+
+class MMAudioTextConditioningStage(PipelineStage):
+    """Encode positive/negative OpenCLIP token sequences and project conditions."""
+
+    def __init__(self, text_encoder, tokenizer, transformer) -> None:
+        super().__init__()
+        self.text_encoder = text_encoder
+        self.tokenizer = tokenizer
+        self.transformer = transformer
+
+    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:
+        device = get_local_torch_device()
+        model_dtype = next(self.transformer.parameters()).dtype
+        self.text_encoder = self.text_encoder.to(device)
+
+        def encode(text: str | list[str] | None) -> torch.Tensor | None:
+            if text is None:
+                return None
+            texts = [text] if isinstance(text, str) else text
+            tokens = self.tokenizer(
+                texts,
+                padding="max_length",
+                truncation=True,
+                max_length=77,
+                return_tensors="pt",
+            )
+            # Hugging Face CLIPTokenizer pads with EOS by default, while
+            # OpenCLIP/MMAudio pads with token id zero.
+            input_ids = tokens.input_ids.masked_fill(tokens.attention_mask == 0, 0).to(device)
+            with set_forward_context(current_timestep=0, attn_metadata=None):
+                return self.text_encoder(input_ids).last_hidden_state.to(model_dtype)
+
+        text_features = encode(batch.prompt)
+        if text_features is None:
+            text_features = self.transformer.get_empty_string_sequence(1)
+        negative_text_features = encode(batch.negative_prompt)
+
+        if fastvideo_args.text_encoder_cpu_offload:
+            self.text_encoder = self.text_encoder.to("cpu")
+
+        conditions = self.transformer.preprocess_conditions(
+            batch.extra["mmaudio_clip_features"],
+            batch.extra["mmaudio_sync_features"],
+            text_features,
+        )
+        empty_conditions = self.transformer.get_empty_conditions(
+            1,
+            negative_text_features=negative_text_features,
+        )
+        batch.extra["mmaudio_conditions"] = conditions
+        batch.extra["mmaudio_empty_conditions"] = empty_conditions
+        return batch
+
+
+class MMAudioLatentPreparationStage(PipelineStage):
+    """Sample the Gaussian flow prior using a device-local seeded generator."""
+
+    def __init__(self, transformer) -> None:
+        super().__init__()
+        self.transformer = transformer
+
+    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:
+        if batch.latents is not None:
+            return batch
+        device = get_local_torch_device()
+        dtype = next(self.transformer.parameters()).dtype
+        generator = torch.Generator(device=device).manual_seed(int(batch.seed))
+        batch.generator = generator
+        batch.latents = torch.randn(
+            (1, self.transformer.latent_seq_len, self.transformer.latent_dim),
+            device=device,
+            dtype=dtype,
+            generator=generator,
+        )
+        return batch
+
+
+class MMAudioDenoisingStage(PipelineStage):
+    """Run MMAudio's forward-time Euler flow with FastVideo's shared scheduler."""
+
+    def __init__(self, transformer, scheduler) -> None:
+        super().__init__()
+        self.transformer = transformer
+        self.scheduler = scheduler
+
+    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:
+        assert batch.latents is not None
+        # Official MMAudio builds ``torch.linspace`` on CPU. A CPU scalar in
+        # the bf16 CUDA update follows scalar-promotion rules, whereas moving
+        # the same float32 scalar to CUDA changes the rounded trajectory.
+        self.scheduler.set_timesteps(batch.num_inference_steps, device="cpu")
+        conditions = batch.extra["mmaudio_conditions"]
+        empty_conditions = batch.extra["mmaudio_empty_conditions"]
+        latents = batch.latents
+        for index, timestep in enumerate(self.scheduler.timesteps):
+            flow = self.transformer.guided_flow(
+                timestep / self.scheduler.config.num_train_timesteps,
+                latents,
+                conditions,
+                empty_conditions,
+                float(batch.guidance_scale),
+            )
+            # The shared FlowMatch scheduler supplies the exact inverted
+            # forward-time sigma schedule. Its generic ``step`` deliberately
+            # upcasts samples to fp32, however, while MMAudio's published
+            # Euler loop performs ``x += dt * flow`` in the model's bf16
+            # dtype. Preserve that operation order/precision here; changing
+            # it shifts the full 25-step trajectory.
+            delta = self.scheduler.sigmas[index + 1] - self.scheduler.sigmas[index]
+            latents = latents + delta * flow
+            batch.step_index = index
+            batch.timestep = timestep
+        batch.latents = self.transformer.unnormalize(latents)
+        return batch
+
+
+class MMAudioDecodingStage(PipelineStage):
+    """Decode MMAudio latents to a mono 44.1 kHz waveform."""
+
+    def __init__(self, audio_vae, vocoder) -> None:
+        super().__init__()
+        self.audio_vae = audio_vae
+        self.vocoder = vocoder
+
+    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:
+        assert batch.latents is not None
+        if fastvideo_args.output_type == "latent":
+            batch.output = batch.latents.detach().cpu()
+            return batch
+
+        device = get_local_torch_device()
+        self.audio_vae = self.audio_vae.to(device)
+        self.vocoder = self.vocoder.to(device)
+        decoder_dtype = next(self.audio_vae.parameters()).dtype
+        mel = self.audio_vae.decode(batch.latents.transpose(1, 2).to(decoder_dtype))
+        audio = self.vocoder(mel.to(next(self.vocoder.parameters()).dtype))
+
+        pc = fastvideo_args.pipeline_config
+        expected_samples = batch.extra["mmaudio_sequence_lengths"][0]
+        expected_samples *= pc.spectrogram_frame_rate * pc.latent_downsample_rate
+        if audio.shape[-1] != expected_samples:
+            raise RuntimeError(f"MMAudio vocoder produced {audio.shape[-1]} samples; expected {expected_samples}.")
+
+        decoded_audio = audio.detach().float().cpu()
+        batch.extra["audio"] = decoded_audio[0].T.contiguous().numpy()
+        batch.extra["audio_sample_rate"] = int(pc.sampling_rate)
+        batch.extra["audio_only"] = True
+        batch.extra["decoded_audio"] = decoded_audio
+        batch.data_type = "audio"
+
+        # VideoGenerator currently transports a tensor for every workload;
+        # audio-only paths never materialize or save these pixels.
+        batch.output = torch.zeros((1, 3, 1, 8, 8), dtype=torch.uint8)
+        return batch

fastvideo/pipelines/composed_pipeline_base.py

@@ -20,6 +20,8 @@ from fastvideo.hooks.activation_trace import attach_activation_trace, detach_act
 from fastvideo.logger import init_logger
 from fastvideo.profiler import get_or_create_profiler
 from fastvideo.models.loader.component_loader import PipelineComponentLoader
+from fastvideo.models.loader.ltx_single_file import (bundle_model_index, is_single_file_bundle, read_ltx_metadata,
+                                                     resolve_text_encoder_root)
 from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
 from fastvideo.pipelines.stages import PipelineStage
 import fastvideo.envs as envs
@@ -79,6 +81,7 @@ class ComposedPipelineBase(ABC):
 
         self.model_path: str = model_path
         self._stages: list[PipelineStage] = []
+        self._bundle_encoder_root: str | None = None
         self._stage_name_mapping: dict[str, PipelineStage] = {}
         self._trace_mgr = None
 
@@ -94,7 +97,6 @@ class ComposedPipelineBase(ABC):
         # FASTVIDEO_TORCH_PROFILER_DIR=/path/to/save/trace
         trace_dir = envs.FASTVIDEO_TORCH_PROFILER_DIR
         self.profiler_controller = get_or_create_profiler(trace_dir)
-        self.profiler = self.profiler_controller.profiler
 
         self.local_rank = get_world_group().local_rank
 
@@ -329,12 +331,50 @@ class ComposedPipelineBase(ABC):
         self.model_path = model_path
         # fastvideo_args.downloaded_model_path = model_path
         logger.info("Model path: %s", model_path)
+        if is_single_file_bundle(model_path):
+            # A bundle is a file, so there is no directory to verify and no
+            # model_index.json to read; it declares the same thing itself.
+            return bundle_model_index(model_path)
         config = verify_model_config_and_directory(
             model_path,
             required_component_dirs=self.get_hf_download_component_dirs(),
         )
         return cast(dict[str, Any], config)
 
+    def _text_encoder_root(self, fastvideo_args: FastVideoArgs) -> str:
+        """Where a bundle's text encoder lives -- it is not in the bundle.
+
+        Declared, never discovered: a per-run environment override first, then
+        the encoder path the pipeline config already carries.
+        """
+        encoder_configs = getattr(fastvideo_args.pipeline_config, "text_encoder_configs", ())
+        return resolve_text_encoder_root(
+            configured_path=(getattr(encoder_configs[0].arch_config, "gemma_model_path", None)
+                             if encoder_configs else None),
+            override_path=os.environ.get("FASTVIDEO_LTX_ENCODER_ROOT") or None,
+            metadata=read_ltx_metadata(self.model_path),
+        )
+
+    def _component_path(self, module_name: str, fastvideo_args: FastVideoArgs) -> str:
+        """Where one component's config and weights live.
+
+        A directory repo puts each component in its own subdirectory. A bundle
+        holds all of them in one file except the text stack, which is declared
+        separately. Subclasses that override ``load_modules`` must route
+        through here rather than joining the path themselves, or a bundle
+        silently falls back to the directory layout and fails looking for a
+        subdirectory of a file.
+        """
+        if not is_single_file_bundle(self.model_path):
+            return os.path.join(self.model_path, module_name)
+        if module_name in ("text_encoder", "tokenizer"):
+            # Resolved once: it raises when undeclared, and doing that per
+            # component would report the same failure several times.
+            if self._bundle_encoder_root is None:
+                self._bundle_encoder_root = self._text_encoder_root(fastvideo_args)
+            return self._bundle_encoder_root
+        return self.model_path
+
     @property
     def required_config_modules(self) -> list[str]:
         """
@@ -466,7 +506,7 @@ class ComposedPipelineBase(ABC):
             else:
                 load_module_name = module_name
 
-            component_model_path = os.path.join(self.model_path, load_module_name)
+            component_model_path = self._component_path(load_module_name, fastvideo_args)
             module = PipelineComponentLoader.load_module(
                 module_name=load_module_name,
                 component_model_path=component_model_path,
@@ -498,17 +538,6 @@ class ComposedPipelineBase(ABC):
         self._stage_name_mapping[stage_name] = stage
         setattr(self, stage_name, stage)
 
-    def profile(self, is_start: bool = True):
-        if self.profiler is None:
-            raise RuntimeError("Profiler is not enabled.")
-        if is_start:
-            self.profiler.start()
-        else:
-            self.profiler.stop()
-            # only print profiler results on rank 0
-            if self.local_rank == 0:
-                print(self.profiler.key_averages().table(sort_by="self_cuda_time_total"))
-
     # TODO(will): don't hardcode no_grad
     @torch.no_grad()
     def forward(

fastvideo/profiler.py

@@ -5,8 +5,9 @@ The profiler is shared across the process; this module adds a light-weight
 controller that gates collection based on named *regions*. Regions may be
 enabled through dedicated environment variables (e.g.
 ``FASTVIDEO_TORCH_PROFILE_MODEL_LOADING=1``) or via the consolidated
-``FASTVIDEO_TORCH_PROFILE_REGIONS`` comma-separated list (e.g.
-``FASTVIDEO_TORCH_PROFILE_REGIONS=model_loading,training_dit``).
+``FASTVIDEO_TORCH_PROFILE_REGIONS`` comma-separated list. Short names work
+(``FASTVIDEO_TORCH_PROFILE_REGIONS=model_loading,training_train`` resolves the
+``profiler_region_`` prefix automatically).
 
 Typical usage from client code::
 
@@ -34,7 +35,6 @@ from fastvideo.logger import init_logger
 
 logger = init_logger(__name__)
 
-_GLOBAL_PROFILER: torch.profiler.profile | None = None
 _GLOBAL_CONTROLLER: TorchProfilerController | None = None
 
 
@@ -44,7 +44,6 @@ class ProfilerRegion:
 
     name: str
     description: str
-    default_enabled: bool = False
 
     def __post_init__(self) -> None:
         if not self.name or self.name.strip() != self.name:
@@ -63,8 +62,6 @@ def _normalize_token(token: str) -> str:
 def register_profiler_region(
     name: str,
     description: str,
-    *,
-    default_enabled: bool = False,
 ) -> None:
     """Register a profiler region so configuration can validate inputs."""
 
@@ -75,16 +72,22 @@ def register_profiler_region(
     region = ProfilerRegion(
         name=canonical,
         description=description,
-        default_enabled=bool(default_enabled),
     )
     _REGISTERED_REGIONS[canonical] = region
 
 
 def resolve_profiler_region(name: str) -> ProfilerRegion | None:
-    """Return the registered region matching ``name`` or ``None`` if absent."""
+    """Return the registered region for ``name`` (long or short form).
+
+    Accepts both the canonical name and the short form without the
+    ``profiler_region_`` prefix, so ``REGIONS=inference_denoising`` works.
+    """
 
     canonical = _normalize_token(name)
-    return _REGISTERED_REGIONS.get(canonical)
+    region = _REGISTERED_REGIONS.get(canonical)
+    if region is None and not canonical.startswith("profiler_region_"):
+        region = _REGISTERED_REGIONS.get(f"profiler_region_{canonical}")
+    return region
 
 
 def list_profiler_regions() -> list[ProfilerRegion]:
@@ -99,17 +102,6 @@ _DEFAULT_ACTIVITIES: tuple[torch.profiler.ProfilerActivity, ...] = (
 )
 
 
-def get_global_profiler() -> torch.profiler.profile | None:
-    """Return the global profiler instance if one was created."""
-
-    return _GLOBAL_PROFILER
-
-
-def set_global_profiler(profiler: torch.profiler.profile | None) -> None:
-    global _GLOBAL_PROFILER
-    _GLOBAL_PROFILER = profiler
-
-
 def get_global_controller() -> TorchProfilerController | None:
     return _GLOBAL_CONTROLLER
 
@@ -122,7 +114,6 @@ def set_global_controller(controller: TorchProfilerController | None) -> None:
 register_profiler_region(
     name="profiler_region_model_loading",
     description="Module/model loading during pipeline initialization.",
-    default_enabled=False,
 )
 # register_profiler_region(
 #     name="profiler_region_inference_pre_denoising",
@@ -149,11 +140,11 @@ register_profiler_region(
 )
 register_profiler_region(
     name="profiler_region_training_train_one_step",
-    description="High-level step orchestration in the training loop.",
+    description="Single optimizer step including forward/backward passes.",
 )
 register_profiler_region(
     name="profiler_region_training_train",
-    description="Single optimizer step including forward/backward passes.",
+    description="High-level step orchestration in the training loop.",
 )
 
 # distillation specific regions
@@ -204,15 +195,18 @@ def get_or_create_profiler(trace_dir: str | None) -> TorchProfilerController:
         profile_memory=envs.FASTVIDEO_TORCH_PROFILER_WITH_PROFILE_MEMORY,
         with_stack=envs.FASTVIDEO_TORCH_PROFILER_WITH_STACK,
         with_flops=envs.FASTVIDEO_TORCH_PROFILER_WITH_FLOPS,
-        schedule=torch.profiler.schedule(
-            wait=envs.FASTVIDEO_TORCH_PROFILER_WAIT_STEPS,
-            warmup=envs.FASTVIDEO_TORCH_PROFILER_WARMUP_STEPS,
-            active=envs.FASTVIDEO_TORCH_PROFILER_ACTIVE_STEPS,
-        ),
+        # No schedule: nothing in the codebase calls profiler.step(), so a
+        # wait/warmup schedule never advances and the profiler records nothing.
+        # Region toggling gates collection; the single trace exports at stop().
         on_trace_ready=torch.profiler.tensorboard_trace_handler(trace_dir, use_gzip=True),
     )
     controller = TorchProfilerController(profiler, _DEFAULT_ACTIVITIES)
+    controller._trace_dir = trace_dir
     controller.start()
+    # The trace only exports at stop(); inference paths have no shutdown hook
+    # that calls it, so register one. stop() is idempotent.
+    import atexit
+    atexit.register(controller.stop)
     logger.info("Torch profiler started")
     return controller
 
@@ -287,16 +281,16 @@ class TorchProfilerController:
         FASTVIDEO_TORCH_PROFILE_REGIONS=model_loading,training_dit \
         python fastvideo/training/wan_training_pipeline.py ...
 
-    Wrapping a code block in a custom region::
+    Wrapping a code block in a registered region::
 
-        controller = TorchProfilerController(profiler, activities)
-        with controller.region("training_validation"):
-            run_validation_epoch()
+        from fastvideo.profiler import profiler_region
 
-    Adding a new region requires three steps:
-      1. Define an env var in ``envs.py``.
-      2. Add a default entry to ``register_profiler_region`` in this module.
-      3. Wrap the target code in :meth:`region` using the new name.
+        with profiler_region("inference_denoising"):
+            run_denoising_loop()
+
+    Adding a new region requires two steps:
+      1. Register it with ``register_profiler_region`` in this module.
+      2. Wrap the target code in :func:`profiler_region` using the new name.
     """
 
     def __init__(
@@ -317,10 +311,13 @@ class TorchProfilerController:
         self._profiler = profiler
         self._activities = activities_tuple
         self._config = config or TorchProfilerConfig.from_env()
-        self._collection_enabled = False
+        # torch.profiler collects from start(); reflect that so the initial
+        # _set_collection(False) in start() actually toggles it off instead of
+        # short-circuiting (which captured everything before the first region).
+        self._collection_enabled = True
         self._active_region_depth = 0
+        self._trace_dir: str | None = None
         logger.info("PROFILER: TorchProfilerController initialized with config: %s", self._config)
-        set_global_profiler(self._profiler)
         set_global_controller(self)
 
     @property
@@ -336,7 +333,10 @@ class TorchProfilerController:
 
         if self._profiler is None:
             return False
-        return self._config.regions.get(region, False)
+        resolved = resolve_profiler_region(region)
+        if resolved is None:
+            return False
+        return self._config.regions.get(resolved.name, False)
 
     def _set_collection(self, enabled: bool) -> None:
         if self._profiler is None:
@@ -348,6 +348,8 @@ class TorchProfilerController:
             self._profiler.toggle_collection_dynamic(enabled, self._activities)
         self._collection_enabled = enabled
 
+    _warned_unregistered: set[str] = set()
+
     @contextlib.contextmanager
     def region(self, region: str):
         """Context manager that enables profiling for ``region`` if configured."""
@@ -356,24 +358,41 @@ class TorchProfilerController:
             yield
             return
 
+        if resolve_profiler_region(region) is None:
+            # a typo here would otherwise silently profile nothing, forever
+            if region not in self._warned_unregistered:
+                self._warned_unregistered.add(region)
+                logger.warning("PROFILER: region %r is not registered (typo?); available: %s", region,
+                               ", ".join(r.name for r in list_profiler_regions()))
+            yield
+            return
+
         if not self.is_region_enabled(region):
             yield
             return
 
-        with torch.profiler.record_function(f"fastvideo.region::{region}"):
-            self._active_region_depth += 1
-            if self._active_region_depth == 1:
-                logger.info("PROFILER: Setting collection to True (depth=%s) for region %s", self._active_region_depth,
-                            region)
-                self._set_collection(True)
-            try:
+        # NVTX range so the same region names are visible in nsys timelines
+        nvtx = torch.cuda.is_available()
+        if nvtx:
+            torch.cuda.nvtx.range_push(f"fastvideo.region::{region}")
+        self._active_region_depth += 1
+        if self._active_region_depth == 1:
+            logger.info("PROFILER: Setting collection to True (depth=%s) for region %s", self._active_region_depth,
+                        region)
+            self._set_collection(True)
+        try:
+            # record_function opens after collection is enabled so the region
+            # marker itself lands in the trace.
+            with torch.profiler.record_function(f"fastvideo.region::{region}"):
                 yield
-            finally:
-                self._active_region_depth -= 1
-                logger.info("PROFILER: Decreasing active region depth to %s", self._active_region_depth)
-                if self._active_region_depth == 0:
-                    logger.info("PROFILER: Setting collection to False upon exiting region %s", region)
-                    self._set_collection(False)
+        finally:
+            self._active_region_depth -= 1
+            logger.info("PROFILER: Decreasing active region depth to %s", self._active_region_depth)
+            if self._active_region_depth == 0:
+                logger.info("PROFILER: Setting collection to False upon exiting region %s", region)
+                self._set_collection(False)
+            if nvtx:
+                torch.cuda.nvtx.range_pop()
 
     def start(self) -> None:
         """Start the profiler and pause collection until a region is entered."""
@@ -388,6 +407,34 @@ class TorchProfilerController:
         self._set_collection(False)
         logger.info("PROFILER: Profiler started with collection disabled")
 
+    def _write_summary(self) -> None:
+        """Compact per-rank op summary next to the trace: a key_averages table
+        and a JSON with input shapes, so operator-split analysis does not
+        require parsing multi-GB chrome traces."""
+        if self._profiler is None or not self._trace_dir:
+            return
+        try:
+            import json as _json
+            import os as _os
+            rank = _os.environ.get("RANK", "0")
+            averages = self._profiler.key_averages(group_by_input_shape=True)
+            stem = _os.path.join(self._trace_dir, f"summary_rank{rank}")
+            with open(f"{stem}.txt", "w") as fh:
+                fh.write(averages.table(sort_by="self_device_time_total", row_limit=60))
+            rows = [{
+                "name": e.key,
+                "shapes": str(e.input_shapes),
+                "self_device_us": e.self_device_time_total,
+                "device_us": e.device_time_total,
+                "count": e.count,
+            } for e in averages]
+            with open(f"{stem}.json", "w") as fh:
+                _json.dump(rows, fh)
+            if rank == "0":
+                logger.info("PROFILER: summary written to %s.txt", stem)
+        except Exception:  # noqa: BLE001 -- summaries must never break shutdown
+            logger.exception("PROFILER: summary generation failed")
+
     def stop(self) -> None:
         """Stop the profiler after disabling collection and clearing state."""
 
@@ -396,9 +443,10 @@ class TorchProfilerController:
 
         logger.info("PROFILER: Stopping profiler...")
         self._profiler.stop()
+        self._write_summary()
+        self._profiler = None  # makes stop() idempotent (atexit may re-enter)
         logger.info("PROFILER: Profiler stopped")
         self._active_region_depth = 0
-        set_global_profiler(None)
         set_global_controller(None)
 
     @property
@@ -411,9 +459,20 @@ class TorchProfilerController:
     def activities(self) -> tuple[torch.profiler.ProfilerActivity, ...]:
         return tuple(self._activities)
 
-    @property
-    def profiler(self) -> torch.profiler.profile | None:
-        return self._profiler
+
+@contextlib.contextmanager
+def profiler_region(region: str):
+    """Module-level region context manager bound to the global controller.
+
+    A no-op when no profiler is configured — stages can use this without the
+    get_global_controller()/nullcontext dance.
+    """
+    controller = get_global_controller()
+    if controller is None:
+        yield
+        return
+    with controller.region(region):
+        yield
 
 
 def profile_region(region: str) -> Callable[[Callable[..., Any]], Callable[..., Any]]:

fastvideo/registry.py

@@ -41,6 +41,7 @@ from fastvideo.configs.pipelines.flux_2 import (
 from fastvideo.configs.pipelines.matrixgame2 import MatrixGame2I2V480PConfig
 from fastvideo.configs.pipelines.matrixgame3 import MatrixGame3I2V720PConfig
 from fastvideo.configs.pipelines.minimax_h3 import MiniMaxH3PipelineConfig
+from fastvideo.configs.pipelines.mmaudio import MMAudioV2AConfig
 from fastvideo.configs.pipelines.turbodiffusion import (
     TurboDiffusionI2V_A14B_Config,
     TurboDiffusionT2V_14B_Config,
@@ -141,6 +142,25 @@ _MODEL_HF_PATH_TO_NAME: dict[str, str] = {}
 # Detectors to identify model families from paths or class names
 _MODEL_NAME_DETECTORS: list[tuple[str, Callable[[str], bool]]] = []
 
+# Single-file checkpoint bundles carry no `model_index.json`, so their pipeline
+# config is resolved from the transformer class the file declares about itself.
+# Keyed on that `_class_name` -- a stable, checkpoint-declared identifier --
+# and NOT on `model_id`, which is a positional index (`str(len(_CONFIG_REGISTRY))`)
+# and would silently rebind whenever registration order changes. The value pins
+# the model family, the pipeline class to build (a bundle names no pipeline
+# itself), and the default preset per training variant; the variant comes from
+# `ltx_single_file.bundle_variant` (declared in the header when present, with a
+# filename-token fallback). The transformer CLASS is never inferred from the
+# file name.
+_BUNDLE_TRANSFORMER_TO_CONFIG: dict[str, tuple[str, str, dict[str, str]]] = {
+    # transformer _class_name -> (model_family, pipeline class,
+    #                             variant -> default_preset)
+    "AVTransformer3DModel": ("ltx2", "LTX2Pipeline", {
+        "distilled": "ltx2_distilled",
+        "base": "ltx2_base",
+    }),
+}
+
 
 def register_configs(
     sampling_param_cls: type[SamplingParam] | None,
@@ -186,6 +206,64 @@ def get_model_short_name(model_id: str) -> str:
     return model_id
 
 
+def _bundle_config_info(model_path: str) -> ConfigInfo:
+    """Resolve the pipeline config for a single-file checkpoint bundle.
+
+    The bundle declares its own component classes in ``__metadata__``; reading
+    that is the only inference done here. The transformer's ``_class_name`` is
+    mapped through :data:`_BUNDLE_TRANSFORMER_TO_CONFIG`, an explicit table,
+    and the bundle's training variant (``bundle_variant``: header-declared,
+    filename token as fallback) selects the ``default_preset`` within that
+    entry, so a distilled bundle gets distilled sampling defaults and a
+    base/sft bundle gets the standard ones. Nothing is matched by wildcard,
+    and the transformer class is never read off the file name -- a file name
+    is incidental, and making class resolution depend on it makes correct
+    loading depend on what a checkpoint happens to be called.
+
+    An unrecognized transformer class raises, naming both the table to extend
+    and the per-run override, rather than falling through to a guess.
+
+    ponytail: LTX is the only family shipping a bundle, so this reads LTX
+    metadata directly. Dispatch on a per-format reader if a second bundle
+    format ever appears.
+    """
+    from fastvideo.models.loader.ltx_single_file import (bundle_variant, read_ltx_metadata)
+
+    try:
+        metadata = read_ltx_metadata(model_path)
+        cls_name = metadata.config.get("transformer", {}).get("_class_name")
+    except KeyError:  # no `config` section: not a bundle we know how to read
+        metadata = None
+        cls_name = None
+
+    entry = _BUNDLE_TRANSFORMER_TO_CONFIG.get(cls_name) if cls_name else None
+    if entry is None or metadata is None:
+        raise ValueError(
+            f"Single-file checkpoint {model_path} declares transformer class {cls_name!r}, which is not in "
+            f"_BUNDLE_TRANSFORMER_TO_CONFIG (known: {sorted(_BUNDLE_TRANSFORMER_TO_CONFIG)}). Add the class to that "
+            "table, or select the pipeline explicitly for this run with `override_pipeline_cls_name`. The pipeline is "
+            "never inferred from the checkpoint's file name.")
+
+    model_family, pipeline_cls_name, variant_presets = entry
+    variant = bundle_variant(metadata, model_path)
+    default_preset = variant_presets.get(variant)
+    if default_preset is None:
+        raise ValueError(f"Single-file checkpoint {model_path} resolves to variant {variant!r}, but "
+                         f"_BUNDLE_TRANSFORMER_TO_CONFIG[{cls_name!r}] only maps {sorted(variant_presets)}. "
+                         "Add the variant to that table.")
+
+    for config_info in _CONFIG_REGISTRY.values():
+        if config_info.model_family == model_family and config_info.default_preset == default_preset:
+            # A bundle names no pipeline itself (there is no model_index.json
+            # to carry a `_class_name`), so the table's pipeline class is
+            # pinned onto the resolved config.
+            return dataclasses.replace(config_info, pipeline_cls_name=pipeline_cls_name)
+
+    raise ValueError(f"_BUNDLE_TRANSFORMER_TO_CONFIG maps {cls_name!r} ({variant!r}) to model_family={model_family!r} "
+                     f"default_preset={default_preset!r}, but no registered config declares that pair. "
+                     "Fix the table or the registration.")
+
+
 def _get_config_info(
     model_path: str,
     *,
@@ -208,7 +286,13 @@ def _get_config_info(
             model_id = _MODEL_HF_PATH_TO_NAME[registered_model_hf_id]
             return _CONFIG_REGISTRY.get(model_id)
 
-    # 3. Use detectors (path or model_index pipeline name).
+    # 3. A single-file bundle is a FILE, so it would take the directory branch
+    #    below and be rejected for having no model_index.json. Its config lives
+    #    in its own `__metadata__` instead.
+    if model_path.endswith(".safetensors"):
+        return _bundle_config_info(model_path)
+
+    # 4. Use detectors (path or model_index pipeline name).
     if os.path.exists(model_path):
         config = verify_model_config_and_directory(model_path, required_component_dirs=[])
     else:
@@ -239,6 +323,22 @@ def _get_config_info(
 
 
 def _register_configs() -> None:
+    # MMAudio large-44k-v2 (video/text-to-audio). The checkpoint is converted
+    # into standard per-component FastVideo/Diffusers-style directories by
+    # scripts/checkpoint_conversion/convert_mmaudio_to_diffusers.py.
+    register_configs(
+        sampling_param_cls=None,
+        pipeline_config_cls=MMAudioV2AConfig,
+        workload_types=(WorkloadType.V2A, WorkloadType.T2A),
+        hf_model_paths=["FastVideo/MMAudio-large-44k-v2-Diffusers"],
+        model_detectors=[
+            lambda path: "mmaudio" in path.lower() or "mmaudiopipeline" in path.lower(),
+        ],
+        model_family="mmaudio",
+        default_preset="mmaudio_large_44k_v2",
+        pipeline_cls_name="MMAudioPipeline",
+    )
+
     # LTX-2 (distilled) — registered FIRST so its detector wins over
     # the base detector when both fire. The detector loop in
     # ``get_model_name_for_path`` ORs the path-based check with a
@@ -250,7 +350,10 @@ def _register_configs() -> None:
     register_configs(
         sampling_param_cls=None,
         pipeline_config_cls=LTX2T2VConfig,
-        workload_types=(WorkloadType.T2V, ),
+        # The LTX-2 family conditions i2v in LATENT SPACE (clean-latent +
+        # denoise mask, ltx2_image_conditioning.py); one checkpoint serves
+        # both workloads, so declaring T2V only under-reported the family.
+        workload_types=(WorkloadType.T2V, WorkloadType.I2V),
         hf_model_paths=[
             "FastVideo/LTX2-Distilled-Diffusers",
             # LTX-2.3 distilled aliases share the distilled pipeline/preset.
@@ -269,7 +372,10 @@ def _register_configs() -> None:
     register_configs(
         sampling_param_cls=None,
         pipeline_config_cls=LTX2T2VConfig,
-        workload_types=(WorkloadType.T2V, ),
+        # The LTX-2 family conditions i2v in LATENT SPACE (clean-latent +
+        # denoise mask, ltx2_image_conditioning.py); one checkpoint serves
+        # both workloads, so declaring T2V only under-reported the family.
+        workload_types=(WorkloadType.T2V, WorkloadType.I2V),
         hf_model_paths=[
             "Lightricks/LTX-2.3",
             "FastVideo/LTX2.3-base",
@@ -290,7 +396,10 @@ def _register_configs() -> None:
     register_configs(
         sampling_param_cls=None,
         pipeline_config_cls=LTX2T2VConfig,
-        workload_types=(WorkloadType.T2V, ),
+        # The LTX-2 family conditions i2v in LATENT SPACE (clean-latent +
+        # denoise mask, ltx2_image_conditioning.py); one checkpoint serves
+        # both workloads, so declaring T2V only under-reported the family.
+        workload_types=(WorkloadType.T2V, WorkloadType.I2V),
         hf_model_paths=[
             "Lightricks/LTX-2",
             "FastVideo/LTX2-base",
@@ -310,12 +419,12 @@ def _register_configs() -> None:
     # 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; use T2V as the
-    # compatibility placeholder until the enum is extended.
     register_configs(
         sampling_param_cls=None,
         pipeline_config_cls=StableAudioT2AConfig,
-        workload_types=(WorkloadType.T2V, ),
+        # Keep T2V as a backward-compatible alias for existing callers while
+        # exposing the semantically correct audio workload to new clients.
+        workload_types=(WorkloadType.T2A, WorkloadType.T2V),
         hf_model_paths=[
             "FastVideo/stable-audio-open-1.0-Diffusers",
         ],
@@ -334,7 +443,7 @@ def _register_configs() -> None:
     register_configs(
         sampling_param_cls=None,
         pipeline_config_cls=StableAudioOpenSmallConfig,
-        workload_types=(WorkloadType.T2V, ),
+        workload_types=(WorkloadType.T2A, WorkloadType.T2V),
         hf_model_paths=[
             "FastVideo/stable-audio-open-small-Diffusers",
         ],
@@ -1220,7 +1329,12 @@ def get_model_info(
         pipeline_name: str | None = override_pipeline_cls_name
         logger.info("Using override pipeline class name %s", pipeline_name)
     else:
-        if os.path.exists(model_path):
+        if model_path.endswith(".safetensors"):
+            # A single-file bundle has no model_index.json and names no
+            # pipeline itself; `_bundle_config_info` above pinned the class
+            # through `pipeline_cls_name`.
+            config = {}
+        elif os.path.exists(model_path):
             config = verify_model_config_and_directory(model_path, required_component_dirs=[])
         else:
             config = maybe_download_model_index(model_path, revision=revision)
@@ -1302,6 +1416,8 @@ def _register_presets() -> None:
         ALL_PRESETS as MATRIXGAME3_PRESETS, )
     from fastvideo.pipelines.basic.minimax_h3.presets import (
         ALL_PRESETS as MINIMAX_H3_PRESETS, )
+    from fastvideo.pipelines.basic.mmaudio.presets import (
+        ALL_PRESETS as MMAUDIO_PRESETS, )
     from fastvideo.pipelines.basic.sd35.presets import (
         ALL_PRESETS as SD35_PRESETS, )
     from fastvideo.pipelines.basic.stable_audio.presets import (
@@ -1333,6 +1449,7 @@ def _register_presets() -> None:
         MATRIXGAME2_PRESETS,
         MATRIXGAME3_PRESETS,
         MINIMAX_H3_PRESETS,
+        MMAUDIO_PRESETS,
         SD35_PRESETS,
         STABLE_AUDIO_PRESETS,
         TURBODIFFUSION_PRESETS,

fastvideo/tests/api/test_ltx_single_file.py

@@ -0,0 +1,400 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Checks for single-file LTX metadata parsing and prefix routing.
+
+Builds a tiny synthetic bundle shaped like the real one -- same metadata keys,
+same top-level tensor prefixes -- so no checkpoint is needed. Run directly
+(``python fastvideo/tests/api/test_ltx_single_file.py``) or under pytest.
+"""
+
+import json
+import tempfile
+from pathlib import Path
+
+import torch
+from safetensors.torch import save_file
+
+from fastvideo.models.loader.ltx_single_file import (
+    LTXCheckpointMetadata,
+    resolve_text_encoder_root,
+    build_dit_config,
+    bundle_model_index,
+    component_weights,
+    read_ltx_metadata,
+)
+
+# Subset of the transformer section as the checkpoint declares it.
+TRANSFORMER_SECTION = {
+    "_class_name": "AVTransformer3DModel",
+    "num_layers": 48,
+    "num_attention_heads": 32,
+    "attention_head_dim": 128,
+    "audio_num_attention_heads": 32,
+    "audio_attention_head_dim": 64,
+    "cross_attention_dim": 4096,
+    "audio_cross_attention_dim": 2048,
+    "caption_channels": 3840,
+    "ff_bias": False,
+    "cross_attention_adaln": True,
+    "caption_proj_before_connector": True,
+    "apply_gated_attention": True,
+    "rope_type": "split",
+    "frequencies_precision": "float64",
+}
+
+
+def _write_bundle(path: Path,
+                  *,
+                  with_gemma_source: bool,
+                  transformer_cls: str = "AVTransformer3DModel",
+                  variant: str | None = None) -> None:
+    tensors = {
+        "model.diffusion_model.patchify_proj.weight": torch.zeros(2, 2),
+        "model.diffusion_model.transformer_blocks.0.ff.net.2.weight": torch.zeros(2, 2),
+        # Stored under the transformer prefix but owned by the text stack.
+        "model.diffusion_model.video_embeddings_connector.x.weight": torch.zeros(2),
+        "model.diffusion_model.audio_embeddings_connector.x.weight": torch.zeros(2),
+        "vae.encoder.conv.weight": torch.zeros(2),
+        "audio_vae.decoder.conv.weight": torch.zeros(2),
+        "vocoder.conv.weight": torch.zeros(2),
+        "text_embedding_projection.video_aggregate_embed.weight": torch.zeros(2),
+        "duration_head.linear.weight": torch.zeros(2),
+    }
+    section = dict(TRANSFORMER_SECTION, _class_name=transformer_cls)
+    config = {
+        "transformer": section,
+        # Declares its class: a component that can be built from the bundle.
+        "vae": {
+            "_class_name": "CausalVideoAutoencoder",
+            "dims": 3
+        },
+        # Present and carrying weights, but naming no class.
+        "audio_vae": {},
+    }
+    metadata = {
+        "config": json.dumps(config),
+        "model_version": "9.9.9",
+        "license": "x" * 128,
+    }
+    if with_gemma_source:
+        metadata["gemma_source_checkpoint"] = json.dumps({"ltx_version": "9.9.9", "gemma_version": "fake-encoder-v0"})
+    if variant is not None:
+        metadata["variant"] = variant
+    save_file(tensors, str(path), metadata=metadata)
+
+
+def test_read_metadata_and_routing() -> None:
+    with tempfile.TemporaryDirectory() as tmp:
+        path = Path(tmp) / "bundle.safetensors"
+        _write_bundle(path, with_gemma_source=True)
+
+        meta = read_ltx_metadata(str(path))
+        assert meta.model_version == "9.9.9"
+        assert meta.gemma_source_checkpoint == {
+            "ltx_version": "9.9.9",
+            "gemma_version": "fake-encoder-v0",
+        }
+        assert set(meta.config) == {"transformer", "vae", "audio_vae"}
+        # frequencies_precision is the checkpoint's name for double_precision_rope.
+        assert meta.config["transformer"]["double_precision_rope"] is True
+
+        # Prefix stripped, connectors excluded, nothing from other components.
+        transformer_keys = {n for n, _ in component_weights(str(path), "transformer")}
+        assert transformer_keys == {
+            "patchify_proj.weight",
+            "transformer_blocks.0.ff.net.2.weight",
+        }
+        for component, expected in (
+            ("vae", {"encoder.conv.weight"}),
+            ("audio_vae", {"decoder.conv.weight"}),
+            ("vocoder", {"conv.weight"}),
+                # The connectors come off the transformer prefix with their
+                # sub-tree name intact; the text encoder renames from there.
+            ("text_encoder", {
+                "video_aggregate_embed.weight",
+                "video_embeddings_connector.x.weight",
+                "audio_embeddings_connector.x.weight",
+            }),
+        ):
+            assert {n for n, _ in component_weights(str(path), component)} == expected
+
+
+def test_bundle_model_index_declares_what_the_metadata_declares() -> None:
+    """The model_index a bundle stands in for: same shape, same libraries."""
+    with tempfile.TemporaryDirectory() as tmp:
+        path = Path(tmp) / "bundle.safetensors"
+        _write_bundle(path, with_gemma_source=True)
+        index = bundle_model_index(str(path))
+
+        # A section that names a class is a component we can build.
+        assert index["transformer"] == ["diffusers", "AVTransformer3DModel"]
+        assert index["vae"] == ["diffusers", "CausalVideoAutoencoder"]
+        # A section that names none stays declared with a null library, so
+        # `load_modules` drops it from the required set on its own. Omitting
+        # the key instead would trip its required-module check.
+        assert index["audio_vae"] == [None, None]
+        # Both live outside the bundle, but the pipeline requires them.
+        assert index["text_encoder"] == ["transformers", "LTX2GemmaTextEncoderModel"]
+        assert index["tokenizer"] == ["transformers", "AutoTokenizer"]
+        # `load_modules` pops both of these without a default.
+        assert "_class_name" in index and "_diffusers_version" in index
+
+
+def test_missing_gemma_source_is_none() -> None:
+    with tempfile.TemporaryDirectory() as tmp:
+        path = Path(tmp) / "bundle.safetensors"
+        _write_bundle(path, with_gemma_source=False)
+        assert read_ltx_metadata(str(path)).gemma_source_checkpoint is None
+
+
+def test_dit_config_takes_ff_bias_from_metadata() -> None:
+    with tempfile.TemporaryDirectory() as tmp:
+        path = Path(tmp) / "bundle.safetensors"
+        _write_bundle(path, with_gemma_source=True)
+        arch = build_dit_config(read_ltx_metadata(str(path))).arch_config
+
+        # ff_bias comes from metadata; audio_ff_bias has no metadata key and
+        # must keep the default, matching the audio_ff.*.bias tensors that the
+        # checkpoint does carry.
+        assert arch.ff_bias is False
+        assert arch.audio_ff_bias is True
+        assert arch.num_layers == 48
+        assert arch.cross_attention_adaln is True
+        assert arch.caption_channels == 3840
+
+
+def test_defaults_unchanged_without_metadata() -> None:
+    from fastvideo.configs.models.dits.ltx2 import LTX2VideoArchConfig
+
+    arch = LTX2VideoArchConfig()
+    assert arch.ff_bias is True
+    assert arch.audio_ff_bias is True
+
+
+def test_encoder_root_requires_an_explicit_declaration() -> None:
+    """Neither source set -> raise, naming both. Never guess from the filesystem."""
+    import pytest
+    with pytest.raises(ValueError, match="must be declared"):
+        resolve_text_encoder_root(configured_path=None, override_path=None)
+
+
+def test_encoder_root_override_beats_config() -> None:
+    assert resolve_text_encoder_root("/from/config", "/from/override") == "/from/override"
+    assert resolve_text_encoder_root("/from/config", None) == "/from/config"
+
+
+def test_encoder_root_is_not_discovered_from_a_sibling_directory() -> None:
+    """A plausible encoder sitting next to the bundle must NOT be picked up."""
+    with tempfile.TemporaryDirectory() as tmp:
+        sibling = Path(tmp) / "gemma"
+        sibling.mkdir()
+        (sibling / "config.json").write_text("{}")
+        import pytest
+        with pytest.raises(ValueError):
+            resolve_text_encoder_root(configured_path=None, override_path=None)
+
+
+def test_declared_root_survives_a_version_mismatch() -> None:
+    """Pairing metadata validates a declared root; it never overrides it."""
+    with tempfile.TemporaryDirectory() as tmp:
+        root = Path(tmp) / "encoder"
+        root.mkdir()
+        (root / "config.json").write_text(json.dumps({"gemma_version": "other-v0"}))
+        meta = LTXCheckpointMetadata(config={},
+                                     model_version="9.9.9",
+                                     gemma_source_checkpoint={"gemma_version": "fake-encoder-v0"})
+        assert resolve_text_encoder_root(str(root), None, meta) == str(root)
+
+
+def test_bundle_resolves_its_pipeline_config_from_declared_transformer_class() -> None:
+    """A bundle path resolves through the alias table, not through model_index.json."""
+    from fastvideo.pipelines.basic.ltx2.pipeline_configs import LTX2T2VConfig
+    from fastvideo.registry import get_pipeline_config_cls_from_name
+
+    with tempfile.TemporaryDirectory() as tmp:
+        # Deliberately uninformative name: resolution must not read the file name.
+        path = Path(tmp) / "anonymous.safetensors"
+        _write_bundle(path, with_gemma_source=True)
+        assert get_pipeline_config_cls_from_name(str(path)) is LTX2T2VConfig
+
+
+def test_unknown_bundle_transformer_class_raises_naming_table_and_override() -> None:
+    """No wildcard fallback: an unmapped class fails loud and says how to fix it."""
+    import pytest
+
+    from fastvideo.registry import get_pipeline_config_cls_from_name
+
+    with tempfile.TemporaryDirectory() as tmp:
+        # Name it after a model the detectors DO know, to prove the file name
+        # is not consulted -- only the class the checkpoint declares.
+        path = Path(tmp) / "ltx2-distilled.safetensors"
+        _write_bundle(path, with_gemma_source=True, transformer_cls="NotAModel")
+        with pytest.raises(ValueError) as excinfo:
+            get_pipeline_config_cls_from_name(str(path))
+        message = str(excinfo.value)
+        assert "_BUNDLE_TRANSFORMER_TO_CONFIG" in message
+        assert "override_pipeline_cls_name" in message
+
+
+def test_preset_resolution_prefers_the_declared_variant() -> None:
+    """A bundle's header-declared `variant` decides its preset, not its name."""
+    from fastvideo.api.sampling_param import SamplingParam
+
+    with tempfile.TemporaryDirectory() as tmp:
+        # Distilled declared, in a file whose name says nothing.
+        distilled = Path(tmp) / "anonymous.safetensors"
+        _write_bundle(distilled, with_gemma_source=False, variant="distilled-rc2")
+        sp = SamplingParam.from_pretrained(str(distilled))
+        assert sp.num_inference_steps == 8
+        assert sp.guidance_scale == 1.0
+
+        # A declared sft variant beats a file name that says "distilled".
+        sft = Path(tmp) / "something-distilled.safetensors"
+        _write_bundle(sft, with_gemma_source=True, variant="sft-rc2")
+        sp = SamplingParam.from_pretrained(str(sft))
+        assert sp.num_inference_steps == 40
+        assert sp.guidance_scale == 3.0
+
+
+def test_preset_resolution_falls_back_to_the_file_name() -> None:
+    """No variant in the header: the "distilled" filename token decides."""
+    from fastvideo.api.sampling_param import SamplingParam
+
+    with tempfile.TemporaryDirectory() as tmp:
+        distilled = Path(tmp) / "some-distilled-bundle.safetensors"
+        _write_bundle(distilled, with_gemma_source=False)
+        sp = SamplingParam.from_pretrained(str(distilled))
+        assert sp.num_inference_steps == 8
+        assert sp.guidance_scale == 1.0
+
+        base = Path(tmp) / "some-sft-bundle.safetensors"
+        _write_bundle(base, with_gemma_source=True)
+        sp = SamplingParam.from_pretrained(str(base))
+        assert sp.num_inference_steps == 40
+        assert sp.guidance_scale == 3.0
+
+
+def test_bundle_preset_table_is_internally_consistent() -> None:
+    """Invariants the bundle->preset mapping relies on, asserted loudly."""
+    from fastvideo.registry import _BUNDLE_TRANSFORMER_TO_CONFIG, _CONFIG_REGISTRY
+
+    registered = {(ci.model_family, ci.default_preset) for ci in _CONFIG_REGISTRY.values()}
+    for cls_name, (family, pipeline_cls_name, variants) in _BUNDLE_TRANSFORMER_TO_CONFIG.items():
+        assert set(variants) == {
+            "distilled", "base"
+        }, (f"{cls_name}: bundle_variant() only ever returns 'distilled' or 'base', "
+            f"but the table maps {sorted(variants)} -- some bundles could not resolve a preset.")
+        assert pipeline_cls_name, (f"{cls_name}: a bundle names no pipeline itself, so the table entry must.")
+        for variant, preset in variants.items():
+            assert (family, preset) in registered, (
+                f"{cls_name}/{variant} -> ({family!r}, {preset!r}) is not a registered config; "
+                "fix _BUNDLE_TRANSFORMER_TO_CONFIG or _register_configs().")
+
+
+def test_bundle_resolves_a_pipeline_class_without_a_model_index() -> None:
+    """`get_model_info` on a bundle pins the pipeline class from the table,
+    so `VideoGenerator.from_pretrained(<bundle>)` needs no override."""
+    from fastvideo.registry import get_model_info
+
+    with tempfile.TemporaryDirectory() as tmp:
+        path = Path(tmp) / "anonymous.safetensors"
+        _write_bundle(path, with_gemma_source=True, variant="distilled-rc2")
+        info = get_model_info(str(path))
+        assert info.pipeline_cls.__name__ == "LTX2Pipeline"
+
+
+_EXAMPLES_DIR = Path(__file__).resolve().parents[3] / "examples" / "inference" / "basic"
+_EXAMPLE_SCRIPTS = ("basic_ltx2_5_distilled.py", "basic_ltx2_5.py")
+
+
+def _load_example(name: str):
+    import importlib.util
+    spec = importlib.util.spec_from_file_location(name[:-3], _EXAMPLES_DIR / name)
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
+
+def test_example_scripts_parse_help_cleanly() -> None:
+    import runpy
+    import sys
+
+    import pytest
+    for name in _EXAMPLE_SCRIPTS:
+        script = str(_EXAMPLES_DIR / name)
+        argv = sys.argv
+        sys.argv = [script, "--help"]
+        try:
+            with pytest.raises(SystemExit) as excinfo:
+                runpy.run_path(script, run_name="__main__")
+            assert excinfo.value.code == 0, f"{name} --help exited nonzero"
+        finally:
+            sys.argv = argv
+
+
+def test_example_sampling_defaults_come_from_the_bundle_preset() -> None:
+    """With no sampling flags the examples pass NO overrides, so what runs is
+    exactly ``SamplingParam.from_pretrained(bundle)`` -- the preset the
+    bundle's variant selects."""
+    from fastvideo.api.sampling_param import SamplingParam
+
+    with tempfile.TemporaryDirectory() as tmp:
+        for name, variant, steps, cfg in (
+            ("basic_ltx2_5_distilled.py", "distilled-rc2", 8, 1.0),
+            ("basic_ltx2_5.py", "sft-rc2", 40, 3.0),
+        ):
+            bundle = Path(tmp) / f"{variant}.safetensors"
+            _write_bundle(bundle, with_gemma_source=True, variant=variant)
+            module = _load_example(name)
+            args = module.parse_args(["--model-path", str(bundle)])
+            assert module.sampling_overrides(args) == {}, (f"{name}: unset flags must not override the preset")
+            sp = SamplingParam.from_pretrained(str(bundle))
+            assert (sp.num_inference_steps, sp.guidance_scale) == (steps, cfg)
+
+
+def test_connector_factorization_must_match_its_stream_width() -> None:
+    """A connector's heads x head_dim must equal the width of the stream it feeds.
+
+    This is the mistake that loads cleanly and computes wrong, so it fails loud
+    and names both numbers. A field the checkpoint does not declare is not
+    checked -- the arch default applies and there is nothing to contradict.
+    """
+    import pytest
+
+    from fastvideo.models.loader.component_loader import _check_connector_widths
+
+    declared = {
+        "connector_num_attention_heads": 32,
+        "connector_attention_head_dim": 128,
+        "cross_attention_dim": 4096,
+        "audio_connector_num_attention_heads": 32,
+        "audio_connector_attention_head_dim": 64,
+        "audio_cross_attention_dim": 2048,
+    }
+    _check_connector_widths(declared)
+    _check_connector_widths({})
+    _check_connector_widths({"connector_num_attention_heads": 32})
+
+    # The audio connector inheriting the video head_dim is the exact failure
+    # this exists to stop: 32 x 128 = 4096, but the audio stream is 2048.
+    with pytest.raises(ValueError, match="audio"):
+        _check_connector_widths(dict(declared, audio_connector_attention_head_dim=128))
+    with pytest.raises(ValueError, match="video"):
+        _check_connector_widths(dict(declared, connector_num_attention_heads=30))
+
+if __name__ == "__main__":
+    test_read_metadata_and_routing()
+    test_bundle_model_index_declares_what_the_metadata_declares()
+    test_missing_gemma_source_is_none()
+    test_dit_config_takes_ff_bias_from_metadata()
+    test_defaults_unchanged_without_metadata()
+    test_encoder_root_override_beats_config()
+    test_bundle_resolves_its_pipeline_config_from_declared_transformer_class()
+    test_unknown_bundle_transformer_class_raises_naming_table_and_override()
+    test_preset_resolution_prefers_the_declared_variant()
+    test_preset_resolution_falls_back_to_the_file_name()
+    test_bundle_preset_table_is_internally_consistent()
+    test_bundle_resolves_a_pipeline_class_without_a_model_index()
+    test_example_scripts_parse_help_cleanly()
+    test_example_sampling_defaults_come_from_the_bundle_preset()
+    test_connector_factorization_must_match_its_stream_width()
+    print("ok")

fastvideo/tests/attention/test_flash_attn_no_pad_custom_op.py

@@ -85,6 +85,50 @@ def test_no_pad_inference_matches_original(no_pad_impls, dtype):
     torch.testing.assert_close(out_test, out_ref, atol=0, rtol=0)
 
 
+def test_backend_translates_explicit_causal_mask(no_pad_impls):
+    """MMAudio's additive mask must route through native causal attention."""
+    from fastvideo.attention.backends.flash_attn import (
+        FlashAttentionImpl,
+        flash_attn_func_compilable,
+    )
+    from fastvideo.attention.backends.sdpa import SDPAMetadata
+
+    torch.manual_seed(0)
+    device = torch.device("cuda")
+    batch, seqlen, heads, head_dim = 1, 32, 2, 64
+    query = torch.randn(batch, seqlen, heads, head_dim, device=device, dtype=torch.bfloat16)
+    key = torch.randn_like(query)
+    value = torch.randn_like(query)
+    causal_mask = torch.full(
+        (1, 1, seqlen, seqlen),
+        float("-inf"),
+        device=device,
+        dtype=query.dtype,
+    ).triu_(1)
+    impl = FlashAttentionImpl(
+        num_heads=heads,
+        head_size=head_dim,
+        causal=False,
+        softmax_scale=head_dim**-0.5,
+    )
+
+    with torch.inference_mode():
+        expected = flash_attn_func_compilable(
+            query,
+            key,
+            value,
+            softmax_scale=head_dim**-0.5,
+            causal=True,
+        )
+        actual = impl.forward(
+            query,
+            key,
+            value,
+            SDPAMetadata(current_timestep=0, attn_mask=causal_mask, is_causal=True),
+        )
+    torch.testing.assert_close(actual, expected, atol=0, rtol=0)
+
+
 @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
 def test_no_pad_training_backward_through_registered_autograd(no_pad_impls, dtype):
     """FA2: grads flow through the registered op and match the original."""

fastvideo/tests/contract/test_profiler_regions.py

@@ -0,0 +1,107 @@
+# SPDX-License-Identifier: Apache-2.0
+"""CPU contract tests for the profiler region system.
+
+The region system's failure mode is silent: a profiler that records nothing
+(the old wait/warmup schedule bug), a trace that never exports (no shutdown
+hook), or a typo'd region name that no-ops forever. Each test here runs a
+child process — the profiler is process-global and its export happens at
+atexit, so in-process tests can't observe the contract.
+"""
+from __future__ import annotations
+
+import glob
+import gzip
+import json
+import os
+import subprocess
+import sys
+
+# Five-window child: ops before any region, inside a region, between regions,
+# inside a second (short-named) region, after the last region. Exits without
+# calling stop() — export must happen via the atexit hook. Each window is
+# labeled with a record_function marker so the parent can check which windows
+# the trace actually captured.
+_CHILD = r"""
+import torch
+from fastvideo.profiler import get_or_create_profiler, profiler_region
+
+controller = get_or_create_profiler("{trace_dir}")
+
+def burn(tag):
+    with torch.profiler.record_function(tag):
+        torch.mm(torch.ones(8, 8), torch.ones(8, 8))
+
+burn("win_pre")
+with profiler_region("profiler_region_inference_denoising"):
+    burn("win_region1")
+burn("win_between")
+with profiler_region("model_loading"):  # short name must resolve
+    burn("win_region2")
+for _ in range(2):  # warn-once: second use must not log again
+    with profiler_region("definitely_not_a_region"):
+        burn("win_typo")
+burn("win_post")
+# no controller.stop(): atexit owns the export
+"""
+
+
+def _run_child(tmp_path):
+    trace_dir = str(tmp_path / "traces")
+    env = os.environ.copy()
+    env["FASTVIDEO_TORCH_PROFILER_DIR"] = trace_dir
+    env["FASTVIDEO_TORCH_PROFILE_REGIONS"] = "inference_denoising,model_loading"
+    proc = subprocess.run(
+        [sys.executable, "-c", _CHILD.format(trace_dir=trace_dir)],
+        env=env,
+        capture_output=True,
+        text=True,
+        timeout=300,
+    )
+    assert proc.returncode == 0, proc.stderr
+    return trace_dir, proc.stdout + proc.stderr
+
+
+def _trace_event_names(trace_dir):
+    traces = glob.glob(os.path.join(trace_dir, "**", "*.json*"), recursive=True)
+    traces = [t for t in traces if "summary" not in os.path.basename(t)]
+    assert traces, f"no trace exported to {trace_dir} (atexit hook missing?)"
+    opener = gzip.open if traces[0].endswith(".gz") else open
+    with opener(traces[0], "rt") as fh:
+        events = json.load(fh).get("traceEvents", [])
+    return {e.get("name", "") for e in events}
+
+
+def test_regions_gate_collection_and_atexit_exports(tmp_path):
+    trace_dir, output = _run_child(tmp_path)
+    names = _trace_event_names(trace_dir)
+
+    assert "win_region1" in names, "op inside an enabled region was not captured"
+    assert "win_region2" in names, "short region name did not resolve/capture"
+    assert "fastvideo.region::profiler_region_inference_denoising" in names
+
+    for leaked in ("win_pre", "win_between", "win_typo", "win_post"):
+        assert leaked not in names, f"op outside any region leaked into trace: {leaked}"
+
+    # unregistered region warns exactly once across repeated uses
+    assert output.count("definitely_not_a_region") >= 1
+    assert output.count("is not registered") == 1
+
+    # per-rank op summary written next to the trace
+    summaries = glob.glob(os.path.join(trace_dir, "summary_rank0.*"))
+    assert sorted(os.path.splitext(s)[1] for s in summaries) == [".json", ".txt"]
+
+
+def test_noop_without_profiler_dir(tmp_path):
+    # profiler_region must be a clean no-op when profiling is not configured
+    child = (
+        "from fastvideo.profiler import profiler_region\n"
+        "with profiler_region('inference_denoising'):\n"
+        "    x = 1\n"
+        "assert x == 1\n"
+    )
+    env = os.environ.copy()
+    env.pop("FASTVIDEO_TORCH_PROFILER_DIR", None)
+    env.pop("FASTVIDEO_TORCH_PROFILE_REGIONS", None)
+    proc = subprocess.run([sys.executable, "-c", child], env=env,
+                          capture_output=True, text=True, timeout=300)
+    assert proc.returncode == 0, proc.stderr

fastvideo/tests/transformers/test_mmaudio.py

@@ -0,0 +1,15 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import torch
+
+from fastvideo.models.dits.mmaudio import TimestepEmbedder
+
+
+def test_timestep_embedder_accepts_integer_timesteps() -> None:
+    embedder = TimestepEmbedder(8, 8, 10_000)
+
+    actual = embedder(torch.tensor([1], dtype=torch.long))
+    expected = embedder(torch.tensor([1.0]))
+
+    torch.testing.assert_close(actual, expected, atol=0, rtol=0)
+    assert actual.dtype == embedder.mlp[0].weight.dtype

fastvideo/third_party/eval/synchformer/ast.py

@@ -4,8 +4,8 @@ from transformers.modeling_outputs import BaseModelOutputWithPooling
 
 # importing modified version of AST
 from fastvideo.third_party.eval.synchformer.hf_src.modeling_ast import (ASTConfig, ASTForAudioClassification)
-from fastvideo.third_party.eval.synchformer.motionformer import (AveragePooling, BaseEncoderLayer,
-                                                                 TemporalTransformerEncoderLayer)
+from fastvideo.third_party.synchformer.motionformer import (AveragePooling, BaseEncoderLayer,
+                                                            TemporalTransformerEncoderLayer)
 
 
 class AST(torch.nn.Module):

fastvideo/third_party/eval/synchformer/motionformer.py

@@ -1,388 +1,3 @@
-import logging
-from pathlib import Path
+"""Backward-compatible import for the shared Synchformer backbone."""
 
-import einops
-import torch
-from omegaconf import OmegaConf
-from timm.layers import trunc_normal_
-from torch import nn
-
-from fastvideo.third_party.eval.synchformer.utils import check_if_file_exists_else_download
-from fastvideo.third_party.eval.synchformer.video_model_builder import VisionTransformer
-
-FILE2URL = {
-    # cfg
-    'motionformer_224_16x4.yaml':
-    'https://raw.githubusercontent.com/facebookresearch/Motionformer/bf43d50/configs/SSV2/motionformer_224_16x4.yaml',
-    'joint_224_16x4.yaml':
-    'https://raw.githubusercontent.com/facebookresearch/Motionformer/bf43d50/configs/SSV2/joint_224_16x4.yaml',
-    'divided_224_16x4.yaml':
-    'https://raw.githubusercontent.com/facebookresearch/Motionformer/bf43d50/configs/SSV2/divided_224_16x4.yaml',
-    # ckpt
-    'ssv2_motionformer_224_16x4.pyth': 'https://dl.fbaipublicfiles.com/motionformer/ssv2_motionformer_224_16x4.pyth',
-    'ssv2_joint_224_16x4.pyth': 'https://dl.fbaipublicfiles.com/motionformer/ssv2_joint_224_16x4.pyth',
-    'ssv2_divided_224_16x4.pyth': 'https://dl.fbaipublicfiles.com/motionformer/ssv2_divided_224_16x4.pyth',
-}
-
-
-class MotionFormer(VisionTransformer):
-    ''' This class serves three puposes:
-            1. Renames the class to MotionFormer.
-            2. Downloads the cfg from the original repo and patches it if needed.
-            3. Takes care of feature extraction by redefining .forward()
-                - if `extract_features=True` and `factorize_space_time=False`,
-                    the output is of shape (B, T, D) where T = 1 + (224 // 16) * (224 // 16) * 8
-                - if `extract_features=True` and `factorize_space_time=True`, the output is of shape (B*S, D)
-                    and spatial and temporal transformer encoder layers are used.
-                - if `extract_features=True` and `factorize_space_time=True` as well as `add_global_repr=True`
-                    the output is of shape (B, D) and spatial and temporal transformer encoder layers
-                    are used as well as the global representation is extracted from segments (extra pos emb
-                    is added).
-    '''
-
-    def __init__(
-        self,
-        extract_features: bool = False,
-        ckpt_path: str = None,
-        factorize_space_time: bool = None,
-        agg_space_module: str = None,
-        agg_time_module: str = None,
-        add_global_repr: bool = True,
-        agg_segments_module: str = None,
-        max_segments: int = None,
-    ):
-        self.extract_features = extract_features
-        self.ckpt_path = ckpt_path
-        self.factorize_space_time = factorize_space_time
-
-        if self.ckpt_path is not None:
-            check_if_file_exists_else_download(self.ckpt_path, FILE2URL)
-            ckpt = torch.load(self.ckpt_path, map_location='cpu')
-            mformer_ckpt2cfg = {
-                'ssv2_motionformer_224_16x4.pyth': 'motionformer_224_16x4.yaml',
-                'ssv2_joint_224_16x4.pyth': 'joint_224_16x4.yaml',
-                'ssv2_divided_224_16x4.pyth': 'divided_224_16x4.yaml',
-            }
-            # init from motionformer ckpt or from our Stage I ckpt
-            # depending on whether the feat extractor was pre-trained on AVCLIPMoCo or not, we need to
-            # load the state dict differently
-            was_pt_on_avclip = self.ckpt_path.endswith('.pt')  # checks if it is a stage I ckpt (FIXME: a bit generic)
-            if self.ckpt_path.endswith(tuple(mformer_ckpt2cfg.keys())):
-                cfg_fname = mformer_ckpt2cfg[Path(self.ckpt_path).name]
-            elif was_pt_on_avclip:
-                # TODO: this is a hack, we should be able to get the cfg from the ckpt (earlier ckpt didn't have it)
-                s1_cfg = ckpt.get('args', None)  # Stage I cfg
-                if s1_cfg is not None:
-                    s1_vfeat_extractor_ckpt_path = s1_cfg.model.params.vfeat_extractor.params.ckpt_path
-                    # if the stage I ckpt was initialized from a motionformer ckpt or train from scratch
-                    if s1_vfeat_extractor_ckpt_path is not None:
-                        cfg_fname = mformer_ckpt2cfg[Path(s1_vfeat_extractor_ckpt_path).name]
-                    else:
-                        cfg_fname = 'divided_224_16x4.yaml'
-                else:
-                    cfg_fname = 'divided_224_16x4.yaml'
-            else:
-                raise ValueError(f'ckpt_path {self.ckpt_path} is not supported.')
-        else:
-            was_pt_on_avclip = False
-            cfg_fname = 'divided_224_16x4.yaml'
-            # logging.info(f'No ckpt_path provided, using {cfg_fname} config.')
-
-        if cfg_fname in ['motionformer_224_16x4.yaml', 'divided_224_16x4.yaml']:
-            pos_emb_type = 'separate'
-        elif cfg_fname == 'joint_224_16x4.yaml':
-            pos_emb_type = 'joint'
-
-        self.mformer_cfg_path = Path(__file__).absolute().parent / cfg_fname
-
-        check_if_file_exists_else_download(self.mformer_cfg_path, FILE2URL)
-        mformer_cfg = OmegaConf.load(self.mformer_cfg_path)
-        logging.info(f'Loading MotionFormer config from {self.mformer_cfg_path.absolute()}')
-
-        # patch the cfg (from the default cfg defined in the repo `Motionformer/slowfast/config/defaults.py`)
-        mformer_cfg.VIT.ATTN_DROPOUT = 0.0
-        mformer_cfg.VIT.POS_EMBED = pos_emb_type
-        mformer_cfg.VIT.USE_ORIGINAL_TRAJ_ATTN_CODE = True
-        mformer_cfg.VIT.APPROX_ATTN_TYPE = 'none'  # guessing
-        mformer_cfg.VIT.APPROX_ATTN_DIM = 64  # from ckpt['cfg']
-
-        # finally init VisionTransformer with the cfg
-        super().__init__(mformer_cfg)
-
-        # load the ckpt now if ckpt is provided and not from AVCLIPMoCo-pretrained ckpt
-        if (self.ckpt_path is not None) and (not was_pt_on_avclip):
-            _ckpt_load_status = self.load_state_dict(ckpt['model_state'], strict=False)
-            if len(_ckpt_load_status.missing_keys) > 0 or len(_ckpt_load_status.unexpected_keys) > 0:
-                logging.warning(f'Loading exact vfeat_extractor ckpt from {self.ckpt_path} failed.' \
-                                f'Missing keys: {_ckpt_load_status.missing_keys}, ' \
-                                f'Unexpected keys: {_ckpt_load_status.unexpected_keys}')
-            else:
-                logging.info(f'Loading vfeat_extractor ckpt from {self.ckpt_path} succeeded.')
-
-        if self.extract_features:
-            assert isinstance(self.norm, nn.LayerNorm), 'early x[:, 1:, :] may not be safe for per-tr weights'
-            # pre-logits are Sequential(nn.Linear(emb, emd), act) and `act` is tanh but see the logger
-            self.pre_logits = nn.Identity()
-            # we don't need the classification head (saving memory)
-            self.head = nn.Identity()
-            self.head_drop = nn.Identity()
-            # avoiding code duplication (used only if agg_*_module is TransformerEncoderLayer)
-            transf_enc_layer_kwargs = dict(
-                d_model=self.embed_dim,
-                nhead=self.num_heads,
-                activation=nn.GELU(),
-                batch_first=True,
-                dim_feedforward=self.mlp_ratio * self.embed_dim,
-                dropout=self.drop_rate,
-                layer_norm_eps=1e-6,
-                norm_first=True,
-            )
-            # define adapters if needed
-            if self.factorize_space_time:
-                if agg_space_module == 'TransformerEncoderLayer':
-                    self.spatial_attn_agg = SpatialTransformerEncoderLayer(**transf_enc_layer_kwargs)
-                elif agg_space_module == 'AveragePooling':
-                    self.spatial_attn_agg = AveragePooling(avg_pattern='BS D t h w -> BS D t',
-                                                           then_permute_pattern='BS D t -> BS t D')
-                if agg_time_module == 'TransformerEncoderLayer':
-                    self.temp_attn_agg = TemporalTransformerEncoderLayer(**transf_enc_layer_kwargs)
-                elif agg_time_module == 'AveragePooling':
-                    self.temp_attn_agg = AveragePooling(avg_pattern='BS t D -> BS D')
-                elif 'Identity' in agg_time_module:
-                    self.temp_attn_agg = nn.Identity()
-            # define a global aggregation layer (aggregarate over segments)
-            self.add_global_repr = add_global_repr
-            if add_global_repr:
-                if agg_segments_module == 'TransformerEncoderLayer':
-                    # we can reuse the same layer as for temporal factorization (B, dim_to_agg, D) -> (B, D)
-                    # we need to add pos emb (PE) because previously we added the same PE for each segment
-                    pos_max_len = max_segments if max_segments is not None else 16  # 16 = 10sec//0.64sec + 1
-                    self.global_attn_agg = TemporalTransformerEncoderLayer(add_pos_emb=True,
-                                                                           pos_emb_drop=mformer_cfg.VIT.POS_DROPOUT,
-                                                                           pos_max_len=pos_max_len,
-                                                                           **transf_enc_layer_kwargs)
-                elif agg_segments_module == 'AveragePooling':
-                    self.global_attn_agg = AveragePooling(avg_pattern='B S D -> B D')
-
-        if was_pt_on_avclip:
-            # we need to filter out the state_dict of the AVCLIP model (has both A and V extractors)
-            # and keep only the state_dict of the feat extractor
-            ckpt_weights = dict()
-            for k, v in ckpt['state_dict'].items():
-                if k.startswith(('module.v_encoder.', 'v_encoder.')):
-                    k = k.replace('module.', '').replace('v_encoder.', '')
-                    ckpt_weights[k] = v
-            _load_status = self.load_state_dict(ckpt_weights, strict=False)
-            if len(_load_status.missing_keys) > 0 or len(_load_status.unexpected_keys) > 0:
-                logging.warning(f'Loading exact vfeat_extractor ckpt from {self.ckpt_path} failed. \n' \
-                                f'Missing keys ({len(_load_status.missing_keys)}): ' \
-                                f'{_load_status.missing_keys}, \n' \
-                                f'Unexpected keys ({len(_load_status.unexpected_keys)}): ' \
-                                f'{_load_status.unexpected_keys} \n' \
-                                f'temp_attn_agg are expected to be missing if ckpt was pt contrastively.')
-            else:
-                logging.info(f'Loading vfeat_extractor ckpt from {self.ckpt_path} succeeded.')
-
-        # patch_embed is not used in MotionFormer, only patch_embed_3d, because cfg.VIT.PATCH_SIZE_TEMP > 1
-        # but it used to calculate the number of patches, so we need to set keep it
-        self.patch_embed.requires_grad_(False)
-
-    def forward(self, x):
-        '''
-        x is of shape (B, S, C, T, H, W) where S is the number of segments.
-        '''
-        # Batch, Segments, Channels, T=frames, Height, Width
-        B, S, C, T, H, W = x.shape
-        # Motionformer expects a tensor of shape (1, B, C, T, H, W).
-        # The first dimension (1) is a dummy dimension to make the input tensor and won't be used:
-        # see `video_model_builder.video_input`.
-        # x = x.unsqueeze(0)  # (1, B, S, C, T, H, W)
-
-        orig_shape = (B, S, C, T, H, W)
-        x = x.view(B * S, C, T, H, W)  # flatten batch and segments
-        x = self.forward_segments(x, orig_shape=orig_shape)
-        # unpack the segments (using rest dimensions to support different shapes e.g. (BS, D) or (BS, t, D))
-        x = x.view(B, S, *x.shape[1:])
-        # x is now of shape (B*S, D) or (B*S, t, D) if `self.temp_attn_agg` is `Identity`
-
-        return x  # x is (B, S, ...)
-
-    def forward_segments(self, x, orig_shape: tuple) -> torch.Tensor:
-        '''x is of shape (1, BS, C, T, H, W) where S is the number of segments.'''
-        x, x_mask = self.forward_features(x)
-
-        assert self.extract_features
-
-        # (BS, T, D) where T = 1 + (224 // 16) * (224 // 16) * 8
-        x = x[:, 1:, :]  # without the CLS token for efficiency (should be safe for LayerNorm and FC)
-        x = self.norm(x)
-        x = self.pre_logits(x)
-        if self.factorize_space_time:
-            x = self.restore_spatio_temp_dims(x, orig_shape)  # (B*S, D, t, h, w) <- (B*S, t*h*w, D)
-
-            x = self.spatial_attn_agg(x, x_mask)  # (B*S, t, D)
-            x = self.temp_attn_agg(x)  # (B*S, D) or (BS, t, D) if `self.temp_attn_agg` is `Identity`
-
-        return x
-
-    def restore_spatio_temp_dims(self, feats: torch.Tensor, orig_shape: tuple) -> torch.Tensor:
-        '''
-            feats are of shape (B*S, T, D) where T = 1 + (224 // 16) * (224 // 16) * 8
-            Our goal is to make them of shape (B*S, t, h, w, D) where h, w are the spatial dimensions.
-            From `self.patch_embed_3d`, it follows that we could reshape feats with:
-                `feats.transpose(1, 2).view(B*S, D, t, h, w)`
-        '''
-        B, S, C, T, H, W = orig_shape
-        D = self.embed_dim
-
-        # num patches in each dimension
-        t = T // self.patch_embed_3d.z_block_size
-        h = self.patch_embed_3d.height
-        w = self.patch_embed_3d.width
-
-        feats = feats.permute(0, 2, 1)  # (B*S, D, T)
-        feats = feats.view(B * S, D, t, h, w)  # (B*S, D, t, h, w)
-
-        return feats
-
-
-class BaseEncoderLayer(nn.TransformerEncoderLayer):
-    '''
-        This is a wrapper around nn.TransformerEncoderLayer that adds a CLS token
-        to the sequence and outputs the CLS token's representation.
-        This base class parents both SpatialEncoderLayer and TemporalEncoderLayer for the RGB stream
-        and the FrequencyEncoderLayer and TemporalEncoderLayer for the audio stream stream.
-        We also, optionally, add a positional embedding to the input sequence which
-        allows to reuse it for global aggregation (of segments) for both streams.
-    '''
-
-    def __init__(self,
-                 add_pos_emb: bool = False,
-                 pos_emb_drop: float = None,
-                 pos_max_len: int = None,
-                 *args_transformer_enc,
-                 **kwargs_transformer_enc):
-        super().__init__(*args_transformer_enc, **kwargs_transformer_enc)
-        self.cls_token = nn.Parameter(torch.zeros(1, 1, self.self_attn.embed_dim))
-        trunc_normal_(self.cls_token, std=.02)
-
-        # add positional embedding
-        self.add_pos_emb = add_pos_emb
-        if add_pos_emb:
-            self.pos_max_len = 1 + pos_max_len  # +1 (for CLS)
-            self.pos_emb = nn.Parameter(torch.zeros(1, self.pos_max_len, self.self_attn.embed_dim))
-            self.pos_drop = nn.Dropout(pos_emb_drop)
-            trunc_normal_(self.pos_emb, std=.02)
-
-        self.apply(self._init_weights)
-
-    def forward(self, x: torch.Tensor, x_mask: torch.Tensor = None):
-        ''' x is of shape (B, N, D); if provided x_mask is of shape (B, N)'''
-        batch_dim = x.shape[0]
-
-        # add CLS token
-        cls_tokens = self.cls_token.expand(batch_dim, -1, -1)  # expanding to match batch dimension
-        x = torch.cat((cls_tokens, x), dim=-2)  # (batch_dim, 1+seq_len, D)
-        if x_mask is not None:
-            cls_mask = torch.ones((batch_dim, 1), dtype=torch.bool, device=x_mask.device)  # 1=keep; 0=mask
-            x_mask_w_cls = torch.cat((cls_mask, x_mask), dim=-1)  # (batch_dim, 1+seq_len)
-            B, N = x_mask_w_cls.shape
-            # torch expects (N, N) or (B*num_heads, N, N) mask (sadness ahead); torch masks
-            x_mask_w_cls = x_mask_w_cls.reshape(B, 1, 1, N)\
-                                       .expand(-1, self.self_attn.num_heads, N, -1)\
-                                       .reshape(B * self.self_attn.num_heads, N, N)
-            assert x_mask_w_cls.dtype == x_mask_w_cls.bool().dtype, 'x_mask_w_cls.dtype != bool'
-            x_mask_w_cls = ~x_mask_w_cls  # invert mask (1=mask)
-        else:
-            x_mask_w_cls = None
-
-        # add positional embedding
-        if self.add_pos_emb:
-            seq_len = x.shape[1]  # (don't even think about moving it before the CLS token concatenation)
-            assert seq_len <= self.pos_max_len, f'Seq len ({seq_len}) > pos_max_len ({self.pos_max_len})'
-            x = x + self.pos_emb[:, :seq_len, :]
-            x = self.pos_drop(x)
-
-        # apply encoder layer (calls nn.TransformerEncoderLayer.forward);
-        x = super().forward(src=x, src_mask=x_mask_w_cls)  # (batch_dim, 1+seq_len, D)
-
-        # CLS token is expected to hold spatial information for each frame
-        x = x[:, 0, :]  # (batch_dim, D)
-
-        return x
-
-    def _init_weights(self, m):
-        if isinstance(m, nn.Linear):
-            trunc_normal_(m.weight, std=.02)
-            if isinstance(m, nn.Linear) and m.bias is not None:
-                nn.init.constant_(m.bias, 0)
-        elif isinstance(m, nn.LayerNorm):
-            nn.init.constant_(m.bias, 0)
-            nn.init.constant_(m.weight, 1.0)
-
-    @torch.jit.ignore
-    def no_weight_decay(self):
-        return {'cls_token', 'pos_emb'}
-
-
-class SpatialTransformerEncoderLayer(BaseEncoderLayer):
-    ''' Aggregates spatial dimensions by applying attention individually to each frame. '''
-
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-
-    def forward(self, x: torch.Tensor, x_mask: torch.Tensor = None) -> torch.Tensor:
-        ''' x is of shape (B*S, D, t, h, w) where S is the number of segments.
-            if specified x_mask (B*S, t, h, w), 0=masked, 1=kept
-            Returns a tensor of shape (B*S, t, D) pooling spatial information for each frame. '''
-        BS, D, t, h, w = x.shape
-
-        # time as a batch dimension and flatten spatial dimensions as sequence
-        x = einops.rearrange(x, 'BS D t h w -> (BS t) (h w) D')
-        # similar to mask
-        if x_mask is not None:
-            x_mask = einops.rearrange(x_mask, 'BS t h w -> (BS t) (h w)')
-
-        # apply encoder layer (BaseEncoderLayer.forward) - it will add CLS token and output its representation
-        x = super().forward(x=x, x_mask=x_mask)  # (B*S*t, D)
-
-        # reshape back to (B*S, t, D)
-        x = einops.rearrange(x, '(BS t) D -> BS t D', BS=BS, t=t)
-
-        # (B*S, t, D)
-        return x
-
-
-class TemporalTransformerEncoderLayer(BaseEncoderLayer):
-    ''' Aggregates temporal dimension with attention. Also used with pos emb as global aggregation
-    in both streams. '''
-
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-
-    def forward(self, x):
-        ''' x is of shape (B*S, t, D) where S is the number of segments.
-            Returns a tensor of shape (B*S, D) pooling temporal information. '''
-        BS, t, D = x.shape
-
-        # apply encoder layer (BaseEncoderLayer.forward) - it will add CLS token and output its representation
-        x = super().forward(x)  # (B*S, D)
-
-        return x  # (B*S, D)
-
-
-class AveragePooling(nn.Module):
-
-    def __init__(self, avg_pattern: str, then_permute_pattern: str = None) -> None:
-        ''' patterns are e.g. "bs t d -> bs d" '''
-        super().__init__()
-        # TODO: need to register them as buffers (but fails because these are strings)
-        self.reduce_fn = 'mean'
-        self.avg_pattern = avg_pattern
-        self.then_permute_pattern = then_permute_pattern
-
-    def forward(self, x: torch.Tensor, x_mask: torch.Tensor = None) -> torch.Tensor:
-        x = einops.reduce(x, self.avg_pattern, self.reduce_fn)
-        if self.then_permute_pattern is not None:
-            x = einops.rearrange(x, self.then_permute_pattern)
-        return x
+from fastvideo.third_party.synchformer.motionformer import *  # noqa: F403

fastvideo/third_party/eval/synchformer/synchformer.py

@@ -5,7 +5,7 @@ import torch
 from torch import nn
 
 from fastvideo.third_party.eval.synchformer.ast import AST
-from fastvideo.third_party.eval.synchformer.motionformer import MotionFormer
+from fastvideo.third_party.synchformer.motionformer import MotionFormer
 from fastvideo.third_party.eval.synchformer.transformer import GlobalTransformer
 
 

fastvideo/third_party/eval/synchformer/utils.py

@@ -1,69 +1,3 @@
-from hashlib import md5
-from pathlib import Path
+"""Backward-compatible import for shared Synchformer utilities."""
 
-import requests
-from tqdm import tqdm
-
-PARENT_LINK = 'https://a3s.fi/swift/v1/AUTH_a235c0f452d648828f745589cde1219a'
-FNAME2LINK = {
-    # S3: Synchability: AudioSet (run 2)
-    '24-01-22T20-34-52.pt': f'{PARENT_LINK}/sync/sync_models/24-01-22T20-34-52/24-01-22T20-34-52.pt',
-    'cfg-24-01-22T20-34-52.yaml': f'{PARENT_LINK}/sync/sync_models/24-01-22T20-34-52/cfg-24-01-22T20-34-52.yaml',
-    # S2: Synchformer: AudioSet (run 2)
-    '24-01-04T16-39-21.pt': f'{PARENT_LINK}/sync/sync_models/24-01-04T16-39-21/24-01-04T16-39-21.pt',
-    'cfg-24-01-04T16-39-21.yaml': f'{PARENT_LINK}/sync/sync_models/24-01-04T16-39-21/cfg-24-01-04T16-39-21.yaml',
-    # S2: Synchformer: AudioSet (run 1)
-    '23-08-28T11-23-23.pt': f'{PARENT_LINK}/sync/sync_models/23-08-28T11-23-23/23-08-28T11-23-23.pt',
-    'cfg-23-08-28T11-23-23.yaml': f'{PARENT_LINK}/sync/sync_models/23-08-28T11-23-23/cfg-23-08-28T11-23-23.yaml',
-    # S2: Synchformer: LRS3 (run 2)
-    '23-12-23T18-33-57.pt': f'{PARENT_LINK}/sync/sync_models/23-12-23T18-33-57/23-12-23T18-33-57.pt',
-    'cfg-23-12-23T18-33-57.yaml': f'{PARENT_LINK}/sync/sync_models/23-12-23T18-33-57/cfg-23-12-23T18-33-57.yaml',
-    # S2: Synchformer: VGS (run 2)
-    '24-01-02T10-00-53.pt': f'{PARENT_LINK}/sync/sync_models/24-01-02T10-00-53/24-01-02T10-00-53.pt',
-    'cfg-24-01-02T10-00-53.yaml': f'{PARENT_LINK}/sync/sync_models/24-01-02T10-00-53/cfg-24-01-02T10-00-53.yaml',
-    # SparseSync: ft VGGSound-Full
-    '22-09-21T21-00-52.pt': f'{PARENT_LINK}/sync/sync_models/22-09-21T21-00-52/22-09-21T21-00-52.pt',
-    'cfg-22-09-21T21-00-52.yaml': f'{PARENT_LINK}/sync/sync_models/22-09-21T21-00-52/cfg-22-09-21T21-00-52.yaml',
-    # SparseSync: ft VGGSound-Sparse
-    '22-07-28T15-49-45.pt': f'{PARENT_LINK}/sync/sync_models/22-07-28T15-49-45/22-07-28T15-49-45.pt',
-    'cfg-22-07-28T15-49-45.yaml': f'{PARENT_LINK}/sync/sync_models/22-07-28T15-49-45/cfg-22-07-28T15-49-45.yaml',
-    # SparseSync: only pt on LRS3
-    '22-07-13T22-25-49.pt': f'{PARENT_LINK}/sync/sync_models/22-07-13T22-25-49/22-07-13T22-25-49.pt',
-    'cfg-22-07-13T22-25-49.yaml': f'{PARENT_LINK}/sync/sync_models/22-07-13T22-25-49/cfg-22-07-13T22-25-49.yaml',
-    # SparseSync: feature extractors
-    'ResNetAudio-22-08-04T09-51-04.pt': f'{PARENT_LINK}/sync/ResNetAudio-22-08-04T09-51-04.pt',  # 2s
-    'ResNetAudio-22-08-03T23-14-49.pt': f'{PARENT_LINK}/sync/ResNetAudio-22-08-03T23-14-49.pt',  # 3s
-    'ResNetAudio-22-08-03T23-14-28.pt': f'{PARENT_LINK}/sync/ResNetAudio-22-08-03T23-14-28.pt',  # 4s
-    'ResNetAudio-22-06-24T08-10-33.pt': f'{PARENT_LINK}/sync/ResNetAudio-22-06-24T08-10-33.pt',  # 5s
-    'ResNetAudio-22-06-24T17-31-07.pt': f'{PARENT_LINK}/sync/ResNetAudio-22-06-24T17-31-07.pt',  # 6s
-    'ResNetAudio-22-06-24T23-57-11.pt': f'{PARENT_LINK}/sync/ResNetAudio-22-06-24T23-57-11.pt',  # 7s
-    'ResNetAudio-22-06-25T04-35-42.pt': f'{PARENT_LINK}/sync/ResNetAudio-22-06-25T04-35-42.pt',  # 8s
-}
-
-
-def check_if_file_exists_else_download(path, fname2link=FNAME2LINK, chunk_size=1024):
-    '''Checks if file exists, if not downloads it from the link to the path'''
-    path = Path(path)
-    if not path.exists():
-        path.parent.mkdir(exist_ok=True, parents=True)
-        link = fname2link.get(path.name, None)
-        if link is None:
-            raise ValueError(f'Cant find the checkpoint file: {path}.',
-                             f'Please download it manually and ensure the path exists.')
-        with requests.get(fname2link[path.name], stream=True) as r:
-            total_size = int(r.headers.get('content-length', 0))
-            with tqdm(total=total_size, unit='B', unit_scale=True) as pbar:
-                with open(path, 'wb') as f:
-                    for data in r.iter_content(chunk_size=chunk_size):
-                        if data:
-                            f.write(data)
-                            pbar.update(chunk_size)
-
-
-def get_md5sum(path):
-    hash_md5 = md5()
-    with open(path, 'rb') as f:
-        for chunk in iter(lambda: f.read(4096 * 8), b''):
-            hash_md5.update(chunk)
-    md5sum = hash_md5.hexdigest()
-    return md5sum
+from fastvideo.third_party.synchformer.utils import *  # noqa: F403

fastvideo/third_party/eval/synchformer/video_model_builder.py

@@ -1,272 +1,3 @@
-#!/usr/bin/env python3
-# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
-# Copyright 2020 Ross Wightman
-# Modified Model definition
+"""Backward-compatible import for shared Synchformer model builders."""
 
-from collections import OrderedDict
-from functools import partial
-
-import torch
-import torch.nn as nn
-from timm.layers import trunc_normal_
-
-from fastvideo.third_party.eval.synchformer import vit_helper
-
-
-class VisionTransformer(nn.Module):
-    """ Vision Transformer with support for patch or hybrid CNN input stage """
-
-    def __init__(self, cfg):
-        super().__init__()
-        self.img_size = cfg.DATA.TRAIN_CROP_SIZE
-        self.patch_size = cfg.VIT.PATCH_SIZE
-        self.in_chans = cfg.VIT.CHANNELS
-        if cfg.TRAIN.DATASET == "Epickitchens":
-            self.num_classes = [97, 300]
-        else:
-            self.num_classes = cfg.MODEL.NUM_CLASSES
-        self.embed_dim = cfg.VIT.EMBED_DIM
-        self.depth = cfg.VIT.DEPTH
-        self.num_heads = cfg.VIT.NUM_HEADS
-        self.mlp_ratio = cfg.VIT.MLP_RATIO
-        self.qkv_bias = cfg.VIT.QKV_BIAS
-        self.drop_rate = cfg.VIT.DROP
-        self.drop_path_rate = cfg.VIT.DROP_PATH
-        self.head_dropout = cfg.VIT.HEAD_DROPOUT
-        self.video_input = cfg.VIT.VIDEO_INPUT
-        self.temporal_resolution = cfg.VIT.TEMPORAL_RESOLUTION
-        self.use_mlp = cfg.VIT.USE_MLP
-        self.num_features = self.embed_dim
-        norm_layer = partial(nn.LayerNorm, eps=1e-6)
-        self.attn_drop_rate = cfg.VIT.ATTN_DROPOUT
-        self.head_act = cfg.VIT.HEAD_ACT
-        self.cfg = cfg
-
-        # Patch Embedding
-        self.patch_embed = vit_helper.PatchEmbed(img_size=224,
-                                                 patch_size=self.patch_size,
-                                                 in_chans=self.in_chans,
-                                                 embed_dim=self.embed_dim)
-
-        # 3D Patch Embedding
-        self.patch_embed_3d = vit_helper.PatchEmbed3D(img_size=self.img_size,
-                                                      temporal_resolution=self.temporal_resolution,
-                                                      patch_size=self.patch_size,
-                                                      in_chans=self.in_chans,
-                                                      embed_dim=self.embed_dim,
-                                                      z_block_size=self.cfg.VIT.PATCH_SIZE_TEMP)
-        self.patch_embed_3d.proj.weight.data = torch.zeros_like(self.patch_embed_3d.proj.weight.data)
-
-        # Number of patches
-        if self.video_input:
-            num_patches = self.patch_embed.num_patches * self.temporal_resolution
-        else:
-            num_patches = self.patch_embed.num_patches
-        self.num_patches = num_patches
-
-        # CLS token
-        self.cls_token = nn.Parameter(torch.zeros(1, 1, self.embed_dim))
-        trunc_normal_(self.cls_token, std=.02)
-
-        # Positional embedding
-        self.pos_embed = nn.Parameter(torch.zeros(1, self.patch_embed.num_patches + 1, self.embed_dim))
-        self.pos_drop = nn.Dropout(p=cfg.VIT.POS_DROPOUT)
-        trunc_normal_(self.pos_embed, std=.02)
-
-        if self.cfg.VIT.POS_EMBED == "joint":
-            self.st_embed = nn.Parameter(torch.zeros(1, num_patches + 1, self.embed_dim))
-            trunc_normal_(self.st_embed, std=.02)
-        elif self.cfg.VIT.POS_EMBED == "separate":
-            self.temp_embed = nn.Parameter(torch.zeros(1, self.temporal_resolution, self.embed_dim))
-
-        # Layer Blocks
-        dpr = [x.item() for x in torch.linspace(0, self.drop_path_rate, self.depth)]
-        if self.cfg.VIT.ATTN_LAYER == "divided":
-            self.blocks = nn.ModuleList([
-                vit_helper.DividedSpaceTimeBlock(
-                    attn_type=cfg.VIT.ATTN_LAYER,
-                    dim=self.embed_dim,
-                    num_heads=self.num_heads,
-                    mlp_ratio=self.mlp_ratio,
-                    qkv_bias=self.qkv_bias,
-                    drop=self.drop_rate,
-                    attn_drop=self.attn_drop_rate,
-                    drop_path=dpr[i],
-                    norm_layer=norm_layer,
-                ) for i in range(self.depth)
-            ])
-        else:
-            self.blocks = nn.ModuleList([
-                vit_helper.Block(attn_type=cfg.VIT.ATTN_LAYER,
-                                 dim=self.embed_dim,
-                                 num_heads=self.num_heads,
-                                 mlp_ratio=self.mlp_ratio,
-                                 qkv_bias=self.qkv_bias,
-                                 drop=self.drop_rate,
-                                 attn_drop=self.attn_drop_rate,
-                                 drop_path=dpr[i],
-                                 norm_layer=norm_layer,
-                                 use_original_code=self.cfg.VIT.USE_ORIGINAL_TRAJ_ATTN_CODE) for i in range(self.depth)
-            ])
-        self.norm = norm_layer(self.embed_dim)
-
-        # MLP head
-        if self.use_mlp:
-            hidden_dim = self.embed_dim
-            if self.head_act == 'tanh':
-                # logging.info("Using TanH activation in MLP")
-                act = nn.Tanh()
-            elif self.head_act == 'gelu':
-                # logging.info("Using GELU activation in MLP")
-                act = nn.GELU()
-            else:
-                # logging.info("Using ReLU activation in MLP")
-                act = nn.ReLU()
-            self.pre_logits = nn.Sequential(OrderedDict([
-                ('fc', nn.Linear(self.embed_dim, hidden_dim)),
-                ('act', act),
-            ]))
-        else:
-            self.pre_logits = nn.Identity()
-
-        # Classifier Head
-        self.head_drop = nn.Dropout(p=self.head_dropout)
-        if isinstance(self.num_classes, (list, )) and len(self.num_classes) > 1:
-            for a, i in enumerate(range(len(self.num_classes))):
-                setattr(self, "head%d" % a, nn.Linear(self.embed_dim, self.num_classes[i]))
-        else:
-            self.head = nn.Linear(self.embed_dim, self.num_classes) if self.num_classes > 0 else nn.Identity()
-
-        # Initialize weights
-        self.apply(self._init_weights)
-
-    def _init_weights(self, m):
-        if isinstance(m, nn.Linear):
-            trunc_normal_(m.weight, std=.02)
-            if isinstance(m, nn.Linear) and m.bias is not None:
-                nn.init.constant_(m.bias, 0)
-        elif isinstance(m, nn.LayerNorm):
-            nn.init.constant_(m.bias, 0)
-            nn.init.constant_(m.weight, 1.0)
-
-    @torch.jit.ignore
-    def no_weight_decay(self):
-        if self.cfg.VIT.POS_EMBED == "joint":
-            return {'pos_embed', 'cls_token', 'st_embed'}
-        else:
-            return {'pos_embed', 'cls_token', 'temp_embed'}
-
-    def get_classifier(self):
-        return self.head
-
-    def reset_classifier(self, num_classes, global_pool=''):
-        self.num_classes = num_classes
-        self.head = (nn.Linear(self.embed_dim, num_classes) if num_classes > 0 else nn.Identity())
-
-    def forward_features(self, x):
-        # if self.video_input:
-        # x = x[0]
-        B = x.shape[0]
-
-        # Tokenize input
-        # if self.cfg.VIT.PATCH_SIZE_TEMP > 1:
-        # for simplicity of mapping between content dimensions (input x) and token dims (after patching)
-        # we use the same trick as for AST (see modeling_ast.ASTModel.forward for the details):
-
-        # apply patching on input
-        x = self.patch_embed_3d(x)
-        tok_mask = None
-
-        # else:
-        #     tok_mask = None
-        #     # 2D tokenization
-        #     if self.video_input:
-        #         x = x.permute(0, 2, 1, 3, 4)
-        #         (B, T, C, H, W) = x.shape
-        #         x = x.reshape(B * T, C, H, W)
-
-        #     x = self.patch_embed(x)
-
-        #     if self.video_input:
-        #         (B2, T2, D2) = x.shape
-        #         x = x.reshape(B, T * T2, D2)
-
-        # Append CLS token
-        cls_tokens = self.cls_token.expand(B, -1, -1)
-        x = torch.cat((cls_tokens, x), dim=1)
-        # if tok_mask is not None:
-        #     # prepend 1(=keep) to the mask to account for the CLS token as well
-        #     tok_mask = torch.cat((torch.ones_like(tok_mask[:, [0]]), tok_mask), dim=1)
-
-        # Interpolate positinoal embeddings
-        # if self.cfg.DATA.TRAIN_CROP_SIZE != 224:
-        #     pos_embed = self.pos_embed
-        #     N = pos_embed.shape[1] - 1
-        #     npatch = int((x.size(1) - 1) / self.temporal_resolution)
-        #     class_emb = pos_embed[:, 0]
-        #     pos_embed = pos_embed[:, 1:]
-        #     dim = x.shape[-1]
-        #     pos_embed = torch.nn.functional.interpolate(
-        #         pos_embed.reshape(1, int(math.sqrt(N)), int(math.sqrt(N)), dim).permute(0, 3, 1, 2),
-        #         scale_factor=math.sqrt(npatch / N),
-        #         mode='bicubic',
-        #     )
-        #     pos_embed = pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
-        #     new_pos_embed = torch.cat((class_emb.unsqueeze(0), pos_embed), dim=1)
-        # else:
-        new_pos_embed = self.pos_embed
-        npatch = self.patch_embed.num_patches
-
-        # Add positional embeddings to input
-        if self.video_input:
-            if self.cfg.VIT.POS_EMBED == "separate":
-                cls_embed = self.pos_embed[:, 0, :].unsqueeze(1)
-                tile_pos_embed = new_pos_embed[:, 1:, :].repeat(1, self.temporal_resolution, 1)
-                tile_temporal_embed = self.temp_embed.repeat_interleave(npatch, 1)
-                total_pos_embed = tile_pos_embed + tile_temporal_embed
-                total_pos_embed = torch.cat([cls_embed, total_pos_embed], dim=1)
-                x = x + total_pos_embed
-            elif self.cfg.VIT.POS_EMBED == "joint":
-                x = x + self.st_embed
-        else:
-            # image input
-            x = x + new_pos_embed
-
-        # Apply positional dropout
-        x = self.pos_drop(x)
-
-        # Encoding using transformer layers
-        for i, blk in enumerate(self.blocks):
-            x = blk(x,
-                    seq_len=npatch,
-                    num_frames=self.temporal_resolution,
-                    approx=self.cfg.VIT.APPROX_ATTN_TYPE,
-                    num_landmarks=self.cfg.VIT.APPROX_ATTN_DIM,
-                    tok_mask=tok_mask)
-
-        ### v-iashin: I moved it to the forward pass
-        # x = self.norm(x)[:, 0]
-        # x = self.pre_logits(x)
-        ###
-        return x, tok_mask
-
-    # def forward(self, x):
-    #     x = self.forward_features(x)
-    #     ### v-iashin: here. This should leave the same forward output as before
-    #     x = self.norm(x)[:, 0]
-    #     x = self.pre_logits(x)
-    #     ###
-    #     x = self.head_drop(x)
-    #     if isinstance(self.num_classes, (list, )) and len(self.num_classes) > 1:
-    #         output = []
-    #         for head in range(len(self.num_classes)):
-    #             x_out = getattr(self, "head%d" % head)(x)
-    #             if not self.training:
-    #                 x_out = torch.nn.functional.softmax(x_out, dim=-1)
-    #             output.append(x_out)
-    #         return output
-    #     else:
-    #         x = self.head(x)
-    #         if not self.training:
-    #             x = torch.nn.functional.softmax(x, dim=-1)
-    #         return x
+from fastvideo.third_party.synchformer.video_model_builder import *  # noqa: F403

fastvideo/third_party/eval/synchformer/vit_helper.py

@@ -1,362 +1,3 @@
-#!/usr/bin/env python3
-# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
-# Copyright 2020 Ross Wightman
-# Modified Model definition
-"""Video models."""
+"""Backward-compatible import for shared Synchformer ViT helpers."""
 
-import math
-
-import torch
-import torch.nn as nn
-from einops import rearrange, repeat
-from timm.layers import to_2tuple
-from torch import einsum
-from torch.nn import functional as F
-
-default_cfgs = {
-    'vit_1k':
-    'https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-vitjx/jx_vit_base_p16_224-80ecf9dd.pth',
-    'vit_1k_large':
-    'https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-vitjx/jx_vit_large_p16_224-4ee7a4dc.pth',
-}
-
-
-def qkv_attn(q, k, v, tok_mask: torch.Tensor = None):
-    sim = einsum('b i d, b j d -> b i j', q, k)
-    # apply masking if provided, tok_mask is (B*S*H, N): 1s - keep; sim is (B*S*H, H, N, N)
-    if tok_mask is not None:
-        BSH, N = tok_mask.shape
-        sim = sim.masked_fill(tok_mask.view(BSH, 1, N) == 0, float('-inf'))  # 1 - broadcasts across N
-    attn = sim.softmax(dim=-1)
-    out = einsum('b i j, b j d -> b i d', attn, v)
-    return out
-
-
-class DividedAttention(nn.Module):
-
-    def __init__(self, dim, num_heads=8, qkv_bias=False, attn_drop=0., proj_drop=0.):
-        super().__init__()
-        self.num_heads = num_heads
-        head_dim = dim // num_heads
-        self.scale = head_dim**-0.5
-        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
-        self.proj = nn.Linear(dim, dim)
-
-        # init to zeros
-        self.qkv.weight.data.fill_(0)
-        self.qkv.bias.data.fill_(0)
-        self.proj.weight.data.fill_(1)
-        self.proj.bias.data.fill_(0)
-
-        self.attn_drop = nn.Dropout(attn_drop)
-        self.proj_drop = nn.Dropout(proj_drop)
-
-    def forward(self, x, einops_from, einops_to, tok_mask: torch.Tensor = None, **einops_dims):
-        # num of heads variable
-        h = self.num_heads
-
-        # project x to q, k, v vaalues
-        q, k, v = self.qkv(x).chunk(3, dim=-1)
-        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (q, k, v))
-        if tok_mask is not None:
-            # replicate token mask across heads (b, n) -> (b, h, n) -> (b*h, n) -- same as qkv but w/o d
-            assert len(tok_mask.shape) == 2
-            tok_mask = tok_mask.unsqueeze(1).expand(-1, h, -1).reshape(-1, tok_mask.shape[1])
-
-        # Scale q
-        q *= self.scale
-
-        # Take out cls_q, cls_k, cls_v
-        (cls_q, q_), (cls_k, k_), (cls_v, v_) = map(lambda t: (t[:, 0:1], t[:, 1:]), (q, k, v))
-        # the same for masking
-        if tok_mask is not None:
-            cls_mask, mask_ = tok_mask[:, 0:1], tok_mask[:, 1:]
-        else:
-            cls_mask, mask_ = None, None
-
-        # let CLS token attend to key / values of all patches across time and space
-        cls_out = qkv_attn(cls_q, k, v, tok_mask=tok_mask)
-
-        # rearrange across time or space
-        q_, k_, v_ = map(lambda t: rearrange(t, f'{einops_from} -> {einops_to}', **einops_dims), (q_, k_, v_))
-
-        # expand CLS token keys and values across time or space and concat
-        r = q_.shape[0] // cls_k.shape[0]
-        cls_k, cls_v = map(lambda t: repeat(t, 'b () d -> (b r) () d', r=r), (cls_k, cls_v))
-
-        k_ = torch.cat((cls_k, k_), dim=1)
-        v_ = torch.cat((cls_v, v_), dim=1)
-
-        # the same for masking (if provided)
-        if tok_mask is not None:
-            # since mask does not have the latent dim (d), we need to remove it from einops dims
-            mask_ = rearrange(mask_, f'{einops_from} -> {einops_to}'.replace(' d', ''), **einops_dims)
-            cls_mask = repeat(cls_mask, 'b () -> (b r) ()', r=r)  # expand cls_mask across time or space
-            mask_ = torch.cat((cls_mask, mask_), dim=1)
-
-        # attention
-        out = qkv_attn(q_, k_, v_, tok_mask=mask_)
-
-        # merge back time or space
-        out = rearrange(out, f'{einops_to} -> {einops_from}', **einops_dims)
-
-        # concat back the cls token
-        out = torch.cat((cls_out, out), dim=1)
-
-        # merge back the heads
-        out = rearrange(out, '(b h) n d -> b n (h d)', h=h)
-
-        ## to out
-        x = self.proj(out)
-        x = self.proj_drop(x)
-        return x
-
-
-class DividedSpaceTimeBlock(nn.Module):
-
-    def __init__(self,
-                 dim=768,
-                 num_heads=12,
-                 attn_type='divided',
-                 mlp_ratio=4.,
-                 qkv_bias=False,
-                 drop=0.,
-                 attn_drop=0.,
-                 drop_path=0.,
-                 act_layer=nn.GELU,
-                 norm_layer=nn.LayerNorm):
-        super().__init__()
-
-        self.einops_from_space = 'b (f n) d'
-        self.einops_to_space = '(b f) n d'
-        self.einops_from_time = 'b (f n) d'
-        self.einops_to_time = '(b n) f d'
-
-        self.norm1 = norm_layer(dim)
-
-        self.attn = DividedAttention(dim, num_heads=num_heads, qkv_bias=qkv_bias, attn_drop=attn_drop, proj_drop=drop)
-
-        self.timeattn = DividedAttention(dim,
-                                         num_heads=num_heads,
-                                         qkv_bias=qkv_bias,
-                                         attn_drop=attn_drop,
-                                         proj_drop=drop)
-
-        # self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
-        self.drop_path = nn.Identity()
-        self.norm2 = norm_layer(dim)
-        mlp_hidden_dim = int(dim * mlp_ratio)
-        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
-        self.norm3 = norm_layer(dim)
-
-    def forward(self, x, seq_len=196, num_frames=8, approx='none', num_landmarks=128, tok_mask: torch.Tensor = None):
-        time_output = self.timeattn(self.norm3(x),
-                                    self.einops_from_time,
-                                    self.einops_to_time,
-                                    n=seq_len,
-                                    tok_mask=tok_mask)
-        time_residual = x + time_output
-
-        space_output = self.attn(self.norm1(time_residual),
-                                 self.einops_from_space,
-                                 self.einops_to_space,
-                                 f=num_frames,
-                                 tok_mask=tok_mask)
-        space_residual = time_residual + self.drop_path(space_output)
-
-        x = space_residual
-        x = x + self.drop_path(self.mlp(self.norm2(x)))
-        return x
-
-
-class Mlp(nn.Module):
-
-    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
-        super().__init__()
-        out_features = out_features or in_features
-        hidden_features = hidden_features or in_features
-        self.fc1 = nn.Linear(in_features, hidden_features)
-        self.act = act_layer()
-        self.fc2 = nn.Linear(hidden_features, out_features)
-        self.drop = nn.Dropout(drop)
-
-    def forward(self, x):
-        x = self.fc1(x)
-        x = self.act(x)
-        x = self.drop(x)
-        x = self.fc2(x)
-        x = self.drop(x)
-        return x
-
-
-class PatchEmbed(nn.Module):
-    """ Image to Patch Embedding
-    """
-
-    def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768):
-        super().__init__()
-        img_size = img_size if type(img_size) is tuple else to_2tuple(img_size)
-        patch_size = img_size if type(patch_size) is tuple else to_2tuple(patch_size)
-        num_patches = (img_size[1] // patch_size[1]) * (img_size[0] // patch_size[0])
-        self.img_size = img_size
-        self.patch_size = patch_size
-        self.num_patches = num_patches
-
-        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
-
-    def forward(self, x):
-        B, C, H, W = x.shape
-        x = self.proj(x).flatten(2).transpose(1, 2)
-        return x
-
-
-class PatchEmbed3D(nn.Module):
-    """ Image to Patch Embedding """
-
-    def __init__(self,
-                 img_size=224,
-                 temporal_resolution=4,
-                 in_chans=3,
-                 patch_size=16,
-                 z_block_size=2,
-                 embed_dim=768,
-                 flatten=True):
-        super().__init__()
-        self.height = (img_size // patch_size)
-        self.width = (img_size // patch_size)
-        ### v-iashin: these two are incorrect
-        # self.frames = (temporal_resolution // z_block_size)
-        # self.num_patches = self.height * self.width * self.frames
-        self.z_block_size = z_block_size
-        ###
-        self.proj = nn.Conv3d(in_chans,
-                              embed_dim,
-                              kernel_size=(z_block_size, patch_size, patch_size),
-                              stride=(z_block_size, patch_size, patch_size))
-        self.flatten = flatten
-
-    def forward(self, x):
-        B, C, T, H, W = x.shape
-        x = self.proj(x)
-        if self.flatten:
-            x = x.flatten(2).transpose(1, 2)
-        return x
-
-
-class HeadMLP(nn.Module):
-
-    def __init__(self, n_input, n_classes, n_hidden=512, p=0.1):
-        super(HeadMLP, self).__init__()
-        self.n_input = n_input
-        self.n_classes = n_classes
-        self.n_hidden = n_hidden
-        if n_hidden is None:
-            # use linear classifier
-            self.block_forward = nn.Sequential(nn.Dropout(p=p), nn.Linear(n_input, n_classes, bias=True))
-        else:
-            # use simple MLP classifier
-            self.block_forward = nn.Sequential(nn.Dropout(p=p), nn.Linear(n_input, n_hidden, bias=True),
-                                               nn.BatchNorm1d(n_hidden), nn.ReLU(inplace=True), nn.Dropout(p=p),
-                                               nn.Linear(n_hidden, n_classes, bias=True))
-        print(f"Dropout-NLP: {p}")
-
-    def forward(self, x):
-        return self.block_forward(x)
-
-
-def _conv_filter(state_dict, patch_size=16):
-    """ convert patch embedding weight from manual patchify + linear proj to conv"""
-    out_dict = {}
-    for k, v in state_dict.items():
-        if 'patch_embed.proj.weight' in k:
-            v = v.reshape((v.shape[0], 3, patch_size, patch_size))
-        out_dict[k] = v
-    return out_dict
-
-
-def adapt_input_conv(in_chans, conv_weight, agg='sum'):
-    conv_type = conv_weight.dtype
-    conv_weight = conv_weight.float()
-    O, I, J, K = conv_weight.shape
-    if in_chans == 1:
-        if I > 3:
-            assert conv_weight.shape[1] % 3 == 0
-            # For models with space2depth stems
-            conv_weight = conv_weight.reshape(O, I // 3, 3, J, K)
-            conv_weight = conv_weight.sum(dim=2, keepdim=False)
-        else:
-            if agg == 'sum':
-                print("Summing conv1 weights")
-                conv_weight = conv_weight.sum(dim=1, keepdim=True)
-            else:
-                print("Averaging conv1 weights")
-                conv_weight = conv_weight.mean(dim=1, keepdim=True)
-    elif in_chans != 3:
-        if I != 3:
-            raise NotImplementedError('Weight format not supported by conversion.')
-        else:
-            if agg == 'sum':
-                print("Summing conv1 weights")
-                repeat = int(math.ceil(in_chans / 3))
-                conv_weight = conv_weight.repeat(1, repeat, 1, 1)[:, :in_chans, :, :]
-                conv_weight *= (3 / float(in_chans))
-            else:
-                print("Averaging conv1 weights")
-                conv_weight = conv_weight.mean(dim=1, keepdim=True)
-                conv_weight = conv_weight.repeat(1, in_chans, 1, 1)
-    conv_weight = conv_weight.to(conv_type)
-    return conv_weight
-
-
-def load_pretrained(model, cfg=None, num_classes=1000, in_chans=3, filter_fn=None, strict=True, progress=False):
-    # Load state dict
-    assert (f"{cfg.VIT.PRETRAINED_WEIGHTS} not in [vit_1k, vit_1k_large]")
-    state_dict = torch.hub.load_state_dict_from_url(url=default_cfgs[cfg.VIT.PRETRAINED_WEIGHTS])
-
-    if filter_fn is not None:
-        state_dict = filter_fn(state_dict)
-
-    input_convs = 'patch_embed.proj'
-    if input_convs is not None and in_chans != 3:
-        if isinstance(input_convs, str):
-            input_convs = (input_convs, )
-        for input_conv_name in input_convs:
-            weight_name = input_conv_name + '.weight'
-            try:
-                state_dict[weight_name] = adapt_input_conv(in_chans, state_dict[weight_name], agg='avg')
-                print(f'Converted input conv {input_conv_name} pretrained weights from 3 to {in_chans} channel(s)')
-            except NotImplementedError as e:
-                del state_dict[weight_name]
-                strict = False
-                print(f'Unable to convert pretrained {input_conv_name} weights, using random init for this layer.')
-
-    classifier_name = 'head'
-    label_offset = cfg.get('label_offset', 0)
-    pretrain_classes = 1000
-    if num_classes != pretrain_classes:
-        # completely discard fully connected if model num_classes doesn't match pretrained weights
-        del state_dict[classifier_name + '.weight']
-        del state_dict[classifier_name + '.bias']
-        strict = False
-    elif label_offset > 0:
-        # special case for pretrained weights with an extra background class in pretrained weights
-        classifier_weight = state_dict[classifier_name + '.weight']
-        state_dict[classifier_name + '.weight'] = classifier_weight[label_offset:]
-        classifier_bias = state_dict[classifier_name + '.bias']
-        state_dict[classifier_name + '.bias'] = classifier_bias[label_offset:]
-
-    loaded_state = state_dict
-    self_state = model.state_dict()
-    all_names = set(self_state.keys())
-    saved_names = set([])
-    for name, param in loaded_state.items():
-        param = param
-        if 'module.' in name:
-            name = name.replace('module.', '')
-        if name in self_state.keys() and param.shape == self_state[name].shape:
-            saved_names.add(name)
-            self_state[name].copy_(param)
-        else:
-            print(f"didnt load: {name} of shape: {param.shape}")
-    print("Missing Keys:")
-    print(all_names - saved_names)
+from fastvideo.third_party.synchformer.vit_helper import *  # noqa: F403

fastvideo/third_party/synchformer/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Vladimir Iashin
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

fastvideo/third_party/synchformer/__init__.py

@@ -0,0 +1,5 @@
+"""Shared vendored Synchformer visual backbone."""
+
+from fastvideo.third_party.synchformer.motionformer import MotionFormer
+
+__all__ = ["MotionFormer"]

fastvideo/third_party/synchformer/divided_224_16x4.yaml

@@ -1,3 +1,4 @@
+# Shared MotionFormer configuration used by Synchformer.
 TRAIN:
   ENABLE: True
   DATASET: Ssv2

fastvideo/third_party/synchformer/motionformer.py

@@ -0,0 +1,400 @@
+import logging
+from pathlib import Path
+
+import einops
+import torch
+from omegaconf import OmegaConf
+from timm.layers import trunc_normal_
+from torch import nn
+
+from fastvideo.third_party.synchformer.utils import check_if_file_exists_else_download
+from fastvideo.third_party.synchformer.video_model_builder import VisionTransformer
+
+FILE2URL = {
+    # cfg
+    'motionformer_224_16x4.yaml':
+    'https://raw.githubusercontent.com/facebookresearch/Motionformer/bf43d50/configs/SSV2/motionformer_224_16x4.yaml',
+    'joint_224_16x4.yaml':
+    'https://raw.githubusercontent.com/facebookresearch/Motionformer/bf43d50/configs/SSV2/joint_224_16x4.yaml',
+    'divided_224_16x4.yaml':
+    'https://raw.githubusercontent.com/facebookresearch/Motionformer/bf43d50/configs/SSV2/divided_224_16x4.yaml',
+    # ckpt
+    'ssv2_motionformer_224_16x4.pyth':
+    'https://dl.fbaipublicfiles.com/motionformer/ssv2_motionformer_224_16x4.pyth',
+    'ssv2_joint_224_16x4.pyth':
+    'https://dl.fbaipublicfiles.com/motionformer/ssv2_joint_224_16x4.pyth',
+    'ssv2_divided_224_16x4.pyth':
+    'https://dl.fbaipublicfiles.com/motionformer/ssv2_divided_224_16x4.pyth',
+}
+
+
+class MotionFormer(VisionTransformer):
+    ''' This class serves three puposes:
+            1. Renames the class to MotionFormer.
+            2. Downloads the cfg from the original repo and patches it if needed.
+            3. Takes care of feature extraction by redefining .forward()
+                - if `extract_features=True` and `factorize_space_time=False`,
+                    the output is of shape (B, T, D) where T = 1 + (224 // 16) * (224 // 16) * 8
+                - if `extract_features=True` and `factorize_space_time=True`, the output is of shape (B*S, D)
+                    and spatial and temporal transformer encoder layers are used.
+                - if `extract_features=True` and `factorize_space_time=True` as well as `add_global_repr=True`
+                    the output is of shape (B, D) and spatial and temporal transformer encoder layers
+                    are used as well as the global representation is extracted from segments (extra pos emb
+                    is added).
+    '''
+
+    def __init__(
+        self,
+        extract_features: bool = False,
+        ckpt_path: str = None,
+        factorize_space_time: bool = None,
+        agg_space_module: str = None,
+        agg_time_module: str = None,
+        add_global_repr: bool = True,
+        agg_segments_module: str = None,
+        max_segments: int = None,
+    ):
+        self.extract_features = extract_features
+        self.ckpt_path = ckpt_path
+        self.factorize_space_time = factorize_space_time
+
+        if self.ckpt_path is not None:
+            check_if_file_exists_else_download(self.ckpt_path, FILE2URL)
+            ckpt = torch.load(self.ckpt_path, map_location='cpu')
+            mformer_ckpt2cfg = {
+                'ssv2_motionformer_224_16x4.pyth': 'motionformer_224_16x4.yaml',
+                'ssv2_joint_224_16x4.pyth': 'joint_224_16x4.yaml',
+                'ssv2_divided_224_16x4.pyth': 'divided_224_16x4.yaml',
+            }
+            # init from motionformer ckpt or from our Stage I ckpt
+            # depending on whether the feat extractor was pre-trained on AVCLIPMoCo or not, we need to
+            # load the state dict differently
+            was_pt_on_avclip = self.ckpt_path.endswith(
+                '.pt')  # checks if it is a stage I ckpt (FIXME: a bit generic)
+            if self.ckpt_path.endswith(tuple(mformer_ckpt2cfg.keys())):
+                cfg_fname = mformer_ckpt2cfg[Path(self.ckpt_path).name]
+            elif was_pt_on_avclip:
+                # TODO: this is a hack, we should be able to get the cfg from the ckpt (earlier ckpt didn't have it)
+                s1_cfg = ckpt.get('args', None)  # Stage I cfg
+                if s1_cfg is not None:
+                    s1_vfeat_extractor_ckpt_path = s1_cfg.model.params.vfeat_extractor.params.ckpt_path
+                    # if the stage I ckpt was initialized from a motionformer ckpt or train from scratch
+                    if s1_vfeat_extractor_ckpt_path is not None:
+                        cfg_fname = mformer_ckpt2cfg[Path(s1_vfeat_extractor_ckpt_path).name]
+                    else:
+                        cfg_fname = 'divided_224_16x4.yaml'
+                else:
+                    cfg_fname = 'divided_224_16x4.yaml'
+            else:
+                raise ValueError(f'ckpt_path {self.ckpt_path} is not supported.')
+        else:
+            was_pt_on_avclip = False
+            cfg_fname = 'divided_224_16x4.yaml'
+            # logging.info(f'No ckpt_path provided, using {cfg_fname} config.')
+
+        if cfg_fname in ['motionformer_224_16x4.yaml', 'divided_224_16x4.yaml']:
+            pos_emb_type = 'separate'
+        elif cfg_fname == 'joint_224_16x4.yaml':
+            pos_emb_type = 'joint'
+
+        self.mformer_cfg_path = Path(__file__).absolute().parent / cfg_fname
+
+        check_if_file_exists_else_download(self.mformer_cfg_path, FILE2URL)
+        mformer_cfg = OmegaConf.load(self.mformer_cfg_path)
+        logging.info(f'Loading MotionFormer config from {self.mformer_cfg_path.absolute()}')
+
+        # patch the cfg (from the default cfg defined in the repo `Motionformer/slowfast/config/defaults.py`)
+        mformer_cfg.VIT.ATTN_DROPOUT = 0.0
+        mformer_cfg.VIT.POS_EMBED = pos_emb_type
+        mformer_cfg.VIT.USE_ORIGINAL_TRAJ_ATTN_CODE = True
+        mformer_cfg.VIT.APPROX_ATTN_TYPE = 'none'  # guessing
+        mformer_cfg.VIT.APPROX_ATTN_DIM = 64  # from ckpt['cfg']
+
+        # finally init VisionTransformer with the cfg
+        super().__init__(mformer_cfg)
+
+        # load the ckpt now if ckpt is provided and not from AVCLIPMoCo-pretrained ckpt
+        if (self.ckpt_path is not None) and (not was_pt_on_avclip):
+            _ckpt_load_status = self.load_state_dict(ckpt['model_state'], strict=False)
+            if len(_ckpt_load_status.missing_keys) > 0 or len(
+                    _ckpt_load_status.unexpected_keys) > 0:
+                logging.warning(f'Loading exact vfeat_extractor ckpt from {self.ckpt_path} failed.' \
+                                f'Missing keys: {_ckpt_load_status.missing_keys}, ' \
+                                f'Unexpected keys: {_ckpt_load_status.unexpected_keys}')
+            else:
+                logging.info(f'Loading vfeat_extractor ckpt from {self.ckpt_path} succeeded.')
+
+        if self.extract_features:
+            assert isinstance(self.norm,
+                              nn.LayerNorm), 'early x[:, 1:, :] may not be safe for per-tr weights'
+            # pre-logits are Sequential(nn.Linear(emb, emd), act) and `act` is tanh but see the logger
+            self.pre_logits = nn.Identity()
+            # we don't need the classification head (saving memory)
+            self.head = nn.Identity()
+            self.head_drop = nn.Identity()
+            # avoiding code duplication (used only if agg_*_module is TransformerEncoderLayer)
+            transf_enc_layer_kwargs = dict(
+                d_model=self.embed_dim,
+                nhead=self.num_heads,
+                activation=nn.GELU(),
+                batch_first=True,
+                dim_feedforward=self.mlp_ratio * self.embed_dim,
+                dropout=self.drop_rate,
+                layer_norm_eps=1e-6,
+                norm_first=True,
+            )
+            # define adapters if needed
+            if self.factorize_space_time:
+                if agg_space_module == 'TransformerEncoderLayer':
+                    self.spatial_attn_agg = SpatialTransformerEncoderLayer(
+                        **transf_enc_layer_kwargs)
+                elif agg_space_module == 'AveragePooling':
+                    self.spatial_attn_agg = AveragePooling(avg_pattern='BS D t h w -> BS D t',
+                                                           then_permute_pattern='BS D t -> BS t D')
+                if agg_time_module == 'TransformerEncoderLayer':
+                    self.temp_attn_agg = TemporalTransformerEncoderLayer(**transf_enc_layer_kwargs)
+                elif agg_time_module == 'AveragePooling':
+                    self.temp_attn_agg = AveragePooling(avg_pattern='BS t D -> BS D')
+                elif 'Identity' in agg_time_module:
+                    self.temp_attn_agg = nn.Identity()
+            # define a global aggregation layer (aggregarate over segments)
+            self.add_global_repr = add_global_repr
+            if add_global_repr:
+                if agg_segments_module == 'TransformerEncoderLayer':
+                    # we can reuse the same layer as for temporal factorization (B, dim_to_agg, D) -> (B, D)
+                    # we need to add pos emb (PE) because previously we added the same PE for each segment
+                    pos_max_len = max_segments if max_segments is not None else 16  # 16 = 10sec//0.64sec + 1
+                    self.global_attn_agg = TemporalTransformerEncoderLayer(
+                        add_pos_emb=True,
+                        pos_emb_drop=mformer_cfg.VIT.POS_DROPOUT,
+                        pos_max_len=pos_max_len,
+                        **transf_enc_layer_kwargs)
+                elif agg_segments_module == 'AveragePooling':
+                    self.global_attn_agg = AveragePooling(avg_pattern='B S D -> B D')
+
+        if was_pt_on_avclip:
+            # we need to filter out the state_dict of the AVCLIP model (has both A and V extractors)
+            # and keep only the state_dict of the feat extractor
+            ckpt_weights = dict()
+            for k, v in ckpt['state_dict'].items():
+                if k.startswith(('module.v_encoder.', 'v_encoder.')):
+                    k = k.replace('module.', '').replace('v_encoder.', '')
+                    ckpt_weights[k] = v
+            _load_status = self.load_state_dict(ckpt_weights, strict=False)
+            if len(_load_status.missing_keys) > 0 or len(_load_status.unexpected_keys) > 0:
+                logging.warning(f'Loading exact vfeat_extractor ckpt from {self.ckpt_path} failed. \n' \
+                                f'Missing keys ({len(_load_status.missing_keys)}): ' \
+                                f'{_load_status.missing_keys}, \n' \
+                                f'Unexpected keys ({len(_load_status.unexpected_keys)}): ' \
+                                f'{_load_status.unexpected_keys} \n' \
+                                f'temp_attn_agg are expected to be missing if ckpt was pt contrastively.')
+            else:
+                logging.info(f'Loading vfeat_extractor ckpt from {self.ckpt_path} succeeded.')
+
+        # patch_embed is not used in MotionFormer, only patch_embed_3d, because cfg.VIT.PATCH_SIZE_TEMP > 1
+        # but it used to calculate the number of patches, so we need to set keep it
+        self.patch_embed.requires_grad_(False)
+
+    def forward(self, x):
+        '''
+        x is of shape (B, S, C, T, H, W) where S is the number of segments.
+        '''
+        # Batch, Segments, Channels, T=frames, Height, Width
+        B, S, C, T, H, W = x.shape
+        # Motionformer expects a tensor of shape (1, B, C, T, H, W).
+        # The first dimension (1) is a dummy dimension to make the input tensor and won't be used:
+        # see `video_model_builder.video_input`.
+        # x = x.unsqueeze(0)  # (1, B, S, C, T, H, W)
+
+        orig_shape = (B, S, C, T, H, W)
+        x = x.view(B * S, C, T, H, W)  # flatten batch and segments
+        x = self.forward_segments(x, orig_shape=orig_shape)
+        # unpack the segments (using rest dimensions to support different shapes e.g. (BS, D) or (BS, t, D))
+        x = x.view(B, S, *x.shape[1:])
+        # x is now of shape (B*S, D) or (B*S, t, D) if `self.temp_attn_agg` is `Identity`
+
+        return x  # x is (B, S, ...)
+
+    def forward_segments(self, x, orig_shape: tuple) -> torch.Tensor:
+        '''x is of shape (1, BS, C, T, H, W) where S is the number of segments.'''
+        x, x_mask = self.forward_features(x)
+
+        assert self.extract_features
+
+        # (BS, T, D) where T = 1 + (224 // 16) * (224 // 16) * 8
+        x = x[:,
+              1:, :]  # without the CLS token for efficiency (should be safe for LayerNorm and FC)
+        x = self.norm(x)
+        x = self.pre_logits(x)
+        if self.factorize_space_time:
+            x = self.restore_spatio_temp_dims(x, orig_shape)  # (B*S, D, t, h, w) <- (B*S, t*h*w, D)
+
+            x = self.spatial_attn_agg(x, x_mask)  # (B*S, t, D)
+            x = self.temp_attn_agg(
+                x)  # (B*S, D) or (BS, t, D) if `self.temp_attn_agg` is `Identity`
+
+        return x
+
+    def restore_spatio_temp_dims(self, feats: torch.Tensor, orig_shape: tuple) -> torch.Tensor:
+        '''
+            feats are of shape (B*S, T, D) where T = 1 + (224 // 16) * (224 // 16) * 8
+            Our goal is to make them of shape (B*S, t, h, w, D) where h, w are the spatial dimensions.
+            From `self.patch_embed_3d`, it follows that we could reshape feats with:
+                `feats.transpose(1, 2).view(B*S, D, t, h, w)`
+        '''
+        B, S, C, T, H, W = orig_shape
+        D = self.embed_dim
+
+        # num patches in each dimension
+        t = T // self.patch_embed_3d.z_block_size
+        h = self.patch_embed_3d.height
+        w = self.patch_embed_3d.width
+
+        feats = feats.permute(0, 2, 1)  # (B*S, D, T)
+        feats = feats.view(B * S, D, t, h, w)  # (B*S, D, t, h, w)
+
+        return feats
+
+
+class BaseEncoderLayer(nn.TransformerEncoderLayer):
+    '''
+        This is a wrapper around nn.TransformerEncoderLayer that adds a CLS token
+        to the sequence and outputs the CLS token's representation.
+        This base class parents both SpatialEncoderLayer and TemporalEncoderLayer for the RGB stream
+        and the FrequencyEncoderLayer and TemporalEncoderLayer for the audio stream stream.
+        We also, optionally, add a positional embedding to the input sequence which
+        allows to reuse it for global aggregation (of segments) for both streams.
+    '''
+
+    def __init__(self,
+                 add_pos_emb: bool = False,
+                 pos_emb_drop: float = None,
+                 pos_max_len: int = None,
+                 *args_transformer_enc,
+                 **kwargs_transformer_enc):
+        super().__init__(*args_transformer_enc, **kwargs_transformer_enc)
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, self.self_attn.embed_dim))
+        trunc_normal_(self.cls_token, std=.02)
+
+        # add positional embedding
+        self.add_pos_emb = add_pos_emb
+        if add_pos_emb:
+            self.pos_max_len = 1 + pos_max_len  # +1 (for CLS)
+            self.pos_emb = nn.Parameter(torch.zeros(1, self.pos_max_len, self.self_attn.embed_dim))
+            self.pos_drop = nn.Dropout(pos_emb_drop)
+            trunc_normal_(self.pos_emb, std=.02)
+
+        self.apply(self._init_weights)
+
+    def forward(self, x: torch.Tensor, x_mask: torch.Tensor = None):
+        ''' x is of shape (B, N, D); if provided x_mask is of shape (B, N)'''
+        batch_dim = x.shape[0]
+
+        # add CLS token
+        cls_tokens = self.cls_token.expand(batch_dim, -1, -1)  # expanding to match batch dimension
+        x = torch.cat((cls_tokens, x), dim=-2)  # (batch_dim, 1+seq_len, D)
+        if x_mask is not None:
+            cls_mask = torch.ones((batch_dim, 1), dtype=torch.bool,
+                                  device=x_mask.device)  # 1=keep; 0=mask
+            x_mask_w_cls = torch.cat((cls_mask, x_mask), dim=-1)  # (batch_dim, 1+seq_len)
+            B, N = x_mask_w_cls.shape
+            # torch expects (N, N) or (B*num_heads, N, N) mask (sadness ahead); torch masks
+            x_mask_w_cls = x_mask_w_cls.reshape(B, 1, 1, N)\
+                                       .expand(-1, self.self_attn.num_heads, N, -1)\
+                                       .reshape(B * self.self_attn.num_heads, N, N)
+            assert x_mask_w_cls.dtype == x_mask_w_cls.bool().dtype, 'x_mask_w_cls.dtype != bool'
+            x_mask_w_cls = ~x_mask_w_cls  # invert mask (1=mask)
+        else:
+            x_mask_w_cls = None
+
+        # add positional embedding
+        if self.add_pos_emb:
+            seq_len = x.shape[
+                1]  # (don't even think about moving it before the CLS token concatenation)
+            assert seq_len <= self.pos_max_len, f'Seq len ({seq_len}) > pos_max_len ({self.pos_max_len})'
+            x = x + self.pos_emb[:, :seq_len, :]
+            x = self.pos_drop(x)
+
+        # apply encoder layer (calls nn.TransformerEncoderLayer.forward);
+        x = super().forward(src=x, src_mask=x_mask_w_cls)  # (batch_dim, 1+seq_len, D)
+
+        # CLS token is expected to hold spatial information for each frame
+        x = x[:, 0, :]  # (batch_dim, D)
+
+        return x
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    @torch.jit.ignore
+    def no_weight_decay(self):
+        return {'cls_token', 'pos_emb'}
+
+
+class SpatialTransformerEncoderLayer(BaseEncoderLayer):
+    ''' Aggregates spatial dimensions by applying attention individually to each frame. '''
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def forward(self, x: torch.Tensor, x_mask: torch.Tensor = None) -> torch.Tensor:
+        ''' x is of shape (B*S, D, t, h, w) where S is the number of segments.
+            if specified x_mask (B*S, t, h, w), 0=masked, 1=kept
+            Returns a tensor of shape (B*S, t, D) pooling spatial information for each frame. '''
+        BS, D, t, h, w = x.shape
+
+        # time as a batch dimension and flatten spatial dimensions as sequence
+        x = einops.rearrange(x, 'BS D t h w -> (BS t) (h w) D')
+        # similar to mask
+        if x_mask is not None:
+            x_mask = einops.rearrange(x_mask, 'BS t h w -> (BS t) (h w)')
+
+        # apply encoder layer (BaseEncoderLayer.forward) - it will add CLS token and output its representation
+        x = super().forward(x=x, x_mask=x_mask)  # (B*S*t, D)
+
+        # reshape back to (B*S, t, D)
+        x = einops.rearrange(x, '(BS t) D -> BS t D', BS=BS, t=t)
+
+        # (B*S, t, D)
+        return x
+
+
+class TemporalTransformerEncoderLayer(BaseEncoderLayer):
+    ''' Aggregates temporal dimension with attention. Also used with pos emb as global aggregation
+    in both streams. '''
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def forward(self, x):
+        ''' x is of shape (B*S, t, D) where S is the number of segments.
+            Returns a tensor of shape (B*S, D) pooling temporal information. '''
+        BS, t, D = x.shape
+
+        # apply encoder layer (BaseEncoderLayer.forward) - it will add CLS token and output its representation
+        x = super().forward(x)  # (B*S, D)
+
+        return x  # (B*S, D)
+
+
+class AveragePooling(nn.Module):
+
+    def __init__(self, avg_pattern: str, then_permute_pattern: str = None) -> None:
+        ''' patterns are e.g. "bs t d -> bs d" '''
+        super().__init__()
+        # TODO: need to register them as buffers (but fails because these are strings)
+        self.reduce_fn = 'mean'
+        self.avg_pattern = avg_pattern
+        self.then_permute_pattern = then_permute_pattern
+
+    def forward(self, x: torch.Tensor, x_mask: torch.Tensor = None) -> torch.Tensor:
+        x = einops.reduce(x, self.avg_pattern, self.reduce_fn)
+        if self.then_permute_pattern is not None:
+            x = einops.rearrange(x, self.then_permute_pattern)
+        return x

fastvideo/third_party/synchformer/utils.py

@@ -0,0 +1,94 @@
+"""Checkpoint utilities shared by Synchformer consumers."""
+
+from hashlib import md5
+from pathlib import Path
+
+import requests
+from tqdm import tqdm
+
+PARENT_LINK = 'https://a3s.fi/swift/v1/AUTH_a235c0f452d648828f745589cde1219a'
+FNAME2LINK = {
+    # S3: Synchability: AudioSet (run 2)
+    '24-01-22T20-34-52.pt':
+    f'{PARENT_LINK}/sync/sync_models/24-01-22T20-34-52/24-01-22T20-34-52.pt',
+    'cfg-24-01-22T20-34-52.yaml':
+    f'{PARENT_LINK}/sync/sync_models/24-01-22T20-34-52/cfg-24-01-22T20-34-52.yaml',
+    # S2: Synchformer: AudioSet (run 2)
+    '24-01-04T16-39-21.pt':
+    f'{PARENT_LINK}/sync/sync_models/24-01-04T16-39-21/24-01-04T16-39-21.pt',
+    'cfg-24-01-04T16-39-21.yaml':
+    f'{PARENT_LINK}/sync/sync_models/24-01-04T16-39-21/cfg-24-01-04T16-39-21.yaml',
+    # S2: Synchformer: AudioSet (run 1)
+    '23-08-28T11-23-23.pt':
+    f'{PARENT_LINK}/sync/sync_models/23-08-28T11-23-23/23-08-28T11-23-23.pt',
+    'cfg-23-08-28T11-23-23.yaml':
+    f'{PARENT_LINK}/sync/sync_models/23-08-28T11-23-23/cfg-23-08-28T11-23-23.yaml',
+    # S2: Synchformer: LRS3 (run 2)
+    '23-12-23T18-33-57.pt':
+    f'{PARENT_LINK}/sync/sync_models/23-12-23T18-33-57/23-12-23T18-33-57.pt',
+    'cfg-23-12-23T18-33-57.yaml':
+    f'{PARENT_LINK}/sync/sync_models/23-12-23T18-33-57/cfg-23-12-23T18-33-57.yaml',
+    # S2: Synchformer: VGS (run 2)
+    '24-01-02T10-00-53.pt':
+    f'{PARENT_LINK}/sync/sync_models/24-01-02T10-00-53/24-01-02T10-00-53.pt',
+    'cfg-24-01-02T10-00-53.yaml':
+    f'{PARENT_LINK}/sync/sync_models/24-01-02T10-00-53/cfg-24-01-02T10-00-53.yaml',
+    # SparseSync: ft VGGSound-Full
+    '22-09-21T21-00-52.pt':
+    f'{PARENT_LINK}/sync/sync_models/22-09-21T21-00-52/22-09-21T21-00-52.pt',
+    'cfg-22-09-21T21-00-52.yaml':
+    f'{PARENT_LINK}/sync/sync_models/22-09-21T21-00-52/cfg-22-09-21T21-00-52.yaml',
+    # SparseSync: ft VGGSound-Sparse
+    '22-07-28T15-49-45.pt':
+    f'{PARENT_LINK}/sync/sync_models/22-07-28T15-49-45/22-07-28T15-49-45.pt',
+    'cfg-22-07-28T15-49-45.yaml':
+    f'{PARENT_LINK}/sync/sync_models/22-07-28T15-49-45/cfg-22-07-28T15-49-45.yaml',
+    # SparseSync: only pt on LRS3
+    '22-07-13T22-25-49.pt':
+    f'{PARENT_LINK}/sync/sync_models/22-07-13T22-25-49/22-07-13T22-25-49.pt',
+    'cfg-22-07-13T22-25-49.yaml':
+    f'{PARENT_LINK}/sync/sync_models/22-07-13T22-25-49/cfg-22-07-13T22-25-49.yaml',
+    # SparseSync: feature extractors
+    'ResNetAudio-22-08-04T09-51-04.pt':
+    f'{PARENT_LINK}/sync/ResNetAudio-22-08-04T09-51-04.pt',  # 2s
+    'ResNetAudio-22-08-03T23-14-49.pt':
+    f'{PARENT_LINK}/sync/ResNetAudio-22-08-03T23-14-49.pt',  # 3s
+    'ResNetAudio-22-08-03T23-14-28.pt':
+    f'{PARENT_LINK}/sync/ResNetAudio-22-08-03T23-14-28.pt',  # 4s
+    'ResNetAudio-22-06-24T08-10-33.pt':
+    f'{PARENT_LINK}/sync/ResNetAudio-22-06-24T08-10-33.pt',  # 5s
+    'ResNetAudio-22-06-24T17-31-07.pt':
+    f'{PARENT_LINK}/sync/ResNetAudio-22-06-24T17-31-07.pt',  # 6s
+    'ResNetAudio-22-06-24T23-57-11.pt':
+    f'{PARENT_LINK}/sync/ResNetAudio-22-06-24T23-57-11.pt',  # 7s
+    'ResNetAudio-22-06-25T04-35-42.pt':
+    f'{PARENT_LINK}/sync/ResNetAudio-22-06-25T04-35-42.pt',  # 8s
+}
+
+
+def check_if_file_exists_else_download(path, fname2link=FNAME2LINK, chunk_size=1024):
+    '''Checks if file exists, if not downloads it from the link to the path'''
+    path = Path(path)
+    if not path.exists():
+        path.parent.mkdir(exist_ok=True, parents=True)
+        link = fname2link.get(path.name, None)
+        if link is None:
+            raise ValueError(f'Cant find the checkpoint file: {path}.',
+                             f'Please download it manually and ensure the path exists.')
+        with requests.get(fname2link[path.name], stream=True) as r:
+            total_size = int(r.headers.get('content-length', 0))
+            with tqdm(total=total_size, unit='B', unit_scale=True) as pbar:
+                with open(path, 'wb') as f:
+                    for data in r.iter_content(chunk_size=chunk_size):
+                        if data:
+                            f.write(data)
+                            pbar.update(chunk_size)
+
+
+def get_md5sum(path):
+    hash_md5 = md5()
+    with open(path, 'rb') as f:
+        for chunk in iter(lambda: f.read(4096 * 8), b''):
+            hash_md5.update(chunk)
+    md5sum = hash_md5.hexdigest()
+    return md5sum

fastvideo/third_party/synchformer/video_model_builder.py

@@ -0,0 +1,277 @@
+#!/usr/bin/env python3
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
+# Copyright 2020 Ross Wightman
+# Modified Model definition
+
+from collections import OrderedDict
+from functools import partial
+
+import torch
+import torch.nn as nn
+from timm.layers import trunc_normal_
+
+from fastvideo.third_party.synchformer import vit_helper
+
+
+class VisionTransformer(nn.Module):
+    """ Vision Transformer with support for patch or hybrid CNN input stage """
+
+    def __init__(self, cfg):
+        super().__init__()
+        self.img_size = cfg.DATA.TRAIN_CROP_SIZE
+        self.patch_size = cfg.VIT.PATCH_SIZE
+        self.in_chans = cfg.VIT.CHANNELS
+        if cfg.TRAIN.DATASET == "Epickitchens":
+            self.num_classes = [97, 300]
+        else:
+            self.num_classes = cfg.MODEL.NUM_CLASSES
+        self.embed_dim = cfg.VIT.EMBED_DIM
+        self.depth = cfg.VIT.DEPTH
+        self.num_heads = cfg.VIT.NUM_HEADS
+        self.mlp_ratio = cfg.VIT.MLP_RATIO
+        self.qkv_bias = cfg.VIT.QKV_BIAS
+        self.drop_rate = cfg.VIT.DROP
+        self.drop_path_rate = cfg.VIT.DROP_PATH
+        self.head_dropout = cfg.VIT.HEAD_DROPOUT
+        self.video_input = cfg.VIT.VIDEO_INPUT
+        self.temporal_resolution = cfg.VIT.TEMPORAL_RESOLUTION
+        self.use_mlp = cfg.VIT.USE_MLP
+        self.num_features = self.embed_dim
+        norm_layer = partial(nn.LayerNorm, eps=1e-6)
+        self.attn_drop_rate = cfg.VIT.ATTN_DROPOUT
+        self.head_act = cfg.VIT.HEAD_ACT
+        self.cfg = cfg
+
+        # Patch Embedding
+        self.patch_embed = vit_helper.PatchEmbed(img_size=224,
+                                                 patch_size=self.patch_size,
+                                                 in_chans=self.in_chans,
+                                                 embed_dim=self.embed_dim)
+
+        # 3D Patch Embedding
+        self.patch_embed_3d = vit_helper.PatchEmbed3D(img_size=self.img_size,
+                                                      temporal_resolution=self.temporal_resolution,
+                                                      patch_size=self.patch_size,
+                                                      in_chans=self.in_chans,
+                                                      embed_dim=self.embed_dim,
+                                                      z_block_size=self.cfg.VIT.PATCH_SIZE_TEMP)
+        self.patch_embed_3d.proj.weight.data = torch.zeros_like(
+            self.patch_embed_3d.proj.weight.data)
+
+        # Number of patches
+        if self.video_input:
+            num_patches = self.patch_embed.num_patches * self.temporal_resolution
+        else:
+            num_patches = self.patch_embed.num_patches
+        self.num_patches = num_patches
+
+        # CLS token
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, self.embed_dim))
+        trunc_normal_(self.cls_token, std=.02)
+
+        # Positional embedding
+        self.pos_embed = nn.Parameter(
+            torch.zeros(1, self.patch_embed.num_patches + 1, self.embed_dim))
+        self.pos_drop = nn.Dropout(p=cfg.VIT.POS_DROPOUT)
+        trunc_normal_(self.pos_embed, std=.02)
+
+        if self.cfg.VIT.POS_EMBED == "joint":
+            self.st_embed = nn.Parameter(torch.zeros(1, num_patches + 1, self.embed_dim))
+            trunc_normal_(self.st_embed, std=.02)
+        elif self.cfg.VIT.POS_EMBED == "separate":
+            self.temp_embed = nn.Parameter(torch.zeros(1, self.temporal_resolution, self.embed_dim))
+
+        # Layer Blocks
+        dpr = [x.item() for x in torch.linspace(0, self.drop_path_rate, self.depth)]
+        if self.cfg.VIT.ATTN_LAYER == "divided":
+            self.blocks = nn.ModuleList([
+                vit_helper.DividedSpaceTimeBlock(
+                    attn_type=cfg.VIT.ATTN_LAYER,
+                    dim=self.embed_dim,
+                    num_heads=self.num_heads,
+                    mlp_ratio=self.mlp_ratio,
+                    qkv_bias=self.qkv_bias,
+                    drop=self.drop_rate,
+                    attn_drop=self.attn_drop_rate,
+                    drop_path=dpr[i],
+                    norm_layer=norm_layer,
+                ) for i in range(self.depth)
+            ])
+        else:
+            self.blocks = nn.ModuleList([
+                vit_helper.Block(attn_type=cfg.VIT.ATTN_LAYER,
+                                 dim=self.embed_dim,
+                                 num_heads=self.num_heads,
+                                 mlp_ratio=self.mlp_ratio,
+                                 qkv_bias=self.qkv_bias,
+                                 drop=self.drop_rate,
+                                 attn_drop=self.attn_drop_rate,
+                                 drop_path=dpr[i],
+                                 norm_layer=norm_layer,
+                                 use_original_code=self.cfg.VIT.USE_ORIGINAL_TRAJ_ATTN_CODE)
+                for i in range(self.depth)
+            ])
+        self.norm = norm_layer(self.embed_dim)
+
+        # MLP head
+        if self.use_mlp:
+            hidden_dim = self.embed_dim
+            if self.head_act == 'tanh':
+                # logging.info("Using TanH activation in MLP")
+                act = nn.Tanh()
+            elif self.head_act == 'gelu':
+                # logging.info("Using GELU activation in MLP")
+                act = nn.GELU()
+            else:
+                # logging.info("Using ReLU activation in MLP")
+                act = nn.ReLU()
+            self.pre_logits = nn.Sequential(
+                OrderedDict([
+                    ('fc', nn.Linear(self.embed_dim, hidden_dim)),
+                    ('act', act),
+                ]))
+        else:
+            self.pre_logits = nn.Identity()
+
+        # Classifier Head
+        self.head_drop = nn.Dropout(p=self.head_dropout)
+        if isinstance(self.num_classes, (list, )) and len(self.num_classes) > 1:
+            for a, i in enumerate(range(len(self.num_classes))):
+                setattr(self, "head%d" % a, nn.Linear(self.embed_dim, self.num_classes[i]))
+        else:
+            self.head = nn.Linear(self.embed_dim,
+                                  self.num_classes) if self.num_classes > 0 else nn.Identity()
+
+        # Initialize weights
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    @torch.jit.ignore
+    def no_weight_decay(self):
+        if self.cfg.VIT.POS_EMBED == "joint":
+            return {'pos_embed', 'cls_token', 'st_embed'}
+        else:
+            return {'pos_embed', 'cls_token', 'temp_embed'}
+
+    def get_classifier(self):
+        return self.head
+
+    def reset_classifier(self, num_classes, global_pool=''):
+        self.num_classes = num_classes
+        self.head = (nn.Linear(self.embed_dim, num_classes) if num_classes > 0 else nn.Identity())
+
+    def forward_features(self, x):
+        # if self.video_input:
+        # x = x[0]
+        B = x.shape[0]
+
+        # Tokenize input
+        # if self.cfg.VIT.PATCH_SIZE_TEMP > 1:
+        # for simplicity of mapping between content dimensions (input x) and token dims (after patching)
+        # we use the same trick as for AST (see modeling_ast.ASTModel.forward for the details):
+
+        # apply patching on input
+        x = self.patch_embed_3d(x)
+        tok_mask = None
+
+        # else:
+        #     tok_mask = None
+        #     # 2D tokenization
+        #     if self.video_input:
+        #         x = x.permute(0, 2, 1, 3, 4)
+        #         (B, T, C, H, W) = x.shape
+        #         x = x.reshape(B * T, C, H, W)
+
+        #     x = self.patch_embed(x)
+
+        #     if self.video_input:
+        #         (B2, T2, D2) = x.shape
+        #         x = x.reshape(B, T * T2, D2)
+
+        # Append CLS token
+        cls_tokens = self.cls_token.expand(B, -1, -1)
+        x = torch.cat((cls_tokens, x), dim=1)
+        # if tok_mask is not None:
+        #     # prepend 1(=keep) to the mask to account for the CLS token as well
+        #     tok_mask = torch.cat((torch.ones_like(tok_mask[:, [0]]), tok_mask), dim=1)
+
+        # Interpolate positinoal embeddings
+        # if self.cfg.DATA.TRAIN_CROP_SIZE != 224:
+        #     pos_embed = self.pos_embed
+        #     N = pos_embed.shape[1] - 1
+        #     npatch = int((x.size(1) - 1) / self.temporal_resolution)
+        #     class_emb = pos_embed[:, 0]
+        #     pos_embed = pos_embed[:, 1:]
+        #     dim = x.shape[-1]
+        #     pos_embed = torch.nn.functional.interpolate(
+        #         pos_embed.reshape(1, int(math.sqrt(N)), int(math.sqrt(N)), dim).permute(0, 3, 1, 2),
+        #         scale_factor=math.sqrt(npatch / N),
+        #         mode='bicubic',
+        #     )
+        #     pos_embed = pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
+        #     new_pos_embed = torch.cat((class_emb.unsqueeze(0), pos_embed), dim=1)
+        # else:
+        new_pos_embed = self.pos_embed
+        npatch = self.patch_embed.num_patches
+
+        # Add positional embeddings to input
+        if self.video_input:
+            if self.cfg.VIT.POS_EMBED == "separate":
+                cls_embed = self.pos_embed[:, 0, :].unsqueeze(1)
+                tile_pos_embed = new_pos_embed[:, 1:, :].repeat(1, self.temporal_resolution, 1)
+                tile_temporal_embed = self.temp_embed.repeat_interleave(npatch, 1)
+                total_pos_embed = tile_pos_embed + tile_temporal_embed
+                total_pos_embed = torch.cat([cls_embed, total_pos_embed], dim=1)
+                x = x + total_pos_embed
+            elif self.cfg.VIT.POS_EMBED == "joint":
+                x = x + self.st_embed
+        else:
+            # image input
+            x = x + new_pos_embed
+
+        # Apply positional dropout
+        x = self.pos_drop(x)
+
+        # Encoding using transformer layers
+        for i, blk in enumerate(self.blocks):
+            x = blk(x,
+                    seq_len=npatch,
+                    num_frames=self.temporal_resolution,
+                    approx=self.cfg.VIT.APPROX_ATTN_TYPE,
+                    num_landmarks=self.cfg.VIT.APPROX_ATTN_DIM,
+                    tok_mask=tok_mask)
+
+        ### v-iashin: I moved it to the forward pass
+        # x = self.norm(x)[:, 0]
+        # x = self.pre_logits(x)
+        ###
+        return x, tok_mask
+
+    # def forward(self, x):
+    #     x = self.forward_features(x)
+    #     ### v-iashin: here. This should leave the same forward output as before
+    #     x = self.norm(x)[:, 0]
+    #     x = self.pre_logits(x)
+    #     ###
+    #     x = self.head_drop(x)
+    #     if isinstance(self.num_classes, (list, )) and len(self.num_classes) > 1:
+    #         output = []
+    #         for head in range(len(self.num_classes)):
+    #             x_out = getattr(self, "head%d" % head)(x)
+    #             if not self.training:
+    #                 x_out = torch.nn.functional.softmax(x_out, dim=-1)
+    #             output.append(x_out)
+    #         return output
+    #     else:
+    #         x = self.head(x)
+    #         if not self.training:
+    #             x = torch.nn.functional.softmax(x, dim=-1)
+    #         return x

fastvideo/third_party/synchformer/vit_helper.py

@@ -0,0 +1,399 @@
+#!/usr/bin/env python3
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
+# Copyright 2020 Ross Wightman
+# Modified Model definition
+"""Video model building blocks shared by Synchformer consumers."""
+
+import math
+
+import torch
+import torch.nn as nn
+from einops import rearrange, repeat
+from timm.layers import to_2tuple
+from torch import einsum
+from torch.nn import functional as F
+
+default_cfgs = {
+    'vit_1k':
+    'https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-vitjx/jx_vit_base_p16_224-80ecf9dd.pth',
+    'vit_1k_large':
+    'https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-vitjx/jx_vit_large_p16_224-4ee7a4dc.pth',
+}
+
+
+def qkv_attn(q, k, v, tok_mask: torch.Tensor = None):
+    sim = einsum('b i d, b j d -> b i j', q, k)
+    # apply masking if provided, tok_mask is (B*S*H, N): 1s - keep; sim is (B*S*H, H, N, N)
+    if tok_mask is not None:
+        BSH, N = tok_mask.shape
+        sim = sim.masked_fill(tok_mask.view(BSH, 1, N) == 0,
+                              float('-inf'))  # 1 - broadcasts across N
+    attn = sim.softmax(dim=-1)
+    out = einsum('b i j, b j d -> b i d', attn, v)
+    return out
+
+
+class DividedAttention(nn.Module):
+
+    def __init__(self, dim, num_heads=8, qkv_bias=False, attn_drop=0., proj_drop=0.):
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = head_dim**-0.5
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.proj = nn.Linear(dim, dim)
+
+        # init to zeros
+        self.qkv.weight.data.fill_(0)
+        self.qkv.bias.data.fill_(0)
+        self.proj.weight.data.fill_(1)
+        self.proj.bias.data.fill_(0)
+
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x, einops_from, einops_to, tok_mask: torch.Tensor = None, **einops_dims):
+        # num of heads variable
+        h = self.num_heads
+
+        # project x to q, k, v vaalues
+        q, k, v = self.qkv(x).chunk(3, dim=-1)
+        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (q, k, v))
+        if tok_mask is not None:
+            # replicate token mask across heads (b, n) -> (b, h, n) -> (b*h, n) -- same as qkv but w/o d
+            assert len(tok_mask.shape) == 2
+            tok_mask = tok_mask.unsqueeze(1).expand(-1, h, -1).reshape(-1, tok_mask.shape[1])
+
+        # Scale q
+        q *= self.scale
+
+        # Take out cls_q, cls_k, cls_v
+        (cls_q, q_), (cls_k, k_), (cls_v, v_) = map(lambda t: (t[:, 0:1], t[:, 1:]), (q, k, v))
+        # the same for masking
+        if tok_mask is not None:
+            cls_mask, mask_ = tok_mask[:, 0:1], tok_mask[:, 1:]
+        else:
+            cls_mask, mask_ = None, None
+
+        # let CLS token attend to key / values of all patches across time and space
+        cls_out = qkv_attn(cls_q, k, v, tok_mask=tok_mask)
+
+        # rearrange across time or space
+        q_, k_, v_ = map(lambda t: rearrange(t, f'{einops_from} -> {einops_to}', **einops_dims),
+                         (q_, k_, v_))
+
+        # expand CLS token keys and values across time or space and concat
+        r = q_.shape[0] // cls_k.shape[0]
+        cls_k, cls_v = map(lambda t: repeat(t, 'b () d -> (b r) () d', r=r), (cls_k, cls_v))
+
+        k_ = torch.cat((cls_k, k_), dim=1)
+        v_ = torch.cat((cls_v, v_), dim=1)
+
+        # the same for masking (if provided)
+        if tok_mask is not None:
+            # since mask does not have the latent dim (d), we need to remove it from einops dims
+            mask_ = rearrange(mask_, f'{einops_from} -> {einops_to}'.replace(' d', ''),
+                              **einops_dims)
+            cls_mask = repeat(cls_mask, 'b () -> (b r) ()',
+                              r=r)  # expand cls_mask across time or space
+            mask_ = torch.cat((cls_mask, mask_), dim=1)
+
+        # attention
+        out = qkv_attn(q_, k_, v_, tok_mask=mask_)
+
+        # merge back time or space
+        out = rearrange(out, f'{einops_to} -> {einops_from}', **einops_dims)
+
+        # concat back the cls token
+        out = torch.cat((cls_out, out), dim=1)
+
+        # merge back the heads
+        out = rearrange(out, '(b h) n d -> b n (h d)', h=h)
+
+        ## to out
+        x = self.proj(out)
+        x = self.proj_drop(x)
+        return x
+
+
+class DividedSpaceTimeBlock(nn.Module):
+
+    def __init__(self,
+                 dim=768,
+                 num_heads=12,
+                 attn_type='divided',
+                 mlp_ratio=4.,
+                 qkv_bias=False,
+                 drop=0.,
+                 attn_drop=0.,
+                 drop_path=0.,
+                 act_layer=nn.GELU,
+                 norm_layer=nn.LayerNorm):
+        super().__init__()
+
+        self.einops_from_space = 'b (f n) d'
+        self.einops_to_space = '(b f) n d'
+        self.einops_from_time = 'b (f n) d'
+        self.einops_to_time = '(b n) f d'
+
+        self.norm1 = norm_layer(dim)
+
+        self.attn = DividedAttention(dim,
+                                     num_heads=num_heads,
+                                     qkv_bias=qkv_bias,
+                                     attn_drop=attn_drop,
+                                     proj_drop=drop)
+
+        self.timeattn = DividedAttention(dim,
+                                         num_heads=num_heads,
+                                         qkv_bias=qkv_bias,
+                                         attn_drop=attn_drop,
+                                         proj_drop=drop)
+
+        # self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+        self.drop_path = nn.Identity()
+        self.norm2 = norm_layer(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = Mlp(in_features=dim,
+                       hidden_features=mlp_hidden_dim,
+                       act_layer=act_layer,
+                       drop=drop)
+        self.norm3 = norm_layer(dim)
+
+    def forward(self,
+                x,
+                seq_len=196,
+                num_frames=8,
+                approx='none',
+                num_landmarks=128,
+                tok_mask: torch.Tensor = None):
+        time_output = self.timeattn(self.norm3(x),
+                                    self.einops_from_time,
+                                    self.einops_to_time,
+                                    n=seq_len,
+                                    tok_mask=tok_mask)
+        time_residual = x + time_output
+
+        space_output = self.attn(self.norm1(time_residual),
+                                 self.einops_from_space,
+                                 self.einops_to_space,
+                                 f=num_frames,
+                                 tok_mask=tok_mask)
+        space_residual = time_residual + self.drop_path(space_output)
+
+        x = space_residual
+        x = x + self.drop_path(self.mlp(self.norm2(x)))
+        return x
+
+
+class Mlp(nn.Module):
+
+    def __init__(self,
+                 in_features,
+                 hidden_features=None,
+                 out_features=None,
+                 act_layer=nn.GELU,
+                 drop=0.):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.act = act_layer()
+        self.fc2 = nn.Linear(hidden_features, out_features)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x
+
+
+class PatchEmbed(nn.Module):
+    """ Image to Patch Embedding
+    """
+
+    def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768):
+        super().__init__()
+        img_size = img_size if type(img_size) is tuple else to_2tuple(img_size)
+        patch_size = img_size if type(patch_size) is tuple else to_2tuple(patch_size)
+        num_patches = (img_size[1] // patch_size[1]) * (img_size[0] // patch_size[0])
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.num_patches = num_patches
+
+        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
+
+    def forward(self, x):
+        B, C, H, W = x.shape
+        x = self.proj(x).flatten(2).transpose(1, 2)
+        return x
+
+
+class PatchEmbed3D(nn.Module):
+    """ Image to Patch Embedding """
+
+    def __init__(self,
+                 img_size=224,
+                 temporal_resolution=4,
+                 in_chans=3,
+                 patch_size=16,
+                 z_block_size=2,
+                 embed_dim=768,
+                 flatten=True):
+        super().__init__()
+        self.height = (img_size // patch_size)
+        self.width = (img_size // patch_size)
+        ### v-iashin: these two are incorrect
+        # self.frames = (temporal_resolution // z_block_size)
+        # self.num_patches = self.height * self.width * self.frames
+        self.z_block_size = z_block_size
+        ###
+        self.proj = nn.Conv3d(in_chans,
+                              embed_dim,
+                              kernel_size=(z_block_size, patch_size, patch_size),
+                              stride=(z_block_size, patch_size, patch_size))
+        self.flatten = flatten
+
+    def forward(self, x):
+        B, C, T, H, W = x.shape
+        x = self.proj(x)
+        if self.flatten:
+            x = x.flatten(2).transpose(1, 2)
+        return x
+
+
+class HeadMLP(nn.Module):
+
+    def __init__(self, n_input, n_classes, n_hidden=512, p=0.1):
+        super(HeadMLP, self).__init__()
+        self.n_input = n_input
+        self.n_classes = n_classes
+        self.n_hidden = n_hidden
+        if n_hidden is None:
+            # use linear classifier
+            self.block_forward = nn.Sequential(nn.Dropout(p=p),
+                                               nn.Linear(n_input, n_classes, bias=True))
+        else:
+            # use simple MLP classifier
+            self.block_forward = nn.Sequential(nn.Dropout(p=p),
+                                               nn.Linear(n_input, n_hidden, bias=True),
+                                               nn.BatchNorm1d(n_hidden), nn.ReLU(inplace=True),
+                                               nn.Dropout(p=p),
+                                               nn.Linear(n_hidden, n_classes, bias=True))
+        print(f"Dropout-NLP: {p}")
+
+    def forward(self, x):
+        return self.block_forward(x)
+
+
+def _conv_filter(state_dict, patch_size=16):
+    """ convert patch embedding weight from manual patchify + linear proj to conv"""
+    out_dict = {}
+    for k, v in state_dict.items():
+        if 'patch_embed.proj.weight' in k:
+            v = v.reshape((v.shape[0], 3, patch_size, patch_size))
+        out_dict[k] = v
+    return out_dict
+
+
+def adapt_input_conv(in_chans, conv_weight, agg='sum'):
+    conv_type = conv_weight.dtype
+    conv_weight = conv_weight.float()
+    O, I, J, K = conv_weight.shape
+    if in_chans == 1:
+        if I > 3:
+            assert conv_weight.shape[1] % 3 == 0
+            # For models with space2depth stems
+            conv_weight = conv_weight.reshape(O, I // 3, 3, J, K)
+            conv_weight = conv_weight.sum(dim=2, keepdim=False)
+        else:
+            if agg == 'sum':
+                print("Summing conv1 weights")
+                conv_weight = conv_weight.sum(dim=1, keepdim=True)
+            else:
+                print("Averaging conv1 weights")
+                conv_weight = conv_weight.mean(dim=1, keepdim=True)
+    elif in_chans != 3:
+        if I != 3:
+            raise NotImplementedError('Weight format not supported by conversion.')
+        else:
+            if agg == 'sum':
+                print("Summing conv1 weights")
+                repeat = int(math.ceil(in_chans / 3))
+                conv_weight = conv_weight.repeat(1, repeat, 1, 1)[:, :in_chans, :, :]
+                conv_weight *= (3 / float(in_chans))
+            else:
+                print("Averaging conv1 weights")
+                conv_weight = conv_weight.mean(dim=1, keepdim=True)
+                conv_weight = conv_weight.repeat(1, in_chans, 1, 1)
+    conv_weight = conv_weight.to(conv_type)
+    return conv_weight
+
+
+def load_pretrained(model,
+                    cfg=None,
+                    num_classes=1000,
+                    in_chans=3,
+                    filter_fn=None,
+                    strict=True,
+                    progress=False):
+    # Load state dict
+    assert (f"{cfg.VIT.PRETRAINED_WEIGHTS} not in [vit_1k, vit_1k_large]")
+    state_dict = torch.hub.load_state_dict_from_url(url=default_cfgs[cfg.VIT.PRETRAINED_WEIGHTS])
+
+    if filter_fn is not None:
+        state_dict = filter_fn(state_dict)
+
+    input_convs = 'patch_embed.proj'
+    if input_convs is not None and in_chans != 3:
+        if isinstance(input_convs, str):
+            input_convs = (input_convs, )
+        for input_conv_name in input_convs:
+            weight_name = input_conv_name + '.weight'
+            try:
+                state_dict[weight_name] = adapt_input_conv(in_chans,
+                                                           state_dict[weight_name],
+                                                           agg='avg')
+                print(
+                    f'Converted input conv {input_conv_name} pretrained weights from 3 to {in_chans} channel(s)'
+                )
+            except NotImplementedError as e:
+                del state_dict[weight_name]
+                strict = False
+                print(
+                    f'Unable to convert pretrained {input_conv_name} weights, using random init for this layer.'
+                )
+
+    classifier_name = 'head'
+    label_offset = cfg.get('label_offset', 0)
+    pretrain_classes = 1000
+    if num_classes != pretrain_classes:
+        # completely discard fully connected if model num_classes doesn't match pretrained weights
+        del state_dict[classifier_name + '.weight']
+        del state_dict[classifier_name + '.bias']
+        strict = False
+    elif label_offset > 0:
+        # special case for pretrained weights with an extra background class in pretrained weights
+        classifier_weight = state_dict[classifier_name + '.weight']
+        state_dict[classifier_name + '.weight'] = classifier_weight[label_offset:]
+        classifier_bias = state_dict[classifier_name + '.bias']
+        state_dict[classifier_name + '.bias'] = classifier_bias[label_offset:]
+
+    loaded_state = state_dict
+    self_state = model.state_dict()
+    all_names = set(self_state.keys())
+    saved_names = set([])
+    for name, param in loaded_state.items():
+        param = param
+        if 'module.' in name:
+            name = name.replace('module.', '')
+        if name in self_state.keys() and param.shape == self_state[name].shape:
+            saved_names.add(name)
+            self_state[name].copy_(param)
+        else:
+            print(f"didnt load: {name} of shape: {param.shape}")
+    print("Missing Keys:")
+    print(all_names - saved_names)

fastvideo/train/trainer.py

@@ -11,10 +11,13 @@ import torch
 from tqdm.auto import tqdm
 
 from fastvideo.distributed import get_sp_group, get_world_group
+from fastvideo.logger import init_logger
 from fastvideo.train.callbacks.callback import CallbackDict
 from fastvideo.train.methods.base import LogScalar, TrainingMethod
 from fastvideo.train.utils.tracking import build_tracker
 
+logger = init_logger(__name__)
+
 if TYPE_CHECKING:
     from fastvideo.train.utils.training_config import (
         TrainingConfig, )
@@ -222,6 +225,10 @@ class Trainer:
             metrics["step_time_sec"] = (time.perf_counter() - t0)
             metrics["vsa_sparsity"] = float(tc.vsa_sparsity)
             if self.global_rank == 0 and metrics:
+                # Console as well as tracker: with trackers disabled the
+                # tracker is a DummyTracker, and this is then the only place a
+                # step's loss is reported at all.
+                logger.info("step %d %s", step, {k: round(v, 6) for k, v in metrics.items()})
                 self.tracker.log(metrics, step)
 
             self.callbacks.on_training_step_end(

fastvideo/train/utils/moduleloader.py

@@ -20,6 +20,9 @@ from fastvideo.utils import (
     maybe_download_model,
     verify_model_config_and_directory,
 )
+from fastvideo.models.loader.ltx_single_file import (
+    model_index_and_component_path,
+)
 from fastvideo.platforms import AttentionBackendEnum
 
 if TYPE_CHECKING:
@@ -106,7 +109,7 @@ def load_module_from_path(
     fastvideo_args: Any = _make_training_args(training_config, model_path=model_path)
 
     local_model_path = maybe_download_model(model_path)
-    config = verify_model_config_and_directory(local_model_path)
+    config, component_path = model_index_and_component_path(local_model_path, module_type)
 
     if module_type not in config:
         raise ValueError(f"Module {module_type!r} not found in "
@@ -120,7 +123,6 @@ def load_module_from_path(
     # Trailing modular-manifest metadata does not change component dispatch;
     # the provider and architecture remain the first two fields.
     transformers_or_diffusers, _architecture = module_info[:2]
-    component_path = os.path.join(local_model_path, module_type)
 
     # fastvideo_args is freshly built above and never escapes this function,
     # so overrides are plain assignments — nothing to save or restore.

fastvideo/training/distillation_pipeline.py

@@ -280,7 +280,8 @@ class DistillationPipeline(TrainingPipeline):
             # Download the model if it's a Hugging Face model ID
             local_model_path = maybe_download_model(model_path)
             logger.info("Model downloaded/found at: %s", local_model_path)
-            config = verify_model_config_and_directory(local_model_path)
+            from fastvideo.models.loader.ltx_single_file import (model_index_and_component_path)
+            config, component_path = model_index_and_component_path(local_model_path, module_type)
 
             if module_type not in config:
                 if hasattr(self, '_extra_config_module_map') and module_type in self._extra_config_module_map:
@@ -298,7 +299,6 @@ class DistillationPipeline(TrainingPipeline):
                 raise ValueError(f"Module {module_type} has null value in config at {local_model_path}")
 
             transformers_or_diffusers, architecture = module_info
-            component_path = os.path.join(local_model_path, module_type)
             module = PipelineComponentLoader.load_module(
                 module_name=module_type,
                 component_model_path=component_path,

pyproject.toml

@@ -208,8 +208,7 @@ dreamverse-mock-server = "dreamverse.mock_server:cli"
 
 [tool.setuptools.packages.find]
 where = [".", "apps/dreamverse"]
-include = ["fastvideo", "fastvideo.*", "dreamverse", "dreamverse.*"]
-exclude = ["dreamverse.tests*"]
+exclude = ["assets*", "docker*", "docs", "scripts*", "apps*", "dreamverse.tests*"]
 
 [tool.setuptools.package-data]
 dreamverse = ["prompts/*.md"]

scripts/checkpoint_conversion/convert_mmaudio_to_diffusers.py

@@ -0,0 +1,344 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Convert MMAudio ``large_44k_v2`` assets into a FastVideo component tree.
+
+The converter is deliberately offline: every large source asset must already
+exist locally. It splits the shared DFN5B OpenCLIP checkpoint into native
+FastVideo text and vision encoders, preserves the exact MMAudio transformer,
+Synchformer, VAE, and BigVGAN weights, and emits the standard
+``model_index.json`` layout consumed by ``ComposedPipelineBase``.
+
+Example::
+
+    python scripts/checkpoint_conversion/convert_mmaudio_to_diffusers.py \
+      --transformer-checkpoint ../MMAudio/weights/mmaudio_large_44k_v2.pth \
+      --audio-vae-checkpoint ../MMAudio/ext_weights/v1-44.pth \
+      --synchformer-checkpoint ../MMAudio/ext_weights/synchformer_state_dict.pth \
+      --dfn5b-dir official_weights/mmaudio/DFN5B-CLIP-ViT-H-14-384 \
+      --bigvgan-dir official_weights/mmaudio/bigvgan_v2_44khz_128band_512x \
+      --output converted_weights/mmaudio/large_44k_v2
+"""
+
+from __future__ import annotations
+
+import argparse
+import gzip
+import json
+from pathlib import Path
+from typing import Any
+
+import torch
+from safetensors.torch import save_file
+
+
+TRANSFORMER_CONFIG = {
+    "_class_name": "MMAudioTransformer",
+    "latent_dim": 40,
+    "clip_dim": 1024,
+    "sync_dim": 768,
+    "text_dim": 1024,
+    "hidden_dim": 896,
+    "depth": 21,
+    "fused_depth": 14,
+    "num_heads": 14,
+    "mlp_ratio": 4.0,
+    "latent_seq_len": 345,
+    "clip_seq_len": 64,
+    "sync_seq_len": 192,
+    "text_seq_len": 77,
+    "v2": True,
+}
+
+TEXT_ENCODER_CONFIG = {
+    "architectures": ["MMAudioDFNCLIPTextEncoder"],
+    "vocab_size": 49408,
+    "hidden_size": 1024,
+    "intermediate_size": 4096,
+    "projection_dim": 1024,
+    "num_hidden_layers": 24,
+    "num_attention_heads": 16,
+    "max_position_embeddings": 77,
+    "text_len": 77,
+    "hidden_act": "quick_gelu",
+    "layer_norm_eps": 1e-5,
+    "pad_token_id": 0,
+    "bos_token_id": 49406,
+    "eos_token_id": 49407,
+}
+
+IMAGE_ENCODER_CONFIG = {
+    "architectures": ["MMAudioDFNCLIPVisionEncoder"],
+    "hidden_size": 1280,
+    "intermediate_size": 5120,
+    "projection_dim": 1024,
+    "num_hidden_layers": 32,
+    "num_attention_heads": 16,
+    "num_channels": 3,
+    "image_size": 378,
+    "patch_size": 14,
+    "hidden_act": "quick_gelu",
+    "layer_norm_eps": 1e-5,
+}
+
+SYNCHFORMER_CONFIG = {
+    "architectures": ["MMAudioSynchformerVisualEncoder"],
+    "image_size": 224,
+    "num_channels": 3,
+    "segment_size": 16,
+    "segment_stride": 8,
+    "hidden_size": 768,
+    "tokens_per_segment": 8,
+}
+
+MODEL_INDEX = {
+    "_class_name": "MMAudioPipeline",
+    "_diffusers_version": "0.36.0",
+    "_fastvideo_model_family": "mmaudio",
+    "_fastvideo_workload_types": ["V2A", "T2A"],
+    "transformer": [
+        "fastvideo.models.dits.mmaudio",
+        "MMAudioTransformer",
+    ],
+    "text_encoder": [
+        "fastvideo.models.encoders.mmaudio_clip",
+        "MMAudioDFNCLIPTextEncoder",
+    ],
+    "tokenizer": ["transformers", "CLIPTokenizer"],
+    "image_encoder": [
+        "fastvideo.models.encoders.mmaudio_clip",
+        "MMAudioDFNCLIPVisionEncoder",
+    ],
+    "image_encoder_2": [
+        "fastvideo.models.encoders.mmaudio_synchformer",
+        "MMAudioSynchformerVisualEncoder",
+    ],
+    "audio_vae": ["fastvideo.models.audio.mmaudio_vae", "MMAudioVAE"],
+    "vocoder": ["fastvideo.models.audio.bigvgan", "BigVGANV2"],
+    "scheduler": ["diffusers", "FlowMatchEulerDiscreteScheduler"],
+}
+
+
+def _write_json(path: Path, value: dict[str, Any]) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    with path.open("w", encoding="utf-8") as handle:
+        json.dump(value, handle, indent=2, sort_keys=True)
+        handle.write("\n")
+
+
+def _write_component(output: Path, name: str, state: dict[str, torch.Tensor], config: dict[str, Any]) -> None:
+    directory = output / name
+    directory.mkdir(parents=True, exist_ok=True)
+    contiguous = {key: tensor.detach().cpu().contiguous() for key, tensor in state.items()}
+    save_file(contiguous, directory / "diffusion_pytorch_model.safetensors", metadata={"format": "pt"})
+    _write_json(directory / "config.json", config)
+
+
+def _load_torch_state(path: Path) -> dict[str, torch.Tensor]:
+    if not path.is_file():
+        raise FileNotFoundError(path)
+    value = torch.load(path, map_location="cpu", weights_only=True)
+    if not isinstance(value, dict):
+        raise TypeError(f"Expected a state dict in {path}, got {type(value)}")
+    for key in ("state_dict", "model", "generator"):
+        nested = value.get(key)
+        if isinstance(nested, dict) and nested:
+            value = nested
+            break
+    if not all(isinstance(tensor, torch.Tensor) for tensor in value.values()):
+        raise TypeError(f"Checkpoint {path} contains non-tensor state entries")
+    return value
+
+
+def map_open_clip_text_state(
+    state: dict[str, torch.Tensor],
+) -> dict[str, torch.Tensor]:
+    mapped: dict[str, torch.Tensor] = {
+        "text_model.embeddings.token_embedding.weight": state["token_embedding.weight"],
+        "text_model.embeddings.position_embedding.weight": state["positional_embedding"],
+        "text_model.final_layer_norm.weight": state["ln_final.weight"],
+        "text_model.final_layer_norm.bias": state["ln_final.bias"],
+    }
+    for name, tensor in state.items():
+        if not name.startswith("transformer.resblocks."):
+            continue
+        target = name.replace("transformer.resblocks.", "text_model.encoder.layers.")
+        target = target.replace(".ln_1.", ".layer_norm1.")
+        target = target.replace(".ln_2.", ".layer_norm2.")
+        target = target.replace(".attn.in_proj_", ".self_attn.qkv_proj.")
+        target = target.replace(".attn.out_proj.", ".self_attn.out_proj.")
+        target = target.replace(".mlp.c_fc.", ".mlp.fc1.")
+        target = target.replace(".mlp.c_proj.", ".mlp.fc2.")
+        mapped[target] = tensor
+    return mapped
+
+
+def map_open_clip_vision_state(
+    state: dict[str, torch.Tensor],
+) -> dict[str, torch.Tensor]:
+    mapped: dict[str, torch.Tensor] = {
+        "vision_model.embeddings.class_embedding": state["visual.class_embedding"],
+        "vision_model.embeddings.patch_embedding.weight": state["visual.conv1.weight"],
+        "vision_model.embeddings.position_embedding.weight": state["visual.positional_embedding"],
+        "vision_model.pre_layrnorm.weight": state["visual.ln_pre.weight"],
+        "vision_model.pre_layrnorm.bias": state["visual.ln_pre.bias"],
+        "vision_model.post_layernorm.weight": state["visual.ln_post.weight"],
+        "vision_model.post_layernorm.bias": state["visual.ln_post.bias"],
+        "visual_projection.weight": state["visual.proj"].t(),
+    }
+    for name, tensor in state.items():
+        if not name.startswith("visual.transformer.resblocks."):
+            continue
+        target = name.replace("visual.transformer.resblocks.", "vision_model.encoder.layers.")
+        target = target.replace(".ln_1.", ".layer_norm1.")
+        target = target.replace(".ln_2.", ".layer_norm2.")
+        target = target.replace(".attn.in_proj_", ".self_attn.qkv_proj.")
+        target = target.replace(".attn.out_proj.", ".self_attn.out_proj.")
+        target = target.replace(".mlp.c_fc.", ".mlp.fc1.")
+        target = target.replace(".mlp.c_proj.", ".mlp.fc2.")
+        mapped[target] = tensor
+    return mapped
+
+
+def write_open_clip_tokenizer(output: Path) -> None:
+    """Write the bundled OpenAI CLIP BPE as an AutoTokenizer component.
+
+    OpenCLIP pads its 77-token tensor with integer zero. ``CLIPTokenizer``
+    cannot use vocabulary ID zero as a special pad token without changing how
+    a literal exclamation mark is tokenized, so the MMAudio text stage zeros
+    positions selected by ``attention_mask`` after tokenization.
+    """
+    from open_clip.tokenizer import bytes_to_unicode, default_bpe
+
+    with gzip.open(default_bpe()) as bpe_file:
+        merges_raw = bpe_file.read().decode("utf-8").split("\n")
+    merges = merges_raw[1 : 49152 - 256 - 2 + 1]
+    merge_pairs = [tuple(merge.split()) for merge in merges]
+    vocab = list(bytes_to_unicode().values())
+    vocab += [token + "</w>" for token in vocab]
+    vocab += ["".join(pair) for pair in merge_pairs]
+    vocab += ["<start_of_text>", "<end_of_text>"]
+    encoder = {token: index for index, token in enumerate(vocab)}
+
+    directory = output / "tokenizer"
+    directory.mkdir(parents=True, exist_ok=True)
+    _write_json(directory / "vocab.json", encoder)
+    with (directory / "merges.txt").open("w", encoding="utf-8") as handle:
+        handle.write("#version: 0.2\n")
+        for first, second in merge_pairs:
+            handle.write(f"{first} {second}\n")
+    _write_json(
+        directory / "tokenizer_config.json",
+        {
+            "tokenizer_class": "CLIPTokenizer",
+            "model_max_length": 77,
+            "bos_token": "<start_of_text>",
+            "eos_token": "<end_of_text>",
+            "unk_token": "<end_of_text>",
+            "pad_token": "<end_of_text>",
+            "do_lower_case": True,
+        },
+    )
+    _write_json(
+        directory / "special_tokens_map.json",
+        {
+            "bos_token": "<start_of_text>",
+            "eos_token": "<end_of_text>",
+            "unk_token": "<end_of_text>",
+            "pad_token": "<end_of_text>",
+        },
+    )
+
+
+def _load_dfn5b_state(directory: Path) -> dict[str, torch.Tensor]:
+    if not directory.is_dir():
+        raise FileNotFoundError(directory)
+    from open_clip import create_model_from_pretrained
+
+    model = create_model_from_pretrained(f"local-dir:{directory}",
+                                         return_transform=False)
+    state = {key: tensor.detach().cpu() for key, tensor in model.state_dict().items()}
+    del model
+    return state
+
+
+def convert(args: argparse.Namespace) -> None:
+    output = args.output.resolve()
+    output.mkdir(parents=True, exist_ok=True)
+
+    transformer_state = _load_torch_state(args.transformer_checkpoint)
+    # Official ``MMAudio.load_weights`` discards this derived buffer. Keeping
+    # it would make a standard strict FastVideo component load fail.
+    transformer_state.pop("t_embed.freqs", None)
+    transformer_state.pop("latent_rot", None)
+    transformer_state.pop("clip_rot", None)
+    _write_component(output, "transformer", transformer_state, TRANSFORMER_CONFIG)
+
+    vae_state = _load_torch_state(args.audio_vae_checkpoint)
+    decoder_state = {
+        key: tensor
+        for key, tensor in vae_state.items()
+        if key.startswith("decoder.") or key in {"data_mean", "data_std"}
+    }
+    if not decoder_state:
+        raise ValueError("Audio VAE checkpoint did not contain decoder weights")
+    _write_component(
+        output,
+        "audio_vae",
+        decoder_state,
+        {"_class_name": "MMAudioVAE", "mode": "44k", "need_encoder": False},
+    )
+
+    synchformer_state = _load_torch_state(args.synchformer_checkpoint)
+    synchformer_visual_state = {
+        name: tensor
+        for name, tensor in synchformer_state.items()
+        if name.startswith("vfeat_extractor.")
+    }
+    if not synchformer_visual_state:
+        raise ValueError(
+            "Synchformer checkpoint did not contain vfeat_extractor weights")
+    _write_component(output, "image_encoder_2", synchformer_visual_state,
+                     SYNCHFORMER_CONFIG)
+
+    dfn_state = _load_dfn5b_state(args.dfn5b_dir)
+    _write_component(output, "text_encoder", map_open_clip_text_state(dfn_state), TEXT_ENCODER_CONFIG)
+    _write_component(output, "image_encoder", map_open_clip_vision_state(dfn_state), IMAGE_ENCODER_CONFIG)
+    write_open_clip_tokenizer(output)
+
+    bigvgan_config_path = args.bigvgan_dir / "config.json"
+    if not bigvgan_config_path.is_file():
+        raise FileNotFoundError(bigvgan_config_path)
+    with bigvgan_config_path.open(encoding="utf-8") as handle:
+        bigvgan_config = json.load(handle)
+    bigvgan_config["_class_name"] = "BigVGANV2"
+    bigvgan_config["weight_norm_removed"] = False
+    bigvgan_state = _load_torch_state(args.bigvgan_dir / "bigvgan_generator.pt")
+    _write_component(output, "vocoder", bigvgan_state, bigvgan_config)
+
+    _write_json(
+        output / "scheduler/scheduler_config.json",
+        {
+            "_class_name": "FlowMatchEulerDiscreteScheduler",
+            "num_train_timesteps": 1000,
+            "shift": 1.0,
+            "invert_sigmas": True,
+            "sigma_min": 0.0,
+            "use_reference_discrete_timesteps": True,
+        },
+    )
+    _write_json(output / "model_index.json", MODEL_INDEX)
+    print(f"Converted MMAudio large_44k_v2 components to {output}")
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
+    parser.add_argument("--transformer-checkpoint", type=Path, required=True)
+    parser.add_argument("--audio-vae-checkpoint", type=Path, required=True)
+    parser.add_argument("--synchformer-checkpoint", type=Path, required=True)
+    parser.add_argument("--dfn5b-dir", type=Path, required=True)
+    parser.add_argument("--bigvgan-dir", type=Path, required=True)
+    parser.add_argument("--output", type=Path, required=True)
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    convert(parse_args())

tests/local_tests/mmaudio/PORT_STATUS.md

@@ -0,0 +1,98 @@
+# MMAudio Port Status
+
+## Summary
+
+- model_family: `mmaudio`
+- workload_types: `V2A`, `T2A`
+- official_ref: `../MMAudio` at `974010a026c731054592d8f777218bd9d85a6c24`
+- first_variant: `large_44k_v2`
+- phase: `native_pipeline_complete`
+- status: `real_weight_parity_pass`
+- last_updated: `2026-08-02`
+
+## Native Components
+
+| Component | FastVideo implementation | Reuse/port decision | Real-weight result |
+|---|---|---|---|
+| MMAudio transformer | `fastvideo/models/dits/mmaudio.py` | Native 1D multimodal DiT | exact |
+| DFN5B text/vision | `fastvideo/models/encoders/mmaudio_clip.py` | Shared native CLIP core, MMAudio adapters | exact |
+| Synchformer visual encoder | `fastvideo/models/encoders/mmaudio_synchformer.py` | Shared backbone under `fastvideo/third_party/synchformer` | exact, including 16-frame/stride-8 usage contract |
+| 44.1 kHz VAE | `fastvideo/models/audio/mmaudio_vae.py` | Native audio component | exact state structure, FP32/BF16 random-weight decoder, and real-weight encode/decode parity |
+| BigVGAN-v2 | `fastvideo/models/audio/bigvgan.py` | Shared native vocoder | exact |
+| Euler flow schedule | shared `FlowMatchEulerDiscreteScheduler` | Reuse schedule; preserve official BF16 scalar update in MMAudio stage | exact |
+
+## Pipeline Integration
+
+- Pipeline: `fastvideo/pipelines/basic/mmaudio/MMAudioPipeline`
+- Config: `fastvideo/configs/pipelines/mmaudio.py::MMAudioV2AConfig`
+- Preset: `mmaudio_large_44k_v2`
+- Registry: resolves both `WorkloadType.V2A` and `WorkloadType.T2A`
+- Required production components: `transformer`, `scheduler`, `text_encoder`,
+  `tokenizer`, `image_encoder`, `image_encoder_2`, `audio_vae`, `vocoder`
+- Output: mono `[B,1,samples]`, 44.1 kHz, exposed through FastVideo's
+  audio-only result contract and saved as WAV by `VideoGenerator`
+- Duration: dynamic sequence lengths, with 8 seconds retained only as the
+  published training/default duration; longer and shorter inference is accepted
+- Existing T2V/I2V/T2I pipelines are not routed through these stages.
+
+The V2A preprocessing contract is identical to official MMAudio:
+
+1. timestamp sampling at 8 FPS for DFN5B and 25 FPS for Synchformer;
+2. DFN path: bicubic resize to `384x384`, float `[0,1]`, CLIP normalization;
+3. sync path: bicubic short-side resize to 224, center crop, normalize to `[-1,1]`;
+4. Synchformer: 16-frame windows, stride 8, `(segment,time)` token flattening.
+
+## Converted Checkpoint
+
+- Converter: `scripts/checkpoint_conversion/convert_mmaudio_to_diffusers.py`
+- Local artifact: `converted_weights/mmaudio/large_44k_v2`
+- Production strict load: pass for all eight components
+- Total local artifact size: about 9.1 GB
+- Converted weights and official assets remain ignored/untracked.
+
+## Parity Evidence
+
+| Scope | Result |
+|---|---|
+| Official/FastVideo video preprocessing | exact (`clip max_abs=0`, `sync max_abs=0`) |
+| Condition features, random latent, projected conditions | exact |
+| First flow prediction and 25-step final latent | exact |
+| Final 2-second V2A waveform (89,088 samples) | exact (`atol=0`, `rtol=0`) |
+| Real 10-second variable-duration V2A | pass (441,344 samples, 10.0078 s) |
+| Default FastVideo offload path | real one-step smoke pass |
+| Local suite | full rerun pending; VAE FP32/BF16 random-weight and real-weight cases pass |
+| VAE component parity | state structure and FP32/BF16 random-weight decoder parity pass on GB200 (`atol=1e-6`, `rtol=1e-6`); real-weight encode/decode parity passes (`atol=1e-5`, `rtol=1e-5`) |
+| Exact-head GB200 V2A examples | two 5-second, 25-step, seed-42 generations pass; each output is mono PCM16 at 44.1 kHz with 221,184 samples |
+
+Commands:
+
+```bash
+pytest -q tests/local_tests/mmaudio
+
+MMAUDIO_RUN_PIPELINE_PARITY=1 \
+MMAUDIO_PARITY_VIDEO=/path/to/video-at-least-2s.mp4 \
+pytest -q tests/local_tests/mmaudio/test_mmaudio_pipeline_parity.py::test_mmaudio_real_v2a_pipeline_waveform_parity -s
+```
+
+The opt-in real pipeline gate passed on an RTX 6000 Ada with the downloaded
+official `large_44k_v2`, DFN5B, Synchformer, VAE, and BigVGAN assets.
+
+## Important Numeric Decisions
+
+- OpenCLIP text uses the explicit additive causal mask used by
+  `nn.MultiheadAttention`; SDPA's `is_causal` shortcut rounds differently in BF16.
+- The Euler time/delta scalars stay on CPU, matching official MMAudio's
+  `torch.linspace` loop. Moving those float32 scalars to CUDA changes BF16 promotion.
+- `t_embed.freqs` is materialized in BF16 after meta loading; dynamic RoPE buffers
+  are rebuilt in FP32 exactly as official `update_seq_lengths` does.
+- MMAudio VAE and BigVGAN weight norm is removed on CPU in FP32 before casting to
+  BF16, matching the official feature utility construction order.
+
+## Deferred Scope
+
+- Publishing the converted checkpoint and immutable source revisions.
+- Optional source-video mux/re-encode helper; the current V2A result is WAV/audio.
+- 16 kHz and small/medium variants.
+- Sequence/tensor-parallel optimization.
+- Training integration. The official repository does not support training the
+  `_v2` variant; any training port should start from a v1 44.1 kHz checkpoint.

tests/local_tests/mmaudio/README.md

@@ -0,0 +1,149 @@
+# MMAudio Local Tests
+
+Local-only parity and smoke tests for the native MMAudio FastVideo port. These
+tests compare FastVideo against the official reference implementation and are
+not expected to run in CI unless explicitly promoted later.
+
+Port progress, open questions, issues, and handoff notes live in
+`tests/local_tests/mmaudio/PORT_STATUS.md`.
+
+## Reference Assets
+
+| Field | Value |
+|---|---|
+| Model family | `mmaudio` |
+| Workload types | `V2A`, `T2A` (new native workload values; no T2V compatibility shim) |
+| Official reference | `https://github.com/hkchengrex/MMAudio` |
+| Local reference dir | `../MMAudio` relative to the FastVideo repository |
+| Official commit/version | `974010a026c731054592d8f777218bd9d85a6c24` |
+| HF weights | `hkchengrex/MMAudio` plus canonical DFN5B CLIP and BigVGAN component repos |
+| HF revision | default; pin immutable revisions before publishing converted artifacts |
+| Local weights dir | `official_weights/mmaudio` (downloaded locally and gitignored) |
+| Source layout | mixed raw official checkpoints and external pretrained components |
+| Needs conversion | yes |
+
+The public reference assets have been downloaded locally. Never write token
+values in this file; use `HF_TOKEN`, `HUGGINGFACE_HUB_TOKEN`, or `HF_API_KEY`
+only in the shell when a future gated asset requires one.
+
+## First Implementation Scope
+
+- Inference parity target: `large_44k_v2`, 44.1 kHz, V2A and T2A.
+- Training parity target: a v1 44.1 kHz checkpoint because the official
+  repository states that `_v2` training is unsupported.
+- Inputs: video plus optional text for V2A, or text-only for T2A.
+- Output: mono waveform and sample rate through FastVideo's audio-only result
+  contract. Source-video muxing is deferred.
+- Duration: 8 seconds is the published training/default duration, but inference
+  uses dynamic sequence lengths and accepts shorter or longer clips. As in the
+  official demo, quality can fall when moving far away from 8 seconds.
+- Deferred until the base parity gate passes: 16 kHz, small/medium variants,
+  sequence/tensor parallel optimization, and quality-reference publication.
+
+## Shared Environment Setup
+
+Run from the FastVideo repository root in the same environment used for
+FastVideo. Do not create a separate upstream-only environment: both sides of a
+parity test must share the same PyTorch/CUDA numeric stack.
+
+The shared environment is `FastVideo/.venv`, created with uv-managed CPython
+3.12.13. FastVideo is installed editable with its development dependencies;
+MMAudio is installed editable with `--no-deps`, followed by the missing
+reference dependencies. NumPy is pinned to 2.0.2, satisfying both MMAudio's
+`numpy<2.1` constraint and FastVideo's installed OpenCV/SciPy constraints.
+
+Recreate the environment with:
+
+```bash
+uv venv --python 3.12 --seed
+source .venv/bin/activate
+UV_TORCH_BACKEND=cu126 uv pip install -e ".[dev]"
+uv pip install --no-deps -e ../MMAudio
+uv pip install cython "gitpython>=3.1" "hydra-core>=1.3.2" \
+  "torchdiffeq>=0.2.5" "librosa>=0.8.1" nitrous-ema hydra_colorlog \
+  "tensordict>=0.6.1" colorlog "open_clip_torch>=2.29.0" "numpy==2.0.2"
+```
+
+## Official Environment Status
+
+```text
+dependency_changes: installed no-deps editable plus missing official dependencies in current env
+official_env_status: imports_ok
+private_dep_stubs: none planned
+blocked_on: none
+```
+
+Run the FastVideo demo with an explicit duration:
+
+```bash
+MMAUDIO_MODEL_PATH=converted_weights/mmaudio/large_44k_v2 \
+python examples/inference/basic/basic_mmaudio.py \
+  --video-path /path/to/video.mp4 \
+  --duration-seconds 10 \
+  --output-path outputs_audio/mmaudio_10s.wav
+```
+
+## Weight Setup
+
+The approved reference set is present: `mmaudio_large_44k_v2.pth`,
+`v1-44.pth`, Synchformer, DFN5B CLIP, and canonical 44.1 kHz BigVGAN-v2.
+
+Converted artifacts will remain untracked under:
+
+```text
+official_weights/mmaudio/
+converted_weights/mmaudio/
+```
+
+## Prototype And Conversion Artifacts
+
+```text
+official_key_dumps:
+  transformer: converted_weights/mmaudio/_mapping/transformer_official_keys.json
+  audio_vae: converted_weights/mmaudio/_mapping/audio_vae_official_keys.json
+  synchformer: converted_weights/mmaudio/_mapping/synchformer_official_keys.json
+fastvideo_key_dumps:
+  transformer: converted_weights/mmaudio/_mapping/transformer_fastvideo_keys.json
+  audio_vae: converted_weights/mmaudio/_mapping/audio_vae_fastvideo_keys.json
+  synchformer: converted_weights/mmaudio/_mapping/synchformer_fastvideo_keys.json
+conversion_script: scripts/checkpoint_conversion/convert_mmaudio_to_diffusers.py
+conversion_source_layout: mixed
+converted_weights_dir: converted_weights/mmaudio
+strict_load_status: pass for all production components
+```
+
+## Expected Parity Tests
+
+| Component | Official files / args | Test | Concerns | Status |
+|---|---|---|---|---|
+| Transformer | `mmaudio/model/networks.py`; `large_44k_v2()` | `tests/local_tests/mmaudio/test_mmaudio_transformer_parity.py` | RoPE scaling, `nearest-exact`, official final-layer `global_c` behavior | exact real-weight parity |
+| DFN5B conditioner | `mmaudio/model/utils/features_utils.py`; `apple/DFN5B-CLIP-ViT-H-14-384` | `tests/local_tests/mmaudio/test_mmaudio_clip_parity.py` | patched tokenwise text output and normalized vision projection | exact real-weight parity |
+| Synchformer | `mmaudio/ext/synchformer/`; 16-frame windows with stride 8 | `tests/local_tests/mmaudio/test_mmaudio_synchformer_parity.py` | shared backbone must stay outside eval-only namespace | exact real-weight model and usage parity |
+| Audio VAE | `mmaudio/ext/autoencoder/`; `v1-44.pth` | `tests/local_tests/mmaudio/test_mmaudio_audio_vae_parity.py` | latent transpose and normalization statistics | exact real-weight parity |
+| BigVGAN | canonical `nvidia/bigvgan_v2_44khz_128band_512x` instantiation used by `AutoEncoderModule` | `tests/local_tests/mmaudio/test_mmaudio_vocoder_parity.py` | exact resampling and weight-norm behavior | exact real-weight parity |
+| Scheduler | `mmaudio/model/flow_matching.py::FlowMatching` | `tests/local_tests/mmaudio/test_mmaudio_scheduler_parity.py` | forward-time Euler convention | existing FastVideo scheduler parity passed |
+| Pipeline | `mmaudio/eval_utils.py::generate` and `demo.py` | `tests/local_tests/mmaudio/test_mmaudio_pipeline_parity.py` | dual-FPS frame sampling, CFG order, Euler integration, waveform output | exact real-weight parity passed |
+
+Planned commands:
+
+```bash
+pytest tests/local_tests/mmaudio/test_mmaudio_transformer_parity.py -v -s
+pytest tests/local_tests/mmaudio/test_mmaudio_clip_parity.py -v -s
+pytest tests/local_tests/mmaudio/test_mmaudio_synchformer_parity.py -v -s
+pytest tests/local_tests/mmaudio/test_mmaudio_audio_vae_parity.py -v -s
+pytest tests/local_tests/mmaudio/test_mmaudio_vocoder_parity.py -v -s
+pytest tests/local_tests/mmaudio/test_mmaudio_scheduler_parity.py -v -s
+pytest tests/local_tests/mmaudio/test_mmaudio_pipeline_parity.py -v -s
+```
+
+## Review Notes
+
+- MMAudio must be a FastVideo-native component/pipeline port; production code
+  must not import `mmaudio.*`.
+- Converted weights and large reference assets must remain ignored.
+- Every stateful component, including reused DFN5B/BigVGAN implementations,
+  requires non-skip numerical parity before final handoff.
+- Pipeline smoke and waveform parity are required before the new model is
+  presented as runnable.
+- MMAudio checkpoints are documented upstream as CC-BY-NC 4.0; converted model
+  publishing must preserve the applicable license and attribution.

tests/local_tests/mmaudio/test_mmaudio_audio_vae_parity.py

@@ -0,0 +1,105 @@
+# SPDX-License-Identifier: Apache-2.0
+"""MMAudio 44.1 kHz audio VAE parity against the official reference.
+
+Coverage scope: implementation_subcomponent. Production-loader coverage is
+added after the converted component directory exists.
+"""
+
+from __future__ import annotations
+
+import gc
+import os
+from pathlib import Path
+
+import pytest
+import torch
+
+
+REPO_ROOT = Path(__file__).resolve().parents[3]
+OFFICIAL_WEIGHTS = Path(
+    os.environ.get(
+        "MMAUDIO_AUDIO_VAE_WEIGHTS",
+        REPO_ROOT.parent / "MMAudio/ext_weights/v1-44.pth",
+    )
+)
+
+
+def _official_model():
+    official_vae = pytest.importorskip("mmaudio.ext.autoencoder.vae")
+    return official_vae.VAE_44k()
+
+
+def _fastvideo_model(*, need_encoder: bool):
+    from fastvideo.models.audio.mmaudio_vae import MMAudioVAE
+
+    return MMAudioVAE(mode="44k", need_encoder=need_encoder)
+
+
+def test_mmaudio_44k_audio_vae_state_structure() -> None:
+    official = _official_model()
+    expected = {name: tensor.shape for name, tensor in official.state_dict().items()}
+    del official
+    gc.collect()
+    fastvideo = _fastvideo_model(need_encoder=True)
+    assert expected == {
+        name: tensor.shape for name, tensor in fastvideo.state_dict().items()
+    }
+
+
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16], ids=["fp32", "bf16"])
+def test_mmaudio_44k_audio_vae_decoder_implementation_parity(dtype: torch.dtype) -> None:
+    if not torch.cuda.is_available():
+        pytest.skip("MMAudio audio VAE implementation parity requires CUDA")
+
+    official = _official_model()
+    del official.encoder
+    gc.collect()
+    fastvideo = _fastvideo_model(need_encoder=False)
+    fastvideo.load_state_dict(official.state_dict(), strict=True)
+    device = torch.device("cuda:0")
+    official.remove_weight_norm().to(device=device, dtype=dtype).eval()
+    fastvideo.remove_weight_norm().to(device=device, dtype=dtype).eval()
+    latent = torch.randn(
+        (1, 40, 4),
+        generator=torch.Generator(device=device).manual_seed(1234),
+        device=device,
+        dtype=dtype,
+    )
+
+    with torch.inference_mode():
+        expected = official.decode(latent)
+        actual = fastvideo.decode(latent)
+
+    torch.testing.assert_close(actual, expected, atol=1e-6, rtol=1e-6)
+
+
+def test_mmaudio_44k_audio_vae_numerical_parity() -> None:
+    if not torch.cuda.is_available():
+        pytest.skip("MMAudio audio VAE parity requires CUDA")
+    if not OFFICIAL_WEIGHTS.is_file():
+        pytest.skip(
+            "Official MMAudio 44.1 kHz VAE weights are absent. Set MMAUDIO_AUDIO_VAE_WEIGHTS or download v1-44.pth."
+        )
+
+    device = torch.device("cuda:0")
+    official = _official_model()
+    fastvideo = _fastvideo_model(need_encoder=True)
+    state = torch.load(OFFICIAL_WEIGHTS, map_location="cpu", weights_only=True)
+    official.load_state_dict(state, strict=True)
+    fastvideo.load_state_dict(state, strict=True)
+    official.remove_weight_norm().to(device).eval()
+    fastvideo.remove_weight_norm().to(device).eval()
+
+    generator = torch.Generator(device=device).manual_seed(1234)
+    mel = torch.randn((1, 128, 128), generator=generator, device=device)
+    latent = torch.randn((1, 40, 64), generator=generator, device=device)
+
+    with torch.inference_mode():
+        expected_posterior = official.encode(mel)
+        actual_posterior = fastvideo.encode(mel)
+        expected_mel = official.decode(latent)
+        actual_mel = fastvideo.decode(latent)
+
+    torch.testing.assert_close(actual_posterior.mean, expected_posterior.mean, atol=1e-5, rtol=1e-5)
+    torch.testing.assert_close(actual_posterior.logvar, expected_posterior.logvar, atol=1e-5, rtol=1e-5)
+    torch.testing.assert_close(actual_mel, expected_mel, atol=1e-5, rtol=1e-5)

tests/local_tests/mmaudio/test_mmaudio_clip_parity.py

@@ -0,0 +1,225 @@
+# SPDX-License-Identifier: Apache-2.0
+"""MMAudio DFN5B OpenCLIP conditioner parity tests."""
+
+from __future__ import annotations
+
+import os
+import importlib.util
+from functools import cache
+from pathlib import Path
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+
+REPO_ROOT = Path(__file__).resolve().parents[3]
+DFN5B_DIR = Path(
+    os.environ.get(
+        "MMAUDIO_DFN5B_DIR",
+        REPO_ROOT / "official_weights/mmaudio/DFN5B-CLIP-ViT-H-14-384",
+    )
+)
+
+
+@cache
+def _conversion_module():
+    path = REPO_ROOT / "scripts/checkpoint_conversion/convert_mmaudio_to_diffusers.py"
+    spec = importlib.util.spec_from_file_location("mmaudio_converter", path)
+    assert spec is not None and spec.loader is not None
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
+
+def test_converted_open_clip_tokenizer_parity(tmp_path: Path) -> None:
+    import open_clip
+    from transformers import CLIPTokenizer
+
+    _conversion_module().write_open_clip_tokenizer(tmp_path)
+    converted = CLIPTokenizer.from_pretrained(tmp_path / "tokenizer")
+    prompts = ["", "A dog runs past a red car!", "雨の日, cinematic"]
+    encoded = converted(
+        prompts,
+        padding="max_length",
+        truncation=True,
+        max_length=77,
+        return_tensors="pt",
+    )
+    actual = encoded.input_ids.masked_fill(encoded.attention_mask == 0, 0)
+    expected = open_clip.get_tokenizer("ViT-H-14-378-quickgelu")(prompts)
+    torch.testing.assert_close(actual, expected)
+
+
+def _map_open_clip_text(state: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
+    return _conversion_module().map_open_clip_text_state(state)
+
+
+def _map_open_clip_vision(state: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
+    return _conversion_module().map_open_clip_vision_state(state)
+
+
+def test_mmaudio_dfn_clip_implementation_parity() -> None:
+    if not torch.cuda.is_available():
+        pytest.skip("DFN CLIP implementation parity requires CUDA")
+
+    from open_clip.model import CLIP, CLIPTextCfg, CLIPVisionCfg
+
+    from fastvideo.configs.models.encoders.mmaudio_clip import (
+        MMAudioDFNCLIPTextArchConfig,
+        MMAudioDFNCLIPTextConfig,
+        MMAudioDFNCLIPVisionArchConfig,
+        MMAudioDFNCLIPVisionConfig,
+    )
+    from fastvideo.models.encoders.mmaudio_clip import (
+        MMAudioDFNCLIPTextEncoder,
+        MMAudioDFNCLIPVisionEncoder,
+    )
+    from fastvideo.distributed import (
+        cleanup_dist_env_and_memory,
+        maybe_init_distributed_environment_and_model_parallel,
+    )
+    from fastvideo.forward_context import set_forward_context
+    from fastvideo.platforms import AttentionBackendEnum
+
+    os.environ.setdefault("MASTER_ADDR", "127.0.0.1")
+    os.environ.setdefault("MASTER_PORT", "29591")
+    maybe_init_distributed_environment_and_model_parallel(1, 1)
+    embed_dim = 16
+    official = CLIP(
+        embed_dim=embed_dim,
+        vision_cfg=CLIPVisionCfg(layers=2, width=32, head_width=8, patch_size=8, image_size=16),
+        text_cfg=CLIPTextCfg(context_length=8, vocab_size=32, width=32, heads=4, layers=2, eos_id=31),
+        quick_gelu=True,
+    )
+    backends = (AttentionBackendEnum.TORCH_SDPA,)
+    text_arch = MMAudioDFNCLIPTextArchConfig(
+        vocab_size=32,
+        hidden_size=32,
+        intermediate_size=128,
+        projection_dim=embed_dim,
+        num_hidden_layers=2,
+        num_attention_heads=4,
+        max_position_embeddings=8,
+        text_len=8,
+        eos_token_id=31,
+        _supported_attention_backends=backends,
+    )
+    vision_arch = MMAudioDFNCLIPVisionArchConfig(
+        hidden_size=32,
+        intermediate_size=128,
+        projection_dim=embed_dim,
+        num_hidden_layers=2,
+        num_attention_heads=4,
+        image_size=16,
+        patch_size=8,
+        _supported_attention_backends=backends,
+    )
+    text_encoder = MMAudioDFNCLIPTextEncoder(MMAudioDFNCLIPTextConfig(arch_config=text_arch))
+    vision_encoder = MMAudioDFNCLIPVisionEncoder(MMAudioDFNCLIPVisionConfig(arch_config=vision_arch))
+    state = official.state_dict()
+    text_encoder.load_state_dict(_map_open_clip_text(state), strict=True)
+    vision_encoder.load_state_dict(_map_open_clip_vision(state), strict=True)
+
+    device = torch.device("cuda:0")
+    official.to(device).eval()
+    text_encoder.to(device).eval()
+    vision_encoder.to(device).eval()
+    tokens = torch.tensor([[1, 2, 3, 31, 0, 0, 0, 0]], device=device)
+    images = torch.randn((1, 3, 16, 16), generator=torch.Generator(device=device).manual_seed(1234), device=device)
+    with torch.inference_mode():
+        expected_text = official.token_embedding(tokens)
+        expected_text = expected_text + official.positional_embedding
+        expected_text = official.transformer(expected_text, attn_mask=official.attn_mask)
+        expected_text = F.normalize(official.ln_final(expected_text), dim=-1)
+        expected_image = official.encode_image(images, normalize=True)
+        with set_forward_context(current_timestep=0, attn_metadata=None):
+            actual_text = text_encoder(tokens).last_hidden_state
+            actual_image = vision_encoder(images).last_hidden_state
+
+    torch.testing.assert_close(actual_text, expected_text, atol=2e-5, rtol=2e-5)
+    torch.testing.assert_close(actual_image, expected_image, atol=2e-5, rtol=2e-5)
+    cleanup_dist_env_and_memory()
+
+
+def test_mmaudio_dfn5b_real_weight_parity() -> None:
+    if not DFN5B_DIR.is_dir():
+        pytest.skip("DFN5B assets are absent. Set MMAUDIO_DFN5B_DIR to a local apple/DFN5B-CLIP-ViT-H-14-384 snapshot.")
+    if not torch.cuda.is_available():
+        pytest.skip("DFN5B parity requires CUDA")
+
+    import open_clip
+
+    from fastvideo.configs.models.encoders.mmaudio_clip import (
+        MMAudioDFNCLIPTextConfig,
+        MMAudioDFNCLIPVisionConfig,
+    )
+    from fastvideo.distributed import (
+        cleanup_dist_env_and_memory,
+        maybe_init_distributed_environment_and_model_parallel,
+    )
+    from fastvideo.forward_context import set_forward_context
+    from fastvideo.models.encoders.mmaudio_clip import (
+        MMAudioDFNCLIPTextEncoder,
+        MMAudioDFNCLIPVisionEncoder,
+    )
+
+    os.environ.setdefault("MASTER_ADDR", "127.0.0.1")
+    os.environ.setdefault("MASTER_PORT", "29592")
+    maybe_init_distributed_environment_and_model_parallel(1, 1)
+    try:
+        official = open_clip.create_model_from_pretrained(
+            f"local-dir:{DFN5B_DIR}", return_transform=False)
+        state = official.state_dict()
+        text_encoder = MMAudioDFNCLIPTextEncoder(
+            MMAudioDFNCLIPTextConfig())
+        vision_encoder = MMAudioDFNCLIPVisionEncoder(
+            MMAudioDFNCLIPVisionConfig())
+        text_encoder.load_state_dict(_map_open_clip_text(state), strict=True)
+        vision_encoder.load_state_dict(_map_open_clip_vision(state),
+                                       strict=True)
+
+        device = torch.device("cuda:0")
+        official.to(device).eval()
+        text_encoder.to(device).eval()
+        vision_encoder.to(device).eval()
+        tokens = open_clip.get_tokenizer("ViT-H-14-378-quickgelu")(
+            ["A dog runs past a red car!"]).to(device)
+        frames = torch.rand(
+            (1, 3, 384, 384),
+            generator=torch.Generator(device=device).manual_seed(9012),
+            device=device,
+        )
+        mean = torch.tensor([0.48145466, 0.4578275, 0.40821073],
+                            device=device).view(1, 3, 1, 1)
+        std = torch.tensor([0.26862954, 0.26130258, 0.27577711],
+                           device=device).view(1, 3, 1, 1)
+        frames = (frames - mean) / std
+
+        with torch.inference_mode():
+            expected_text = official.token_embedding(tokens)
+            expected_text = expected_text + official.positional_embedding
+            expected_text = official.transformer(
+                expected_text, attn_mask=official.attn_mask)
+            expected_text = F.normalize(official.ln_final(expected_text),
+                                        dim=-1)
+            expected_image = official.encode_image(frames, normalize=True)
+            with set_forward_context(current_timestep=0,
+                                     attn_metadata=None):
+                actual_text = text_encoder(tokens).last_hidden_state
+                actual_image = vision_encoder(frames).last_hidden_state
+
+        text_error = (actual_text.float() - expected_text.float()).abs()
+        image_error = (actual_image.float() - expected_image.float()).abs()
+        print("text_max_abs", text_error.max().item())
+        print("image_max_abs", image_error.max().item())
+        torch.testing.assert_close(actual_text,
+                                   expected_text,
+                                   atol=2e-5,
+                                   rtol=2e-5)
+        torch.testing.assert_close(actual_image,
+                                   expected_image,
+                                   atol=2e-5,
+                                   rtol=2e-5)
+    finally:
+        cleanup_dist_env_and_memory()

tests/local_tests/mmaudio/test_mmaudio_component_loader.py

@@ -0,0 +1,62 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Shared component-loader extensions required by V2A pipelines."""
+
+from __future__ import annotations
+
+import json
+from types import SimpleNamespace
+
+
+def test_indexed_image_encoder_uses_matching_config_and_precision(tmp_path) -> None:
+    from fastvideo.configs.models.encoders.mmaudio_clip import (
+        MMAudioDFNCLIPVisionConfig,
+    )
+    from fastvideo.configs.models.encoders.mmaudio_synchformer import (
+        MMAudioSynchformerConfig,
+    )
+    from fastvideo.models.loader.component_loader import ImageEncoderLoader
+
+    component = tmp_path / "image_encoder_2"
+    component.mkdir()
+    with (component / "config.json").open("w", encoding="utf-8") as handle:
+        json.dump(
+            {
+                "architectures": ["MMAudioSynchformerVisualEncoder"],
+                "segment_stride": 4,
+            },
+            handle,
+        )
+
+    class CaptureLoader(ImageEncoderLoader):
+        def load_model(
+            self,
+            model_path,
+            model_config,
+            target_device,
+            fastvideo_args,
+            dtype="fp16",
+            use_text_encoder_override=False,
+            cpu_offload=None,
+        ):
+            del model_path, target_device, fastvideo_args, use_text_encoder_override
+            return model_config, dtype, cpu_offload
+
+    vision_config = MMAudioDFNCLIPVisionConfig()
+    sync_config = MMAudioSynchformerConfig()
+    pipeline_config = SimpleNamespace(
+        image_encoder_config=vision_config,
+        image_encoder_precision="fp32",
+        image_encoder_configs=(vision_config, sync_config),
+        image_encoder_precisions=("bf16", "fp16"),
+    )
+    args = SimpleNamespace(
+        pipeline_config=pipeline_config,
+        image_encoder_cpu_offload=True,
+    )
+
+    selected, precision, cpu_offload = CaptureLoader().load(str(component), args)
+    assert selected is sync_config
+    assert selected.arch_config.segment_stride == 4
+    assert vision_config.arch_config.image_size == 378
+    assert precision == "fp16"
+    assert cpu_offload is True

tests/local_tests/mmaudio/test_mmaudio_pipeline_parity.py

@@ -0,0 +1,245 @@
+# SPDX-License-Identifier: Apache-2.0
+"""MMAudio pipeline contracts and opt-in real end-to-end parity."""
+
+from __future__ import annotations
+
+import gc
+import math
+import os
+from pathlib import Path
+
+import pytest
+import torch
+
+
+REPO_ROOT = Path(__file__).resolve().parents[3]
+OFFICIAL_REPO = REPO_ROOT.parent / "MMAudio"
+CONVERTED_MODEL = REPO_ROOT / "converted_weights/mmaudio/large_44k_v2"
+DFN5B_DIR = REPO_ROOT / "official_weights/mmaudio/DFN5B-CLIP-ViT-H-14-384"
+BIGVGAN_DIR = REPO_ROOT / "official_weights/mmaudio/bigvgan_v2_44khz_128band_512x"
+
+
+def test_mmaudio_pipeline_config_registry_and_preset() -> None:
+    from fastvideo.configs.pipelines.mmaudio import MMAudioV2AConfig
+    from fastvideo.fastvideo_args import WorkloadType
+    from fastvideo.registry import get_model_info, get_preset_selection
+
+    assert WorkloadType.from_string("v2a") is WorkloadType.V2A
+    assert WorkloadType.from_string("t2a") is WorkloadType.T2A
+    config = MMAudioV2AConfig()
+    assert config.dit_config.prefix == "MMAudio"
+    assert len(config.image_encoder_configs or ()) == 2
+    assert config.text_encoder_precisions == ("bf16",)
+    assert config.image_encoder_precisions == ("bf16", "bf16")
+    assert config.audio_decoder_precision == "bf16"
+    assert config.vocoder_precision == "bf16"
+    assert config.duration_s == 8.0
+    assert config.max_audio_duration_s is None
+    assert get_preset_selection("FastVideo/MMAudio-large-44k-v2-Diffusers") == (
+        "mmaudio_large_44k_v2",
+        "mmaudio",
+    )
+
+    if CONVERTED_MODEL.is_dir():
+        info = get_model_info(str(CONVERTED_MODEL), workload_type=WorkloadType.V2A)
+        assert info.pipeline_cls.__name__ == "MMAudioPipeline"
+        assert info.pipeline_config_cls is MMAudioV2AConfig
+
+
+def test_mmaudio_published_sequence_lengths() -> None:
+    from fastvideo.configs.pipelines.mmaudio import MMAudioV2AConfig
+    from fastvideo.pipelines.basic.mmaudio.stages import mmaudio_sequence_lengths
+
+    assert mmaudio_sequence_lengths(8.0, MMAudioV2AConfig()) == (345, 64, 192)
+    assert mmaudio_sequence_lengths(10.0, MMAudioV2AConfig()) == (431, 80, 240)
+    assert mmaudio_sequence_lengths(2.0, MMAudioV2AConfig()) == (87, 16, 40)
+
+
+def test_mmaudio_bf16_euler_stage_matches_official() -> None:
+    if not torch.cuda.is_available():
+        pytest.skip("MMAudio BF16 scheduler parity requires CUDA")
+
+    from mmaudio.model.flow_matching import FlowMatching
+
+    from fastvideo.configs.pipelines.mmaudio import MMAudioV2AConfig
+    from fastvideo.fastvideo_args import FastVideoArgs
+    from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import (
+        FlowMatchEulerDiscreteScheduler,
+    )
+    from fastvideo.pipelines.basic.mmaudio.stages import MMAudioDenoisingStage
+    from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+
+    class ToyTransformer:
+        @staticmethod
+        def guided_flow(timestep, latent, conditions, empty_conditions, guidance_scale):
+            del conditions, empty_conditions, guidance_scale
+            return torch.tanh(latent * 0.125 + timestep)
+
+        @staticmethod
+        def unnormalize(latent):
+            return latent
+
+    device = torch.device("cuda:0")
+    initial = torch.randn(
+        (1, 87, 40),
+        device=device,
+        dtype=torch.bfloat16,
+        generator=torch.Generator(device=device).manual_seed(42),
+    )
+    official = FlowMatching(min_sigma=0.0, inference_mode="euler", num_steps=25)
+    expected = official.to_data(lambda time, latent: torch.tanh(latent * 0.125 + time), initial.clone())
+
+    scheduler = FlowMatchEulerDiscreteScheduler(
+        shift=1.0,
+        invert_sigmas=True,
+        sigma_min=0.0,
+        use_reference_discrete_timesteps=True,
+    )
+    batch = ForwardBatch(data_type="audio", latents=initial.clone(), num_inference_steps=25, guidance_scale=4.5)
+    batch.extra["mmaudio_conditions"] = object()
+    batch.extra["mmaudio_empty_conditions"] = object()
+    actual = MMAudioDenoisingStage(ToyTransformer(), scheduler)(
+        batch,
+        FastVideoArgs(model_path="test", pipeline_config=MMAudioV2AConfig()),
+    ).latents
+    torch.testing.assert_close(actual, expected, atol=0, rtol=0)
+
+
+@pytest.mark.skipif(
+    os.environ.get("MMAUDIO_RUN_PIPELINE_PARITY") != "1",
+    reason="Set MMAUDIO_RUN_PIPELINE_PARITY=1 and MMAUDIO_PARITY_VIDEO to run the 25-step real-weight gate.",
+)
+def test_mmaudio_real_v2a_pipeline_waveform_parity(monkeypatch: pytest.MonkeyPatch) -> None:
+    """Compare official and FastVideo waveforms with the same 2s video/seed."""
+    video_value = os.environ.get("MMAUDIO_PARITY_VIDEO")
+    if not video_value:
+        pytest.skip("Set MMAUDIO_PARITY_VIDEO to a video containing at least two seconds.")
+    video_path = Path(video_value)
+    required = (
+        OFFICIAL_REPO / "weights/mmaudio_large_44k_v2.pth",
+        OFFICIAL_REPO / "ext_weights/v1-44.pth",
+        OFFICIAL_REPO / "ext_weights/synchformer_state_dict.pth",
+        CONVERTED_MODEL / "model_index.json",
+        DFN5B_DIR / "open_clip_pytorch_model.bin",
+        BIGVGAN_DIR / "bigvgan_generator.pt",
+    )
+    missing = [str(path) for path in required if not path.exists()]
+    if missing:
+        pytest.skip(f"MMAudio real pipeline assets are missing: {missing}")
+    if not torch.cuda.is_available():
+        pytest.skip("MMAudio real pipeline parity requires CUDA")
+
+    import open_clip
+    import mmaudio.ext.autoencoder.autoencoder as official_autoencoder
+    import mmaudio.model.utils.features_utils as official_features_module
+    from mmaudio.eval_utils import generate as official_generate
+    from mmaudio.eval_utils import load_video as official_load_video
+    from mmaudio.ext.bigvgan_v2.bigvgan import BigVGAN as OfficialBigVGAN
+    from mmaudio.model.flow_matching import FlowMatching
+    from mmaudio.model.networks import get_my_mmaudio
+    from mmaudio.model.utils.features_utils import FeaturesUtils
+
+    monkeypatch.setattr(
+        official_features_module,
+        "create_model_from_pretrained",
+        lambda *args, **kwargs: open_clip.create_model_from_pretrained(
+            f"local-dir:{DFN5B_DIR}", return_transform=False
+        ),
+    )
+
+    class LocalBigVGAN:
+        @classmethod
+        def from_pretrained(cls, model_id, **kwargs):
+            del model_id
+            return OfficialBigVGAN.from_pretrained(str(BIGVGAN_DIR), **kwargs)
+
+    monkeypatch.setattr(official_autoencoder, "BigVGANv2", LocalBigVGAN)
+    torch.backends.cuda.matmul.allow_tf32 = True
+    torch.backends.cudnn.allow_tf32 = True
+    device = torch.device("cuda:0")
+    dtype = torch.bfloat16
+    with torch.inference_mode():
+        official_model = get_my_mmaudio("large_44k_v2").to(device, dtype).eval()
+        official_model.load_weights(
+            torch.load(required[0], map_location=device, weights_only=True)
+        )
+        official_features = FeaturesUtils(
+            tod_vae_ckpt=required[1],
+            synchformer_ckpt=required[2],
+            enable_conditions=True,
+            mode="44k",
+            need_vae_encoder=False,
+        ).to(device, dtype).eval()
+        video = official_load_video(video_path, 2.0, load_all_frames=False)
+        duration = video.duration_sec
+        official_model.update_seq_lengths(
+            math.ceil(duration * 44100 / 512 / 2),
+            int(duration * 8),
+            int((((duration * 25) - 16) // 8 + 1) * 16 / 2),
+        )
+        expected = official_generate(
+            video.clip_frames[None],
+            video.sync_frames[None],
+            ["A dog runs past a red car!"],
+            negative_text=[""],
+            feature_utils=official_features,
+            net=official_model,
+            fm=FlowMatching(min_sigma=0, inference_mode="euler", num_steps=25),
+            rng=torch.Generator(device=device).manual_seed(42),
+            cfg_strength=4.5,
+        ).float().cpu()
+    del official_features, official_model
+    gc.collect()
+    torch.cuda.empty_cache()
+
+    from fastvideo.configs.pipelines.mmaudio import MMAudioV2AConfig
+    from fastvideo.distributed import cleanup_dist_env_and_memory
+    from fastvideo.fastvideo_args import FastVideoArgs, WorkloadType
+    from fastvideo.pipelines.basic.mmaudio import MMAudioPipeline
+    from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+
+    args = FastVideoArgs(
+        model_path=str(CONVERTED_MODEL),
+        workload_type=WorkloadType.V2A,
+        pipeline_config=MMAudioV2AConfig(),
+        tp_size=1,
+        sp_size=1,
+        hsdp_shard_dim=1,
+        num_gpus=1,
+        dit_cpu_offload=False,
+        dit_layerwise_offload=False,
+        text_encoder_cpu_offload=False,
+        image_encoder_cpu_offload=False,
+        vae_cpu_offload=False,
+        pin_cpu_memory=False,
+    )
+    try:
+        pipeline = MMAudioPipeline(str(CONVERTED_MODEL), args)
+        pipeline.post_init()
+        output = pipeline.forward(
+            ForwardBatch(
+                data_type="video",
+                video_path=str(video_path),
+                prompt="A dog runs past a red car!",
+                negative_prompt="",
+                audio_start_in_s=0.0,
+                audio_end_in_s=2.0,
+                num_inference_steps=25,
+                guidance_scale=4.5,
+                seed=42,
+                num_videos_per_prompt=1,
+                height=8,
+                width=8,
+                num_frames=1,
+                save_video=False,
+                return_frames=False,
+            ),
+            args,
+        )
+        actual = output.extra["decoded_audio"]
+        assert output.extra["audio_sample_rate"] == 44100
+        assert output.extra["audio_only"] is True
+        torch.testing.assert_close(actual, expected, atol=0, rtol=0)
+        pipeline.close()
+    finally:
+        cleanup_dist_env_and_memory()

tests/local_tests/mmaudio/test_mmaudio_scheduler_parity.py

@@ -0,0 +1,37 @@
+# SPDX-License-Identifier: Apache-2.0
+"""MMAudio flow-matching scheduler reuse parity."""
+
+from __future__ import annotations
+
+import torch
+
+
+def test_mmaudio_euler_schedule_matches_fastvideo_flow_scheduler() -> None:
+    from mmaudio.model.flow_matching import FlowMatching
+
+    from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import (
+        FlowMatchEulerDiscreteScheduler,
+    )
+
+    num_steps = 25
+    official = FlowMatching(min_sigma=0.0, inference_mode="euler", num_steps=num_steps)
+    fastvideo = FlowMatchEulerDiscreteScheduler(
+        shift=1.0,
+        invert_sigmas=True,
+        sigma_min=0.0,
+        use_reference_discrete_timesteps=True,
+    )
+    fastvideo.set_timesteps(num_steps, device="cpu")
+    initial = torch.randn((2, 7, 4), generator=torch.Generator().manual_seed(1234))
+
+    def flow(time: torch.Tensor, sample: torch.Tensor) -> torch.Tensor:
+        time = torch.as_tensor(time, dtype=sample.dtype, device=sample.device)
+        return torch.tanh(sample * 0.125 + time)
+
+    expected = official.to_data(flow, initial.clone())
+    actual = initial.clone()
+    for timestep in fastvideo.timesteps:
+        model_output = flow(timestep / fastvideo.config.num_train_timesteps, actual)
+        actual = fastvideo.step(model_output, timestep, actual).prev_sample
+
+    torch.testing.assert_close(actual, expected, atol=1e-6, rtol=1e-6)

tests/local_tests/mmaudio/test_mmaudio_synchformer_parity.py

@@ -0,0 +1,146 @@
+# SPDX-License-Identifier: Apache-2.0
+"""MMAudio Synchformer parity against the official reference.
+
+Coverage scope: implementation_subcomponent. Production-loader coverage is
+added after the converted component directory exists.
+"""
+
+from __future__ import annotations
+
+import os
+import gc
+from pathlib import Path
+
+import pytest
+import torch
+
+
+REPO_ROOT = Path(__file__).resolve().parents[3]
+OFFICIAL_WEIGHTS = Path(
+    os.environ.get(
+        "MMAUDIO_SYNCHFORMER_WEIGHTS",
+        REPO_ROOT.parent / "MMAudio/ext_weights/synchformer_state_dict.pth",
+    )
+)
+
+
+def _build_models(device: torch.device, dtype: torch.dtype):
+    from mmaudio.ext.synchformer import Synchformer
+
+    from fastvideo.configs.models.encoders.mmaudio_synchformer import (
+        MMAudioSynchformerConfig,
+    )
+    from fastvideo.models.encoders.mmaudio_synchformer import (
+        MMAudioSynchformerVisualEncoder,
+    )
+    from fastvideo.models.loader.utils import set_default_torch_dtype
+
+    with torch.device(device), set_default_torch_dtype(dtype):
+        official = Synchformer()
+        fastvideo = MMAudioSynchformerVisualEncoder(MMAudioSynchformerConfig())
+    return official, fastvideo
+
+
+def test_mmaudio_synchformer_state_structure() -> None:
+    # MotionFormer calls ``Tensor.item`` while building its stochastic-depth
+    # schedule, so it cannot be constructed on the meta device. Instantiate
+    # the two large models sequentially to keep peak host memory bounded.
+    from mmaudio.ext.synchformer import Synchformer
+
+    official = Synchformer()
+    official_state = official.state_dict()
+    official_shapes = {name: tensor.shape for name, tensor in official_state.items()}
+    del official_state, official
+    gc.collect()
+
+    from fastvideo.configs.models.encoders.mmaudio_synchformer import (
+        MMAudioSynchformerConfig,
+    )
+    from fastvideo.models.encoders.mmaudio_synchformer import (
+        MMAudioSynchformerVisualEncoder,
+    )
+
+    fastvideo = MMAudioSynchformerVisualEncoder(MMAudioSynchformerConfig())
+    fastvideo_shapes = {name: tensor.shape for name, tensor in fastvideo.state_dict().items()}
+
+    assert official_shapes == fastvideo_shapes
+
+
+def test_mmaudio_synchformer_numerical_parity() -> None:
+    if not torch.cuda.is_available():
+        pytest.skip("MMAudio Synchformer parity requires CUDA")
+    if not OFFICIAL_WEIGHTS.is_file():
+        pytest.skip(
+            "Official MMAudio Synchformer weights are absent. Set "
+            "MMAUDIO_SYNCHFORMER_WEIGHTS or download the released checkpoint."
+        )
+
+    device = torch.device("cuda:0")
+    dtype = torch.bfloat16
+    official, fastvideo = _build_models(device, dtype)
+    state = torch.load(OFFICIAL_WEIGHTS, map_location="cpu", weights_only=True)
+    official.load_state_dict(state, strict=True)
+    visual_state = {
+        name: tensor
+        for name, tensor in state.items()
+        if name.startswith("vfeat_extractor.")
+    }
+    fastvideo.load_state_dict(visual_state, strict=True)
+    official.eval()
+    fastvideo.eval()
+
+    generator = torch.Generator(device=device).manual_seed(1234)
+    segments = torch.randn((1, 1, 16, 3, 224, 224), generator=generator, device=device, dtype=dtype)
+
+    with torch.inference_mode(), torch.autocast("cuda", dtype=dtype):
+        expected = official(segments)
+        actual = fastvideo.forward_segmented(segments)
+
+    difference = (actual.float() - expected.float()).abs()
+    print("max_abs", difference.max().item())
+    print("mean_abs", difference.mean().item())
+    torch.testing.assert_close(actual, expected, atol=1e-3, rtol=1e-3)
+
+
+def test_mmaudio_synchformer_feature_contract_parity() -> None:
+    """Cover official windowing, batching, flatten order, and model forward."""
+    if not torch.cuda.is_available():
+        pytest.skip("MMAudio Synchformer feature parity requires CUDA")
+    if not OFFICIAL_WEIGHTS.is_file():
+        pytest.skip("Official MMAudio Synchformer weights are absent")
+
+    from mmaudio.model.utils.features_utils import FeaturesUtils
+
+    device = torch.device("cuda:0")
+    dtype = torch.bfloat16
+    official, fastvideo = _build_models(device, dtype)
+    state = torch.load(OFFICIAL_WEIGHTS, map_location="cpu", weights_only=True)
+    official.load_state_dict(state, strict=True)
+    fastvideo.load_state_dict(
+        {
+            name: tensor
+            for name, tensor in state.items()
+            if name.startswith("vfeat_extractor.")
+        },
+        strict=True,
+    )
+    official.eval()
+    fastvideo.eval()
+
+    feature_utils = FeaturesUtils(enable_conditions=False,
+                                  mode="44k",
+                                  need_vae_encoder=False)
+    feature_utils.synchformer = official
+    video = torch.randn(
+        (2, 24, 3, 224, 224),
+        generator=torch.Generator(device=device).manual_seed(5678),
+        device=device,
+        dtype=dtype,
+    ).clamp_(-1, 1)
+
+    with torch.inference_mode(), torch.autocast("cuda", dtype=dtype):
+        expected = feature_utils.encode_video_with_sync(video, batch_size=2)
+        actual = fastvideo(video).last_hidden_state
+
+    assert expected.shape == actual.shape == (2, 16, 768)
+    torch.testing.assert_close(actual, expected, atol=1e-3, rtol=1e-3)

tests/local_tests/mmaudio/test_mmaudio_transformer_parity.py

@@ -0,0 +1,126 @@
+# SPDX-License-Identifier: Apache-2.0
+"""MMAudio transformer parity against the official reference.
+
+Coverage scope: implementation_subcomponent. Production-loader coverage is
+added after the converted Diffusers component directory exists.
+"""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+import pytest
+import torch
+
+
+REPO_ROOT = Path(__file__).resolve().parents[3]
+OFFICIAL_WEIGHTS = Path(
+    os.environ.get(
+        "MMAUDIO_TRANSFORMER_WEIGHTS",
+        REPO_ROOT.parent / "MMAudio/weights/mmaudio_large_44k_v2.pth",
+    )
+)
+
+
+def _require_cuda_and_weights() -> None:
+    if not torch.cuda.is_available():
+        pytest.skip("MMAudio full-transformer parity requires CUDA")
+    if not OFFICIAL_WEIGHTS.is_file():
+        pytest.skip(
+            "Official MMAudio transformer weights are absent. Set "
+            "MMAUDIO_TRANSFORMER_WEIGHTS or download large_44k_v2 assets."
+        )
+
+
+def test_mmaudio_transformer_implementation_parity() -> None:
+    from mmaudio.model.networks import MMAudio
+
+    from fastvideo.configs.models.dits.mmaudio import (
+        MMAudioArchConfig,
+        MMAudioTransformerConfig,
+    )
+    from fastvideo.models.dits.mmaudio import MMAudioTransformer
+
+    kwargs = {
+        "latent_dim": 8,
+        "clip_dim": 16,
+        "sync_dim": 12,
+        "text_dim": 16,
+        "hidden_dim": 64,
+        "depth": 3,
+        "fused_depth": 2,
+        "num_heads": 4,
+        "mlp_ratio": 4.0,
+        "latent_seq_len": 5,
+        "clip_seq_len": 3,
+        "sync_seq_len": 8,
+        "text_seq_len": 4,
+        "v2": True,
+    }
+    official = MMAudio(**kwargs)
+    fastvideo = MMAudioTransformer(MMAudioTransformerConfig(arch_config=MMAudioArchConfig(**kwargs)), hf_config={})
+    assert {name: tensor.shape for name, tensor in official.state_dict().items()} == {
+        name: tensor.shape for name, tensor in fastvideo.state_dict().items()
+    }
+    fastvideo.load_state_dict(official.state_dict(), strict=True)
+    generator = torch.Generator().manual_seed(1234)
+    latent = torch.randn((1, 5, 8), generator=generator)
+    clip = torch.randn((1, 3, 16), generator=generator)
+    sync = torch.randn((1, 8, 12), generator=generator)
+    text = torch.randn((1, 4, 16), generator=generator)
+    timestep = torch.tensor([0.375])
+    official.eval()
+    fastvideo.eval()
+    with torch.inference_mode():
+        expected = official(latent, clip, sync, text, timestep)
+        actual = fastvideo(latent, (clip, sync, text), timestep)
+    torch.testing.assert_close(actual, expected, atol=1e-6, rtol=1e-6)
+
+
+def test_mmaudio_large_44k_v2_transformer_parity() -> None:
+    _require_cuda_and_weights()
+    from mmaudio.model.networks import large_44k_v2
+
+    from fastvideo.configs.models.dits.mmaudio import MMAudioTransformerConfig
+    from fastvideo.models.dits.mmaudio import MMAudioTransformer
+    from fastvideo.models.loader.utils import set_default_torch_dtype
+
+    device = torch.device("cuda:0")
+    dtype = torch.bfloat16
+    state = torch.load(OFFICIAL_WEIGHTS, map_location="cpu", weights_only=True)
+
+    with torch.device(device), set_default_torch_dtype(dtype):
+        official = large_44k_v2()
+        fastvideo = MMAudioTransformer(MMAudioTransformerConfig(), hf_config={})
+    official.load_weights(dict(state))
+    # The released checkpoint contains a stale derived buffer which the
+    # official ``load_weights`` explicitly discards. Validate it before
+    # applying the same canonicalization used by the converter.
+    checkpoint_freqs = state.pop("t_embed.freqs")
+    torch.testing.assert_close(checkpoint_freqs,
+                               fastvideo.t_embed.freqs.cpu(),
+                               atol=1e-6,
+                               rtol=1e-6)
+    missing, unexpected = fastvideo.load_state_dict(state, strict=True)
+    assert missing == []
+    assert unexpected == []
+    del state
+    official.eval()
+    fastvideo.eval()
+
+    generator = torch.Generator(device=device).manual_seed(1234)
+    latent = torch.randn((1, 345, 40), generator=generator, device=device, dtype=dtype)
+    clip = torch.randn((1, 64, 1024), generator=generator, device=device, dtype=dtype)
+    sync = torch.randn((1, 192, 768), generator=generator, device=device, dtype=dtype)
+    text = torch.randn((1, 77, 1024), generator=generator, device=device, dtype=dtype)
+    timestep = torch.tensor([0.375], device=device, dtype=dtype)
+
+    with torch.inference_mode(), torch.autocast("cuda", dtype=dtype):
+        expected = official(latent, clip, sync, text, timestep)
+        actual = fastvideo(latent, (clip, sync, text), timestep)
+
+    difference = (actual.float() - expected.float()).abs()
+    print("max_abs", difference.max().item())
+    print("mean_abs", difference.mean().item())
+    torch.testing.assert_close(actual, expected, atol=1e-3, rtol=1e-3)

tests/local_tests/mmaudio/test_mmaudio_vocoder_parity.py

@@ -0,0 +1,96 @@
+# SPDX-License-Identifier: Apache-2.0
+"""MMAudio BigVGAN-v2 parity against the canonical NVIDIA checkpoint."""
+
+from __future__ import annotations
+
+import json
+import os
+from pathlib import Path
+
+import pytest
+import torch
+
+
+REPO_ROOT = Path(__file__).resolve().parents[3]
+BIGVGAN_DIR = Path(
+    os.environ.get(
+        "MMAUDIO_BIGVGAN_DIR",
+        REPO_ROOT / "official_weights/mmaudio/bigvgan_v2_44khz_128band_512x",
+    )
+)
+
+
+def _require_assets() -> None:
+    required = (BIGVGAN_DIR / "config.json", BIGVGAN_DIR / "bigvgan_generator.pt")
+    if not all(path.is_file() for path in required):
+        pytest.skip(
+            "Canonical BigVGAN-v2 assets are absent. Set MMAUDIO_BIGVGAN_DIR "
+            "to a local nvidia/bigvgan_v2_44khz_128band_512x snapshot."
+        )
+
+
+def test_bigvgan_v2_implementation_parity() -> None:
+    from mmaudio.ext.bigvgan_v2.bigvgan import BigVGAN
+    from mmaudio.ext.bigvgan_v2.env import AttrDict
+
+    from fastvideo.models.audio.bigvgan import BigVGANV2
+
+    config = {
+        "num_mels": 4,
+        "upsample_initial_channel": 16,
+        "resblock": "1",
+        "resblock_kernel_sizes": [3],
+        "resblock_dilation_sizes": [[1, 3, 5]],
+        "upsample_rates": [2],
+        "upsample_kernel_sizes": [4],
+        "activation": "snakebeta",
+        "snake_logscale": True,
+        "use_bias_at_final": True,
+        "use_tanh_at_final": True,
+    }
+    official = BigVGAN(AttrDict(config), use_cuda_kernel=False)
+    fastvideo = BigVGANV2(config)
+    assert {name: tensor.shape for name, tensor in official.state_dict().items()} == {
+        name: tensor.shape for name, tensor in fastvideo.state_dict().items()
+    }
+    fastvideo.load_state_dict(official.state_dict(), strict=True)
+    official.remove_weight_norm()
+    fastvideo.remove_weight_norm()
+    mel = torch.randn((1, 4, 8), generator=torch.Generator().manual_seed(1234))
+    with torch.inference_mode():
+        expected = official(mel)
+        actual = fastvideo(mel)
+    torch.testing.assert_close(actual, expected, atol=1e-6, rtol=1e-6)
+
+
+def test_mmaudio_bigvgan_v2_parity() -> None:
+    _require_assets()
+    if not torch.cuda.is_available():
+        pytest.skip("MMAudio BigVGAN parity requires CUDA")
+
+    from mmaudio.ext.bigvgan_v2.bigvgan import BigVGAN, load_hparams_from_json
+
+    from fastvideo.models.audio.bigvgan import BigVGANV2
+
+    config_path = BIGVGAN_DIR / "config.json"
+    with config_path.open(encoding="utf-8") as handle:
+        config = json.load(handle)
+    official = BigVGAN(load_hparams_from_json(config_path), use_cuda_kernel=False)
+    fastvideo = BigVGANV2(config)
+    state = torch.load(BIGVGAN_DIR / "bigvgan_generator.pt", map_location="cpu", weights_only=True)["generator"]
+    official.load_state_dict(state, strict=True)
+    fastvideo.load_state_dict(state, strict=True)
+    assert {name: tensor.shape for name, tensor in official.state_dict().items()} == {
+        name: tensor.shape for name, tensor in fastvideo.state_dict().items()
+    }
+
+    device = torch.device("cuda:0")
+    official.remove_weight_norm()
+    fastvideo.remove_weight_norm()
+    official.to(device).eval()
+    fastvideo.to(device).eval()
+    mel = torch.randn((1, 128, 8), generator=torch.Generator(device=device).manual_seed(1234), device=device)
+    with torch.inference_mode():
+        expected = official(mel)
+        actual = fastvideo(mel)
+    torch.testing.assert_close(actual, expected, atol=1e-5, rtol=1e-5)