diff --git a/examples/inference/eval/mmaudio/README.md b/examples/inference/eval/mmaudio/README.md new file mode 100644 index 000000000..1c549c83e --- /dev/null +++ b/examples/inference/eval/mmaudio/README.md @@ -0,0 +1,237 @@ +# MMAudio dataset inference and evaluation + +This example evaluates a trained MMAudio checkpoint without importing the +upstream MMAudio package. Generation uses FastVideo's native +`MMAudioPipeline`; the default metrics use `fastvideo.eval`. + +## Relation to official MMAudio + +Official `batch_eval.py` launches one complete model replica per GPU with +`torchrun`, uses a `DistributedSampler`, runs batched bf16 inference, and +writes audio without composing videos. Metric calculation is a separate step +through `av_bench.extract` and `av_bench.evaluate`. + +This FastVideo path keeps that high-level separation: + +1. one complete MMAudio pipeline replica per GPU; +2. data-parallel dataset inference to WAV; +3. a canonical JSONL manifest that binds id, source video, generated audio, + caption, seed, and reference-audio source; +4. independent multi-GPU metric workers over the complete corpus. + +There are two intentional differences. Rank-strided indices avoid the padded +duplicates that PyTorch `DistributedSampler(drop_last=False)` can introduce, +and cached CLIP/Synchformer/text features skip video decoding and conditioning +encoder work already completed during FastVideo preprocessing. The current +cached-feature stage accepts one sample per pipeline call; four GPUs still +process four samples concurrently. + +## Environment + +Use the same FastVideo environment as training: + +```bash +cd /mnt/lustre/vlm-kai/FastVideo +source .venv/bin/activate +uv pip install --python .venv/bin/python hear21passt pyloudnorm +``` + +The four default native metrics are: + +| FastVideo metric | Protocol | +| --- | --- | +| `audio.frechet_distance` | corpus PaSST FAD / FD_PaSST | +| `audio.kl_divergence` | paired PaSST KL, `KL(gt || pred)` | +| `audio.clap_score` | text-audio CLAP similarity | +| `audio.desync` | source-video/generated-audio Synchformer DeSync | + +FAD and KL port the av-benchmark mathematics. DeSync vendors the same +Synchformer family. FastVideo CLAP uses the Hugging Face LAION CLAP checkpoint +and is not byte-identical to the external package's CLAP implementation; use +the optional official backend below when exact paper-table reproduction is +required. + +## 1. Export the final EMA + +The training result is a transformer state dict. Inference also needs the +frozen audio VAE, vocoder, scheduler, and (for raw-video inference) the frozen +conditioning encoders. Export a complete FastVideo model once: + +```bash +bash examples/inference/eval/mmaudio/export_final_ema.sh +``` + +Default input: + +```text +outputs/mmaudio_small_44k_ddp_from_scratch/posthoc_ema/official_ddp/ + mmaudio_ema_final_sigma_0p05_step_000300000.pth +``` + +Default output: + +```text +converted_weights/mmaudio/small_44k_ema_300000/ +``` + +Override `EMA_CHECKPOINT`, `OUTPUT_MODEL`, or `ASSET_ROOT` before the command +for another run or variant. + +## 2. Multi-GPU inference + +The launcher defaults to the VGGSound test feature cache and four GPUs: + +```bash +bash examples/inference/eval/mmaudio/run_inference_vggsound.sh +``` + +For a four-sample smoke test: + +```bash +MAX_SAMPLES=4 \ +OUTPUT_DIR=outputs/mmaudio_small_44k_ema_300000_vggsound_smoke \ +bash examples/inference/eval/mmaudio/run_inference_vggsound.sh +``` + +For a background full run with a persistent log: + +```bash +mkdir -p logs +nohup bash examples/inference/eval/mmaudio/run_inference_vggsound.sh \ + > logs/mmaudio_vggsound_inference.log 2>&1 & +tail -f logs/mmaudio_vggsound_inference.log +``` + +The output directory contains: + +```text +audio/.wav +eval_manifest.jsonl +failures.jsonl +summary.json +manifest_rank_*.jsonl +failures_rank_*.jsonl +``` + +Existing WAV files are resumed rather than regenerated. Add `--overwrite` to +the Python runner when intentional regeneration is required. Seeds are +`base_seed + global_index`, so results do not change with GPU count. + +## 3. Native FastVideo evaluation + +Run FAD, KL, CLAP, and DeSync with four metric replicas: + +```bash +bash examples/inference/eval/mmaudio/run_evaluate_vggsound.sh +``` + +The first run extracts ground-truth audio from each source MP4 into +`/reference_audio` and caches metric weights. Later runs reuse +both. Results are written to `fastvideo_eval_results.json`; per-sample metrics +are under `samples`, while FAD is under `corpus`. + +The underlying generic command is: + +```bash +fastvideo eval run \ + --manifest outputs/mmaudio_small_44k_ema_300000_vggsound_test/eval_manifest.jsonl \ + --metrics audio.frechet_distance,audio.kl_divergence,audio.clap_score,audio.desync \ + --num-gpus 4 \ + --extract-audio outputs/mmaudio_small_44k_ema_300000_vggsound_test/reference_audio \ + --extract-workers 16 \ + --output-format full \ + --output outputs/mmaudio_small_44k_ema_300000_vggsound_test/fastvideo_eval_results.json +``` + +Do not evaluate FAD by invoking the evaluator once per file. FAD is a +set-vs-set statistic; this CLI submits the complete manifest in one call and +serializes `EvalResults.corpus`. + +## Exact external av-benchmark backend + +For strict comparison to official MMAudio numbers, FastVideo provides the +dedicated `fastvideo eval v2a` command. It launches the exact +`av_bench.extract` and `av_bench.evaluate` functions with a Python interpreter +from an isolated environment. This keeps av-benchmark's CLAP, ImageBind, and +PyTorch dependencies out of the main FastVideo environment. + +Create the isolated environment once: + +```bash +cd /mnt/lustre/vlm-kai +git clone https://github.com/hkchengrex/av-benchmark.git +uv venv av-benchmark/.venv --python 3.12 +uv pip install --python av-benchmark/.venv/bin/python -e av-benchmark +# PyTorch 2.9+ routes torchaudio decoding through TorchCodec. Transformers 5 +# removed an API used by the vendored Synchformer AST, while SciPy 1.17 removed +# an argument used by the official FAD implementation. Keep these constraints +# in this isolated environment; select the wheel matching its PyTorch CUDA build. +UV_TORCH_BACKEND=cu130 uv pip install \ + --python av-benchmark/.venv/bin/python \ + torchcodec==0.15.0 'transformers<5' 'scipy<1.17' + +mkdir -p av-benchmark/weights +curl -L https://huggingface.co/lukewys/laion_clap/resolve/main/music_speech_audioset_epoch_15_esc_89.98.pt \ + -o av-benchmark/weights/music_speech_audioset_epoch_15_esc_89.98.pt +curl -L https://github.com/hkchengrex/MMAudio/releases/download/v0.1/synchformer_state_dict.pth \ + -o av-benchmark/weights/synchformer_state_dict.pth +``` + +Download only the official VGGSound ground-truth cache (about 16 GB): + +```bash +cd /mnt/lustre/vlm-kai/FastVideo +.venv/bin/hf download hkchengrex/MMAudio-precomputed-results \ + --repo-type dataset \ + --include "vggsound-test-eval-cache/*" \ + --local-dir /mnt/lustre/vlm-kai/datasets/VGGSound/av_benchmark +``` + +Then run the dedicated launcher: + +```bash +bash examples/inference/eval/mmaudio/run_evaluate_vggsound_av_benchmark.sh +``` + +For a background run: + +```bash +nohup bash examples/inference/eval/mmaudio/run_evaluate_vggsound_av_benchmark.sh \ + > logs/mmaudio_vggsound_av_benchmark.log 2>&1 & +tail -f logs/mmaudio_vggsound_av_benchmark.log +``` + +The official VGGSound cache produces the complete paper-style metric set: + +```text +FD-VGG, FD-PANN, FD-PASST +KL-PANNS-softmax, KL-PASST-softmax +ISC-PANNS-mean/std, ISC-PASST-mean/std +IB-Score, DeSync +``` + +The official VGGSound cache does not contain CLAP text features, so the launcher +defaults to `SKIP_CLAP=1`; this saves prediction feature passes without changing +the metric set. Set `SKIP_CLAP=0` for a custom GT cache that includes CLAP text +features. Prediction features are cached and reused; set `RECOMPUTE=1` to force +a clean extraction. + +The launcher also defaults to `ALIGN_PREDICTION_KEYS=1`. The official +VGGSound cache prepends underscores to some sanitized ids; FastVideo keeps the +original ids. The backend changes only unique matches in prediction feature +caches, leaves WAV filenames untouched, and records exact/remapped/ambiguous/ +unmatched counts under `prediction_key_alignment` in the result JSON. + +The launcher defaults to `NUM_WORKERS=0`: this avoids both the machine's 64 MB +`/dev/shm` limit and forking workers after the official extractor initializes +CUDA. Raise it only after validating the container and backend combination. + +## Evaluation protocol warning + +The current FastVideo test cache contains the filtered natural-language +descriptions used in this training run. Official MMAudio's VGGSound +`batch_eval.py` reads class-label captions from the original VGGSound CSV. +FAD, KL, and DeSync formulas do not read captions directly, but changing the +generation caption changes the generated waveform and can therefore change all +scores. CLAP additionally consumes the caption during evaluation. Record the +generation manifest, test split, and ground-truth cache with reported results. diff --git a/examples/inference/eval/mmaudio/eval_mmaudio_dataset.py b/examples/inference/eval/mmaudio/eval_mmaudio_dataset.py new file mode 100644 index 000000000..f4ccd830e --- /dev/null +++ b/examples/inference/eval/mmaudio/eval_mmaudio_dataset.py @@ -0,0 +1,304 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Distributed MMAudio inference over FastVideo preprocessing caches. + +Launch with ``torchrun``. Each process owns one complete MMAudio pipeline and +processes a non-padded, rank-strided subset of the dataset. The output manifest +is directly consumable by ``fastvideo eval run --manifest ...``. +""" + +from __future__ import annotations + +import argparse +import json +import os +import re +from pathlib import Path +from typing import Any + +import numpy as np +import torch +from scipy.io import wavfile +from torch.utils.data import DataLoader, Dataset, Sampler +from tqdm import tqdm + +from fastvideo.dataset.mmaudio_feature_dataset import build_mmaudio_feature_dataset +from fastvideo.distributed import cleanup_dist_env_and_memory +from fastvideo.pipelines.basic.mmaudio import MMAudioPipeline +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch + + +class _RankStrideSampler(Sampler[int]): + """Assign every global index exactly once without tail padding.""" + + def __init__(self, size: int, rank: int, world_size: int) -> None: + self.indices = range(rank, size, world_size) + + def __iter__(self): + return iter(self.indices) + + def __len__(self) -> int: + return len(self.indices) + + +class _IndexedDataset(Dataset): + + def __init__(self, dataset: Dataset) -> None: + self.dataset = dataset + + def __len__(self) -> int: + return len(self.dataset) + + def __getitem__(self, index: int) -> tuple[int, dict[str, Any]]: + return index, self.dataset[index] + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model-path", type=Path, required=True) + parser.add_argument("--feature-root", type=Path, required=True) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--duration-seconds", type=float, default=8.0) + parser.add_argument("--num-inference-steps", type=int, default=25) + parser.add_argument("--guidance-scale", type=float, default=4.5) + parser.add_argument("--seed", type=int, default=14159265) + parser.add_argument("--negative-prompt", default="") + parser.add_argument("--num-workers", type=int, default=2) + parser.add_argument("--max-samples", type=int, default=0) + parser.add_argument("--overwrite", action="store_true") + parser.add_argument("--compile", action="store_true") + return parser.parse_args() + + +def _safe_name(value: str) -> str: + value = re.sub(r"[^A-Za-z0-9._-]+", "_", value).strip("._") + return value or "sample" + + +def _feature_shapes(transformer: torch.nn.Module) -> dict[str, int]: + arch = transformer.config.arch_config + return { + "latent_seq_len": int(transformer.latent_seq_len), + "latent_dim": int(transformer.latent_dim), + "clip_seq_len": int(transformer.clip_seq_len), + "clip_dim": int(arch.clip_dim), + "sync_seq_len": int(transformer.sync_seq_len), + "sync_dim": int(arch.sync_dim), + "text_seq_len": int(arch.text_seq_len), + "text_dim": int(arch.text_dim), + } + + +class _IgnoredComponent: + """Placeholder for a required module the direct-feature path never touches. + + ``MMAudioPipeline`` requires the CLIP text/vision and Synchformer encoders, + but its direct-feature conditioning path returns before dereferencing them. + Callers that pass ``video_path`` (or omit the cached features) must supply + the real modules instead. + """ + + def __repr__(self) -> str: + return "" + + +def _build_pipeline(args: argparse.Namespace, world_size: int) -> MMAudioPipeline: + # Cached features make the three conditioning encoders unnecessary. The + # stages still exist, but consume direct tensors and never dereference + # these sentinels. + ignored_component = _IgnoredComponent() + return MMAudioPipeline.from_pretrained( + str(args.model_path), + inference_mode=True, + loaded_modules={ + "text_encoder": ignored_component, + "tokenizer": ignored_component, + "image_encoder": ignored_component, + "image_encoder_2": ignored_component, + }, + workload_type="v2a", + num_gpus=world_size, + tp_size=1, + sp_size=1, + hsdp_replicate_dim=1, + hsdp_shard_dim=1, + dit_cpu_offload=False, + dit_layerwise_offload=False, + text_encoder_cpu_offload=False, + image_encoder_cpu_offload=False, + vae_cpu_offload=False, + enable_torch_compile=args.compile, + ) + + +def _forward_batch( + sample: dict[str, Any], + *, + duration_seconds: float, + num_inference_steps: int, + guidance_scale: float, + seed: int, + negative_prompt: str, +) -> ForwardBatch: + return ForwardBatch( + data_type="video", + prompt="", + negative_prompt=negative_prompt, + audio_start_in_s=0.0, + audio_end_in_s=duration_seconds, + num_inference_steps=num_inference_steps, + guidance_scale=guidance_scale, + seed=seed, + num_videos_per_prompt=1, + height=8, + width=8, + num_frames=1, + save_video=False, + return_frames=False, + extra={ + "mmaudio_clip_features": sample["clip_features"].unsqueeze(0), + "mmaudio_sync_features": sample["sync_features"].unsqueeze(0), + "mmaudio_text_features": sample["text_features"].unsqueeze(0), + }, + ) + + +def _write_jsonl(path: Path, rows: list[dict[str, Any]]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + temporary = path.with_suffix(path.suffix + ".tmp") + with temporary.open("w", encoding="utf-8") as handle: + for row in rows: + handle.write(json.dumps(row, sort_keys=True) + "\n") + temporary.replace(path) + + +def _merge_rank_outputs(output_dir: Path, world_size: int) -> None: + rows: list[dict[str, Any]] = [] + failures: list[dict[str, Any]] = [] + for rank in range(world_size): + manifest = output_dir / f"manifest_rank_{rank:05d}.jsonl" + failure_path = output_dir / f"failures_rank_{rank:05d}.jsonl" + if manifest.is_file(): + rows.extend(json.loads(line) for line in manifest.read_text(encoding="utf-8").splitlines() if line.strip()) + if failure_path.is_file(): + failures.extend( + json.loads(line) for line in failure_path.read_text(encoding="utf-8").splitlines() if line.strip()) + rows.sort(key=lambda row: int(row["index"])) + failures.sort(key=lambda row: int(row["index"])) + ids = [str(row["id"]) for row in rows] + if len(ids) != len(set(ids)): + raise RuntimeError("Duplicate sample ids found while merging inference manifests") + _write_jsonl(output_dir / "eval_manifest.jsonl", rows) + _write_jsonl(output_dir / "failures.jsonl", failures) + summary = { + "num_succeeded": len(rows), + "num_failed": len(failures), + "world_size": world_size, + } + (output_dir / "summary.json").write_text( + json.dumps(summary, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + + +@torch.inference_mode() +def main() -> None: + args = parse_args() + rank = int(os.environ.get("RANK", "0")) + local_rank = int(os.environ.get("LOCAL_RANK", "0")) + world_size = int(os.environ.get("WORLD_SIZE", "1")) + if not torch.cuda.is_available(): + raise RuntimeError("MMAudio dataset inference requires CUDA") + torch.cuda.set_device(local_rank) + + output_dir = args.output_dir.expanduser().resolve() + audio_dir = output_dir / "audio" + audio_dir.mkdir(parents=True, exist_ok=True) + pipeline = _build_pipeline(args, world_size) + transformer = pipeline.get_module("transformer") + dataset = build_mmaudio_feature_dataset( + args.feature_root, + feature_shapes=_feature_shapes(transformer), + include_metadata=True, + ) + dataset_size = len(dataset) + if args.max_samples > 0: + dataset_size = min(dataset_size, args.max_samples) + sampler = _RankStrideSampler(dataset_size, rank, world_size) + loader = DataLoader( + _IndexedDataset(dataset), + batch_size=None, + sampler=sampler, + num_workers=max(0, args.num_workers), + pin_memory=True, + persistent_workers=args.num_workers > 0, + ) + + manifest_rows: list[dict[str, Any]] = [] + failure_rows: list[dict[str, Any]] = [] + progress = tqdm( + loader, + total=len(sampler), + desc=f"MMAudio inference rank {rank}", + position=local_rank, + ) + for index, sample in progress: + index = int(index) + sample_id = _safe_name(str(sample.get("sample_id", index))) + source_path = str(sample.get("source_path", "")) + caption = str(sample.get("caption", "")) + output_path = audio_dir / f"{sample_id}.wav" + sample_seed = args.seed + index + try: + if not source_path or not Path(source_path).is_file(): + raise FileNotFoundError(f"source video is missing for sample {sample_id}: {source_path}") + if args.overwrite or not output_path.is_file(): + batch = _forward_batch( + sample, + duration_seconds=args.duration_seconds, + num_inference_steps=args.num_inference_steps, + guidance_scale=args.guidance_scale, + seed=sample_seed, + negative_prompt=args.negative_prompt, + ) + result = pipeline.forward(batch, pipeline.fastvideo_args) + audio = result.extra.get("audio") + sample_rate = result.extra.get("audio_sample_rate") + if not isinstance(audio, np.ndarray) or not isinstance(sample_rate, int): + raise RuntimeError("MMAudio pipeline did not return decoded audio") + temporary = output_path.with_suffix(".tmp.wav") + wavfile.write(temporary, sample_rate, audio.astype(np.float32)) + temporary.replace(output_path) + manifest_rows.append({ + "id": sample_id, + "index": index, + "video": str(Path(source_path).resolve()), + "audio": str(output_path.resolve()), + "reference_audio_source": str(Path(source_path).resolve()), + "text_prompt": caption, + "seed": sample_seed, + }) + except Exception as error: # noqa: BLE001 - isolate corrupt benchmark samples + failure_rows.append({ + "id": sample_id, + "index": index, + "source": source_path, + "error_type": type(error).__name__, + "error": str(error), + }) + + _write_jsonl(output_dir / f"manifest_rank_{rank:05d}.jsonl", manifest_rows) + _write_jsonl(output_dir / f"failures_rank_{rank:05d}.jsonl", failure_rows) + if torch.distributed.is_initialized(): + torch.distributed.barrier() + if rank == 0: + _merge_rank_outputs(output_dir, world_size) + print(f"Merged eval manifest: {output_dir / 'eval_manifest.jsonl'}") + print(f"Inference summary: {output_dir / 'summary.json'}") + if torch.distributed.is_initialized(): + torch.distributed.barrier() + pipeline.close() + cleanup_dist_env_and_memory() + + +if __name__ == "__main__": + main() diff --git a/examples/inference/eval/mmaudio/evaluate_official_av_benchmark.py b/examples/inference/eval/mmaudio/evaluate_official_av_benchmark.py new file mode 100644 index 000000000..e3296247a --- /dev/null +++ b/examples/inference/eval/mmaudio/evaluate_official_av_benchmark.py @@ -0,0 +1,60 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Run the exact external av-benchmark API used by official MMAudio.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--audio-dir", type=Path, required=True) + parser.add_argument("--gt-cache", type=Path, required=True) + parser.add_argument("--prediction-cache", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--device", default="cuda") + parser.add_argument("--batch-size", type=int, default=32) + parser.add_argument("--audio-length", type=float, default=8.0) + return parser.parse_args() + + +def main() -> None: + args = parse_args() + try: + from av_bench.evaluate import evaluate + from av_bench.extract import extract + except ImportError as error: + raise ImportError("The exact official backend is optional. Install hkchengrex/" + "av-benchmark in a separate environment before running this file.") from error + + args.prediction_cache.mkdir(parents=True, exist_ok=True) + extract( + audio_path=args.audio_dir, + output_path=args.prediction_cache, + device=args.device, + batch_size=args.batch_size, + audio_length=args.audio_length, + ) + metrics = evaluate( + gt_audio_cache=args.gt_cache, + pred_audio_cache=args.prediction_cache, + ) + payload = { + "backend": "hkchengrex/av-benchmark", + "audio_dir": str(args.audio_dir.resolve()), + "gt_cache": str(args.gt_cache.resolve()), + "prediction_cache": str(args.prediction_cache.resolve()), + "metrics": metrics, + } + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text( + json.dumps(payload, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + print(json.dumps(payload, indent=2, sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/examples/inference/eval/mmaudio/export_final_ema.sh b/examples/inference/eval/mmaudio/export_final_ema.sh new file mode 100755 index 000000000..5b9b278ec --- /dev/null +++ b/examples/inference/eval/mmaudio/export_final_ema.sh @@ -0,0 +1,28 @@ +#!/usr/bin/env bash +set -euo pipefail + +# One-time export of the synthesized PostHocEMA transformer plus the frozen +# official conditioning/decoder components into a complete FastVideo model. +ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/../../../.." && pwd)" +cd "${ROOT_DIR}" + +EMA_CHECKPOINT="${EMA_CHECKPOINT:-outputs/mmaudio_small_44k_ddp_from_scratch/posthoc_ema/official_ddp/mmaudio_ema_final_sigma_0p05_step_000300000.pth}" +OUTPUT_MODEL="${OUTPUT_MODEL:-converted_weights/mmaudio/small_44k_ema_300000}" +ASSET_ROOT="${ASSET_ROOT:-official_weights/mmaudio}" + +if [[ -e "${OUTPUT_MODEL}/model_index.json" ]]; then + echo "Model already exported: ${OUTPUT_MODEL}" + exit 0 +fi + +source .venv/bin/activate +python scripts/checkpoint_conversion/convert_mmaudio_to_diffusers.py \ + --variant small_44k \ + --transformer-checkpoint "${EMA_CHECKPOINT}" \ + --audio-vae-checkpoint "${ASSET_ROOT}/raw/ext_weights/v1-44.pth" \ + --synchformer-checkpoint "${ASSET_ROOT}/raw/ext_weights/synchformer_state_dict.pth" \ + --dfn5b-dir "${ASSET_ROOT}/DFN5B-CLIP-ViT-H-14-384" \ + --bigvgan-dir "${ASSET_ROOT}/bigvgan_v2_44khz_128band_512x" \ + --output "${OUTPUT_MODEL}" + +echo "Exported FastVideo MMAudio model: ${OUTPUT_MODEL}" diff --git a/examples/inference/eval/mmaudio/run_evaluate_vggsound.sh b/examples/inference/eval/mmaudio/run_evaluate_vggsound.sh new file mode 100755 index 000000000..b6157c728 --- /dev/null +++ b/examples/inference/eval/mmaudio/run_evaluate_vggsound.sh @@ -0,0 +1,25 @@ +#!/usr/bin/env bash +set -euo pipefail + +ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/../../../.." && pwd)" +cd "${ROOT_DIR}" +source .venv/bin/activate + +OUTPUT_DIR="${OUTPUT_DIR:-outputs/mmaudio_small_44k_ema_300000_vggsound_test}" +MANIFEST="${MANIFEST:-${OUTPUT_DIR}/eval_manifest.jsonl}" +REFERENCE_AUDIO_CACHE="${REFERENCE_AUDIO_CACHE:-${OUTPUT_DIR}/reference_audio}" +RESULTS="${RESULTS:-${OUTPUT_DIR}/fastvideo_eval_results.json}" +NUM_GPUS="${NUM_GPUS:-4}" +EXTRACT_WORKERS="${EXTRACT_WORKERS:-16}" +METRICS="${METRICS:-audio.frechet_distance,audio.kl_divergence,audio.clap_score,audio.desync}" + +fastvideo eval run \ + --manifest "${MANIFEST}" \ + --metrics "${METRICS}" \ + --num-gpus "${NUM_GPUS}" \ + --extract-audio "${REFERENCE_AUDIO_CACHE}" \ + --extract-workers "${EXTRACT_WORKERS}" \ + --output-format full \ + --output "${RESULTS}" + +echo "FastVideo evaluation results: ${RESULTS}" diff --git a/examples/inference/eval/mmaudio/run_evaluate_vggsound_av_benchmark.sh b/examples/inference/eval/mmaudio/run_evaluate_vggsound_av_benchmark.sh new file mode 100755 index 000000000..1fd46e6e1 --- /dev/null +++ b/examples/inference/eval/mmaudio/run_evaluate_vggsound_av_benchmark.sh @@ -0,0 +1,65 @@ +#!/usr/bin/env bash +set -euo pipefail + +ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/../../../.." && pwd)" +cd "${ROOT_DIR}" + +OUTPUT_DIR="${OUTPUT_DIR:-outputs/mmaudio_small_44k_ema_300000_vggsound_test}" +AUDIO_DIR="${AUDIO_DIR:-${OUTPUT_DIR}/audio}" +GT_CACHE="${GT_CACHE:-/mnt/lustre/vlm-kai/datasets/VGGSound/av_benchmark/vggsound-test-eval-cache}" +PREDICTION_CACHE="${PREDICTION_CACHE:-${OUTPUT_DIR}/av_benchmark_cache}" +RESULTS="${RESULTS:-${OUTPUT_DIR}/av_benchmark_results.json}" +AV_BENCH_PYTHON="${AV_BENCH_PYTHON:-/mnt/lustre/vlm-kai/av-benchmark/.venv/bin/python}" +BATCH_SIZE="${BATCH_SIZE:-32}" +# The official extractor initializes CUDA models before constructing its +# DataLoader. Zero avoids CUDA-after-fork hangs and this container's 64 MB shm. +NUM_WORKERS="${NUM_WORKERS:-0}" +DEVICE="${DEVICE:-cuda}" +RECOMPUTE="${RECOMPUTE:-0}" +SKIP_VIDEO_RELATED="${SKIP_VIDEO_RELATED:-0}" +# The official VGGSound GT cache has no CLAP text features, so CLAP is not an +# official VGGSound output metric. Skip its prediction passes by default. +SKIP_CLAP="${SKIP_CLAP:-1}" +# Official VGGSound precomputed cache sanitizes a subset of sample ids by +# prefixing underscores. Align only unique matches inside prediction caches. +ALIGN_PREDICTION_KEYS="${ALIGN_PREDICTION_KEYS:-1}" + +if [[ ! -x "${AV_BENCH_PYTHON}" ]]; then + echo "Missing isolated av-benchmark Python: ${AV_BENCH_PYTHON}" >&2 + echo "Set AV_BENCH_PYTHON=/path/to/av-benchmark/.venv/bin/python" >&2 + exit 2 +fi +if [[ ! -d "${GT_CACHE}" ]]; then + echo "Missing official VGGSound GT cache: ${GT_CACHE}" >&2 + echo "Set GT_CACHE=/path/to/vggsound-test-eval-cache" >&2 + exit 2 +fi + +EXTRA_ARGS=() +if [[ "${RECOMPUTE}" == "1" ]]; then + EXTRA_ARGS+=(--recompute) +fi +if [[ "${SKIP_VIDEO_RELATED}" == "1" ]]; then + EXTRA_ARGS+=(--skip-video-related) +fi +if [[ "${SKIP_CLAP}" == "1" ]]; then + EXTRA_ARGS+=(--skip-clap) +fi +if [[ "${ALIGN_PREDICTION_KEYS}" == "1" ]]; then + EXTRA_ARGS+=(--align-prediction-keys) +fi + +.venv/bin/fastvideo eval v2a \ + --backend av-benchmark \ + --audio-dir "${AUDIO_DIR}" \ + --gt-cache "${GT_CACHE}" \ + --prediction-cache "${PREDICTION_CACHE}" \ + --output "${RESULTS}" \ + --python-executable "${AV_BENCH_PYTHON}" \ + --device "${DEVICE}" \ + --batch-size "${BATCH_SIZE}" \ + --num-workers "${NUM_WORKERS}" \ + --audio-length 8 \ + "${EXTRA_ARGS[@]}" + +echo "Official av-benchmark results: ${RESULTS}" diff --git a/examples/inference/eval/mmaudio/run_inference_vggsound.sh b/examples/inference/eval/mmaudio/run_inference_vggsound.sh new file mode 100755 index 000000000..e8024b45e --- /dev/null +++ b/examples/inference/eval/mmaudio/run_inference_vggsound.sh @@ -0,0 +1,42 @@ +#!/usr/bin/env bash +set -euo pipefail + +ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/../../../.." && pwd)" +cd "${ROOT_DIR}" +source .venv/bin/activate + +MODEL_PATH="${MODEL_PATH:-converted_weights/mmaudio/small_44k_ema_300000}" +FEATURE_ROOT="${FEATURE_ROOT:-/mnt/lustre/vlm-kai/datasets/VGGSound/mmaudio_features_torio/test}" +OUTPUT_DIR="${OUTPUT_DIR:-outputs/mmaudio_small_44k_ema_300000_vggsound_test}" +NUM_GPUS="${NUM_GPUS:-4}" +NUM_WORKERS="${NUM_WORKERS:-2}" +MAX_SAMPLES="${MAX_SAMPLES:-0}" +COMPILE="${COMPILE:-1}" +MASTER_PORT="${MASTER_PORT:-29513}" +export OMP_NUM_THREADS="${OMP_NUM_THREADS:-4}" + +EXTRA_ARGS=() +if [[ "${MAX_SAMPLES}" -gt 0 ]]; then + EXTRA_ARGS+=(--max-samples "${MAX_SAMPLES}") +fi +if [[ "${COMPILE}" == "1" ]]; then + EXTRA_ARGS+=(--compile) +fi + +torchrun \ + --standalone \ + --nproc_per_node="${NUM_GPUS}" \ + --master_port="${MASTER_PORT}" \ + examples/inference/eval/mmaudio/eval_mmaudio_dataset.py \ + --model-path "${MODEL_PATH}" \ + --feature-root "${FEATURE_ROOT}" \ + --output-dir "${OUTPUT_DIR}" \ + --duration-seconds 8 \ + --num-inference-steps 25 \ + --guidance-scale 4.5 \ + --seed 14159265 \ + --num-workers "${NUM_WORKERS}" \ + "${EXTRA_ARGS[@]}" + +echo "Generated audio: ${OUTPUT_DIR}/audio" +echo "Evaluation manifest: ${OUTPUT_DIR}/eval_manifest.jsonl" diff --git a/examples/train/configs/example.yaml b/examples/train/configs/example.yaml index 824fb0987..d79cd317b 100644 --- a/examples/train/configs/example.yaml +++ b/examples/train/configs/example.yaml @@ -105,6 +105,7 @@ training: optimizer: learning_rate: 2.0e-6 # default: 0.0 betas: [0.9, 0.999] # default: [0.9, 0.999] + eps: 1.0e-8 # default: 1.0e-8 weight_decay: 0.01 # default: 0.0 lr_scheduler: constant # default: "constant" lr_warmup_steps: 0 # default: 0 diff --git a/examples/train/configs/fine_tuning/mmaudio/large_44k_from_scratch.yaml b/examples/train/configs/fine_tuning/mmaudio/large_44k_from_scratch.yaml new file mode 100644 index 000000000..d5fe66ced --- /dev/null +++ b/examples/train/configs/fine_tuning/mmaudio/large_44k_from_scratch.yaml @@ -0,0 +1,100 @@ +# Official MMAudio large_44k transformer training from random initialization. +# Requires 44.1 kHz TensorDict features with audio latents [345, 40]. + +models: + student: + _target_: fastvideo.train.models.mmaudio.MMAudioModel + variant: large_44k + from_scratch: true + # Official standalone empty_string.pth. No DiT weights are loaded. + empty_string_features_path: "" + # Optional .pt cache. When empty, stats are recomputed at startup. + latent_statistics_path: "" + latent_statistics_chunk_size: 32 + trainable: true + +method: + _target_: fastvideo.train.methods.fine_tuning.flow_matching.FlowMatchingFineTuneMethod + +training: + # Matching component-tree template used only by DCP-to-Diffusers export. + model_path: converted_weights/mmaudio/large_44k + distributed: + num_gpus: 4 + sp_size: 1 + tp_size: 1 + hsdp_replicate_dim: 1 + hsdp_shard_dim: 4 + pin_cpu_memory: false + + data: + # Set with --training.data.data_path or run_train_vggsound_from_scratch.sh. + data_path: "" + preprocessed_data_type: mmaudio_features + dataloader_num_workers: 2 + train_batch_size: 8 + training_cfg_rate: 0.1 + seed: 14159265 + + optimizer: + learning_rate: 1.0e-4 + betas: [0.9, 0.95] + eps: 1.0e-6 + weight_decay: 1.0e-6 + lr_scheduler: multistep_with_warmup + lr_warmup_steps: 1000 + lr_milestones: [240000, 270000] + lr_gamma: 0.1 + + loop: + max_train_steps: 300000 + # 4 GPUs x 8 samples/GPU x 16 rounds = effective global batch 512. + gradient_accumulation_steps: 16 + + checkpoint: + output_dir: outputs/mmaudio_large_44k_from_scratch + training_state_checkpointing_steps: 10000 + checkpoints_total_limit: 3 + resume_from_checkpoint: "" + + tracker: + project_name: fastvideo_mmaudio + run_name: mmaudio_large_44k_from_scratch + + model: + logit_mean: 0.0 + logit_std: 1.0 + precondition_outputs: false + + # FP32 master weights; MMAudioModel runs transformer forward under BF16 autocast. + dit_precision: fp32 + +callbacks: + grad_clip: + _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback + max_grad_norm: 1.0 + posthoc_ema: + _target_: fastvideo.train.callbacks.posthoc_ema.PostHocEMACallback + sigma_rels: [0.05, 0.1] + update_every: 1 + checkpoint_every: 5000 + start_iter: 0 + default_output_sigma: 0.05 + step_size_correction: true + validation: + _target_: fastvideo.train.callbacks.mmaudio_validation.MMAudioValidationCallback + data_path: "" + every_steps: 5000 + max_batches: 0 + batch_size: 8 + num_data_workers: 2 + run_at_start: false + use_ema: false + inference_every_steps: 20000 + inference_model_path: converted_weights/mmaudio/large_44k_v2 + # Global total across all ranks; fixed by validation seed and sampler. + inference_num_samples: 16 + inference_num_steps: 25 + inference_guidance_scale: 4.5 + inference_save_video: true + inference_log_to_tracker: true diff --git a/examples/train/configs/fine_tuning/mmaudio/medium_44k_from_scratch.yaml b/examples/train/configs/fine_tuning/mmaudio/medium_44k_from_scratch.yaml new file mode 100644 index 000000000..65ddf740f --- /dev/null +++ b/examples/train/configs/fine_tuning/mmaudio/medium_44k_from_scratch.yaml @@ -0,0 +1,100 @@ +# Official MMAudio medium_44k transformer training from random initialization. +# Requires 44.1 kHz TensorDict features with audio latents [345, 40]. + +models: + student: + _target_: fastvideo.train.models.mmaudio.MMAudioModel + variant: medium_44k + from_scratch: true + # Official standalone empty_string.pth. No DiT weights are loaded. + empty_string_features_path: "" + # Optional .pt cache. When empty, stats are recomputed at startup. + latent_statistics_path: "" + latent_statistics_chunk_size: 32 + trainable: true + +method: + _target_: fastvideo.train.methods.fine_tuning.flow_matching.FlowMatchingFineTuneMethod + +training: + # Matching component-tree template used only by DCP-to-Diffusers export. + model_path: converted_weights/mmaudio/medium_44k + distributed: + num_gpus: 4 + sp_size: 1 + tp_size: 1 + hsdp_replicate_dim: 1 + hsdp_shard_dim: 4 + pin_cpu_memory: false + + data: + # Set with --training.data.data_path or run_train_vggsound_from_scratch.sh. + data_path: "" + preprocessed_data_type: mmaudio_features + dataloader_num_workers: 2 + train_batch_size: 8 + training_cfg_rate: 0.1 + seed: 14159265 + + optimizer: + learning_rate: 1.0e-4 + betas: [0.9, 0.95] + eps: 1.0e-6 + weight_decay: 1.0e-6 + lr_scheduler: multistep_with_warmup + lr_warmup_steps: 1000 + lr_milestones: [240000, 270000] + lr_gamma: 0.1 + + loop: + max_train_steps: 300000 + # 4 GPUs x 8 samples/GPU x 16 rounds = effective global batch 512. + gradient_accumulation_steps: 16 + + checkpoint: + output_dir: outputs/mmaudio_medium_44k_from_scratch + training_state_checkpointing_steps: 10000 + checkpoints_total_limit: 3 + resume_from_checkpoint: "" + + tracker: + project_name: fastvideo_mmaudio + run_name: mmaudio_medium_44k_from_scratch + + model: + logit_mean: 0.0 + logit_std: 1.0 + precondition_outputs: false + + # FP32 master weights; MMAudioModel runs transformer forward under BF16 autocast. + dit_precision: fp32 + +callbacks: + grad_clip: + _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback + max_grad_norm: 1.0 + posthoc_ema: + _target_: fastvideo.train.callbacks.posthoc_ema.PostHocEMACallback + sigma_rels: [0.05, 0.1] + update_every: 1 + checkpoint_every: 5000 + start_iter: 0 + default_output_sigma: 0.05 + step_size_correction: true + validation: + _target_: fastvideo.train.callbacks.mmaudio_validation.MMAudioValidationCallback + data_path: "" + every_steps: 5000 + max_batches: 0 + batch_size: 8 + num_data_workers: 2 + run_at_start: false + use_ema: false + inference_every_steps: 20000 + inference_model_path: converted_weights/mmaudio/large_44k_v2 + # Global total across all ranks; fixed by validation seed and sampler. + inference_num_samples: 16 + inference_num_steps: 25 + inference_guidance_scale: 4.5 + inference_save_video: true + inference_log_to_tracker: true diff --git a/examples/train/configs/fine_tuning/mmaudio/small_16k_from_scratch.yaml b/examples/train/configs/fine_tuning/mmaudio/small_16k_from_scratch.yaml new file mode 100644 index 000000000..e0bb4e34a --- /dev/null +++ b/examples/train/configs/fine_tuning/mmaudio/small_16k_from_scratch.yaml @@ -0,0 +1,100 @@ +# Official MMAudio small_16k transformer training from random initialization. +# Requires 16 kHz TensorDict features with audio latents [250, 20]. + +models: + student: + _target_: fastvideo.train.models.mmaudio.MMAudioModel + variant: small_16k + from_scratch: true + # Official standalone empty_string.pth. No DiT weights are loaded. + empty_string_features_path: "" + # Optional .pt cache. When empty, stats are recomputed at startup. + latent_statistics_path: "" + latent_statistics_chunk_size: 32 + trainable: true + +method: + _target_: fastvideo.train.methods.fine_tuning.flow_matching.FlowMatchingFineTuneMethod + +training: + # Matching component-tree template used only by DCP-to-Diffusers export. + model_path: converted_weights/mmaudio/small_16k + distributed: + num_gpus: 4 + sp_size: 1 + tp_size: 1 + hsdp_replicate_dim: 1 + hsdp_shard_dim: 4 + pin_cpu_memory: false + + data: + # Set with --training.data.data_path or run_train_vggsound_from_scratch.sh. + data_path: "" + preprocessed_data_type: mmaudio_features + dataloader_num_workers: 2 + train_batch_size: 8 + training_cfg_rate: 0.1 + seed: 14159265 + + optimizer: + learning_rate: 1.0e-4 + betas: [0.9, 0.95] + eps: 1.0e-6 + weight_decay: 1.0e-6 + lr_scheduler: multistep_with_warmup + lr_warmup_steps: 1000 + lr_milestones: [240000, 270000] + lr_gamma: 0.1 + + loop: + max_train_steps: 300000 + # 4 GPUs x 8 samples/GPU x 16 rounds = effective global batch 512. + gradient_accumulation_steps: 16 + + checkpoint: + output_dir: outputs/mmaudio_small_16k_from_scratch + training_state_checkpointing_steps: 10000 + checkpoints_total_limit: 3 + resume_from_checkpoint: "" + + tracker: + project_name: fastvideo_mmaudio + run_name: mmaudio_small_16k_from_scratch + + model: + logit_mean: 0.0 + logit_std: 1.0 + precondition_outputs: false + + # FP32 master weights; MMAudioModel runs transformer forward under BF16 autocast. + dit_precision: fp32 + +callbacks: + grad_clip: + _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback + max_grad_norm: 1.0 + posthoc_ema: + _target_: fastvideo.train.callbacks.posthoc_ema.PostHocEMACallback + sigma_rels: [0.05, 0.1] + update_every: 1 + checkpoint_every: 5000 + start_iter: 0 + default_output_sigma: 0.05 + step_size_correction: true + validation: + _target_: fastvideo.train.callbacks.mmaudio_validation.MMAudioValidationCallback + data_path: "" + every_steps: 5000 + max_batches: 0 + batch_size: 8 + num_data_workers: 2 + run_at_start: false + use_ema: false + # Set this after preparing a complete 16 kHz inference component tree. + inference_every_steps: 0 + inference_model_path: "" + inference_num_samples: 16 + inference_num_steps: 25 + inference_guidance_scale: 4.5 + inference_save_video: true + inference_log_to_tracker: true diff --git a/examples/train/configs/fine_tuning/mmaudio/small_44k.yaml b/examples/train/configs/fine_tuning/mmaudio/small_44k.yaml new file mode 100644 index 000000000..b4a17fa94 --- /dev/null +++ b/examples/train/configs/fine_tuning/mmaudio/small_44k.yaml @@ -0,0 +1,69 @@ +# MMAudio v1 small_44k conditional flow-matching fine-tuning. +# +# This config intentionally leaves data_path empty. Point it at an official- +# compatible TensorDict memmap feature cache before starting training. + +models: + student: + _target_: fastvideo.train.models.mmaudio.MMAudioModel + init_from: converted_weights/mmaudio/small_44k + trainable: true + +method: + _target_: fastvideo.train.methods.fine_tuning.flow_matching.FlowMatchingFineTuneMethod + +training: + distributed: + num_gpus: 1 + sp_size: 1 + tp_size: 1 + hsdp_replicate_dim: 1 + hsdp_shard_dim: 1 + pin_cpu_memory: false + + data: + # Fill this after the precomputed MMAudio feature cache is ready. + data_path: "" + preprocessed_data_type: mmaudio_features + dataloader_num_workers: 4 + train_batch_size: 1 + training_cfg_rate: 0.1 + seed: 14159265 + + # These match the official MMAudio optimizer defaults. FastVideo currently + # uses constant-with-warmup here; the original recipe adds LR drops at + # iterations 240k and 270k. + optimizer: + learning_rate: 1.0e-4 + betas: [0.9, 0.95] + eps: 1.0e-6 + weight_decay: 1.0e-6 + lr_scheduler: constant_with_warmup + lr_warmup_steps: 1000 + + loop: + max_train_steps: 300000 + gradient_accumulation_steps: 1 + + checkpoint: + output_dir: outputs/mmaudio_small_44k_finetune + training_state_checkpointing_steps: 10000 + checkpoints_total_limit: 3 + resume_from_checkpoint: latest + + tracker: + project_name: fastvideo_mmaudio + run_name: mmaudio_small_44k_finetune + + model: + logit_mean: 0.0 + logit_std: 1.0 + precondition_outputs: false + + # Official MMAudio trains FP32 weights under BF16 autocast. + dit_precision: fp32 + +callbacks: + grad_clip: + _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback + max_grad_norm: 1.0 diff --git a/examples/train/configs/fine_tuning/mmaudio/small_44k_ddp_from_scratch.yaml b/examples/train/configs/fine_tuning/mmaudio/small_44k_ddp_from_scratch.yaml new file mode 100644 index 000000000..838483714 --- /dev/null +++ b/examples/train/configs/fine_tuning/mmaudio/small_44k_ddp_from_scratch.yaml @@ -0,0 +1,108 @@ +# Official MMAudio small_44k training with one full DDP replica per GPU. +# This is independent from the existing FSDP/HSDP scratch config. + +models: + student: + _target_: fastvideo.train.models.mmaudio.MMAudioModel + variant: small_44k + from_scratch: true + empty_string_features_path: "" + latent_statistics_path: "" + latent_statistics_chunk_size: 32 + trainable: true + +method: + _target_: fastvideo.train.methods.fine_tuning.flow_matching.FlowMatchingFineTuneMethod + +training: + model_path: converted_weights/mmaudio/small_44k + distributed: + strategy: ddp + num_gpus: 4 + sp_size: 1 + tp_size: 1 + # DDP replicates the full model and does not use HSDP meshes. + hsdp_replicate_dim: 1 + hsdp_shard_dim: 1 + pin_cpu_memory: false + + data: + data_path: "" + preprocessed_data_type: mmaudio_features + dataloader_num_workers: 0 + train_batch_size: 128 + training_cfg_rate: 0.1 + seed: 14159265 + + optimizer: + learning_rate: 1.0e-4 + betas: [0.9, 0.95] + eps: 1.0e-6 + weight_decay: 1.0e-6 + lr_scheduler: multistep_with_warmup + lr_warmup_steps: 1000 + lr_milestones: [240000, 270000] + lr_gamma: 0.1 + # Matches the upstream Runner's AdamW(..., fused=True). + fused: true + + loop: + max_train_steps: 300000 + # 4 GPUs x 128 samples/GPU = official effective global batch 512. + gradient_accumulation_steps: 1 + + checkpoint: + output_dir: outputs/mmaudio_small_44k_ddp_from_scratch + # Keep this divisible by the 5,000-step PostHocEMA interval so resume + # always has an exactly aligned pair of official EMA snapshots. + training_state_checkpointing_steps: 10000 + checkpoints_total_limit: 3 + resume_from_checkpoint: "" + + tracker: + project_name: fastvideo_mmaudio + run_name: mmaudio_small_44k_ddp_from_scratch + + model: + logit_mean: 0.0 + logit_std: 1.0 + precondition_outputs: false + # Matches official MMAudio's torch.compile(self.train_fn). Set false to + # use eager execution when debugging or on an unsupported compiler stack. + compile_train_fn: true + torch_compile_kwargs: + fullgraph: true + + # Official MMAudio uses FP32 master weights with BF16 autocast forward. + dit_precision: fp32 + +callbacks: + grad_clip: + _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback + max_grad_norm: 1.0 + posthoc_ema: + _target_: fastvideo.train.callbacks.ddp_posthoc_ema.DDPPostHocEMACallback + sigma_rels: [0.05, 0.1] + update_every: 1 + checkpoint_every: 5000 + start_iter: 0 + default_output_sigma: 0.05 + step_size_correction: true + validation: + _target_: fastvideo.train.callbacks.mmaudio_validation.MMAudioValidationCallback + data_path: "" + every_steps: 5000 + max_batches: 0 + batch_size: 8 + num_data_workers: 0 + run_at_start: false + # Official periodic validation samples the online model; EMA is + # synthesized after training. + use_ema: false + inference_every_steps: 20000 + inference_model_path: converted_weights/mmaudio/large_44k_v2 + inference_num_samples: 16 + inference_num_steps: 25 + inference_guidance_scale: 4.5 + inference_save_video: true + inference_log_to_tracker: true diff --git a/examples/train/configs/fine_tuning/mmaudio/small_44k_from_scratch.yaml b/examples/train/configs/fine_tuning/mmaudio/small_44k_from_scratch.yaml new file mode 100644 index 000000000..4162746cd --- /dev/null +++ b/examples/train/configs/fine_tuning/mmaudio/small_44k_from_scratch.yaml @@ -0,0 +1,100 @@ +# Official MMAudio small_44k transformer training from random initialization. +# Requires 44.1 kHz TensorDict features with audio latents [345, 40]. + +models: + student: + _target_: fastvideo.train.models.mmaudio.MMAudioModel + variant: small_44k + from_scratch: true + # Official standalone empty_string.pth. No DiT weights are loaded. + empty_string_features_path: "" + # Optional .pt cache. When empty, stats are recomputed at startup. + latent_statistics_path: "" + latent_statistics_chunk_size: 32 + trainable: true + +method: + _target_: fastvideo.train.methods.fine_tuning.flow_matching.FlowMatchingFineTuneMethod + +training: + # Matching component-tree template used only by DCP-to-Diffusers export. + model_path: converted_weights/mmaudio/small_44k + distributed: + num_gpus: 4 + sp_size: 1 + tp_size: 1 + hsdp_replicate_dim: 1 + hsdp_shard_dim: 4 + pin_cpu_memory: false + + data: + # Set with --training.data.data_path or run_train_vggsound_from_scratch.sh. + data_path: "" + preprocessed_data_type: mmaudio_features + dataloader_num_workers: 2 + train_batch_size: 8 + training_cfg_rate: 0.1 + seed: 14159265 + + optimizer: + learning_rate: 1.0e-4 + betas: [0.9, 0.95] + eps: 1.0e-6 + weight_decay: 1.0e-6 + lr_scheduler: multistep_with_warmup + lr_warmup_steps: 1000 + lr_milestones: [240000, 270000] + lr_gamma: 0.1 + + loop: + max_train_steps: 300000 + # 4 GPUs x 8 samples/GPU x 16 rounds = effective global batch 512. + gradient_accumulation_steps: 16 + + checkpoint: + output_dir: outputs/mmaudio_small_44k_from_scratch + training_state_checkpointing_steps: 10000 + checkpoints_total_limit: 3 + resume_from_checkpoint: "" + + tracker: + project_name: fastvideo_mmaudio + run_name: mmaudio_small_44k_from_scratch + + model: + logit_mean: 0.0 + logit_std: 1.0 + precondition_outputs: false + + # FP32 master weights; MMAudioModel runs transformer forward under BF16 autocast. + dit_precision: fp32 + +callbacks: + grad_clip: + _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback + max_grad_norm: 1.0 + posthoc_ema: + _target_: fastvideo.train.callbacks.posthoc_ema.PostHocEMACallback + sigma_rels: [0.05, 0.1] + update_every: 1 + checkpoint_every: 5000 + start_iter: 0 + default_output_sigma: 0.05 + step_size_correction: true + validation: + _target_: fastvideo.train.callbacks.mmaudio_validation.MMAudioValidationCallback + data_path: "" + every_steps: 5000 + max_batches: 0 + batch_size: 8 + num_data_workers: 2 + run_at_start: false + use_ema: false + inference_every_steps: 20000 + inference_model_path: converted_weights/mmaudio/large_44k_v2 + # Global total across all ranks; fixed by validation seed and sampler. + inference_num_samples: 16 + inference_num_steps: 25 + inference_guidance_scale: 4.5 + inference_save_video: true + inference_log_to_tracker: true diff --git a/examples/training/finetune/mmaudio/preprocess_vggsound.sh b/examples/training/finetune/mmaudio/preprocess_vggsound.sh new file mode 100755 index 000000000..0b8c8a979 --- /dev/null +++ b/examples/training/finetune/mmaudio/preprocess_vggsound.sh @@ -0,0 +1,49 @@ +#!/usr/bin/env bash + +set -euo pipefail + +GPU_NUM="${GPU_NUM:-4}" +MODEL_PATH="${MODEL_PATH:-converted_weights/mmaudio/preprocess_44k}" +DATASET_PATH="${DATASET_PATH:?Set DATASET_PATH to the extracted VGGSound directory}" +SPLIT="${SPLIT:-train}" +CAPTION_PATH="${CAPTION_PATH:-${DATASET_PATH}/sets/filtered_caption/vgg-${SPLIT}-filtered-caption.tsv}" +OUTPUT_DIR="${OUTPUT_DIR:-${DATASET_PATH}/mmaudio_features/${SPLIT}}" +MASTER_PORT="${MASTER_PORT:-29513}" +BATCH_SIZE="${BATCH_SIZE:-1}" +DATALOADER_NUM_WORKERS="${DATALOADER_NUM_WORKERS:-2}" +SAMPLES_PER_FILE="${SAMPLES_PER_FILE:-256}" +VIDEO_LOADER_TYPE="${VIDEO_LOADER_TYPE:-torchcodec}" +export OMP_NUM_THREADS="${OMP_NUM_THREADS:-4}" + +if [[ ! -f "${CAPTION_PATH}" ]]; then + echo "Caption manifest does not exist: ${CAPTION_PATH}" >&2 + exit 1 +fi + +MODEL_LOAD_ARGS=() +if [[ "${VIDEO_LOADER_TYPE}" == "torio" ]]; then + # PyTorch 2.7 torio environments can have a broken NCCL tensor broadcast. + # Loading on CPU makes every rank read its own local weights, matching the + # official MMAudio extractor; the feature stage moves modules to its GPU. + MODEL_LOAD_ARGS+=(--image-encoder-cpu-offload true) + MODEL_LOAD_ARGS+=(--text-encoder-cpu-offload true) + MODEL_LOAD_ARGS+=(--vae-cpu-offload true) +fi + +torchrun --standalone \ + --nproc_per_node="${GPU_NUM}" \ + --master_port="${MASTER_PORT}" \ + -m fastvideo.pipelines.preprocess.v1_preprocessing_new \ + --model-path "${MODEL_PATH}" \ + --mode preprocess \ + --workload-type v2a \ + "${MODEL_LOAD_ARGS[@]}" \ + --preprocess.dataset-type vggsound \ + --preprocess.dataset-path "${DATASET_PATH}" \ + --preprocess.dataset-metadata-path "${CAPTION_PATH}" \ + --preprocess.dataset-split "${SPLIT}" \ + --preprocess.dataset-output-dir "${OUTPUT_DIR}" \ + --preprocess.video-loader-type "${VIDEO_LOADER_TYPE}" \ + --preprocess.preprocess-video-batch-size "${BATCH_SIZE}" \ + --preprocess.dataloader-num-workers "${DATALOADER_NUM_WORKERS}" \ + --preprocess.samples-per-file "${SAMPLES_PER_FILE}" diff --git a/examples/training/finetune/mmaudio/run_preprocess_vggsound.sh b/examples/training/finetune/mmaudio/run_preprocess_vggsound.sh new file mode 100755 index 000000000..4716010ad --- /dev/null +++ b/examples/training/finetune/mmaudio/run_preprocess_vggsound.sh @@ -0,0 +1,119 @@ +#!/usr/bin/env bash + +# One-command VGGSound feature extraction for native FastVideo MMAudio training. +# Every setting can be overridden with an environment variable; the defaults +# match the four-GPU GB200 machine used for the validated smoke tests. + +set -euo pipefail + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" +REPO_ROOT="$(cd -- "${SCRIPT_DIR}/../../../.." && pwd)" + +if [[ "${1:-}" == "--help" || "${1:-}" == "-h" ]]; then + cat <<'EOF' +Usage: + ./examples/training/finetune/mmaudio/run_preprocess_vggsound.sh + +Optional environment variables: + SPLITS="train val test" Dataset splits to process sequentially + GPU_NUM=4 Number of GPUs/ranks + CUDA_VISIBLE_DEVICES=0,1,2,3 + BATCH_SIZE=16 Per-GPU batch size + DATALOADER_NUM_WORKERS=16 Per-rank workers (torio: background decode threads) + DATASET_PATH=... Extracted VGGSound root + MODEL_PATH=... Converted MMAudio preprocessing model + OUTPUT_ROOT=... Root for split feature caches + VIDEO_LOADER_TYPE=torchcodec Media reader (torchcodec or torio) + SAMPLES_PER_FILE=256 Samples per TensorDict shard + MASTER_PORT=29513 First torchrun rendezvous port + +Examples: + SPLITS=train ./examples/training/finetune/mmaudio/run_preprocess_vggsound.sh + BATCH_SIZE=8 GPU_NUM=2 CUDA_VISIBLE_DEVICES=0,1 \ + ./examples/training/finetune/mmaudio/run_preprocess_vggsound.sh +EOF + exit 0 +fi + +if [[ $# -ne 0 ]]; then + echo "Unknown argument: $1 (use --help)" >&2 + exit 2 +fi + +VENV_PATH="${VENV_PATH:-${REPO_ROOT}/.venv}" +if [[ ! -f "${VENV_PATH}/bin/activate" ]]; then + echo "FastVideo virtual environment does not exist: ${VENV_PATH}" >&2 + exit 1 +fi +# shellcheck disable=SC1091 +source "${VENV_PATH}/bin/activate" + +export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0,1,2,3}" +export OMP_NUM_THREADS="${OMP_NUM_THREADS:-4}" +export PYTHONUNBUFFERED=1 + +GPU_NUM="${GPU_NUM:-4}" +BATCH_SIZE="${BATCH_SIZE:-16}" +DATALOADER_NUM_WORKERS="${DATALOADER_NUM_WORKERS:-16}" +SAMPLES_PER_FILE="${SAMPLES_PER_FILE:-256}" +MASTER_PORT="${MASTER_PORT:-29513}" +SPLITS="${SPLITS:-train val test}" +DATASET_PATH="${DATASET_PATH:?Set DATASET_PATH to the extracted VGGSound directory}" +MODEL_PATH="${MODEL_PATH:-${REPO_ROOT}/converted_weights/mmaudio/preprocess_44k}" +OUTPUT_ROOT="${OUTPUT_ROOT:-${DATASET_PATH}/mmaudio_features}" +VIDEO_LOADER_TYPE="${VIDEO_LOADER_TYPE:-torchcodec}" + +if [[ ! -d "${DATASET_PATH}/videos" ]]; then + echo "VGGSound video directory does not exist: ${DATASET_PATH}/videos" >&2 + exit 1 +fi +if [[ ! -f "${MODEL_PATH}/model_index.json" ]]; then + echo "Converted preprocessing model does not exist: ${MODEL_PATH}" >&2 + exit 1 +fi + +echo "FastVideo MMAudio VGGSound preprocessing" +echo " repository: ${REPO_ROOT}" +echo " dataset: ${DATASET_PATH}" +echo " model: ${MODEL_PATH}" +echo " output: ${OUTPUT_ROOT}" +echo " splits: ${SPLITS}" +echo " GPUs: ${CUDA_VISIBLE_DEVICES} (${GPU_NUM} ranks)" +echo " batch: ${BATCH_SIZE} per GPU" +echo " workers: ${DATALOADER_NUM_WORKERS} per GPU" +echo " reader: ${VIDEO_LOADER_TYPE}" + +cd "${REPO_ROOT}" + +port="${MASTER_PORT}" +for split in ${SPLITS}; do + caption_path="${DATASET_PATH}/sets/filtered_caption/vgg-${split}-filtered-caption.tsv" + output_dir="${OUTPUT_ROOT}/${split}" + if [[ ! -f "${caption_path}" ]]; then + echo "Caption manifest does not exist: ${caption_path}" >&2 + exit 1 + fi + + echo + echo "===== Starting ${split} split =====" + echo "Progress is reported below in batches; existing cached IDs are skipped." + + GPU_NUM="${GPU_NUM}" \ + BATCH_SIZE="${BATCH_SIZE}" \ + DATALOADER_NUM_WORKERS="${DATALOADER_NUM_WORKERS}" \ + SAMPLES_PER_FILE="${SAMPLES_PER_FILE}" \ + MASTER_PORT="${port}" \ + MODEL_PATH="${MODEL_PATH}" \ + DATASET_PATH="${DATASET_PATH}" \ + CAPTION_PATH="${caption_path}" \ + OUTPUT_DIR="${output_dir}" \ + SPLIT="${split}" \ + VIDEO_LOADER_TYPE="${VIDEO_LOADER_TYPE}" \ + bash "${SCRIPT_DIR}/preprocess_vggsound.sh" + + echo "===== Finished ${split}: ${output_dir} =====" + port="$((port + 1))" +done + +echo +echo "All requested VGGSound splits finished successfully." diff --git a/examples/training/finetune/mmaudio/run_preprocess_vggsound_torio.sh b/examples/training/finetune/mmaudio/run_preprocess_vggsound_torio.sh new file mode 100755 index 000000000..f261693d0 --- /dev/null +++ b/examples/training/finetune/mmaudio/run_preprocess_vggsound_torio.sh @@ -0,0 +1,32 @@ +#!/usr/bin/env bash + +# Run the FastVideo MMAudio preprocessing pipeline with the optional torio +# reference media reader in an isolated environment. + +set -euo pipefail + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" +REPO_ROOT="$(cd -- "${SCRIPT_DIR}/../../../.." && pwd)" + +if [[ "${1:-}" == "--help" || "${1:-}" == "-h" ]]; then + exec "${SCRIPT_DIR}/run_preprocess_vggsound.sh" "$@" +fi + +TORIO_VENV_PATH="${TORIO_VENV_PATH:-${REPO_ROOT}/../envs/fastvideo-mmaudio-torio}" +DATASET_PATH="${DATASET_PATH:?Set DATASET_PATH to the extracted VGGSound directory}" +OUTPUT_ROOT="${OUTPUT_ROOT:-${DATASET_PATH}/mmaudio_features_torio}" +DATALOADER_NUM_WORKERS="${DATALOADER_NUM_WORKERS:-4}" + +if [[ ! -x "${TORIO_VENV_PATH}/bin/torchrun" ]]; then + echo "FastVideo torio environment does not exist: ${TORIO_VENV_PATH}" >&2 + exit 1 +fi + +export PYTHONNOUSERSITE=1 + +VENV_PATH="${TORIO_VENV_PATH}" \ +DATASET_PATH="${DATASET_PATH}" \ +OUTPUT_ROOT="${OUTPUT_ROOT}" \ +DATALOADER_NUM_WORKERS="${DATALOADER_NUM_WORKERS}" \ +VIDEO_LOADER_TYPE=torio \ + exec "${SCRIPT_DIR}/run_preprocess_vggsound.sh" "$@" diff --git a/examples/training/finetune/mmaudio/run_train_vggsound_ddp_from_scratch.sh b/examples/training/finetune/mmaudio/run_train_vggsound_ddp_from_scratch.sh new file mode 100755 index 000000000..6f5590c7c --- /dev/null +++ b/examples/training/finetune/mmaudio/run_train_vggsound_ddp_from_scratch.sh @@ -0,0 +1,109 @@ +#!/usr/bin/env bash +# Train MMAudio small_44k with the official full-replica DDP strategy. +# This launcher is independent from run_train_vggsound_from_scratch.sh. + +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +REPO_ROOT="$(cd "${SCRIPT_DIR}/../../../.." && pwd)" + +# --------------------------------------------------------------------------- +# Edit the DDP run here. Do not launch while another job owns the same GPUs. +# --------------------------------------------------------------------------- +FEATURE_DIR="/mnt/lustre/vlm-kai/datasets/VGGSound/mmaudio_features_torio/train" +VALIDATION_FEATURE_DIR="/mnt/lustre/vlm-kai/datasets/VGGSound/mmaudio_features_torio/val" +EMPTY_STRING_FEATURES="${REPO_ROOT}/official_weights/mmaudio/raw/ext_weights/empty_string.pth" +EMPTY_STRING_URL="https://github.com/hkchengrex/MMAudio/releases/download/v0.1/empty_string.pth" + +NUM_GPUS=4 +MASTER_PORT=29502 +PER_GPU_BATCH_SIZE=128 +GRAD_ACCUM_STEPS=1 +# /dev/shm is currently 64 MB, so subprocess workers can exhaust shared memory. +DATALOADER_NUM_WORKERS=0 +VALIDATION_NUM_WORKERS=0 + +LEARNING_RATE="1.0e-4" +MAX_TRAIN_STEPS=300000 +COMPILE_TRAIN_FN=true +CHECKPOINT_EVERY_STEPS=10000 +VALIDATION_EVERY_STEPS=5000 +VALIDATION_MAX_BATCHES=0 +INFERENCE_EVERY_STEPS=20000 +INFERENCE_NUM_SAMPLES=16 +INFERENCE_NUM_STEPS=25 +INFERENCE_SAVE_VIDEO=true +INFERENCE_LOG_TO_TRACKER=true +INFERENCE_MODEL_PATH="${REPO_ROOT}/converted_weights/mmaudio/large_44k_v2" + +OUTPUT_DIR="${REPO_ROOT}/outputs/mmaudio_small_44k_ddp_from_scratch" +RESUME_FROM_CHECKPOINT="${OUTPUT_DIR}/checkpoint-20000" +TRACKER="wandb" +TRACKER_PROJECT="fastvideo_mmaudio" + +# These W&B settings are used only when resuming an existing checkpoint. +# For a new run, set RESUME_FROM_CHECKPOINT=""; the variables are then +# explicitly unset so training cannot accidentally append to the old page. +WANDB_RESUME_RUN_ID="vabj9rjb" +WANDB_RESUME_ENTITY="alan-wang-university-of-toronto" + +CONFIG="${REPO_ROOT}/examples/train/configs/fine_tuning/mmaudio/small_44k_ddp_from_scratch.yaml" +LATENT_STATISTICS="${FEATURE_DIR}/latent_statistics_44k.pt" + +if [[ -x "${REPO_ROOT}/.venv/bin/python" ]]; then + export PATH="${REPO_ROOT}/.venv/bin:${PATH}" +fi + +for required_dir in "${FEATURE_DIR}" "${VALIDATION_FEATURE_DIR}"; do + if [[ ! -d "${required_dir}" ]]; then + echo "MMAudio feature directory does not exist: ${required_dir}" >&2 + exit 2 + fi +done + +if [[ ! -f "${EMPTY_STRING_FEATURES}" ]]; then + mkdir -p "$(dirname "${EMPTY_STRING_FEATURES}")" + echo "Downloading official empty-string encoding to ${EMPTY_STRING_FEATURES}" + curl --fail --location --continue-at - \ + --output "${EMPTY_STRING_FEATURES}" \ + "${EMPTY_STRING_URL}" +fi + +if [[ -n "${RESUME_FROM_CHECKPOINT}" ]]; then + export WANDB_RUN_ID="${WANDB_RESUME_RUN_ID}" + export WANDB_RESUME="must" + export WANDB_ENTITY="${WANDB_RESUME_ENTITY}" +else + unset WANDB_RUN_ID WANDB_RESUME WANDB_ENTITY +fi + +export NUM_GPUS MASTER_PORT +cd "${REPO_ROOT}" +exec bash examples/train/run.sh "${CONFIG}" \ + --models.student.empty_string_features_path "${EMPTY_STRING_FEATURES}" \ + --models.student.latent_statistics_path "${LATENT_STATISTICS}" \ + --training.data.data_path "${FEATURE_DIR}" \ + --training.data.train_batch_size "${PER_GPU_BATCH_SIZE}" \ + --training.data.dataloader_num_workers "${DATALOADER_NUM_WORKERS}" \ + --training.optimizer.learning_rate "${LEARNING_RATE}" \ + --training.loop.max_train_steps "${MAX_TRAIN_STEPS}" \ + --training.loop.gradient_accumulation_steps "${GRAD_ACCUM_STEPS}" \ + --training.model.compile_train_fn "${COMPILE_TRAIN_FN}" \ + --training.distributed.num_gpus "${NUM_GPUS}" \ + --training.checkpoint.output_dir "${OUTPUT_DIR}" \ + --training.checkpoint.training_state_checkpointing_steps "${CHECKPOINT_EVERY_STEPS}" \ + --training.checkpoint.resume_from_checkpoint "${RESUME_FROM_CHECKPOINT}" \ + --training.tracker.trackers "[${TRACKER}]" \ + --training.tracker.project_name "${TRACKER_PROJECT}" \ + --training.tracker.run_name "mmaudio_small_44k_ddp_from_scratch" \ + --callbacks.validation.data_path "${VALIDATION_FEATURE_DIR}" \ + --callbacks.validation.every_steps "${VALIDATION_EVERY_STEPS}" \ + --callbacks.validation.max_batches "${VALIDATION_MAX_BATCHES}" \ + --callbacks.validation.num_data_workers "${VALIDATION_NUM_WORKERS}" \ + --callbacks.validation.inference_every_steps "${INFERENCE_EVERY_STEPS}" \ + --callbacks.validation.inference_num_samples "${INFERENCE_NUM_SAMPLES}" \ + --callbacks.validation.inference_num_steps "${INFERENCE_NUM_STEPS}" \ + --callbacks.validation.inference_save_video "${INFERENCE_SAVE_VIDEO}" \ + --callbacks.validation.inference_log_to_tracker "${INFERENCE_LOG_TO_TRACKER}" \ + --callbacks.validation.inference_model_path "${INFERENCE_MODEL_PATH}" \ + "$@" diff --git a/examples/training/finetune/mmaudio/run_train_vggsound_from_scratch.sh b/examples/training/finetune/mmaudio/run_train_vggsound_from_scratch.sh new file mode 100755 index 000000000..d4b88e4ab --- /dev/null +++ b/examples/training/finetune/mmaudio/run_train_vggsound_from_scratch.sh @@ -0,0 +1,126 @@ +#!/usr/bin/env bash +# Train an official MMAudio v1 architecture from random DiT initialization. + +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +REPO_ROOT="$(cd "${SCRIPT_DIR}/../../../.." && pwd)" + +# --------------------------------------------------------------------------- +# Edit the training run here. No environment variables are required. +# --------------------------------------------------------------------------- +VARIANT="small_44k" + +FEATURE_DIR="/mnt/lustre/vlm-kai/datasets/VGGSound/mmaudio_features_torio/train" +VALIDATION_FEATURE_DIR="/mnt/lustre/vlm-kai/datasets/VGGSound/mmaudio_features_torio/val" +EMPTY_STRING_FEATURES="${REPO_ROOT}/official_weights/mmaudio/raw/ext_weights/empty_string.pth" +EMPTY_STRING_URL="https://github.com/hkchengrex/MMAudio/releases/download/v0.1/empty_string.pth" + +NUM_GPUS=4 +# Validated for the default small_44k model on 4 x GB200 (about 30 GB/GPU). +# Re-profile before switching VARIANT to medium_44k or large_44k. +PER_GPU_BATCH_SIZE=128 +GRAD_ACCUM_STEPS=1 +# This container has a 64 MB /dev/shm. Keep these at zero unless the job is +# launched with a larger shared-memory mount. +DATALOADER_NUM_WORKERS=0 +VALIDATION_NUM_WORKERS=0 + +LEARNING_RATE="1.0e-4" +MAX_TRAIN_STEPS=300000 +CHECKPOINT_EVERY_STEPS=10000 +VALIDATION_EVERY_STEPS=5000 +# 0 evaluates the complete validation split. +VALIDATION_MAX_BATCHES=0 +INFERENCE_EVERY_STEPS=20000 +INFERENCE_NUM_SAMPLES=16 +INFERENCE_NUM_STEPS=25 +INFERENCE_SAVE_VIDEO=true +INFERENCE_LOG_TO_TRACKER=true +INFERENCE_MODEL_PATH="${REPO_ROOT}/converted_weights/mmaudio/large_44k_v2" + +OUTPUT_DIR="${REPO_ROOT}/outputs/mmaudio_${VARIANT}_from_scratch" +RESUME_FROM_CHECKPOINT="" +# The API key is stored by `wandb login` in the current user's credentials; +# never put it in this tracked launcher. +TRACKER="wandb" +TRACKER_PROJECT="fastvideo_mmaudio" + +# Prefer the repository environment so the script also works from a fresh +# non-activated shell. An already activated external environment remains the +# fallback when .venv is absent. +if [[ -x "${REPO_ROOT}/.venv/bin/python" ]]; then + export PATH="${REPO_ROOT}/.venv/bin:${PATH}" +fi + +for required_dir in "${FEATURE_DIR}" "${VALIDATION_FEATURE_DIR}"; do + if [[ ! -d "${required_dir}" ]]; then + echo "MMAudio feature directory does not exist: ${required_dir}" >&2 + exit 2 + fi +done + +# Official TRAINING.md publishes this small standalone CLIP encoding. It is +# fixed conditioning data, not a pretrained MMAudio transformer checkpoint. +if [[ ! -f "${EMPTY_STRING_FEATURES}" ]]; then + mkdir -p "$(dirname "${EMPTY_STRING_FEATURES}")" + echo "Downloading official empty-string encoding to ${EMPTY_STRING_FEATURES}" + curl --fail --location --continue-at - \ + --output "${EMPTY_STRING_FEATURES}" \ + "${EMPTY_STRING_URL}" +fi + +case "${VARIANT}" in + small_16k) + CONFIG_NAME="small_16k_from_scratch.yaml" + STATS_NAME="latent_statistics_16k.pt" + INFERENCE_EVERY_STEPS=0 + INFERENCE_MODEL_PATH="" + ;; + small_44k|medium_44k|large_44k) + CONFIG_NAME="${VARIANT}_from_scratch.yaml" + STATS_NAME="latent_statistics_44k.pt" + ;; + large_44k_v2) + echo "Official MMAudio training does not support large_44k_v2." >&2 + exit 2 + ;; + *) + echo "Unknown VARIANT=${VARIANT}" >&2 + exit 2 + ;; +esac + +CONFIG="${REPO_ROOT}/examples/train/configs/fine_tuning/mmaudio/${CONFIG_NAME}" +LATENT_STATISTICS="${FEATURE_DIR}/${STATS_NAME}" + +export NUM_GPUS +cd "${REPO_ROOT}" +exec bash examples/train/run.sh "${CONFIG}" \ + --models.student.empty_string_features_path "${EMPTY_STRING_FEATURES}" \ + --models.student.latent_statistics_path "${LATENT_STATISTICS}" \ + --training.data.data_path "${FEATURE_DIR}" \ + --training.data.train_batch_size "${PER_GPU_BATCH_SIZE}" \ + --training.data.dataloader_num_workers "${DATALOADER_NUM_WORKERS}" \ + --training.optimizer.learning_rate "${LEARNING_RATE}" \ + --training.loop.max_train_steps "${MAX_TRAIN_STEPS}" \ + --training.loop.gradient_accumulation_steps "${GRAD_ACCUM_STEPS}" \ + --training.distributed.num_gpus "${NUM_GPUS}" \ + --training.distributed.hsdp_shard_dim "${NUM_GPUS}" \ + --training.checkpoint.output_dir "${OUTPUT_DIR}" \ + --training.checkpoint.training_state_checkpointing_steps "${CHECKPOINT_EVERY_STEPS}" \ + --training.checkpoint.resume_from_checkpoint "${RESUME_FROM_CHECKPOINT}" \ + --training.tracker.trackers "[${TRACKER}]" \ + --training.tracker.project_name "${TRACKER_PROJECT}" \ + --training.tracker.run_name "mmaudio_${VARIANT}_from_scratch" \ + --callbacks.validation.data_path "${VALIDATION_FEATURE_DIR}" \ + --callbacks.validation.every_steps "${VALIDATION_EVERY_STEPS}" \ + --callbacks.validation.max_batches "${VALIDATION_MAX_BATCHES}" \ + --callbacks.validation.num_data_workers "${VALIDATION_NUM_WORKERS}" \ + --callbacks.validation.inference_every_steps "${INFERENCE_EVERY_STEPS}" \ + --callbacks.validation.inference_num_samples "${INFERENCE_NUM_SAMPLES}" \ + --callbacks.validation.inference_num_steps "${INFERENCE_NUM_STEPS}" \ + --callbacks.validation.inference_save_video "${INFERENCE_SAVE_VIDEO}" \ + --callbacks.validation.inference_log_to_tracker "${INFERENCE_LOG_TO_TRACKER}" \ + --callbacks.validation.inference_model_path "${INFERENCE_MODEL_PATH}" \ + "$@" diff --git a/fastvideo/configs/configs.py b/fastvideo/configs/configs.py index 08d654f51..577d937a1 100644 --- a/fastvideo/configs/configs.py +++ b/fastvideo/configs/configs.py @@ -15,6 +15,7 @@ class DatasetType(str, Enum): """ HF = "hf" MERGED = "merged" + VGGSOUND = "vggsound" @classmethod def from_string(cls, value: str) -> "DatasetType": @@ -37,6 +38,7 @@ class VideoLoaderType(str, Enum): """ TORCHCODEC = "torchcodec" TORCHVISION = "torchvision" + TORIO = "torio" @classmethod def from_string(cls, value: str) -> "VideoLoaderType": @@ -60,6 +62,8 @@ class PreprocessConfig: # Model and dataset configuration model_path: str = "" dataset_path: str = "" + dataset_metadata_path: str = "" + dataset_split: str = "train" dataset_type: DatasetType = DatasetType.HF dataset_output_dir: str = "./output" @@ -104,6 +108,14 @@ class PreprocessConfig: type=str, default=PreprocessConfig.dataset_path, help="Path to the dataset directory for preprocessing") + preprocess_args.add_argument(f"--{prefix_with_dot}dataset-metadata-path", + type=str, + default=PreprocessConfig.dataset_metadata_path, + help="Optional path to the dataset metadata or caption manifest") + preprocess_args.add_argument(f"--{prefix_with_dot}dataset-split", + type=str, + default=PreprocessConfig.dataset_split, + help="Dataset split to preprocess (for example train, val, or test)") preprocess_args.add_argument(f"--{prefix_with_dot}dataset-type", type=str, choices=DatasetType.choices(), @@ -207,6 +219,8 @@ class PreprocessConfig: def check_preprocess_config(self) -> None: if self.dataset_path == "": raise ValueError("dataset_path must be set for preprocess mode") + if not self.dataset_split.strip(): + raise ValueError("dataset_split must not be empty") if self.samples_per_file <= 0: raise ValueError("samples_per_file must be greater than 0") if self.flush_frequency <= 0: diff --git a/fastvideo/configs/models/audio/mmaudio_vae.py b/fastvideo/configs/models/audio/mmaudio_vae.py index 94cc5441b..151c28948 100644 --- a/fastvideo/configs/models/audio/mmaudio_vae.py +++ b/fastvideo/configs/models/audio/mmaudio_vae.py @@ -8,11 +8,15 @@ from fastvideo.configs.models.base import ArchConfig, ModelConfig @dataclass class MMAudioVAEArchConfig(ArchConfig): + """MMAudio mel-spectrogram VAE architecture. + + ``mode`` is the single source of truth: ``MMAudioVAE`` derives the mel + dimension, the latent dimension, and the hidden width from it, so they are + not duplicated here. + """ + 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 diff --git a/fastvideo/configs/models/dits/__init__.py b/fastvideo/configs/models/dits/__init__.py index 50b13c299..7c9586241 100644 --- a/fastvideo/configs/models/dits/__init__.py +++ b/fastvideo/configs/models/dits/__init__.py @@ -12,7 +12,14 @@ 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.mmaudio import ( + MMAUDIO_44K_TRAINING_VARIANTS, + MMAUDIO_TRAINING_VARIANTS, + MMAUDIO_VARIANT_ARCHITECTURES, + MMAudioArchConfig, + MMAudioTransformerConfig, + get_mmaudio_transformer_config, +) from fastvideo.configs.models.dits.stable_audio import StableAudioConfig from fastvideo.models.wan.config import WanVideoConfig from fastvideo.configs.models.dits.zimage import ZImageDiTConfig @@ -30,6 +37,7 @@ __all__ = [ "LongCatVideoConfig", "LTX2VideoConfig", "HYWorldConfig", "Kandinsky5VideoConfig", "Kandinsky6ArchConfig", "Kandinsky6VideoAudioConfig", "Kandinsky6SRArchConfig", "Kandinsky6SRConfig", "MagiHumanVideoConfig", "StableAudioConfig", "GlmImageDiTConfig", "HeliosConfig", "LingBotWorld2CausalFastVideoConfig", - "LingBotWorldFastVideoConfig", - "LingBotVideoConfig", "MiniMaxH3Config", "ZImageDiTConfig", "MMAudioArchConfig", "MMAudioTransformerConfig" + "LingBotWorldFastVideoConfig", "LingBotVideoConfig", "MiniMaxH3Config", "ZImageDiTConfig", "MMAudioArchConfig", + "MMAudioTransformerConfig", "MMAUDIO_VARIANT_ARCHITECTURES", "MMAUDIO_TRAINING_VARIANTS", + "MMAUDIO_44K_TRAINING_VARIANTS", "get_mmaudio_transformer_config" ] diff --git a/fastvideo/configs/models/dits/mmaudio.py b/fastvideo/configs/models/dits/mmaudio.py index aba3ec82b..6f1835cb0 100644 --- a/fastvideo/configs/models/dits/mmaudio.py +++ b/fastvideo/configs/models/dits/mmaudio.py @@ -2,11 +2,83 @@ from __future__ import annotations from dataclasses import dataclass, field +from typing import TypedDict from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig from fastvideo.platforms import AttentionBackendEnum +class MMAudioVariantArchitecture(TypedDict): + latent_dim: int + hidden_dim: int + depth: int + fused_depth: int + num_heads: int + latent_seq_len: int + v2: bool + + +MMAUDIO_VARIANT_ARCHITECTURES: dict[str, MMAudioVariantArchitecture] = { + "small_16k": { + "latent_dim": 20, + "hidden_dim": 448, + "depth": 12, + "fused_depth": 8, + "num_heads": 7, + "latent_seq_len": 250, + "v2": False, + }, + "small_44k": { + "latent_dim": 40, + "hidden_dim": 448, + "depth": 12, + "fused_depth": 8, + "num_heads": 7, + "latent_seq_len": 345, + "v2": False, + }, + "medium_44k": { + "latent_dim": 40, + "hidden_dim": 896, + "depth": 12, + "fused_depth": 8, + "num_heads": 14, + "latent_seq_len": 345, + "v2": False, + }, + "large_44k": { + "latent_dim": 40, + "hidden_dim": 896, + "depth": 21, + "fused_depth": 14, + "num_heads": 14, + "latent_seq_len": 345, + "v2": False, + }, + "large_44k_v2": { + "latent_dim": 40, + "hidden_dim": 896, + "depth": 21, + "fused_depth": 14, + "num_heads": 14, + "latent_seq_len": 345, + "v2": True, + }, +} + +MMAUDIO_TRAINING_VARIANTS = ( + "small_16k", + "small_44k", + "medium_44k", + "large_44k", +) +MMAUDIO_44K_TRAINING_VARIANTS = ( + "small_44k", + "medium_44k", + "large_44k", +) + + def _is_mmaudio_transformer_block(name: str, module) -> bool: del module parts = name.split(".") @@ -47,3 +119,13 @@ class MMAudioArchConfig(DiTArchConfig): class MMAudioTransformerConfig(DiTConfig): arch_config: DiTArchConfig = field(default_factory=MMAudioArchConfig) prefix: str = "MMAudio" + + +def get_mmaudio_transformer_config(variant: str) -> MMAudioTransformerConfig: + """Build the exact official architecture config for ``variant``.""" + try: + architecture = MMAUDIO_VARIANT_ARCHITECTURES[variant] + except KeyError as exc: + supported = ", ".join(MMAUDIO_VARIANT_ARCHITECTURES) + raise ValueError(f"Unknown MMAudio variant {variant!r}; expected one of: {supported}") from exc + return MMAudioTransformerConfig(arch_config=MMAudioArchConfig(**architecture), ) diff --git a/fastvideo/configs/pipelines/__init__.py b/fastvideo/configs/pipelines/__init__.py index aaf732a7f..456d33231 100644 --- a/fastvideo/configs/pipelines/__init__.py +++ b/fastvideo/configs/pipelines/__init__.py @@ -28,6 +28,6 @@ __all__ = [ "DreamXWorld5BCamPipelineConfig", "DreamXWorld5BARPipelineConfig", "HeliosPipelineConfig", "HYWorldConfig", "Kandinsky5T2VConfig", "Kandinsky5I2VConfig", "Kandinsky5DMDConfig", "Kandinsky6TI2VAConfig", "Kandinsky6SRPipelineConfig", "LingBotWorld2CausalFastI2V480PConfig", "LingBotWorldFastI2V480PConfig", - "LingBotVideoT2VConfig", - "MatrixGame2I2V480PConfig", "MatrixGame3I2V720PConfig", "MMAudioV2AConfig", "get_pipeline_config_cls_from_name" + "LingBotVideoT2VConfig", "MatrixGame2I2V480PConfig", "MatrixGame3I2V720PConfig", "MMAudioV2AConfig", + "get_pipeline_config_cls_from_name" ] diff --git a/fastvideo/configs/pipelines/mmaudio.py b/fastvideo/configs/pipelines/mmaudio.py index e340ffe26..64b35ccfd 100644 --- a/fastvideo/configs/pipelines/mmaudio.py +++ b/fastvideo/configs/pipelines/mmaudio.py @@ -6,8 +6,9 @@ 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.audio import (BigVGANV2ArchConfig, BigVGANV2Config, MMAudioVAEArchConfig, + MMAudioVAEConfig) +from fastvideo.configs.models.dits import MMAudioTransformerConfig, get_mmaudio_transformer_config from fastvideo.configs.models.encoders import ( MMAudioDFNCLIPTextConfig, MMAudioDFNCLIPVisionConfig, @@ -39,6 +40,10 @@ class MMAudioV2AConfig(PipelineConfig): audio_decoder_config: ModelConfig = field(default_factory=MMAudioVAEConfig) audio_decoder_precision: str = "bf16" + # Offline training preprocessing uses the VAE encoder in fp32, matching + # the published MMAudio feature extraction recipe. Inference remains + # decoder-only bf16 through ``audio_decoder_precision`` above. + audio_encoder_precision: str = "fp32" vocoder_config: ModelConfig = field(default_factory=BigVGANV2Config) vocoder_precision: str = "bf16" @@ -46,6 +51,8 @@ class MMAudioV2AConfig(PipelineConfig): # supports other durations, although quality can drop far away from the # eight-second training duration. duration_s: float = 8.0 + # ``None`` means no hard cap; set it to reject requests beyond the + # checkpoint's validated maximum. max_audio_duration_s: float | None = None sampling_rate: int = 44100 spectrogram_frame_rate: int = 512 @@ -62,5 +69,53 @@ class MMAudioV2AConfig(PipelineConfig): num_inference_steps: int = 25 guidance_scale: float = 4.5 + # Inherited from ``PipelineConfig``; MMAudio decodes audio through its own + # stage, so the video-VAE paths that read these are not exercised. vae_tiling: bool = False vae_sp: bool = False + + +@dataclass +class MMAudioSmall16kV2AConfig(MMAudioV2AConfig): + """Official small 16 kHz MMAudio inference configuration.""" + + dit_config: DiTConfig = field(default_factory=lambda: get_mmaudio_transformer_config("small_16k")) + # ``MMAudioVAEArchConfig.mode`` is the single source of truth for the VAE + # shapes; ``MMAudioVAE`` re-derives them from it. + audio_decoder_config: ModelConfig = field( + default_factory=lambda: MMAudioVAEConfig(arch_config=MMAudioVAEArchConfig(mode="16k", ))) + vocoder_config: ModelConfig = field(default_factory=lambda: BigVGANV2Config(arch_config=BigVGANV2ArchConfig( + sample_rate=16000, + num_mels=80, + ))) + sampling_rate: int = 16000 + spectrogram_frame_rate: int = 256 + + +@dataclass +class MMAudioSmall44kV2AConfig(MMAudioV2AConfig): + dit_config: DiTConfig = field(default_factory=lambda: get_mmaudio_transformer_config("small_44k")) + + +@dataclass +class MMAudioMedium44kV2AConfig(MMAudioV2AConfig): + dit_config: DiTConfig = field(default_factory=lambda: get_mmaudio_transformer_config("medium_44k")) + + +@dataclass +class MMAudioLarge44kV2AConfig(MMAudioV2AConfig): + dit_config: DiTConfig = field(default_factory=lambda: get_mmaudio_transformer_config("large_44k")) + + +@dataclass +class MMAudioLarge44kV2PipelineConfig(MMAudioV2AConfig): + dit_config: DiTConfig = field(default_factory=lambda: get_mmaudio_transformer_config("large_44k_v2")) + + +MMAUDIO_PIPELINE_CONFIGS = { + "small_16k": MMAudioSmall16kV2AConfig, + "small_44k": MMAudioSmall44kV2AConfig, + "medium_44k": MMAudioMedium44kV2AConfig, + "large_44k": MMAudioLarge44kV2AConfig, + "large_44k_v2": MMAudioLarge44kV2PipelineConfig, +} diff --git a/fastvideo/dataset/__init__.py b/fastvideo/dataset/__init__.py index 99ac81efc..2bc877c9b 100644 --- a/fastvideo/dataset/__init__.py +++ b/fastvideo/dataset/__init__.py @@ -4,6 +4,9 @@ from torchvision.transforms import Lambda from fastvideo.dataset.parquet_dataset_map_style import (build_parquet_map_style_dataloader) from fastvideo.dataset.ltx2_precomputed_dataset import (build_ltx2_precomputed_dataloader, LTX2PrecomputedDataset) +from fastvideo.dataset.mmaudio_feature_dataset import (MMAudioFeatureDataset, build_mmaudio_feature_dataloader) +from fastvideo.dataset.v2a_feature_cache import V2AFeatureShardWriter +from fastvideo.dataset.vggsound import VGGSoundDataset from fastvideo.dataset.preprocessing_datasets import VideoCaptionMergedDataset, TextDataset from fastvideo.dataset.transform import (CenterCropResizeVideo, Normalize255, TemporalRandomCrop) from fastvideo.dataset.validation_dataset import ValidationDataset @@ -43,6 +46,10 @@ __all__ = [ "build_parquet_map_style_dataloader", "build_ltx2_precomputed_dataloader", "LTX2PrecomputedDataset", + "MMAudioFeatureDataset", + "build_mmaudio_feature_dataloader", + "V2AFeatureShardWriter", + "VGGSoundDataset", "ValidationDataset", "VideoCaptionMergedDataset", "TextDataset", diff --git a/fastvideo/dataset/mmaudio_feature_dataset.py b/fastvideo/dataset/mmaudio_feature_dataset.py new file mode 100644 index 000000000..4736a582a --- /dev/null +++ b/fastvideo/dataset/mmaudio_feature_dataset.py @@ -0,0 +1,327 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Memory-mapped precomputed feature dataset for audio generation training.""" + +from __future__ import annotations + +import json +from collections.abc import Sequence +from pathlib import Path +from typing import Any + +import torch +from torch.utils.data import ConcatDataset, Dataset +from torchdata.stateful_dataloader import StatefulDataLoader + +from fastvideo.dataset.parquet_dataset_map_style import DP_SP_BatchSampler +from fastvideo.distributed import get_sp_world_size, get_world_rank, get_world_size +from fastvideo.logger import init_logger + +logger = init_logger(__name__) + + +class MMAudioFeatureDataset(Dataset): + """Read MMAudio-compatible TensorDict memmaps without upstream imports. + + Required tensors are ``mean``, ``std``, and ``text_features``. Video + caches additionally contain ``clip_features`` and ``sync_features``. + Audio-only caches omit both video tensors; fixed-size placeholders and a + false ``video_exists`` flag are returned so video and audio datasets can be + mixed by the same dataloader. + """ + + def __init__( + self, + root: str | Path, + *, + latent_seq_len: int, + latent_dim: int, + clip_seq_len: int, + clip_dim: int, + sync_seq_len: int, + sync_dim: int, + text_seq_len: int, + text_dim: int, + include_metadata: bool = False, + ) -> None: + super().__init__() + self.root = Path(root).expanduser().resolve() + if not self.root.is_dir(): + raise FileNotFoundError(f"MMAudio feature cache does not exist: {self.root}") + + try: + from tensordict import TensorDict + except ImportError as exc: + raise ImportError( + "MMAudio training caches require TensorDict. Install with " + "`uv pip install -e '.[mmaudio-train]'`." + ) from exc + + self.tensors = TensorDict.load_memmap(self.root) + keys = set(self.tensors.keys()) + required = {"mean", "std", "text_features"} + missing = sorted(required - keys) + if missing: + raise ValueError(f"MMAudio feature cache {self.root} is missing tensors: {missing}") + + self.mean = self.tensors["mean"] + self.std = self.tensors["std"] + self.text_features = self.tensors["text_features"] + self.clip_features = self.tensors.get("clip_features") + self.sync_features = self.tensors.get("sync_features") + if (self.clip_features is None) != (self.sync_features is None): + raise ValueError( + "MMAudio feature cache must contain both clip_features and " + "sync_features, or neither for audio-only data." + ) + + self._length = int(self.mean.shape[0]) + for name in required: + if int(self.tensors[name].shape[0]) != self._length: + raise ValueError(f"MMAudio cache tensor {name!r} has a mismatched sample count") + if self.clip_features is not None and int(self.clip_features.shape[0]) != self._length: + raise ValueError("MMAudio cache clip_features has a mismatched sample count") + if self.sync_features is not None and int(self.sync_features.shape[0]) != self._length: + raise ValueError("MMAudio cache sync_features has a mismatched sample count") + + self._require_shape("mean", self.mean, (latent_seq_len, latent_dim)) + self._require_shape("std", self.std, (latent_seq_len, latent_dim)) + self._require_shape("text_features", self.text_features, (text_seq_len, text_dim)) + if self.clip_features is not None: + self._require_shape("clip_features", self.clip_features, (clip_seq_len, clip_dim)) + assert self.sync_features is not None + self._require_shape("sync_features", self.sync_features, (sync_seq_len, sync_dim)) + + feature_dtype = self.text_features.dtype + self.empty_clip = torch.zeros(clip_seq_len, clip_dim, dtype=feature_dtype) + self.empty_sync = torch.zeros(sync_seq_len, sync_dim, dtype=feature_dtype) + self.has_video = self.clip_features is not None + self.metadata: list[dict[str, Any]] | None = None + if include_metadata: + metadata_path = self.root / "samples.jsonl" + if metadata_path.is_file(): + with metadata_path.open(encoding="utf-8") as handle: + self.metadata = [json.loads(line) for line in handle if line.strip()] + if len(self.metadata) != self._length: + raise ValueError( + f"MMAudio cache {metadata_path} has {len(self.metadata)} " + f"metadata rows for {self._length} feature samples" + ) + else: + logger.warning( + "MMAudio cache %s has no samples.jsonl; validation " + "audio cannot be associated with source videos", + self.root, + ) + self.metadata = [{} for _ in range(self._length)] + logger.info("Loaded %d MMAudio feature samples from %s", self._length, self.root) + + @staticmethod + def _require_shape( + name: str, + tensor: torch.Tensor, + expected: tuple[int, ...], + ) -> None: + actual = tuple(tensor.shape[1:]) + if actual != expected: + raise ValueError( + f"MMAudio cache tensor {name!r} must have per-sample shape " + f"{expected}, got {actual}" + ) + + def __len__(self) -> int: + return self._length + + def __getitem__(self, index: int) -> dict[str, Any]: + clip_features = self.empty_clip + sync_features = self.empty_sync + if self.has_video: + assert self.clip_features is not None and self.sync_features is not None + clip_features = self.clip_features[index] + sync_features = self.sync_features[index] + sample: dict[str, Any] = { + "audio_latent_mean": self.mean[index], + "audio_latent_std": self.std[index], + "clip_features": clip_features, + "sync_features": sync_features, + "text_features": self.text_features[index], + "video_exists": torch.tensor(self.has_video, dtype=torch.bool), + "text_exists": torch.tensor(True, dtype=torch.bool), + } + if self.metadata is not None: + metadata = self.metadata[index] + sample.update({ + "sample_id": str(metadata.get("id", index)), + "caption": str(metadata.get("caption", "")), + "source_path": str(metadata.get("source", "")), + }) + return sample + + +def _data_specs(data_path: str | Sequence[str] | dict[str, int]) -> list[tuple[str, int]]: + if isinstance(data_path, dict): + return [ + (str(path), int(repeat)) + for path, repeat in data_path.items() + if int(repeat) > 0 + ] + if isinstance(data_path, Sequence) and not isinstance(data_path, str): + return [(str(path), 1) for path in data_path if str(path).strip()] + return [(part.strip(), 1) for part in str(data_path).split(",") if part.strip()] + + +def _expand_cache_root(root: str | Path) -> list[Path]: + """Resolve one direct TensorDict cache or a directory of cache shards.""" + path = Path(root).expanduser().resolve() + if not path.is_dir(): + raise FileNotFoundError(f"MMAudio feature cache does not exist: {path}") + if (path / "meta.json").is_file(): + return [path] + shards = sorted(metadata.parent for metadata in path.rglob("meta.json")) + if not shards: + raise FileNotFoundError( + f"No TensorDict cache shards (meta.json) were found under {path}") + return shards + + +def compute_mmaudio_latent_stats( + data_path: str | Sequence[str] | dict[str, int], + *, + latent_seq_len: int, + latent_dim: int, + chunk_size: int = 32, +) -> tuple[torch.Tensor, torch.Tensor]: + """Compute official MMAudio normalization stats from the first cache. + + The reference trainer uses the posterior means from its first video + dataset and reduces over sample and sequence dimensions. This chunked + implementation preserves that contract without materializing the complete + VGGSound tensor in RAM. + """ + specs = _data_specs(data_path) + if not specs: + raise ValueError( + "training.data.data_path is empty; latent statistics require a " + "video feature cache." + ) + root = specs[0][0] + chunk_size = max(1, int(chunk_size)) + total = 0 + value_sum = torch.zeros(latent_dim, dtype=torch.float64) + square_sum = torch.zeros(latent_dim, dtype=torch.float64) + + try: + from tensordict import TensorDict + except ImportError as exc: + raise ImportError( + "MMAudio latent statistics require TensorDict. Install with " + "`uv pip install -e '.[mmaudio-train]'`." + ) from exc + + for shard_root in _expand_cache_root(root): + tensors = TensorDict.load_memmap(shard_root) + mean = tensors.get("mean") + if not isinstance(mean, torch.Tensor): + raise ValueError( + f"MMAudio cache {shard_root} is missing tensor 'mean'" + ) + if tuple(mean.shape[1:]) != (latent_seq_len, latent_dim): + raise ValueError( + f"MMAudio cache {shard_root} mean must have per-sample shape " + f"{(latent_seq_len, latent_dim)}, got {tuple(mean.shape[1:])}" + ) + for start in range(0, int(mean.shape[0]), chunk_size): + values = mean[start:start + chunk_size].to(torch.float64) + flattened = values.reshape(-1, latent_dim) + value_sum += flattened.sum(dim=0) + square_sum += flattened.square().sum(dim=0) + total += int(flattened.shape[0]) + + if total < 2: + raise ValueError( + f"MMAudio latent statistics need at least two values, got {total}" + ) + latent_mean = value_sum / total + variance = ( + square_sum - value_sum.square() / total + ).clamp_min(0.0) / (total - 1) + latent_std = variance.sqrt() + return ( + latent_mean.to(torch.float32).reshape(1, 1, latent_dim), + latent_std.to(torch.float32).reshape(1, 1, latent_dim), + ) + + +def build_mmaudio_feature_dataset( + data_path: str | Sequence[str] | dict[str, int], + *, + feature_shapes: dict[str, int], + include_metadata: bool = False, +) -> Dataset: + """Build one map-style dataset over MMAudio TensorDict cache shards. + + Unlike :func:`build_mmaudio_feature_dataloader`, this helper does not add + a distributed sampler. Dataset-scale inference can therefore assign + indices explicitly (for example ``range(rank, len(dataset), world_size)``) + without ``DistributedSampler`` padding the tail with duplicate samples. + """ + specs = _data_specs(data_path) + if not specs: + raise ValueError( + "data_path is empty. Set it to one or more precomputed MMAudio " + "TensorDict mmap directories." + ) + + datasets: list[Dataset] = [] + for root, repeat in specs: + for shard_root in _expand_cache_root(root): + dataset = MMAudioFeatureDataset( + shard_root, + include_metadata=include_metadata, + **feature_shapes, + ) + datasets.extend([dataset] * repeat) + return datasets[0] if len(datasets) == 1 else ConcatDataset(datasets) + + +def build_mmaudio_feature_dataloader( + data_path: str | Sequence[str] | dict[str, int], + *, + batch_size: int, + num_data_workers: int, + seed: int, + pin_memory: bool, + feature_shapes: dict[str, int], + include_metadata: bool = False, +) -> StatefulDataLoader: + """Build a distributed/stateful loader over one or more feature caches.""" + combined = build_mmaudio_feature_dataset( + data_path, + feature_shapes=feature_shapes, + include_metadata=include_metadata, + ) + + sp_world_size = get_sp_world_size() + sampler = DP_SP_BatchSampler( + batch_size=int(batch_size), + dataset_size=len(combined), + num_sp_groups=get_world_size() // sp_world_size, + sp_world_size=sp_world_size, + global_rank=get_world_rank(), + drop_last=True, + seed=int(seed), + ) + return StatefulDataLoader( + combined, + batch_sampler=sampler, + num_workers=int(num_data_workers), + pin_memory=bool(pin_memory), + persistent_workers=int(num_data_workers) > 0, + ) + + +__all__ = [ + "MMAudioFeatureDataset", + "build_mmaudio_feature_dataset", + "build_mmaudio_feature_dataloader", + "compute_mmaudio_latent_stats", +] diff --git a/fastvideo/dataset/v2a_feature_cache.py b/fastvideo/dataset/v2a_feature_cache.py new file mode 100644 index 000000000..adbed1fcf --- /dev/null +++ b/fastvideo/dataset/v2a_feature_cache.py @@ -0,0 +1,126 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Shared sharded feature-cache writer for video-to-audio models.""" + +from __future__ import annotations + +import json +from collections import defaultdict +from pathlib import Path +from typing import Any + +import torch + +from fastvideo.logger import init_logger + +logger = init_logger(__name__) + + +class V2AFeatureShardWriter: + """Write model-specific feature dictionaries as resumable TensorDict shards. + + Pipelines publish a ``dict[str, Tensor]`` with a common batch dimension; + this writer handles buffering, per-rank shard names, metadata, and resume + discovery without knowing what any feature means. + """ + + def __init__( + self, + output_root: str | Path, + *, + rank: int, + samples_per_shard: int, + ) -> None: + if samples_per_shard <= 0: + raise ValueError("samples_per_shard must be positive") + self.output_root = Path(output_root).expanduser().resolve() + self.worker_root = self.output_root / f"worker_{rank:05d}" + self.worker_root.mkdir(parents=True, exist_ok=True) + self.samples_per_shard = int(samples_per_shard) + self._features: dict[str, list[torch.Tensor]] = defaultdict(list) + self._metadata: list[dict[str, Any]] = [] + self.existing_ids = self._read_existing_ids(self.output_root) + existing_indices = [] + for path in self.worker_root.glob("shard_*"): + try: + existing_indices.append(int(path.name.split("_")[-1])) + except ValueError: + continue + self._next_shard = max(existing_indices, default=-1) + 1 + + @staticmethod + def _read_existing_ids(root: Path) -> set[str]: + existing: set[str] = set() + if not root.is_dir(): + return existing + for metadata_path in root.rglob("samples.jsonl"): + with metadata_path.open(encoding="utf-8") as handle: + for line in handle: + if not line.strip(): + continue + sample_id = json.loads(line).get("id") + if sample_id: + existing.add(str(sample_id)) + return existing + + def contains(self, sample_id: str) -> bool: + return sample_id in self.existing_ids + + @property + def buffered_samples(self) -> int: + return len(self._metadata) + + def append( + self, + features: dict[str, torch.Tensor], + metadata: list[dict[str, Any]], + ) -> None: + if not features: + return + batch_sizes = {int(tensor.shape[0]) for tensor in features.values()} + if len(batch_sizes) != 1: + raise ValueError(f"V2A feature batch dimensions do not match: {batch_sizes}") + batch_size = batch_sizes.pop() + if batch_size != len(metadata): + raise ValueError(f"V2A feature batch has {batch_size} samples but {len(metadata)} metadata rows") + if self._features and set(features) != set(self._features): + raise ValueError( + f"V2A feature keys changed within one shard: {sorted(self._features)} vs {sorted(features)}") + + for index, row in enumerate(metadata): + sample_id = str(row["id"]) + if sample_id in self.existing_ids: + continue + for name, tensor in features.items(): + self._features[name].append(tensor[index].detach().cpu().contiguous()) + self._metadata.append(dict(row)) + self.existing_ids.add(sample_id) + if self.buffered_samples >= self.samples_per_shard: + self.flush() + + def flush(self) -> Path | None: + if not self._metadata: + return None + try: + from tensordict import TensorDict + except ImportError as exc: + raise ImportError("V2A feature caches require TensorDict; install fastvideo[mmaudio-train].") from exc + + shard_path = self.worker_root / f"shard_{self._next_shard:06d}" + if shard_path.exists(): + raise FileExistsError(f"Refusing to overwrite existing V2A cache shard: {shard_path}") + stacked = {name: torch.stack(values) for name, values in self._features.items()} + TensorDict(stacked, batch_size=[len(self._metadata)]).memmap_(str(shard_path)) + with (shard_path / "samples.jsonl").open("w", encoding="utf-8") as handle: + for row in self._metadata: + handle.write(json.dumps(row, ensure_ascii=False) + "\n") + logger.info("Wrote %d V2A feature samples to %s", len(self._metadata), shard_path) + self._features.clear() + self._metadata.clear() + self._next_shard += 1 + return shard_path + + def close(self) -> None: + self.flush() + + +__all__ = ["V2AFeatureShardWriter"] diff --git a/fastvideo/dataset/vggsound.py b/fastvideo/dataset/vggsound.py new file mode 100644 index 000000000..669864681 --- /dev/null +++ b/fastvideo/dataset/vggsound.py @@ -0,0 +1,118 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Raw VGGSound metadata adapter for V2A preprocessing.""" + +from __future__ import annotations + +import csv +import os +from pathlib import Path + +from torch.utils.data import Dataset + +from fastvideo.logger import init_logger + +logger = init_logger(__name__) + + +class VGGSoundDataset(Dataset): + """Map VGGSound metadata rows to local MP4 paths and captions. + + ``Loie/VGGSound`` stores clips as ``_.mp4``. + The downloaded tar archives must be extracted before random-access GPU + preprocessing; repeatedly seeking inside gzip archives is prohibitively + expensive for a shuffled training dataset. + """ + + def __init__( + self, + root: str | Path, + *, + split: str = "train", + metadata_path: str | Path | None = None, + ) -> None: + super().__init__() + self.root = Path(root).expanduser().resolve() + metadata = (Path(metadata_path).expanduser().resolve() if metadata_path is not None else self.root / + "vggsound.csv") + if not metadata.is_file(): + raise FileNotFoundError(f"VGGSound metadata file does not exist: {metadata}") + + candidates = (self.root / "videos", self.root / "video", self.root) + self.video_root = next((path for path in candidates if path.is_dir()), self.root) + self.samples: list[tuple[str, str, Path]] = [] + if metadata.suffix.lower() == ".tsv": + self._read_caption_manifest(metadata, split=split) + else: + self._read_vggsound_csv(metadata, split=split) + + if not self.samples: + raise ValueError(f"No VGGSound samples found for split {split!r} in {metadata}") + + def _read_caption_manifest(self, metadata: Path, *, split: str) -> None: + seen_ids: set[str] = set() + missing_videos = 0 + available_videos = { + Path(name).stem + for name in os.listdir(self.video_root) + if Path(name).suffix.lower() == ".mp4" + } + with metadata.open(newline="", encoding="utf-8-sig") as handle: + reader = csv.DictReader(handle, delimiter="\t") + required_columns = {"id", "label"} + columns = set(reader.fieldnames or ()) + if not required_columns.issubset(columns): + raise ValueError( + f"VGGSound caption manifest {metadata} must contain columns {sorted(required_columns)}, " + f"got {sorted(columns)}") + for row_number, row in enumerate(reader, start=2): + sample_id = str(row.get("id") or "").strip() + caption = row.get("label") + if not sample_id: + raise ValueError(f"VGGSound caption manifest has an empty id at row {row_number}") + if caption is None or not caption.strip(): + raise ValueError(f"VGGSound caption manifest has an empty label at row {row_number}") + if sample_id in seen_ids: + raise ValueError(f"Duplicate VGGSound id {sample_id!r} at row {row_number}") + seen_ids.add(sample_id) + video_path = self.video_root / f"{sample_id}.mp4" + if sample_id not in available_videos: + missing_videos += 1 + continue + self.samples.append((sample_id, caption, video_path)) + logger.info( + "Loaded %d VGGSound %s captions from %s (%d videos missing)", + len(self.samples), + split, + metadata, + missing_videos, + ) + + def _read_vggsound_csv(self, metadata: Path, *, split: str) -> None: + with metadata.open(newline="", encoding="utf-8") as handle: + reader = csv.reader(handle) + for row_number, row in enumerate(reader, start=1): + if len(row) != 4: + raise ValueError(f"VGGSound CSV row {row_number} must have four columns, got {row!r}") + youtube_id, start_seconds, caption, row_split = row + if row_split.strip().lower() != split.lower(): + continue + try: + start = int(start_seconds) + except ValueError as exc: + raise ValueError(f"Invalid VGGSound start time at row {row_number}: {start_seconds!r}") from exc + sample_id = f"{youtube_id}_{start:06d}" + self.samples.append((sample_id, caption, self.video_root / f"{sample_id}.mp4")) + + def __len__(self) -> int: + return len(self.samples) + + def __getitem__(self, index: int) -> dict[str, str]: + sample_id, caption, path = self.samples[index] + return { + "id": sample_id, + "caption": caption, + "video_path": str(path), + } + + +__all__ = ["VGGSoundDataset"] diff --git a/fastvideo/entrypoints/cli/eval.py b/fastvideo/entrypoints/cli/eval.py index 50cb4376b..a35dcb958 100644 --- a/fastvideo/entrypoints/cli/eval.py +++ b/fastvideo/entrypoints/cli/eval.py @@ -43,20 +43,30 @@ class EvalSubcommand(CLISubcommand): _cmd_list(args) elif action == "run": _cmd_run(args) + elif action == "v2a": + _cmd_v2a(args) else: # Re-print help if no action was given. self._parser.print_help() # type: ignore[attr-defined] def validate(self, args: argparse.Namespace) -> None: action = getattr(args, "eval_action", None) - if action == "run" and not args.videos: - raise SystemExit("`fastvideo eval run` requires --videos") + if action == "v2a": + if args.backend != "av-benchmark": + raise SystemExit(f"Unsupported V2A evaluation backend: {args.backend}") + return + if action != "run": + return + if not args.manifest and not args.videos and not args.audios: + raise SystemExit("`fastvideo eval run` requires --manifest, --videos, or --audios") + if args.manifest and (args.videos or args.audios or args.reference or args.reference_audio): + raise SystemExit("--manifest cannot be combined with media path arguments") def subparser_init(self, subparsers: argparse._SubParsersAction) -> FlexibleArgumentParser: eval_parser = subparsers.add_parser( "eval", - help="Run video-gen evaluation metrics", - usage="fastvideo eval {list,run} [...]", + help="Run video and V2A evaluation metrics", + usage="fastvideo eval {list,run,v2a} [...]", ) sub = eval_parser.add_subparsers(dest="eval_action", required=False) @@ -66,27 +76,101 @@ class EvalSubcommand(CLISubcommand): # `eval run` run_p = sub.add_parser("run", help="Evaluate videos against one or more metrics") + run_p.add_argument( + "--manifest", + type=str, + default=None, + help="JSON/JSONL canonical sample manifest (recommended for benchmarks).", + ) run_p.add_argument("--videos", type=str, nargs="+", required=False, help="Path, glob, or directory of generated videos.") + run_p.add_argument("--audios", + type=str, + nargs="+", + required=False, + help="Path, glob, or directory of generated audio files.") run_p.add_argument("--reference", type=str, default=None, help="Path / glob / dir of reference videos (for paired metrics).") + run_p.add_argument( + "--reference-audio", + type=str, + nargs="+", + default=None, + help="Path, glob, or directory of reference audio files.", + ) run_p.add_argument("--metrics", type=str, default="all", help="Comma-separated metric names, or 'all'.") run_p.add_argument("--device", type=str, default="cuda", help="Torch device (e.g. 'cuda', 'cuda:0', 'cpu').") - run_p.add_argument("--text-prompt", - type=str, - nargs="*", - default=None, - help="Prompt(s) for text-conditioned metrics. One per video.") + run_p.add_argument("--num-gpus", type=int, default=1, help="Number of metric replicas for dataset evaluation.") + run_p.add_argument("--compile", action="store_true", help="Compile supported metric models.") + run_p.add_argument( + "--skip-missing-deps", + action="store_true", + help="Skip explicitly requested metrics whose optional dependencies are missing.", + ) + run_p.add_argument( + "--extract-audio", + nargs="?", + const=True, + default=False, + help="Extract missing audio tracks; optionally provide a persistent cache directory.", + ) + run_p.add_argument("--extract-workers", type=int, default=4) + run_p.add_argument( + "--text-prompt", + type=str, + nargs="*", + default=None, + help="Prompt(s) for text-conditioned metrics. One per video.", + ) run_p.add_argument("--fps", type=float, default=None, help="Frame-rate annotation passed to fps-aware metrics.") run_p.add_argument("--output", type=str, default=None, help="Write results as JSON to this path (default: stdout).") + run_p.add_argument( + "--output-format", + choices=("legacy", "full"), + default="legacy", + help="'full' includes corpus metrics and sample identities.", + ) + + # `eval v2a`: isolated corpus-level backends for video-to-audio. + v2a_p = sub.add_parser("v2a", help="Evaluate V2A audio with a corpus benchmark backend") + v2a_p.add_argument("--backend", choices=("av-benchmark", ), default="av-benchmark") + v2a_p.add_argument("--audio-dir", required=True, help="Directory containing generated WAV/FLAC files.") + v2a_p.add_argument("--gt-cache", required=True, help="Official ground-truth feature cache directory.") + v2a_p.add_argument( + "--prediction-cache", + default=None, + help="Generated-audio feature cache (default: /../av_benchmark_cache).", + ) + v2a_p.add_argument( + "--output", + default=None, + help="Result JSON (default: /../av_benchmark_results.json).", + ) + v2a_p.add_argument( + "--python-executable", + default=None, + help="Python from an isolated av-benchmark environment.", + ) + v2a_p.add_argument("--device", default="cuda") + v2a_p.add_argument("--batch-size", type=int, default=32) + v2a_p.add_argument("--num-workers", type=int, default=0) + v2a_p.add_argument("--audio-length", type=float, default=8.0) + v2a_p.add_argument("--recompute", action="store_true", help="Re-extract the prediction cache.") + v2a_p.add_argument("--skip-video-related", action="store_true", help="Skip ImageBind and DeSync.") + v2a_p.add_argument("--skip-clap", action="store_true", help="Skip LAION/MS CLAP extraction and scores.") + v2a_p.add_argument( + "--align-prediction-keys", + action="store_true", + help="Align uniquely underscore-prefixed prediction ids to official GT cache ids.", + ) # Stash the parser so cmd() can re-print help on no-action. self._parser = eval_parser # type: ignore[attr-defined] @@ -95,6 +179,7 @@ class EvalSubcommand(CLISubcommand): def _cmd_list(args: argparse.Namespace) -> None: from fastvideo.eval import list_metrics + names = list_metrics() if args.group: prefix = args.group.rstrip(".") + "." @@ -107,38 +192,73 @@ def _cmd_list(args: argparse.Namespace) -> None: def _cmd_run(args: argparse.Namespace) -> None: - from fastvideo.eval import create_evaluator - from fastvideo.eval.io import load_video + from fastvideo.eval import create_evaluator, samples_from, samples_from_manifest - video_paths = _expand_paths(args.videos) - if not video_paths: - raise SystemExit(f"No videos matched: {args.videos}") - ref_paths = _expand_paths([args.reference]) if args.reference else None + if args.manifest: + samples = samples_from_manifest( + args.manifest, + extract_audio=args.extract_audio, + extract_workers=args.extract_workers, + ) + else: + video_paths = _expand_paths(args.videos, _VIDEO_EXTS) if args.videos else None + audio_paths = _expand_paths(args.audios, _AUDIO_EXTS) if args.audios else None + if video_paths is not None and not video_paths: + raise SystemExit(f"No videos matched: {args.videos}") + if audio_paths is not None and not audio_paths: + raise SystemExit(f"No audio files matched: {args.audios}") + generated_count = len(video_paths if video_paths is not None else audio_paths or []) + ref_paths = (_fill_reference_paths(_expand_paths([args.reference], _VIDEO_EXTS), generated_count) + if args.reference else None) + ref_audio_paths = (_fill_reference_paths(_expand_paths(args.reference_audio, _AUDIO_EXTS), generated_count) + if args.reference_audio else None) + text_prompts = args.text_prompt + if text_prompts is not None and len(text_prompts) == 1: + text_prompts = text_prompts * generated_count + samples = samples_from( + video=video_paths, + audio=audio_paths, + reference=ref_paths, + reference_audio=ref_audio_paths, + text_prompts=text_prompts, + fps=args.fps, + extract_audio=args.extract_audio, + extract_workers=args.extract_workers, + ) metrics_arg: list[str] | str = ("all" if args.metrics == "all" else [m.strip() for m in args.metrics.split(",") if m.strip()]) - evaluator = create_evaluator(metrics=metrics_arg, device=args.device) + evaluator = create_evaluator( + metrics=metrics_arg, + device=args.device, + num_gpus=max(1, args.num_gpus), + compile=args.compile, + skip_missing_deps=args.skip_missing_deps, + ) + logger.info("Evaluating %d samples on %d GPU worker(s)", len(samples), evaluator.num_gpus) + results = evaluator.evaluate(samples=samples) + corpus = _serialize_results(results.corpus) + output_value: dict | list + if args.output_format == "full": + sample_results = [{ + **_sample_identity(sample, index), + "scores": _serialize_results(result), + } for index, (sample, result) in enumerate(zip(samples, results, strict=True))] + output_value = { + "samples": sample_results, + "corpus": corpus, + "num_samples": len(samples), + } + else: + if corpus: + logger.warning("Corpus results are omitted in legacy output format; pass --output-format full to keep them") + output_value = [{ + **_legacy_sample_identity(sample), + "scores": _serialize_results(result), + } for sample, result in zip(samples, results, strict=True)] - all_results: list[dict] = [] - for i, vp in enumerate(video_paths): - logger.info("Evaluating %s (%d/%d)", vp, i + 1, len(video_paths)) - kwargs: dict = {"video": load_video(vp)} - if ref_paths is not None: - ref = ref_paths[i] if i < len(ref_paths) else ref_paths[0] - kwargs["reference"] = load_video(ref) - if args.text_prompt is not None: - kwargs["text_prompt"] = (args.text_prompt[i] if i < len(args.text_prompt) else args.text_prompt[0]) - if args.fps is not None: - kwargs["fps"] = args.fps - - results = evaluator.evaluate(**kwargs) - all_results.append({ - "video": str(vp), - "scores": _serialize_results(results), - }) - - payload = json.dumps(all_results, indent=2, default=_jsonable) + payload = json.dumps(output_value, indent=2, default=_jsonable) if args.output: Path(args.output).write_text(payload) logger.info("Wrote results to %s", args.output) @@ -146,13 +266,38 @@ def _cmd_run(args: argparse.Namespace) -> None: print(payload) -def _expand_paths(patterns: list[str]) -> list[str]: +def _cmd_v2a(args: argparse.Namespace) -> None: + """Run an isolated corpus-level backend for generated V2A audio.""" + from fastvideo.eval.v2a import run_av_benchmark + + output = run_av_benchmark( + audio_dir=args.audio_dir, + gt_cache=args.gt_cache, + prediction_cache=args.prediction_cache, + output=args.output, + python_executable=args.python_executable, + device=args.device, + batch_size=args.batch_size, + num_workers=args.num_workers, + audio_length=args.audio_length, + recompute=args.recompute, + skip_video_related=args.skip_video_related, + skip_clap=args.skip_clap, + align_prediction_keys=args.align_prediction_keys, + ) + logger.info("Official av-benchmark results: %s", output) + + +_VIDEO_EXTS = (".mp4", ".avi", ".mov", ".mkv", ".gif", ".webm") +_AUDIO_EXTS = (".wav", ".mp3", ".flac", ".m4a", ".ogg", ".opus") + + +def _expand_paths(patterns: list[str], extensions: tuple[str, ...]) -> list[str]: out: list[str] = [] for pat in patterns: p = Path(pat) if p.is_dir(): - for ext in (".mp4", ".avi", ".mov", ".mkv", ".gif"): - out.extend(sorted(str(f) for f in p.iterdir() if f.suffix.lower() == ext)) + out.extend(sorted(str(f) for f in p.iterdir() if f.suffix.lower() in extensions)) elif any(c in pat for c in "*?["): out.extend(sorted(glob.glob(pat))) else: @@ -167,6 +312,43 @@ def _expand_paths(patterns: list[str]) -> list[str]: return deduped +def _fill_reference_paths(paths: list[str], count: int) -> list[str]: + if not paths: + return paths + if len(paths) >= count: + return paths + return paths + [paths[0]] * (count - len(paths)) + + +def _legacy_sample_identity(sample: dict) -> dict: + """Preserve the original ``fastvideo eval run`` output schema.""" + from fastvideo.eval.types import Video + + video = sample.get("video") + if isinstance(video, Video) and video.source is not None: + return {"video": str(video.source)} + audio = sample.get("audio") + if isinstance(audio, str | Path): + return {"audio": str(audio)} + return {} + + +def _sample_identity(sample: dict, index: int) -> dict: + from fastvideo.eval.types import Video + + identity: dict = {"index": index} + if sample.get("id") is not None: + identity["id"] = sample["id"] + for key in ("video", "audio", "reference", "reference_audio", "text_prompt"): + value = sample.get(key) + if isinstance(value, Video): + if value.source is not None: + identity[key] = str(value.source) + elif isinstance(value, str | Path): + identity[key] = str(value) + return identity + + def _serialize_results(results) -> dict | list: """Turn evaluator output (dict or list-of-dicts) into JSON-friendly form.""" if isinstance(results, list): diff --git a/fastvideo/eval/README.md b/fastvideo/eval/README.md index b618c7065..2dbcf4504 100644 --- a/fastvideo/eval/README.md +++ b/fastvideo/eval/README.md @@ -157,6 +157,13 @@ pairs. Per-sample attachments (`text_prompt(s)`, `fps`, `auxiliary_info`, `extras=`) attach by kwarg; `extract_audio=True` pulls audio tracks off video sources for audio metrics. +Samples whose `video` is a path (or a `Video` without frames) are decoded in +the worker that picks them up. When the caller did not attach `fps`, the pool +probes the container's frame rate and attaches it as `sample["fps"]` so +fps-aware metrics see real timing; an explicit `fps` (for example `--fps` on +`fastvideo eval run`) always wins and skips the probe. Scores recorded before +this probe existed can therefore shift unless `--fps` is passed. + `evaluate` also accepts a pre-loaded `(T, C, H, W)` tensor or a path string under `video` / `reference`. Paths are decoded inside the worker that picks up the sample, so peak memory stays bounded by `num_gpus` @@ -429,6 +436,55 @@ translation); rigid whole-frame motion (e.g. pure camera rotation) is harder. To sweep several baselines into a table, see `examples/inference/eval/eval_third_person_separation.py`. +## Exact V2A corpus evaluation + +`fastvideo eval v2a` is a dedicated corpus-level entry point for +video-to-audio benchmarks. Unlike `fastvideo eval run`, which composes native +per-sample and set metrics through the registry, the V2A command can launch an +upstream benchmark as one protocol-preserving unit. + +The `av-benchmark` backend calls the official `av_bench.extract` and +`av_bench.evaluate` functions. Its dependency stack is intentionally not part +of any standard FastVideo extra. Point the command at a Python interpreter from +an isolated av-benchmark environment: + +```bash +fastvideo eval v2a \ + --backend av-benchmark \ + --audio-dir outputs/v2a/audio \ + --gt-cache /path/to/official/ground-truth-cache \ + --prediction-cache outputs/v2a/av_benchmark_cache \ + --output outputs/v2a/av_benchmark_results.json \ + --python-executable /path/to/av-benchmark/.venv/bin/python \ + --audio-length 8 \ + --batch-size 32 \ + --num-workers 0 +``` + +The backend is model-independent: generated filenames only need to match the +keys in the selected official cache. It can therefore evaluate MMAudio or any +other V2A model. A complete prediction cache is reused automatically; +`--recompute` forces extraction again. + +Some official dataset caches sanitize filenames. For such a cache, +`--align-prediction-keys` atomically remaps only prediction keys having a +unique match after leading underscores are ignored. The source audio is never +renamed, ambiguous keys remain unmatched, and alignment counts are serialized +in the result JSON. + +The official VGGSound cache reports: + +```text +FD-VGG, FD-PANN, FD-PASST +KL-PANNS-softmax, KL-PASST-softmax +ISC-PANNS-mean/std, ISC-PASST-mean/std +IB-Score, DeSync +``` + +This command runs in a subprocess by design. It does not import av-benchmark, +alter CUDA/PyTorch packages, or register metrics in ordinary FastVideo +generation and evaluation processes. + ## Out of scope (follow-up PRs) - **MIND** metrics. Depend on a separate `vipe` upstream submodule. diff --git a/fastvideo/eval/__init__.py b/fastvideo/eval/__init__.py index 6c989bf04..949cbac70 100644 --- a/fastvideo/eval/__init__.py +++ b/fastvideo/eval/__init__.py @@ -26,7 +26,7 @@ from fastvideo.eval.metrics.base import BaseMetric # noqa: E402 from fastvideo.eval.registry import register, list_metrics, get_metric # noqa: E402 from fastvideo.eval.api import evaluate # noqa: E402 from fastvideo.eval.evaluator import Evaluator, create_evaluator # noqa: E402 -from fastvideo.eval.io.inputs import as_video, samples_from # noqa: E402 +from fastvideo.eval.io.inputs import as_video, samples_from, samples_from_manifest # noqa: E402 # Trigger metric auto-discovery import fastvideo.eval.metrics # noqa: F401, E402 @@ -46,4 +46,5 @@ __all__ = [ "get_cache_dir", "as_video", "samples_from", + "samples_from_manifest", ] diff --git a/fastvideo/eval/io/__init__.py b/fastvideo/eval/io/__init__.py index c2c526084..48f4c6d21 100644 --- a/fastvideo/eval/io/__init__.py +++ b/fastvideo/eval/io/__init__.py @@ -1,6 +1,6 @@ from fastvideo.eval.io.audio import NoAudioStreamError, extract_audio_track -from fastvideo.eval.io.inputs import as_video, samples_from -from fastvideo.eval.io.paths import (build_eval_kwargs, default_filename, glob_videos, sanitize_prompt) +from fastvideo.eval.io.inputs import as_video, samples_from, samples_from_manifest +from fastvideo.eval.io.paths import build_eval_kwargs, default_filename, glob_videos, sanitize_prompt from fastvideo.eval.io.video import extract_frames, load_video __all__ = [ @@ -12,6 +12,7 @@ __all__ = [ "build_eval_kwargs", "as_video", "samples_from", + "samples_from_manifest", "extract_audio_track", "NoAudioStreamError", ] diff --git a/fastvideo/eval/io/inputs.py b/fastvideo/eval/io/inputs.py index dc9e54590..95e2aeccf 100644 --- a/fastvideo/eval/io/inputs.py +++ b/fastvideo/eval/io/inputs.py @@ -23,6 +23,7 @@ Design rules: set samples for corpus-shaped metrics (FVD / FAD) without disturbing per-sample paired metrics (LPIPS / PSNR / SSIM / gt_optical_flow). """ + from __future__ import annotations import json @@ -199,6 +200,68 @@ def samples_from( return samples +def samples_from_manifest( + manifest: str | Path, + *, + extract_audio: bool | str | Path = False, + extract_workers: int = 4, +) -> list[dict]: + """Load canonical eval samples from a JSON or JSONL manifest. + + Each row may contain ``video``, ``audio``, ``reference``, + ``reference_audio``, ``reference_audio_source``, ``text_prompt`` and + arbitrary metric metadata. ``reference_audio_source`` is a media file + whose audio track is extracted when ``extract_audio`` is enabled. Paths are + resolved relative to the manifest directory. Video paths are wrapped as + :class:`Video` handles for lazy decoding; all other keys are preserved. + This is the recommended interface for benchmark generation, where matching + by sample id is safer than independently sorting folders. + """ + path = Path(manifest).expanduser().resolve() + if not path.is_file(): + raise FileNotFoundError(path) + if path.suffix == ".jsonl": + rows = [json.loads(line) for line in path.read_text(encoding="utf-8").splitlines() if line.strip()] + elif path.suffix == ".json": + rows = json.loads(path.read_text(encoding="utf-8")) + if not isinstance(rows, list): + raise ValueError(f"{path}: expected a top-level JSON list") + else: + raise ValueError(f"{path}: eval manifest must end in .json or .jsonl") + + samples: list[dict[str, Any]] = [] + for index, row in enumerate(rows): + if not isinstance(row, dict): + raise ValueError(f"{path}: row {index + 1} must be an object, got {type(row).__name__}") + sample = dict(row) + for key in ( + "video", + "reference", + "audio", + "reference_audio", + "reference_audio_source", + ): + value = sample.get(key) + if value is None: + continue + if not isinstance(value, str | Path): + raise ValueError(f"{path}: row {index + 1} key {key!r} must be a path") + media_path = Path(value).expanduser() + if not media_path.is_absolute(): + media_path = path.parent / media_path + media_path = media_path.resolve() + sample[key] = as_video(media_path) if key in {"video", "reference"} else str(media_path) + if "video" not in sample and "audio" not in sample: + raise ValueError(f"{path}: row {index + 1} needs at least video or audio") + samples.append(sample) + + if not samples: + raise ValueError(f"{path}: eval manifest is empty") + if extract_audio: + _attach_extracted_audio(samples, extract_audio, extract_workers) + return samples + + # ------------------------------------------------------------------ # Internal helpers # ------------------------------------------------------------------ @@ -280,12 +343,18 @@ def _attach_extracted_audio( # Build the work list — (sample_idx, dest_audio_key, source_path) work: list[tuple[int, str, str]] = [] for i, s in enumerate(samples): - for vkey, akey in (("video", "audio"), ("reference", "reference_audio")): + for vkey, akey in ( + ("video", "audio"), + ("reference", "reference_audio"), + ("reference_audio_source", "reference_audio"), + ): if akey in s: continue v = s.get(vkey) if isinstance(v, Video) and v.source is not None: work.append((i, akey, str(v.source))) + elif vkey == "reference_audio_source" and isinstance(v, str | Path): + work.append((i, akey, str(v))) if not work: return diff --git a/fastvideo/eval/io/video.py b/fastvideo/eval/io/video.py index 134770f32..7d4d2205c 100644 --- a/fastvideo/eval/io/video.py +++ b/fastvideo/eval/io/video.py @@ -40,6 +40,29 @@ def extract_frames(video: torch.Tensor, n_frames: int | None = None) -> torch.Te return video[indices] +def probe_video_fps(source: str | Path) -> float | None: + """Read a path-backed video's native frame rate without decoding frames.""" + try: + import av + + with av.open(str(source)) as container: + if not container.streams.video: + return None + stream = container.streams.video[0] + rate = stream.average_rate or stream.guessed_rate + if rate is not None and float(rate) > 0: + return float(rate) + except (ImportError, OSError, ValueError): + pass + try: + from decord import VideoReader, cpu + + rate = float(VideoReader(str(source), ctx=cpu(0)).get_avg_fps()) + return rate if rate > 0 else None + except (ImportError, OSError, RuntimeError, ValueError): + return None + + # --- Internal helpers --- diff --git a/fastvideo/eval/metrics/audio/frechet_distance/metric.py b/fastvideo/eval/metrics/audio/frechet_distance/metric.py index aa091dd7d..e5ac6d80a 100644 --- a/fastvideo/eval/metrics/audio/frechet_distance/metric.py +++ b/fastvideo/eval/metrics/audio/frechet_distance/metric.py @@ -47,7 +47,13 @@ def _frechet_distance(mu1: np.ndarray, sigma2 = sigma2 + eps * np.eye(sigma2.shape[0]) diff = mu1 - mu2 - covmean, _ = linalg.sqrtm(sigma1.dot(sigma2), disp=False) + covariance_product = sigma1.dot(sigma2) + try: + # SciPy < 1.18 returns ``(root, error_estimate)`` with disp=False. + covmean, _ = linalg.sqrtm(covariance_product, disp=False) + except TypeError: + # SciPy 1.18 removed ``disp`` and now returns only the matrix root. + covmean = linalg.sqrtm(covariance_product) if np.iscomplexobj(covmean): if not np.allclose(np.diagonal(covmean).imag, 0, atol=1e-3): raise ValueError(f"imaginary component {np.max(np.abs(covmean.imag))}") @@ -108,6 +114,7 @@ class FrechetAudioDistanceMetric(BaseMetric): def setup(self) -> None: if self._model is None: from hear21passt.base import get_basic_model + model = get_basic_model(mode="all") model.eval() model = model.to(self.device) @@ -127,13 +134,11 @@ class FrechetAudioDistanceMetric(BaseMetric): ref = ref.numpy() ref = np.asarray(ref) if ref.ndim != 2 or ref.shape[1] != PASST_EMBED_DIM: - raise ValueError(f"Expected a 2-D tensor of shape (N, {PASST_EMBED_DIM}) " - f"at {path}; got shape {ref.shape}.") + raise ValueError(f"Expected a 2-D tensor of shape (N, {PASST_EMBED_DIM}) at {path}; got shape {ref.shape}.") finite = np.isfinite(ref).all(axis=1) ref = ref[finite] if ref.shape[0] < 2: - raise ValueError(f"Cached reference features at {path} have only " - f"{ref.shape[0]} finite rows; need >= 2.") + raise ValueError(f"Cached reference features at {path} have only {ref.shape[0]} finite rows; need >= 2.") self._cached_ref_mu = ref.mean(axis=0) self._cached_ref_sigma = np.cov(ref, rowvar=False) self._n_cached_ref = int(ref.shape[0]) @@ -209,8 +214,7 @@ class FrechetAudioDistanceMetric(BaseMetric): name=self.name, score=None, details={ - "skipped": f"FAD needs >=2 finite-embed gen samples " - f"(got {n_gen} valid of {len(self._gen_buf)})", + "skipped": f"FAD needs >=2 finite-embed gen samples (got {n_gen} valid of {len(self._gen_buf)})", "n_gen": n_gen, "n_ref": n_ref, "n_gen_dropped_nonfinite": n_gen_dropped, diff --git a/fastvideo/eval/pool.py b/fastvideo/eval/pool.py index 0faef8677..28be2fa7c 100644 --- a/fastvideo/eval/pool.py +++ b/fastvideo/eval/pool.py @@ -9,6 +9,7 @@ write back into a result list in input order. Pool sizing: ``max_size = prefetch_factor * num_workers``. """ + from __future__ import annotations import queue @@ -144,11 +145,16 @@ class VideoPool: :func:`fastvideo.eval.io.as_video` / :func:`samples_from`) to make them decode-eligible. """ - from fastvideo.eval.io.video import load_video + from fastvideo.eval.io.video import load_video, probe_video_fps out = dict(sample) for key, val in sample.items(): if isinstance(val, Video) and val.frames is None and val.source is not None: + if key == "video" and "fps" not in out: + if val.fps is None: + val.fps = probe_video_fps(val.source) + if val.fps is not None: + out["fps"] = float(val.fps) val.frames = load_video(val.source) out[key] = val return out diff --git a/fastvideo/eval/v2a/__init__.py b/fastvideo/eval/v2a/__init__.py new file mode 100644 index 000000000..67fcb7881 --- /dev/null +++ b/fastvideo/eval/v2a/__init__.py @@ -0,0 +1,5 @@ +"""Video-to-audio corpus evaluation backends.""" + +from fastvideo.eval.v2a.av_benchmark import run_av_benchmark + +__all__ = ["run_av_benchmark"] diff --git a/fastvideo/eval/v2a/_av_benchmark_runner.py b/fastvideo/eval/v2a/_av_benchmark_runner.py new file mode 100644 index 000000000..c9de64847 --- /dev/null +++ b/fastvideo/eval/v2a/_av_benchmark_runner.py @@ -0,0 +1,265 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Standalone process runner for the official av-benchmark API. + +This file is executed directly by an interpreter from an isolated environment. +Keep it free of ``fastvideo`` imports so the benchmark environment does not +need FastVideo's runtime dependency stack. +""" + +from __future__ import annotations + +import argparse +import json +from importlib import metadata +from pathlib import Path +from typing import Any +from collections.abc import Callable + +_BASE_CACHE_FILES = ( + "pann_features.pth", + "vggish_features.pth", + "passt_features_embed.pth", + "passt_logits.pth", +) +_VIDEO_CACHE_FILES = ( + "imagebind_audio.pth", + "synchformer_audio.pth", +) +_CLAP_CACHE_FILES = ( + "clap_laion_audio.pth", + "clap_ms_audio.pth", +) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--audio-dir", type=Path, required=True) + parser.add_argument("--gt-cache", type=Path, required=True) + parser.add_argument("--prediction-cache", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--device", default="cuda") + parser.add_argument("--batch-size", type=int, default=32) + parser.add_argument("--num-workers", type=int, default=0) + parser.add_argument("--audio-length", type=float, default=8.0) + parser.add_argument("--recompute", action="store_true") + parser.add_argument("--skip-video-related", action="store_true") + parser.add_argument("--skip-clap", action="store_true") + parser.add_argument("--align-prediction-keys", action="store_true") + return parser.parse_args() + + +def _load_backend() -> tuple[Callable[..., Any], Callable[..., Any]]: + try: + from av_bench.evaluate import evaluate + from av_bench.extract import extract + except ImportError as error: + raise ImportError("Official av_bench is not installed in this Python environment. " + "Install https://github.com/hkchengrex/av-benchmark and pass its " + "Python through --python-executable.") from error + return extract, evaluate + + +def _required_prediction_files(*, skip_video_related: bool, skip_clap: bool) -> tuple[str, ...]: + required = list(_BASE_CACHE_FILES) + if not skip_video_related: + required.extend(_VIDEO_CACHE_FILES) + if not skip_clap: + required.extend(_CLAP_CACHE_FILES) + return tuple(required) + + +def _validate_gt_cache(gt_cache: Path) -> None: + missing = [name for name in _BASE_CACHE_FILES if not (gt_cache / name).is_file()] + if missing: + raise FileNotFoundError(f"Ground-truth cache {gt_cache} is missing required files: {missing}") + + +def _jsonable(value: Any) -> Any: + if isinstance(value, Path): + return str(value) + if hasattr(value, "detach"): + value = value.detach().cpu() + if hasattr(value, "item"): + try: + return value.item() + except ValueError: + pass + if hasattr(value, "tolist"): + return value.tolist() + raise TypeError(f"Object of type {type(value).__name__} is not JSON serializable") + + +def _align_prediction_cache_keys( + *, + gt_cache: Path, + prediction_cache: Path, + required_files: tuple[str, ...], +) -> dict[str, int]: + """Align filename-sanitized prediction keys to unique GT cache keys. + + Official VGGSound caches prefix one or more underscores to some stems. + Generated audio keeps the original stem, so paired metrics otherwise drop + those samples silently. Only unique leading-underscore matches are changed. + """ + import torch + + gt_logits = torch.load(gt_cache / "passt_logits.pth", map_location="cpu", weights_only=True) + pred_logits = torch.load(prediction_cache / "passt_logits.pth", map_location="cpu", weights_only=True) + if not isinstance(gt_logits, dict) or not isinstance(pred_logits, dict): + raise TypeError("PaSST caches must be dictionaries keyed by sample id") + + gt_keys = {str(key) for key in gt_logits} + pred_keys = {str(key) for key in pred_logits} + candidates_by_stem: dict[str, list[str]] = {} + for key in gt_keys: + candidates_by_stem.setdefault(key.lstrip("_"), []).append(key) + + occupied = pred_keys & gt_keys + mapping: dict[str, str] = {} + ambiguous = 0 + for key in sorted(pred_keys - gt_keys): + candidates = [ + candidate for candidate in candidates_by_stem.get(key.lstrip("_"), []) if candidate not in occupied + ] + if len(candidates) == 1: + mapping[key] = candidates[0] + occupied.add(candidates[0]) + elif len(candidates) > 1: + ambiguous += 1 + + if mapping: + for filename in required_files: + path = prediction_cache / filename + values = torch.load(path, map_location="cpu", weights_only=True) + if not isinstance(values, dict): + raise TypeError(f"Prediction cache must be a dictionary: {path}") + aligned = {mapping.get(str(key), str(key)): value for key, value in values.items()} + if len(aligned) != len(values): + raise RuntimeError(f"Prediction key alignment created a collision in {path}") + temporary = path.with_suffix(path.suffix + ".tmp") + torch.save(aligned, temporary) + temporary.replace(path) + + return { + "exact_before_alignment": len(pred_keys & gt_keys), + "remapped": len(mapping), + "ambiguous": ambiguous, + "unmatched_after_alignment": len(pred_keys) - len(pred_keys & gt_keys) - len(mapping), + } + + +def run_benchmark( + *, + audio_dir: Path, + gt_cache: Path, + prediction_cache: Path, + output: Path, + device: str, + batch_size: int, + num_workers: int, + audio_length: float, + recompute: bool, + skip_video_related: bool, + skip_clap: bool, + align_prediction_keys: bool = False, +) -> dict[str, Any]: + audio_dir = audio_dir.expanduser().resolve() + gt_cache = gt_cache.expanduser().resolve() + prediction_cache = prediction_cache.expanduser().resolve() + output = output.expanduser().resolve() + if not audio_dir.is_dir(): + raise FileNotFoundError(f"Generated audio directory does not exist: {audio_dir}") + if not gt_cache.is_dir(): + raise FileNotFoundError(f"Ground-truth cache directory does not exist: {gt_cache}") + if batch_size <= 0: + raise ValueError(f"batch_size must be positive, got {batch_size}") + if num_workers < 0: + raise ValueError(f"num_workers must be non-negative, got {num_workers}") + if audio_length <= 0: + raise ValueError(f"audio_length must be positive, got {audio_length}") + + audio_files = sorted((*audio_dir.glob("*.wav"), *audio_dir.glob("*.flac"))) + if not audio_files: + raise FileNotFoundError(f"No WAV or FLAC files found in {audio_dir}") + _validate_gt_cache(gt_cache) + extract, evaluate = _load_backend() + + required = _required_prediction_files( + skip_video_related=skip_video_related, + skip_clap=skip_clap, + ) + cache_complete = all((prediction_cache / name).is_file() for name in required) + extraction_ran = recompute or not cache_complete + if extraction_ran: + prediction_cache.mkdir(parents=True, exist_ok=True) + extract( + audio_path=audio_dir, + output_path=prediction_cache, + device=device, + batch_size=batch_size, + num_workers=num_workers, + audio_length=audio_length, + skip_video_related=skip_video_related, + skip_clap=skip_clap, + ) + + key_alignment = None + if align_prediction_keys: + key_alignment = _align_prediction_cache_keys( + gt_cache=gt_cache, + prediction_cache=prediction_cache, + required_files=required, + ) + + metrics = evaluate( + gt_audio_cache=gt_cache, + pred_audio_cache=prediction_cache, + skip_video_related=skip_video_related, + skip_clap=skip_clap, + ) + try: + backend_version = metadata.version("av_bench") + except metadata.PackageNotFoundError: + backend_version = "editable-or-unknown" + payload = { + "backend": "hkchengrex/av-benchmark", + "backend_version": backend_version, + "audio_dir": str(audio_dir), + "gt_cache": str(gt_cache), + "prediction_cache": str(prediction_cache), + "num_generated_audio": len(audio_files), + "audio_length": audio_length, + "extraction_ran": extraction_ran, + "skip_video_related": skip_video_related, + "skip_clap": skip_clap, + "prediction_key_alignment": key_alignment, + "metrics": metrics, + } + output.parent.mkdir(parents=True, exist_ok=True) + temporary = output.with_suffix(output.suffix + ".tmp") + temporary.write_text(json.dumps(payload, indent=2, sort_keys=True, default=_jsonable) + "\n", encoding="utf-8") + temporary.replace(output) + return payload + + +def main() -> None: + args = parse_args() + payload = run_benchmark( + audio_dir=args.audio_dir, + gt_cache=args.gt_cache, + prediction_cache=args.prediction_cache, + output=args.output, + device=args.device, + batch_size=args.batch_size, + num_workers=args.num_workers, + audio_length=args.audio_length, + recompute=args.recompute, + skip_video_related=args.skip_video_related, + skip_clap=args.skip_clap, + align_prediction_keys=args.align_prediction_keys, + ) + print(json.dumps(payload, indent=2, sort_keys=True, default=_jsonable)) + + +if __name__ == "__main__": + main() diff --git a/fastvideo/eval/v2a/av_benchmark.py b/fastvideo/eval/v2a/av_benchmark.py new file mode 100644 index 000000000..8455c76ec --- /dev/null +++ b/fastvideo/eval/v2a/av_benchmark.py @@ -0,0 +1,97 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Isolated launcher for the official ``hkchengrex/av-benchmark`` backend. + +The benchmark has a large dependency stack that can conflict with FastVideo's +runtime. This module intentionally does not import ``av_bench``. It launches a +small runner with a caller-selected Python interpreter, so exact official +metrics can live in a separate virtual environment. +""" + +from __future__ import annotations + +import subprocess +import sys +from pathlib import Path + +from fastvideo.logger import init_logger + +logger = init_logger(__name__) + + +def run_av_benchmark( + *, + audio_dir: str | Path, + gt_cache: str | Path, + prediction_cache: str | Path | None = None, + output: str | Path | None = None, + python_executable: str | Path | None = None, + device: str = "cuda", + batch_size: int = 32, + num_workers: int = 0, + audio_length: float = 8.0, + recompute: bool = False, + skip_video_related: bool = False, + skip_clap: bool = False, + align_prediction_keys: bool = False, +) -> Path: + """Extract prediction features and evaluate with official av-benchmark. + + ``python_executable`` should point to an isolated environment containing + the official ``av_bench`` package. Omitting it uses the current process, + which is useful only when that package is already installed. + """ + audio_path = Path(audio_dir).expanduser().resolve() + gt_path = Path(gt_cache).expanduser().resolve() + if not audio_path.is_dir(): + raise FileNotFoundError(f"Generated audio directory does not exist: {audio_path}") + if not gt_path.is_dir(): + raise FileNotFoundError(f"Ground-truth cache directory does not exist: {gt_path}") + if batch_size <= 0: + raise ValueError(f"batch_size must be positive, got {batch_size}") + if num_workers < 0: + raise ValueError(f"num_workers must be non-negative, got {num_workers}") + if audio_length <= 0: + raise ValueError(f"audio_length must be positive, got {audio_length}") + + root = audio_path.parent + prediction_path = (Path(prediction_cache).expanduser().resolve() if prediction_cache is not None else root / + "av_benchmark_cache") + output_path = (Path(output).expanduser().resolve() if output is not None else root / "av_benchmark_results.json") + runner = Path(__file__).with_name("_av_benchmark_runner.py") + python = str(Path(python_executable).expanduser()) if python_executable is not None else sys.executable + + command = [ + python, + str(runner), + "--audio-dir", + str(audio_path), + "--gt-cache", + str(gt_path), + "--prediction-cache", + str(prediction_path), + "--output", + str(output_path), + "--device", + device, + "--batch-size", + str(batch_size), + "--num-workers", + str(num_workers), + "--audio-length", + str(audio_length), + ] + if recompute: + command.append("--recompute") + if skip_video_related: + command.append("--skip-video-related") + if skip_clap: + command.append("--skip-clap") + if align_prediction_keys: + command.append("--align-prediction-keys") + + logger.info("Launching official av-benchmark with isolated Python: %s", python) + logger.info("Prediction cache: %s", prediction_path) + subprocess.run(command, check=True) + if not output_path.is_file(): + raise RuntimeError(f"av-benchmark completed without writing expected output: {output_path}") + return output_path diff --git a/fastvideo/models/audio/__init__.py b/fastvideo/models/audio/__init__.py index 9ba265f8e..57be9eaa7 100644 --- a/fastvideo/models/audio/__init__.py +++ b/fastvideo/models/audio/__init__.py @@ -5,5 +5,15 @@ from fastvideo.models.audio.ltx2_audio_vae import ( LTX2AudioEncoder, LTX2Vocoder, ) +from fastvideo.models.audio.mmaudio_processing import ( + MMAudioMelConverter, + build_mmaudio_mel_converter, +) -__all__ = ["LTX2AudioEncoder", "LTX2AudioDecoder", "LTX2Vocoder"] +__all__ = [ + "LTX2AudioEncoder", + "LTX2AudioDecoder", + "LTX2Vocoder", + "MMAudioMelConverter", + "build_mmaudio_mel_converter", +] diff --git a/fastvideo/models/audio/mmaudio_processing.py b/fastvideo/models/audio/mmaudio_processing.py new file mode 100644 index 000000000..44510d94c --- /dev/null +++ b/fastvideo/models/audio/mmaudio_processing.py @@ -0,0 +1,103 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Audio preprocessing used by native MMAudio training feature extraction.""" + +from __future__ import annotations + +from typing import Literal + +import torch +import torch.nn as nn + + +class MMAudioMelConverter(nn.Module): + """Convert waveforms to MMAudio's log-mel representation. + + The FFT, padding, mel filter, and logarithm match the published MMAudio + preprocessing contract. Keeping this component in FastVideo avoids a + runtime import from the upstream training repository. + """ + + def __init__( + self, + *, + sampling_rate: int, + n_fft: int, + num_mels: int, + hop_size: int, + win_size: int, + fmin: float, + fmax: float, + log_base: Literal["e", "10"], + ) -> None: + super().__init__() + from librosa.filters import mel as librosa_mel_fn + + mel = librosa_mel_fn( + sr=sampling_rate, + n_fft=n_fft, + n_mels=num_mels, + fmin=fmin, + fmax=fmax, + ) + self.n_fft = n_fft + self.hop_size = hop_size + self.win_size = win_size + self.log_base = log_base + self.register_buffer("mel_basis", torch.from_numpy(mel).float()) + self.register_buffer("hann_window", torch.hann_window(win_size)) + + def forward(self, waveform: torch.Tensor, center: bool = False) -> torch.Tensor: + if waveform.ndim != 2: + raise ValueError(f"MMAudio waveforms must have shape [batch, samples], got {tuple(waveform.shape)}") + waveform = waveform.clamp(min=-1.0, max=1.0).to(self.mel_basis.device) + padding = int((self.n_fft - self.hop_size) / 2) + waveform = torch.nn.functional.pad( + waveform.unsqueeze(1), + [padding, padding], + mode="reflect", + ).squeeze(1) + spectrum = torch.stft( + waveform, + self.n_fft, + hop_length=self.hop_size, + win_length=self.win_size, + window=self.hann_window, + center=center, + pad_mode="reflect", + normalized=False, + onesided=True, + return_complex=True, + ) + magnitude = torch.sqrt(torch.view_as_real(spectrum).pow(2).sum(-1) + 1e-9) + mel = torch.matmul(self.mel_basis, magnitude) + mel = torch.clamp(mel, min=1e-5) + return torch.log10(mel) if self.log_base == "10" else torch.log(mel) + + +def build_mmaudio_mel_converter(mode: Literal["16k", "44k"] = "44k", ) -> MMAudioMelConverter: + if mode == "16k": + return MMAudioMelConverter( + sampling_rate=16_000, + n_fft=1024, + num_mels=80, + hop_size=256, + win_size=1024, + fmin=0, + fmax=8_000, + log_base="10", + ) + if mode == "44k": + return MMAudioMelConverter( + sampling_rate=44_100, + n_fft=2048, + num_mels=128, + hop_size=512, + win_size=2048, + fmin=0, + fmax=44_100 / 2, + log_base="e", + ) + raise ValueError(f"Unknown MMAudio mel mode: {mode}") + + +__all__ = ["MMAudioMelConverter", "build_mmaudio_mel_converter"] diff --git a/fastvideo/models/dits/mmaudio.py b/fastvideo/models/dits/mmaudio.py index 5b5e9fa3d..8b2c629e6 100644 --- a/fastvideo/models/dits/mmaudio.py +++ b/fastvideo/models/dits/mmaudio.py @@ -10,6 +10,7 @@ separately verified optimization. from __future__ import annotations +import math from dataclasses import dataclass from typing import Any @@ -282,7 +283,15 @@ class MMAudioTransformer(BaseDiT): 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: + def __init__( + self, + config: MMAudioTransformerConfig, + hf_config: dict[str, Any], + latent_mean: torch.Tensor | None = None, + latent_std: torch.Tensor | None = None, + empty_string_feat: torch.Tensor | None = None, + **kwargs, + ) -> None: del kwargs super().__init__(config=config, hf_config=hf_config) arch = config.arch_config @@ -348,9 +357,57 @@ class MMAudioTransformer(BaseDiT): ] ) - 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) + def fixed_tensor( + value: torch.Tensor | None, + shape: tuple[int, ...], + *, + fill_value: float, + name: str, + ) -> torch.Tensor: + if value is None: + return torch.full( + shape, + fill_value, + device=self.sync_pos_emb.device, + dtype=self.sync_pos_emb.dtype, + ) + if value.numel() != math.prod(shape): + raise ValueError( + f"MMAudio {name} must have {math.prod(shape)} values, " + f"got {value.numel()}." + ) + return value.to( + device=self.sync_pos_emb.device, + dtype=self.sync_pos_emb.dtype, + ).reshape(shape).clone() + + self.latent_mean = nn.Parameter( + fixed_tensor( + latent_mean, + (1, 1, arch.latent_dim), + fill_value=float("nan"), + name="latent_mean", + ), + requires_grad=False, + ) + self.latent_std = nn.Parameter( + fixed_tensor( + latent_std, + (1, 1, arch.latent_dim), + fill_value=float("nan"), + name="latent_std", + ), + requires_grad=False, + ) + self.empty_string_feat = nn.Parameter( + fixed_tensor( + empty_string_feat, + (arch.text_seq_len, arch.text_dim), + fill_value=0.0, + name="empty_string_feat", + ), + 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) @@ -421,7 +478,12 @@ class MMAudioTransformer(BaseDiT): self._latent_seq_len = latent_seq_len self._clip_seq_len = clip_seq_len self._sync_seq_len = sync_seq_len - self.initialize_rotations() + # This method is also called by inference pipeline stages. Persistent + # buffers created under inference_mode become inference tensors and + # cannot later be saved for backward if a live transformer resumes + # training. Build the derived buffers as ordinary no-grad tensors. + with torch.inference_mode(False), torch.no_grad(): + self.initialize_rotations() def initialize_weights(self) -> None: diff --git a/fastvideo/models/loader/component_loader.py b/fastvideo/models/loader/component_loader.py index 642d4cd2e..9ace6f93f 100644 --- a/fastvideo/models/loader/component_loader.py +++ b/fastvideo/models/loader/component_loader.py @@ -1066,7 +1066,15 @@ class AudioDecoderLoader(ComponentLoader): audio_vae.load_state_dict(loaded, strict=True) return audio_vae.eval() - precision = getattr(fastvideo_args.pipeline_config, "audio_decoder_precision", "bf16") + # A preprocessing component may carry the MMAudio encoder in addition + # to the inference decoder. Feature extraction is defined in fp32 by + # the published recipe; do not quantize it to decoder inference dtype. + needs_encoder = bool(config.get("need_encoder", False)) + precision_name = ( + "audio_encoder_precision" if needs_encoder else + "audio_decoder_precision") + precision = getattr(fastvideo_args.pipeline_config, precision_name, + "fp32" if needs_encoder else "bf16") # 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 diff --git a/fastvideo/models/loader/fsdp_load.py b/fastvideo/models/loader/fsdp_load.py index ed6e59adf..c9849cdb2 100644 --- a/fastvideo/models/loader/fsdp_load.py +++ b/fastvideo/models/loader/fsdp_load.py @@ -494,6 +494,75 @@ def maybe_load_fsdp_model( return model +def build_fsdp_model_from_scratch( + model_cls: type[nn.Module], + init_params: dict[str, Any], + device: torch.device, + hsdp_replicate_dim: int, + hsdp_shard_dim: int, + default_dtype: torch.dtype, + param_dtype: torch.dtype, + reduce_dtype: torch.dtype, + *, + seed: int, + cpu_offload: bool = False, + output_dtype: torch.dtype | None = None, + pin_cpu_memory: bool = True, +) -> torch.nn.Module: + """Deterministically initialize and FSDP-shard a model without weights.""" + + mp_policy = MixedPrecisionPolicy( + param_dtype, + reduce_dtype, + output_dtype, + cast_forward_inputs=False, + ) + set_mixed_precision_policy( + param_dtype=param_dtype, + reduce_dtype=reduce_dtype, + output_dtype=output_dtype, + mp_policy=mp_policy, + ) + + fork_devices: list[int] = [] + if device.type == "cuda": + fork_devices = [ + device.index if device.index is not None else torch.cuda.current_device() + ] + with torch.random.fork_rng(devices=fork_devices): + torch.manual_seed(int(seed)) + with set_default_dtype(default_dtype), torch.device(device): + model = model_cls(**init_params) + + from fastvideo.platforms import current_platform + + if not current_platform.is_mps(): + pin_cpu_memory = pin_cpu_memory and is_pin_memory_available() + if current_platform.is_npu(): + device_mesh = init_device_mesh( + "npu", + mesh_shape=(hsdp_replicate_dim, hsdp_shard_dim), + mesh_dim_names=("replicate", "shard"), + ) + else: + device_mesh = init_device_mesh( + "cuda", + mesh_shape=(hsdp_replicate_dim, hsdp_shard_dim), + mesh_dim_names=("replicate", "shard"), + ) + shard_model( + model, + cpu_offload=cpu_offload, + reshard_after_forward=True, + mp_policy=mp_policy, + mesh=device_mesh, + fsdp_shard_conditions=model._fsdp_shard_conditions, + pin_cpu_memory=pin_cpu_memory, + ) + + for parameter in model.parameters(): + parameter.requires_grad = False + return model def _regional_compile_unsupported_reason( init_params: dict[str, Any], *, diff --git a/fastvideo/pipelines/basic/mmaudio/README.md b/fastvideo/pipelines/basic/mmaudio/README.md index 7cbc81367..059ce4896 100644 --- a/fastvideo/pipelines/basic/mmaudio/README.md +++ b/fastvideo/pipelines/basic/mmaudio/README.md @@ -1,8 +1,9 @@ -# MMAudio Video-to-Audio Inference +# MMAudio Video-to-Audio 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. +video-to-audio (V2A) and text-to-audio (T2A) generation. All five published +MMAudio model architectures are registered independently so their sequence +length, latent width, sample rate, VAE, and vocoder cannot be mixed silently. The production pipeline is implemented with FastVideo components and does not import the upstream `mmaudio` Python package. It loads the transformer, DFN5B @@ -10,20 +11,41 @@ CLIP text/vision encoders, Synchformer visual encoder, audio VAE, BigVGAN-v2, tokenizer, and scheduler from a standard FastVideo/Diffusers-style checkpoint directory. +## Official model variants + +| Variant | Audio | Hidden / layers / heads | Inference | Official training recipe | +| --- | ---: | --- | --- | --- | +| `small_16k` | 16 kHz | 448 / 12 / 7 | implemented | supported | +| `small_44k` | 44.1 kHz | 448 / 12 / 7 | implemented | supported | +| `medium_44k` | 44.1 kHz | 896 / 12 / 14 | implemented | supported | +| `large_44k` | 44.1 kHz | 896 / 21 / 14 | implemented | supported | +| `large_44k_v2` | 44.1 kHz | 896 / 21 / 14 | implemented and validated | not supported upstream | + +`small_16k` uses latent shape `[250, 20]`, the 16 kHz VAE, and the original +80-mel BigVGAN. The 44.1 kHz variants use latent shape `[345, 40]`, the 44 kHz +VAE, and 128-mel BigVGAN-v2. `_v2` changes transformer activations and timestep +conditioning but not the large model dimensions. + ## Status -- `large_44k_v2` V2A inference: supported +- Five variant-specific presets, pipeline configs, registry entries, and + converter choices: supported - T2A pipeline routing: supported; the provided example currently targets V2A -- Output: mono WAV, 44.1 kHz +- Output: mono WAV at 16 kHz or 44.1 kHz, according to the checkpoint - Default inference duration: 8 seconds - Variable-duration inference: supported - Single-GPU inference with FastVideo's default offloading: supported +- `small_16k`, `small_44k`, `medium_44k`, and `large_44k` from-scratch training: + supported through FastVideo's modular trainer and official-compatible features - Source-video/audio muxing: not yet part of the pipeline output -- Training and small/medium/16 kHz variants: not included in this port +- `large_44k_v2` training: intentionally rejected because the upstream recipe + does not support training the v2 checkpoint 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. +test with `large_44k_v2`. The other four variants have architecture/config +coverage; real-weight waveform parity still requires converting their official +checkpoints. ## Requirements @@ -49,6 +71,12 @@ source .venv/bin/activate UV_TORCH_BACKEND=cu126 uv pip install -e . ``` +Training additionally needs TensorDict for the memory-mapped feature cache: + +```bash +UV_TORCH_BACKEND=cu126 uv pip install -e '.[mmaudio-train]' +``` + Use `UV_TORCH_BACKEND=cu130` instead on CUDA 13. Conda is not required. ## Checkpoint Layout @@ -149,6 +177,352 @@ python scripts/checkpoint_conversion/convert_mmaudio_to_diffusers.py \ The converted checkpoint is approximately 9 GB. +## Train MMAudio v1 + +Training uses the native `fastvideo/train/` stack. FastVideo owns the training +loop, HSDP/FSDP wrapping, optimizer and LR scheduler, gradient accumulation, +gradient clipping, tracking, and distributed checkpoints. The reusable +`FlowMatchingFineTuneMethod` implements shape-agnostic velocity supervision; +`MMAudioModel` supplies MMAudio's audio posterior sampling, latent +normalization, independent video/text CFG dropout, and multimodal forward. + +The first smoke-test target should be `small_44k`. The upstream documentation +does not support training checkpoints ending in `_v2`, so the inference +checkpoint `large_44k_v2` must not be used here. + +### Convert a trainable v1 transformer + +Download the official v1 checkpoint: + +```bash +mkdir -p official_weights/mmaudio/raw/weights +curl -L --continue-at - \ + https://huggingface.co/hkchengrex/MMAudio/resolve/main/weights/mmaudio_small_44k.pth \ + -o official_weights/mmaudio/raw/weights/mmaudio_small_44k.pth +``` + +Training only loads the transformer, so it does not need another copy of the +VAE, DFN5B, Synchformer, or vocoder components: + +```bash +python scripts/checkpoint_conversion/convert_mmaudio_to_diffusers.py \ + --variant small_44k \ + --transformer-only \ + --transformer-checkpoint official_weights/mmaudio/raw/weights/mmaudio_small_44k.pth \ + --output converted_weights/mmaudio/small_44k +``` + +Replace `small_44k` with `medium_44k` or `large_44k` in both arguments to train +a larger v1 model. `small_16k` is also supported, but requires 16 kHz features. + +### Train the DiT from scratch + +From scratch means all trainable MMAudio transformer weights are randomly +initialized. The VAE, DFN5B CLIP, and Synchformer remain pretrained feature +extractors, as in the official recipe. The fixed CLIP embedding for an empty +string is loaded either from upstream `empty_string.pth` or a converted official +checkpoint, but is not trainable and does not initialize any DiT layer. To avoid +downloading a DiT checkpoint for a pure scratch run: + +```bash +curl -L --continue-at - \ + https://github.com/hkchengrex/MMAudio/releases/download/v0.1/empty_string.pth \ + -o official_weights/mmaudio/raw/ext_weights/empty_string.pth +``` + +Four configs mirror the official v1 architectures: + +```text +examples/train/configs/fine_tuning/mmaudio/small_16k_from_scratch.yaml +examples/train/configs/fine_tuning/mmaudio/small_44k_from_scratch.yaml +examples/train/configs/fine_tuning/mmaudio/medium_44k_from_scratch.yaml +examples/train/configs/fine_tuning/mmaudio/large_44k_from_scratch.yaml +``` + +They use AdamW (`1e-4`, betas `0.9/0.95`, epsilon `1e-6`, weight decay +`1e-6`), 1,000 warmup steps, MultiStepLR drops at 240k and 270k, gradient clip +1.0, and 300k steps. The FSDP-aware `PostHocEMACallback` mirrors upstream +MMAudio's `nitrous_ema` settings: sigma profiles `[0.05, 0.1]`, an update every +step, snapshots every 5,000 steps, and final synthesis at sigma `0.05`. It is an +additional callback; FastVideo's existing conventional EMA remains unchanged +for other models. + +Edit dataset paths, variant, batch size, learning rate, validation intervals, +output directory, resume path, and tracker at the top of the launcher. No +shell environment variables or command-line overrides are required. The +launcher also downloads the official standalone `empty_string.pth` if it is +missing. Then run the current 44.1 kHz VGGSound cache on four GPUs: + +```bash +bash examples/training/finetune/mmaudio/run_train_vggsound_from_scratch.sh +``` + +### Optional full-replica DDP for `small_44k` + +MMAudio's upstream runner uses native PyTorch DDP with one complete model per +GPU and `broadcast_buffers=False`. FastVideo exposes that behavior as the +opt-in `training.distributed.strategy: ddp`; the default remains `fsdp`, so +existing model configs and the launcher above are unchanged. The initial DDP +plugin is intentionally limited to MMAudio from-scratch training. + +The independent four-GPU config and launcher are: + +```text +examples/train/configs/fine_tuning/mmaudio/small_44k_ddp_from_scratch.yaml +examples/training/finetune/mmaudio/run_train_vggsound_ddp_from_scratch.sh +``` + +Run it only after the current GPU job has finished: + +```bash +bash examples/training/finetune/mmaudio/run_train_vggsound_ddp_from_scratch.sh +``` + +The script uses port `29502` and writes to +`outputs/mmaudio_small_44k_ddp_from_scratch`, so it does not reuse the FSDP +job's rendezvous or output directory. DDP can be faster for this small model +when every GPU can hold the full model, optimizer, activations, and rank 0's +two EMA copies. The DDP config also enables upstream's fused AdamW. It uses +more per-GPU model memory than FSDP. + +The DDP config also enables `training.model.compile_train_fn: true`. +Following upstream MMAudio, compilation happens after DDP wrapping. For +PyTorch 2.12 compatibility, FastVideo keeps explicit-Generator posterior, +flow-time, prior, and CFG-mask sampling in eager code, applies masks with +broadcasted `torch.where`, and compiles the inner transformer's tensor-only +forward with `fullgraph: true`. The outer DDP reducer/control flow and the +lightweight MSE loss remain eager. This preserves the official RNG order while +making any graph break inside the expensive transformer graph a hard error +instead of a silent fallback. + +Set `COMPILE_TRAIN_FN=false` near the top of the DDP launcher to disable it. +The typed default is `false`, and the option is currently rejected for FSDP, +so existing FastVideo pipelines are unaffected. Validation loss reuses the +compiled transformer forward; periodic generation keeps its separate eager +`guided_flow` path. Optional `torch.compile` arguments can be placed under +`training.model.torch_compile_kwargs`. + +For parity with upstream MMAudio, rank 0 runs the exact +`nitrous_ema.PostHocEMA` implementation with sigma profiles `[0.05, 0.1]`, +updates every optimizer step, and snapshots every 5,000 steps. FastVideo's +distributed checkpoint interval is 10,000 steps so every resumable training +checkpoint has a matching EMA snapshot. Periodic validation/inference still +uses online weights, and the final sigma-0.05 synthesis is saved under +`posthoc_ema/official_ddp/`. + +Set `VARIANT` to `medium_44k` or `large_44k` without changing the feature +cache. The launcher computes official latent normalization statistics once, +saves `latent_statistics_44k.pt` inside the feature directory, and reuses it +across those three variants. `small_16k` needs a separately preprocessed 16 kHz +cache because its latent shape is `[250, 20]`. `large_44k_v2` is rejected before +weights or data are loaded because the published training implementation does +not support v2. + +### Validation and training-time inference + +`MMAudioValidationCallback` evaluates the cached validation split every 5,000 +steps by default. It follows the official validation loss construction: +posterior latent sampling, logit-normal flow time, prior noise, and independent +video/text CFG masks. The validation RNG is reset on every pass so losses are +comparable between checkpoints. Set `VALIDATION_MAX_BATCHES=0` in the launcher +to use the full split, or a positive value for a faster subset. + +Every 20,000 steps the same callback runs the native FastVideo MMAudio pipeline +with the live training transformer and cached CLIP, Synchformer, and text +features. It generates a fixed global set of 16 validation outputs (four per +rank on a four-GPU job) spanning the validation batch's source rows, not the +complete validation split. The validation seed, +sampler, and per-sample noise seeds stay fixed across training steps. +`INFERENCE_MODEL_PATH` supplies only the frozen scheduler, audio VAE, and +vocoder needed to decode a waveform; its transformer weights and feature +encoders are not loaded. Samples are written to: + +```text +outputs/mmaudio__from_scratch/validation_audio/step_/ +``` + +Both a standalone WAV and an MP4 containing the original validation video plus +generated audio are saved. MP4 composition mirrors official MMAudio demo's +`reencode_with_audio`: decode source frames through the requested endpoint, +re-encode H.264 at the source's guessed frame rate and 10 Mbps in `yuv420p`, +then feed the mono generated waveform to the AAC encoder at the model sample +rate. It uses PyAV and does not import Torio. As in the official 44.1 kHz path, +the encoded container can be slightly longer than 8 seconds because 353,280 +audio samples and AAC frame padding do not end at exactly 8.0 seconds. If +composition or source-video loading fails, the WAV is retained and training +continues. + +Post-hoc EMA rank-local profiles and the final synthesized sigma-0.05 shards +are written under `posthoc_ema/rank_/`. Periodic validation uses online +weights by default, matching the official training loop; official MMAudio only +synthesizes sigma-0.05 EMA weights after training. Set `use_ema: true` in the +YAML only when intentionally evaluating an available post-hoc snapshot. The +current callback reports flow-matching validation loss and generates media +samples; dataset-level FAD/CLAP/synchronization benchmark metrics remain a +separate evaluation job. + +Set `TRACKER="wandb"` near the top of the launcher after `wandb login` to log +`total_loss`, `flow_matching_loss`, gradient norm, step time, validation loss, +and the 16 composed validation MP4 files. Set +`INFERENCE_LOG_TO_TRACKER=false` to keep media local while retaining scalar +logging. + +Each YAML's `training.model_path` names a matching converted component tree. +Training never reads its DiT weights; the path is retained so FastVideo can +copy the component configs and replace the transformer when exporting a DCP +checkpoint: + +```bash +python scripts/checkpoint_conversion/convert_mmaudio_to_diffusers.py \ + --variant small_44k \ + --transformer-config-only \ + --output converted_weights/mmaudio/small_44k +``` + +This export skeleton contains no weights and is sufficient for transformer-only +checkpoint export. Then run: + +```bash +python -m fastvideo.train.entrypoint.dcp_to_diffusers \ + --checkpoint outputs/mmaudio_small_44k_from_scratch/checkpoint-10000 \ + --output-dir converted_weights/mmaudio/small_44k_trained \ + --overwrite +``` + +For a directly runnable inference export, the template at +`converted_weights/mmaudio/small_44k` must contain all pipeline components. +A transformer-only template produces a transformer-only export, which can be +combined with the shared 44 kHz encoders/VAE/vocoder component tree. + +### Preprocess VGGSound with FastVideo + +Feature extraction is a native FastVideo preprocessing workflow. The upstream +MMAudio repository is useful as a numerical reference, but is not imported by +the preprocessing or training command. + +The reference training loader imports the now-removed +`torio.io.StreamingMediaDecoder`. Current FastVideo uses `torchaudio.load` for +audio and PyAV timestamp sampling for the two video frame rates, while retaining +the reference normalization, resampling, transforms, mel, and encoder contracts. + +First create a preprocessing-only component tree. It contains the 44.1 kHz +VAE encoder, DFN5B CLIP text/vision encoders, Synchformer, and tokenizer; it +does not duplicate the trainable DiT or inference vocoder: + +```bash +python scripts/checkpoint_conversion/convert_mmaudio_to_diffusers.py \ + --preprocessor-only \ + --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 \ + --output converted_weights/mmaudio/preprocess_44k +``` + +For the `Loie/VGGSound` release, extract the gzip shards once for random-access +decoding. The archives contain a seven-component directory prefix: + +```bash +mkdir -p /path/to/VGGSound/videos +for shard in /path/to/VGGSound/vggsound_*.tar.gz; do + tar -xzf "$shard" -C /path/to/VGGSound/videos --strip-components=7 +done +``` + +Run the shared FastVideo preprocessing entrypoint. Each rank gets disjoint +samples and writes resumable TensorDict shards. Existing sample IDs are skipped +when the same output directory is resumed. A filtered caption manifest must +have the same `idlabel` schema used by MMAudio; `label` is the full text +description encoded by DFN5B, not the original VGGSound class label. + +```bash +torchrun --standalone --nproc_per_node=4 \ + -m fastvideo.pipelines.preprocess.v1_preprocessing_new \ + --model-path converted_weights/mmaudio/preprocess_44k \ + --mode preprocess \ + --workload-type v2a \ + --preprocess.dataset-type vggsound \ + --preprocess.dataset-path /path/to/VGGSound \ + --preprocess.dataset-metadata-path /path/to/VGGSound/sets/filtered_caption/vgg-train-filtered-caption.tsv \ + --preprocess.dataset-split train \ + --preprocess.dataset-output-dir /path/to/VGGSound/mmaudio_features/train \ + --preprocess.preprocess-video-batch-size 1 \ + --preprocess.dataloader-num-workers 2 \ + --preprocess.samples-per-file 256 +``` + +The configurable launcher defaults to four GPUs and derives the manifest from +`SPLIT`: + +```bash +DATASET_PATH=/path/to/VGGSound SPLIT=train \ + bash examples/training/finetune/mmaudio/preprocess_vggsound.sh +``` + +Use the same pipeline for every split by changing only `SPLIT`; write each one +to its own cache directory: + +```bash +for split in train val test; do + DATASET_PATH=/path/to/VGGSound SPLIT="$split" GPU_NUM=4 \ + bash examples/training/finetune/mmaudio/preprocess_vggsound.sh +done +``` + +MMAudio's split contract peak-normalizes train audio to 0.95 and leaves val/test +audio amplitudes unchanged. Corrupt, silent, or short train inputs and invalid +val/test inputs are skipped and recorded in `failures_rank_*.jsonl`. + +### Precomputed feature contract + +The online training step deliberately does not run the audio VAE, DFN5B, or +Synchformer. Point `training.data.data_path` at either one TensorDict +memory-mapped directory or the parent directory of FastVideo's feature shards: + +```text +video + audio cache +├── mean [N, 345, 40] +├── std [N, 345, 40] +├── clip_features [N, 64, 1024] +├── sync_features [N, 192, 768] +└── text_features [N, 77, 1024] + +audio-only cache +├── mean [N, 345, 40] +├── std [N, 345, 40] +└── text_features [N, 77, 1024] +``` + +Video and audio-only caches can be mixed by using a mapping of cache paths to +repeat counts in YAML. Missing video conditions are replaced by MMAudio's +learned null-video tokens. The `text_features` values are normalized per-token +DFN5B CLIP hidden states computed from the dataset captions. The cache loader +has no runtime import dependency on the upstream `mmaudio` package. + +The example config intentionally contains an empty data path until a dataset is +prepared: + +```text +examples/train/configs/fine_tuning/mmaudio/small_44k.yaml +``` + +Start a single-GPU run while supplying the cache path on the command line: + +```bash +NUM_GPUS=1 bash examples/train/run.sh \ + examples/train/configs/fine_tuning/mmaudio/small_44k.yaml \ + --training.data.data_path /path/to/VGGSound/mmaudio_features/train +``` + +Leaving the path empty produces an explicit error instead of silently starting +with the wrong data. The pretrained single-GPU FSDP path and +scratch-initialized four-GPU FSDP path have both been validated. Sequence and +tensor parallelism are not implemented, so keep `sp_size: 1` and `tp_size: 1`. The +example preserves the official learning rate, AdamW betas/epsilon, weight +decay, 1,000-step warmup, CFG dropout, logit-normal timestep sampling, and +300,000-step duration. Its per-GPU batch size is intentionally one for initial +smoke testing and should be scaled with gradient accumulation or additional +data-parallel GPUs. + ## Run V2A Inference The runnable example is diff --git a/fastvideo/pipelines/basic/mmaudio/mmaudio_pipeline.py b/fastvideo/pipelines/basic/mmaudio/mmaudio_pipeline.py index a11a9b6c0..2a2e27b26 100644 --- a/fastvideo/pipelines/basic/mmaudio/mmaudio_pipeline.py +++ b/fastvideo/pipelines/basic/mmaudio/mmaudio_pipeline.py @@ -14,7 +14,16 @@ from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase class MMAudioPipeline(ComposedPipelineBase): - """MMAudio large-44k-v2 V2A/T2A inference pipeline.""" + """MMAudio large-44k-v2 V2A/T2A inference pipeline. + + ``text_encoder``, ``tokenizer``, ``image_encoder`` and ``image_encoder_2`` + are required modules, but the direct-feature path of the conditioning + stages returns before dereferencing them. Callers that pass cached + ``mmaudio_clip_features``/``mmaudio_sync_features`` (and + ``mmaudio_text_features``) may therefore inject a placeholder for those + four modules; callers that pass ``video_path`` or omit the cached features + must supply the real modules. + """ _required_config_modules = [ "transformer", diff --git a/fastvideo/pipelines/basic/mmaudio/presets.py b/fastvideo/pipelines/basic/mmaudio/presets.py index 08155837e..331b35089 100644 --- a/fastvideo/pipelines/basic/mmaudio/presets.py +++ b/fastvideo/pipelines/basic/mmaudio/presets.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -"""Published MMAudio large-44k-v2 inference defaults.""" +"""Published inference defaults for all official MMAudio variants.""" from fastvideo.api.presets import InferencePreset, PresetStageSpec @@ -26,15 +26,52 @@ _DEFAULTS = { "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), + +def _mmaudio_preset( + variant: str, + description: str, +) -> InferencePreset: + return InferencePreset( + name=f"mmaudio_{variant}", + version=1, + model_family="mmaudio", + description=description, + workload_type="v2a", + stage_schemas=(_DENOISE_STAGE, ), + defaults=dict(_DEFAULTS), + ) + + +MMAUDIO_SMALL_16K = _mmaudio_preset( + "small_16k", + ("MMAudio small 16 kHz video-to-audio generation with DFN5B CLIP, " + "Synchformer, the 16 kHz audio VAE, and the 16 kHz BigVGAN vocoder."), +) +MMAUDIO_SMALL_44K = _mmaudio_preset( + "small_44k", + ("MMAudio small 44.1 kHz video-to-audio generation with DFN5B CLIP, " + "Synchformer, the 44.1 kHz audio VAE, and BigVGAN-v2."), +) +MMAUDIO_MEDIUM_44K = _mmaudio_preset( + "medium_44k", + ("MMAudio medium 44.1 kHz video-to-audio generation with DFN5B CLIP, " + "Synchformer, the 44.1 kHz audio VAE, and BigVGAN-v2."), +) +MMAUDIO_LARGE_44K = _mmaudio_preset( + "large_44k", + ("MMAudio large 44.1 kHz video-to-audio generation with DFN5B CLIP, " + "Synchformer, the 44.1 kHz audio VAE, and BigVGAN-v2."), +) +MMAUDIO_LARGE_44K_V2 = _mmaudio_preset( + "large_44k_v2", + ("MMAudio large-44k-v2 video-to-audio generation with DFN5B CLIP, " + "Synchformer, the 44.1 kHz audio VAE, and BigVGAN-v2."), ) -ALL_PRESETS = (MMAUDIO_LARGE_44K_V2, ) +ALL_PRESETS = ( + MMAUDIO_SMALL_16K, + MMAUDIO_SMALL_44K, + MMAUDIO_MEDIUM_44K, + MMAUDIO_LARGE_44K, + MMAUDIO_LARGE_44K_V2, +) diff --git a/fastvideo/pipelines/basic/mmaudio/stages.py b/fastvideo/pipelines/basic/mmaudio/stages.py index f41870138..599951699 100644 --- a/fastvideo/pipelines/basic/mmaudio/stages.py +++ b/fastvideo/pipelines/basic/mmaudio/stages.py @@ -1,9 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 -"""Inference stages for MMAudio V2A/T2A. +"""Inference stages and shared video preprocessing 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. + +Training preprocessing can instead select the FFmpeg ``fps`` filter sampler, +matching MMAudio's original ``StreamingMediaDecoder`` training path. """ from __future__ import annotations @@ -68,6 +71,90 @@ def _read_frames_at_fps( return [np.stack(frames) for frames in outputs] +def _read_frames_with_fps_filters( + video_path: str | Path, + frame_rates: tuple[float, ...], + *, + start_s: float, + end_s: float, +) -> list[np.ndarray]: + """Decode RGB frames using one FFmpeg graph with an ``fps`` branch per rate. + + MMAudio's original training loader configures + ``StreamingMediaDecoder.add_basic_video_stream(frame_rate=...)``, which + delegates resampling to FFmpeg's ``fps`` filter. A single split graph keeps + both output rates on the same decode and avoids sharing a VideoFrame between + multiple filter graphs. + """ + import av + + if not frame_rates or any(fps <= 0 for fps in frame_rates): + raise ValueError("MMAudio video frame rates must be positive") + if start_s < 0 or end_s <= start_s: + raise ValueError(f"Invalid MMAudio video interval [{start_s}, {end_s}] seconds") + + expected_counts = [int((end_s - start_s) * fps) for fps in frame_rates] + if any(count <= 0 for count in expected_counts): + raise ValueError(f"MMAudio video interval [{start_s}, {end_s}] is too short for " + f"frame rates {frame_rates}") + outputs: list[list[np.ndarray]] = [[] for _ in frame_rates] + + with av.open(str(video_path)) as container: + stream = container.streams.video[0] + stream.thread_type = "AUTO" + + graph = av.filter.Graph() + source = graph.add_buffer(template=stream) + trim = graph.add("trim", f"start={start_s}:end={end_s}") + source.link_to(trim) + setpts = graph.add("setpts", "PTS-STARTPTS") + trim.link_to(setpts) + + if len(frame_rates) > 1: + split = graph.add("split", str(len(frame_rates))) + setpts.link_to(split) + branch_source = split + else: + branch_source = setpts + + sinks = [] + for index, fps in enumerate(frame_rates): + fps_filter = graph.add("fps", f"fps={fps}") + if len(frame_rates) > 1: + branch_source.link_to(fps_filter, index, 0) + else: + branch_source.link_to(fps_filter) + rgb_filter = graph.add("format", "pix_fmts=rgb24") + fps_filter.link_to(rgb_filter) + sink = graph.add("buffersink") + rgb_filter.link_to(sink) + sinks.append(sink) + graph.configure() + + def drain_outputs() -> None: + for index, (sink, expected_count) in enumerate(zip(sinks, expected_counts, strict=True)): + while len(outputs[index]) < expected_count: + try: + output_frame = sink.pull() + except (av.error.BlockingIOError, av.error.EOFError): + break + outputs[index].append(output_frame.to_ndarray(format="rgb24")) + + for frame in container.decode(stream): + source.push(frame) + drain_outputs() + if all(len(frames) >= expected for frames, expected in zip(outputs, expected_counts, strict=True)): + break + else: + source.push(None) + drain_outputs() + + if any(not frames for frames in outputs): + raise ValueError(f"Could not decode 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, *, @@ -76,13 +163,19 @@ def preprocess_mmaudio_video( sync_fps: int = 25, clip_size: int = 384, sync_size: int = 224, + use_ffmpeg_fps_filter: bool = False, ) -> 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]``. + + The default timestamp sampler preserves inference parity. Training callers + should enable ``use_ffmpeg_fps_filter`` to match MMAudio's original + StreamingMediaDecoder/FFmpeg resampling behavior. """ - clip_array, sync_array = _read_frames_at_fps( + frame_reader = (_read_frames_with_fps_filters if use_ffmpeg_fps_filter else _read_frames_at_fps) + clip_array, sync_array = frame_reader( video_path, (float(clip_fps), float(sync_fps)), start_s=0.0, @@ -144,8 +237,10 @@ class MMAudioInputValidationStage(PipelineStage): 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) + direct_video = ((batch.extra.get("mmaudio_clip_frames") is not None + and batch.extra.get("mmaudio_sync_frames") is not None) + or (batch.extra.get("mmaudio_clip_features") is not None + and batch.extra.get("mmaudio_sync_features") 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.") @@ -187,6 +282,51 @@ class MMAudioVideoConditioningStage(PipelineStage): def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch: pc = fastvideo_args.pipeline_config duration_s = float(batch.extra["mmaudio_duration_s"]) + direct_clip_features = batch.extra.get("mmaudio_clip_features") + direct_sync_features = batch.extra.get("mmaudio_sync_features") + if (direct_clip_features is None) != (direct_sync_features is None): + raise ValueError("MMAudio direct conditioning requires both CLIP and sync features.") + if direct_clip_features is not None: + if direct_clip_features.ndim != 3 or direct_sync_features.ndim != 3: + raise ValueError("MMAudio direct CLIP/sync features must have shape [B,T,C].") + if direct_clip_features.shape[0] != 1 or direct_sync_features.shape[0] != 1: + raise ValueError("MMAudio direct features currently require batch size 1.") + latent_length, clip_length, sync_length = mmaudio_sequence_lengths( + duration_s, + pc, + ) + expected_clip = (1, clip_length, pc.dit_config.arch_config.clip_dim) + expected_sync = (1, sync_length, pc.dit_config.arch_config.sync_dim) + if tuple(direct_clip_features.shape) != expected_clip: + raise ValueError(f"MMAudio direct CLIP features must have shape {expected_clip}, " + f"got {tuple(direct_clip_features.shape)}.") + if tuple(direct_sync_features.shape) != expected_sync: + raise ValueError(f"MMAudio direct sync features must have shape {expected_sync}, " + f"got {tuple(direct_sync_features.shape)}.") + self.transformer.update_seq_lengths( + latent_length, + clip_length, + sync_length, + ) + device = get_local_torch_device() + dtype = next(self.transformer.parameters()).dtype + batch.extra["mmaudio_sequence_lengths"] = ( + latent_length, + clip_length, + sync_length, + ) + batch.extra["mmaudio_clip_features"] = direct_clip_features.to( + device=device, + dtype=dtype, + non_blocking=True, + ) + batch.extra["mmaudio_sync_features"] = direct_sync_features.to( + device=device, + dtype=dtype, + non_blocking=True, + ) + return batch + 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) @@ -291,42 +431,73 @@ class MMAudioTextConditioningStage(PipelineStage): 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) + direct_text_features = batch.extra.get("mmaudio_text_features") - 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", + if direct_text_features is not None: + if direct_text_features.ndim != 3 or direct_text_features.shape[0] != 1: + raise ValueError("MMAudio direct text features must have shape [1,T,C].") + expected_shape = ( + 1, + self.transformer.config.arch_config.text_seq_len, + self.transformer.config.arch_config.text_dim, ) - # 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) + if tuple(direct_text_features.shape) != expected_shape: + raise ValueError(f"MMAudio direct text features must have shape {expected_shape}, " + f"got {tuple(direct_text_features.shape)}.") + text_features = direct_text_features.to( + device=device, + dtype=model_dtype, + non_blocking=True, + ) + negative_text_features = None + else: + self.text_encoder = self.text_encoder.to(device) - 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) + 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) - if fastvideo_args.text_encoder_cpu_offload: - self.text_encoder = self.text_encoder.to("cpu") + 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) - 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, - ) + if fastvideo_args.text_encoder_cpu_offload: + self.text_encoder = self.text_encoder.to("cpu") + + # A live training transformer can be FSDP2-sharded. Condition + # preprocessing reads positional-embedding parameters before the + # transformer's normal pre-forward hook would all-gather them. + unshard = getattr(self.transformer, "unshard", None) + if callable(unshard): + unshard() + with torch.autocast( + device_type=device.type, + dtype=torch.bfloat16, + enabled=device.type == "cuda", + ): + 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 @@ -390,13 +561,22 @@ class MMAudioDenoisingStage(PipelineStage): # Stop if interrupted if getattr(self, "interrupt", False): break - flow = self.transformer.guided_flow( - timestep / self.scheduler.config.num_train_timesteps, - latents, - conditions, - empty_conditions, - float(batch.guidance_scale), - ) + # Converted inference checkpoints carry bf16 parameters directly, + # while a live FSDP training model keeps fp32 master parameters and + # runs its forward pass under bf16 autocast. Use the same forward + # precision here so training-time validation can reuse this stage. + with torch.autocast( + device_type=latents.device.type, + dtype=torch.bfloat16, + enabled=latents.device.type == "cuda", + ): + 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 diff --git a/fastvideo/pipelines/pipeline_batch_info.py b/fastvideo/pipelines/pipeline_batch_info.py index 2eb029ac5..00a368799 100644 --- a/fastvideo/pipelines/pipeline_batch_info.py +++ b/fastvideo/pipelines/pipeline_batch_info.py @@ -353,6 +353,7 @@ class TrainingBatch: timesteps: torch.Tensor | None = None sigmas: torch.Tensor | None = None noise: torch.Tensor | None = None + training_target: torch.Tensor | None = None # MiniMax H3 reuses the packed row boundaries from batch preparation to # split the transformer's joint sequence back into video and audio outputs. diff --git a/fastvideo/pipelines/pipeline_registry.py b/fastvideo/pipelines/pipeline_registry.py index dc00d35b4..8980d7b7a 100644 --- a/fastvideo/pipelines/pipeline_registry.py +++ b/fastvideo/pipelines/pipeline_registry.py @@ -19,6 +19,7 @@ logger = init_logger(__name__) _PREPROCESS_WORKLOAD_TYPE_TO_PIPELINE_NAME: dict[WorkloadType, str] = { WorkloadType.I2V: "PreprocessPipelineI2V", WorkloadType.T2V: "PreprocessPipelineT2V", + WorkloadType.V2A: "PreprocessPipelineV2A", } @@ -57,9 +58,24 @@ class _PipelineRegistry: """Get supported pipelines for the given pipeline type.""" return set(self.pipelines.get(pipeline_type.value, {}).keys()) - def _load_preprocess_pipeline_cls(self, workload_type: WorkloadType) -> type[ComposedPipelineBase] | None: - pipeline_name = _PREPROCESS_WORKLOAD_TYPE_TO_PIPELINE_NAME[workload_type] - return self.pipelines.get(PipelineType.PREPROCESS.value, {}).get(pipeline_name) + def _load_preprocess_pipeline_cls( + self, + pipeline_name_in_config: str, + workload_type: WorkloadType, + ) -> type[ComposedPipelineBase] | None: + pipelines = self.pipelines.get(PipelineType.PREPROCESS.value, {}) + + # Existing T2V/I2V integrations retain their exact workload-level + # lookup. New V2A families may define different feature contracts and + # therefore use a family-specific preprocessing class. + if (workload_type is WorkloadType.V2A and pipeline_name_in_config.endswith("Pipeline")): + family_pipeline_name = (pipeline_name_in_config.removesuffix("Pipeline") + "PreprocessPipeline") + family_pipeline = pipelines.get(family_pipeline_name) + if family_pipeline is not None: + return family_pipeline + + generic_name = _PREPROCESS_WORKLOAD_TYPE_TO_PIPELINE_NAME[workload_type] + return pipelines.get(generic_name) def _try_load_pipeline_cls(self, pipeline_name_in_config: str, pipeline_type: PipelineType, workload_type: WorkloadType) -> type[ComposedPipelineBase] | type[LoRAPipeline] | None: @@ -68,7 +84,10 @@ class _PipelineRegistry: return None if pipeline_type == PipelineType.PREPROCESS: - return self._load_preprocess_pipeline_cls(workload_type) + return self._load_preprocess_pipeline_cls( + pipeline_name_in_config, + workload_type, + ) elif pipeline_type == PipelineType.BASIC or pipeline_type == PipelineType.TRAINING: return self.pipelines[pipeline_type.value].get(pipeline_name_in_config) else: diff --git a/fastvideo/pipelines/preprocess/mmaudio/__init__.py b/fastvideo/pipelines/preprocess/mmaudio/__init__.py new file mode 100644 index 000000000..16e995fbf --- /dev/null +++ b/fastvideo/pipelines/preprocess/mmaudio/__init__.py @@ -0,0 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 + +from fastvideo.pipelines.preprocess.mmaudio.mmaudio_preprocess_pipeline import ( + MMAudioPreprocessPipeline, ) + +__all__ = ["MMAudioPreprocessPipeline"] diff --git a/fastvideo/pipelines/preprocess/mmaudio/mmaudio_preprocess_pipeline.py b/fastvideo/pipelines/preprocess/mmaudio/mmaudio_preprocess_pipeline.py new file mode 100644 index 000000000..ae562f49c --- /dev/null +++ b/fastvideo/pipelines/preprocess/mmaudio/mmaudio_preprocess_pipeline.py @@ -0,0 +1,55 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Native FastVideo preprocessing pipeline for MMAudio training.""" + +from __future__ import annotations + +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.models.audio.mmaudio_processing import build_mmaudio_mel_converter +from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase +from fastvideo.pipelines.preprocess.mmaudio.stages import MMAudioFeatureExtractionStage + + +class MMAudioPreprocessPipeline(ComposedPipelineBase): + """Encode raw audio, video, and captions into MMAudio training features.""" + + _required_config_modules = [ + "audio_vae", + "text_encoder", + "tokenizer", + "image_encoder", + "image_encoder_2", + ] + + def __init__(self, model_path: str, fastvideo_args: FastVideoArgs, **kwargs) -> None: + # Official training features are extracted in fp32. These overrides + # are local to preprocess mode and do not change inference defaults. + config = fastvideo_args.pipeline_config + config.text_encoder_precisions = ("fp32", ) + config.image_encoder_precisions = ("fp32", "fp32") + config.audio_encoder_precision = "fp32" + super().__init__(model_path, fastvideo_args, **kwargs) + + def initialize_pipeline(self, fastvideo_args: FastVideoArgs) -> None: + audio_vae = self.get_module("audio_vae") + if not hasattr(audio_vae, "encoder"): + raise ValueError("MMAudio preprocessing requires an audio_vae component converted " + "with need_encoder=true. Use the converter's --preprocessor-only mode.") + if getattr(audio_vae, "mode", None) != "44k": + raise ValueError("The current MMAudio training integration requires the 44k audio VAE") + + def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None: + mel_converter = build_mmaudio_mel_converter("44k") + self.add_stage( + "feature_extraction_stage", + MMAudioFeatureExtractionStage( + audio_vae=self.get_module("audio_vae"), + mel_converter=mel_converter, + image_encoder=self.get_module("image_encoder"), + sync_encoder=self.get_module("image_encoder_2"), + text_encoder=self.get_module("text_encoder"), + tokenizer=self.get_module("tokenizer"), + ), + ) + + +EntryClass = MMAudioPreprocessPipeline diff --git a/fastvideo/pipelines/preprocess/mmaudio/stages.py b/fastvideo/pipelines/preprocess/mmaudio/stages.py new file mode 100644 index 000000000..9c9e099be --- /dev/null +++ b/fastvideo/pipelines/preprocess/mmaudio/stages.py @@ -0,0 +1,267 @@ +# SPDX-License-Identifier: Apache-2.0 +"""MMAudio-specific feature stages for the shared FastVideo V2A workflow.""" + +from __future__ import annotations + +import math +from pathlib import Path +from typing import Any, cast + +import torch +import torchaudio + +from fastvideo.configs.configs import VideoLoaderType +from fastvideo.distributed import get_local_torch_device +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.forward_context import set_forward_context +from fastvideo.logger import init_logger +from fastvideo.pipelines.basic.mmaudio.stages import ( + _CLIP_MEAN, + _CLIP_STD, + preprocess_mmaudio_video, +) +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch, PreprocessBatch +from fastvideo.pipelines.preprocess.mmaudio.torio_media_reader import (PREPROCESSED_MEDIA_KEY, PREPROCESS_ERROR_KEY, + preprocess_mmaudio_media_with_torio) +from fastvideo.pipelines.stages.base import PipelineStage +from fastvideo.pipelines.stages.validators import VerificationResult + +logger = init_logger(__name__) + + +class MMAudioFeatureExtractionStage(PipelineStage): + """Extract VAE posterior statistics plus CLIP/Synchformer features.""" + + def __init__( + self, + *, + audio_vae: torch.nn.Module, + mel_converter: torch.nn.Module, + image_encoder: torch.nn.Module, + sync_encoder: torch.nn.Module, + text_encoder: torch.nn.Module, + tokenizer: Any, + ) -> None: + super().__init__() + self.audio_vae = audio_vae.eval() + self.mel_converter = mel_converter.eval() + self.image_encoder = image_encoder.eval() + self.sync_encoder = sync_encoder.eval() + self.text_encoder = text_encoder.eval() + self.tokenizer = tokenizer + + def verify_input(self, batch, fastvideo_args): + return VerificationResult() + + def verify_output(self, batch, fastvideo_args): + return VerificationResult() + + @staticmethod + def _load_audio( + path: str | Path, + *, + target_sample_rate: int, + target_samples: int, + normalize_audio: bool, + ) -> torch.Tensor: + waveform, sample_rate = torchaudio.load(str(path)) + if waveform.ndim != 2 or waveform.shape[0] == 0: + raise ValueError(f"Decoded audio has an invalid shape: {tuple(waveform.shape)}") + waveform = waveform.mean(dim=0) + if normalize_audio: + abs_max = waveform.abs().max() + if abs_max <= 1e-6: + raise ValueError("Decoded audio is silent") + waveform = waveform / abs_max * 0.95 + if sample_rate != target_sample_rate: + waveform = torchaudio.functional.resample( + waveform, + sample_rate, + target_sample_rate, + lowpass_filter_width=64, + rolloff=0.9475937167399596, + resampling_method="sinc_interp_kaiser", + beta=14.769656459379492, + ) + if waveform.shape[0] < target_samples: + raise ValueError( + f"Audio is too short: need {target_samples} samples at {target_sample_rate} Hz, got {waveform.shape[0]}" + ) + return waveform[:target_samples].contiguous() + + @torch.inference_mode() + def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch: + batch = cast(PreprocessBatch, batch) + if not isinstance(batch.prompt, list): + raise ValueError("MMAudio preprocessing expects a list of captions") + if len(batch.video_loader) != len(batch.prompt): + raise ValueError("MMAudio preprocessing paths and captions have different batch sizes") + + pc = fastvideo_args.pipeline_config + duration_s = float(pc.duration_s) + sample_rate = int(pc.sampling_rate) + alignment = int(pc.spectrogram_frame_rate * pc.latent_downsample_rate) + target_samples = math.ceil(duration_s * sample_rate / alignment) * alignment + expected_clip_frames = int(duration_s * pc.clip_frame_rate) + expected_sync_frames = int(duration_s * pc.sync_frame_rate) + preprocess_config = fastvideo_args.preprocess_config + dataset_split = "train" if preprocess_config is None else preprocess_config.dataset_split.lower() + normalize_audio = dataset_split == "train" + use_torio = (preprocess_config is not None and preprocess_config.video_loader_type == VideoLoaderType.TORIO) + + preprocessed_media = batch.extra.get(PREPROCESSED_MEDIA_KEY) + preprocess_errors = batch.extra.get(PREPROCESS_ERROR_KEY) + for name, values in ( + (PREPROCESSED_MEDIA_KEY, preprocessed_media), + (PREPROCESS_ERROR_KEY, preprocess_errors), + ): + if values is not None and (not isinstance(values, list) or len(values) != len(batch.video_loader)): + raise ValueError(f"MMAudio {name} must align with the input batch") + + valid_ids: list[str] = [] + valid_paths: list[str] = [] + valid_captions: list[str] = [] + audio_batch: list[torch.Tensor] = [] + clip_batch: list[torch.Tensor] = [] + sync_batch: list[torch.Tensor] = [] + failures: list[dict[str, str]] = [] + + for index, (sample_id, path, + caption) in enumerate(zip( + batch.video_file_name, + batch.video_loader, + batch.prompt, + strict=True, + )): + try: + path_str = str(path) + if not Path(path_str).is_file(): + raise FileNotFoundError(path_str) + preprocess_error = (None if preprocess_errors is None else preprocess_errors[index]) + if preprocess_error is not None: + raise ValueError(str(preprocess_error)) + media = (None if preprocessed_media is None else preprocessed_media[index]) + if media is not None: + if not isinstance(media, dict): + raise ValueError("Preprocessed MMAudio media must be a dictionary") + audio = media["audio"] + clip_frames = media["clip_frames"] + sync_frames = media["sync_frames"] + effective_duration = float(media["effective_duration"]) + elif use_torio: + audio, clip_frames, sync_frames, effective_duration = (preprocess_mmaudio_media_with_torio( + path_str, + duration_s=duration_s, + target_sample_rate=sample_rate, + target_samples=target_samples, + normalize_audio=normalize_audio, + clip_fps=pc.clip_frame_rate, + sync_fps=pc.sync_frame_rate, + clip_size=pc.clip_image_size, + sync_size=pc.sync_image_size, + )) + else: + audio = self._load_audio( + path_str, + target_sample_rate=sample_rate, + target_samples=target_samples, + normalize_audio=normalize_audio, + ) + clip_frames, sync_frames, effective_duration = preprocess_mmaudio_video( + path_str, + 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, + use_ffmpeg_fps_filter=True, + ) + if effective_duration < duration_s: + raise ValueError(f"Video is too short: need {duration_s}s, got {effective_duration:.3f}s") + if clip_frames.shape[0] != expected_clip_frames: + raise ValueError(f"Expected {expected_clip_frames} CLIP frames, got {clip_frames.shape[0]}") + if sync_frames.shape[0] != expected_sync_frames: + raise ValueError(f"Expected {expected_sync_frames} sync frames, got {sync_frames.shape[0]}") + except Exception as exc: + logger.warning("Skipping MMAudio preprocessing sample %s: %s", sample_id, exc) + failures.append({"id": str(sample_id), "error": str(exc)}) + continue + valid_ids.append(str(sample_id)) + valid_paths.append(path_str) + valid_captions.append(str(caption)) + audio_batch.append(audio) + clip_batch.append(clip_frames) + sync_batch.append(sync_frames) + + batch.extra["failed_samples"] = failures + if not valid_ids: + batch.video_file_name = [] + batch.video_loader = [] + batch.prompt = [] + batch.extra["precomputed_features"] = {} + batch.extra["precomputed_metadata"] = [] + return batch + + device = get_local_torch_device() + self.audio_vae = self.audio_vae.to(device=device, dtype=torch.float32) + self.mel_converter = self.mel_converter.to(device=device, dtype=torch.float32) + waveforms = torch.stack(audio_batch).to(device=device, dtype=torch.float32) + posterior = self.audio_vae.encode(self.mel_converter(waveforms)) + latent_mean = posterior.mean.transpose(1, 2).detach().cpu() + latent_std = posterior.std.transpose(1, 2).detach().cpu() + + self.image_encoder = self.image_encoder.to(device=device, dtype=torch.float32) + clip_video = torch.stack(clip_batch).to(device=device, dtype=torch.float32) + mean = torch.tensor(_CLIP_MEAN, device=device).view(1, 1, 3, 1, 1) + std = torch.tensor(_CLIP_STD, device=device).view(1, 1, 3, 1, 1) + clip_video = ((clip_video - mean) / std).flatten(0, 1) + clip_outputs: list[torch.Tensor] = [] + chunk_size = int(pc.clip_batch_size_multiplier) + with set_forward_context(current_timestep=0, attn_metadata=None): + for start in range(0, clip_video.shape[0], 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).reshape(len(valid_ids), expected_clip_frames, -1).detach().cpu() + + self.sync_encoder = self.sync_encoder.to(device=device, dtype=torch.float32) + sync_video = torch.stack(sync_batch).to(device=device, dtype=torch.float32) + with set_forward_context(current_timestep=0, attn_metadata=None): + sync_features = self.sync_encoder(sync_video).last_hidden_state.detach().cpu() + + self.text_encoder = self.text_encoder.to(device=device, dtype=torch.float32) + tokens = self.tokenizer( + valid_captions, + padding="max_length", + truncation=True, + max_length=77, + return_tensors="pt", + ) + input_ids = tokens.input_ids.masked_fill(tokens.attention_mask == 0, 0).to(device) + with set_forward_context(current_timestep=0, attn_metadata=None): + text_features = self.text_encoder(input_ids).last_hidden_state.detach().cpu() + + batch.video_file_name = valid_ids + batch.video_loader = valid_paths + batch.prompt = valid_captions + batch.extra["precomputed_features"] = { + "mean": latent_mean, + "std": latent_std, + "clip_features": clip_features, + "sync_features": sync_features, + "text_features": text_features, + } + batch.extra["precomputed_metadata"] = [{ + "id": sample_id, + "caption": caption, + "source": path, + "split": dataset_split, + } for sample_id, caption, path in zip( + valid_ids, + valid_captions, + valid_paths, + strict=True, + )] + return batch + + +__all__ = ["MMAudioFeatureExtractionStage"] diff --git a/fastvideo/pipelines/preprocess/mmaudio/torio_media_reader.py b/fastvideo/pipelines/preprocess/mmaudio/torio_media_reader.py new file mode 100644 index 000000000..e775240c6 --- /dev/null +++ b/fastvideo/pipelines/preprocess/mmaudio/torio_media_reader.py @@ -0,0 +1,169 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Optional torio media reader for MMAudio training preprocessing.""" + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import torch +import torchaudio +from torchvision.transforms import v2 + +PREPROCESSED_MEDIA_KEY = "preprocessed_media" +PREPROCESS_ERROR_KEY = "preprocess_error" + + +def preprocess_mmaudio_media_with_torio( + media_path: str | Path, + *, + duration_s: float, + target_sample_rate: int, + target_samples: int, + normalize_audio: bool, + clip_fps: int = 8, + sync_fps: int = 25, + clip_size: int = 384, + sync_size: int = 224, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, float]: + """Decode audio and dual-rate video with StreamingMediaDecoder. + + The dependency is imported lazily so normal FastVideo environments do not + require the deprecated torio package. This backend is intended for offline + reference-compatible training preprocessing only. + """ + try: + from torio.io import StreamingMediaDecoder + except ImportError as exc: + raise ImportError("MMAudio torio preprocessing requires the isolated torio " + "environment. Use run_preprocess_vggsound_torio.sh.") from exc + + expected_clip_frames = int(clip_fps * duration_s) + expected_sync_frames = int(sync_fps * duration_s) + reader = StreamingMediaDecoder(str(media_path)) + reader.add_basic_video_stream( + frames_per_chunk=expected_clip_frames, + frame_rate=float(clip_fps), + format="rgb24", + ) + reader.add_basic_video_stream( + frames_per_chunk=expected_sync_frames, + frame_rate=float(sync_fps), + format="rgb24", + ) + reader.add_basic_audio_stream(frames_per_chunk=2**30) + reader.fill_buffer() + clip_frames, sync_frames, audio_frames = reader.pop_chunks() + + if clip_frames is None: + raise ValueError("CLIP video stream returned no frames") + if clip_frames.shape[0] < expected_clip_frames: + raise ValueError(f"CLIP video is too short: expected {expected_clip_frames} frames, " + f"got {clip_frames.shape[0]}") + if sync_frames is None: + raise ValueError("Synchformer video stream returned no frames") + if sync_frames.shape[0] < expected_sync_frames: + raise ValueError(f"Synchformer video is too short: expected {expected_sync_frames} " + f"frames, got {sync_frames.shape[0]}") + if audio_frames is None or audio_frames.ndim != 2: + shape = None if audio_frames is None else tuple(audio_frames.shape) + raise ValueError(f"Decoded audio has an invalid shape: {shape}") + + waveform = audio_frames.transpose(0, 1).mean(dim=0) + if normalize_audio: + abs_max = waveform.abs().max() + if abs_max <= 1e-6: + raise ValueError("Decoded audio is silent") + waveform = waveform / abs_max * 0.95 + + source_sample_rate = int(reader.get_out_stream_info(2).sample_rate) + if source_sample_rate != target_sample_rate: + waveform = torchaudio.transforms.Resample( + source_sample_rate, + target_sample_rate, + lowpass_filter_width=64, + rolloff=0.9475937167399596, + resampling_method="sinc_interp_kaiser", + beta=14.769656459379492, + )(waveform) + if waveform.shape[0] < target_samples: + raise ValueError(f"Audio is too short: need {target_samples} samples at " + f"{target_sample_rate} Hz, got {waveform.shape[0]}") + waveform = waveform[:target_samples].contiguous() + + 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[:expected_clip_frames]) + sync_frames = sync_transform(sync_frames[:expected_sync_frames]) + effective_duration = min( + duration_s, + clip_frames.shape[0] / clip_fps, + sync_frames.shape[0] / sync_fps, + ) + return waveform, clip_frames, sync_frames, effective_duration + + +@dataclass(frozen=True) +class MMAudioTorioRowPreprocessor: + """Decode one V2A metadata row on a background CPU thread.""" + + duration_s: float + target_sample_rate: int + target_samples: int + normalize_audio: bool + clip_fps: int = 8 + sync_fps: int = 25 + clip_size: int = 384 + sync_size: int = 224 + + def __call__(self, row: dict[str, Any]) -> dict[str, Any]: + output = dict(row) + media_path = Path(str(output["video_path"])) + try: + if not media_path.is_file(): + raise FileNotFoundError(media_path) + audio, clip, sync, duration = preprocess_mmaudio_media_with_torio( + media_path, + duration_s=self.duration_s, + target_sample_rate=self.target_sample_rate, + target_samples=self.target_samples, + normalize_audio=self.normalize_audio, + clip_fps=self.clip_fps, + sync_fps=self.sync_fps, + clip_size=self.clip_size, + sync_size=self.sync_size, + ) + output[PREPROCESSED_MEDIA_KEY] = { + "audio": audio, + "clip_frames": clip, + "sync_frames": sync, + "effective_duration": duration, + } + except Exception as exc: + output[PREPROCESS_ERROR_KEY] = f"{type(exc).__name__}: {exc}" + return output + + +__all__ = [ + "MMAudioTorioRowPreprocessor", + "PREPROCESSED_MEDIA_KEY", + "PREPROCESS_ERROR_KEY", + "preprocess_mmaudio_media_with_torio", +] diff --git a/fastvideo/registry.py b/fastvideo/registry.py index 365bb3ea9..467bf9605 100644 --- a/fastvideo/registry.py +++ b/fastvideo/registry.py @@ -45,7 +45,14 @@ 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.mmaudio import ( + MMAudioLarge44kV2PipelineConfig, + MMAudioLarge44kV2AConfig, + MMAudioMedium44kV2AConfig, + MMAudioSmall16kV2AConfig, + MMAudioSmall44kV2AConfig, + MMAudioV2AConfig, +) from fastvideo.configs.pipelines.turbodiffusion import ( TurboDiffusionI2V_A14B_Config, TurboDiffusionT2V_14B_Config, @@ -291,6 +298,31 @@ def _get_config_info( return None +def _mmaudio_variant_detector(expected: str) -> Callable[[str], bool]: + + def detector(path: str) -> bool: + normalized = path.lower().replace("-", "_") + return expected in normalized and (expected == "large_44k_v2" or "large_44k_v2" not in normalized) + + return detector + + +def _unversioned_mmaudio_detector(*, mode: str) -> Callable[[str], bool]: + variants = ("small_16k", "small_44k", "medium_44k", "large_44k") + + def detector(path: str) -> bool: + normalized = path.lower().replace("-", "_") + # This fallback is path-only. Matching ``mmaudiopipeline`` would also + # fire for every variant-specific model_index and create ambiguity. + if "mmaudio" not in normalized or "mmaudiopipeline" in normalized: + return False + if any(variant in normalized for variant in variants): + return False + return ("16k" in normalized) if mode == "16k" else ("16k" not in normalized) + + return detector + + def _register_wan_configs(definitions: tuple[WanModelDefinition, ...]) -> None: for definition in definitions: register_configs( @@ -316,17 +348,48 @@ def _register_configs() -> None: pipeline_cls_name="HeliosPyramidPipeline", ) - # 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. + # Official MMAudio video/text-to-audio variants. Converted directories keep + # the variant in their path, allowing the registry to select the matching + # 16 kHz or 44.1 kHz component configuration before any modules are loaded. + mmaudio_variants = ( + ("small_16k", MMAudioSmall16kV2AConfig), + ("small_44k", MMAudioSmall44kV2AConfig), + ("medium_44k", MMAudioMedium44kV2AConfig), + ("large_44k_v2", MMAudioLarge44kV2PipelineConfig), + ("large_44k", MMAudioLarge44kV2AConfig), + ) + for variant, pipeline_config_cls in mmaudio_variants: + register_configs( + sampling_param_cls=None, + pipeline_config_cls=pipeline_config_cls, + workload_types=(WorkloadType.V2A, WorkloadType.T2A), + hf_model_paths=[ + f"FastVideo/MMAudio-{variant.replace('_', '-')}-Diffusers", + ], + model_detectors=[_mmaudio_variant_detector(variant)], + model_family="mmaudio", + default_preset=f"mmaudio_{variant}", + pipeline_cls_name="MMAudioPipeline", + ) + + # Preserve converted preprocessing trees and older local checkpoint names + # that predate explicit variant tags. These detectors intentionally do not + # inspect ``_class_name=MMAudioPipeline`` so they cannot overlap the five + # exact variant entries above. + register_configs( + sampling_param_cls=None, + pipeline_config_cls=MMAudioSmall16kV2AConfig, + workload_types=(WorkloadType.V2A, WorkloadType.T2A), + model_detectors=[_unversioned_mmaudio_detector(mode="16k")], + model_family="mmaudio", + default_preset="mmaudio_small_16k", + pipeline_cls_name="MMAudioPipeline", + ) 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_detectors=[_unversioned_mmaudio_detector(mode="44k")], model_family="mmaudio", default_preset="mmaudio_large_44k_v2", pipeline_cls_name="MMAudioPipeline", diff --git a/fastvideo/tests/api/test_mmaudio.py b/fastvideo/tests/api/test_mmaudio.py new file mode 100644 index 000000000..89420f80e --- /dev/null +++ b/fastvideo/tests/api/test_mmaudio.py @@ -0,0 +1,62 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Registry coverage for every official MMAudio model variant.""" + +import json +from pathlib import Path + +import pytest + +from fastvideo.configs.pipelines.mmaudio import ( + MMAudioLarge44kV2PipelineConfig, + MMAudioLarge44kV2AConfig, + MMAudioMedium44kV2AConfig, + MMAudioSmall16kV2AConfig, + MMAudioSmall44kV2AConfig, + MMAudioV2AConfig, +) +from fastvideo.registry import get_default_preset, get_pipeline_config_cls_from_name + + +@pytest.mark.parametrize( + ("variant", "config_cls"), + [ + ("small-16k", MMAudioSmall16kV2AConfig), + ("small-44k", MMAudioSmall44kV2AConfig), + ("medium-44k", MMAudioMedium44kV2AConfig), + ("large-44k", MMAudioLarge44kV2AConfig), + ("large-44k-v2", MMAudioLarge44kV2PipelineConfig), + ], +) +def test_mmaudio_variant_registry(variant: str, config_cls: type) -> None: + model_path = f"FastVideo/MMAudio-{variant}-Diffusers" + + assert get_pipeline_config_cls_from_name(model_path) is config_cls + assert get_default_preset(model_path) == f"mmaudio_{variant.replace('-', '_')}" + + +@pytest.mark.parametrize( + ("directory_name", "config_cls", "preset"), + [ + ("preprocess_44k", MMAudioV2AConfig, "mmaudio_large_44k_v2"), + ("preprocess_16k", MMAudioSmall16kV2AConfig, "mmaudio_small_16k"), + ], +) +def test_unversioned_mmaudio_preprocessor_paths_remain_compatible( + tmp_path: Path, + directory_name: str, + config_cls: type, + preset: str, +) -> None: + model_path = tmp_path / "mmaudio" / directory_name + model_path.mkdir(parents=True) + (model_path / "model_index.json").write_text( + json.dumps({ + "_class_name": "MMAudioPipeline", + "_diffusers_version": "0.36.0", + "_fastvideo_preprocessor_only": True, + }), + encoding="utf-8", + ) + + assert get_pipeline_config_cls_from_name(str(model_path)) is config_cls + assert get_default_preset(str(model_path)) == preset diff --git a/fastvideo/tests/api/test_mmaudio_conversion.py b/fastvideo/tests/api/test_mmaudio_conversion.py new file mode 100644 index 000000000..0f25fe410 --- /dev/null +++ b/fastvideo/tests/api/test_mmaudio_conversion.py @@ -0,0 +1,101 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Lightweight conversion coverage for official MMAudio variants.""" + +from __future__ import annotations + +import argparse +import importlib.util +import json +from pathlib import Path + +import torch + + +def _converter_module(): + path = Path(__file__).resolve().parents[3] / ( + "scripts/checkpoint_conversion/convert_mmaudio_to_diffusers.py" + ) + spec = importlib.util.spec_from_file_location("mmaudio_variant_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_small_16k_transformer_conversion_writes_variant_contract( + tmp_path: Path, +) -> None: + converter = _converter_module() + checkpoint = tmp_path / "mmaudio_small_16k.pth" + output = tmp_path / "converted" + torch.save({"empty_string_feat": torch.zeros(77, 1024)}, checkpoint) + + converter.convert( + argparse.Namespace( + variant="small_16k", + transformer_checkpoint=checkpoint, + audio_vae_checkpoint=None, + synchformer_checkpoint=None, + dfn5b_dir=None, + bigvgan_dir=None, + bigvgan_checkpoint=None, + transformer_only=True, + transformer_config_only=False, + preprocessor_only=False, + output=output, + ) + ) + + transformer_config = json.loads( + (output / "transformer/config.json").read_text(encoding="utf-8") + ) + model_index = json.loads( + (output / "model_index.json").read_text(encoding="utf-8") + ) + assert transformer_config["latent_dim"] == 20 + assert transformer_config["latent_seq_len"] == 250 + assert transformer_config["hidden_dim"] == 448 + assert transformer_config["v2"] is False + assert model_index["_fastvideo_mmaudio_variant"] == "small_16k" + assert model_index["_fastvideo_audio_sample_rate"] == 16000 + + +def test_converter_lists_all_five_official_variants() -> None: + converter = _converter_module() + assert tuple(converter.TRANSFORMER_VARIANTS) == ( + "small_16k", + "small_44k", + "medium_44k", + "large_44k", + "large_44k_v2", + ) + assert converter.BIGVGAN_16K_CONFIG["num_mels"] == 80 + assert converter.BIGVGAN_16K_CONFIG["upsample_rates"] == [4, 4, 2, 2, 2, 2] + + +def test_weight_free_transformer_export_skeleton(tmp_path: Path) -> None: + converter = _converter_module() + output = tmp_path / "skeleton" + + converter.convert( + argparse.Namespace( + variant="medium_44k", + transformer_checkpoint=None, + audio_vae_checkpoint=None, + synchformer_checkpoint=None, + dfn5b_dir=None, + bigvgan_dir=None, + bigvgan_checkpoint=None, + transformer_only=False, + transformer_config_only=True, + preprocessor_only=False, + output=output, + ) + ) + + config = json.loads( + (output / "transformer/config.json").read_text(encoding="utf-8") + ) + assert config["hidden_dim"] == 896 + assert config["depth"] == 12 + assert not list((output / "transformer").glob("*.safetensors")) diff --git a/fastvideo/tests/dataset/test_mmaudio_feature_dataset.py b/fastvideo/tests/dataset/test_mmaudio_feature_dataset.py new file mode 100644 index 000000000..796455d86 --- /dev/null +++ b/fastvideo/tests/dataset/test_mmaudio_feature_dataset.py @@ -0,0 +1,137 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Tests for the MMAudio-compatible precomputed feature cache.""" + +from __future__ import annotations + +import json +from pathlib import Path + +import pytest +import torch + +from fastvideo.dataset.mmaudio_feature_dataset import ( + MMAudioFeatureDataset, + _expand_cache_root, + compute_mmaudio_latent_stats, +) + +tensordict = pytest.importorskip("tensordict") + + +def _write_cache(path: Path, *, with_video: bool) -> None: + tensors = { + "mean": torch.randn(2, 4, 3), + "std": torch.rand(2, 4, 3), + "text_features": torch.randn(2, 5, 7), + } + if with_video: + tensors["clip_features"] = torch.randn(2, 6, 8) + tensors["sync_features"] = torch.randn(2, 9, 10) + tensordict.TensorDict(tensors, batch_size=[2]).memmap_(str(path)) + + +def _load(path: Path, *, include_metadata: bool = False) -> MMAudioFeatureDataset: + return MMAudioFeatureDataset( + path, + latent_seq_len=4, + latent_dim=3, + clip_seq_len=6, + clip_dim=8, + sync_seq_len=9, + sync_dim=10, + text_seq_len=5, + text_dim=7, + include_metadata=include_metadata, + ) + + +@pytest.mark.parametrize("with_video", [False, True]) +def test_mmaudio_feature_cache_contract( + tmp_path: Path, + with_video: bool, +) -> None: + cache = tmp_path / "cache" + _write_cache(cache, with_video=with_video) + + dataset = _load(cache) + sample = dataset[0] + + assert len(dataset) == 2 + assert sample["audio_latent_mean"].shape == (4, 3) + assert sample["audio_latent_std"].shape == (4, 3) + assert sample["clip_features"].shape == (6, 8) + assert sample["sync_features"].shape == (9, 10) + assert sample["text_features"].shape == (5, 7) + assert sample["video_exists"].item() is with_video + assert sample["text_exists"].item() is True + if not with_video: + assert torch.count_nonzero(sample["clip_features"]) == 0 + assert torch.count_nonzero(sample["sync_features"]) == 0 + + +def test_mmaudio_feature_cache_rejects_wrong_shape(tmp_path: Path) -> None: + cache = tmp_path / "cache" + _write_cache(cache, with_video=True) + + with pytest.raises(ValueError, match="mean.*per-sample shape"): + MMAudioFeatureDataset( + cache, + latent_seq_len=99, + latent_dim=3, + clip_seq_len=6, + clip_dim=8, + sync_seq_len=9, + sync_dim=10, + text_seq_len=5, + text_dim=7, + ) + + +def test_mmaudio_feature_cache_loads_validation_metadata(tmp_path: Path) -> None: + cache = tmp_path / "cache" + _write_cache(cache, with_video=True) + rows = [ + {"id": "sample-a", "caption": "first", "source": "/video/a.mp4"}, + {"id": "sample-b", "caption": "second", "source": "/video/b.mp4"}, + ] + with (cache / "samples.jsonl").open("w", encoding="utf-8") as handle: + for row in rows: + handle.write(json.dumps(row) + "\n") + + sample = _load(cache, include_metadata=True)[1] + + assert sample["sample_id"] == "sample-b" + assert sample["caption"] == "second" + assert sample["source_path"] == "/video/b.mp4" + + +def test_mmaudio_feature_cache_discovers_nested_shards(tmp_path: Path) -> None: + first = tmp_path / "worker_00000" / "shard_000000" + second = tmp_path / "worker_00001" / "shard_000000" + _write_cache(first, with_video=True) + _write_cache(second, with_video=True) + + assert _expand_cache_root(tmp_path) == [first, second] + + +def test_mmaudio_latent_stats_match_official_reduction(tmp_path: Path) -> None: + first = tmp_path / "worker_00000" / "shard_000000" + second = tmp_path / "worker_00001" / "shard_000000" + _write_cache(first, with_video=True) + _write_cache(second, with_video=True) + + first_tensors = tensordict.TensorDict.load_memmap(first)["mean"] + second_tensors = tensordict.TensorDict.load_memmap(second)["mean"] + complete = torch.cat([first_tensors, second_tensors], dim=0) + expected_mean = complete.mean(dim=(0, 1), keepdim=True) + expected_std = complete.std(dim=(0, 1), keepdim=True) + + actual_mean, actual_std = compute_mmaudio_latent_stats( + tmp_path, + latent_seq_len=4, + latent_dim=3, + chunk_size=1, + ) + + torch.testing.assert_close(actual_mean, expected_mean) + torch.testing.assert_close(actual_std, expected_std) diff --git a/fastvideo/tests/dataset/test_v2a_preprocessing.py b/fastvideo/tests/dataset/test_v2a_preprocessing.py new file mode 100644 index 000000000..da25da068 --- /dev/null +++ b/fastvideo/tests/dataset/test_v2a_preprocessing.py @@ -0,0 +1,410 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Tests for shared raw V2A preprocessing infrastructure.""" + +from __future__ import annotations + +import json +import math +import sys +import types +from fractions import Fraction +from pathlib import Path + +import numpy as np +import pytest +import torch + +from fastvideo.configs.configs import VideoLoaderType +from fastvideo.dataset.mmaudio_feature_dataset import _expand_cache_root +from fastvideo.dataset.v2a_feature_cache import V2AFeatureShardWriter +from fastvideo.dataset.vggsound import VGGSoundDataset +from fastvideo.fastvideo_args import WorkloadType +from fastvideo.pipelines.basic.mmaudio import stages as basic_mmaudio_stages +from fastvideo.pipelines.pipeline_registry import PipelineType, _PipelineRegistry +from fastvideo.pipelines.preprocess.mmaudio import stages as mmaudio_stages +from fastvideo.pipelines.preprocess.mmaudio import torio_media_reader +from fastvideo.pipelines.preprocess.mmaudio.torio_media_reader import ( + MMAudioTorioRowPreprocessor, + PREPROCESSED_MEDIA_KEY, + PREPROCESS_ERROR_KEY, + preprocess_mmaudio_media_with_torio, +) +from fastvideo.workflow.preprocess.preprocess_workflow_v2a import ( + PreprocessWorkflowV2A, V2AForwardBatchBuilder) + +pytest.importorskip("tensordict") + + +class _GenericT2VPreprocessPipeline: + pass + + +class _FamilyPreprocessPipeline: + pass + + +def _write_fractional_fps_video( + path: Path, + *, + duration_s: float, + frame_rate: Fraction = Fraction(30_000, 1_001), +) -> None: + av = pytest.importorskip("av") + with av.open(str(path), mode="w") as container: + stream = container.add_stream("mpeg4", rate=frame_rate) + stream.width = 32 + stream.height = 32 + stream.pix_fmt = "yuv420p" + time_base = Fraction(frame_rate.denominator, frame_rate.numerator) + frame_count = math.ceil(duration_s * float(frame_rate)) + for index in range(frame_count): + pixels = np.full((32, 32, 3), index % 255, dtype=np.uint8) + frame = av.VideoFrame.from_ndarray(pixels, format="rgb24") + frame.pts = index + frame.time_base = time_base + for packet in stream.encode(frame): + container.mux(packet) + for packet in stream.encode(): + container.mux(packet) + + +def test_vggsound_dataset_maps_hf_metadata_to_clip_name(tmp_path: Path) -> None: + (tmp_path / "videos").mkdir() + (tmp_path / "vggsound.csv").write_text( + "abc123,7,a train sound,train\nxyz789,42,a test sound,test\n", + encoding="utf-8", + ) + + dataset = VGGSoundDataset(tmp_path, split="train") + + assert len(dataset) == 1 + assert dataset[0] == { + "id": "abc123_000007", + "caption": "a train sound", + "video_path": str(tmp_path / "videos" / "abc123_000007.mp4"), + } + + +def test_vggsound_dataset_reads_filtered_caption_manifest(tmp_path: Path) -> None: + videos = tmp_path / "videos" + videos.mkdir() + (videos / "video-one_000010.mp4").touch() + manifest = tmp_path / "vgg-val-filtered-caption.tsv" + manifest.write_text( + "id\tlabel\n" + "video-one_000010\tA detailed description of the audible event.\n" + "missing_000020\tThis row has no corresponding video.\n", + encoding="utf-8", + ) + + dataset = VGGSoundDataset(tmp_path, split="val", metadata_path=manifest) + + assert len(dataset) == 1 + assert dataset[0] == { + "id": "video-one_000010", + "caption": "A detailed description of the audible event.", + "video_path": str(videos / "video-one_000010.mp4"), + } + + +def test_mmaudio_audio_normalization_matches_split_contract(monkeypatch: pytest.MonkeyPatch) -> None: + waveform = torch.tensor([[0.25, -0.5, 0.125, 0.0]], dtype=torch.float32) + monkeypatch.setattr( + mmaudio_stages.torchaudio, + "load", + lambda _path: (waveform.clone(), 44_100), + ) + + train_audio = mmaudio_stages.MMAudioFeatureExtractionStage._load_audio( + "sample.mp4", + target_sample_rate=44_100, + target_samples=4, + normalize_audio=True, + ) + eval_audio = mmaudio_stages.MMAudioFeatureExtractionStage._load_audio( + "sample.mp4", + target_sample_rate=44_100, + target_samples=4, + normalize_audio=False, + ) + + torch.testing.assert_close(train_audio.abs().max(), torch.tensor(0.95)) + torch.testing.assert_close(eval_audio, waveform[0]) + + +def test_mmaudio_torio_loader_is_an_explicit_config_choice() -> None: + assert VideoLoaderType.from_string("torio") is VideoLoaderType.TORIO + + +def test_mmaudio_torio_reader_preserves_reference_media_contract( + monkeypatch: pytest.MonkeyPatch, +) -> None: + calls: list[tuple[str, int, float | None]] = [] + + class _StreamInfo: + sample_rate = 44_100 + + class _FakeStreamingMediaDecoder: + def __init__(self, path: str) -> None: + assert path == "sample.mp4" + + def add_basic_video_stream( + self, + *, + frames_per_chunk: int, + frame_rate: float, + format: str, + ) -> None: + assert format == "rgb24" + calls.append(("video", frames_per_chunk, frame_rate)) + + def add_basic_audio_stream(self, *, frames_per_chunk: int) -> None: + calls.append(("audio", frames_per_chunk, None)) + + def fill_buffer(self) -> None: + pass + + def pop_chunks(self): + clip = torch.arange(8 * 3 * 16 * 20, dtype=torch.uint8).reshape(8, 3, 16, 20) + sync = torch.arange(25 * 3 * 16 * 20, dtype=torch.uint8).reshape(25, 3, 16, 20) + audio = torch.stack([ + torch.linspace(-0.5, 0.5, 44_100), + torch.linspace(-0.25, 0.25, 44_100), + ], dim=1) + return clip, sync, audio + + def get_out_stream_info(self, index: int) -> _StreamInfo: + assert index == 2 + return _StreamInfo() + + torio_module = types.ModuleType("torio") + torio_module.__path__ = [] + io_module = types.ModuleType("torio.io") + io_module.StreamingMediaDecoder = _FakeStreamingMediaDecoder + torio_module.io = io_module + monkeypatch.setitem(sys.modules, "torio", torio_module) + monkeypatch.setitem(sys.modules, "torio.io", io_module) + + audio, clip, sync, duration = preprocess_mmaudio_media_with_torio( + "sample.mp4", + duration_s=1.0, + target_sample_rate=44_100, + target_samples=44_100, + normalize_audio=True, + clip_size=32, + sync_size=24, + ) + + assert calls == [ + ("video", 8, 8.0), + ("video", 25, 25.0), + ("audio", 2**30, None), + ] + assert audio.shape == (44_100,) + torch.testing.assert_close(audio.abs().max(), torch.tensor(0.95)) + assert clip.shape == (8, 3, 32, 32) + assert sync.shape == (25, 3, 24, 24) + assert duration == 1.0 + + +def test_mmaudio_torio_row_preprocessor_returns_worker_ready_media( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + video_path = tmp_path / "sample.mp4" + video_path.touch() + expected = ( + torch.ones(16), + torch.ones(8, 3, 4, 4), + torch.ones(25, 3, 4, 4), + 1.0, + ) + monkeypatch.setattr( + torio_media_reader, + "preprocess_mmaudio_media_with_torio", + lambda *_args, **_kwargs: expected, + ) + preprocessor = MMAudioTorioRowPreprocessor( + duration_s=1.0, + target_sample_rate=16, + target_samples=16, + normalize_audio=True, + clip_size=4, + sync_size=4, + ) + + output = preprocessor({ + "id": "sample", + "caption": "sound", + "video_path": str(video_path), + }) + + assert PREPROCESS_ERROR_KEY not in output + media = output[PREPROCESSED_MEDIA_KEY] + assert media["audio"] is expected[0] + assert media["clip_frames"] is expected[1] + assert media["sync_frames"] is expected[2] + assert media["effective_duration"] == 1.0 + + +def test_v2a_batch_builder_preserves_preprocessed_media_alignment() -> None: + media = { + "audio": torch.ones(1), + "clip_frames": torch.ones(1), + "sync_frames": torch.ones(1), + "effective_duration": 1.0, + } + rows = [{ + "id": "good", + "caption": "good sound", + "video_path": "good.mp4", + PREPROCESSED_MEDIA_KEY: media, + }, { + "id": "bad", + "caption": "bad sound", + "video_path": "bad.mp4", + PREPROCESS_ERROR_KEY: "decode failed", + }] + + batch = V2AForwardBatchBuilder(seed=123)(rows) + + assert batch.extra[PREPROCESSED_MEDIA_KEY] == [media, None] + assert batch.extra[PREPROCESS_ERROR_KEY] == [None, "decode failed"] + + +def test_v2a_torio_thread_prefetch_preserves_batch_order() -> None: + workflow = object.__new__(PreprocessWorkflowV2A) + workflow.training_dataloader = [ + [{"id": "0"}, {"id": "1"}], + [{"id": "2"}, {"id": "3"}], + [{"id": "4"}], + ] + workflow.torio_decode_workers = 2 + workflow.torio_row_preprocessor = lambda row: { + **row, + "decoded": True, + } + + batches = list(workflow._iter_preprocessed_rows()) + + assert [[row["id"] for row in batch] for batch in batches] == [ + ["0", "1"], + ["2", "3"], + ["4"], + ] + assert all(row["decoded"] for batch in batches for row in batch) + + +def test_mmaudio_ffmpeg_sampler_recovers_fractional_fps_boundary( + tmp_path: Path, +) -> None: + video_path = tmp_path / "fractional-fps.mp4" + _write_fractional_fps_video(video_path, duration_s=10.0) + + clip_frames, sync_frames = basic_mmaudio_stages._read_frames_with_fps_filters( + video_path, + (8.0, 25.0), + start_s=0.0, + end_s=8.0, + ) + + assert clip_frames.shape == (64, 32, 32, 3) + assert sync_frames.shape == (200, 32, 32, 3) + + +def test_mmaudio_ffmpeg_sampler_does_not_pad_short_video( + tmp_path: Path, +) -> None: + video_path = tmp_path / "short.mp4" + _write_fractional_fps_video(video_path, duration_s=7.0) + + clip_frames, sync_frames = basic_mmaudio_stages._read_frames_with_fps_filters( + video_path, + (8.0, 25.0), + start_s=0.0, + end_s=8.0, + ) + + assert clip_frames.shape[0] < 64 + assert sync_frames.shape[0] < 200 + + +def test_mmaudio_inference_keeps_timestamp_sampler_by_default( + monkeypatch: pytest.MonkeyPatch, +) -> None: + calls: list[str] = [] + sampled_frames = [ + np.zeros((8, 8, 8, 3), dtype=np.uint8), + np.zeros((25, 8, 8, 3), dtype=np.uint8), + ] + + def timestamp_sampler(*_args, **_kwargs): + calls.append("timestamp") + return sampled_frames + + def ffmpeg_sampler(*_args, **_kwargs): + calls.append("ffmpeg") + return sampled_frames + + monkeypatch.setattr( + basic_mmaudio_stages, + "_read_frames_at_fps", + timestamp_sampler, + ) + monkeypatch.setattr( + basic_mmaudio_stages, + "_read_frames_with_fps_filters", + ffmpeg_sampler, + ) + + basic_mmaudio_stages.preprocess_mmaudio_video( + "unused.mp4", + duration_s=1.0, + clip_size=8, + sync_size=8, + ) + basic_mmaudio_stages.preprocess_mmaudio_video( + "unused.mp4", + duration_s=1.0, + clip_size=8, + sync_size=8, + use_ffmpeg_fps_filter=True, + ) + + assert calls == ["timestamp", "ffmpeg"] + + +def test_v2a_writer_creates_discoverable_resumable_shards(tmp_path: Path) -> None: + writer = V2AFeatureShardWriter(tmp_path, rank=0, samples_per_shard=1) + features = { + "mean": torch.randn(1, 4, 3), + "std": torch.rand(1, 4, 3), + "clip_features": torch.randn(1, 6, 8), + "sync_features": torch.randn(1, 9, 10), + "text_features": torch.randn(1, 5, 7), + } + writer.append(features, [{"id": "sample-1", "caption": "sound"}]) + writer.close() + + shards = _expand_cache_root(tmp_path) + assert shards == [tmp_path / "worker_00000" / "shard_000000"] + metadata = json.loads((shards[0] / "samples.jsonl").read_text().strip()) + assert metadata["id"] == "sample-1" + + resumed = V2AFeatureShardWriter(tmp_path, rank=0, samples_per_shard=1) + assert resumed.contains("sample-1") + resumed.append(features, [{"id": "sample-1", "caption": "duplicate"}]) + resumed.close() + assert len(_expand_cache_root(tmp_path)) == 1 + + +def test_family_preprocessing_is_enabled_only_for_v2a() -> None: + registry = _PipelineRegistry({ + PipelineType.PREPROCESS.value: { + "PreprocessPipelineT2V": _GenericT2VPreprocessPipeline, + "ExamplePreprocessPipeline": _FamilyPreprocessPipeline, + } + }) + + assert registry._load_preprocess_pipeline_cls("ExamplePipeline", WorkloadType.V2A) is _FamilyPreprocessPipeline + assert registry._load_preprocess_pipeline_cls("ExamplePipeline", WorkloadType.T2V) is _GenericT2VPreprocessPipeline diff --git a/fastvideo/tests/eval/test_audio_fad_math.py b/fastvideo/tests/eval/test_audio_fad_math.py new file mode 100644 index 000000000..8cbec9d1c --- /dev/null +++ b/fastvideo/tests/eval/test_audio_fad_math.py @@ -0,0 +1,17 @@ +import numpy as np +import pytest + +from fastvideo.eval.metrics.audio.frechet_distance.metric import ( + _frechet_distance, ) + + +def test_frechet_distance_supports_current_scipy_sqrtm_api(): + mean_a = np.array([0.0, 1.0]) + mean_b = np.array([1.0, 3.0]) + covariance = np.array([[2.0, 0.25], [0.25, 1.0]]) + + same = _frechet_distance(mean_a, covariance, mean_a, covariance) + shifted = _frechet_distance(mean_a, covariance, mean_b, covariance) + + assert same == pytest.approx(0.0, abs=1e-10) + assert shifted == pytest.approx(5.0, abs=1e-10) diff --git a/fastvideo/tests/eval/test_inputs.py b/fastvideo/tests/eval/test_inputs.py new file mode 100644 index 000000000..8b2d895a5 --- /dev/null +++ b/fastvideo/tests/eval/test_inputs.py @@ -0,0 +1,97 @@ +import json + +import pytest +import torch + +from fastvideo.eval import Video, samples_from_manifest +from fastvideo.eval.pool import VideoPool + + +def test_samples_from_manifest_resolves_media_paths(tmp_path): + media = tmp_path / "media" + media.mkdir() + video = media / "sample.mp4" + audio = media / "sample.wav" + video.touch() + audio.touch() + manifest = tmp_path / "samples.jsonl" + manifest.write_text( + json.dumps({ + "id": "sample", + "video": "media/sample.mp4", + "audio": "media/sample.wav", + "text_prompt": "A bell rings.", + }) + "\n", + encoding="utf-8", + ) + + samples = samples_from_manifest(manifest) + + assert len(samples) == 1 + assert samples[0]["id"] == "sample" + assert isinstance(samples[0]["video"], Video) + assert samples[0]["video"].source == str(video) + assert samples[0]["audio"] == str(audio) + assert samples[0]["text_prompt"] == "A bell rings." + + +def test_samples_from_manifest_rejects_rows_without_generated_media(tmp_path): + manifest = tmp_path / "samples.jsonl" + manifest.write_text( + json.dumps({ + "id": "sample", + "text_prompt": "A bell rings." + }) + "\n", + encoding="utf-8", + ) + + with pytest.raises(ValueError, match="needs at least video or audio"): + samples_from_manifest(manifest) + + +def test_samples_from_manifest_rejects_empty_manifest(tmp_path): + manifest = tmp_path / "samples.jsonl" + manifest.write_text("", encoding="utf-8") + + with pytest.raises(ValueError, match="manifest is empty"): + samples_from_manifest(manifest) + + +def test_video_pool_keeps_explicit_fps_over_probe(monkeypatch): + frames = torch.zeros(2, 3, 4, 4) + probes = [] + + def _probe(source): + probes.append(source) + return 29.97 + + monkeypatch.setattr("fastvideo.eval.io.video.probe_video_fps", _probe) + monkeypatch.setattr( + "fastvideo.eval.io.video.load_video", + lambda _source: frames, + ) + video = Video(source="sample.mp4") + + decoded = VideoPool([])._decode({"video": video, "fps": 8.0}) + + assert decoded["fps"] == pytest.approx(8.0) + assert probes == [] + assert decoded["video"].frames is frames + + +def test_video_pool_propagates_path_video_fps(monkeypatch): + frames = torch.zeros(2, 3, 4, 4) + monkeypatch.setattr( + "fastvideo.eval.io.video.probe_video_fps", + lambda _source: 29.97, + ) + monkeypatch.setattr( + "fastvideo.eval.io.video.load_video", + lambda _source: frames, + ) + video = Video(source="sample.mp4") + + decoded = VideoPool([])._decode({"video": video}) + + assert decoded["fps"] == pytest.approx(29.97) + assert decoded["video"].frames is frames diff --git a/fastvideo/tests/eval/test_v2a_av_benchmark.py b/fastvideo/tests/eval/test_v2a_av_benchmark.py new file mode 100644 index 000000000..4fe9f2ed7 --- /dev/null +++ b/fastvideo/tests/eval/test_v2a_av_benchmark.py @@ -0,0 +1,163 @@ +import json +import subprocess +from pathlib import Path + +from fastvideo.eval.v2a import run_av_benchmark +from fastvideo.eval.v2a import _av_benchmark_runner as runner + + +def _make_gt_cache(path: Path) -> None: + path.mkdir() + for name in runner._BASE_CACHE_FILES: + (path / name).touch() + + +def test_run_av_benchmark_uses_selected_python(monkeypatch, tmp_path): + audio_dir = tmp_path / "audio" + gt_cache = tmp_path / "gt" + audio_dir.mkdir() + gt_cache.mkdir() + selected_python = tmp_path / "av-benchmark-venv" / "python" + captured = {} + + def fake_subprocess_run(command, check): + captured["command"] = command + captured["check"] = check + output = Path(command[command.index("--output") + 1]) + output.parent.mkdir(parents=True, exist_ok=True) + output.write_text("{}", encoding="utf-8") + return subprocess.CompletedProcess(command, 0) + + monkeypatch.setattr(subprocess, "run", fake_subprocess_run) + + output = run_av_benchmark( + audio_dir=audio_dir, + gt_cache=gt_cache, + python_executable=selected_python, + batch_size=8, + num_workers=3, + skip_clap=True, + ) + + command = captured["command"] + assert command[0] == str(selected_python) + assert command[1].endswith("_av_benchmark_runner.py") + assert command[command.index("--batch-size") + 1] == "8" + assert command[command.index("--num-workers") + 1] == "3" + assert "--skip-clap" in command + assert captured["check"] is True + assert output == audio_dir.parent / "av_benchmark_results.json" + + +def test_official_runner_extracts_and_serializes_metrics(monkeypatch, tmp_path): + audio_dir = tmp_path / "audio" + gt_cache = tmp_path / "gt" + prediction_cache = tmp_path / "pred-cache" + output = tmp_path / "results.json" + audio_dir.mkdir() + (audio_dir / "sample.wav").touch() + _make_gt_cache(gt_cache) + calls = {} + + def fake_extract(**kwargs): + calls["extract"] = kwargs + prediction_cache.mkdir(parents=True, exist_ok=True) + for name in runner._required_prediction_files(skip_video_related=False, skip_clap=True): + (prediction_cache / name).touch() + + def fake_evaluate(**kwargs): + calls["evaluate"] = kwargs + return { + "FD-PASST": 12.5, + "KL-PASST-softmax": 0.25, + "DeSync": 0.4, + } + + monkeypatch.setattr(runner, "_load_backend", lambda: (fake_extract, fake_evaluate)) + + payload = runner.run_benchmark( + audio_dir=audio_dir, + gt_cache=gt_cache, + prediction_cache=prediction_cache, + output=output, + device="cuda", + batch_size=4, + num_workers=2, + audio_length=8.0, + recompute=False, + skip_video_related=False, + skip_clap=True, + ) + + assert calls["extract"]["audio_path"] == audio_dir.resolve() + assert calls["extract"]["batch_size"] == 4 + assert calls["extract"]["num_workers"] == 2 + assert calls["evaluate"]["gt_audio_cache"] == gt_cache.resolve() + assert payload["num_generated_audio"] == 1 + assert payload["extraction_ran"] is True + assert payload["metrics"]["FD-PASST"] == 12.5 + assert json.loads(output.read_text(encoding="utf-8"))["metrics"]["DeSync"] == 0.4 + + +def test_official_runner_reuses_complete_prediction_cache(monkeypatch, tmp_path): + audio_dir = tmp_path / "audio" + gt_cache = tmp_path / "gt" + prediction_cache = tmp_path / "pred-cache" + output = tmp_path / "results.json" + audio_dir.mkdir() + prediction_cache.mkdir() + (audio_dir / "sample.flac").touch() + _make_gt_cache(gt_cache) + for name in runner._required_prediction_files(skip_video_related=True, skip_clap=True): + (prediction_cache / name).touch() + + def fail_extract(**_kwargs): + raise AssertionError("complete prediction cache should be reused") + + monkeypatch.setattr(runner, "_load_backend", lambda: (fail_extract, lambda **_kwargs: {"FD-PASST": 1.0})) + + payload = runner.run_benchmark( + audio_dir=audio_dir, + gt_cache=gt_cache, + prediction_cache=prediction_cache, + output=output, + device="cuda", + batch_size=4, + num_workers=0, + audio_length=8.0, + recompute=False, + skip_video_related=True, + skip_clap=True, + ) + + assert payload["extraction_ran"] is False + assert payload["metrics"] == {"FD-PASST": 1.0} + + +def test_prediction_keys_align_only_unique_leading_underscore_matches(tmp_path): + import torch + + gt_cache = tmp_path / "gt" + prediction_cache = tmp_path / "pred" + gt_cache.mkdir() + prediction_cache.mkdir() + torch.save({"_123": torch.tensor(1), "normal": torch.tensor(2)}, gt_cache / "passt_logits.pth") + torch.save( + {"123": torch.tensor(3), "normal": torch.tensor(4), "unknown": torch.tensor(5)}, + prediction_cache / "passt_logits.pth", + ) + + stats = runner._align_prediction_cache_keys( + gt_cache=gt_cache, + prediction_cache=prediction_cache, + required_files=("passt_logits.pth",), + ) + + aligned = torch.load(prediction_cache / "passt_logits.pth", weights_only=True) + assert set(aligned) == {"_123", "normal", "unknown"} + assert stats == { + "exact_before_alignment": 1, + "remapped": 1, + "ambiguous": 0, + "unmatched_after_alignment": 1, + } diff --git a/fastvideo/tests/train/callbacks/test_ddp_posthoc_ema.py b/fastvideo/tests/train/callbacks/test_ddp_posthoc_ema.py new file mode 100644 index 000000000..487931e46 --- /dev/null +++ b/fastvideo/tests/train/callbacks/test_ddp_posthoc_ema.py @@ -0,0 +1,73 @@ +# SPDX-License-Identifier: Apache-2.0 +"""CPU-only lifecycle test for official MMAudio DDP PostHocEMA.""" + +from __future__ import annotations + +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch +import torch.distributed as dist + +from fastvideo.train.callbacks.ddp_posthoc_ema import DDPPostHocEMACallback +from fastvideo.train.utils.distributed_strategy import wrap_module_ddp + + +def test_official_ddp_posthoc_ema_snapshots_and_resumes( + tmp_path: Path, +) -> None: + pytest.importorskip("nitrous_ema") + rendezvous = tmp_path / "ema_rendezvous" + dist.init_process_group( + "gloo", + init_method=f"file://{rendezvous}", + rank=0, + world_size=1, + ) + try: + transformer = wrap_module_ddp( + torch.nn.Linear(2, 2, bias=False), + device=torch.device("cpu"), + broadcast_buffers=False, + ) + method = SimpleNamespace( + student=SimpleNamespace( + transformer=transformer, + device=torch.device("cpu"), + ), + ) + callback = DDPPostHocEMACallback( + checkpoint_every=2, + checkpoint_folder=str(tmp_path / "snapshots"), + ) + callback.training_config = SimpleNamespace( + checkpoint=SimpleNamespace(output_dir=str(tmp_path)), + ) + callback.on_train_start(method) + callback.on_training_step_end(method, {}, iteration=1) + with torch.no_grad(): + transformer.module.weight.add_(1) + callback.on_training_step_end(method, {}, iteration=2) + + snapshot_root = tmp_path / "snapshots" + assert (snapshot_root / "0.2.pt").is_file() + assert (snapshot_root / "1.2.pt").is_file() + + restored = DDPPostHocEMACallback( + checkpoint_every=2, + checkpoint_folder=str(snapshot_root), + ) + restored.training_config = callback.training_config + restored.on_train_start(method) + restored.load_state_dict(callback.state_dict()) + assert restored._ema is not None + assert int(restored._ema.step) == 2 + + callback.on_train_end(method, iteration=2) + final = list(snapshot_root.glob( + "mmaudio_ema_final_sigma_0p05_step_*.pth", + )) + assert len(final) == 1 + finally: + dist.destroy_process_group() diff --git a/fastvideo/tests/train/callbacks/test_mmaudio_validation.py b/fastvideo/tests/train/callbacks/test_mmaudio_validation.py new file mode 100644 index 000000000..a3073d176 --- /dev/null +++ b/fastvideo/tests/train/callbacks/test_mmaudio_validation.py @@ -0,0 +1,39 @@ +# SPDX-License-Identifier: Apache-2.0 + +from fastvideo.train.callbacks.mmaudio_validation import ( + _global_inference_indices, + _inference_sample_index, ) + + +def test_mmaudio_inference_uses_fixed_global_budget() -> None: + assignments = [_global_inference_indices(16, 4, rank) for rank in range(4)] + + assert all(len(rank_indices) == 4 for rank_indices in assignments) + assert sorted(index for rank_indices in assignments + for index in rank_indices) == list(range(16)) + + +def test_mmaudio_inference_rows_cover_validation_batch() -> None: + available = 8 + assignments = [_global_inference_indices(16, 4, rank) for rank in range(4)] + + rows = { + _inference_sample_index(global_index, local_index, available) + for rank_indices in assignments + for local_index, global_index in enumerate(rank_indices) + } + + assert rows == set(range(min(16, available))) + + +def test_mmaudio_inference_pads_collective_calls() -> None: + assignments = [_global_inference_indices(16, 6, rank) for rank in range(6)] + + assert all(len(rank_indices) == 3 for rank_indices in assignments) + valid = [ + index for rank_indices in assignments for index in rank_indices + if index is not None + ] + assert sorted(valid) == list(range(16)) + assert sum(index is None for rank_indices in assignments + for index in rank_indices) == 2 diff --git a/fastvideo/tests/train/callbacks/test_posthoc_ema.py b/fastvideo/tests/train/callbacks/test_posthoc_ema.py new file mode 100644 index 000000000..8f8c84628 --- /dev/null +++ b/fastvideo/tests/train/callbacks/test_posthoc_ema.py @@ -0,0 +1,102 @@ +# SPDX-License-Identifier: Apache-2.0 + +from pathlib import Path +from types import SimpleNamespace + +import torch + +from fastvideo.train.callbacks.posthoc_ema import ( + PostHocEMACallback, + sigma_rel_to_gamma, + solve_posthoc_weights, +) + + +class _FakeMethod: + + def __init__(self) -> None: + self.student = SimpleNamespace( + transformer=torch.nn.Linear(2, 2, bias=False), + ) + + +def test_posthoc_math_is_finite() -> None: + gamma_005 = sigma_rel_to_gamma(0.05) + gamma_01 = sigma_rel_to_gamma(0.1) + weights = solve_posthoc_weights( + torch.tensor([5, 5, 10, 10]), + torch.tensor([gamma_005, gamma_01, gamma_005, gamma_01]), + 10, + gamma_005, + ) + assert weights.shape == (4, ) + assert torch.isfinite(weights).all() + + +def test_posthoc_update_matches_nitrous_ema_step_order(tmp_path: Path) -> None: + callback = PostHocEMACallback( + checkpoint_every=100, + checkpoint_folder=str(tmp_path), + ) + callback.training_config = SimpleNamespace( + checkpoint=SimpleNamespace(output_dir=str(tmp_path)), + ) + method = _FakeMethod() + callback.on_train_start(method) + + first_weight = method.student.transformer.weight.detach().float().clone() + callback.on_training_step_end(method, {}, iteration=1) + torch.testing.assert_close(callback._ema_models[0].shadow["weight"], first_weight) + + with torch.no_grad(): + method.student.transformer.weight.add_(1) + second_weight = method.student.transformer.weight.detach().float().clone() + callback.on_training_step_end(method, {}, iteration=2) + + gamma = sigma_rel_to_gamma(0.05) + # nitrous-ema increments KarrasEMA.step before reading beta. + official_beta = ((1 - 1 / 3)**(1 + gamma) * + (1 - 1 / 4)**(1 + gamma))**0.5 + expected = first_weight * official_beta + second_weight * (1 - official_beta) + torch.testing.assert_close(callback._ema_models[0].shadow["weight"], expected) + + +def test_posthoc_callback_snapshots_and_synthesizes(tmp_path: Path) -> None: + callback = PostHocEMACallback( + checkpoint_every=2, + checkpoint_folder=str(tmp_path), + ) + callback.training_config = SimpleNamespace( + checkpoint=SimpleNamespace(output_dir=str(tmp_path)), + ) + method = _FakeMethod() + callback.on_train_start(method) + + callback.on_training_step_end(method, {}, iteration=1) + with torch.no_grad(): + method.student.transformer.weight.add_(1) + callback.on_training_step_end(method, {}, iteration=2) + + assert len(list((tmp_path / "rank_00000").glob("*.pt"))) == 2 + synthesized = callback.synthesize_local_shard() + assert synthesized is not None + assert set(synthesized) == {"weight"} + assert torch.isfinite(synthesized["weight"]).all() + + state = callback.state_dict() + restored = PostHocEMACallback( + checkpoint_every=2, + checkpoint_folder=str(tmp_path), + ) + restored.training_config = callback.training_config + restored.on_train_start(method) + restored.load_state_dict(state) + assert restored.state_dict()["calls"] == 2 + + callback.on_train_end(method, iteration=2) + final_paths = list((tmp_path / "rank_00000").glob( + "synthesized_sigma_0p05_step_*.pt", + )) + assert len(final_paths) == 1 + final_state = torch.load(final_paths[0], weights_only=True) + assert set(final_state) == {"weight"} diff --git a/fastvideo/tests/train/methods/test_mmaudio_flow_matching.py b/fastvideo/tests/train/methods/test_mmaudio_flow_matching.py new file mode 100644 index 000000000..c344ce14d --- /dev/null +++ b/fastvideo/tests/train/methods/test_mmaudio_flow_matching.py @@ -0,0 +1,491 @@ +# SPDX-License-Identifier: Apache-2.0 +"""CPU unit tests for the MMAudio flow-matching training adapter.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch +from torch import nn + +from fastvideo.configs.models.dits import MMAUDIO_VARIANT_ARCHITECTURES +from fastvideo.configs.pipelines.mmaudio import MMAudioSmall44kV2AConfig +from fastvideo.models.dits.mmaudio import MMAudioTransformer +from fastvideo.train.methods.fine_tuning.flow_matching import ( + FlowMatchingFineTuneMethod, +) +from fastvideo.train.models.mmaudio import MMAudioModel +from fastvideo.train.models.mmaudio.mmaudio import _load_or_compute_latent_stats +from fastvideo.train.utils.config import RunConfig +from fastvideo.train.utils.training_config import ( + DataConfig, + DistributedConfig, + ModelTrainingConfig, + OptimizerConfig, + TrainingConfig, + TrainingLoopConfig, +) + + +class _TinyMMAudioTransformer(nn.Module): + latent_seq_len = 4 + latent_dim = 3 + clip_seq_len = 6 + sync_seq_len = 9 + + def __init__(self, *, v2: bool = False) -> None: + super().__init__() + self.v2 = v2 + self.unshard_calls = 0 + arch = SimpleNamespace( + clip_dim=8, + sync_dim=10, + text_seq_len=5, + text_dim=7, + ) + self.config = SimpleNamespace(arch_config=arch) + self.proj = nn.Linear(self.latent_dim, self.latent_dim, bias=False) + self.latent_mean = nn.Parameter( + torch.tensor([[[0.5, -0.25, 1.0]]]), + requires_grad=False, + ) + self.latent_std = nn.Parameter( + torch.tensor([[[2.0, 0.5, 4.0]]]), + requires_grad=False, + ) + self.empty_string_feat = nn.Parameter( + torch.full((5, 7), 13.0), + requires_grad=False, + ) + self.empty_clip_feat = nn.Parameter(torch.full((1, 8), 11.0)) + self.empty_sync_feat = nn.Parameter(torch.full((1, 10), 12.0)) + + def unshard(self) -> None: + self.unshard_calls += 1 + + def normalize(self, value: torch.Tensor) -> torch.Tensor: + return value.sub(self.latent_mean).div(self.latent_std) + + 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_string_sequence(self, batch_size: int) -> torch.Tensor: + return self.empty_string_feat.unsqueeze(0).expand(batch_size, -1, -1) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: dict[str, torch.Tensor], + timestep: torch.Tensor, + ) -> torch.Tensor: + del timestep + condition_dependency = ( + encoder_hidden_states["clip_features"].sum() + + encoder_hidden_states["sync_features"].sum() + ) * 0.0 + return self.proj(hidden_states) + condition_dependency + + +def _training_config() -> TrainingConfig: + return TrainingConfig( + data=DataConfig( + data_path="", + preprocessed_data_type="mmaudio_features", + train_batch_size=2, + training_cfg_rate=0.4, + seed=123, + ), + optimizer=OptimizerConfig( + learning_rate=1e-4, + betas=(0.9, 0.95), + eps=1e-6, + ), + loop=TrainingLoopConfig(max_train_steps=2), + ) + + +def _raw_batch() -> dict[str, torch.Tensor]: + return { + "audio_latent_mean": ( + torch.arange(24, dtype=torch.float32).reshape(2, 4, 3) / 10 + ), + "audio_latent_std": torch.full((2, 4, 3), 0.2), + "clip_features": torch.zeros(2, 6, 8), + "sync_features": torch.zeros(2, 9, 10), + "text_features": torch.zeros(2, 5, 7), + "video_exists": torch.tensor([False, True]), + "text_exists": torch.tensor([True, False]), + } + + +def test_prepare_batch_matches_official_rng_and_flow_contract() -> None: + config = _training_config() + transformer = _TinyMMAudioTransformer() + model = MMAudioModel( + init_from="unused", + training_config=config, + transformer=transformer, + ) + raw = _raw_batch() + actual_generator = torch.Generator().manual_seed(7) + expected_generator = torch.Generator().manual_seed(7) + + batch = model.prepare_batch(raw, generator=actual_generator) + + posterior_noise = torch.empty_like(raw["audio_latent_mean"]).normal_( + generator=expected_generator + ) + clean = raw["audio_latent_mean"] + raw["audio_latent_std"] * posterior_noise + clean = (clean - transformer.latent_mean) / transformer.latent_std + timestep = torch.sigmoid( + torch.randn(2, generator=expected_generator) + * config.model.logit_std + + config.model.logit_mean + ) + prior = torch.empty_like(clean).normal_(generator=expected_generator) + noisy = (1 - timestep[:, None, None]) * prior + timestep[:, None, None] * clean + null_video = torch.rand(2, generator=expected_generator) < 0.4 + null_text = torch.rand(2, generator=expected_generator) < 0.4 + + torch.testing.assert_close(batch.latents, clean) + torch.testing.assert_close(batch.timesteps, timestep) + torch.testing.assert_close(batch.noise, prior) + torch.testing.assert_close(batch.noisy_model_input, noisy) + torch.testing.assert_close(batch.training_target, clean - prior) + assert transformer.unshard_calls == 1 + + assert batch.conditional_dict is not None + expected_video_null = torch.tensor([True, False]) | null_video + expected_text_null = torch.tensor([False, True]) | null_text + clip = batch.conditional_dict["clip_features"] + sync = batch.conditional_dict["sync_features"] + text = batch.conditional_dict["text_features"] + torch.testing.assert_close( + clip[expected_video_null], + torch.full_like(clip[expected_video_null], 11.0), + ) + torch.testing.assert_close( + sync[expected_video_null], + torch.full_like(sync[expected_video_null], 12.0), + ) + torch.testing.assert_close( + text[expected_text_null], + torch.full_like(text[expected_text_null], 13.0), + ) + + assert transformer.latent_mean.requires_grad is False + assert transformer.latent_std.requires_grad is False + assert transformer.empty_string_feat.requires_grad is False + assert transformer.empty_clip_feat.requires_grad is True + assert transformer.empty_sync_feat.requires_grad is True + + +def test_sequence_length_update_in_inference_mode_keeps_normal_buffers() -> None: + owner = SimpleNamespace() + inference_modes: list[bool] = [] + + def initialize_rotations() -> None: + inference_modes.append(torch.is_inference_mode_enabled()) + owner.latent_rot = torch.ones(1) + owner.clip_rot = torch.ones(1) + + owner.initialize_rotations = initialize_rotations + + with torch.inference_mode(): + MMAudioTransformer.update_seq_lengths(owner, 4, 6, 9) + + assert owner._latent_seq_len == 4 + assert owner._clip_seq_len == 6 + assert owner._sync_seq_len == 9 + assert inference_modes == [False] + assert not torch.is_inference(owner.latent_rot) + assert not torch.is_inference(owner.clip_rot) + + +def test_flow_matching_method_runs_loss_and_backward( + monkeypatch: pytest.MonkeyPatch, +) -> None: + config = _training_config() + monkeypatch.setattr( + "fastvideo.dataset.mmaudio_feature_dataset." + "build_mmaudio_feature_dataloader", + lambda *args, **kwargs: None, + ) + model = MMAudioModel( + init_from="unused", + training_config=config, + transformer=_TinyMMAudioTransformer(), + ) + cfg = RunConfig( + models={}, + method={ + "_target_": ( + "fastvideo.train.methods.fine_tuning.flow_matching." + "FlowMatchingFineTuneMethod" + ) + }, + training=config, + callbacks={}, + raw={}, + ) + method = FlowMatchingFineTuneMethod( + cfg=cfg, + role_models={"student": model}, + ) + method.cuda_generator = torch.Generator().manual_seed(17) + + eager_parameters = model.eager_optimizer_state_parameters() + assert len(eager_parameters) == 2 + for parameter in eager_parameters: + assert set(method._student_optimizer.state[parameter]) == { + "step", + "exp_avg", + "exp_avg_sq", + } + + loss_map, outputs, metrics = method.single_train_step(_raw_batch(), 0) + + assert torch.isfinite(loss_map["total_loss"]) + assert loss_map["flow_matching_loss"] is loss_map["total_loss"] + assert metrics == {} + method.backward(loss_map, outputs) + assert model.transformer.proj.weight.grad is not None + assert torch.count_nonzero(model.transformer.proj.weight.grad) > 0 + assert model.transformer.empty_clip_feat.grad is not None + assert model.transformer.empty_sync_feat.grad is not None + + +def test_flow_matching_method_compiles_inner_transformer_forward( + monkeypatch: pytest.MonkeyPatch, +) -> None: + config = _training_config() + config.model = ModelTrainingConfig( + compile_train_fn=True, + torch_compile_kwargs={"backend": "eager"}, + ) + monkeypatch.setattr( + "fastvideo.dataset.mmaudio_feature_dataset." + "build_mmaudio_feature_dataloader", + lambda *args, **kwargs: None, + ) + compile_calls: list[tuple[object, dict[str, object]]] = [] + + def _compile(fn, **kwargs): + compile_calls.append((fn, kwargs)) + return fn + + monkeypatch.setattr(torch, "compile", _compile) + model = MMAudioModel( + init_from="unused", + training_config=config, + transformer=_TinyMMAudioTransformer(), + ) + config.distributed = DistributedConfig(strategy="ddp") + cfg = RunConfig( + models={}, + method={ + "_target_": ( + "fastvideo.train.methods.fine_tuning.flow_matching." + "FlowMatchingFineTuneMethod" + ) + }, + training=config, + callbacks={}, + raw={}, + ) + method = FlowMatchingFineTuneMethod( + cfg=cfg, + role_models={"student": model}, + ) + method.cuda_generator = torch.Generator().manual_seed(19) + + loss_map, _, _ = method.single_train_step(_raw_batch(), 0) + + assert len(compile_calls) == 1 + assert compile_calls[0][1] == {"backend": "eager", "fullgraph": True} + assert torch.isfinite(loss_map["total_loss"]) + + +def test_compile_train_fn_rejects_fsdp( + monkeypatch: pytest.MonkeyPatch, +) -> None: + config = _training_config() + config.model = ModelTrainingConfig(compile_train_fn=True) + monkeypatch.setattr( + "fastvideo.dataset.mmaudio_feature_dataset." + "build_mmaudio_feature_dataloader", + lambda *args, **kwargs: None, + ) + model = MMAudioModel( + init_from="unused", + training_config=config, + transformer=_TinyMMAudioTransformer(), + ) + cfg = RunConfig( + models={}, + method={ + "_target_": ( + "fastvideo.train.methods.fine_tuning.flow_matching." + "FlowMatchingFineTuneMethod" + ) + }, + training=config, + callbacks={}, + raw={}, + ) + + with pytest.raises(ValueError, match="strategy=ddp"): + FlowMatchingFineTuneMethod( + cfg=cfg, + role_models={"student": model}, + ) + + +def test_v2_training_is_rejected_by_default() -> None: + with pytest.raises(ValueError, match="does not support `_v2`"): + MMAudioModel( + init_from="unused", + training_config=_training_config(), + transformer=_TinyMMAudioTransformer(v2=True), + ) + + +def test_official_mmaudio_variant_architectures() -> None: + expected = { + "small_16k": (20, 448, 12, 8, 7, 250, False), + "small_44k": (40, 448, 12, 8, 7, 345, False), + "medium_44k": (40, 896, 12, 8, 14, 345, False), + "large_44k": (40, 896, 21, 14, 14, 345, False), + "large_44k_v2": (40, 896, 21, 14, 14, 345, True), + } + actual = { + name: ( + values["latent_dim"], + values["hidden_dim"], + values["depth"], + values["fused_depth"], + values["num_heads"], + values["latent_seq_len"], + values["v2"], + ) + for name, values in MMAUDIO_VARIANT_ARCHITECTURES.items() + } + assert actual == expected + + +def test_from_scratch_builds_variant_without_loading_dit_checkpoint( + monkeypatch: pytest.MonkeyPatch, +) -> None: + captured: dict[str, object] = {} + + monkeypatch.setattr( + "fastvideo.train.models.mmaudio.mmaudio._load_or_compute_latent_stats", + lambda *args, **kwargs: ( + torch.zeros(1, 1, 40), + torch.ones(1, 1, 40), + ), + ) + monkeypatch.setattr( + "fastvideo.train.models.mmaudio.mmaudio._load_empty_string_features", + lambda *args, **kwargs: torch.zeros(77, 1024), + ) + + def _build(**kwargs): + captured.update(kwargs) + return _TinyMMAudioTransformer() + + monkeypatch.setattr( + "fastvideo.train.models.mmaudio.mmaudio.build_fsdp_model_from_scratch", + _build, + ) + model = MMAudioModel( + training_config=_training_config(), + variant="small_44k", + from_scratch=True, + empty_string_features_path="unused.safetensors", + ) + + assert model._init_from == "scratch:small_44k" + assert isinstance(model.training_config.pipeline_config, MMAudioSmall44kV2AConfig) + assert captured["model_cls"].__name__ == "MMAudioTransformer" + assert captured["init_params"]["hf_config"]["v2"] is False + + +def test_from_scratch_rejects_v2_variant_before_loading_assets() -> None: + with pytest.raises(ValueError, match="does not train 'large_44k_v2'"): + MMAudioModel( + training_config=_training_config(), + variant="large_44k_v2", + from_scratch=True, + empty_string_features_path="unused.safetensors", + ) + + +def test_44k_variants_share_cached_latent_statistics( + monkeypatch: pytest.MonkeyPatch, +) -> None: + payload = { + "variant": "small_44k", + "latent_mean": torch.zeros(1, 1, 40), + "latent_std": torch.ones(1, 1, 40), + } + + class _ReceiverGroup: + is_first_rank = False + + @staticmethod + def broadcast_object(value, src=0): + del value, src + return payload + + monkeypatch.setattr( + "fastvideo.train.models.mmaudio.mmaudio.get_world_group", + lambda: _ReceiverGroup(), + ) + mean, std = _load_or_compute_latent_stats( + _training_config(), + variant="medium_44k", + cache_path=None, + chunk_size=1, + ) + + torch.testing.assert_close(mean, payload["latent_mean"]) + torch.testing.assert_close(std, payload["latent_std"]) + + +def test_16k_rejects_44k_cached_latent_statistics( + monkeypatch: pytest.MonkeyPatch, +) -> None: + payload = { + "variant": "small_44k", + "latent_mean": torch.zeros(1, 1, 40), + "latent_std": torch.ones(1, 1, 40), + } + + class _ReceiverGroup: + is_first_rank = False + + @staticmethod + def broadcast_object(value, src=0): + del value, src + return payload + + monkeypatch.setattr( + "fastvideo.train.models.mmaudio.mmaudio.get_world_group", + lambda: _ReceiverGroup(), + ) + with pytest.raises(ValueError, match="small_44k.*small_16k"): + _load_or_compute_latent_stats( + _training_config(), + variant="small_16k", + cache_path=None, + chunk_size=1, + ) diff --git a/fastvideo/tests/train/utils/test_config.py b/fastvideo/tests/train/utils/test_config.py index a15c52027..380a012df 100644 --- a/fastvideo/tests/train/utils/test_config.py +++ b/fastvideo/tests/train/utils/test_config.py @@ -54,6 +54,7 @@ def test_minimal_yaml_applies_all_defaults(tmp_path: Path) -> None: cfg = load_run_config(_write_yaml(tmp_path, _minimal_yaml())) t = cfg.training + assert t.distributed.strategy == "fsdp" assert t.distributed.num_gpus == 1 assert t.distributed.tp_size == 1 assert t.distributed.sp_size == 1 @@ -67,9 +68,13 @@ def test_minimal_yaml_applies_all_defaults(tmp_path: Path) -> None: assert t.optimizer.learning_rate == 0.0 assert t.optimizer.betas == (0.9, 0.999) + assert t.optimizer.eps == pytest.approx(1e-8) assert t.optimizer.weight_decay == 0.0 assert t.optimizer.lr_scheduler == "constant" assert t.optimizer.min_lr_ratio == 0.5 + assert t.optimizer.lr_milestones == () + assert t.optimizer.lr_gamma == pytest.approx(0.1) + assert t.optimizer.fused is False assert t.loop.max_train_steps == 0 assert t.loop.gradient_accumulation_steps == 1 @@ -87,6 +92,8 @@ def test_minimal_yaml_applies_all_defaults(tmp_path: Path) -> None: assert t.model.weighting_scheme == "uniform" assert t.model.precondition_outputs is False assert t.model.moba_config == {} + assert t.model.compile_train_fn is False + assert t.model.torch_compile_kwargs == {} assert t.dit_precision == "fp32" assert t.vsa_sparsity == 0.0 @@ -106,6 +113,7 @@ def test_full_yaml_populates_all_training_fields(tmp_path: Path) -> None: }, "data": { "data_path": "/some/path", + "preprocessed_data_type": "mmaudio_features", "train_batch_size": 2, "dataloader_num_workers": 4, "training_cfg_rate": 0.1, @@ -118,10 +126,13 @@ def test_full_yaml_populates_all_training_fields(tmp_path: Path) -> None: "optimizer": { "learning_rate": 1e-4, "betas": [0.9, 0.95], + "eps": 1e-6, "weight_decay": 0.01, "lr_scheduler": "cosine", "lr_warmup_steps": 100, "min_lr_ratio": 0.1, + "lr_milestones": [240, 270], + "lr_gamma": 0.2, }, "loop": { "max_train_steps": 1000, @@ -149,6 +160,11 @@ def test_full_yaml_populates_all_training_fields(tmp_path: Path) -> None: "logit_mean": 0.5, "logit_std": 1.5, "precondition_outputs": True, + "compile_train_fn": True, + "torch_compile_kwargs": { + "backend": "eager", + "dynamic": False, + }, }, "dit_precision": "bf16", } @@ -161,14 +177,18 @@ def test_full_yaml_populates_all_training_fields(tmp_path: Path) -> None: assert t.data.train_batch_size == 2 assert t.data.data_path == "/some/path" + assert t.data.preprocessed_data_type == "mmaudio_features" assert t.data.num_frames == 33 assert t.data.seed == 42 assert t.optimizer.learning_rate == pytest.approx(1e-4) assert t.optimizer.betas == (0.9, 0.95) + assert t.optimizer.eps == pytest.approx(1e-6) assert t.optimizer.lr_scheduler == "cosine" assert t.optimizer.lr_warmup_steps == 100 assert t.optimizer.min_lr_ratio == pytest.approx(0.1) + assert t.optimizer.lr_milestones == (240, 270) + assert t.optimizer.lr_gamma == pytest.approx(0.2) assert t.loop.max_train_steps == 1000 assert t.loop.gradient_accumulation_steps == 4 @@ -185,6 +205,11 @@ def test_full_yaml_populates_all_training_fields(tmp_path: Path) -> None: assert t.vsa_sparsity == pytest.approx(0.5) assert t.model.weighting_scheme == "logit_normal" assert t.model.precondition_outputs is True + assert t.model.compile_train_fn is True + assert t.model.torch_compile_kwargs == { + "backend": "eager", + "dynamic": False, + } assert t.dit_precision == "bf16" @@ -260,6 +285,20 @@ def test_betas_parses_list_and_string_forms( assert cfg.training.optimizer.betas == expected +def test_distributed_strategy_parses_ddp(tmp_path: Path) -> None: + data = _minimal_yaml() + data["training"] = {"distributed": {"strategy": "DDP"}} + cfg = load_run_config(_write_yaml(tmp_path, data)) + assert cfg.training.distributed.strategy == "ddp" + + +def test_distributed_strategy_rejects_unknown_value(tmp_path: Path) -> None: + data = _minimal_yaml() + data["training"] = {"distributed": {"strategy": "zero3"}} + with pytest.raises(ValueError, match="strategy must be one of"): + load_run_config(_write_yaml(tmp_path, data)) + + def test_pipeline_quant_config_resolves_for_load_and_export(tmp_path: Path) -> None: from fastvideo.layers.quantization.nvfp4_qat_train_config import ( NVFP4QATTrainConfig, ) diff --git a/fastvideo/tests/train/utils/test_distributed_strategy.py b/fastvideo/tests/train/utils/test_distributed_strategy.py new file mode 100644 index 000000000..a73d50cb5 --- /dev/null +++ b/fastvideo/tests/train/utils/test_distributed_strategy.py @@ -0,0 +1,99 @@ +# SPDX-License-Identifier: Apache-2.0 +"""CPU-only coverage for the modular trainer's DDP strategy.""" + +from __future__ import annotations + +from pathlib import Path + +import pytest +import torch +import torch.distributed as dist + +from fastvideo.train.utils.distributed_strategy import ( + build_replicated_model_from_scratch, + normalize_distributed_strategy, + unwrap_ddp_module, + wrap_module_ddp, +) +from fastvideo.training.checkpointing_utils import ModelWrapper + + +class _TinyModel(torch.nn.Module): + + def __init__(self, width: int = 2) -> None: + super().__init__() + self.linear = torch.nn.Linear(width, width) + self.public_value = 17 + + def forward(self, value: torch.Tensor) -> torch.Tensor: + return self.linear(value) + + def helper(self) -> int: + return self.public_value + + +@pytest.fixture +def local_process_group(tmp_path: Path): + owns_group = not dist.is_initialized() + if owns_group: + rendezvous = tmp_path / "ddp_rendezvous" + dist.init_process_group( + "gloo", + init_method=f"file://{rendezvous}", + rank=0, + world_size=1, + ) + try: + yield + finally: + if owns_group: + dist.destroy_process_group() + + +def test_strategy_normalization() -> None: + assert normalize_distributed_strategy(None) == "fsdp" + assert normalize_distributed_strategy(" DDP ") == "ddp" + with pytest.raises(ValueError, match="strategy must be one of"): + normalize_distributed_strategy("deepspeed") + + +def test_replicated_initialization_is_seeded() -> None: + first = build_replicated_model_from_scratch( + _TinyModel, + {"width": 3}, + device=torch.device("cpu"), + default_dtype=torch.float32, + seed=123, + ) + second = build_replicated_model_from_scratch( + _TinyModel, + {"width": 3}, + device=torch.device("cpu"), + default_dtype=torch.float32, + seed=123, + ) + for left, right in zip(first.parameters(), second.parameters(), strict=True): + torch.testing.assert_close(left, right) + + +def test_ddp_adapter_forward_delegation_and_checkpoint( + local_process_group, +) -> None: + del local_process_group + wrapped = wrap_module_ddp( + _TinyModel(), + device=torch.device("cpu"), + broadcast_buffers=False, + ) + assert wrapped.helper() == 17 + assert unwrap_ddp_module(wrapped) is wrapped.module + + output = wrapped(torch.ones(2, 2)).sum() + output.backward() + assert all(parameter.grad is not None for parameter in wrapped.parameters()) + + # torch.distributed.checkpoint canonicalizes DDP's ``module.`` prefix. + # ModelWrapper must apply the same canonicalization when filtering out + # frozen parameters, otherwise it silently emits an empty checkpoint. + state = ModelWrapper(wrapped).state_dict() + assert set(state) == {"linear.weight", "linear.bias"} diff --git a/fastvideo/tests/train/utils/test_optimizer.py b/fastvideo/tests/train/utils/test_optimizer.py new file mode 100644 index 000000000..740d76750 --- /dev/null +++ b/fastvideo/tests/train/utils/test_optimizer.py @@ -0,0 +1,106 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Tests for FastVideo modular trainer optimizer construction.""" + +import pytest +import torch + +from fastvideo.train.utils.optimizer import ( + build_optimizer_and_scheduler, + seed_adamw_parameter_state, +) +from fastvideo.train.utils.training_config import OptimizerConfig, TrainingLoopConfig +from fastvideo.training.checkpointing_utils import SchedulerWrapper + + +def test_multistep_with_warmup_matches_mmaudio_scheduler_composition() -> None: + parameter = torch.nn.Parameter(torch.zeros(())) + config = OptimizerConfig( + learning_rate=1e-4, + lr_scheduler="multistep_with_warmup", + lr_warmup_steps=2, + lr_milestones=(3, 5), + lr_gamma=0.1, + ) + optimizer, scheduler = build_optimizer_and_scheduler( + params=[parameter], + optimizer_config=config, + loop_config=TrainingLoopConfig(max_train_steps=10), + learning_rate=config.learning_rate, + betas=config.betas, + scheduler_name=config.lr_scheduler, + ) + + assert isinstance(scheduler, torch.optim.lr_scheduler.SequentialLR) + initial_lr = optimizer.param_groups[0]["lr"] + assert initial_lr == pytest.approx(1e-4 / 3) + + values = [] + for _ in range(8): + optimizer.step() + scheduler.step() + values.append(optimizer.param_groups[0]["lr"]) + + assert values[1] == pytest.approx(1e-4) + assert values[4] == pytest.approx(1e-5) + assert values[6] == pytest.approx(1e-6) + + +def test_multistep_with_warmup_requires_milestones() -> None: + parameter = torch.nn.Parameter(torch.zeros(())) + with pytest.raises(ValueError, match="lr_milestones"): + build_optimizer_and_scheduler( + params=[parameter], + optimizer_config=OptimizerConfig( + learning_rate=1e-4, + lr_scheduler="multistep_with_warmup", + lr_warmup_steps=2, + ), + loop_config=TrainingLoopConfig(max_train_steps=10), + learning_rate=1e-4, + betas=(0.9, 0.95), + scheduler_name="multistep_with_warmup", + ) + + +@pytest.mark.parametrize("fused", [False, True]) +def test_seed_adamw_state_uses_expected_step_device(fused: bool) -> None: + parameter = torch.nn.Parameter(torch.zeros(())) + optimizer = torch.optim.AdamW([parameter], fused=fused) + seed_adamw_parameter_state(optimizer, parameter) + expected = parameter.device if fused else torch.device("cpu") + assert optimizer.state[parameter]["step"].device == expected + + +def test_seed_adamw_state_rejects_unowned_parameter() -> None: + owned = torch.nn.Parameter(torch.zeros(())) + unowned = torch.nn.Parameter(torch.zeros(())) + optimizer = torch.optim.AdamW([owned]) + with pytest.raises(ValueError, match="unowned"): + seed_adamw_parameter_state(optimizer, unowned) + + +def test_scheduler_wrapper_stabilizes_multistep_milestone_keys() -> None: + parameter = torch.nn.Parameter(torch.zeros(())) + config = OptimizerConfig( + learning_rate=1e-4, + lr_scheduler="multistep_with_warmup", + lr_warmup_steps=2, + lr_milestones=(3, 5), + ) + _, scheduler = build_optimizer_and_scheduler( + params=[parameter], + optimizer_config=config, + loop_config=TrainingLoopConfig(max_train_steps=10), + learning_rate=config.learning_rate, + betas=config.betas, + scheduler_name=config.lr_scheduler, + ) + wrapper = SchedulerWrapper(scheduler) + state = wrapper.state_dict() + milestones = state["scheduler"]["_schedulers"][1]["milestones"] + assert milestones == {"3": 1, "5": 1} + + # Model what DCP returns after flattening mapping keys to strings. + wrapper.load_state_dict(state) + restored = scheduler.state_dict()["_schedulers"][1]["milestones"] + assert restored == {3: 1, 5: 1} diff --git a/fastvideo/train/callbacks/ddp_posthoc_ema.py b/fastvideo/train/callbacks/ddp_posthoc_ema.py new file mode 100644 index 000000000..88758c3ec --- /dev/null +++ b/fastvideo/train/callbacks/ddp_posthoc_ema.py @@ -0,0 +1,173 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Official nitrous PostHocEMA lifecycle for native DDP training.""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any, TYPE_CHECKING + +import torch +import torch.distributed as dist +from torch.nn.parallel import DistributedDataParallel + +from fastvideo.logger import init_logger +from fastvideo.train.callbacks.callback import Callback +from fastvideo.train.utils.distributed_strategy import unwrap_ddp_module + +if TYPE_CHECKING: + from fastvideo.train.methods.base import TrainingMethod + +logger = init_logger(__name__) + + +class DDPPostHocEMACallback(Callback): + """Run the exact ``nitrous_ema.PostHocEMA`` used by MMAudio. + + Official MMAudio owns two complete CUDA EMA models on local rank zero and + checkpoints them every 5,000 optimizer steps. Other DDP ranks do not hold + EMA copies. This callback preserves that lifecycle while the ordinary + FastVideo FSDP callback continues to maintain rank-local CPU shards. + """ + + def __init__( + self, + *, + sigma_rels: list[float] | tuple[float, ...] = (0.05, 0.1), + update_every: int = 1, + checkpoint_every: int = 5000, + checkpoint_folder: str | None = None, + start_iter: int = 0, + default_output_sigma: float = 0.05, + step_size_correction: bool = True, + ) -> None: + self.sigma_rels = tuple(float(value) for value in sigma_rels) + self.update_every = max(1, int(update_every)) + self.checkpoint_every = max(1, int(checkpoint_every)) + self.checkpoint_folder = checkpoint_folder + self.start_iter = int(start_iter) + self.default_output_sigma = float(default_output_sigma) + self.step_size_correction = bool(step_size_correction) + + self._ema: Any | None = None + self._calls = 0 + self._rank = 0 + self._checkpoint_root: Path | None = None + + def on_train_start( + self, + method: TrainingMethod, + iteration: int = 0, + ) -> None: + del iteration + transformer = method.student.transformer + if not isinstance(transformer, DistributedDataParallel): + raise TypeError("DDPPostHocEMACallback requires a native DDP transformer") + + self._rank = dist.get_rank() if dist.is_initialized() else 0 + root = self.checkpoint_folder + if not root: + root = str(Path(self.training_config.checkpoint.output_dir) / "posthoc_ema" / "official_ddp") + self._checkpoint_root = Path(root).expanduser().resolve() + + if self._rank == 0: + try: + from nitrous_ema import PostHocEMA + except ImportError as exc: + raise ImportError("DDP MMAudio training requires nitrous-ema. Install " + "FastVideo with the mmaudio-train extra.") from exc + + module = unwrap_ddp_module(transformer) + self._ema = PostHocEMA( + module, + sigma_rels=self.sigma_rels, + update_every=self.update_every, + checkpoint_every_num_steps=self.checkpoint_every, + checkpoint_folder=str(self._checkpoint_root), + step_size_correction=self.step_size_correction, + ).to(method.student.device) + logger.info( + "Official DDP PostHocEMA enabled on rank 0: " + "sigma_rels=%s checkpoint_every=%d folder=%s", + self.sigma_rels, + self.checkpoint_every, + self._checkpoint_root, + ) + + def on_training_step_end( + self, + method: TrainingMethod, + loss_dict: dict[str, Any], + iteration: int = 0, + ) -> None: + del method, loss_dict + if iteration < self.start_iter: + return + self._calls += 1 + if self._rank == 0: + if self._ema is None: + raise RuntimeError("DDP PostHocEMA was not initialized") + self._ema.update() + if int(self._ema.step) != self._calls: + raise RuntimeError("DDP PostHocEMA step diverged from the trainer step") + + def _restore_checkpoint(self) -> None: + if self._rank != 0 or self._ema is None or self._calls <= 0: + return + if self._checkpoint_root is None: + raise RuntimeError("DDP PostHocEMA checkpoint folder is unset") + for index, ema_model in enumerate(self._ema.ema_models): + path = self._checkpoint_root / f"{index}.{self._calls}.pt" + if not path.is_file(): + raise FileNotFoundError("Cannot resume DDP PostHocEMA: expected matching snapshot " + f"{path}. Keep training checkpoints aligned with " + "checkpoint_every.") + state = torch.load( + path, + map_location=method_device(ema_model), + weights_only=True, + ) + ema_model.load_state_dict(state) + logger.info("Restored official DDP PostHocEMA at step %d", self._calls) + + def state_dict(self) -> dict[str, Any]: + # Full nitrous states already live in rank-0 snapshots. Keeping only + # the common step in DCP gives every rank the same callback schema. + return {"calls": self._calls} + + def load_state_dict( + self, + state_dict: dict[str, Any], + ) -> None: + self._calls = int(state_dict.get("calls", 0)) + self._restore_checkpoint() + + def on_train_end( + self, + method: TrainingMethod, + iteration: int = 0, + ) -> None: + del method, iteration + if self._rank != 0 or self._ema is None or self._calls <= 0: + return + if self._calls % self.checkpoint_every != 0: + self._ema.checkpoint() + synthesized = self._ema.synthesize_ema_model( + sigma_rel=self.default_output_sigma, + step=self._calls, + device="cpu", + ) + if self._checkpoint_root is None: + raise RuntimeError("DDP PostHocEMA checkpoint folder is unset") + sigma_name = str(self.default_output_sigma).replace(".", "p") + path = self._checkpoint_root / (f"mmaudio_ema_final_sigma_{sigma_name}_step_" + f"{self._calls:09d}.pth") + torch.save(synthesized.ema_model.state_dict(), path) + logger.info("Saved official synthesized DDP PostHocEMA to %s", path) + + +def method_device(module: torch.nn.Module) -> torch.device: + """Return the device of a nitrous EMA module.""" + return next(module.parameters()).device + + +__all__ = ["DDPPostHocEMACallback"] diff --git a/fastvideo/train/callbacks/mmaudio_validation.py b/fastvideo/train/callbacks/mmaudio_validation.py new file mode 100644 index 000000000..7b5400772 --- /dev/null +++ b/fastvideo/train/callbacks/mmaudio_validation.py @@ -0,0 +1,509 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Validation loss and audio sampling for MMAudio feature training.""" + +from __future__ import annotations + +import contextlib +import math +import re +from pathlib import Path +from typing import Any, TYPE_CHECKING + +import numpy as np +import torch +import torch.nn.functional as F +from scipy.io import wavfile + +from fastvideo.distributed import get_local_torch_device, get_world_group +from fastvideo.logger import init_logger +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.train.callbacks.callback import Callback + +if TYPE_CHECKING: + from fastvideo.train.methods.base import TrainingMethod + +logger = init_logger(__name__) + + +def _global_inference_indices( + total_samples: int, + world_size: int, + rank: int, +) -> list[int | None]: + """Assign a fixed global sample budget with equal calls on every rank.""" + if total_samples <= 0 or world_size <= 0 or not 0 <= rank < world_size: + raise ValueError("Invalid distributed MMAudio inference assignment") + calls_per_rank = math.ceil(total_samples / world_size) + return [ + global_index if global_index < total_samples else None for local_index in range(calls_per_rank) + for global_index in (local_index * world_size + rank, ) + ] + + +def _inference_sample_index( + global_index: int | None, + local_index: int, + available: int, +) -> int: + """Wrap an assigned global sample onto a row of this rank's feature batch. + + ``available`` is the per-rank validation batch size, which is usually + smaller than the global sample budget. + """ + return (global_index if global_index is not None else local_index) % available + + +def _safe_sample_name(value: str) -> str: + value = re.sub(r"[^A-Za-z0-9._-]+", "_", value).strip("._") + return value or "sample" + + +class _IgnoredComponent: + """Placeholder for a required pipeline module that is never dereferenced. + + ``MMAudioPipeline`` requires the CLIP text/vision and Synchformer encoders, + but its direct-feature conditioning path returns before touching them. + """ + + def __repr__(self) -> str: + return "" + + +class MMAudioValidationCallback(Callback): + """Evaluate cached val features and periodically run native V2A inference. + + Validation loss follows the official MMAudio ``val_fn``: sample a VAE + posterior latent, a logit-normal flow time, prior noise, and independent + video/text CFG masks. The RNG is reset for every pass, making values at + different training steps directly comparable. + + Optional inference reuses the live FSDP transformer in + :class:`MMAudioPipeline`, while frozen VAE/vocoder weights are loaded from + ``inference_model_path``. Precomputed CLIP, Synchformer, and text features + go directly into the pipeline, so validation does not decode source video + or run feature encoders again. + """ + + def __init__( + self, + *, + data_path: str, + every_steps: int = 5000, + max_batches: int = 0, + batch_size: int = 8, + num_data_workers: int = 2, + run_at_start: bool = False, + use_ema: bool = False, + inference_every_steps: int = 20000, + inference_model_path: str = "", + inference_num_samples: int = 16, + inference_num_steps: int = 25, + inference_guidance_scale: float = 4.5, + inference_seed: int = 14159265, + inference_save_video: bool = True, + inference_log_to_tracker: bool = True, + output_dir: str | None = None, + ) -> None: + self.data_path = str(data_path) + self.every_steps = int(every_steps) + self.max_batches = max(0, int(max_batches)) + self.batch_size = max(1, int(batch_size)) + self.num_data_workers = max(0, int(num_data_workers)) + self.run_at_start = bool(run_at_start) + self.use_ema = bool(use_ema) + self.inference_every_steps = max(0, int(inference_every_steps)) + self.inference_model_path = str(inference_model_path) + self.inference_num_samples = max(1, int(inference_num_samples)) + self.inference_num_steps = max(1, int(inference_num_steps)) + self.inference_guidance_scale = float(inference_guidance_scale) + self.inference_seed = int(inference_seed) + self.inference_save_video = bool(inference_save_video) + self.inference_log_to_tracker = bool(inference_log_to_tracker) + self.output_dir = output_dir + + self._dataloader: Any = None + self._pipeline: Any = None + self._rank = 0 + + def on_train_start( + self, + method: TrainingMethod, + iteration: int = 0, + ) -> None: + del iteration + if not self.data_path: + raise ValueError("MMAudio validation requires callbacks.validation.data_path") + self._rank = get_world_group().rank + transformer = method.student.transformer + from fastvideo.dataset.mmaudio_feature_dataset import ( + build_mmaudio_feature_dataloader, ) + + feature_shapes = { + "latent_seq_len": int(transformer.latent_seq_len), + "latent_dim": int(transformer.latent_dim), + "clip_seq_len": int(transformer.clip_seq_len), + "clip_dim": int(transformer.config.arch_config.clip_dim), + "sync_seq_len": int(transformer.sync_seq_len), + "sync_dim": int(transformer.config.arch_config.sync_dim), + "text_seq_len": int(transformer.config.arch_config.text_seq_len), + "text_dim": int(transformer.config.arch_config.text_dim), + } + self._dataloader = build_mmaudio_feature_dataloader( + self.data_path, + batch_size=self.batch_size, + num_data_workers=self.num_data_workers, + seed=int(self.training_config.data.seed), + pin_memory=self.training_config.distributed.pin_cpu_memory, + feature_shapes=feature_shapes, + include_metadata=True, + ) + logger.info("Initialized MMAudio validation features from %s", self.data_path) + + def on_validation_begin( + self, + method: TrainingMethod, + iteration: int = 0, + ) -> None: + if iteration == 0 and not self.run_at_start: + return + run_loss = self.every_steps > 0 and iteration % self.every_steps == 0 + run_inference = (self.inference_every_steps > 0 and bool(self.inference_model_path) + and iteration % self.inference_every_steps == 0) + if not run_loss and not run_inference: + return + + transformer = method.student.transformer + was_training = transformer.training + transformer.eval() + try: + ema_context = self._ema_context(transformer) if self.use_ema else contextlib.nullcontext(transformer) + with ema_context: + validation_batch = None + if run_loss: + validation_batch = self._run_loss_validation(method, iteration) + if run_inference: + if validation_batch is None: + validation_batch = next(iter(self._dataloader)) + self._run_inference(method, validation_batch, iteration) + finally: + if was_training: + transformer.train() + + def _ema_context(self, transformer: torch.nn.Module) -> Any: + callback_dict = getattr(self, "_callback_dict", None) + if callback_dict is not None: + for callback in callback_dict._callbacks.values(): + if callback is self: + continue + context = getattr(callback, "ema_context", None) + if callable(context): + return context(transformer) + logger.warning("MMAudio validation requested EMA weights but no EMA callback was found") + return contextlib.nullcontext(transformer) + + @torch.inference_mode() + def _run_loss_validation( + self, + method: TrainingMethod, + iteration: int, + ) -> dict[str, Any] | None: + if self._dataloader is None: + raise RuntimeError("MMAudio validation dataloader is not initialized") + device = get_local_torch_device() + generator = torch.Generator(device=device).manual_seed(int(self.training_config.data.seed) + self._rank, ) + loss_sum = torch.zeros((), device=device, dtype=torch.float64) + sample_count = torch.zeros((), device=device, dtype=torch.float64) + first_batch: dict[str, Any] | None = None + + for batch_index, raw_batch in enumerate(self._dataloader): + if self.max_batches and batch_index >= self.max_batches: + break + if first_batch is None: + first_batch = raw_batch + training_batch = method.student.prepare_batch( + raw_batch, + generator=generator, + latents_source="data", + ) + noisy_latents = training_batch.noisy_model_input + timesteps = training_batch.timesteps + target = training_batch.training_target + if not all(isinstance(value, torch.Tensor) for value in (noisy_latents, timesteps, target)): + raise RuntimeError("MMAudio validation batch is incomplete") + assert isinstance(noisy_latents, torch.Tensor) + assert isinstance(timesteps, torch.Tensor) + assert isinstance(target, torch.Tensor) + prediction = method.student.predict_noise( + noisy_latents, + timesteps, + training_batch, + conditional=True, + attn_kind="dense", + ) + batch_size = int(noisy_latents.shape[0]) + loss_sum += F.mse_loss( + prediction.float(), + target.float(), + reduction="mean", + ).double() * batch_size + sample_count += batch_size + + if torch.distributed.is_initialized(): + torch.distributed.all_reduce(loss_sum) + torch.distributed.all_reduce(sample_count) + if sample_count.item() == 0: + raise RuntimeError("MMAudio validation did not produce any samples") + loss = float((loss_sum / sample_count).item()) + if self._rank == 0: + tracker = getattr(method, "tracker", None) + if tracker is not None: + tracker.log({"validation/flow_matching_loss": loss}, iteration) + logger.info("MMAudio validation step=%d loss=%.8f", iteration, loss) + return first_batch + + def _build_pipeline(self, transformer: torch.nn.Module) -> Any: + if self._pipeline is not None: + return self._pipeline + from fastvideo.pipelines.basic.mmaudio import MMAudioPipeline + + ignored_component = _IgnoredComponent() + self._pipeline = MMAudioPipeline.from_pretrained( + self.inference_model_path, + inference_mode=True, + loaded_modules={ + "transformer": transformer, + "text_encoder": ignored_component, + "tokenizer": ignored_component, + "image_encoder": ignored_component, + "image_encoder_2": ignored_component, + }, + workload_type="v2a", + num_gpus=int(self.training_config.distributed.num_gpus), + tp_size=int(self.training_config.distributed.tp_size), + sp_size=int(self.training_config.distributed.sp_size), + hsdp_shard_dim=int(self.training_config.distributed.hsdp_shard_dim), + pin_cpu_memory=bool(self.training_config.distributed.pin_cpu_memory), + dit_cpu_offload=False, + dit_layerwise_offload=False, + vae_cpu_offload=False, + ) + self._pipeline.fastvideo_args.pipeline_config = self.training_config.pipeline_config + for module_name in ("text_encoder", "tokenizer", "image_encoder", "image_encoder_2"): + module = self._pipeline.get_module(module_name) + assert isinstance(module, + _IgnoredComponent), (f"MMAudio validation placeholder for {module_name} was not used") + return self._pipeline + + @torch.inference_mode() + def _run_inference( + self, + method: TrainingMethod, + raw_batch: dict[str, Any], + iteration: int, + ) -> None: + pipeline = self._build_pipeline(method.student.transformer) + pc = self.training_config.pipeline_config + if pc is None: + raise RuntimeError("MMAudio inference validation requires pipeline_config") + + text_features = raw_batch.get("text_features") + if not isinstance(text_features, torch.Tensor): + raise RuntimeError("MMAudio validation inference requires text features") + available = int(text_features.shape[0]) + if available <= 0: + raise RuntimeError("MMAudio validation inference batch is empty") + world_size = (torch.distributed.get_world_size() if torch.distributed.is_initialized() else 1) + assignments = _global_inference_indices( + self.inference_num_samples, + world_size, + self._rank, + ) + output_root = Path(self.output_dir or + self.training_config.checkpoint.output_dir, ) / "validation_audio" / f"step_{iteration:09d}" + output_root.mkdir(parents=True, exist_ok=True) + + for local_index, global_index in enumerate(assignments): + # Every FSDP rank must execute the same number of forwards. Ranks + # outside a non-divisible global budget run one padded call but do + # not save it. + sample_index = _inference_sample_index(global_index, local_index, available) + output_index = (global_index if global_index is not None else self.inference_num_samples + self._rank) + batch = ForwardBatch( + data_type="video", + prompt="", + negative_prompt="", + audio_start_in_s=0.0, + audio_end_in_s=float(pc.duration_s), + num_inference_steps=self.inference_num_steps, + guidance_scale=self.inference_guidance_scale, + seed=self.inference_seed + output_index, + num_videos_per_prompt=1, + height=8, + width=8, + num_frames=1, + save_video=False, + return_frames=False, + extra={ + "mmaudio_clip_features": raw_batch["clip_features"][sample_index:sample_index + 1], + "mmaudio_sync_features": raw_batch["sync_features"][sample_index:sample_index + 1], + "mmaudio_text_features": raw_batch["text_features"][sample_index:sample_index + 1], + }, + ) + result = pipeline.forward(batch, pipeline.fastvideo_args) + audio = result.extra.get("audio") + sample_rate = result.extra.get("audio_sample_rate") + if not isinstance(audio, np.ndarray) or not isinstance(sample_rate, int): + raise RuntimeError("MMAudio validation inference did not produce decoded audio") + if global_index is None: + continue + sample_id = self._metadata_value(raw_batch, "sample_id", sample_index) + source_path = self._metadata_value(raw_batch, "source_path", sample_index) + stem = f"sample_{global_index:03d}_{_safe_sample_name(sample_id)}" + audio_path = output_root / f"{stem}.wav" + wavfile.write(audio_path, sample_rate, audio.astype(np.float32)) + logger.info("Saved MMAudio validation audio to %s", audio_path) + if self.inference_save_video and source_path: + self._mux_source_video( + source_path=Path(source_path), + output_path=output_root / f"{stem}.mp4", + audio=audio, + sample_rate=sample_rate, + duration_s=float(pc.duration_s), + ) + + if torch.distributed.is_initialized(): + torch.distributed.barrier() + if self._rank == 0 and self.inference_log_to_tracker: + self._log_inference_media(method, output_root, iteration) + + @staticmethod + def _metadata_value( + raw_batch: dict[str, Any], + key: str, + index: int, + ) -> str: + values = raw_batch.get(key) + if isinstance(values, list | tuple) and index < len(values): + return str(values[index]) + return "" + + @staticmethod + def _mux_source_video( + *, + source_path: Path, + output_path: Path, + audio: np.ndarray, + sample_rate: int, + duration_s: float, + ) -> Path | None: + if not source_path.is_file(): + logger.warning("Cannot compose MMAudio validation video; missing %s", source_path) + return None + try: + import av + from av import AudioFrame + + # Match official mmaudio.data.av_utils.read_frames: keep every + # decoded source frame through duration_s at the stream's guessed + # frame rate, including a frame exactly on the endpoint. + all_frames: list[np.ndarray] = [] + with av.open(str(source_path)) as source: + input_stream = source.streams.video[0] + frame_rate = input_stream.guessed_rate + input_stream.thread_type = "AUTO" + for packet in source.demux(input_stream): + reached_end = False + for frame in packet.decode(): + if frame.time is not None and frame.time > duration_s: + reached_end = True + break + all_frames.append(frame.to_ndarray(format="rgb24")) + if reached_end: + break + if not all_frames: + raise ValueError("source video produced no decoded frames") + + # Match official reencode_with_audio: H.264 at the original + # guessed FPS, 10 Mbps, yuv420p, followed by mono float audio + # encoded as AAC at the model sample rate. + output = av.open(str(output_path), "w") + try: + video_stream = output.add_stream("h264", frame_rate) + video_stream.codec_context.bit_rate = int(10 * 1e6) + video_stream.width = int(all_frames[0].shape[1]) + video_stream.height = int(all_frames[0].shape[0]) + video_stream.pix_fmt = "yuv420p" + audio_stream = output.add_stream("aac", sample_rate) + + for frame_array in all_frames: + video_frame = av.VideoFrame.from_ndarray(frame_array) + for packet in video_stream.encode(video_frame): + output.mux(packet) + for packet in video_stream.encode(): + output.mux(packet) + + audio_np = np.asarray(audio, dtype=np.float32) + if audio_np.ndim == 1: + audio_np = audio_np[None, :] + elif audio_np.ndim == 2 and audio_np.shape[1] == 1: + audio_np = audio_np.T + if audio_np.ndim != 2 or audio_np.shape[0] != 1: + raise ValueError(f"Expected mono audio, got {audio_np.shape}") + audio_frame = AudioFrame.from_ndarray( + np.ascontiguousarray(audio_np), + format="flt", + layout="mono", + ) + audio_frame.sample_rate = sample_rate + for packet in audio_stream.encode(audio_frame): + output.mux(packet) + for packet in audio_stream.encode(): + output.mux(packet) + finally: + output.close() + except Exception as exc: + logger.warning( + "Failed to compose MMAudio validation video %s: %s", + output_path, + exc, + ) + return None + logger.info("Saved MMAudio validation video to %s", output_path) + return output_path + + def _log_inference_media( + self, + method: TrainingMethod, + output_root: Path, + iteration: int, + ) -> None: + tracker = getattr(method, "tracker", None) + if tracker is None: + return + artifacts = [] + for path in sorted(output_root.glob("sample_*.mp4"))[:self.inference_num_samples]: + artifact = tracker.video( + str(path), + caption=path.stem, + format="mp4", + ) + if artifact is not None: + artifacts.append(artifact) + if artifacts: + tracker.log_artifacts( + {"validation/mmaudio_samples": artifacts}, + iteration, + ) + + def on_train_end( + self, + method: TrainingMethod, + iteration: int = 0, + ) -> None: + del method, iteration + if self._pipeline is not None: + self._pipeline.close() + self._pipeline = None + + +__all__ = ["MMAudioValidationCallback"] diff --git a/fastvideo/train/callbacks/posthoc_ema.py b/fastvideo/train/callbacks/posthoc_ema.py new file mode 100644 index 000000000..03952bf91 --- /dev/null +++ b/fastvideo/train/callbacks/posthoc_ema.py @@ -0,0 +1,296 @@ +# SPDX-License-Identifier: Apache-2.0 +"""FSDP2-compatible post-hoc EMA used by the official MMAudio recipe.""" + +from __future__ import annotations + +import contextlib +import os +from collections.abc import Generator +from pathlib import Path +from typing import Any, TYPE_CHECKING + +import numpy as np +import torch + +from fastvideo.logger import init_logger +from fastvideo.train.callbacks.callback import Callback +from fastvideo.training.training_utils import EMA_FSDP + +if TYPE_CHECKING: + from fastvideo.train.methods.base import TrainingMethod + +logger = init_logger(__name__) + + +def sigma_rel_to_gamma(sigma_rel: float) -> float: + """Algorithm 2 from Karras et al., matching ``nitrous-ema``.""" + if sigma_rel <= 0: + raise ValueError("Post-hoc EMA sigma_rel must be positive") + t = sigma_rel**-2 + return float(np.roots([1, 7, 16 - t, 12 - t]).real.max()) + + +def _p_dot_p( + t_a: torch.Tensor, + gamma_a: torch.Tensor, + t_b: torch.Tensor, + gamma_b: torch.Tensor, +) -> torch.Tensor: + t_ratio = t_a / t_b + t_exp = torch.where(t_a < t_b, gamma_b, -gamma_a) + t_max = torch.maximum(t_a, t_b) + numerator = (gamma_a + 1) * (gamma_b + 1) * t_ratio**t_exp + denominator = (gamma_a + gamma_b + 1) * t_max + return numerator / denominator + + +def solve_posthoc_weights( + timesteps: torch.Tensor, + gammas: torch.Tensor, + target_timestep: int, + target_gamma: float, +) -> torch.Tensor: + """Algorithm 3 from Karras et al., matching ``nitrous-ema``.""" + t_i = timesteps.double().reshape(-1, 1) + gamma_i = gammas.double().reshape(-1, 1) + t_j = timesteps.double().reshape(1, -1) + gamma_j = gammas.double().reshape(1, -1) + matrix = _p_dot_p(t_i, gamma_i, t_j, gamma_j) + target_t = torch.tensor([[target_timestep]], dtype=torch.float64) + target_g = torch.tensor([[target_gamma]], dtype=torch.float64) + rhs = _p_dot_p(t_i, gamma_i, target_t, target_g) + return torch.linalg.solve(matrix, rhs).squeeze(-1) + + +class PostHocEMACallback(Callback): + """Maintain and checkpoint multiple Karras EMA profiles on local shards. + + The upstream ``nitrous-ema`` implementation deep-copies a complete model + on rank zero. That is appropriate for MMAudio's DDP trainer but not for a + FastVideo FSDP2/HSDP transformer. This callback applies the same update and + synthesis equations independently to every local FSDP shard. Together the + rank-local snapshots represent the same full EMA model without gathering + it on every optimizer step. + """ + + def __init__( + self, + *, + sigma_rels: list[float] | tuple[float, ...] = (0.05, 0.1), + update_every: int = 1, + checkpoint_every: int = 5000, + checkpoint_folder: str | None = None, + start_iter: int = 0, + default_output_sigma: float = 0.05, + step_size_correction: bool = True, + ) -> None: + if len(sigma_rels) < 2: + raise ValueError("Post-hoc EMA requires at least two sigma profiles") + self.sigma_rels = tuple(float(value) for value in sigma_rels) + self.gammas = tuple(sigma_rel_to_gamma(value) for value in self.sigma_rels) + if len(set(self.gammas)) != len(self.gammas): + raise ValueError("Post-hoc EMA sigma profiles must be unique") + self.update_every = max(1, int(update_every)) + self.checkpoint_every = max(1, int(checkpoint_every)) + self.checkpoint_folder = checkpoint_folder + self.start_iter = int(start_iter) + self.default_output_sigma = float(default_output_sigma) + self.step_size_correction = bool(step_size_correction) + + self._ema_models: list[EMA_FSDP] = [] + self._calls = 0 + self._initted = False + self._rank = 0 + self._snapshot_root: Path | None = None + + def on_train_start( + self, + method: TrainingMethod, + iteration: int = 0, + ) -> None: + del iteration + transformer = method.student.transformer + self._ema_models = [EMA_FSDP(transformer, decay=0.0, mode="local_shard") for _ in self.gammas] + if torch.distributed.is_initialized(): + self._rank = torch.distributed.get_rank() + root = self.checkpoint_folder + if not root: + root = str(Path(self.training_config.checkpoint.output_dir) / "posthoc_ema") + self._snapshot_root = Path(root).expanduser().resolve() + (self._snapshot_root / f"rank_{self._rank:05d}").mkdir( + parents=True, + exist_ok=True, + ) + logger.info( + "PostHocEMA enabled: sigma_rels=%s, checkpoint_every=%d, folder=%s", + self.sigma_rels, + self.checkpoint_every, + self._snapshot_root, + ) + + def _decay(self, gamma: float) -> float: + step = self._calls + if not self.step_size_correction: + return (1.0 - 1.0 / (step + 1))**(1.0 + gamma) + first = (1.0 - 1.0 / (step + 1))**(1.0 + gamma) + second = (1.0 - 1.0 / (step + 1 + self.update_every))**(1.0 + gamma) + return (first * second)**(self.update_every / 2) + + def on_training_step_end( + self, + method: TrainingMethod, + loss_dict: dict[str, Any], + iteration: int = 0, + ) -> None: + del loss_dict + if iteration < self.start_iter or not self._ema_models: + return + + previous_calls = self._calls + self._calls += 1 + if previous_calls % self.update_every != 0: + return + + transformer = method.student.transformer + if not self._initted: + for ema in self._ema_models: + ema._init_shadow(transformer) + self._initted = True + for gamma, ema in zip(self.gammas, self._ema_models, strict=True): + ema.decay = self._decay(gamma) + ema.update(transformer) + + if self._calls % self.checkpoint_every == 0: + self._save_snapshots() + + def _save_snapshots(self) -> None: + if self._snapshot_root is None: + raise RuntimeError("PostHocEMA callback has not been initialized") + rank_dir = self._snapshot_root / f"rank_{self._rank:05d}" + for index, ema in enumerate(self._ema_models): + path = rank_dir / f"{index}.{self._calls}.pt" + temporary = path.with_suffix(".pt.tmp") + torch.save(ema.state_dict(), temporary) + os.replace(temporary, path) + logger.info("Saved PostHocEMA shard snapshots at step %d", self._calls) + + def _snapshot_records( + self, + *, + max_step: int | None, + ) -> list[tuple[Path, int, float]]: + if self._snapshot_root is None: + return [] + rank_dir = self._snapshot_root / f"rank_{self._rank:05d}" + records: list[tuple[Path, int, float]] = [] + for path in rank_dir.glob("*.pt"): + try: + profile_index, timestep = map(int, path.stem.split(".")) + gamma = self.gammas[profile_index] + except (ValueError, IndexError): + continue + if max_step is None or timestep <= max_step: + records.append((path, timestep, gamma)) + return sorted(records, key=lambda item: (item[1], item[0].name)) + + def synthesize_local_shard( + self, + *, + sigma_rel: float | None = None, + step: int | None = None, + ) -> dict[str, torch.Tensor] | None: + target_sigma = self.default_output_sigma if sigma_rel is None else float(sigma_rel) + records = self._snapshot_records(max_step=step) + if not records: + return None + target_step = max(record[1] for record in records) if step is None else int(step) + if target_step > max(record[1] for record in records): + raise ValueError("Cannot synthesize PostHocEMA beyond the newest snapshot") + + weights = solve_posthoc_weights( + torch.tensor([record[1] for record in records]), + torch.tensor([record[2] for record in records]), + target_step, + sigma_rel_to_gamma(target_sigma), + ) + synthesized: dict[str, torch.Tensor] = {} + for (path, _, _), weight in zip(records, weights.tolist(), strict=True): + state = torch.load(path, map_location="cpu", weights_only=True) + if not isinstance(state, dict): + raise ValueError(f"Invalid PostHocEMA snapshot: {path}") + for name, tensor in state.items(): + if not isinstance(tensor, torch.Tensor): + continue + if name not in synthesized: + synthesized[name] = tensor.float().mul(float(weight)) + else: + synthesized[name].add_(tensor.float(), alpha=float(weight)) + return synthesized + + @contextlib.contextmanager + def ema_context( + self, + transformer: torch.nn.Module, + ) -> Generator[torch.nn.Module, None, None]: + synthesized = self.synthesize_local_shard() + if synthesized is None: + yield transformer + return + temporary_ema = EMA_FSDP(transformer, decay=0.0, mode="local_shard") + temporary_ema.load_state_dict(synthesized) + with temporary_ema.apply_to_model(transformer): + yield transformer + + def on_train_end( + self, + method: TrainingMethod, + iteration: int = 0, + ) -> None: + del method, iteration + if not self._initted or self._snapshot_root is None: + return + + # Upstream MMAudio synthesizes the sigma=0.05 model after training. + # Save the latest profiles even when a short/early-stopped run did not + # end exactly on checkpoint_every, then synthesize each FSDP shard. + self._save_snapshots() + synthesized = self.synthesize_local_shard( + sigma_rel=self.default_output_sigma, + step=self._calls, + ) + if synthesized is None: + return + rank_dir = self._snapshot_root / f"rank_{self._rank:05d}" + sigma_name = str(self.default_output_sigma).replace(".", "p") + path = rank_dir / (f"synthesized_sigma_{sigma_name}_step_{self._calls:09d}.pt") + temporary = path.with_suffix(".pt.tmp") + torch.save(synthesized, temporary) + os.replace(temporary, path) + logger.info("Saved synthesized PostHocEMA shard to %s", path) + + def state_dict(self) -> dict[str, Any]: + return { + "calls": self._calls, + "initted": self._initted, + "profiles": { + str(index): ema.state_dict() + for index, ema in enumerate(self._ema_models) + }, + } + + def load_state_dict(self, state_dict: dict[str, Any]) -> None: + self._calls = int(state_dict.get("calls", 0)) + self._initted = bool(state_dict.get("initted", False)) + profiles = state_dict.get("profiles", {}) + if isinstance(profiles, dict): + for index, ema in enumerate(self._ema_models): + profile = profiles.get(str(index)) + if isinstance(profile, dict): + ema.load_state_dict(profile) + + +__all__ = [ + "PostHocEMACallback", + "sigma_rel_to_gamma", + "solve_posthoc_weights", +] diff --git a/fastvideo/train/methods/__init__.py b/fastvideo/train/methods/__init__.py index 0be93931e..aca975980 100644 --- a/fastvideo/train/methods/__init__.py +++ b/fastvideo/train/methods/__init__.py @@ -6,6 +6,7 @@ __all__ = [ "TrainingMethod", "DMD2Method", "FineTuneMethod", + "FlowMatchingFineTuneMethod", "KDMethod", "SelfForcingMethod", "DiffusionForcingSFTMethod", @@ -21,6 +22,10 @@ def __getattr__(name: str) -> object: if name == "FineTuneMethod": from fastvideo.train.methods.fine_tuning.finetune import FineTuneMethod return FineTuneMethod + if name == "FlowMatchingFineTuneMethod": + from fastvideo.train.methods.fine_tuning.flow_matching import ( + FlowMatchingFineTuneMethod, ) + return FlowMatchingFineTuneMethod if name == "KDMethod": from fastvideo.train.methods.knowledge_distillation.kd import KDMethod return KDMethod diff --git a/fastvideo/train/methods/base.py b/fastvideo/train/methods/base.py index 92faf4b02..6a9fe48b3 100644 --- a/fastvideo/train/methods/base.py +++ b/fastvideo/train/methods/base.py @@ -12,6 +12,7 @@ from fastvideo import envs from fastvideo.logger import init_logger from fastvideo.train.models.base import ModelBase from fastvideo.train.utils.checkpoint import _RoleModuleContainer +from fastvideo.train.utils.optimizer import seed_adamw_parameter_state from fastvideo.training.checkpointing_utils import ( ModelWrapper, OptimizerWrapper, @@ -189,11 +190,7 @@ class TrainingMethod(torch.nn.Module, ABC): continue if len(opt.state.get(p, {})) > 0: continue - opt.state[p] = { - "step": torch.tensor(0.0), - "exp_avg": torch.zeros_like(p), - "exp_avg_sq": torch.zeros_like(p), - } + seed_adamw_parameter_state(opt, p) # -- Shared hooks (override in subclasses as needed) -- diff --git a/fastvideo/train/methods/fine_tuning/__init__.py b/fastvideo/train/methods/fine_tuning/__init__.py index 96da179f7..e98c4b2c9 100644 --- a/fastvideo/train/methods/fine_tuning/__init__.py +++ b/fastvideo/train/methods/fine_tuning/__init__.py @@ -7,11 +7,13 @@ from typing import TYPE_CHECKING if TYPE_CHECKING: from fastvideo.train.methods.fine_tuning.dfsft import DiffusionForcingSFTMethod from fastvideo.train.methods.fine_tuning.finetune import FineTuneMethod + from fastvideo.train.methods.fine_tuning.flow_matching import FlowMatchingFineTuneMethod from fastvideo.train.methods.fine_tuning.tfsft import TeacherForcingSFTMethod __all__ = [ "DiffusionForcingSFTMethod", "FineTuneMethod", + "FlowMatchingFineTuneMethod", "TeacherForcingSFTMethod", ] @@ -27,6 +29,11 @@ def __getattr__(name: str) -> object: from fastvideo.train.methods.fine_tuning.finetune import FineTuneMethod return FineTuneMethod + if name == "FlowMatchingFineTuneMethod": + from fastvideo.train.methods.fine_tuning.flow_matching import ( + FlowMatchingFineTuneMethod, ) + + return FlowMatchingFineTuneMethod if name == "TeacherForcingSFTMethod": from fastvideo.train.methods.fine_tuning.tfsft import ( TeacherForcingSFTMethod, ) diff --git a/fastvideo/train/methods/fine_tuning/finetune.py b/fastvideo/train/methods/fine_tuning/finetune.py index 7cd4cd36c..c02c325a2 100644 --- a/fastvideo/train/methods/fine_tuning/finetune.py +++ b/fastvideo/train/methods/fine_tuning/finetune.py @@ -11,7 +11,9 @@ import torch.nn.functional as F from fastvideo.train.methods.base import TrainingMethod, LogScalar from fastvideo.train.models.base import ModelBase, NoisePrediction from fastvideo.train.utils.optimizer import ( - build_optimizer_and_scheduler, ) + build_optimizer_and_scheduler, + seed_adamw_parameter_state, +) def _compute_finetune_loss_map( @@ -254,3 +256,16 @@ class FineTuneMethod(TrainingMethod): betas=student_betas, scheduler_name=student_sched, ) + optimizer_parameters = { + parameter + for group in self._student_optimizer.param_groups + for parameter in group["params"] + } + for parameter in self.student.eager_optimizer_state_parameters(): + if parameter not in optimizer_parameters: + raise ValueError("An eager optimizer-state parameter is not owned by the " + "student optimizer") + seed_adamw_parameter_state( + self._student_optimizer, + parameter, + ) diff --git a/fastvideo/train/methods/fine_tuning/flow_matching.py b/fastvideo/train/methods/fine_tuning/flow_matching.py new file mode 100644 index 000000000..8b63bd999 --- /dev/null +++ b/fastvideo/train/methods/fine_tuning/flow_matching.py @@ -0,0 +1,144 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Shape-agnostic conditional flow-matching fine-tuning.""" + +from __future__ import annotations + +from collections.abc import Callable +from typing import Any + +import torch +import torch.nn.functional as F + +from fastvideo.logger import init_logger +from fastvideo.train.methods.base import LogScalar +from fastvideo.train.methods.fine_tuning.finetune import FineTuneMethod +from fastvideo.train.models.base import ModelBase +from fastvideo.train.utils.distributed_strategy import is_ddp_strategy + +logger = init_logger(__name__) + +TrainStepResult = tuple[ + dict[str, torch.Tensor], + dict[str, Any], + dict[str, LogScalar], +] + +TensorTrainFn = Callable[..., torch.Tensor] + + +class FlowMatchingFineTuneMethod(FineTuneMethod): + """Fine-tune a model against a velocity target supplied by its adapter. + + Unlike :class:`FineTuneMethod`, this method does not assume a five- + dimensional video latent layout. The model plugin owns sampling the flow + path and stores its target in ``TrainingBatch.training_target``. + """ + + def __init__( + self, + *, + cfg: Any, + role_models: dict[str, ModelBase], + ) -> None: + super().__init__(cfg=cfg, role_models=role_models) + flow_matching_forward = getattr(self.student, "flow_matching_forward", None) + if not callable(flow_matching_forward): + raise TypeError("FlowMatchingFineTuneMethod requires the student model to " + "implement flow_matching_forward()") + self._flow_matching_forward: TensorTrainFn = flow_matching_forward + self._train_fn: TensorTrainFn = self._eager_train_fn + + if bool(self.training_config.model.compile_train_fn): + if not is_ddp_strategy(self.training_config): + raise ValueError("training.model.compile_train_fn is currently supported " + "only with training.distributed.strategy=ddp") + compile_kwargs = dict(self.training_config.model.torch_compile_kwargs) + compile_kwargs.setdefault("fullgraph", True) + compile_training_forward = getattr(self.student, "compile_training_forward", None) + if not callable(compile_training_forward): + raise TypeError("compile_train_fn requires the student model to implement " + "compile_training_forward()") + logger.info( + "Enabling fullgraph transformer compilation with kwargs=%s", + compile_kwargs, + ) + compile_training_forward(compile_kwargs) + + def _eager_train_fn( + self, + noisy_latents: torch.Tensor, + timesteps: torch.Tensor, + target: torch.Tensor, + clip_features: torch.Tensor, + sync_features: torch.Tensor, + text_features: torch.Tensor, + ) -> torch.Tensor: + prediction = self._flow_matching_forward( + noisy_latents, + timesteps, + clip_features, + sync_features, + text_features, + ) + return F.mse_loss(prediction.float(), target.float()) + + def single_train_step( + self, + batch: dict[str, Any], + iteration: int, + ) -> TrainStepResult: + del iteration + if self.cuda_generator is None: + raise RuntimeError("on_train_start() must run before training") + + # RNG, existence masks, and CFG masks stay eager so the explicit + # per-rank CUDA Generator advances exactly as in official MMAudio. + training_batch = self.student.prepare_batch( + batch, + generator=self.cuda_generator, + latents_source="data", + ) + noisy_latents = training_batch.noisy_model_input + timesteps = training_batch.timesteps + target = training_batch.training_target + if not all(isinstance(value, torch.Tensor) for value in (noisy_latents, timesteps, target)): + raise RuntimeError("prepare_batch() returned incomplete flow inputs") + assert isinstance(noisy_latents, torch.Tensor) + assert isinstance(timesteps, torch.Tensor) + assert isinstance(target, torch.Tensor) + if target.shape != noisy_latents.shape: + raise RuntimeError("Flow-matching target and model input shapes must match, got " + f"{tuple(target.shape)} and {tuple(noisy_latents.shape)}") + + conditions = training_batch.conditional_dict + if not isinstance(conditions, dict): + raise RuntimeError("prepare_batch() must set conditional_dict") + condition_values = tuple(conditions.get(key) for key in ("clip_features", "sync_features", "text_features")) + if not all(isinstance(value, torch.Tensor) for value in condition_values): + raise RuntimeError("MMAudio flow conditions must all be tensors") + clip_features, sync_features, text_features = condition_values + assert isinstance(clip_features, torch.Tensor) + assert isinstance(sync_features, torch.Tensor) + assert isinstance(text_features, torch.Tensor) + + loss = self._train_fn( + noisy_latents, + timesteps, + target, + clip_features, + sync_features, + text_features, + ) + attn_metadata = (training_batch.attn_metadata_vsa if self._attn_kind == "vsa" else training_batch.attn_metadata) + loss_map = { + "total_loss": loss, + "flow_matching_loss": loss, + } + outputs: dict[str, Any] = { + "_fv_backward": (timesteps, attn_metadata), + } + metrics: dict[str, LogScalar] = {} + return loss_map, outputs, metrics + + +__all__ = ["FlowMatchingFineTuneMethod"] diff --git a/fastvideo/train/models/base.py b/fastvideo/train/models/base.py index 790e9585c..0cf9e41be 100644 --- a/fastvideo/train/models/base.py +++ b/fastvideo/train/models/base.py @@ -117,6 +117,15 @@ class ModelBase(ABC): def on_train_start(self) -> None: # noqa: B027 """Called once before the training loop begins.""" + def eager_optimizer_state_parameters(self, ) -> tuple[torch.nn.Parameter, ...]: + """Parameters whose optimizer state must exist before their first grad. + + Most models return an empty tuple. Models with intermittently used + trainable parameters can override this so strict distributed + checkpoints retain a stable optimizer-state schema across steps. + """ + return () + def decode_latents( self, latents_b_t_c_h_w: torch.Tensor, diff --git a/fastvideo/train/models/mmaudio/__init__.py b/fastvideo/train/models/mmaudio/__init__.py new file mode 100644 index 000000000..97b10d6b2 --- /dev/null +++ b/fastvideo/train/models/mmaudio/__init__.py @@ -0,0 +1,5 @@ +# SPDX-License-Identifier: Apache-2.0 + +from fastvideo.train.models.mmaudio.mmaudio import MMAudioModel + +__all__ = ["MMAudioModel"] diff --git a/fastvideo/train/models/mmaudio/mmaudio.py b/fastvideo/train/models/mmaudio/mmaudio.py new file mode 100644 index 000000000..6aa513ca9 --- /dev/null +++ b/fastvideo/train/models/mmaudio/mmaudio.py @@ -0,0 +1,609 @@ +# SPDX-License-Identifier: Apache-2.0 +"""MMAudio model plugin for the modular FastVideo trainer.""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any, Literal, TYPE_CHECKING + +import torch + +from fastvideo.configs.models.dits import ( + MMAUDIO_44K_TRAINING_VARIANTS, + MMAUDIO_TRAINING_VARIANTS, + MMAUDIO_VARIANT_ARCHITECTURES, + get_mmaudio_transformer_config, +) +from fastvideo.configs.pipelines.mmaudio import ( + MMAUDIO_PIPELINE_CONFIGS, + MMAudioV2AConfig, +) +from fastvideo.distributed import get_local_torch_device, get_world_group +from fastvideo.forward_context import set_forward_context +from fastvideo.logger import init_logger +from fastvideo.models.dits.mmaudio import MMAudioTransformer +from fastvideo.models.loader.fsdp_load import build_fsdp_model_from_scratch +from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import ( + FlowMatchEulerDiscreteScheduler, ) +from fastvideo.pipelines import TrainingBatch +from fastvideo.platforms import AttentionBackendEnum +from fastvideo.train.models.base import ModelBase +from fastvideo.train.utils.distributed_strategy import ( + build_replicated_model_from_scratch, + is_ddp_strategy, + unwrap_ddp_module, + wrap_module_ddp, +) +from fastvideo.train.utils.module_state import apply_trainable +from fastvideo.train.utils.moduleloader import load_module_from_path +from fastvideo.utils import PRECISION_TO_TYPE + +if TYPE_CHECKING: + from fastvideo.train.utils.lora import LoraConfig + from fastvideo.train.utils.training_config import TrainingConfig + +logger = init_logger(__name__) + + +def _load_empty_string_features(path: str | Path) -> torch.Tensor: + """Load only the fixed CLIP empty-string embedding from an official asset.""" + source = Path(path).expanduser().resolve() + if not source.exists(): + raise FileNotFoundError(f"MMAudio empty-string feature source does not exist: {source}") + + if source.is_dir(): + preferred = ( + source / "transformer" / "diffusion_pytorch_model.safetensors", + source / "diffusion_pytorch_model.safetensors", + ) + candidates = [candidate for candidate in preferred if candidate.is_file()] + if not candidates: + candidates = sorted(source.rglob("*.safetensors")) + for candidate in candidates: + from safetensors import safe_open + + with safe_open(candidate, framework="pt", device="cpu") as handle: + if "empty_string_feat" in handle: + return handle.get_tensor("empty_string_feat").float() + raise ValueError(f"No 'empty_string_feat' tensor was found in safetensors under {source}") + + if source.suffix == ".safetensors": + from safetensors import safe_open + + with safe_open(source, framework="pt", device="cpu") as handle: + if "empty_string_feat" not in handle: + raise ValueError(f"MMAudio safetensors {source} has no 'empty_string_feat' tensor") + return handle.get_tensor("empty_string_feat").float() + + state = torch.load(source, map_location="cpu", weights_only=True) + if isinstance(state, torch.Tensor): + tensor = state + elif isinstance(state, dict): + tensor = state.get("empty_string_feat") + if tensor is None and isinstance(state.get("state_dict"), dict): + tensor = state["state_dict"].get("empty_string_feat") + if tensor is None: + for key, value in state.items(): + if str(key).endswith("empty_string_feat") and isinstance(value, torch.Tensor): + tensor = value + break + else: + tensor = None + if not isinstance(tensor, torch.Tensor): + raise ValueError(f"MMAudio asset {source} does not contain an empty-string tensor") + if tensor.ndim == 3 and tensor.shape[0] == 1: + tensor = tensor.squeeze(0) + return tensor.float() + + +def _load_or_compute_latent_stats( + training_config: TrainingConfig, + *, + variant: str, + cache_path: str | None, + chunk_size: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Compute/cache official latent statistics once, then share across ranks.""" + world_group = get_world_group() + payload: dict[str, Any] | None = None + if world_group.is_first_rank: + stats_path = Path(cache_path).expanduser().resolve() if cache_path else None + if stats_path is not None and stats_path.is_file(): + loaded = torch.load(stats_path, map_location="cpu", weights_only=True) + if not isinstance(loaded, dict): + raise ValueError(f"Invalid MMAudio latent statistics file: {stats_path}") + payload = loaded + logger.info("Loaded MMAudio latent statistics from %s", stats_path) + else: + from fastvideo.dataset.mmaudio_feature_dataset import ( + compute_mmaudio_latent_stats, ) + + architecture = MMAUDIO_VARIANT_ARCHITECTURES[variant] + logger.info("Computing MMAudio latent statistics from the first feature cache") + latent_mean, latent_std = compute_mmaudio_latent_stats( + training_config.data.data_path, + latent_seq_len=int(architecture["latent_seq_len"]), + latent_dim=int(architecture["latent_dim"]), + chunk_size=chunk_size, + ) + payload = { + "variant": variant, + "latent_mean": latent_mean, + "latent_std": latent_std, + } + if stats_path is not None: + stats_path.parent.mkdir(parents=True, exist_ok=True) + torch.save(payload, stats_path) + logger.info("Saved MMAudio latent statistics to %s", stats_path) + + payload = world_group.broadcast_object(payload, src=0) + if not isinstance(payload, dict): + raise RuntimeError("Failed to broadcast MMAudio latent statistics") + cached_variant = payload.get("variant") + same_44k_contract = (cached_variant in MMAUDIO_44K_TRAINING_VARIANTS and variant in MMAUDIO_44K_TRAINING_VARIANTS) + if cached_variant is not None and cached_variant != variant and not same_44k_contract: + raise ValueError(f"MMAudio latent statistics are for {cached_variant!r}, not {variant!r}") + latent_mean = payload.get("latent_mean") + latent_std = payload.get("latent_std") + if not isinstance(latent_mean, torch.Tensor) or not isinstance(latent_std, torch.Tensor): + raise ValueError("MMAudio latent statistics must contain tensor mean/std values") + expected_shape = (1, 1, int(MMAUDIO_VARIANT_ARCHITECTURES[variant]["latent_dim"])) + if tuple(latent_mean.shape) != expected_shape or tuple(latent_std.shape) != expected_shape: + raise ValueError(f"MMAudio latent statistics for {variant!r} must have shape " + f"{expected_shape}, got {tuple(latent_mean.shape)} and " + f"{tuple(latent_std.shape)}") + return latent_mean, latent_std + + +class MMAudioModel(ModelBase): + """Official-compatible conditional flow-matching training adapter. + + Only the MMAudio transformer participates in the training graph. Audio VAE, + DFN5B, and Synchformer outputs are read from an offline feature cache. + """ + + _transformer_cls_name = "MMAudioTransformer" + + def __init__( + self, + *, + init_from: str | None = None, + training_config: TrainingConfig, + variant: str | None = None, + from_scratch: bool = False, + empty_string_features_path: str | None = None, + latent_statistics_path: str | None = None, + latent_statistics_chunk_size: int = 32, + trainable: bool = True, + allow_v2_training: bool = False, + lora: LoraConfig | dict[str, Any] | None = None, + attention_backend: AttentionBackendEnum | str | None = None, + transformer: torch.nn.Module | None = None, + ) -> None: + super().__init__( + trainable=trainable, + lora=lora, + attention_backend=attention_backend, + ) + if int(training_config.distributed.sp_size or 1) != 1: + raise ValueError("MMAudio training does not yet support sequence parallelism; set sp_size=1") + if int(training_config.distributed.tp_size or 1) != 1: + raise ValueError("MMAudio training does not yet support tensor parallelism; set tp_size=1") + use_ddp = is_ddp_strategy(training_config) + if use_ddp and (int(training_config.distributed.hsdp_replicate_dim) != 1 + or int(training_config.distributed.hsdp_shard_dim) not in {-1, 1}): + raise ValueError("MMAudio DDP strategy does not use HSDP dimensions; " + "set hsdp_replicate_dim=1 and hsdp_shard_dim=1") + + self.training_config = training_config + if variant is not None: + try: + pipeline_config_cls = MMAUDIO_PIPELINE_CONFIGS[variant] + except KeyError as exc: + supported = ", ".join(MMAUDIO_PIPELINE_CONFIGS) + raise ValueError(f"Unknown MMAudio variant {variant!r}; expected one of: {supported}") from exc + self.training_config.pipeline_config = pipeline_config_cls() + elif self.training_config.pipeline_config is None: + self.training_config.pipeline_config = MMAudioV2AConfig() + + if transformer is None: + if from_scratch: + if init_from: + raise ValueError("MMAudio from_scratch=true cannot also set init_from; use " + "empty_string_features_path for the fixed CLIP embedding") + if variant is None: + raise ValueError("MMAudio from-scratch training requires variant") + if variant not in MMAUDIO_TRAINING_VARIANTS: + supported = ", ".join(MMAUDIO_TRAINING_VARIANTS) + raise ValueError(f"The official MMAudio recipe does not train {variant!r}; " + f"supported from-scratch variants: {supported}") + if not empty_string_features_path: + raise ValueError("MMAudio from-scratch training requires " + "empty_string_features_path") + latent_mean, latent_std = _load_or_compute_latent_stats( + self.training_config, + variant=variant, + cache_path=latent_statistics_path, + chunk_size=latent_statistics_chunk_size, + ) + empty_string_feat = _load_empty_string_features(empty_string_features_path) + architecture = MMAUDIO_VARIANT_ARCHITECTURES[variant] + expected_text_shape = (77, 1024) + if tuple(empty_string_feat.shape) != expected_text_shape: + raise ValueError("MMAudio empty-string features must have shape " + f"{expected_text_shape}, got {tuple(empty_string_feat.shape)}") + model_config = get_mmaudio_transformer_config(variant) + hf_config = { + "_class_name": self._transformer_cls_name, + **architecture, + "clip_dim": 1024, + "clip_seq_len": 64, + "sync_dim": 768, + "sync_seq_len": 192, + "text_dim": 1024, + "text_seq_len": 77, + "mlp_ratio": 4.0, + } + default_dtype = PRECISION_TO_TYPE[self.training_config.dit_precision] + init_params = { + "config": model_config, + "hf_config": hf_config, + "latent_mean": latent_mean, + "latent_std": latent_std, + "empty_string_feat": empty_string_feat, + } + if use_ddp: + transformer = build_replicated_model_from_scratch( + MMAudioTransformer, + init_params, + device=get_local_torch_device(), + default_dtype=default_dtype, + seed=self.training_config.data.seed, + ) + else: + transformer = build_fsdp_model_from_scratch( + model_cls=MMAudioTransformer, + init_params=init_params, + device=get_local_torch_device(), + hsdp_replicate_dim=self.training_config.distributed.hsdp_replicate_dim, + hsdp_shard_dim=self.training_config.distributed.hsdp_shard_dim, + default_dtype=default_dtype, + param_dtype=torch.bfloat16, + reduce_dtype=torch.float32, + seed=self.training_config.data.seed, + pin_cpu_memory=self.training_config.distributed.pin_cpu_memory, + ) + self._init_from = f"scratch:{variant}" + else: + if not init_from: + raise ValueError("MMAudio pretrained training requires init_from, or set " + "from_scratch=true with a variant") + if use_ddp: + raise NotImplementedError("MMAudio DDP currently supports from_scratch=true only; " + "pretrained component loading remains on the FSDP path") + self._init_from = str(init_from) + transformer = load_module_from_path( + model_path=self._init_from, + module_type="transformer", + training_config=self.training_config, + override_transformer_cls_name=self._transformer_cls_name, + attention_backend=self.attention_backend, + ) + else: + self._init_from = str(init_from or f"provided:{variant or 'unknown'}") + self.transformer = transformer + if bool(getattr(self.transformer, "v2", False)) and not allow_v2_training: + raise ValueError("The official MMAudio training recipe does not support `_v2` " + "checkpoints. Use small_44k, medium_44k, or large_44k.") + + if not self._enable_lora_if_configured(self.transformer): + self.transformer = apply_trainable(self.transformer, trainable=self._trainable) + if self._trainable: + # These are checkpoint statistics/fixed text features in the + # official model, not optimization variables. ``apply_trainable`` + # intentionally enables a whole module, so restore their contract. + for name in ("latent_mean", "latent_std", "empty_string_feat"): + parameter = getattr(self.transformer, name, None) + if isinstance(parameter, torch.nn.Parameter): + parameter.requires_grad_(False) + # The two learned null-video tokens are trained by MMAudio. + for name in ("empty_clip_feat", "empty_sync_feat"): + parameter = getattr(self.transformer, name, None) + if isinstance(parameter, torch.nn.Parameter): + parameter.requires_grad_(True) + if use_ddp: + # Official MMAudio wraps the fully initialized FP32 model after + # selecting trainable parameters. It disables buffer broadcasts; + # the fixed latent statistics and positional buffers are identical + # because every rank uses the same initialization seed. + self.transformer = wrap_module_ddp( + self.transformer, + device=get_local_torch_device(), + broadcast_buffers=False, + ) + + self.noise_scheduler = FlowMatchEulerDiscreteScheduler( + shift=1.0, + invert_sigmas=True, + sigma_min=0.0, + use_reference_discrete_timesteps=True, + ) + self.dataloader: Any = None + self.start_step = 0 + + @property + def device(self) -> torch.device: + return next(self.transformer.parameters()).device + + def _unshard_transformer(self) -> None: + unshard = getattr(self.transformer, "unshard", None) + if callable(unshard): + unshard() + + def eager_optimizer_state_parameters(self, ) -> tuple[torch.nn.Parameter, ...]: + parameters = ( + getattr(self.transformer, "empty_clip_feat", None), + getattr(self.transformer, "empty_sync_feat", None), + ) + return tuple(parameter for parameter in parameters + if isinstance(parameter, torch.nn.Parameter) and parameter.requires_grad) + + def init_preprocessors(self, training_config: TrainingConfig) -> None: + if training_config.data.preprocessed_data_type != "mmaudio_features": + raise ValueError("MMAudioModel requires " + "training.data.preprocessed_data_type=mmaudio_features") + from fastvideo.dataset.mmaudio_feature_dataset import ( + build_mmaudio_feature_dataloader, ) + + feature_shapes = { + "latent_seq_len": int(self.transformer.latent_seq_len), + "latent_dim": int(self.transformer.latent_dim), + "clip_seq_len": int(self.transformer.clip_seq_len), + "clip_dim": int(self.transformer.config.arch_config.clip_dim), + "sync_seq_len": int(self.transformer.sync_seq_len), + "sync_dim": int(self.transformer.config.arch_config.sync_dim), + "text_seq_len": int(self.transformer.config.arch_config.text_seq_len), + "text_dim": int(self.transformer.config.arch_config.text_dim), + } + self.dataloader = build_mmaudio_feature_dataloader( + training_config.data.data_path, + batch_size=training_config.data.train_batch_size, + num_data_workers=training_config.data.dataloader_num_workers, + seed=training_config.data.seed, + pin_memory=training_config.distributed.pin_cpu_memory, + feature_shapes=feature_shapes, + ) + + @staticmethod + def _batch_tensor(raw_batch: dict[str, Any], primary: str, alias: str | None = None) -> torch.Tensor: + value = raw_batch.get(primary) + if value is None and alias is not None: + value = raw_batch.get(alias) + if not isinstance(value, torch.Tensor): + names = primary if alias is None else f"{primary!r} (or {alias!r})" + raise ValueError(f"MMAudio training batch is missing tensor {names}") + return value + + @staticmethod + def _existence_mask( + raw_batch: dict[str, Any], + primary: str, + alias: str, + *, + batch_size: int, + device: torch.device, + ) -> torch.Tensor: + value = raw_batch.get(primary, raw_batch.get(alias)) + if value is None: + return torch.ones(batch_size, device=device, dtype=torch.bool) + if not isinstance(value, torch.Tensor): + value = torch.as_tensor(value) + value = value.to(device=device, dtype=torch.bool).reshape(-1) + if value.numel() != batch_size: + raise ValueError(f"{primary} must have {batch_size} values, got {value.numel()}") + return value + + def prepare_batch( + self, + raw_batch: dict[str, Any], + *, + generator: torch.Generator, + latents_source: Literal["data", "zeros"] = "data", + ) -> TrainingBatch: + # MMAudio uses transformer-owned latent statistics and learned null + # tokens while constructing the batch, before ``transformer.forward`` + # would normally trigger FSDP2's pre-forward all-gather. Explicitly + # unshard here so these tensors can be broadcast/expanded safely. The + # following transformer forward performs the normal reshard. + self._unshard_transformer() + + device = self.device + latent_mean = self._batch_tensor(raw_batch, "audio_latent_mean", "a_mean").to(device=device) + latent_std = self._batch_tensor(raw_batch, "audio_latent_std", "a_std").to(device=device) + clip_features = self._batch_tensor(raw_batch, "clip_features").to(device=device).clone() + sync_features = self._batch_tensor(raw_batch, "sync_features").to(device=device).clone() + text_features = self._batch_tensor(raw_batch, "text_features").to(device=device).clone() + + if latent_mean.shape != latent_std.shape: + raise ValueError("MMAudio audio latent mean/std shapes must match, got " + f"{tuple(latent_mean.shape)} and {tuple(latent_std.shape)}") + batch_size = int(latent_mean.shape[0]) + expected_latent_shape = ( + batch_size, + int(self.transformer.latent_seq_len), + int(self.transformer.latent_dim), + ) + if tuple(latent_mean.shape) != expected_latent_shape: + raise ValueError(f"MMAudio audio latents must have shape {expected_latent_shape}, " + f"got {tuple(latent_mean.shape)}") + + video_exists = self._existence_mask( + raw_batch, + "video_exists", + "video_exist", + batch_size=batch_size, + device=device, + ) + text_exists = self._existence_mask( + raw_batch, + "text_exists", + "text_exist", + batch_size=batch_size, + device=device, + ) + empty_clip = self.transformer.get_empty_clip_sequence(1).to(dtype=clip_features.dtype) + empty_sync = self.transformer.get_empty_sync_sequence(1).to(dtype=sync_features.dtype) + empty_text = self.transformer.get_empty_string_sequence(1).to(dtype=text_features.dtype) + video_exists_expanded = video_exists[:, None, None] + text_exists_expanded = text_exists[:, None, None] + clip_features = torch.where(video_exists_expanded, clip_features, empty_clip) + sync_features = torch.where(video_exists_expanded, sync_features, empty_sync) + text_features = torch.where(text_exists_expanded, text_features, empty_text) + + if latents_source == "data": + posterior_noise = torch.empty_like(latent_mean).normal_(generator=generator) + clean_latents = latent_mean + latent_std * posterior_noise + clean_latents = self.transformer.normalize(clean_latents) + elif latents_source == "zeros": + clean_latents = torch.zeros_like(latent_mean) + else: + raise ValueError(f"Unknown latents_source: {latents_source!r}") + + logit_mean = float(self.training_config.model.logit_mean) + logit_std = float(self.training_config.model.logit_std) + timestep = torch.sigmoid(torch.randn(batch_size, device=device, generator=generator) * logit_std + logit_mean) + prior_noise = torch.empty_like(clean_latents).normal_(generator=generator) + timestep_expanded = timestep[:, None, None] + noisy_latents = (1.0 - timestep_expanded) * prior_noise + timestep_expanded * clean_latents + + null_probability = float(self.training_config.data.training_cfg_rate) + if not 0.0 <= null_probability <= 1.0: + raise ValueError("training.data.training_cfg_rate must be in [0, 1]") + null_video = torch.rand(batch_size, device=device, generator=generator) < null_probability + null_video_expanded = null_video[:, None, None] + clip_features = torch.where(null_video_expanded, empty_clip, clip_features) + sync_features = torch.where(null_video_expanded, empty_sync, sync_features) + null_text = torch.rand(batch_size, device=device, generator=generator) < null_probability + null_text_expanded = null_text[:, None, None] + text_features = torch.where(null_text_expanded, empty_text, text_features) + + # Keep learned null-video tokens in every step's autograd graph even + # when this batch samples no null conditions. AdamW then creates a + # stable optimizer-state schema, which is required for strict DCP + # checkpoint resume. Multiplication by zero leaves forward values + # unchanged; selected null conditions still receive their real grads. + clip_features = clip_features + empty_clip.sum() * 0.0 + sync_features = sync_features + empty_sync.sum() * 0.0 + + batch = TrainingBatch() + batch.latents = clean_latents + batch.noisy_model_input = noisy_latents + batch.noise = prior_noise + batch.timesteps = timestep + batch.sigmas = timestep_expanded + batch.training_target = clean_latents - prior_noise + batch.raw_latent_shape = tuple(clean_latents.shape) + batch.conditional_dict = { + "clip_features": clip_features, + "sync_features": sync_features, + "text_features": text_features, + } + return batch + + def add_noise( + self, + clean_latents: torch.Tensor, + noise: torch.Tensor, + timestep: torch.Tensor, + ) -> torch.Tensor: + while timestep.ndim < clean_latents.ndim: + timestep = timestep.unsqueeze(-1) + return (1.0 - timestep) * noise + timestep * clean_latents + + def compile_training_forward( + self, + compile_kwargs: dict[str, Any], + ) -> None: + """Compile transformer math while leaving DDP control flow eager.""" + module = unwrap_ddp_module(self.transformer) + compiled_forward = torch.compile(module.forward, **compile_kwargs) + module.forward = compiled_forward + logger.info( + "Compiled inner MMAudio transformer forward with kwargs=%s", + compile_kwargs, + ) + + def flow_matching_forward( + self, + noisy_latents: torch.Tensor, + timestep: torch.Tensor, + clip_features: torch.Tensor, + sync_features: torch.Tensor, + text_features: torch.Tensor, + ) -> torch.Tensor: + """Tensor-only adapter around the DDP-wrapped training forward.""" + device_type = noisy_latents.device.type + with torch.autocast( + device_type=device_type, + dtype=torch.bfloat16, + enabled=device_type == "cuda", + ): + return self.transformer( + hidden_states=noisy_latents, + encoder_hidden_states={ + "clip_features": clip_features, + "sync_features": sync_features, + "text_features": text_features, + }, + timestep=timestep, + ) + + def predict_noise( + self, + noisy_latents: torch.Tensor, + timestep: torch.Tensor, + batch: TrainingBatch, + *, + conditional: bool, + cfg_uncond: dict[str, Any] | None = None, + attn_kind: Literal["dense", "vsa"] = "dense", + ) -> torch.Tensor: + if attn_kind != "dense": + raise ValueError("MMAudio training currently supports dense Torch SDPA only") + if conditional: + conditions = batch.conditional_dict + if conditions is None: + raise RuntimeError("MMAudio TrainingBatch is missing conditional_dict") + elif cfg_uncond is not None: + conditions = cfg_uncond + else: + self._unshard_transformer() + batch_size = noisy_latents.shape[0] + conditions = { + "clip_features": self.transformer.get_empty_clip_sequence(batch_size), + "sync_features": self.transformer.get_empty_sync_sequence(batch_size), + "text_features": self.transformer.get_empty_string_sequence(batch_size), + } + + with set_forward_context(current_timestep=timestep, attn_metadata=None): + return self.flow_matching_forward( + noisy_latents, + timestep, + conditions["clip_features"], + conditions["sync_features"], + conditions["text_features"], + ) + + def backward( + self, + loss: torch.Tensor, + ctx: Any, + *, + grad_accum_rounds: int, + ) -> None: + timestep, attn_metadata = ctx + with set_forward_context(current_timestep=timestep, attn_metadata=attn_metadata): + (loss / max(1, int(grad_accum_rounds))).backward() + + +__all__ = ["MMAudioModel"] diff --git a/fastvideo/train/trainer.py b/fastvideo/train/trainer.py index adb5cd7e5..3ac85179d 100644 --- a/fastvideo/train/trainer.py +++ b/fastvideo/train/trainer.py @@ -161,6 +161,7 @@ class Trainer: progress = tqdm( range(start_step + 1, max_steps + 1), initial=start_step, + total=max_steps, desc="Steps", disable=self.local_rank > 0, ) diff --git a/fastvideo/train/utils/config.py b/fastvideo/train/utils/config.py index c34fcb0a0..ea21fa2a1 100644 --- a/fastvideo/train/utils/config.py +++ b/fastvideo/train/utils/config.py @@ -344,7 +344,8 @@ def _build_training_config( ``preprocessed_data_type`` selects the dataloader's tensor contract: ``t2v`` carries video and text, ``t2va`` carries synchronized video, - audio, and text, and ``text_only`` carries prompt conditioning. + audio, and text, ``text_only`` carries prompt conditioning, and + ``mmaudio_features`` carries cached MMAudio V2A features. """ d = dict(t.get("distributed", {}) or {}) da = dict(t.get("data", {}) or {}) @@ -379,9 +380,9 @@ def _build_training_config( data_path = str(raw_data_path) preprocessed_data_type = str(da.get("preprocessed_data_type", "t2v") or "t2v").strip().lower() - if preprocessed_data_type not in {"t2v", "t2va", "text_only"}: + if preprocessed_data_type not in {"t2v", "t2va", "text_only", "mmaudio_features"}: raise ValueError("training.data.preprocessed_data_type must be one of " - "{'t2v', 't2va', 'text_only'}, got " + "{'t2v', 't2va', 'text_only', 'mmaudio_features'}, got " f"{preprocessed_data_type!r}") peak_tflops = get_optional_float( @@ -393,8 +394,12 @@ def _build_training_config( raise ValueError("training.performance.peak_tflops_per_gpu must be " f"> 0, got {peak_tflops}") + from fastvideo.train.utils.distributed_strategy import ( + normalize_distributed_strategy, ) + return TrainingConfig( distributed=DistributedConfig( + strategy=normalize_distributed_strategy(str(d.get("strategy", "fsdp") or "fsdp")), num_gpus=num_gpus, tp_size=int(d.get("tp_size", 1) or 1), sp_size=int(d.get("sp_size", num_gpus) or num_gpus), @@ -417,12 +422,16 @@ def _build_training_config( optimizer=OptimizerConfig( learning_rate=float(o.get("learning_rate", 0.0) or 0.0), betas=betas, + eps=float(o.get("eps", 1e-8) or 1e-8), weight_decay=float(o.get("weight_decay", 0.0) or 0.0), lr_scheduler=str(o.get("lr_scheduler", "constant") or "constant"), lr_warmup_steps=int(o.get("lr_warmup_steps", 0) or 0), lr_num_cycles=int(o.get("lr_num_cycles", 0) or 0), lr_power=float(o.get("lr_power", 0.0) or 0.0), min_lr_ratio=float(o.get("min_lr_ratio", 0.5) or 0.5), + lr_milestones=tuple(int(step) for step in (o.get("lr_milestones", []) or [])), + lr_gamma=float(o.get("lr_gamma", 0.1) or 0.1), + fused=bool(o.get("fused", False)), ), loop=TrainingLoopConfig( max_train_steps=int(lo.get("max_train_steps", 0) or 0), @@ -454,6 +463,8 @@ def _build_training_config( precondition_outputs=bool(m.get("precondition_outputs", False)), moba_config=dict(m.get("moba_config", {}) or {}), enable_gradient_checkpointing_type=(m.get("enable_gradient_checkpointing_type")), + compile_train_fn=bool(m.get("compile_train_fn", False)), + torch_compile_kwargs=dict(m.get("torch_compile_kwargs", {}) or {}), ), pipeline_config=pipeline_config, model_path=model_path, diff --git a/fastvideo/train/utils/distributed_strategy.py b/fastvideo/train/utils/distributed_strategy.py new file mode 100644 index 000000000..c876c0cef --- /dev/null +++ b/fastvideo/train/utils/distributed_strategy.py @@ -0,0 +1,140 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Distributed model strategies for the modular trainer.""" + +from __future__ import annotations + +from contextlib import contextmanager +from typing import Any, TypeVar +from collections.abc import Generator + +import torch +import torch.distributed as dist +from torch.nn.parallel import DistributedDataParallel + +from fastvideo.logger import init_logger + +logger = init_logger(__name__) + +_ModelT = TypeVar("_ModelT", bound=torch.nn.Module) +SUPPORTED_DISTRIBUTED_STRATEGIES = frozenset({"fsdp", "ddp"}) + + +def normalize_distributed_strategy(value: str | None) -> str: + """Normalize and validate a modular-training strategy name.""" + strategy = str(value or "fsdp").strip().lower() + if strategy not in SUPPORTED_DISTRIBUTED_STRATEGIES: + choices = ", ".join(sorted(SUPPORTED_DISTRIBUTED_STRATEGIES)) + raise ValueError(f"training.distributed.strategy must be one of {{{choices}}}, " + f"got {strategy!r}") + return strategy + + +def is_ddp_strategy(training_config: Any) -> bool: + distributed = getattr(training_config, "distributed", None) + return normalize_distributed_strategy(getattr(distributed, "strategy", "fsdp")) == "ddp" + + +@contextmanager +def _fork_model_rng( + device: torch.device, + seed: int, +) -> Generator[None, None, None]: + fork_devices: list[int] = [] + if device.type == "cuda": + fork_devices = [device.index if device.index is not None else torch.cuda.current_device()] + with torch.random.fork_rng(devices=fork_devices): + torch.manual_seed(int(seed)) + yield + + +def build_replicated_model_from_scratch( + model_cls: type[_ModelT], + init_params: dict[str, Any], + *, + device: torch.device, + default_dtype: torch.dtype, + seed: int, +) -> _ModelT: + """Initialize one deterministic, full model replica on the local device. + + Model plugins call :func:`wrap_module_ddp` only after applying their + trainable/frozen parameter policy. This is required because DDP builds its + gradient reducer from the parameters that require gradients at wrap time. + """ + old_dtype = torch.get_default_dtype() + try: + torch.set_default_dtype(default_dtype) + # Official MMAudio constructs on CPU and then calls ``.cuda()`` before + # DDP. Preserve that RNG/initialization order for scratch parity. + with _fork_model_rng(device, seed), torch.device("cpu"): + model = model_cls(**init_params) + finally: + torch.set_default_dtype(old_dtype) + return model.to(device=device) + + +class DelegatingDistributedDataParallel(DistributedDataParallel): + """DDP wrapper that preserves a model plugin's public attributes. + + FastVideo model plugins access architecture attributes and helper methods + such as ``latent_seq_len`` and ``normalize`` through ``transformer``. + Native PyTorch DDP exposes these on ``module`` only; this adapter delegates + unknown attributes while retaining the standard DDP forward/reducer. + """ + + def __getattr__(self, name: str) -> Any: + try: + return super().__getattr__(name) + except AttributeError: + module = super().__getattr__("module") + return getattr(module, name) + + +def wrap_module_ddp( + module: _ModelT, + *, + device: torch.device, + broadcast_buffers: bool = False, +) -> DelegatingDistributedDataParallel: + """Wrap a full model replica like official MMAudio's DDP runner.""" + if not dist.is_available() or not dist.is_initialized(): + raise RuntimeError("DDP strategy requires an initialized process group") + if not any(parameter.requires_grad for parameter in module.parameters()): + raise ValueError("DDP strategy requires at least one trainable parameter") + + kwargs: dict[str, Any] = { + "broadcast_buffers": bool(broadcast_buffers), + } + if device.type == "cuda": + if device.index is None: + raise ValueError("CUDA DDP requires a concrete local device index") + kwargs.update({ + "device_ids": [device.index], + "output_device": device.index, + }) + wrapped = DelegatingDistributedDataParallel(module, **kwargs) + logger.info( + "Wrapped %s with native DDP on %s (broadcast_buffers=%s)", + type(module).__name__, + device, + broadcast_buffers, + ) + return wrapped + + +def unwrap_ddp_module(module: _ModelT) -> torch.nn.Module: + """Return the underlying module for DDP, otherwise return the input.""" + if isinstance(module, DistributedDataParallel): + return module.module + return module + + +__all__ = [ + "DelegatingDistributedDataParallel", + "SUPPORTED_DISTRIBUTED_STRATEGIES", + "build_replicated_model_from_scratch", + "is_ddp_strategy", + "normalize_distributed_strategy", + "unwrap_ddp_module", + "wrap_module_ddp", +] diff --git a/fastvideo/train/utils/optimizer.py b/fastvideo/train/utils/optimizer.py index 43a79d98d..064f7e0ea 100644 --- a/fastvideo/train/utils/optimizer.py +++ b/fastvideo/train/utils/optimizer.py @@ -36,28 +36,80 @@ def build_optimizer_and_scheduler( raise ValueError("No trainable parameters passed to " "build_optimizer_and_scheduler") + optimizer_kwargs = {} + if bool(getattr(optimizer_config, "fused", False)): + optimizer_kwargs["fused"] = True optimizer = torch.optim.AdamW( params, lr=float(learning_rate), betas=betas, weight_decay=float(optimizer_config.weight_decay), - eps=1e-8, + eps=float(optimizer_config.eps), + **optimizer_kwargs, ) - scheduler = get_scheduler( - str(scheduler_name), - optimizer=optimizer, - num_warmup_steps=int(optimizer_config.lr_warmup_steps), - num_training_steps=int(loop_config.max_train_steps), - num_cycles=int(optimizer_config.lr_num_cycles), - power=float(optimizer_config.lr_power), - min_lr_ratio=float(optimizer_config.min_lr_ratio), - last_epoch=-1, - ) + if str(scheduler_name) == "multistep_with_warmup": + milestones = list(optimizer_config.lr_milestones) + if not milestones: + raise ValueError("multistep_with_warmup requires training.optimizer.lr_milestones") + warmup_steps = int(optimizer_config.lr_warmup_steps) + if warmup_steps <= 0: + raise ValueError("multistep_with_warmup requires lr_warmup_steps > 0") + + def warmup(current_step: int) -> float: + return (current_step + 1) / (warmup_steps + 1) + + warmup_scheduler = torch.optim.lr_scheduler.LambdaLR( + optimizer, + lr_lambda=warmup, + ) + step_scheduler = torch.optim.lr_scheduler.MultiStepLR( + optimizer, + milestones=milestones, + gamma=float(optimizer_config.lr_gamma), + ) + scheduler = torch.optim.lr_scheduler.SequentialLR( + optimizer, + schedulers=[warmup_scheduler, step_scheduler], + milestones=[warmup_steps], + ) + else: + scheduler = get_scheduler( + str(scheduler_name), + optimizer=optimizer, + num_warmup_steps=int(optimizer_config.lr_warmup_steps), + num_training_steps=int(loop_config.max_train_steps), + num_cycles=int(optimizer_config.lr_num_cycles), + power=float(optimizer_config.lr_power), + min_lr_ratio=float(optimizer_config.min_lr_ratio), + last_epoch=-1, + ) return optimizer, scheduler +def seed_adamw_parameter_state( + optimizer: torch.optim.Optimizer, + parameter: torch.nn.Parameter, +) -> None: + """Create AdamW state using PyTorch's fused/capturable device contract.""" + if optimizer.state.get(parameter): + return + owner_group = next( + (group for group in optimizer.param_groups if any(candidate is parameter for candidate in group["params"])), + None, + ) + if owner_group is None: + raise ValueError("Cannot seed optimizer state for an unowned parameter") + step_on_parameter = bool(owner_group.get("capturable", False) or owner_group.get("fused", False)) + step_device = parameter.device if step_on_parameter else torch.device("cpu") + optimizer.state[parameter] = { + "step": torch.zeros((), dtype=torch.float32, device=step_device), + "exp_avg": torch.zeros_like(parameter), + "exp_avg_sq": torch.zeros_like(parameter), + } + + def clip_grad_norm_if_needed( module: torch.nn.Module, max_grad_norm: float, diff --git a/fastvideo/train/utils/training_config.py b/fastvideo/train/utils/training_config.py index 85e838ec8..00645ba63 100644 --- a/fastvideo/train/utils/training_config.py +++ b/fastvideo/train/utils/training_config.py @@ -18,6 +18,8 @@ class DistributedConfig: hsdp_replicate_dim: int = 1 hsdp_shard_dim: int = -1 pin_cpu_memory: bool = False + # Keep this last so existing positional construction remains compatible. + strategy: str = "fsdp" @dataclass(slots=True) @@ -38,12 +40,18 @@ class DataConfig: class OptimizerConfig: learning_rate: float = 0.0 betas: tuple[float, float] = (0.9, 0.999) + eps: float = 1e-8 weight_decay: float = 0.0 lr_scheduler: str = "constant" lr_warmup_steps: int = 0 lr_num_cycles: int = 0 lr_power: float = 0.0 min_lr_ratio: float = 0.5 + lr_milestones: tuple[int, ...] = () + lr_gamma: float = 0.1 + # ``False`` preserves PyTorch's existing optimizer selection. When true, + # FastVideo explicitly requests CUDA fused AdamW. + fused: bool = False @dataclass(slots=True) @@ -87,6 +95,10 @@ class ModelTrainingConfig: precondition_outputs: bool = False moba_config: dict = field(default_factory=dict) enable_gradient_checkpointing_type: str | None = None + # Compile the model adapter's tensor-only training forward. This is + # separate from loader-level model compilation and defaults off. + compile_train_fn: bool = False + torch_compile_kwargs: dict = field(default_factory=dict) @dataclass(slots=True) diff --git a/fastvideo/training/checkpointing_utils.py b/fastvideo/training/checkpointing_utils.py index 9c0da9e23..c3dc77b2e 100644 --- a/fastvideo/training/checkpointing_utils.py +++ b/fastvideo/training/checkpointing_utils.py @@ -1,5 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 import random +from collections import Counter +from collections.abc import Mapping from typing import Any import numpy as np @@ -21,6 +23,8 @@ class ModelWrapper(torch.distributed.checkpoint.stateful.Stateful): k.replace("._checkpoint_wrapped_module.", ".") for k, v in self.model.named_parameters() if v.requires_grad } + if isinstance(self.model, torch.nn.parallel.DistributedDataParallel): + param_requires_grad = {key.removeprefix("module.") for key in param_requires_grad} filtered_state_dict = {k: v for k, v in state_dict.items() if k in param_requires_grad} @@ -62,10 +66,35 @@ class SchedulerWrapper(torch.distributed.checkpoint.stateful.Stateful): self.scheduler = scheduler def state_dict(self) -> dict[str, Any]: - return {"scheduler": self.scheduler.state_dict()} + return {"scheduler": _scheduler_state_for_dcp(self.scheduler.state_dict())} def load_state_dict(self, state_dict: dict[str, Any]) -> None: - self.scheduler.load_state_dict(state_dict["scheduler"]) + self.scheduler.load_state_dict(_scheduler_state_from_dcp(state_dict["scheduler"])) + + +def _scheduler_state_for_dcp(value: Any, *, name: str = "") -> Any: + """Use stable string keys for scheduler mappings flattened by DCP.""" + if isinstance(value, Mapping): + if name == "milestones": + return {str(key): item for key, item in value.items()} + return {key: _scheduler_state_for_dcp(item, name=str(key)) for key, item in value.items()} + if isinstance(value, list): + return [_scheduler_state_for_dcp(item) for item in value] + return value + + +def _scheduler_state_from_dcp(value: Any, *, name: str = "") -> Any: + """Restore MultiStepLR milestones after DCP stringifies mapping keys.""" + if isinstance(value, Mapping): + if name == "milestones": + milestones: dict[int, Any] = {} + for key, item in value.items(): + milestones[int(key)] = item + return Counter(milestones) + return {key: _scheduler_state_from_dcp(item, name=str(key)) for key, item in value.items()} + if isinstance(value, list): + return [_scheduler_state_from_dcp(item) for item in value] + return value class RandomStateWrapper(torch.distributed.checkpoint.stateful.Stateful): diff --git a/fastvideo/utils.py b/fastvideo/utils.py index 5c6ba6128..3b7d29c92 100644 --- a/fastvideo/utils.py +++ b/fastvideo/utils.py @@ -657,8 +657,11 @@ def verify_model_config_and_directory( else: # Full snapshots keep the historical invariant that transformer/ is # present and every active manifest component exists locally. + # Explicitly marked offline preprocessor trees contain only frozen + # encoders and avoid duplicating a DiT that will never be loaded. transformer_dir = os.path.join(model_path, "transformer") - if not os.path.exists(transformer_dir): + preprocessor_only = config.get("_fastvideo_preprocessor_only") is True + if not preprocessor_only and not os.path.exists(transformer_dir): raise ValueError(f"Model directory {model_path} does not contain a transformer/ directory.") # Diffusers convention: component entries start with [library, class]. diff --git a/fastvideo/workflow/preprocess/preprocess_workflow.py b/fastvideo/workflow/preprocess/preprocess_workflow.py index b21d714b3..6d84f2837 100644 --- a/fastvideo/workflow/preprocess/preprocess_workflow.py +++ b/fastvideo/workflow/preprocess/preprocess_workflow.py @@ -117,6 +117,9 @@ class PreprocessWorkflow(WorkflowBase): elif fastvideo_args.workload_type == WorkloadType.I2V: from fastvideo.workflow.preprocess.preprocess_workflow_i2v import (PreprocessWorkflowI2V) return cast(PreprocessWorkflow, PreprocessWorkflowI2V) + elif fastvideo_args.workload_type == WorkloadType.V2A: + from fastvideo.workflow.preprocess.preprocess_workflow_v2a import (PreprocessWorkflowV2A) + return cast(PreprocessWorkflow, PreprocessWorkflowV2A) else: raise ValueError( f"Workload type: {fastvideo_args.workload_type} is not supported in preprocessing workflow.") diff --git a/fastvideo/workflow/preprocess/preprocess_workflow_v2a.py b/fastvideo/workflow/preprocess/preprocess_workflow_v2a.py new file mode 100644 index 000000000..0b2944dc0 --- /dev/null +++ b/fastvideo/workflow/preprocess/preprocess_workflow_v2a.py @@ -0,0 +1,246 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Shared raw-media workflow for video-to-audio feature preprocessing.""" + +from __future__ import annotations + +import json +import math +from collections import deque +from concurrent.futures import Future, ThreadPoolExecutor +from pathlib import Path +from typing import TYPE_CHECKING, Any +from collections.abc import Iterator + +import torch +from torch.utils.data import DataLoader, Subset +from tqdm import tqdm + +from fastvideo.configs.configs import DatasetType, VideoLoaderType +from fastvideo.dataset.v2a_feature_cache import V2AFeatureShardWriter +from fastvideo.dataset.vggsound import VGGSoundDataset +from fastvideo.distributed.parallel_state import get_world_rank, get_world_size +from fastvideo.logger import init_logger +from fastvideo.pipelines.pipeline_batch_info import PreprocessBatch +from fastvideo.pipelines.preprocess.mmaudio.torio_media_reader import (MMAudioTorioRowPreprocessor, + PREPROCESSED_MEDIA_KEY, PREPROCESS_ERROR_KEY) +from fastvideo.workflow.preprocess.preprocess_workflow import PreprocessWorkflow + +if TYPE_CHECKING: + from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase + +logger = init_logger(__name__) + + +def _identity_collate(rows: list[dict[str, Any]]) -> list[dict[str, Any]]: + return rows + + +def build_v2a_dataset( + dataset_type: DatasetType, + dataset_path: str, + *, + dataset_metadata_path: str = "", + dataset_split: str = "train", +): + """Build a raw-media dataset for a V2A preprocessing workflow.""" + if isinstance(dataset_type, str): + dataset_type = DatasetType.from_string(dataset_type) + if dataset_type is DatasetType.VGGSOUND: + metadata_path = dataset_metadata_path or None + return VGGSoundDataset( + dataset_path, + split=dataset_split, + metadata_path=metadata_path, + ) + raise ValueError(f"V2A preprocessing currently supports dataset_type=vggsound; got {dataset_type.value!r}") + + +class V2AForwardBatchBuilder: + """Translate standardized raw-media rows into a FastVideo batch.""" + + def __init__(self, seed: int) -> None: + self.seed = seed + + def __call__(self, rows: list[dict[str, Any]]) -> PreprocessBatch: + extra: dict[str, Any] = {} + preprocessed_media = [row.get(PREPROCESSED_MEDIA_KEY) for row in rows] + preprocess_errors = [row.get(PREPROCESS_ERROR_KEY) for row in rows] + if any(media is not None for media in preprocessed_media): + extra[PREPROCESSED_MEDIA_KEY] = preprocessed_media + if any(error is not None for error in preprocess_errors): + extra[PREPROCESS_ERROR_KEY] = preprocess_errors + return PreprocessBatch( + data_type="video", + video_loader=[str(row["video_path"]) for row in rows], + video_file_name=[str(row["id"]) for row in rows], + prompt=[str(row["caption"]) for row in rows], + batch_size=len(rows), + generator=torch.Generator("cpu").manual_seed(self.seed), + extra=extra, + ) + + +class PreprocessWorkflowV2A(PreprocessWorkflow): + """Run model-specific V2A encoders with shared data/cache orchestration.""" + + training_dataloader: DataLoader + preprocess_pipeline: ComposedPipelineBase + forward_batch_builder: V2AForwardBatchBuilder + feature_writer: V2AFeatureShardWriter + + def register_components(self) -> None: + self.torio_row_preprocessor: MMAudioTorioRowPreprocessor | None = None + self.torio_decode_workers = 0 + self.initial_skipped = 0 + assert self.fastvideo_args.preprocess_config is not None + config = self.fastvideo_args.preprocess_config + dataset = build_v2a_dataset( + config.dataset_type, + config.dataset_path, + dataset_metadata_path=config.dataset_metadata_path, + dataset_split=config.dataset_split, + ) + + rank = get_world_rank() + world_size = get_world_size() + feature_writer = V2AFeatureShardWriter( + config.dataset_output_dir, + rank=rank, + samples_per_shard=config.samples_per_file, + ) + rank_indices = list(range(rank, len(dataset), world_size)) + pending_indices = [index for index in rank_indices if not feature_writer.contains(str(dataset[index]["id"]))] + self.initial_skipped = len(rank_indices) - len(pending_indices) + rank_indices = pending_indices + rank_dataset = Subset(dataset, rank_indices) + + dataloader_workers = config.dataloader_num_workers + if config.video_loader_type == VideoLoaderType.TORIO: + pc = self.fastvideo_args.pipeline_config + duration_s = float(pc.duration_s) + sample_rate = int(pc.sampling_rate) + alignment = int(pc.spectrogram_frame_rate * pc.latent_downsample_rate) + target_samples = math.ceil(duration_s * sample_rate / alignment) * alignment + self.torio_row_preprocessor = MMAudioTorioRowPreprocessor( + duration_s=duration_s, + target_sample_rate=sample_rate, + target_samples=target_samples, + normalize_audio=config.dataset_split.lower() == "train", + clip_fps=pc.clip_frame_rate, + sync_fps=pc.sync_frame_rate, + clip_size=pc.clip_image_size, + sync_size=pc.sync_image_size, + ) + self.torio_decode_workers = max(0, dataloader_workers) + # Large video tensors cannot pass safely through a small /dev/shm. + # Threads also avoid tensor serialization while overlapping the + # next decode batch with GPU feature extraction. + dataloader_workers = 0 + logger.info( + "Torio preprocessing uses %d background decode threads on rank %d", + self.torio_decode_workers, + rank, + ) + + dataloader = DataLoader( + rank_dataset, + batch_size=config.preprocess_video_batch_size, + num_workers=dataloader_workers, + collate_fn=_identity_collate, + shuffle=False, + persistent_workers=dataloader_workers > 0, + ) + self.add_component("training_dataloader", dataloader) + self.add_component("forward_batch_builder", V2AForwardBatchBuilder(config.seed)) + self.add_component("feature_writer", feature_writer) + + def _iter_preprocessed_rows(self) -> Iterator[list[dict[str, Any]]]: + preprocessor = self.torio_row_preprocessor + if preprocessor is None: + yield from self.training_dataloader + return + + if self.torio_decode_workers <= 0: + for rows in self.training_dataloader: + yield [preprocessor(row) for row in rows] + return + + raw_batches = iter(self.training_dataloader) + queued_batches: deque[list[Future[dict[str, Any]]]] = deque() + with ThreadPoolExecutor( + max_workers=self.torio_decode_workers, + thread_name_prefix="mmaudio-torio", + ) as executor: + + def submit_next() -> bool: + try: + rows = next(raw_batches) + except StopIteration: + return False + queued_batches.append([executor.submit(preprocessor, row) for row in rows]) + return True + + submit_next() + submit_next() + while queued_batches: + futures = queued_batches.popleft() + rows = [future.result() for future in futures] + submit_next() + yield rows + + def prepare_system_environment(self) -> None: + assert self.fastvideo_args.preprocess_config is not None + output_root = Path(self.fastvideo_args.preprocess_config.dataset_output_dir).expanduser().resolve() + output_root.mkdir(parents=True, exist_ok=True) + self.output_root = output_root + self.failure_log = output_root / f"failures_rank_{get_world_rank():05d}.jsonl" + logger.info("V2A feature cache output: %s", output_root) + + def _write_failures(self, failures: list[dict[str, str]]) -> None: + if not failures: + return + with self.failure_log.open("a", encoding="utf-8") as handle: + for failure in failures: + handle.write(json.dumps(failure, ensure_ascii=False) + "\n") + + def run(self) -> None: + processed = 0 + skipped = self.initial_skipped + failed = 0 + try: + for rows in tqdm( + self._iter_preprocessed_rows(), + total=len(self.training_dataloader), + desc="Preprocessing V2A training data", + unit="batch", + ): + pending_rows = [row for row in rows if not self.feature_writer.contains(str(row["id"]))] + skipped += len(rows) - len(pending_rows) + if not pending_rows: + continue + batch = self.forward_batch_builder(pending_rows) + batch = self.preprocess_pipeline.forward(batch, self.fastvideo_args) + failures = batch.extra.get("failed_samples", []) + failed += len(failures) + self._write_failures(failures) + features = batch.extra.get("precomputed_features", {}) + metadata = batch.extra.get("precomputed_metadata", []) + self.feature_writer.append(features, metadata) + processed += len(metadata) + finally: + self.feature_writer.close() + logger.info( + "Finished V2A preprocessing on rank %s: processed=%s skipped=%s failed=%s output=%s", + get_world_rank(), + processed, + skipped, + failed, + self.output_root, + ) + + +__all__ = [ + "PreprocessWorkflowV2A", + "V2AForwardBatchBuilder", + "build_v2a_dataset", +] diff --git a/pyproject.toml b/pyproject.toml index e687e649e..9bf3439b3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -147,6 +147,12 @@ swanlab = [ "swanlab>=0.6.7", ] +mmaudio-train = [ + "librosa>=0.10.0", + "nitrous-ema==0.0.1", + "tensordict>=0.6.1", +] + lint = [ "pre-commit==4.0.1", "types-aiofiles", diff --git a/scripts/checkpoint_conversion/convert_mmaudio_to_diffusers.py b/scripts/checkpoint_conversion/convert_mmaudio_to_diffusers.py index 2ef85052c..932644604 100644 --- a/scripts/checkpoint_conversion/convert_mmaudio_to_diffusers.py +++ b/scripts/checkpoint_conversion/convert_mmaudio_to_diffusers.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -"""Convert MMAudio ``large_44k_v2`` assets into a FastVideo component tree. +"""Convert every official MMAudio variant 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 @@ -16,6 +16,10 @@ Example:: --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 + +Use ``--variant small_16k`` with ``v1-16.pth`` and ``best_netG.pt`` for +the 16 kHz model. The four 44.1 kHz variants share ``v1-44.pth`` and the +NVIDIA BigVGAN-v2 directory. """ from __future__ import annotations @@ -30,24 +34,80 @@ import torch from safetensors.torch import save_file -TRANSFORMER_CONFIG = { +BASE_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, } +TRANSFORMER_VARIANTS = { + "small_16k": { + "latent_dim": 20, + "hidden_dim": 448, + "depth": 12, + "fused_depth": 8, + "num_heads": 7, + "latent_seq_len": 250, + "v2": False, + }, + "small_44k": { + "hidden_dim": 448, + "depth": 12, + "fused_depth": 8, + "num_heads": 7, + "v2": False, + }, + "medium_44k": { + "hidden_dim": 896, + "depth": 12, + "fused_depth": 8, + "num_heads": 14, + "v2": False, + }, + "large_44k": { + "hidden_dim": 896, + "depth": 21, + "fused_depth": 14, + "num_heads": 14, + "v2": False, + }, + "large_44k_v2": { + "hidden_dim": 896, + "depth": 21, + "fused_depth": 14, + "num_heads": 14, + "v2": True, + }, +} + +BIGVGAN_16K_CONFIG = { + "_class_name": "BigVGANV2", + "resblock": "1", + "num_mels": 80, + "upsample_rates": [4, 4, 2, 2, 2, 2], + "upsample_kernel_sizes": [8, 8, 4, 4, 4, 4], + "upsample_initial_channel": 1536, + "resblock_kernel_sizes": [3, 7, 11], + "resblock_dilation_sizes": [ + [1, 3, 5], + [1, 3, 5], + [1, 3, 5], + ], + "activation": "snakebeta", + "snake_logscale": True, + "use_bias_at_final": True, + "use_tanh_at_final": True, + "weight_norm_removed": False, +} + + TEXT_ENCODER_CONFIG = { "architectures": ["MMAudioDFNCLIPTextEncoder"], "vocab_size": 49408, @@ -263,28 +323,112 @@ def _load_dfn5b_state(directory: Path) -> dict[str, torch.Tensor]: def convert(args: argparse.Namespace) -> None: output = args.output.resolve() output.mkdir(parents=True, exist_ok=True) + mode = "16k" if args.variant.endswith("16k") else "44k" + model_index = dict(MODEL_INDEX) + model_index["_fastvideo_mmaudio_variant"] = args.variant + model_index["_fastvideo_audio_sample_rate"] = ( + 16000 if mode == "16k" else 44100 + ) - 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) + transformer_config_only = bool(getattr(args, "transformer_config_only", False)) + if args.transformer_only and args.preprocessor_only: + raise ValueError( + "--transformer-only and --preprocessor-only are mutually exclusive") + if transformer_config_only and (args.transformer_only or args.preprocessor_only): + raise ValueError( + "--transformer-config-only cannot be combined with other partial conversion flags" + ) + + transformer_config = { + **BASE_TRANSFORMER_CONFIG, + **TRANSFORMER_VARIANTS[args.variant], + } + if transformer_config_only: + transformer_dir = output / "transformer" + transformer_dir.mkdir(parents=True, exist_ok=True) + _write_json(transformer_dir / "config.json", transformer_config) + transformer_model_index = { + key: value + for key, value in model_index.items() + if key.startswith("_") or key == "transformer" + } + _write_json(output / "model_index.json", transformer_model_index) + print(f"Wrote MMAudio {args.variant} transformer export skeleton to {output}") + return + + if not args.preprocessor_only: + if args.transformer_checkpoint is None: + raise ValueError( + "--transformer-checkpoint is required unless --preprocessor-only is used") + transformer_state = _load_torch_state(args.transformer_checkpoint) + # Official ``MMAudio.load_weights`` discards these derived buffers. + 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) + + if args.transformer_only: + transformer_model_index = { + key: value + for key, value in model_index.items() + if key.startswith("_") or key == "transformer" + } + _write_json(output / "model_index.json", transformer_model_index) + print(f"Converted MMAudio {args.variant} transformer to {output}") + return + + required_assets = { + "--audio-vae-checkpoint": args.audio_vae_checkpoint, + "--synchformer-checkpoint": args.synchformer_checkpoint, + "--dfn5b-dir": args.dfn5b_dir, + } + if not args.preprocessor_only: + if mode == "16k": + required_assets["--bigvgan-checkpoint"] = args.bigvgan_checkpoint + else: + required_assets["--bigvgan-dir"] = args.bigvgan_dir + missing_assets = [name for name, path in required_assets.items() if path is None] + if missing_assets: + raise ValueError( + "Full pipeline conversion requires " + ", ".join(missing_assets) + ) + + assert args.audio_vae_checkpoint is not None + assert args.synchformer_checkpoint is not None + assert args.dfn5b_dir is not None + if not args.preprocessor_only: + if mode == "16k": + assert args.bigvgan_checkpoint is not None + else: + assert args.bigvgan_dir is not None 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") + if args.preprocessor_only: + audio_vae_state = vae_state + need_encoder = True + if not any(key.startswith("encoder.") for key in audio_vae_state): + raise ValueError( + "Audio VAE checkpoint did not contain encoder weights") + else: + audio_vae_state = { + key: tensor + for key, tensor in vae_state.items() + if key.startswith("decoder.") or key in {"data_mean", "data_std"} + } + need_encoder = False + if not audio_vae_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}, + audio_vae_state, + { + "_class_name": "MMAudioVAE", + "mode": mode, + "need_encoder": need_encoder + }, ) synchformer_state = _load_torch_state(args.synchformer_checkpoint) @@ -304,14 +448,38 @@ def convert(args: argparse.Namespace) -> None: _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") + if args.preprocessor_only: + preprocess_components = { + key: value + for key, value in model_index.items() + if key.startswith("_") or key in { + "audio_vae", + "text_encoder", + "tokenizer", + "image_encoder", + "image_encoder_2", + } + } + preprocess_components["_fastvideo_preprocessor_only"] = True + _write_json(output / "model_index.json", preprocess_components) + print(f"Converted MMAudio {mode} preprocessing components to {output}") + return + + if mode == "16k": + assert args.bigvgan_checkpoint is not None + bigvgan_config = dict(BIGVGAN_16K_CONFIG) + bigvgan_state = _load_torch_state(args.bigvgan_checkpoint) + else: + assert args.bigvgan_dir is not None + 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( @@ -325,17 +493,53 @@ def convert(args: argparse.Namespace) -> None: "use_reference_discrete_timesteps": True, }, ) - _write_json(output / "model_index.json", MODEL_INDEX) - print(f"Converted MMAudio large_44k_v2 components to {output}") + _write_json(output / "model_index.json", model_index) + print(f"Converted MMAudio {args.variant} 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( + "--variant", + choices=tuple(TRANSFORMER_VARIANTS), + default="large_44k_v2", + help="Transformer architecture. v1 variants are required for training.", + ) + parser.add_argument("--transformer-checkpoint", type=Path) + parser.add_argument("--audio-vae-checkpoint", type=Path) + parser.add_argument("--synchformer-checkpoint", type=Path) + parser.add_argument("--dfn5b-dir", type=Path) + parser.add_argument( + "--bigvgan-dir", + type=Path, + help="NVIDIA BigVGAN-v2 directory used by 44.1 kHz variants.", + ) + parser.add_argument( + "--bigvgan-checkpoint", + type=Path, + help="Official best_netG.pt used by the small_16k variant.", + ) + parser.add_argument( + "--transformer-only", + action="store_true", + help="Write only the transformer component needed by training.", + ) + parser.add_argument( + "--transformer-config-only", + action="store_true", + help=( + "Write a weight-free transformer component tree used as the " + "DCP export template for from-scratch training." + ), + ) + parser.add_argument( + "--preprocessor-only", + action="store_true", + help=( + "Write only the VAE encoder, DFN5B text/vision encoders, " + "Synchformer, and tokenizer needed for offline feature extraction." + ), + ) parser.add_argument("--output", type=Path, required=True) return parser.parse_args() diff --git a/tests/local_tests/mmaudio/PORT_STATUS.md b/tests/local_tests/mmaudio/PORT_STATUS.md index 329fc7736..e838ed414 100644 --- a/tests/local_tests/mmaudio/PORT_STATUS.md +++ b/tests/local_tests/mmaudio/PORT_STATUS.md @@ -6,8 +6,8 @@ - workload_types: `V2A`, `T2A` - official_ref: `../MMAudio` at `974010a026c731054592d8f777218bd9d85a6c24` - first_variant: `large_44k_v2` -- phase: `native_pipeline_complete` -- status: `real_weight_parity_pass` +- phase: `all_official_variants_and_scratch_training_integrated` +- status: `large_v2_inference_parity_and_4gpu_scratch_training_pass` - last_updated: `2026-08-02` ## Native Components @@ -24,8 +24,9 @@ ## Pipeline Integration - Pipeline: `fastvideo/pipelines/basic/mmaudio/MMAudioPipeline` -- Config: `fastvideo/configs/pipelines/mmaudio.py::MMAudioV2AConfig` -- Preset: `mmaudio_large_44k_v2` +- Config: five variant-specific classes in `fastvideo/configs/pipelines/mmaudio.py` +- Presets: `mmaudio_small_16k`, `mmaudio_small_44k`, + `mmaudio_medium_44k`, `mmaudio_large_44k`, `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` @@ -35,7 +36,7 @@ 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: +The inference 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; @@ -50,16 +51,85 @@ The V2A preprocessing contract is identical to official MMAudio: - Total local artifact size: about 9.1 GB - Converted weights and official assets remain ignored/untracked. +The v1 training smoke additionally uses a transformer-only component tree at +`converted_weights/mmaudio/small_44k` (about 601 MB). The converter supports +`small_16k`, `small_44k`, `medium_44k`, `large_44k`, and the inference-default +`large_44k_v2`; `--transformer-only` avoids duplicating inference-only assets. + +## Training Integration + +- Model adapter: `fastvideo/train/models/mmaudio/MMAudioModel` +- Shared method: shape-agnostic + `fastvideo/train/methods/fine_tuning/flow_matching.py` +- Data: official-compatible TensorDict memmaps through + `fastvideo/dataset/mmaudio_feature_dataset.py` +- Config: four `*_from_scratch.yaml` v1 configs plus the pretrained + `small_44k.yaml` fine-tuning config +- Numerical recipe: posterior sampling, latent normalization, logit-normal + time, linear flow interpolation, independent video/text CFG dropout, and + velocity MSE match upstream `Runner.train_fn` +- Distributed contract: real single-GPU pretrained FSDP step and real 4-GPU + scratch-initialized FSDP step passed; DCP save/restore passed on the pretrained path +- Scratch contract: rank-0 latent statistics reduction/cache, fixed CLIP empty + string loading, random DiT initialization, flow loss/backward, AdamW, + official warmup/MultiStepLR composition, gradient clipping, and EMA are wired +- Production code imports no `mmaudio.*` modules. + +## Native VGGSound Preprocessing + +- Shared orchestration: `fastvideo/workflow/preprocess/preprocess_workflow_v2a.py` +- Dataset adapter: `fastvideo/dataset/vggsound.py` +- MMAudio feature stage: `fastvideo/pipelines/preprocess/mmaudio/stages.py` +- Sharded resumable cache: `fastvideo/dataset/v2a_feature_cache.py` +- Input manifests use the official `idlabel` contract; `label` is passed as + the full caption to the DFN5B text encoder. +- `dataset_metadata_path` explicitly selects one manifest and `dataset_split` + preserves the official train-only peak normalization behavior. +- Train, val, and test use one pipeline class and separate output roots. +- Production preprocessing imports no modules from the reference checkout. +- The reference training loader uses the removed + `torio.io.StreamingMediaDecoder`. The default FastVideo backend uses + `torchaudio.load` and a PyAV-hosted FFmpeg `fps` graph. An explicit `torio` + backend is available in a separate PyTorch 2.7 environment for exact + reference-media audits. Both feed the same native FastVideo feature stage; + inference keeps its separate timestamp sampler and remains unchanged. +- The FFmpeg training sampler fixes fractional-FPS boundary clips that decoded + as 64 DFN frames but only 199 Synchformer frames with timestamp comparison. + Three previously failed VGGSound clips now produce exactly 64/200 frames, and + one passed the full FastVideo feature extraction and shard-writing path. +- The isolated torio backend passed exact audio, CLIP-frame, and sync-frame + parity against the official VGGSound adapter plus a full FastVideo shard smoke. +- Torio preprocessing now decodes with a bounded background thread pool per + rank and keeps two batches queued, overlapping CPU media decode with GPU + feature extraction without DataLoader shared-memory tensor transfer. Cached + IDs are filtered before decode. Four threads per rank were fastest in a + 16-video decode check; a 4-GPU/16-video end-to-end smoke wrote 4/4 samples on + every rank and reduced the feature loop from 19.25 seconds on one GPU to a + 6.65-second slowest rank on the four-GPU run. +- Timestamp, FFmpeg, and torio caches must use separate output roots. + ## Parity Evidence | Scope | Result | |---|---| -| Official/FastVideo video preprocessing | exact (`clip max_abs=0`, `sync max_abs=0`) | +| Official/FastVideo inference video preprocessing | exact (`clip max_abs=0`, `sync max_abs=0`) | +| Training FFmpeg fps-filter boundary regression | pass (64 DFN frames, 200 Synchformer frames) | +| Previously failed real VGGSound clips | 3/3 recovered; end-to-end feature extraction 1/1 | +| Official/FastVideo torio training media preprocessing | exact audio/CLIP/sync tensors (`atol=0`, `rtol=0`) | +| Torio native FastVideo feature/shard path | pass; all five feature tensors written with expected shapes | | 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 | +| Real `small_44k` v1 transformer forward | exact BF16 parity | +| Real `small_44k` v1 FastVideo train step | pass (forward/backward/AdamW/grad clip) | +| Real `small_44k` 4-GPU from-scratch train step | pass (5.87 s, one feature shard) | +| Real DCP checkpoint and resume to next step | pass | +| Local MMAudio parity suite | `21 passed, 1 skipped` (the skipped test is the opt-in full gate) | +| Combined parity/dataset/training regression | `33 passed, 1 skipped` | +| Real filtered-caption manifests | exact ID/order/caption/count parity: train 180,062; val 2,047; test 15,221 | +| Filtered-caption tokenizer lengths | 0/197,330 exceed the 77-token contract; maximum 74 | | 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 | @@ -92,7 +162,9 @@ official `large_44k_v2`, DFN5B, Synchformer, VAE, and BigVGAN assets. - 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. +- Real-weight waveform parity for `small_16k`, `small_44k`, `medium_44k`, and + `large_44k`; their configs/converter/presets are implemented. - 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. +- Full-dataset loss/quality convergence validation. +- The training config intentionally keeps `data_path` empty until the train + feature cache has been produced. diff --git a/tests/local_tests/mmaudio/README.md b/tests/local_tests/mmaudio/README.md index 9fe242dc9..b2f7c381d 100644 --- a/tests/local_tests/mmaudio/README.md +++ b/tests/local_tests/mmaudio/README.md @@ -37,8 +37,10 @@ only in the shell when a future gated asset requires one. - 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. +- All five official architecture variants now have FastVideo configs, presets, + registry entries, and conversion support. Real-weight parity is complete for + `large_44k_v2`; the other four remain local follow-up gates. Sequence/tensor + parallel optimization and quality-reference publication are deferred. ## Shared Environment Setup @@ -147,3 +149,93 @@ pytest tests/local_tests/mmaudio/test_mmaudio_pipeline_parity.py -v -s presented as runnable. - MMAudio checkpoints are documented upstream as CC-BY-NC 4.0; converted model publishing must preserve the applicable license and attribution. + +## Native VGGSound Feature Extraction + +FastVideo consumes the same `idlabel` split manifests as the reference +`mmaudio.data.extraction.vgg_sound.VGGSound` adapter. `label` is a full caption, +not the original VGGSound category. The reference checkout is used only for +local parity review and is not imported by production preprocessing. + +Training video frames are resampled with one PyAV-hosted FFmpeg filter graph: +8 FPS for DFN5B and 25 FPS for Synchformer. This matches the reference +`StreamingMediaDecoder.add_basic_video_stream(frame_rate=...)` behavior while +remaining compatible with FastVideo's current PyTorch environment. The +inference timestamp sampler is a separate default path and is not changed. + +```bash +DATASET_PATH=/path/to/VGGSound SPLIT=train GPU_NUM=4 \ + bash examples/training/finetune/mmaudio/preprocess_vggsound.sh +``` + +The launcher supports `train`, `val`, and `test` through the same FastVideo +pipeline. The corresponding default manifests are +`sets/filtered_caption/vgg-{split}-filtered-caption.tsv`, and caches are written +under `mmaudio_features/{split}`. The cache is resumable: rerunning the command +skips IDs already written and retries pending/failed IDs. Failure logs are +append-only, so deduplicate by ID when auditing multiple runs. + +Caches produced with the earlier timestamp-based training sampler are valid +TensorDict data but are not preprocessing-compatible with the FFmpeg sampler. +Do not resume into such a cache. Rebuild into a new output root instead: + +```bash +OUTPUT_ROOT=/path/to/VGGSound/mmaudio_features_ffmpeg \ + ./examples/training/finetune/mmaudio/run_preprocess_vggsound.sh +``` + +### Optional official-compatible torio reader + +For audits that require the original training media reader, FastVideo also +provides an opt-in backend in +`fastvideo/pipelines/preprocess/mmaudio/torio_media_reader.py`. It reproduces +the official `StreamingMediaDecoder` streams, audio normalization/resampling, +and both torchvision transform chains. The rest of the operation remains the +native FastVideo preprocessing pipeline: VAE, DFN5B, Synchformer, text encoder, +distributed workflow, progress bar, resumable writer, and failure logging. +Production code does not import `mmaudio.*` or execute the reference checkout. + +Torio was removed from current PyTorch releases, so keep its PyTorch 2.7 CUDA +12.8 stack in a sibling environment rather than changing FastVideo's `.venv`: + +```text +/path/to/workspace/FastVideo +/path/to/workspace/envs/fastvideo-mmaudio-torio +``` + +The verified environment uses Python 3.12, PyTorch 2.7.0+cu128, +torchvision/torchaudio 0.22.0/2.7.0, NumPy 1.26.4, and an editable +`fastvideo --no-deps` install. Select it only through the dedicated launcher: + +```bash +DATASET_PATH=/path/to/VGGSound \ +TORIO_VENV_PATH=/path/to/workspace/envs/fastvideo-mmaudio-torio \ +CUDA_VISIBLE_DEVICES=0,1,2,3 GPU_NUM=4 BATCH_SIZE=16 \ +DATALOADER_NUM_WORKERS=4 \ + ./examples/training/finetune/mmaudio/run_preprocess_vggsound_torio.sh +``` + +By default this processes `train val test` with their existing split manifests +and writes a separate cache under `mmaudio_features_torio/{split}`. Override +`SPLITS`, `OUTPUT_ROOT`, or the usual batching variables exactly as with the +standard launcher. Never resume a PyAV/FFmpeg cache into the torio output root. + +For torio, `DATALOADER_NUM_WORKERS` is the number of background decode threads +per GPU rank, not DataLoader subprocesses. The default is 4 based on local +VGGSound measurements. FastVideo keeps two batches queued, so CPU decoding of +the next batch overlaps GPU feature extraction without sending large tensors +through `/dev/shm`. IDs already present in the feature cache are filtered before +decode and are not read again on resume. + +Exact real-media parity against the official VGGSound dataset adapter is +covered by an opt-in local test: + +```bash +PYTHONNOUSERSITE=1 \ +MMAUDIO_TORIO_TEST_VIDEO=/path/to/an-8s-or-longer-video.mp4 \ +/path/to/workspace/envs/fastvideo-mmaudio-torio/bin/pytest -q \ + tests/local_tests/mmaudio/test_mmaudio_torio_preprocessing.py -s +``` + +This test imports the official dataset adapter only on the test side; it is not +a runtime dependency of FastVideo. diff --git a/tests/local_tests/mmaudio/test_mmaudio_audio_vae_parity.py b/tests/local_tests/mmaudio/test_mmaudio_audio_vae_parity.py index 9dd5005db..d570a3d57 100644 --- a/tests/local_tests/mmaudio/test_mmaudio_audio_vae_parity.py +++ b/tests/local_tests/mmaudio/test_mmaudio_audio_vae_parity.py @@ -46,6 +46,18 @@ def test_mmaudio_44k_audio_vae_state_structure() -> None: } +def test_mmaudio_44k_mel_converter_parity() -> None: + from mmaudio.ext.mel_converter import get_mel_converter + + from fastvideo.models.audio.mmaudio_processing import build_mmaudio_mel_converter + + waveform = torch.randn((2, 353_280), generator=torch.Generator().manual_seed(1234)) + expected = get_mel_converter("44k")(waveform) + actual = build_mmaudio_mel_converter("44k")(waveform) + + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + + @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(): diff --git a/tests/local_tests/mmaudio/test_mmaudio_torio_preprocessing.py b/tests/local_tests/mmaudio/test_mmaudio_torio_preprocessing.py new file mode 100644 index 000000000..63433a32b --- /dev/null +++ b/tests/local_tests/mmaudio/test_mmaudio_torio_preprocessing.py @@ -0,0 +1,78 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Opt-in parity for FastVideo's isolated torio preprocessing backend.""" + +from __future__ import annotations + +import os +import sys +import types +from pathlib import Path + +import pytest +import torch + + +REPO_ROOT = Path(__file__).resolve().parents[3] +OFFICIAL_REPO = Path( + os.environ.get("MMAUDIO_OFFICIAL_REPO", REPO_ROOT.parent / "MMAudio") +) + + +@pytest.mark.skipif( + not os.environ.get("MMAUDIO_TORIO_TEST_VIDEO"), + reason="Set MMAUDIO_TORIO_TEST_VIDEO to run real media parity.", +) +def test_fastvideo_torio_reader_matches_official_vggsound_dataset( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + pytest.importorskip("torio.io") + if not OFFICIAL_REPO.is_dir(): + pytest.skip(f"Official MMAudio checkout is missing: {OFFICIAL_REPO}") + + video_path = Path(os.environ["MMAUDIO_TORIO_TEST_VIDEO"]) + if not video_path.is_file(): + pytest.skip(f"Test video is missing: {video_path}") + + monkeypatch.syspath_prepend(str(OFFICIAL_REPO)) + dist_utils = types.ModuleType("mmaudio.utils.dist_utils") + dist_utils.local_rank = 0 + monkeypatch.setitem(sys.modules, "mmaudio.utils.dist_utils", dist_utils) + + from mmaudio.data.extraction.vgg_sound import VGGSound + + video_id = video_path.stem + videos = tmp_path / "videos" + videos.mkdir() + (videos / f"{video_id}.mp4").symlink_to(video_path.resolve()) + manifest = tmp_path / "captions.tsv" + manifest.write_text( + f"id\tlabel\n{video_id}\tA reference preprocessing sample.\n", + encoding="utf-8", + ) + + official = VGGSound( + videos, + tsv_path=manifest, + sample_rate=44_100, + duration_sec=8.0, + audio_samples=353_280, + normalize_audio=True, + ).sample(0) + + from fastvideo.pipelines.preprocess.mmaudio.torio_media_reader import ( + preprocess_mmaudio_media_with_torio, + ) + + audio, clip, sync, duration = preprocess_mmaudio_media_with_torio( + video_path, + duration_s=8.0, + target_sample_rate=44_100, + target_samples=353_280, + normalize_audio=True, + ) + + torch.testing.assert_close(audio, official["audio"], atol=0, rtol=0) + torch.testing.assert_close(clip, official["clip_video"], atol=0, rtol=0) + torch.testing.assert_close(sync, official["sync_video"], atol=0, rtol=0) + assert duration == 8.0 diff --git a/tests/local_tests/mmaudio/test_mmaudio_transformer_parity.py b/tests/local_tests/mmaudio/test_mmaudio_transformer_parity.py index d06d2d717..374729a1f 100644 --- a/tests/local_tests/mmaudio/test_mmaudio_transformer_parity.py +++ b/tests/local_tests/mmaudio/test_mmaudio_transformer_parity.py @@ -21,6 +21,12 @@ OFFICIAL_WEIGHTS = Path( REPO_ROOT.parent / "MMAudio/weights/mmaudio_large_44k_v2.pth", ) ) +SMALL_44K_OFFICIAL_WEIGHTS = Path( + os.environ.get( + "MMAUDIO_SMALL_44K_TRANSFORMER_WEIGHTS", + REPO_ROOT.parent / "MMAudio/weights/mmaudio_small_44k.pth", + ) +) def _require_cuda_and_weights() -> None: @@ -33,7 +39,8 @@ def _require_cuda_and_weights() -> None: ) -def test_mmaudio_transformer_implementation_parity() -> None: +@pytest.mark.parametrize("v2", [False, True]) +def test_mmaudio_transformer_implementation_parity(v2: bool) -> None: from mmaudio.model.networks import MMAudio from fastvideo.configs.models.dits.mmaudio import ( @@ -56,7 +63,7 @@ def test_mmaudio_transformer_implementation_parity() -> None: "clip_seq_len": 3, "sync_seq_len": 8, "text_seq_len": 4, - "v2": True, + "v2": v2, } official = MMAudio(**kwargs) fastvideo = MMAudioTransformer(MMAudioTransformerConfig(arch_config=MMAudioArchConfig(**kwargs)), hf_config={}) @@ -78,6 +85,75 @@ def test_mmaudio_transformer_implementation_parity() -> None: torch.testing.assert_close(actual, expected, atol=1e-6, rtol=1e-6) +def test_mmaudio_small_44k_v1_transformer_parity() -> None: + if not torch.cuda.is_available(): + pytest.skip("MMAudio full-transformer parity requires CUDA") + if not SMALL_44K_OFFICIAL_WEIGHTS.is_file(): + pytest.skip( + "Official MMAudio small_44k weights are absent. Set " + "MMAUDIO_SMALL_44K_TRANSFORMER_WEIGHTS." + ) + + from mmaudio.model.networks import small_44k + + from fastvideo.configs.models.dits.mmaudio import ( + MMAudioArchConfig, + 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( + SMALL_44K_OFFICIAL_WEIGHTS, + map_location="cpu", + weights_only=True, + ) + arch = MMAudioArchConfig( + hidden_dim=448, + depth=12, + fused_depth=8, + num_heads=7, + v2=False, + ) + with torch.device(device), set_default_torch_dtype(dtype): + official = small_44k() + fastvideo = MMAudioTransformer( + MMAudioTransformerConfig(arch_config=arch), + hf_config={}, + ) + official.load_weights(dict(state)) + state.pop("latent_rot", None) + state.pop("clip_rot", None) + missing, unexpected = fastvideo.load_state_dict(state, strict=True) + assert missing == [] + assert unexpected == [] + official.eval() + fastvideo.eval() + + generator = torch.Generator(device=device).manual_seed(5678) + 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) + + torch.testing.assert_close(actual, expected, atol=1e-3, rtol=1e-3) + + def test_mmaudio_large_44k_v2_transformer_parity() -> None: _require_cuda_and_weights() from mmaudio.model.networks import large_44k_v2