diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 31805cdb2..49d8bc864 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -18,6 +18,7 @@ exclude: | fastvideo/train\.py| fastvideo/utils/.*| examples/.*| + \.agents/.*| .github/workflows/publish-fastvideo.yml| .github/workflows/_template-build-image.yml| docs/source/inference/support_matrix.md diff --git a/examples/inference/basic/basic_stable_audio.py b/examples/inference/basic/basic_stable_audio.py new file mode 100644 index 000000000..b44bb0484 --- /dev/null +++ b/examples/inference/basic/basic_stable_audio.py @@ -0,0 +1,77 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Stable Audio Open 1.0 — text-to-audio (baseline) example. + +User story (game-audio designer, prototyping): + "I'm prototyping a level and I need 6 seconds of background + ambience — gentle wind, distant thunder, a hint of birdsong. I + don't want to dig through a sound library; I want to type what I + hear in my head and get a wav back. If it's wrong I'll iterate + on the prompt. This is the first stop." + +User story (musician sketching ideas): + "I want to bounce a 30s lo-fi drum loop to use as a placeholder + bed while I build the rest of the track. Type prompt, get audio, + drop into the DAW. The actual production beat I'll record + myself, but I need *something* to write the chords against." + +User story (researcher exploring the model): + "First time touching Stable Audio Open — what does it sound + like at default settings? This is the smallest amount of code + that goes from prompt to mp4." + +How it works: + Pure text-to-audio (T2A). The pipeline runs: + T5 + NumberConditioner -> StableAudioDiT -> Oobleck VAE + via the `dpmpp-3m-sde` k-diffusion sampler. All components are + FastVideo-native — no diffusers / transformers model imports at + runtime (see REVIEW item 30). Mirrors upstream + `stable_audio_tools.inference.generation.generate_diffusion_cond` + bit-for-bit (~0.2% abs_mean drift on 25 steps). + +Tunable knobs (the "creative dials"): + audio_end_in_s + 1–6 — quick ideation (sub-10s wall clock at 100 steps) + 10–30 — full musical phrase / loop length (the README example + uses 30s) + 47.5 — model maximum (full sample_size = 2097152 / 44100 Hz) + num_inference_steps + 25 — fast preview, occasional artifacts + 100 — preset default (matches the HF model card) + 250 — diminishing returns past here + guidance_scale + 3 — looser, more variation per seed + 7 — preset default; matches README + 12+ — sharper but can sound "fried" + +Prerequisites: + 1. Accept the terms on https://huggingface.co/stabilityai/stable-audio-open-1.0 + and export your HF token in the shell: + export HF_TOKEN=hf_... + 2. Install optional inference deps (one-time): + pip install k_diffusion einops_exts alias_free_torch torchsde +""" +from fastvideo import VideoGenerator + +PROMPT = "Lo-fi hip hop instrumental with vinyl crackle and gentle piano." + + +def main() -> None: + generator = VideoGenerator.from_pretrained( + "FastVideo/stable-audio-open-1.0-Diffusers", + num_gpus=1, + ) + output_path = "outputs_audio/stable_audio_basic/output_stable_audio.wav" + generator.generate_video( + prompt=PROMPT, + output_path=output_path, + save_video=True, + # 6-second clip; the model max is ~47.5s. + audio_end_in_s=6.0, + # The registered preset gives 100 steps + CFG=7.0 by default; + # override num_inference_steps / guidance_scale here for QA. + ) + generator.shutdown() + + +if __name__ == "__main__": + main() diff --git a/examples/inference/basic/basic_stable_audio_a2a.py b/examples/inference/basic/basic_stable_audio_a2a.py new file mode 100644 index 000000000..7b7dd00e7 --- /dev/null +++ b/examples/inference/basic/basic_stable_audio_a2a.py @@ -0,0 +1,77 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Stable Audio Open 1.0 — audio-to-audio variation example. + +User story (musician, late at night): + "I generated this 12-second lo-fi loop earlier and I love the chord + progression and overall vibe, but the snare hit at 0:08 sounds wrong + and the rhythm feels stiff. I don't want to start over from scratch + and lose what's working — I want the model to keep the harmony and + mood but reroll the percussion + groove." + +User story (sound designer, on a deadline): + "I have one good 'sword clang' SFX. The art director wants 8 sibling + variations that all feel like the same sword from different angles — + same metal, same weight, slightly different impact. I'd rather + refine my one good take than text-prompt my way through 50 misses." + +Pass `init_audio=path/to/clip` (any wav/mp3/mp4/m4a/flac the standard +deps decode) and the model will use it as a starting point for the +text prompt instead of pure noise. + +Picking `init_audio_strength` (0.0 to 1.0): + + Higher = closer to the source clip. Lower = more transformation. + (Same convention as the "Input Audio Strength" slider in + Stability's commercial Stable Audio web UI, so values transfer + directly.) + + | strength | what you get | + |----------|----------------------------------------------------| + | 1.00 | Output ≈ reference. No transformation. | + | 0.85 | Texture micro-variation only. | + | 0.70 | Light reroll, same instruments. | + | 0.60 | Default. Instrument identity is replaceable | + | | (cello can take over from piano on the same notes).| + | 0.50 | Heavy — only melody / chord progression survives. | + | 0.30 | Reference acts as a loose mood prompt. | + | 0.00 | Plain T2A — reference ignored. | + + Rule of thumb by intent: + * "Fix one part of this clip" -> 0.75 .. 0.85 + * "Same notes, different instrument" -> 0.55 .. 0.65 + * "Same chord progression, new content" -> 0.40 .. 0.55 + * "Use this as a loose mood prompt" -> 0.20 .. 0.35 + + If the reference timbre is bleeding through more than you want, + lower it; if the structure is gone, raise it. + +Prerequisites: same as `basic_stable_audio.py`. +""" +from fastvideo import VideoGenerator + +PROMPT = "Change the piano to a cello playing the same notes" +# Path to any audio-bearing file (wav, mp3, mp4, m4a, flac, ...). +# Set to `None` to skip A2A and run plain T2A. +INIT_AUDIO_PATH: str | None = None +# Reference fidelity in [0, 1] -- higher = closer to source. +INIT_AUDIO_STRENGTH = 0.6 + + +def main() -> None: + generator = VideoGenerator.from_pretrained( + "FastVideo/stable-audio-open-1.0-Diffusers", + num_gpus=1, + ) + generator.generate_video( + prompt=PROMPT, + output_path="outputs_audio/stable_audio_a2a/output_a2a.wav", + save_video=True, + audio_end_in_s=6.0, + init_audio=INIT_AUDIO_PATH, + init_audio_strength=INIT_AUDIO_STRENGTH, + ) + generator.shutdown() + + +if __name__ == "__main__": + main() diff --git a/examples/inference/basic/basic_stable_audio_inpaint.py b/examples/inference/basic/basic_stable_audio_inpaint.py new file mode 100644 index 000000000..126bbc54a --- /dev/null +++ b/examples/inference/basic/basic_stable_audio_inpaint.py @@ -0,0 +1,84 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Stable Audio Open 1.0 — inpainting / outpainting (loop extension) example. + +User story (loop extension — the killer app): + "I have a 6-second drum loop my client likes. They want it as + background bed for a 30-second ad. I need it to loop seamlessly, + but a hard cut every 6s sounds bad. Let me extend it to 30s, + keeping the first 6s exactly as-is and letting the model continue + the groove for the remaining 24s." + +User story (audio repair): + "There's a microphone bump at 0:14 in this 30-second field + recording — really obvious in headphones. Mask out 0:13 to 0:15 + and let the model regenerate plausible ambience that blends in. + Everything else stays exactly as I recorded it." + +User story (transition smoothing): + "I have two 10-second clips I want to crossfade. Mask out a 1s + overlap region in the middle and let the model invent a coherent + transition between the two." + +How it works (RePaint-style blending): + Stable Audio Open 1.0 wasn't trained as an inpainting model + (`model_type=diffusion_cond`, not `diffusion_cond_inpaint`), so we + can't use the upstream's mask-conditioned approach directly. We + use the RePaint trick instead, which works on any v-prediction + diffusion model: + + 1. Encode the reference clip into latent space. + 2. At every denoising step `i`, replace the kept region of the + in-flight latent (where mask == 1) with the reference + re-noised to the next timestep's sigma. Only the unkept + region (mask == 0) is freely denoised. + 3. After the loop, the kept region is exactly the reference; + the unkept region is freshly generated content. + + This is approximate compared to a properly trained inpainting + checkpoint — the seam between kept/unkept can have slight EQ + discontinuity — but it works on the existing public model. + +Tunable: the mask is a 1-D tensor in {0, 1} at the model's sample +rate. Conventions: + 1.0 = keep this sample from the reference + 0.0 = regenerate this sample + +Prerequisites: same as `basic_stable_audio.py`. +""" +import os + +from fastvideo import VideoGenerator + +PROMPT = "Steady lo-fi hip hop drum loop with vinyl crackle." +# Required: path to the reference audio file (wav, mp3, mp4, m4a, flac, +# ...) you want to extend or repair. The pipeline raises if a mask is +# passed without a reference, so this must be a real path. +REFERENCE_AUDIO_PATH = "path/to/your/loop.wav" +KEEP_SECONDS = 6.0 # first KEEP_SECONDS preserved exactly +TOTAL_SECONDS = 12.0 # extend the loop to this duration + + +def main() -> None: + if not os.path.isfile(REFERENCE_AUDIO_PATH): + raise FileNotFoundError( + f"REFERENCE_AUDIO_PATH={REFERENCE_AUDIO_PATH!r} does not exist. " + "Edit this script to point at a real audio file (wav/mp3/mp4/" + "m4a/flac) before running.") + generator = VideoGenerator.from_pretrained( + "FastVideo/stable-audio-open-1.0-Diffusers", + num_gpus=1, + ) + generator.generate_video( + prompt=PROMPT, + output_path="outputs_audio/stable_audio_inpaint/output_inpaint.wav", + save_video=True, + audio_end_in_s=TOTAL_SECONDS, + inpaint_audio=REFERENCE_AUDIO_PATH, + # Tuple form: keep first KEEP_SECONDS, regenerate the rest. + inpaint_mask=(KEEP_SECONDS, TOTAL_SECONDS), + ) + generator.shutdown() + + +if __name__ == "__main__": + main() diff --git a/examples/inference/basic/basic_stable_audio_small.py b/examples/inference/basic/basic_stable_audio_small.py new file mode 100644 index 000000000..6a4306872 --- /dev/null +++ b/examples/inference/basic/basic_stable_audio_small.py @@ -0,0 +1,53 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Stable Audio Open Small — fast / lightweight T2A example. + +User story (interactive UI builder): + "I'm building a sound-design UI where the user types a prompt and + we want sub-2-second feedback so the experience feels like + autocomplete, not a render queue. The full Stable Audio Open 1.0 + takes ~8s on a single GPU; the small variant takes a fraction of + that — quality is lower but completely usable for real-time + iteration." + +User story (overnight batch jobs): + "I'm generating 10,000 short SFX variants for a procedural game. + Wall-clock matters more than per-clip polish — give me the small + model so I can fit the run in one night instead of a week." + +How it works: + The small variant is a separate Stability AI checkpoint + (`stabilityai/stable-audio-open-small`) that ships the same Oobleck + VAE as the 1.0 base model but a smaller / faster DiT (`embed_dim=1024`, + `depth=16`, `qk_norm="ln"`) and only one duration conditioner + (`seconds_total`, no `seconds_start`). FastVideo loads from the + converted Diffusers-format repo `FastVideo/stable-audio-open-small-Diffusers` + via the standard component loader; per-variant arch fields come + from `transformer/config.json` and `conditioner/config.json`. + +Prerequisites: same as `basic_stable_audio.py`. The converted repo is +public so no gated-access flow is required. +""" +from fastvideo import VideoGenerator + +PROMPT = "Lo-fi hip hop instrumental with vinyl crackle and gentle piano." + + +def main() -> None: + generator = VideoGenerator.from_pretrained( + "FastVideo/stable-audio-open-small-Diffusers", + num_gpus=1, + ) + output_path = "outputs_audio/stable_audio_small/output_stable_audio_small.wav" + generator.generate_video( + prompt=PROMPT, + output_path=output_path, + save_video=True, + # Small variant trains on a ~11.9s window — keep `audio_end_in_s` + # at or below that. + audio_end_in_s=6.0, + ) + generator.shutdown() + + +if __name__ == "__main__": + main() diff --git a/fastvideo/api/results.py b/fastvideo/api/results.py index 533148f84..dd52b03f1 100644 --- a/fastvideo/api/results.py +++ b/fastvideo/api/results.py @@ -15,6 +15,7 @@ class GenerationResult: samples: Any | None = None frames: Any | None = None audio: Any | None = None + audio_sample_rate: int | None = None size: tuple[int, int, int] | None = None generation_time: float | None = None logging_info: Any | None = None @@ -44,6 +45,7 @@ class GenerationResult: "samples", "frames", "audio", + "audio_sample_rate", "size", "generation_time", "logging_info", @@ -62,6 +64,7 @@ class GenerationResult: samples=result.get("samples"), frames=result.get("frames"), audio=result.get("audio"), + audio_sample_rate=result.get("audio_sample_rate"), size=result.get("size"), generation_time=result.get("generation_time"), logging_info=result.get("logging_info"), @@ -80,6 +83,7 @@ class GenerationResult: "samples": self.samples, "frames": self.frames, "audio": self.audio, + "audio_sample_rate": self.audio_sample_rate, "size": self.size, "generation_time": self.generation_time, "logging_info": self.logging_info, diff --git a/fastvideo/api/sampling_param.py b/fastvideo/api/sampling_param.py index b18e35326..f89d97e20 100644 --- a/fastvideo/api/sampling_param.py +++ b/fastvideo/api/sampling_param.py @@ -112,6 +112,33 @@ class SamplingParam: ltx2_stg_blocks_video: list[int] = field(default_factory=lambda: [29]) ltx2_stg_blocks_audio: list[int] = field(default_factory=lambda: [29]) + # Stable Audio (T2A): clip start/end in seconds. Honored by + # `StableAudioConditioningStage` + `StableAudioDecodingStage`. Other + # families ignore them. + audio_start_in_s: float | None = None + audio_end_in_s: float | None = None + + # Stable Audio audio-to-audio (variation): + # `init_audio` -- a path or `[B, C, samples]` waveform at the model + # sample rate; the pipeline encodes it via the VAE + # and uses it as the starting latent. + # `init_audio_strength` -- 0..1, higher = closer to the reference + # (matches the convention of Stability's + # commercial Stable Audio 2.0 UI). 1.0 ~= + # VAE round-trip, 0.0 ~= plain T2A. + # `init_noise_level` -- legacy raw `sigma_max` override (0.3..500, + # higher = more freedom). Kept for callers + # that already use it; prefer `init_audio_strength`. + init_audio: Any = None + init_audio_strength: float | None = None + init_noise_level: float | None = None + + # Stable Audio inpainting (RePaint-style): `inpaint_audio` is the + # reference clip, `inpaint_mask` is a [samples] tensor in {0, 1} where + # 1 means *keep the reference* and 0 means *regenerate*. + inpaint_audio: Any = None + inpaint_mask: Any = None + # Continuation state carried across streaming/multi-segment calls. continuation_state: ContinuationState | None = None # When True, the pipeline returns a ContinuationState on the result so diff --git a/fastvideo/configs/models/dits/__init__.py b/fastvideo/configs/models/dits/__init__.py index 89ae14938..b6817e254 100644 --- a/fastvideo/configs/models/dits/__init__.py +++ b/fastvideo/configs/models/dits/__init__.py @@ -5,11 +5,13 @@ from fastvideo.configs.models.dits.hunyuanvideo import HunyuanVideoConfig from fastvideo.configs.models.dits.hunyuanvideo15 import HunyuanVideo15Config from fastvideo.configs.models.dits.longcat import LongCatVideoConfig from fastvideo.configs.models.dits.ltx2 import LTX2VideoConfig +from fastvideo.configs.models.dits.stable_audio import StableAudioConfig from fastvideo.configs.models.dits.wanvideo import WanVideoConfig from fastvideo.configs.models.dits.hyworld import HYWorldConfig from fastvideo.configs.models.dits.kandinsky5 import Kandinsky5VideoConfig __all__ = [ "HunyuanVideoConfig", "HunyuanVideo15Config", "HunyuanGameCraftConfig", "WanVideoConfig", "CosmosVideoConfig", - "Cosmos25VideoConfig", "LongCatVideoConfig", "LTX2VideoConfig", "HYWorldConfig", "Kandinsky5VideoConfig" + "Cosmos25VideoConfig", "LongCatVideoConfig", "LTX2VideoConfig", "HYWorldConfig", "Kandinsky5VideoConfig", + "StableAudioConfig" ] diff --git a/fastvideo/configs/models/dits/stable_audio.py b/fastvideo/configs/models/dits/stable_audio.py new file mode 100644 index 000000000..659036f79 --- /dev/null +++ b/fastvideo/configs/models/dits/stable_audio.py @@ -0,0 +1,76 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Config for the Stable Audio Open 1.0 DiT. + +Note: the SA pipeline bypasses the standard `ComposedPipelineBase` +component loader because the published HF repo ships a single monolithic +`model.safetensors` (no Diffusers-style `model_index.json` or +per-subfolder layout). The arch fields and `param_names_mapping` here +document the architecture and key remap so the same conventions used by +the rest of the DiT family apply (FSDP shard conditions, supported +attention backends, future loader integrations) — they are not currently +consumed by `fastvideo/models/loader/fsdp_load.py` for SA. +""" +from __future__ import annotations + +from dataclasses import dataclass, field + +from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig +from fastvideo.platforms import AttentionBackendEnum + + +def _is_transformer_layer(n: str, m) -> bool: + # Matches `transformer.layers.{i}` in the SA DiT module tree. + parts = n.split(".") + return (len(parts) >= 3 and parts[-3] == "transformer" and parts[-2] == "layers" and parts[-1].isdigit()) + + +@dataclass +class StableAudioArchConfig(DiTArchConfig): + _fsdp_shard_conditions: list = field(default_factory=lambda: [_is_transformer_layer]) + + # SA's checkpoint is `stable_audio_tools` raw format (not Diffusers), + # so the only remaps are: strip the `model.model.` host-pipeline + # prefix, and rename `nn.LayerNorm`'s `gamma`/`beta` to torch's + # canonical `weight`/`bias`. Linear / cross-attention naming already + # matches FastVideo's conventions, so no further remap is needed. + param_names_mapping: dict = field( + default_factory=lambda: { + r"^model\.model\.(.*?)\.gamma$": r"\1.weight", + r"^model\.model\.(.*?)\.beta$": r"\1.bias", + r"^model\.model\.(.*)$": r"\1", + }) + + # SA only supports backends compatible with single-GPU LocalAttention. + _supported_attention_backends: tuple[AttentionBackendEnum, ...] = ( + AttentionBackendEnum.FLASH_ATTN, + AttentionBackendEnum.TORCH_SDPA, + ) + + # Architecture constants (from the published `model_config.json` for + # `stabilityai/stable-audio-open-1.0`). + io_channels: int = 64 + embed_dim: int = 1536 + depth: int = 24 + num_attention_heads: int = 24 + cond_token_dim: int = 768 + global_cond_dim: int = 1536 + project_cond_tokens: bool = False + project_global_cond: bool = True + # Set to "ln" to wrap attention Q/K in LayerNorm (used by + # `stable-audio-open-small`; absent in the 1.0 base). + qk_norm: str | None = None + + def __post_init__(self) -> None: + super().__post_init__() + self.hidden_size = self.embed_dim + self.in_channels = self.io_channels + self.out_channels = self.io_channels + self.num_channels_latents = self.io_channels + self.attention_head_dim = self.embed_dim // self.num_attention_heads + + +@dataclass +class StableAudioConfig(DiTConfig): + arch_config: DiTArchConfig = field(default_factory=StableAudioArchConfig) + + prefix: str = "StableAudio" diff --git a/fastvideo/configs/models/encoders/__init__.py b/fastvideo/configs/models/encoders/__init__.py index ede12e603..485ac7612 100644 --- a/fastvideo/configs/models/encoders/__init__.py +++ b/fastvideo/configs/models/encoders/__init__.py @@ -7,9 +7,12 @@ from fastvideo.configs.models.encoders.qwen2_5 import Qwen2_5_VLConfig from fastvideo.configs.models.encoders.siglip import SiglipVisionConfig from fastvideo.configs.models.encoders.reason1 import Reason1ArchConfig, Reason1Config from fastvideo.configs.models.encoders.gemma import LTX2GemmaConfig +from fastvideo.configs.models.encoders.stable_audio_conditioner import (StableAudioConditionerArchConfig, + StableAudioConditionerConfig) __all__ = [ "EncoderConfig", "TextEncoderConfig", "ImageEncoderConfig", "BaseEncoderOutput", "CLIPTextConfig", "CLIPVisionConfig", "WAN2_1ControlCLIPVisionConfig", "LlamaConfig", "T5Config", "T5LargeConfig", "Qwen2_5_VLConfig", - "Reason1ArchConfig", "Reason1Config", "LTX2GemmaConfig", "SiglipVisionConfig" + "Reason1ArchConfig", "Reason1Config", "LTX2GemmaConfig", "SiglipVisionConfig", "StableAudioConditionerArchConfig", + "StableAudioConditionerConfig" ] diff --git a/fastvideo/configs/models/encoders/stable_audio_conditioner.py b/fastvideo/configs/models/encoders/stable_audio_conditioner.py new file mode 100644 index 000000000..39bbb4a6f --- /dev/null +++ b/fastvideo/configs/models/encoders/stable_audio_conditioner.py @@ -0,0 +1,80 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Config for the Stable Audio Open 1.0 multi-conditioner. + +The conditioner bundles three sub-conditioners — a T5 text encoder +(prompt) and two NumberConditioners (`seconds_start` / `seconds_total`) +— into the (cross_attn_cond, cross_attn_mask, global_embed) triple the +DiT consumes. The architecture is fully specified by the official +`stable_audio_tools` `MultiConditioner` config; the constants here +mirror that. +""" +from __future__ import annotations + +from dataclasses import dataclass, field + +from fastvideo.configs.models.base import ArchConfig +from fastvideo.configs.models.encoders.base import (EncoderArchConfig, EncoderConfig) + + +def _default_configs() -> list[dict]: + """Default = `stable-audio-open-1.0`'s three sub-conditioners.""" + return [ + { + "id": "prompt", + "type": "t5", + "config": { + "t5_model_name": "t5-base", + "max_length": 128 + } + }, + { + "id": "seconds_start", + "type": "number", + "config": { + "min_val": 0, + "max_val": 512 + } + }, + { + "id": "seconds_total", + "type": "number", + "config": { + "min_val": 0, + "max_val": 512 + } + }, + ] + + +@dataclass +class StableAudioConditionerArchConfig(EncoderArchConfig): + architectures: list[str] = field(default_factory=lambda: ["StableAudioMultiConditioner"]) + + # Shared embedding width across all sub-conditioners (T5 last-hidden + # dim and NumberEmbedder feature dim both = `cond_dim`). + cond_dim: int = 768 + + # Sub-conditioner identifiers. Order in `cross_attention_cond_ids` + # is the concat order for the cross-attn token sequence; order in + # `global_cond_ids` is the concat order for the global FiLM-style + # embedding. + cross_attention_cond_ids: tuple[str, ...] = ("prompt", "seconds_start", "seconds_total") + global_cond_ids: tuple[str, ...] = ("seconds_start", "seconds_total") + + # Per-sub-conditioner spec list (mirrors upstream + # `model_config.json.model.conditioning.configs`). Each entry is + # `{"id": ..., "type": "t5"|"number", "config": {...}}`. The default + # matches `stable-audio-open-1.0`; SA-small overrides via the + # `conditioner/config.json` shipped in the converted repo. + configs: list = field(default_factory=_default_configs) + + # Match official `stable_audio_tools/models/conditioners.py:334`: + # T5 is loaded directly in fp16. + t5_dtype: str = "float16" + + +@dataclass +class StableAudioConditionerConfig(EncoderConfig): + arch_config: ArchConfig = field(default_factory=StableAudioConditionerArchConfig) + + prefix: str = "stable_audio_conditioner" diff --git a/fastvideo/configs/models/encoders/t5.py b/fastvideo/configs/models/encoders/t5.py index 555009fa9..a96ce7f8f 100644 --- a/fastvideo/configs/models/encoders/t5.py +++ b/fastvideo/configs/models/encoders/t5.py @@ -41,6 +41,14 @@ class T5ArchConfig(TextEncoderArchConfig): text_len: int = 512 dtype: str | None = None gradient_checkpointing: bool = False + # Extra fields present in upstream HF T5Config but unused by FastVideo's + # encoder. Declared here so `update_model_arch` doesn't reject them when + # loading repos like `stabilityai/stable-audio-open-1.0` that ship the + # full HF config. + n_positions: int = 512 + decoder_start_token_id: int = 0 + output_past: bool = True + task_specific_params: dict | None = None stacked_params_mapping: list[tuple[str, str, str]] = field(default_factory=lambda: [ # (param_name, shard_name, shard_id) (".qkv_proj", ".q", "q"), diff --git a/fastvideo/configs/models/vaes/__init__.py b/fastvideo/configs/models/vaes/__init__.py index 8e5a41082..2111fd254 100644 --- a/fastvideo/configs/models/vaes/__init__.py +++ b/fastvideo/configs/models/vaes/__init__.py @@ -5,6 +5,7 @@ from fastvideo.configs.models.vaes.gen3cvae import Gen3CVAEConfig from fastvideo.configs.models.vaes.hunyuanvae import HunyuanVAEConfig from fastvideo.configs.models.vaes.hunyuan15vae import Hunyuan15VAEConfig from fastvideo.configs.models.vaes.ltx2vae import LTX2VAEConfig +from fastvideo.configs.models.vaes.oobleck import OobleckVAEArchConfig, OobleckVAEConfig from fastvideo.configs.models.vaes.wanvae import WanVAEConfig __all__ = [ @@ -16,4 +17,6 @@ __all__ = [ "Gen3CVAEConfig", "Hunyuan15VAEConfig", "LTX2VAEConfig", + "OobleckVAEArchConfig", + "OobleckVAEConfig", ] diff --git a/fastvideo/configs/models/vaes/oobleck.py b/fastvideo/configs/models/vaes/oobleck.py new file mode 100644 index 000000000..0db6b7efb --- /dev/null +++ b/fastvideo/configs/models/vaes/oobleck.py @@ -0,0 +1,68 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Config for the Stable Audio Open 1.0 "Oobleck" VAE. + +Mirrors the per-channel `vae/config.json` shipped in +`stabilityai/stable-audio-open-1.0` 1:1 (see +`fastvideo/models/vaes/oobleck.py::OobleckVAE.from_pretrained`, which +constructs the VAE from these fields). Inherits the FastVideo VAEConfig +base so the standard `load_encoder` / `load_decoder` flags + tiling +knobs apply. + +Naming: the VAE architecture is officially "Oobleck" (per Stability +AI's stable-audio-tools) — the surrounding model family is "Stable +Audio Open 1.0". This config is named after the architecture +(`OobleckVAEConfig`) since the same VAE is shared across Stable Audio +checkpoints; downstream pipelines reference it by its arch name, not +by a host-pipeline name. +""" +from __future__ import annotations + +from dataclasses import dataclass, field + +from fastvideo.configs.models.vaes.base import VAEArchConfig, VAEConfig + + +@dataclass +class OobleckVAEArchConfig(VAEArchConfig): + """Stable Audio Open 1.0 VAE architecture constants.""" + + architectures: list[str] = field(default_factory=lambda: ["AutoencoderOobleck"]) + + # From stabilityai/stable-audio-open-1.0/vae/config.json. + encoder_hidden_size: int = 128 + downsampling_ratios: list[int] = field(default_factory=lambda: [2, 4, 4, 8, 8]) + channel_multiples: list[int] = field(default_factory=lambda: [1, 2, 4, 8, 16]) + decoder_channels: int = 128 + decoder_input_channels: int = 64 + audio_channels: int = 2 # stereo + sampling_rate: int = 44100 + + +@dataclass +class OobleckVAEConfig(VAEConfig): + """FastVideo VAE config wrapping the Oobleck arch. + + Audio VAEs don't use the temporal/spatial tiling defaults that the + base VAEConfig is shaped for (those exist for video VAEs); they are + retained but irrelevant for audio. + """ + + arch_config: VAEArchConfig = field(default_factory=OobleckVAEArchConfig) + + # Audio is 1-D, so the video-VAE tiling defaults are inert. Disable + # them so callers don't accidentally trip on tile-stride math built + # for spatial tensors. + use_tiling: bool = False + use_temporal_tiling: bool = False + use_parallel_tiling: bool = False + + # Where the FastVideo loader / pipeline-glue wrapper should fetch + # weights from when no local path is supplied. Gated repo — caller's + # HF token must have accepted terms on + # https://huggingface.co/stabilityai/stable-audio-open-1.0. + pretrained_path: str = "stabilityai/stable-audio-open-1.0" + pretrained_subfolder: str = "vae" + # Match official `stable_audio_tools`: VAE runs in fp16 (the + # `pretransform.model_half` path in + # `stable_audio_tools/models/pretransforms.py`). + pretrained_dtype: str = "float16" diff --git a/fastvideo/configs/pipelines/stable_audio.py b/fastvideo/configs/pipelines/stable_audio.py new file mode 100644 index 000000000..2751833ca --- /dev/null +++ b/fastvideo/configs/pipelines/stable_audio.py @@ -0,0 +1,67 @@ +# SPDX-License-Identifier: Apache-2.0 +"""`PipelineConfig` for Stable Audio Open 1.0.""" +from __future__ import annotations + +from dataclasses import dataclass, field + +from fastvideo.configs.models import DiTConfig, VAEConfig +from fastvideo.configs.models.dits import StableAudioConfig +from fastvideo.configs.models.vaes import OobleckVAEConfig +from fastvideo.configs.pipelines.base import PipelineConfig + + +@dataclass +class StableAudioT2AConfig(PipelineConfig): + """Stable Audio Open 1.0 pipeline config.""" + + dit_config: DiTConfig = field(default_factory=StableAudioConfig) + # Standard `TransformerLoader` reads `dit_precision`; default in + # `PipelineConfig` is bf16, but we want fp16 to match official. + dit_precision: str = "fp16" + + vae_config: VAEConfig = field(default_factory=OobleckVAEConfig) + vae_tiling: bool = False + vae_sp: bool = False + + # `StableAudioMultiConditioner` owns its own T5; zero out the + # parent's text-encoder slots so the length-equality validator passes. + text_encoder_configs: tuple = field(default_factory=tuple) + preprocess_text_funcs: tuple = field(default_factory=tuple) + postprocess_text_funcs: tuple = field(default_factory=tuple) + + num_inference_steps: int = 100 + guidance_scale: float = 7.0 + audio_end_in_s: float = 10.0 # short-clip default + audio_start_in_s: float = 0.0 + sampling_rate: int = 44100 + audio_channels: int = 2 + # Stable Audio Open 1.0 was trained at a fixed 2,097,152-sample + # window (= 2097152 / 44100 ≈ 47.55s). Anything past this is + # silently truncated by the post-decode slice — validate up-front. + sample_size: int = 2097152 + max_audio_duration_s: float = 2097152 / 44100 + + # Match the official `stable_audio_tools` defaults (`model_half=True` + # in `run_gradio.py`), which loads the DiT, VAE, and T5 in fp16 and + # wraps T5 forward in `autocast(fp16)`. fp16 is also a hard + # requirement for FlashAttention-2 / FA-3. + precision: str = "fp16" + vae_precision: str = "fp16" + text_encoder_precisions: tuple[str, ...] = field(default_factory=tuple) + + def __post_init__(self) -> None: + # A2A needs encode; load both halves for either path. + self.vae_config.load_encoder = True + self.vae_config.load_decoder = True + + +@dataclass +class StableAudioOpenSmallConfig(StableAudioT2AConfig): + """`stable-audio-open-small` overrides: shorter training window + (524288 samples ≈ 11.89s @ 44.1 kHz) and a faster default sampler + config carried by the small preset. + """ + + sample_size: int = 524288 + max_audio_duration_s: float = 524288 / 44100 + audio_end_in_s: float = 6.0 # short-clip default suitable for the small window diff --git a/fastvideo/entrypoints/video_generator.py b/fastvideo/entrypoints/video_generator.py index 91ba1837a..fafc5efe1 100644 --- a/fastvideo/entrypoints/video_generator.py +++ b/fastvideo/entrypoints/video_generator.py @@ -627,18 +627,32 @@ class VideoGenerator: gen_time = time.perf_counter() - start_time logger.info("Generated successfully in %.2f seconds", gen_time) - # Process outputs - videos = rearrange(samples, "b c t h w -> t b c h w") - frames = [] - for x in videos: - x = torchvision.utils.make_grid(x, nrow=6) - x = x.permute(1, 2, 0).squeeze(-1) - x = (x * 255).to(torch.uint8) - frames.append(x.cpu().numpy()) + # Process outputs (skip the make_grid loop for audio-only, where + # `samples` is a 1×3×1×8×8 placeholder no caller will use). + audio_only = bool(output_batch.extra.get("audio_only")) + frames: list[np.ndarray] = [] + if not audio_only: + videos = rearrange(samples, "b c t h w -> t b c h w") + for x in videos: + x = torchvision.utils.make_grid(x, nrow=6) + x = x.permute(1, 2, 0).squeeze(-1) + x = (x * 255).to(torch.uint8) + frames.append(x.cpu().numpy()) # Save output if requested if batch.save_video: - if self._is_image_workload(): + if output_batch.extra.get("audio_only"): + # Audio-only workload: write a standalone .wav rather than + # muxing the audio into a placeholder mp4 (which forces + # ffmpeg to round 8x8 placeholder frames up to 16x16). + output_path = self._rewrite_extension(output_path, ".wav") + self._write_pcm_wav( + output_path, + output_batch.extra["audio"], + int(output_batch.extra["audio_sample_rate"]), + ) + logger.info("Saved audio to %s", output_path) + elif self._is_image_workload(): # Image workloads (t2i, i2i, …): save the first frame as PNG. imageio.imwrite(output_path, frames[0]) logger.info("Saved image to %s", output_path) @@ -655,7 +669,11 @@ class VideoGenerator: "prompts": prompt, "samples": samples if batch.return_frames else None, "frames": frames if batch.return_frames else None, - "audio": output_batch.extra.get("audio") if batch.return_frames else None, + # Audio is the primary output for audio workloads — return it + # whenever the pipeline produced one, regardless of + # `return_frames` (which gates the video-shaped buffers). + "audio": output_batch.extra.get("audio"), + "audio_sample_rate": output_batch.extra.get("audio_sample_rate"), "size": (target_height, target_width, batch.num_frames), "generation_time": gen_time, "logging_info": logging_info, @@ -683,7 +701,55 @@ class VideoGenerator: return result.to_legacy_dict() @staticmethod + def _rewrite_extension(path: str, new_ext: str) -> str: + root, old_ext = os.path.splitext(path) + new_path = root + new_ext + if old_ext and old_ext.lower() != new_ext.lower(): + logger.info("Rewriting output extension %s -> %s.", old_ext, new_ext) + return new_path + + @staticmethod + def _audio_to_int16(audio: torch.Tensor | np.ndarray, ) -> tuple[np.ndarray, int]: + """Normalize `[samples]` / `[samples, channels]` / `[channels, + samples]` audio in roughly [-1, 1] to a `(int16 [samples, + channels], num_channels)` pair. Raises `ValueError` for shapes + we can't classify. + """ + if torch.is_tensor(audio): + audio_np = audio.detach().cpu().float().numpy() + else: + audio_np = np.asarray(audio, dtype=np.float32) + if audio_np.ndim == 1: + audio_np = audio_np[:, None] + elif audio_np.ndim == 2: + if audio_np.shape[0] <= 8 and audio_np.shape[1] > audio_np.shape[0]: + audio_np = audio_np.T + else: + raise ValueError(f"Unexpected audio shape {audio_np.shape}.") + audio_np = np.clip(audio_np, -1.0, 1.0) + audio_int16 = (audio_np * 32767.0).astype(np.int16) + return audio_int16, audio_int16.shape[1] + + @classmethod + def _write_pcm_wav( + cls, + wav_path: str, + audio: torch.Tensor | np.ndarray, + sample_rate: int, + ) -> int: + """Write 16-bit PCM WAV; returns the channel count.""" + import wave + audio_int16, num_channels = cls._audio_to_int16(audio) + with wave.open(wav_path, "wb") as f: + f.setnchannels(num_channels) + f.setsampwidth(2) + f.setframerate(sample_rate) + f.writeframes(audio_int16.tobytes()) + return num_channels + + @classmethod def _mux_audio( + cls, video_path: str, audio: torch.Tensor | np.ndarray, sample_rate: int, @@ -696,37 +762,13 @@ class VideoGenerator: "Install with: pip install av") return False - if torch.is_tensor(audio): - audio_np = audio.detach().cpu().float().numpy() - else: - audio_np = np.asarray(audio, dtype=np.float32) - - if audio_np.ndim == 1: - audio_np = audio_np[:, None] - elif audio_np.ndim == 2: - if audio_np.shape[0] <= 8 and audio_np.shape[1] > audio_np.shape[0]: - audio_np = audio_np.T - else: - logger.warning("Unexpected audio shape %s; skipping mux.", audio_np.shape) - return False - - audio_np = np.clip(audio_np, -1.0, 1.0) - audio_int16 = (audio_np * 32767.0).astype(np.int16) - num_channels = audio_int16.shape[1] - layout = "stereo" if num_channels == 2 else "mono" - try: - import wave with tempfile.TemporaryDirectory() as tmpdir: out_path = os.path.join(tmpdir, "muxed.mp4") wav_path = os.path.join(tmpdir, "audio.wav") - # Write audio to WAV file - with wave.open(wav_path, "wb") as wav_file: - wav_file.setnchannels(num_channels) - wav_file.setsampwidth(2) - wav_file.setframerate(sample_rate) - wav_file.writeframes(audio_int16.tobytes()) + num_channels = cls._write_pcm_wav(wav_path, audio, sample_rate) + layout = "stereo" if num_channels == 2 else "mono" # Open input video and audio input_video = av.open(video_path) diff --git a/fastvideo/models/dits/stable_audio.py b/fastvideo/models/dits/stable_audio.py new file mode 100644 index 000000000..693f2cad5 --- /dev/null +++ b/fastvideo/models/dits/stable_audio.py @@ -0,0 +1,389 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Stable Audio Open 1.0 DiT. + +Continuous transformer with rotary self-attention, GQA cross-attention, +and prepend global conditioning. 24 layers, embed_dim=1536, head_dim=64. +""" +from __future__ import annotations + +import math +from typing import Any + +import torch +from einops import rearrange +from torch import nn + +from fastvideo.attention import LocalAttention +from fastvideo.configs.models.dits import StableAudioConfig +from fastvideo.layers.layernorm import FP32LayerNorm +from fastvideo.layers.linear import ReplicatedLinear +from fastvideo.models.dits.base import BaseDiT +from fastvideo.models.loader.utils import get_param_names_mapping + +# Single import-time snapshot — re-reading via `StableAudioConfig()` per +# `Attention.__init__` would rebuild the nested dataclass + regex map ~48 +# times during a single DiT construction. Reused for the class-level +# attribute defaults below. +_DEFAULT_CONFIG = StableAudioConfig() +_SUPPORTED_BACKENDS = _DEFAULT_CONFIG.arch_config._supported_attention_backends + + +class FourierFeatures(nn.Module): + """Random-Fourier learned-frequency timestep encoder.""" + + def __init__(self, in_features: int, out_features: int, std: float = 1.0) -> None: + super().__init__() + assert out_features % 2 == 0 + self.weight = nn.Parameter(torch.randn([out_features // 2, in_features]) * std) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + f = 2 * math.pi * x @ self.weight.T + return torch.cat([f.cos(), f.sin()], dim=-1) + + +# Partial-rotary with halves-swap (`unbind(-2)`, `[-x2, x1]`). Different +# from FastVideo's `_apply_rotary_emb` (interleaved pairs, `unbind(-1)`), +# so kept local. + + +class RotaryEmbedding(nn.Module): + + def __init__(self, dim: int, base: float = 10000.0) -> None: + super().__init__() + inv_freq = 1.0 / (base**(torch.arange(0, dim, 2).float() / dim)) + self.register_buffer("inv_freq", inv_freq) + self.register_buffer("scale", None) + + def forward_from_seq_len(self, seq_len: int): + t = torch.arange(seq_len, device=self.inv_freq.device, dtype=torch.float32) + freqs = torch.einsum("i , j -> i j", t, self.inv_freq) + freqs = torch.cat((freqs, freqs), dim=-1) + return freqs, 1.0 + + +def _rotate_half(x: torch.Tensor) -> torch.Tensor: + x = rearrange(x, "... (j d) -> ... j d", j=2) + x1, x2 = x.unbind(dim=-2) + return torch.cat((-x2, x1), dim=-1) + + +def _apply_rotary_pos_emb(t: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor: + out_dtype = t.dtype + rot_dim, seq_len = freqs.shape[-1], t.shape[-2] + freqs = freqs.to(torch.float32)[-seq_len:, :] + t = t.to(torch.float32) + if t.ndim == 4 and freqs.ndim == 3: + freqs = rearrange(freqs, "b n d -> b 1 n d") + t_rot, t_unrot = t[..., :rot_dim], t[..., rot_dim:] + t_rot = (t_rot * freqs.cos()) + (_rotate_half(t_rot) * freqs.sin()) + return torch.cat((t_rot.to(out_dtype), t_unrot.to(out_dtype)), dim=-1) + + +# SwiGLU FF — local because `fastvideo.layers.mlp.MLP` is non-gated. + + +class _GLU(nn.Module): + + def __init__(self, dim_in: int, dim_out: int, activation: nn.Module) -> None: + super().__init__() + self.act = activation + self.proj = ReplicatedLinear(dim_in, dim_out * 2, bias=True) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x, _ = self.proj(x) + x, gate = x.chunk(2, dim=-1) + return x * self.act(gate) + + +class FeedForward(nn.Module): + # Sequential layout `(GLU, Identity, Linear, Identity)` keeps the + # checkpoint keys at indices 0 and 2. + + def __init__(self, dim: int, mult: int = 4, zero_init_output: bool = True) -> None: + super().__init__() + inner_dim = int(dim * mult) + linear_in = _GLU(dim, inner_dim, nn.SiLU()) + linear_out = ReplicatedLinear(inner_dim, dim, bias=True) + if zero_init_output: + nn.init.zeros_(linear_out.weight) + nn.init.zeros_(linear_out.bias) + self.ff = nn.Sequential(linear_in, nn.Identity(), linear_out, nn.Identity()) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + for mod in self.ff: + if isinstance(mod, ReplicatedLinear): + x, _ = mod(x) + else: + x = mod(x) + return x + + +# Cross-attention is GQA (24 query heads, 12 KV heads); both backends +# (FlashAttn, SDPA with `enable_gqa=True`) handle it. + + +class Attention(nn.Module): + + def __init__(self, dim: int, dim_heads: int = 64, dim_context: int | None = None, + zero_init_output: bool = True, qk_norm: str | None = None) -> None: + super().__init__() + self.dim = dim + self.dim_heads = dim_heads + dim_kv = dim_context if dim_context is not None else dim + self.num_heads = dim // dim_heads + self.kv_heads = dim_kv // dim_heads + if dim_context is not None: + self.to_q = ReplicatedLinear(dim, dim, bias=False) + self.to_kv = ReplicatedLinear(dim_kv, dim_kv * 2, bias=False) + else: + self.to_qkv = ReplicatedLinear(dim, dim * 3, bias=False) + self.to_out = ReplicatedLinear(dim, dim, bias=False) + if zero_init_output: + nn.init.zeros_(self.to_out.weight) + + # `stable-audio-open-small` wraps Q/K in LayerNorm before attn + # (`attn_kwargs.qk_norm = "ln"` in its `model_config.json`); the + # 1.0 base does not. Names match upstream (`q_norm`/`k_norm`) + # so the converted state dict loads strict. + if qk_norm == "ln": + self.q_norm = nn.LayerNorm(dim_heads) + self.k_norm = nn.LayerNorm(dim_heads) + elif qk_norm is None: + self.q_norm = nn.Identity() + self.k_norm = nn.Identity() + else: + raise ValueError(f"Unsupported qk_norm={qk_norm!r}; expected 'ln' or None.") + + self.attn = LocalAttention(num_heads=self.num_heads, head_size=dim_heads, + num_kv_heads=self.kv_heads, causal=False, + supported_attention_backends=_SUPPORTED_BACKENDS) + + def forward(self, x: torch.Tensor, context: torch.Tensor | None = None, + rotary_pos_emb: tuple[torch.Tensor, float] | None = None) -> torch.Tensor: + h, kv_h, has_context = self.num_heads, self.kv_heads, context is not None + kv_input = context if has_context else x + if has_context: + q, _ = self.to_q(x) + kv, _ = self.to_kv(kv_input) + k, v = kv.chunk(2, dim=-1) + else: + qkv, _ = self.to_qkv(x) + q, k, v = qkv.chunk(3, dim=-1) + # LocalAttention expects [batch, seq_len, num_heads, head_dim]. + q = rearrange(q, "b n (h d) -> b n h d", h=h) + k = rearrange(k, "b n (h d) -> b n h d", h=kv_h) + v = rearrange(v, "b n (h d) -> b n h d", h=kv_h) + q = self.q_norm(q) + k = self.k_norm(k) + + if rotary_pos_emb is not None: + freqs, _ = rotary_pos_emb + v_dtype = v.dtype + # Partial rotary (rot_dim < head_dim) with halves-swap, so + # apply outside LocalAttention. q,k come in as [B, S, H, D]; + # transpose to [B, H, S, D] for the helper. + q_t = q.transpose(1, 2) + k_t = k.transpose(1, 2) + if q_t.shape[-2] >= k_t.shape[-2]: + ratio = q_t.shape[-2] / k_t.shape[-2] + q_freqs, k_freqs = freqs, ratio * freqs + else: + ratio = k_t.shape[-2] / q_t.shape[-2] + q_freqs, k_freqs = ratio * freqs, freqs + q = _apply_rotary_pos_emb(q_t, q_freqs).to(v_dtype).transpose(1, 2) + k = _apply_rotary_pos_emb(k_t, k_freqs).to(v_dtype).transpose(1, 2) + + out = self.attn(q, k, v) + out = rearrange(out, "b n h d -> b n (h d)") + out, _ = self.to_out(out) + return out + + +class TransformerBlock(nn.Module): + + def __init__(self, dim: int, dim_heads: int = 64, cross_attend: bool = False, + dim_context: int | None = None, zero_init_branch_outputs: bool = True, + qk_norm: str | None = None) -> None: + super().__init__() + self.dim = dim + self.dim_heads = min(dim_heads, dim) + self.cross_attend = cross_attend + self.pre_norm = FP32LayerNorm(dim, elementwise_affine=True) + self.self_attn = Attention(dim, dim_heads=self.dim_heads, + zero_init_output=zero_init_branch_outputs, + qk_norm=qk_norm) + if cross_attend: + self.cross_attend_norm = FP32LayerNorm(dim, elementwise_affine=True) + self.cross_attn = Attention(dim, dim_heads=self.dim_heads, dim_context=dim_context, + zero_init_output=zero_init_branch_outputs, + qk_norm=qk_norm) + self.ff_norm = FP32LayerNorm(dim, elementwise_affine=True) + self.ff = FeedForward(dim, zero_init_output=zero_init_branch_outputs) + + def forward(self, x: torch.Tensor, context: torch.Tensor | None = None, + rotary_pos_emb: tuple[torch.Tensor, float] | None = None) -> torch.Tensor: + x = x + self.self_attn(self.pre_norm(x), rotary_pos_emb=rotary_pos_emb) + if context is not None and self.cross_attend: + x = x + self.cross_attn(self.cross_attend_norm(x), context=context) + x = x + self.ff(self.ff_norm(x)) + return x + + +class ContinuousTransformer(nn.Module): + + def __init__(self, dim: int, depth: int, *, dim_heads: int = 64, dim_in: int | None = None, + dim_out: int | None = None, cross_attend: bool = False, + cond_token_dim: int | None = None, zero_init_branch_outputs: bool = True, + qk_norm: str | None = None) -> None: + super().__init__() + self.dim = dim + self.depth = depth + self.project_in = (ReplicatedLinear(dim_in, dim, bias=False) if dim_in is not None + else nn.Identity()) + self.project_out = (ReplicatedLinear(dim, dim_out, bias=False) if dim_out is not None + else nn.Identity()) + self.rotary_pos_emb = RotaryEmbedding(max(dim_heads // 2, 32)) + self.layers = nn.ModuleList([ + TransformerBlock(dim, dim_heads=dim_heads, cross_attend=cross_attend, + dim_context=cond_token_dim, + zero_init_branch_outputs=zero_init_branch_outputs, + qk_norm=qk_norm) for _ in range(depth) + ]) + + def forward(self, x: torch.Tensor, prepend_embeds: torch.Tensor | None = None, + context: torch.Tensor | None = None) -> torch.Tensor: + if isinstance(self.project_in, ReplicatedLinear): + x, _ = self.project_in(x) + if prepend_embeds is not None: + assert prepend_embeds.shape[-1] == x.shape[-1] + x = torch.cat((prepend_embeds, x), dim=-2) + rotary = self.rotary_pos_emb.forward_from_seq_len(x.shape[1]) + for layer in self.layers: + x = layer(x, context=context, rotary_pos_emb=rotary) + if isinstance(self.project_out, ReplicatedLinear): + x, _ = self.project_out(x) + return x + + +class StableAudioDiT(BaseDiT): + """Stable Audio Open 1.0 diffusion transformer.""" + + _fsdp_shard_conditions = _DEFAULT_CONFIG.arch_config._fsdp_shard_conditions + _compile_conditions = _DEFAULT_CONFIG.arch_config._compile_conditions + param_names_mapping = _DEFAULT_CONFIG.arch_config.param_names_mapping + reverse_param_names_mapping: dict = {} + + def __init__(self, config: StableAudioConfig | None = None, + hf_config: dict[str, Any] | None = None) -> None: + if config is None: + config = StableAudioConfig() + super().__init__(config=config, hf_config=hf_config or {}) + arch = config.arch_config + self.hidden_size = arch.hidden_size + self.num_attention_heads = arch.num_attention_heads + self.num_channels_latents = arch.num_channels_latents + io_channels = arch.io_channels + embed_dim = arch.embed_dim + depth = arch.depth + num_heads = arch.num_attention_heads + cond_token_dim = arch.cond_token_dim + global_cond_dim = arch.global_cond_dim + project_cond_tokens = arch.project_cond_tokens + project_global_cond = arch.project_global_cond + qk_norm = arch.qk_norm + + self.cond_token_dim = cond_token_dim + timestep_features_dim = 256 + self.timestep_features = FourierFeatures(1, timestep_features_dim) + self.to_timestep_embed = nn.Sequential( + ReplicatedLinear(timestep_features_dim, embed_dim, bias=True), + nn.SiLU(), + ReplicatedLinear(embed_dim, embed_dim, bias=True), + ) + self.diffusion_objective = "v" + + cond_embed_dim = cond_token_dim if not project_cond_tokens else embed_dim + self.to_cond_embed = nn.Sequential( + ReplicatedLinear(cond_token_dim, cond_embed_dim, bias=False), + nn.SiLU(), + ReplicatedLinear(cond_embed_dim, cond_embed_dim, bias=False), + ) + + global_embed_dim = global_cond_dim if not project_global_cond else embed_dim + self.to_global_embed = nn.Sequential( + ReplicatedLinear(global_cond_dim, global_embed_dim, bias=False), + nn.SiLU(), + ReplicatedLinear(global_embed_dim, global_embed_dim, bias=False), + ) + + self.transformer = ContinuousTransformer( + dim=embed_dim, depth=depth, dim_heads=embed_dim // num_heads, dim_in=io_channels, + dim_out=io_channels, cross_attend=True, cond_token_dim=cond_embed_dim, + qk_norm=qk_norm, + ) + + self.preprocess_conv = nn.Conv1d(io_channels, io_channels, 1, bias=False) + nn.init.zeros_(self.preprocess_conv.weight) + self.postprocess_conv = nn.Conv1d(io_channels, io_channels, 1, bias=False) + nn.init.zeros_(self.postprocess_conv.weight) + + self.io_channels = io_channels + self.embed_dim = embed_dim + self.depth = depth + self.num_heads = num_heads + self.__post_init__() + + @staticmethod + def _seq_apply(seq: nn.Sequential, x: torch.Tensor) -> torch.Tensor: + for mod in seq: + if isinstance(mod, ReplicatedLinear): + x, _ = mod(x) + else: + x = mod(x) + return x + + def forward(self, x: torch.Tensor, t: torch.Tensor, *, cross_attn_cond: torch.Tensor, + global_embed: torch.Tensor) -> torch.Tensor: + """Forward over a single batch. CFG batching is the caller's job.""" + model_dtype = next(self.parameters()).dtype + x = x.to(model_dtype) + t = t.to(model_dtype) + cross_attn_cond = cross_attn_cond.to(model_dtype) + global_embed = global_embed.to(model_dtype) + + cross_attn_cond = self._seq_apply(self.to_cond_embed, cross_attn_cond) + global_embed = self._seq_apply(self.to_global_embed, global_embed) + timestep_embed = self._seq_apply(self.to_timestep_embed, self.timestep_features(t[:, None])) + global_embed = global_embed + timestep_embed + prepend_inputs = global_embed.unsqueeze(1) + + x = self.preprocess_conv(x) + x + x = rearrange(x, "b c t -> b t c") + out = self.transformer(x, prepend_embeds=prepend_inputs, context=cross_attn_cond) + out = rearrange(out, "b t c -> b c t")[:, :, prepend_inputs.shape[1]:] + return self.postprocess_conv(out) + out + + @classmethod + def from_official_state_dict(cls, state_dict: dict[str, torch.Tensor], + prefix: str = "model.model.") -> "StableAudioDiT": + """Load from a raw `stable_audio_tools` monolithic state dict. + Kept for tests / older checkpoints; production loads go through + the standard `TransformerLoader` against the converted Diffusers + repo. + """ + model = cls() + mapping_fn = get_param_names_mapping(model.config.arch_config.param_names_mapping) + remapped: dict[str, torch.Tensor] = {} + for k, v in state_dict.items(): + if not k.startswith(prefix): + continue + new_key, _, _ = mapping_fn(k) + remapped[new_key] = v + missing, unexpected = model.load_state_dict(remapped, strict=True) + if missing or unexpected: + raise RuntimeError( + f"StableAudioDiT load mismatch — missing={missing[:5]} unexpected={unexpected[:5]}") + return model + + +EntryClass = StableAudioDiT diff --git a/fastvideo/models/encoders/stable_audio_conditioner.py b/fastvideo/models/encoders/stable_audio_conditioner.py new file mode 100644 index 000000000..c4aba0105 --- /dev/null +++ b/fastvideo/models/encoders/stable_audio_conditioner.py @@ -0,0 +1,214 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Stable Audio Open 1.0 conditioner. + +T5-base text encoder + two NumberConditioners (`seconds_start`, +`seconds_total`), wrapped by `StableAudioMultiConditioner` which +produces the cross-attention and global-conditioning tensors the DiT +expects. +""" +from __future__ import annotations + +import math + +import torch +import torch.nn as nn +from einops import rearrange + +from fastvideo.configs.models.encoders import StableAudioConditionerConfig + + +class _LearnedPositionalEmbedding(nn.Module): + + def __init__(self, dim: int) -> None: + super().__init__() + assert (dim % 2) == 0 + self.weights = nn.Parameter(torch.randn(dim // 2)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = rearrange(x, "b -> b 1") + freqs = x * rearrange(self.weights, "d -> 1 d") * 2 * math.pi + fouriered = torch.cat((freqs.sin(), freqs.cos()), dim=-1) + return torch.cat((x, fouriered), dim=-1) + + +def _time_positional_embedding(dim: int, out_features: int) -> nn.Sequential: + return nn.Sequential(_LearnedPositionalEmbedding(dim), + nn.Linear(in_features=dim + 1, out_features=out_features)) + + +class NumberEmbedder(nn.Module): + + def __init__(self, features: int, dim: int = 256) -> None: + super().__init__() + self.features = features + self.embedding = _time_positional_embedding(dim=dim, out_features=features) + + def forward(self, x: torch.Tensor | list[float]) -> torch.Tensor: + if not torch.is_tensor(x): + device = next(self.embedding.parameters()).device + x = torch.tensor(x, device=device) + shape = x.shape + x = rearrange(x, "... -> (...)") + out = self.embedding(x) + return out.view(*shape, self.features) + + +class _Conditioner(nn.Module): + + def __init__(self, dim: int, output_dim: int, project_out: bool = False) -> None: + super().__init__() + self.dim = dim + self.output_dim = output_dim + self.proj_out = (nn.Linear(dim, output_dim) if dim != output_dim or project_out + else nn.Identity()) + + +class T5Conditioner(_Conditioner): + """T5 text conditioner. Pads to `model_max_length` (=128 for the SA + repo's tokenizer, NOT the standard 512) and emits a masked + last-hidden-state. + """ + + T5_MODEL_DIMS = {"t5-base": 768} + + def __init__(self, output_dim: int, t5_model_name: str = "t5-base", + max_length: int = 128, dtype: str = "float16") -> None: + super().__init__(self.T5_MODEL_DIMS[t5_model_name], output_dim, project_out=False) + from transformers import AutoTokenizer, T5EncoderModel + self.max_length = max_length + self.tokenizer = AutoTokenizer.from_pretrained(t5_model_name) + # T5 loaded directly in fp16 (config-driven) to match official + # `stable_audio_tools/models/conditioners.py:334`. Registered as + # a normal submodule so `.to(device)` / `torch.compile` track it; + # `from_official_state_dict` filters `conditioners.prompt.*` from + # the missing-key check (T5 weights are absent from the SA + # checkpoint by design). + # Explicit lookup so a typo (e.g. "fp16" instead of "float16") errors + # at load time rather than silently falling back to a wrong dtype. + torch_dtype = getattr(torch, dtype) + if not isinstance(torch_dtype, torch.dtype): + raise ValueError(f"T5Conditioner dtype={dtype!r} is not a torch.dtype.") + self._t5_dtype = torch_dtype + self.model = (T5EncoderModel.from_pretrained(t5_model_name).eval().requires_grad_(False).to(torch_dtype)) + + def forward(self, texts: list[str], device: torch.device | str) -> tuple[torch.Tensor, torch.Tensor]: + encoded = self.tokenizer(texts, truncation=True, max_length=self.max_length, + padding="max_length", return_tensors="pt") + input_ids = encoded["input_ids"].to(device) + attention_mask = encoded["attention_mask"].to(device).to(torch.bool) + # Mirror official's `autocast(fp16)` wrap on T5 forward. + with torch.no_grad(), torch.autocast(device_type="cuda", dtype=self._t5_dtype): + embeddings = self.model(input_ids=input_ids, + attention_mask=attention_mask)["last_hidden_state"] + embeddings = self.proj_out(embeddings) * attention_mask.unsqueeze(-1).float() + return embeddings, attention_mask + + +class NumberConditioner(_Conditioner): + """Float-valued conditioner with min/max clamping + NumberEmbedder.""" + + def __init__(self, output_dim: int, min_val: float = 0, max_val: float = 1) -> None: + super().__init__(output_dim, output_dim) + self.min_val = min_val + self.max_val = max_val + self.embedder = NumberEmbedder(features=output_dim) + + def forward(self, floats: list[float], device: torch.device | str) -> tuple[torch.Tensor, torch.Tensor]: + floats = [float(x) for x in floats] + floats_t = torch.tensor(floats, device=device).clamp(self.min_val, self.max_val) + normalized = (floats_t - self.min_val) / (self.max_val - self.min_val) + emb_dtype = next(self.embedder.parameters()).dtype + normalized = normalized.to(emb_dtype) + float_embeds = self.embedder(normalized).unsqueeze(1) + return float_embeds, torch.ones(float_embeds.shape[0], 1, device=device) + + +class StableAudioMultiConditioner(nn.Module): + """SA-Open-1.0 conditioner: T5 prompt + duration NumberConditioners. + + All hardcoded constants (cond_dim, sub-conditioner ids, T5 model + name + max_length, NumberConditioner ranges) live on + `StableAudioConditionerConfig` — see + `fastvideo/configs/models/encoders/stable_audio_conditioner.py`. + """ + + def __init__(self, config: StableAudioConditionerConfig | None = None) -> None: + super().__init__() + self.config = config or StableAudioConditionerConfig() + arch = self.config.arch_config + # Build sub-conditioners from the `configs` list (mirrors + # upstream's `MultiConditioner` factory). + sub: dict[str, nn.Module] = {} + for spec in arch.configs: + sid = spec["id"] + stype = spec["type"] + scfg = spec["config"] + if stype == "t5": + sub[sid] = T5Conditioner(output_dim=arch.cond_dim, + t5_model_name=scfg["t5_model_name"], + max_length=scfg["max_length"], + dtype=arch.t5_dtype) + elif stype == "number": + sub[sid] = NumberConditioner(output_dim=arch.cond_dim, + min_val=scfg["min_val"], max_val=scfg["max_val"]) + else: + raise ValueError(f"Unknown sub-conditioner type {stype!r} for id {sid!r}.") + self.conditioners = nn.ModuleDict(sub) + self.cross_attention_cond_ids = tuple(arch.cross_attention_cond_ids) + self.global_cond_ids = tuple(arch.global_cond_ids) + + def forward(self, batch_metadata: list[dict], + device: torch.device | str) -> dict[str, tuple[torch.Tensor, torch.Tensor]]: + out: dict[str, tuple[torch.Tensor, torch.Tensor]] = {} + for key, conditioner in self.conditioners.items(): + inputs = [x[key] for x in batch_metadata] + out[key] = conditioner(inputs, device) + return out + + def get_conditioning_inputs( + self, cond: dict[str, tuple[torch.Tensor, torch.Tensor]] + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Pack conditioner outputs into the (cross_attn_cond, + cross_attn_mask, global_embed) triple the DiT consumes. Order + is driven by `cross_attention_cond_ids` / `global_cond_ids` + from the config — SA-1.0 uses three sub-conditioners + (prompt + seconds_start + seconds_total); SA-small uses two + (prompt + seconds_total). + """ + x_embs = [cond[i][0] for i in self.cross_attention_cond_ids] + x_masks = [cond[i][1] for i in self.cross_attention_cond_ids] + cross_attn_cond = torch.cat(x_embs, dim=1) + cross_attn_mask = torch.cat(x_masks, dim=1) + global_embed = torch.cat([cond[i][0][:, 0] for i in self.global_cond_ids], dim=-1) + return cross_attn_cond, cross_attn_mask, global_embed + + @classmethod + def from_official_state_dict(cls, state_dict: dict[str, torch.Tensor], + prefix: str = "conditioner.") -> "StableAudioMultiConditioner": + """Load NumberConditioner weights from a raw `stable_audio_tools` + monolithic state dict. Kept for tests / older checkpoints; + production loads go through the standard `ConditionerLoader` + against the converted Diffusers repo. + """ + mc = cls() + own_state = mc.state_dict() + loaded: dict[str, torch.Tensor] = {} + for k, v in state_dict.items(): + if not k.startswith(prefix): + continue + stripped = k[len(prefix):] + if stripped in own_state: + loaded[stripped] = v + # T5 keys are intentionally absent from the checkpoint. + missing = [k for k in own_state.keys() if k not in loaded + and not k.startswith("conditioners.prompt.")] + unexpected = [k for k in loaded.keys() if k not in own_state] + if missing or unexpected: + raise RuntimeError( + f"StableAudioMultiConditioner load mismatch — missing={missing[:5]} unexpected={unexpected[:5]}" + ) + mc.load_state_dict(loaded, strict=False) + return mc + + +EntryClass = StableAudioMultiConditioner diff --git a/fastvideo/models/loader/component_loader.py b/fastvideo/models/loader/component_loader.py index 9dfde7cb2..5a2145fe8 100644 --- a/fastvideo/models/loader/component_loader.py +++ b/fastvideo/models/loader/component_loader.py @@ -95,6 +95,10 @@ class ComponentLoader(ABC): "image_encoder": (ImageEncoderLoader, "transformers"), "upsampler": (UpsamplerLoader, "diffusers"), "upsampler_2": (UpsamplerLoader, "diffusers"), + # Stable Audio's `StableAudioMultiConditioner` bundles T5 + + # NumberConditioners; not a pure text encoder, so it gets + # its own loader. + "conditioner": (ConditionerLoader, "fastvideo"), } if module_type in module_loaders: @@ -998,6 +1002,61 @@ class SchedulerLoader(ComponentLoader): return scheduler +class ConditionerLoader(ComponentLoader): + """Loader for multi-conditioner components (e.g. Stable Audio's + `StableAudioMultiConditioner`, which bundles T5 + NumberConditioners + and is neither a pure text encoder nor a Diffusers-shaped module). + Reads `/config.json` to resolve the class via + `ModelRegistry`, instantiates with no args (the class pulls its own + defaults from its FastVideo config), then loads + `diffusion_pytorch_model.safetensors` non-strictly so externally + fetched sub-encoders (T5) don't trip the missing-key check. + """ + + def load(self, model_path: str, fastvideo_args: FastVideoArgs): + config = get_diffusers_config(model=model_path) + class_name = config.pop("_class_name", None) + config.pop("_name_or_path", None) + if class_name is None: + raise ValueError( + f"Conditioner config at {model_path} is missing the " + f"`_class_name` attribute required to resolve a model class.") + model_cls, _ = ModelRegistry.resolve_model_cls(class_name) + + target_device = get_local_torch_device() + precision = getattr(fastvideo_args.pipeline_config, "precision", "fp16") + target_dtype = PRECISION_TO_TYPE.get(precision, torch.float16) + + # Without this merge the model falls back to its dataclass + # defaults (e.g. SA-1.0's 3-conditioner spec — wrong for SA-small). + from dataclasses import fields as _fields + from fastvideo.configs.models.encoders import ( + StableAudioConditionerConfig, ) + if model_cls.__name__ == "StableAudioMultiConditioner": + cond_config = StableAudioConditionerConfig() + # `update_model_arch` is strict (raises on unknown keys); the + # converter writes a few non-arch keys (`_class_name`, + # `_diffusers_version`, `_name_or_path`) that must be filtered + # out first. + valid = {f.name for f in _fields(cond_config.arch_config)} + cond_config.update_model_arch({k: v for k, v in config.items() if k in valid}) + with set_default_torch_dtype(target_dtype): + model = model_cls(cond_config) + else: + with set_default_torch_dtype(target_dtype): + model = model_cls() + + weights = os.path.join(str(model_path), "diffusion_pytorch_model.safetensors") + if not os.path.isfile(weights): + raise FileNotFoundError( + f"Conditioner weights not found: {weights}") + state = safetensors_load_file(weights) + # Non-strict: T5 weights live outside this checkpoint (fetched in + # the conditioner's `__init__` from the standard HF repo). + model.load_state_dict(state, strict=False) + return model.to(device=target_device, dtype=target_dtype).eval() + + class UpsamplerLoader(ComponentLoader): """Loader for upsamplers.""" diff --git a/fastvideo/models/registry.py b/fastvideo/models/registry.py index da296fa40..a298cf667 100644 --- a/fastvideo/models/registry.py +++ b/fastvideo/models/registry.py @@ -86,6 +86,9 @@ _VAE_MODELS = { ("vaes", "gen3c_tokenizer_vae", "AutoencoderKLGen3CTokenizer"), "AutoencoderKLStepvideo": ("vaes", "stepvideovae", "AutoencoderKLStepvideo"), "CausalVideoAutoencoder": ("vaes", "ltx2vae", "LTX2CausalVideoAutoencoder"), + # `stable-audio-open-1.0/vae/config.json` ships `_class_name="AutoencoderOobleck"` + # (Diffusers' name); FastVideo's class is `OobleckVAE`. + "AutoencoderOobleck": ("vaes", "oobleck", "OobleckVAE"), } _AUDIO_MODELS = { diff --git a/fastvideo/models/vaes/oobleck.py b/fastvideo/models/vaes/oobleck.py new file mode 100644 index 000000000..8ba2fd7c8 --- /dev/null +++ b/fastvideo/models/vaes/oobleck.py @@ -0,0 +1,376 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Stable Audio Open 1.0 "Oobleck" VAE. + +5-stage Conv1d autoencoder with Snake activations + diagonal-Gaussian +bottleneck. Loads `stabilityai/stable-audio-open-1.0/vae/` weights +directly via `OobleckVAE.from_pretrained(...)`. + + vae = OobleckVAE.from_pretrained("stabilityai/stable-audio-open-1.0", subfolder="vae") + waveform = vae.decode(latent) # (B, audio_channels, samples) + latent = vae.encode(waveform).sample() # or .mode() +""" +from __future__ import annotations + +import json +import math +import os +from dataclasses import dataclass + +import numpy as np +import torch +import torch.nn as nn +from torch.nn.utils import weight_norm + +from fastvideo.logger import init_logger + +logger = init_logger(__name__) + + +class Snake1d(nn.Module): + """A 1D Snake activation with learnable per-channel alpha/beta.""" + + def __init__(self, hidden_dim: int, logscale: bool = True): + super().__init__() + self.alpha = nn.Parameter(torch.zeros(1, hidden_dim, 1)) + self.beta = nn.Parameter(torch.zeros(1, hidden_dim, 1)) + self.alpha.requires_grad = True + self.beta.requires_grad = True + self.logscale = logscale + + def forward(self, x: torch.Tensor) -> torch.Tensor: + shape = x.shape + alpha = self.alpha if not self.logscale else torch.exp(self.alpha) + beta = self.beta if not self.logscale else torch.exp(self.beta) + x = x.reshape(shape[0], shape[1], -1) + x = x + (beta + 1e-9).reciprocal() * torch.sin(alpha * x).pow(2) + return x.reshape(shape) + + +class OobleckResidualUnit(nn.Module): + def __init__(self, dimension: int = 16, dilation: int = 1): + super().__init__() + pad = ((7 - 1) * dilation) // 2 + self.snake1 = Snake1d(dimension) + self.conv1 = weight_norm(nn.Conv1d( + dimension, dimension, kernel_size=7, dilation=dilation, padding=pad, + )) + self.snake2 = Snake1d(dimension) + self.conv2 = weight_norm(nn.Conv1d(dimension, dimension, kernel_size=1)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + out = self.conv1(self.snake1(x)) + out = self.conv2(self.snake2(out)) + pad = (x.shape[-1] - out.shape[-1]) // 2 + if pad > 0: + x = x[..., pad:-pad] + return x + out + + +class OobleckEncoderBlock(nn.Module): + def __init__(self, input_dim: int, output_dim: int, stride: int = 1): + super().__init__() + self.res_unit1 = OobleckResidualUnit(input_dim, dilation=1) + self.res_unit2 = OobleckResidualUnit(input_dim, dilation=3) + self.res_unit3 = OobleckResidualUnit(input_dim, dilation=9) + self.snake1 = Snake1d(input_dim) + self.conv1 = weight_norm(nn.Conv1d( + input_dim, output_dim, + kernel_size=2 * stride, stride=stride, + padding=math.ceil(stride / 2), + )) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.res_unit1(x) + x = self.res_unit2(x) + x = self.snake1(self.res_unit3(x)) + return self.conv1(x) + + +class OobleckDecoderBlock(nn.Module): + def __init__(self, input_dim: int, output_dim: int, stride: int = 1): + super().__init__() + self.snake1 = Snake1d(input_dim) + self.conv_t1 = weight_norm(nn.ConvTranspose1d( + input_dim, output_dim, + kernel_size=2 * stride, stride=stride, + padding=math.ceil(stride / 2), + )) + self.res_unit1 = OobleckResidualUnit(output_dim, dilation=1) + self.res_unit2 = OobleckResidualUnit(output_dim, dilation=3) + self.res_unit3 = OobleckResidualUnit(output_dim, dilation=9) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.snake1(x) + x = self.conv_t1(x) + x = self.res_unit1(x) + x = self.res_unit2(x) + return self.res_unit3(x) + + +class OobleckDiagonalGaussianDistribution: + """Diagonal-Gaussian VAE posterior with `softplus(scale) + 1e-4` std.""" + + def __init__(self, parameters: torch.Tensor, deterministic: bool = False): + self.parameters = parameters + self.mean, self.scale = parameters.chunk(2, dim=1) + self.std = nn.functional.softplus(self.scale) + 1e-4 + self.var = self.std * self.std + self.logvar = torch.log(self.var) + self.deterministic = deterministic + + def sample(self, generator: torch.Generator | None = None) -> torch.Tensor: + noise = torch.randn( + self.mean.shape, generator=generator, + device=self.parameters.device, dtype=self.parameters.dtype, + ) + return self.mean + self.std * noise + + def mode(self) -> torch.Tensor: + return self.mean + + +@dataclass +class OobleckDecoderOutput: + sample: torch.Tensor + + +class OobleckEncoder(nn.Module): + def __init__( + self, + encoder_hidden_size: int, + audio_channels: int, + downsampling_ratios: list[int], + channel_multiples: list[int], + ): + super().__init__() + strides = downsampling_ratios + channel_multiples = [1] + list(channel_multiples) + self.conv1 = weight_norm(nn.Conv1d( + audio_channels, encoder_hidden_size, kernel_size=7, padding=3, + )) + self.block = nn.ModuleList([ + OobleckEncoderBlock( + input_dim=encoder_hidden_size * channel_multiples[i], + output_dim=encoder_hidden_size * channel_multiples[i + 1], + stride=s, + ) + for i, s in enumerate(strides) + ]) + d_model = encoder_hidden_size * channel_multiples[-1] + self.snake1 = Snake1d(d_model) + self.conv2 = weight_norm(nn.Conv1d( + d_model, encoder_hidden_size, kernel_size=3, padding=1, + )) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.conv1(x) + for m in self.block: + x = m(x) + x = self.snake1(x) + return self.conv2(x) + + +class OobleckDecoder(nn.Module): + def __init__( + self, + channels: int, + input_channels: int, + audio_channels: int, + upsampling_ratios: list[int], + channel_multiples: list[int], + ): + super().__init__() + strides = upsampling_ratios + channel_multiples = [1] + list(channel_multiples) + self.conv1 = weight_norm(nn.Conv1d( + input_channels, channels * channel_multiples[-1], + kernel_size=7, padding=3, + )) + self.block = nn.ModuleList([ + OobleckDecoderBlock( + input_dim=channels * channel_multiples[len(strides) - i], + output_dim=channels * channel_multiples[len(strides) - i - 1], + stride=s, + ) + for i, s in enumerate(strides) + ]) + self.snake1 = Snake1d(channels) + self.conv2 = weight_norm(nn.Conv1d( + channels, audio_channels, kernel_size=7, padding=3, bias=False, + )) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.conv1(x) + for layer in self.block: + x = layer(x) + x = self.snake1(x) + return self.conv2(x) + + +# --------------------------------------------------------------------------- +# Top-level VAE +# --------------------------------------------------------------------------- + + +class OobleckVAE(nn.Module): + """Stable Audio Open 1.0 VAE. + + Constructed either from an `OobleckVAEConfig` (the standard + `VAELoader` path) or from explicit kwargs (back-compat for tests + and `from_pretrained` callers). + """ + + def __init__( + self, + config=None, # type: OobleckVAEConfig | None + *, + encoder_hidden_size: int = 128, + downsampling_ratios: list[int] | None = None, + channel_multiples: list[int] | None = None, + decoder_channels: int = 128, + decoder_input_channels: int = 64, + audio_channels: int = 2, + sampling_rate: int = 44100, + ): + super().__init__() + if config is not None: + arch = config.arch_config + encoder_hidden_size = arch.encoder_hidden_size + downsampling_ratios = list(arch.downsampling_ratios) + channel_multiples = list(arch.channel_multiples) + decoder_channels = arch.decoder_channels + decoder_input_channels = arch.decoder_input_channels + audio_channels = arch.audio_channels + sampling_rate = arch.sampling_rate + if downsampling_ratios is None: + downsampling_ratios = [2, 4, 4, 8, 8] + if channel_multiples is None: + channel_multiples = [1, 2, 4, 8, 16] + self.encoder_hidden_size = encoder_hidden_size + self.downsampling_ratios = downsampling_ratios + self.decoder_channels = decoder_channels + self.upsampling_ratios = list(reversed(downsampling_ratios)) + self.hop_length = int(np.prod(downsampling_ratios)) + self.sampling_rate = sampling_rate + self.audio_channels = audio_channels + self.decoder_input_channels = decoder_input_channels + + self.encoder = OobleckEncoder( + encoder_hidden_size=encoder_hidden_size, + audio_channels=audio_channels, + downsampling_ratios=downsampling_ratios, + channel_multiples=channel_multiples, + ) + self.decoder = OobleckDecoder( + channels=decoder_channels, + input_channels=decoder_input_channels, + audio_channels=audio_channels, + upsampling_ratios=self.upsampling_ratios, + channel_multiples=channel_multiples, + ) + + def encode( + self, x: torch.Tensor, + ) -> OobleckDiagonalGaussianDistribution: + return OobleckDiagonalGaussianDistribution(self.encoder(x)) + + def decode(self, z: torch.Tensor) -> OobleckDecoderOutput: + return OobleckDecoderOutput(sample=self.decoder(z)) + + def forward( + self, sample: torch.Tensor, sample_posterior: bool = False, + ) -> OobleckDecoderOutput: + posterior = self.encode(sample) + z = posterior.sample() if sample_posterior else posterior.mode() + return self.decode(z) + + # ------------------------------------------------------------------- + # Loader + # ------------------------------------------------------------------- + @classmethod + def from_pretrained( + cls, + model_path: str, + *, + subfolder: str | None = None, + torch_dtype: torch.dtype | None = None, + ) -> "OobleckVAE": + """Instantiate and load weights from a Stable Audio VAE dir. + + `model_path` may be: + * a HF repo id (e.g. `stabilityai/stable-audio-open-1.0`), + * a local directory containing `config.json` + safetensors, + * a local directory whose `subfolder="vae"` holds those files. + + For gated repos, the HF token is read from `HF_TOKEN` / + `HUGGINGFACE_HUB_TOKEN` / `HF_API_KEY` (see `resolve_hf_token`). + """ + import inspect + from safetensors.torch import load_file + + from fastvideo.utils import resolve_hf_token + + # Resolve to a local directory. + if os.path.isdir(model_path): + root = model_path + else: + from huggingface_hub import snapshot_download + allow = ["vae/*"] if subfolder else ["*"] + root = snapshot_download( + repo_id=model_path, token=resolve_hf_token(), allow_patterns=allow, + ) + if subfolder: + root = os.path.join(root, subfolder) + if not os.path.isdir(root): + raise FileNotFoundError(f"Not a directory: {root}") + + cfg_path = os.path.join(root, "config.json") + if not os.path.isfile(cfg_path): + raise FileNotFoundError( + f"Expected config.json at {cfg_path}. If using a HF repo, " + f"pass subfolder='vae'." + ) + with open(cfg_path) as f: + cfg = json.load(f) + cfg_fields = {k: v for k, v in cfg.items() if not k.startswith("_")} + # Diffusers configs commonly carry extra fields (`scaling_factor`, + # `_diffusers_version`, ...) the bare `OobleckVAE` ctor doesn't accept. + init_params = inspect.signature(cls.__init__).parameters + cfg_fields = {k: v for k, v in cfg_fields.items() if k in init_params} + + model = cls(**cfg_fields) + + weights_path = os.path.join(root, "diffusion_pytorch_model.safetensors") + if not os.path.isfile(weights_path): + # Allow `model.safetensors` as a fallback. + alt = os.path.join(root, "model.safetensors") + if os.path.isfile(alt): + weights_path = alt + else: + raise FileNotFoundError( + f"No safetensors weights under {root}. Expected " + f"diffusion_pytorch_model.safetensors." + ) + state = load_file(weights_path) + missing, unexpected = model.load_state_dict(state, strict=False) + if missing: + raise RuntimeError( + f"OobleckVAE missing {len(missing)} keys from {weights_path}: " + f"{missing[:5]}" + ) + if unexpected: + # Non-critical: some checkpoints embed the VAE inside a larger + # container (e.g. `pretransform.model.*`). Log the count so + # genuine loader regressions don't go unnoticed. + logger.debug( + "OobleckVAE: ignored %d unexpected keys from %s " + "(first 3: %s)", len(unexpected), weights_path, unexpected[:3], + ) + + if torch_dtype is not None: + model = model.to(dtype=torch_dtype) + model.eval() + return model + + +EntryClass = OobleckVAE diff --git a/fastvideo/models/vaes/sa_audio.py b/fastvideo/models/vaes/sa_audio.py new file mode 100644 index 000000000..1031f9533 --- /dev/null +++ b/fastvideo/models/vaes/sa_audio.py @@ -0,0 +1,122 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Lazy-loading pipeline wrapper around `OobleckVAE`. + +Two reasons this exists rather than using `OobleckVAE` directly: + + 1. The underlying VAE is fetched on first `encode`/`decode` call, not + at construction — lets pipelines build the module tree on CPU + before knowing the target device. + 2. The lazy VAE's params are hidden from `named_parameters()` so the + FastVideo pipeline-component loader doesn't try to match Oobleck's + safetensors against the host pipeline's converted-repo state dict. + +For standalone use prefer `OobleckVAE.from_pretrained(...)` directly. +""" +from __future__ import annotations + +import os + +import torch +from torch import nn + +from fastvideo.configs.models.vaes import OobleckVAEConfig + + +class SAAudioVAEModel(nn.Module): + """Pipeline-glue lazy loader around `OobleckVAE`.""" + + def __init__(self, config: OobleckVAEConfig) -> None: + super().__init__() + self.config = config + arch = config.arch_config + self.pretrained_path: str = config.pretrained_path + self.pretrained_subfolder: str | None = config.pretrained_subfolder + self.pretrained_dtype: str = config.pretrained_dtype + self.sampling_rate: int = arch.sampling_rate + self.audio_channels: int = arch.audio_channels + self.decoder_input_channels: int = arch.decoder_input_channels + self._oobleck_vae = None + + def named_parameters(self, prefix: str = "", recurse: bool = True): + # Hide the lazy-loaded VAE — its weights are fetched separately + # and shouldn't appear in the host pipeline's loader sweep. + for name, param in super().named_parameters(prefix=prefix, recurse=recurse): + if name.startswith("_oobleck_vae.") or name == "_oobleck_vae": + continue + yield name, param + + def _build(self, device: torch.device | None = None): + from fastvideo.models.vaes.oobleck import OobleckVAE + + path = self.pretrained_path + if not path: + raise ValueError( + "OobleckVAEConfig.pretrained_path must be set; expected " + "`stabilityai/stable-audio-open-1.0` or a local path." + ) + dtype = getattr(torch, self.pretrained_dtype, torch.float32) + # If the caller already pointed us at the VAE dir directly, drop + # the subfolder. Otherwise pass through (default "vae"). + subfolder: str | None = self.pretrained_subfolder + if subfolder and os.path.isdir(path) and os.path.isfile(os.path.join(path, "config.json")): + subfolder = None + model = OobleckVAE.from_pretrained(path, subfolder=subfolder, torch_dtype=dtype) + if device is not None: + model = model.to(device=device) + model.eval() + return model + + @property + def oobleck_vae(self): + if self._oobleck_vae is None: + self._oobleck_vae = self._build() + return self._oobleck_vae + + # Back-compat alias: callers that imported this earlier referred to + # the underlying VAE as `sa_audio_vae_model`. Both names point at the + # same object. + @property + def sa_audio_vae_model(self): + return self.oobleck_vae + + @property + def hop_length(self) -> int: + return int(self.oobleck_vae.hop_length) + + def _move_to_input_device(self, model, ref: torch.Tensor): + if ref is None: + return model + first_param = next(model.parameters(), None) + if first_param is not None and first_param.device != ref.device: + model = model.to(device=ref.device) + self._oobleck_vae = model + return model + + def decode(self, latent: torch.Tensor) -> torch.Tensor: + """Decode an audio latent (`[B, C_latent, L]`) -> waveform + (`[B, audio_channels, samples]`). + """ + model = self.oobleck_vae + model = self._move_to_input_device(model, latent) + with torch.no_grad(): + out = model.decode(latent.to(next(model.parameters()).dtype)) + if hasattr(out, "sample"): + return out.sample + return out + + def encode(self, waveform: torch.Tensor, sample_posterior: bool = False) -> torch.Tensor: + """Encode `[B, C_audio, samples]` -> latent `[B, C_latent, L]`. + + `sample_posterior=False` (default): deterministic mean. + `sample_posterior=True`: stochastic sample (`mean + softplus(scale) * randn`). + """ + model = self.oobleck_vae + model = self._move_to_input_device(model, waveform) + with torch.no_grad(): + out = model.encode(waveform.to(next(model.parameters()).dtype)) + if hasattr(out, "latent_dist"): + out = out.latent_dist + return out.sample() if sample_posterior else out.mode() + + +EntryClass = SAAudioVAEModel diff --git a/fastvideo/pipelines/basic/stable_audio/__init__.py b/fastvideo/pipelines/basic/stable_audio/__init__.py new file mode 100644 index 000000000..988131360 --- /dev/null +++ b/fastvideo/pipelines/basic/stable_audio/__init__.py @@ -0,0 +1 @@ +# SPDX-License-Identifier: Apache-2.0 diff --git a/fastvideo/pipelines/basic/stable_audio/presets.py b/fastvideo/pipelines/basic/stable_audio/presets.py new file mode 100644 index 000000000..4f2a964bb --- /dev/null +++ b/fastvideo/pipelines/basic/stable_audio/presets.py @@ -0,0 +1,63 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Stable Audio presets. + +Sampling defaults track the published HF model card +(https://huggingface.co/stabilityai/stable-audio-open-1.0): +100 steps, CFG=7, dpmpp-3m-sde, sigma_min=0.3, sigma_max=500, rho=1.0. +""" +from fastvideo.api.presets import InferencePreset, PresetStageSpec + +_DENOISE_STAGE = PresetStageSpec( + name="denoise", + kind="denoising", + description="Stable Audio Cosine-DPM++ denoising with text + duration CFG.", + allowed_overrides=frozenset({ + "num_inference_steps", + "guidance_scale", + }), +) + +# `audio_start_in_s` / `audio_end_in_s` are call-kwargs, kept off here. +# `height`/`width` are pinned to 8 (the shared `InputValidationStage` +# rejects values that aren't divisible by 8) and `num_frames` to 1 so +# the video-shaped preallocation in `VideoGenerator` stays tiny — the +# real output is the audio waveform on `result["audio"]`, not the +# placeholder frame tensor. +_SHARED_DEFAULTS = { + "seed": 0, + "guidance_scale": 7.0, + "num_inference_steps": 100, + "negative_prompt": "", + "height": 8, + "width": 8, + "num_frames": 1, +} + +STABLE_AUDIO_OPEN_1_0_BASE = InferencePreset( + name="stable_audio_open_1_0_base", + version=1, + model_family="stable_audio", + description=("Stability AI Stable Audio Open 1.0 text-to-audio. Generates up " + "to ~47.5s of stereo 44.1 kHz audio per call. Default duration " + "is 10s; raise via `audio_end_in_s` up to the model max."), + workload_type="t2v", # NOTE: WorkloadType has no T2A variant yet (REVIEW item 28) + stage_schemas=(_DENOISE_STAGE, ), + defaults=dict(_SHARED_DEFAULTS), +) + +# Smaller / faster checkpoint with the same Oobleck VAE but a 1024-dim +# 16-layer DiT with `qk_norm="ln"`. Sampling defaults match the official +# `stable-audio-open-small` model card. +STABLE_AUDIO_OPEN_SMALL = InferencePreset( + name="stable_audio_open_small", + version=1, + model_family="stable_audio", + description=("Stability AI Stable Audio Open Small text-to-audio. Faster than " + "the 1.0 base; supports up to ~11.9s of stereo 44.1 kHz audio per " + "call (smaller training window)."), + workload_type="t2v", + stage_schemas=(_DENOISE_STAGE, ), + defaults=dict(_SHARED_DEFAULTS), +) + +ALL_PRESETS = (STABLE_AUDIO_OPEN_1_0_BASE, STABLE_AUDIO_OPEN_SMALL) diff --git a/fastvideo/pipelines/basic/stable_audio/stable_audio_pipeline.py b/fastvideo/pipelines/basic/stable_audio/stable_audio_pipeline.py new file mode 100644 index 000000000..7cb319c5c --- /dev/null +++ b/fastvideo/pipelines/basic/stable_audio/stable_audio_pipeline.py @@ -0,0 +1,125 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Stable Audio Open 1.0 pipeline (T2A + A2A + RePaint inpainting). + +Components are loaded via the standard +`ComposedPipelineBase.load_modules` against the FastVideo-curated +Diffusers-format repo `FastVideo/stable-audio-open-1.0-Diffusers` +(produced by +`scripts/checkpoint_conversion/stable_audio_to_diffusers.py`). The DiT +is a `BaseDiT` subclass loaded by `TransformerLoader`; the VAE is +loaded by `VAELoader`; the multi-conditioner (T5 + NumberConditioners) +is loaded by `ConditionerLoader` (a Stable Audio-specific addition). + +Stages: + + InputValidationStage + → StableAudioConditioningStage (T5 + NumberConditioner -> cross-attn + global cond, with CFG) + → StableAudioLatentPreparationStage (initial Gaussian noise; encodes A2A / inpaint refs) + → StableAudioDenoisingStage (k-diffusion `dpmpp-3m-sde` over the DiT) + → StableAudioDecodingStage (OobleckVAE -> waveform) +""" +from __future__ import annotations + +import functools + +import torch + +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.logger import init_logger +from fastvideo.pipelines.basic.stable_audio.stages import ( + StableAudioConditioningStage, + StableAudioDecodingStage, + StableAudioDenoisingStage, + StableAudioLatentPreparationStage, +) +from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase +from fastvideo.pipelines.stages import InputValidationStage + +logger = init_logger(__name__) + + +@functools.lru_cache(maxsize=1) +def _warn_tf32_disabled_for_stable_audio() -> None: + logger.warning("Stable Audio pipeline is disabling process-global " + "torch.backends.{cuda.matmul.allow_tf32, cudnn.allow_tf32, " + "cuda.matmul.allow_fp16_reduced_precision_reduction, " + "cudnn.benchmark} for A2A renoise determinism. Other models " + "loaded into this process will inherit these settings.") + + +def _disable_tf32_for_stable_audio() -> None: + """Disable TF32 / cuDNN nondeterminism — A2A renoise-then-denoise SDE + amplifies per-element drift, and the published parity bounds were + set with these off. Process-global; the first call logs a warning. + """ + _warn_tf32_disabled_for_stable_audio() + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cudnn.allow_tf32 = False + torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = False + torch.backends.cudnn.benchmark = False + + +class StableAudioPipeline(ComposedPipelineBase): + """Stable Audio Open 1.0 pipeline. + + Mode is kwargs-driven on `generate_video()`: + + * Text-to-audio (default) -- `prompt=...`, `audio_end_in_s=...` + * Audio-to-audio variation -- add `init_audio=ref` (and optionally + `init_noise_level`, lower = closer to reference) + * RePaint inpainting / outpainting -- add `inpaint_audio=ref` and + `inpaint_mask` (1-D, 1 = keep / 0 = regenerate) + + See `examples/inference/basic/basic_stable_audio*.py` for runnable + examples of each mode. + """ + + _required_config_modules = [ + "vae", + "transformer", + "conditioner", + ] + + def initialize_pipeline(self, fastvideo_args: FastVideoArgs) -> None: + """Apply Stable Audio's process-global numerics overrides BEFORE + the standard component loaders run (TF32 off for A2A renoise + determinism).""" + _disable_tf32_for_stable_audio() + + def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None: + pc = fastvideo_args.pipeline_config + + self.add_stage(stage_name="input_validation_stage", stage=InputValidationStage()) + + self.add_stage( + stage_name="conditioning_stage", + stage=StableAudioConditioningStage(conditioner=self.get_module("conditioner")), + ) + + self.add_stage( + stage_name="latent_preparation_stage", + stage=StableAudioLatentPreparationStage( + io_channels=64, + # Per-variant training window: 2,097,152 (~47.5s) for + # SA-1.0; 524,288 (~11.9s) for SA-small. Pulled from the + # pipeline config so each variant gets its own latent + # length. + sample_size=pc.sample_size, + vae=self.get_module("vae"), + sample_rate=pc.sampling_rate, + audio_channels=pc.audio_channels, + ), + ) + + self.add_stage( + stage_name="denoising_stage", + stage=StableAudioDenoisingStage(transformer=self.get_module("transformer")), + ) + + self.add_stage( + stage_name="decoding_stage", + stage=StableAudioDecodingStage(vae=self.get_module("vae")), + ) + + +EntryClass = StableAudioPipeline diff --git a/fastvideo/pipelines/basic/stable_audio/stages/__init__.py b/fastvideo/pipelines/basic/stable_audio/stages/__init__.py new file mode 100644 index 000000000..8068754fe --- /dev/null +++ b/fastvideo/pipelines/basic/stable_audio/stages/__init__.py @@ -0,0 +1,12 @@ +# SPDX-License-Identifier: Apache-2.0 +from fastvideo.pipelines.basic.stable_audio.stages.conditioning import StableAudioConditioningStage +from fastvideo.pipelines.basic.stable_audio.stages.decoding import StableAudioDecodingStage +from fastvideo.pipelines.basic.stable_audio.stages.denoising import StableAudioDenoisingStage +from fastvideo.pipelines.basic.stable_audio.stages.latent_preparation import StableAudioLatentPreparationStage + +__all__ = [ + "StableAudioConditioningStage", + "StableAudioDecodingStage", + "StableAudioDenoisingStage", + "StableAudioLatentPreparationStage", +] diff --git a/fastvideo/pipelines/basic/stable_audio/stages/conditioning.py b/fastvideo/pipelines/basic/stable_audio/stages/conditioning.py new file mode 100644 index 000000000..a70593f9c --- /dev/null +++ b/fastvideo/pipelines/basic/stable_audio/stages/conditioning.py @@ -0,0 +1,98 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Stable Audio conditioning stage.""" +from __future__ import annotations + +import torch + +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.pipelines.stages.base import PipelineStage +from fastvideo.pipelines.stages.validators import VerificationResult + + +class StableAudioConditioningStage(PipelineStage): + """Run the conditioner over the prompt + duration and stash the + DiT-ready (cross_attn_cond, cross_attn_mask, global_embed) triple + on `batch.extra` (plus the negative-prompt triple when CFG is on). + """ + + def __init__(self, conditioner) -> None: + super().__init__() + self.conditioner = conditioner + + def verify_input(self, batch, fastvideo_args): + return VerificationResult() + + def verify_output(self, batch, fastvideo_args): + return VerificationResult() + + @torch.inference_mode() + def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch: + pc = fastvideo_args.pipeline_config + device = next(self.conditioner.parameters()).device + + start_attr = getattr(batch, "audio_start_in_s", None) + end_attr = getattr(batch, "audio_end_in_s", None) + audio_start_in_s = float(start_attr if start_attr is not None else pc.audio_start_in_s) + audio_end_in_s = float(end_attr if end_attr is not None else pc.audio_end_in_s) + max_duration = float(getattr(pc, "max_audio_duration_s", 2097152 / 44100)) + if audio_start_in_s < 0: + raise ValueError(f"audio_start_in_s must be >= 0, got {audio_start_in_s}.") + if audio_end_in_s <= audio_start_in_s: + raise ValueError(f"audio_end_in_s ({audio_end_in_s}) must be > audio_start_in_s " + f"({audio_start_in_s}).") + if audio_end_in_s > max_duration: + raise ValueError(f"audio_end_in_s ({audio_end_in_s}s) exceeds the model's fixed " + f"window of {max_duration:.4f}s. Stable Audio Open 1.0 always " + f"samples a 2,097,152-frame latent and slices to " + f"[start, end] after decode; values past the window are silently " + f"truncated. Lower audio_end_in_s or split the request.") + guidance_scale = float(batch.guidance_scale or pc.guidance_scale) + do_cfg = guidance_scale > 1.0 + + if isinstance(batch.prompt, str): + prompt = batch.prompt + elif isinstance(batch.prompt, list): + if len(batch.prompt) > 1: + raise ValueError(f"Stable Audio does not support batched prompts; got " + f"{len(batch.prompt)} entries. Pass a single string or a " + f"single-element list.") + prompt = batch.prompt[0] if batch.prompt else "" + else: + raise TypeError(f"`prompt` must be a string or a list of strings, got " + f"{type(batch.prompt).__name__}.") + # Send only the keys the conditioner declares (per-variant). + all_cond_values = { + "prompt": prompt, + "seconds_start": audio_start_in_s, + "seconds_total": audio_end_in_s, + } + active_ids = self.conditioner.cross_attention_cond_ids + cond_meta = [{k: all_cond_values[k] for k in active_ids if k in all_cond_values}] + cond = self.conditioner(cond_meta, device) + cross_attn_cond, cross_attn_mask, global_embed = self.conditioner.get_conditioning_inputs(cond) + + neg_cross_attn_cond = None + neg_cross_attn_mask = None + neg_global_embed = None + if do_cfg: + neg_prompt = batch.negative_prompt or "" + if isinstance(neg_prompt, list): + neg_prompt = neg_prompt[0] if neg_prompt else "" + neg_values = dict(all_cond_values, prompt=neg_prompt) + neg_meta = [{k: neg_values[k] for k in active_ids if k in neg_values}] + neg = self.conditioner(neg_meta, device) + neg_cross_attn_cond, neg_cross_attn_mask, neg_global_embed = (self.conditioner.get_conditioning_inputs(neg)) + + if batch.extra is None: + batch.extra = {} + batch.extra["cross_attn_cond"] = cross_attn_cond + batch.extra["cross_attn_mask"] = cross_attn_mask + batch.extra["global_embed"] = global_embed + batch.extra["negative_cross_attn_cond"] = neg_cross_attn_cond + batch.extra["negative_cross_attn_mask"] = neg_cross_attn_mask + batch.extra["negative_global_embed"] = neg_global_embed + batch.extra["do_cfg"] = do_cfg + batch.extra["audio_start_in_s"] = audio_start_in_s + batch.extra["audio_end_in_s"] = audio_end_in_s + return batch diff --git a/fastvideo/pipelines/basic/stable_audio/stages/decoding.py b/fastvideo/pipelines/basic/stable_audio/stages/decoding.py new file mode 100644 index 000000000..63cbd1bd6 --- /dev/null +++ b/fastvideo/pipelines/basic/stable_audio/stages/decoding.py @@ -0,0 +1,67 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Stable Audio decoding: latent -> waveform via OobleckVAE. + +Slices the output to `[audio_start_in_s, audio_end_in_s]` and stashes +the result on `batch.extra["audio"]` + `["audio_sample_rate"]` for +`VideoGenerator._mux_audio` to pick up. +""" +from __future__ import annotations + +import torch + +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.pipelines.stages.base import PipelineStage +from fastvideo.pipelines.stages.validators import VerificationResult + + +class StableAudioDecodingStage(PipelineStage): + """Decode latent → audio waveform + slice to [start, end].""" + + def __init__(self, vae) -> None: + super().__init__() + self.vae = vae + + def verify_input(self, batch, fastvideo_args): + return VerificationResult() + + def verify_output(self, batch, fastvideo_args): + return VerificationResult() + + @torch.inference_mode() + def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch: + pc = fastvideo_args.pipeline_config + latents = batch.latents + + # VAE may be CPU-parked under `vae_cpu_offload=True`. + from fastvideo.distributed.parallel_state import get_local_torch_device + self.vae = self.vae.to(get_local_torch_device()) + decoded = self.vae.decode(latents) + if hasattr(decoded, "sample"): # tolerate tensor or dataclass + decoded = decoded.sample + + sr = int(getattr(self.vae, "sampling_rate", pc.sampling_rate)) + start_in_s = float(batch.extra.get("audio_start_in_s", pc.audio_start_in_s)) + end_in_s = float(batch.extra.get("audio_end_in_s", pc.audio_end_in_s)) + decoded = decoded[:, :, int(start_in_s * sr):int(end_in_s * sr)] + + if batch.extra is None: + batch.extra = {} + # `_mux_audio` / `_write_pcm_wav` want `[samples, channels]`. + batch.extra["audio"] = decoded.squeeze(0).T.detach().float().cpu().numpy() + batch.extra["audio_sample_rate"] = sr + batch.extra["audio_only"] = True + # Raw tensor for parity tests. + batch.extra["decoded_audio"] = decoded.detach().cpu() + + # `VideoGenerator.generate_video` is video-shaped (asserts + # `output_batch.output is not None`); fill with a placeholder of + # the expected `[B, 3, num_frames, H, W]` shape — the real audio + # is on `batch.extra` above. Pure-audio workload support tracked + # in REVIEW item 28. + b = decoded.shape[0] + n_frames = int(getattr(batch, "num_frames", 1) or 1) + h = int(getattr(batch, "height", 1) or 1) + w = int(getattr(batch, "width", 1) or 1) + batch.output = torch.zeros((b, 3, n_frames, h, w), dtype=torch.uint8) + return batch diff --git a/fastvideo/pipelines/basic/stable_audio/stages/denoising.py b/fastvideo/pipelines/basic/stable_audio/stages/denoising.py new file mode 100644 index 000000000..e99374de8 --- /dev/null +++ b/fastvideo/pipelines/basic/stable_audio/stages/denoising.py @@ -0,0 +1,207 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Stable Audio denoising — k-diffusion `dpmpp-3m-sde` over the DiT. + +CFG-batched conditioning is built once outside the sampler loop so the +adapter only does `cat([x, x])` + DiT call per step. +""" +from __future__ import annotations + +import math + +import torch +import torch.nn as nn + +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.forward_context import set_forward_context +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.pipelines.stages.base import PipelineStage +from fastvideo.pipelines.stages.validators import VerificationResult + + +class _DiTAdapter(nn.Module): + """`StableAudioDiT` -> `K.external.VDenoiser` adapter. + + `batch_cond` / `batch_global` are precomputed CFG-batched tensors + (`[2, ...]` for CFG, `[1, ...]` otherwise); building them once + outside the sampler loop saves ~3 cats × 100 steps per call. + """ + + def __init__(self, dit, *, batch_cond: torch.Tensor, batch_global: torch.Tensor, cfg_scale: float) -> None: + super().__init__() + self.dit = dit + self.batch_cond = batch_cond + self.batch_global = batch_global + self.cfg_scale = cfg_scale + self.do_cfg = cfg_scale != 1.0 + + def forward(self, x: torch.Tensor, t: torch.Tensor, **_unused) -> torch.Tensor: + if not self.do_cfg: + return self.dit(x, t, cross_attn_cond=self.batch_cond, global_embed=self.batch_global) + batch_x = torch.cat([x, x], dim=0) + batch_t = torch.cat([t, t], dim=0) + out = self.dit(batch_x, batch_t, cross_attn_cond=self.batch_cond, global_embed=self.batch_global) + cond_out, uncond_out = torch.chunk(out, 2, dim=0) + return uncond_out + (cond_out - uncond_out) * self.cfg_scale + + +class StableAudioDenoisingStage(PipelineStage): + """k-diffusion `dpmpp-3m-sde` sampling loop.""" + + # Sampler defaults from the published model card. + _SIGMA_MIN = 0.3 + _SIGMA_MAX = 500.0 + _RHO = 1.0 + _LOG_SIGMA_MIN = math.log(_SIGMA_MIN) + _LOG_SIGMA_MAX = math.log(_SIGMA_MAX) + + def __init__(self, transformer) -> None: + super().__init__() + self.transformer = transformer + + def _resolve_sigma_max(self, batch) -> float: + """Map A2A intent to `sigma_max`. + + Public knob is `init_audio_strength` (0..1, higher = closer to + source), log-interpolated between SIGMA_MIN (= preservation) and + SIGMA_MAX (= full T2A). Raw `init_noise_level` is the legacy + sigma_max override; passing both is an error. + """ + raw = getattr(batch, "init_noise_level", None) + strength = getattr(batch, "init_audio_strength", None) + if raw is not None and strength is not None: + raise ValueError("Pass `init_audio_strength` (0..1) OR `init_noise_level` " + "(raw sigma_max), not both.") + if raw is not None: + return float(raw) + s = max(0.0, min(1.0, float(strength) if strength is not None else 0.6)) + return float(math.exp(self._LOG_SIGMA_MAX - s * (self._LOG_SIGMA_MAX - self._LOG_SIGMA_MIN))) + + def verify_input(self, batch, fastvideo_args): + return VerificationResult() + + def verify_output(self, batch, fastvideo_args): + return VerificationResult() + + @torch.inference_mode() + def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch: + pc = fastvideo_args.pipeline_config + ext = batch.extra + device = batch.latents.device + guidance_scale = float(batch.guidance_scale or pc.guidance_scale) + steps = int(batch.num_inference_steps) + + import k_diffusion as K + + init_latent = ext.get("init_latent") + sigma_max = self._resolve_sigma_max(batch) if init_latent is not None else self._SIGMA_MAX + + sigmas = K.sampling.get_sigmas_polyexponential(steps, self._SIGMA_MIN, sigma_max, self._RHO, device=device) + # Cast noise + conditioning to the DiT's dtype before sampling + # (matches `stable_audio_tools/inference/generation.py:185-187`). + model_dtype = next(self.transformer.parameters()).dtype + + def _cast(t: torch.Tensor | None) -> torch.Tensor | None: + return t.to(model_dtype) if t is not None else None + + x = (batch.latents * sigmas[0]).to(model_dtype) + if init_latent is not None: + x = x + init_latent.to(model_dtype) + + batch_cond, batch_global = _build_cfg_conditioning( + cross_attn_cond=ext["cross_attn_cond"].to(model_dtype), + global_embed=ext["global_embed"].to(model_dtype), + negative_cross_attn_cond=_cast(ext.get("negative_cross_attn_cond")), + negative_cross_attn_mask=ext.get("negative_cross_attn_mask"), + negative_global_embed=_cast(ext.get("negative_global_embed")), + do_cfg=guidance_scale != 1.0, + ) + adapter = _DiTAdapter(self.transformer, + batch_cond=batch_cond, + batch_global=batch_global, + cfg_scale=guidance_scale) + denoiser = K.external.VDenoiser(adapter) + + # RePaint blending hook — works on any v-prediction model, no + # inpaint-trained checkpoint needed. + inpaint_mask = ext.get("inpaint_mask_latent") + inpaint_ref = ext.get("inpaint_reference_latent") + if inpaint_mask is not None and inpaint_ref is not None: + inpaint_mask = inpaint_mask.to(model_dtype) + inpaint_ref = inpaint_ref.to(model_dtype) + callback = _make_inpaint_callback(inpaint_ref, inpaint_mask, sigmas) + else: + callback = None + + # `LocalAttention` (in `StableAudioDiT`) reads `get_forward_context()` + # for `attn_metadata`; wrap the whole loop. + with set_forward_context(current_timestep=0, attn_metadata=None): + sampled = K.sampling.sample_dpmpp_3m_sde(denoiser, + x, + sigmas, + disable=False, + extra_args={}, + callback=callback) + + # Final blend so the kept region of the inpaint reference is exact. + if inpaint_mask is not None and inpaint_ref is not None: + sampled = inpaint_ref * inpaint_mask + sampled * (1 - inpaint_mask) + batch.latents = sampled + return batch + + +def _build_cfg_conditioning( + *, + cross_attn_cond: torch.Tensor, + global_embed: torch.Tensor, + negative_cross_attn_cond: torch.Tensor | None, + negative_cross_attn_mask: torch.Tensor | None, + negative_global_embed: torch.Tensor | None, + do_cfg: bool, +) -> tuple[torch.Tensor, torch.Tensor]: + """Build the CFG-batched `(cond, global)` tensors once. + + Cond ordering is `[conditioned, unconditioned]` (the adapter splits + with the same convention). Masked negative cond is zero-filled where + `mask == 0`. + """ + if not do_cfg: + return cross_attn_cond, global_embed + if negative_cross_attn_cond is not None: + if negative_cross_attn_mask is not None: + neg_mask = negative_cross_attn_mask.to(torch.bool).unsqueeze(2) + null_embed = torch.zeros_like(cross_attn_cond) + negative_cross_attn_cond = torch.where(neg_mask, negative_cross_attn_cond, null_embed) + batch_cond = torch.cat([cross_attn_cond, negative_cross_attn_cond], dim=0) + else: + batch_cond = torch.cat([cross_attn_cond, torch.zeros_like(cross_attn_cond)], dim=0) + other_global = global_embed if negative_global_embed is None else negative_global_embed + batch_global = torch.cat([global_embed, other_global], dim=0) + return batch_cond, batch_global + + +def _make_inpaint_callback(reference_latent: torch.Tensor, mask: torch.Tensor, sigmas: torch.Tensor): + """RePaint blending callback for the k-diffusion sampler. + + At every step, replaces the kept region (`mask == 1`) of the in- + flight latent with the reference re-noised to the next sigma — + pulls the kept region back onto the trajectory the model expects, + so RePaint-style inpainting converges on non-inpaint-trained models. + + Pre-allocates the noise buffer so the ~100 sampler steps don't churn + ~25 MB of fresh allocations per call. + """ + noise_buf = torch.empty_like(reference_latent) + inv_mask = 1 - mask + + def cb(info: dict) -> None: + i = int(info["i"]) + next_i = min(i + 1, len(sigmas) - 1) + sigma_next = float(sigmas[next_i]) + noise_buf.normal_() + # `state["x"]` is the live latent; the dpmpp-3m-sde sampler picks + # up our in-place mutation between steps (verified against + # k_diffusion 0.1.1.post1). + x = info["x"] + x.copy_((reference_latent + noise_buf * sigma_next) * mask + x * inv_mask) + + return cb diff --git a/fastvideo/pipelines/basic/stable_audio/stages/latent_preparation.py b/fastvideo/pipelines/basic/stable_audio/stages/latent_preparation.py new file mode 100644 index 000000000..cad3e1340 --- /dev/null +++ b/fastvideo/pipelines/basic/stable_audio/stages/latent_preparation.py @@ -0,0 +1,174 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Stable Audio latent preparation. + +Seeds + samples the initial Gaussian noise; encodes `init_audio` (A2A +variation) or `inpaint_audio` + `inpaint_mask` (RePaint inpainting) into +latent-space tensors on `batch.extra` for the denoising stage. +""" +from __future__ import annotations + +import os + +import torch +import torch.nn.functional as F + +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.pipelines.stages.base import PipelineStage +from fastvideo.pipelines.stages.validators import VerificationResult + + +class StableAudioLatentPreparationStage(PipelineStage): + + def __init__(self, + io_channels: int = 64, + sample_size: int = 2097152, + vae=None, + sample_rate: int = 44100, + audio_channels: int = 2) -> None: + super().__init__() + self.io_channels = io_channels + # Audio-domain length the model was trained for; latent length + # = sample_size // vae.hop_length (= 2097152 / 2048 = 1024). + self.sample_size = sample_size + self.vae = vae # used to encode init_audio / inpaint_audio + self.sample_rate = sample_rate + self.audio_channels = audio_channels + + def verify_input(self, batch, fastvideo_args): + return VerificationResult() + + def verify_output(self, batch, fastvideo_args): + return VerificationResult() + + @torch.inference_mode() + def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch: + ext = batch.extra or {} + device = ext["cross_attn_cond"].device + latent_sample_size = self.sample_size // self._hop_length() + + seed = int(batch.seed) if batch.seed is not None else 0 + torch.manual_seed(seed) + latents = torch.randn((1, self.io_channels, latent_sample_size), device=device) + + batch.latents = latents + if batch.extra is None: + batch.extra = {} + + init_audio = getattr(batch, "init_audio", None) + inpaint_audio = getattr(batch, "inpaint_audio", None) + inpaint_mask = getattr(batch, "inpaint_mask", None) + + # Loud-fail rather than silently falling through to T2A. + if inpaint_audio is not None and inpaint_mask is None: + raise ValueError("Stable Audio inpainting requires both `inpaint_audio` and " + "`inpaint_mask` (1-D tensor in {0, 1} at the model sample rate, " + "1 = keep, 0 = regenerate). Got `inpaint_audio` without `inpaint_mask`.") + if inpaint_mask is not None and inpaint_audio is None: + raise ValueError("Stable Audio inpainting requires both `inpaint_audio` and " + "`inpaint_mask`. Got `inpaint_mask` without `inpaint_audio` — " + "did you mean to pass `init_audio` (audio-to-audio variation)?") + if init_audio is not None and inpaint_audio is not None: + raise ValueError("Stable Audio cannot do A2A variation and inpainting in the " + "same call. Pass either `init_audio` (variation) or " + "`inpaint_audio` + `inpaint_mask` (inpainting), not both.") + + if init_audio is not None: + batch.extra["init_latent"] = self._encode_audio_reference(init_audio, device) + + if inpaint_audio is not None and inpaint_mask is not None: + batch.extra["inpaint_reference_latent"] = self._encode_audio_reference(inpaint_audio, device) + batch.extra["inpaint_mask_latent"] = self._prepare_mask(inpaint_mask, latent_sample_size, device) + return batch + + def _hop_length(self) -> int: + return int(self.vae.hop_length) + + def _encode_audio_reference(self, audio, device: torch.device) -> torch.Tensor: + """Pad/truncate to `sample_size` and encode via the VAE. + + `audio` may be a tensor (`[samples]`, `[C, samples]`, or + `[B, C, samples]`) at the model's sample rate, or a path to any + audio-bearing file (`.wav` / `.mp3` / `.mp4` / `.m4a` / `.flac`, + ...) that PyAV can decode — we resample on load so callers don't + have to. + """ + assert self.vae is not None, "VAE required for init_audio / inpaint_audio encoding" + # VAE may be CPU-parked under `vae_cpu_offload=True`. + self.vae = self.vae.to(device) + if isinstance(audio, str | os.PathLike): + audio = _decode_audio_file(audio, target_sr=self.sample_rate) + audio = audio.to(device=device, dtype=torch.float32) + if audio.dim() == 1: + audio = audio.unsqueeze(0).unsqueeze(0) + elif audio.dim() == 2: + audio = audio.unsqueeze(0) + # Match expected channel count (mono → repeat to stereo). + if audio.shape[1] == 1 and self.audio_channels == 2: + audio = audio.repeat(1, 2, 1) + elif audio.shape[1] == 2 and self.audio_channels == 1: + audio = audio.mean(dim=1, keepdim=True) + # Pad/truncate to model sample_size. + cur_len = audio.shape[-1] + if cur_len < self.sample_size: + audio = F.pad(audio, (0, self.sample_size - cur_len)) + elif cur_len > self.sample_size: + audio = audio[..., :self.sample_size] + # Stochastic sample (the next random draw after the latent + # `randn` above), so encode-noise stays on the seeded sequence. + return self.vae.encode(audio.to(next(self.vae.parameters()).dtype)).sample() + + def _prepare_mask(self, mask, latent_len: int, device: torch.device) -> torch.Tensor: + """Pad/truncate a binary mask to `sample_size`, then + nearest-resample to `[1, 1, latent_len]`. Convention: 1 = keep + the reference, 0 = regenerate. + + `mask` may be a `[samples]` tensor at the model sample rate or a + `(keep_seconds, total_seconds)` tuple — the tuple form builds + "keep first K seconds, regenerate the rest" automatically. + """ + if isinstance(mask, tuple) and len(mask) == 2: + keep_s, total_s = (float(x) for x in mask) + keep_n = int(keep_s * self.sample_rate) + total_n = int(total_s * self.sample_rate) + mask = torch.zeros(total_n, dtype=torch.float32) + mask[:keep_n] = 1.0 + m = mask.to(device=device, dtype=torch.float32) + if m.dim() == 1: + m = m.unsqueeze(0) + cur_len = m.shape[-1] + if cur_len < self.sample_size: + m = F.pad(m, (0, self.sample_size - cur_len)) + elif cur_len > self.sample_size: + m = m[..., :self.sample_size] + return F.interpolate(m.unsqueeze(1), size=latent_len, mode="nearest") + + +def _decode_audio_file(path, target_sr: int) -> torch.Tensor: + """Decode any audio-bearing file (wav, mp3, mp4, m4a, flac, ...) via + PyAV and resample to `target_sr`. Returns `[channels, samples]` + float32 in roughly [-1, 1]. + + PyAV is already a FastVideo dep (used for muxing in + `VideoGenerator._mux_audio`). `torchaudio.load` on container + formats (mp4 / m4a) routes through `torchcodec`, which pulls in a + full CUDA NVRTC stack we don't otherwise need. + """ + import av + import numpy as np + container = av.open(str(path)) + audio_stream = next(s for s in container.streams if s.type == "audio") + resampler = av.AudioResampler(format="fltp", layout="stereo", rate=target_sr) + chunks: list = [] + for frame in container.decode(audio_stream): + for resampled in resampler.resample(frame): + chunks.append(resampled.to_ndarray()) + for resampled in resampler.resample(None): + chunks.append(resampled.to_ndarray()) + container.close() + if not chunks: + raise RuntimeError(f"No audio frames decoded from {path}") + waveform = np.concatenate(chunks, axis=-1) + if waveform.ndim == 1: + waveform = waveform[None, :] + return torch.from_numpy(waveform).float() diff --git a/fastvideo/pipelines/pipeline_batch_info.py b/fastvideo/pipelines/pipeline_batch_info.py index 18c243dd5..9a0b2479f 100644 --- a/fastvideo/pipelines/pipeline_batch_info.py +++ b/fastvideo/pipelines/pipeline_batch_info.py @@ -192,6 +192,21 @@ class ForwardBatch: ltx2_stg_blocks_video: list[int] = field(default_factory=list) ltx2_stg_blocks_audio: list[int] = field(default_factory=list) + # Stable Audio (T2A): clip start/end in seconds. Parallels the + # `SamplingParam` fields of the same name; the + # `StableAudioConditioningStage` / `DecodingStage` read them. + audio_start_in_s: float | None = None + audio_end_in_s: float | None = None + + # Stable Audio A2A variation + inpainting payloads (parallel to + # `SamplingParam`). `Any` because we accept torch tensors or numpy + # arrays the user supplies; the latent-prep stage normalises shapes. + init_audio: Any = None + init_audio_strength: float | None = None + init_noise_level: float | None = None + inpaint_audio: Any = None + inpaint_mask: Any = None + n_tokens: int | None = None # Other parameters that may be needed by specific schedulers diff --git a/fastvideo/registry.py b/fastvideo/registry.py index 66c86820a..9fc4607e4 100644 --- a/fastvideo/registry.py +++ b/fastvideo/registry.py @@ -48,6 +48,7 @@ from fastvideo.configs.pipelines.wan import ( WanT2V720PConfig, ) from fastvideo.configs.pipelines.sd35 import SD35Config +from fastvideo.configs.pipelines.stable_audio import (StableAudioOpenSmallConfig, StableAudioT2AConfig) from fastvideo.api.sampling_param import SamplingParam from fastvideo.fastvideo_args import WorkloadType @@ -242,6 +243,47 @@ def _register_configs() -> None: default_preset="ltx2_distilled", ) + # Stable Audio Open (text-to-audio). Both variants must be loaded + # from the FastVideo-curated converted Diffusers-format repos — + # the upstream `stabilityai/stable-audio-open-{1.0,small}` repos + # ship `model.safetensors` as a single monolithic checkpoint with + # no per-component subfolders our standard loader can consume. See + # `scripts/checkpoint_conversion/stable_audio_to_diffusers.py`. + # NOTE: WorkloadType has no T2A variant yet (REVIEW item 28); using + # T2V as the placeholder until the enum is extended. + register_configs( + sampling_param_cls=None, + pipeline_config_cls=StableAudioT2AConfig, + workload_types=(WorkloadType.T2V, ), + hf_model_paths=[ + "FastVideo/stable-audio-open-1.0-Diffusers", + ], + # Substring match against HF cache snapshot paths (the lookup + # runs on the resolved local directory, which uses `--` between + # org and repo: `models--FastVideo--stable-audio-open-1.0-Diffusers`). + model_detectors=[ + lambda path: "stable-audio-open-1" in path.lower(), + ], + model_family="stable_audio", + default_preset="stable_audio_open_1_0_base", + ) + # Small variant uses its own `pipeline_config_cls` so it picks up + # the smaller (524288-sample) training window in `sample_size` / + # `max_audio_duration_s`. + register_configs( + sampling_param_cls=None, + pipeline_config_cls=StableAudioOpenSmallConfig, + workload_types=(WorkloadType.T2V, ), + hf_model_paths=[ + "FastVideo/stable-audio-open-small-Diffusers", + ], + model_detectors=[ + lambda path: "stable-audio-open-small" in path.lower(), + ], + model_family="stable_audio", + default_preset="stable_audio_open_small", + ) + # Hunyuan 1.5 (specific) register_configs( sampling_param_cls=None, @@ -786,6 +828,8 @@ def _register_presets() -> None: ALL_PRESETS as MATRIXGAME_PRESETS, ) from fastvideo.pipelines.basic.sd35.presets import ( ALL_PRESETS as SD35_PRESETS, ) + from fastvideo.pipelines.basic.stable_audio.presets import ( + ALL_PRESETS as STABLE_AUDIO_PRESETS, ) from fastvideo.pipelines.basic.turbodiffusion.presets import ( ALL_PRESETS as TURBODIFFUSION_PRESETS, ) from fastvideo.pipelines.basic.wan.presets import ( @@ -803,6 +847,7 @@ def _register_presets() -> None: LTX2_PRESETS, MATRIXGAME_PRESETS, SD35_PRESETS, + STABLE_AUDIO_PRESETS, TURBODIFFUSION_PRESETS, WAN_PRESETS, ) diff --git a/fastvideo/utils.py b/fastvideo/utils.py index d7094c7b6..884bd8ca9 100644 --- a/fastvideo/utils.py +++ b/fastvideo/utils.py @@ -636,6 +636,23 @@ def maybe_download_model_index(model_name_or_path: str) -> dict[str, Any]: raise ValueError(f"Failed to download or parse model_index.json for {model_name_or_path}: {e}") from e +_HF_TOKEN_ENV_VARS = ("HF_TOKEN", "HUGGINGFACE_HUB_TOKEN", "HF_API_KEY") + + +def resolve_hf_token() -> str | None: + """Return the first non-empty HF token from the standard env vars. + + Order: `HF_TOKEN`, `HUGGINGFACE_HUB_TOKEN`, `HF_API_KEY` (the last is + a FastVideo convention; `huggingface_hub` itself doesn't read it). + Does not mutate `os.environ`. + """ + for src in _HF_TOKEN_ENV_VARS: + v = os.environ.get(src) + if v: + return v + return None + + def update_environment_variables(envs: dict[str, str]): for k, v in envs.items(): if k in os.environ and os.environ[k] != v: diff --git a/scripts/checkpoint_conversion/stable_audio_to_diffusers.py b/scripts/checkpoint_conversion/stable_audio_to_diffusers.py new file mode 100644 index 000000000..eee392299 --- /dev/null +++ b/scripts/checkpoint_conversion/stable_audio_to_diffusers.py @@ -0,0 +1,295 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Convert a `stable_audio_tools` raw checkpoint into a Diffusers-style +per-component layout consumable by FastVideo's standard +`ComposedPipelineBase` loader. + +Why: the published Stability AI Stable Audio Open repos ship a single +monolithic `model.safetensors` containing DiT (`model.model.*`), +VAE / pretransform (`pretransform.model.*`), and NumberConditioner +(`conditioner.conditioners.*`) state dicts under shared prefixes. +FastVideo's standard loader expects the Diffusers per-subfolder layout +(`model_index.json` + per-component subfolders). This script splits the +monolithic checkpoint, drops the host-pipeline prefixes, normalizes +LayerNorm `gamma`/`beta` keys to `weight`/`bias`, and writes a +Diffusers-style tree. + +Usage:: + + python scripts/checkpoint_conversion/stable_audio_to_diffusers.py \ + --src stabilityai/stable-audio-open-1.0 \ + --dst converted_weights/stable_audio_open_1_0_diffusers + + # Or with a local snapshot directory: + python scripts/checkpoint_conversion/stable_audio_to_diffusers.py \ + --src /path/to/stable-audio-open-1.0 \ + --dst converted_weights/stable_audio_open_small_diffusers + +The output tree is:: + + / + ├── model_index.json + ├── transformer/{config.json, diffusion_pytorch_model.safetensors} + ├── vae/{config.json, diffusion_pytorch_model.safetensors} + ├── conditioner/{config.json, diffusion_pytorch_model.safetensors} + ├── text_encoder/ (copied from src if present) + ├── tokenizer/ (copied from src if present) + └── scheduler/ (copied from src if present) + +Push to the FastVideo HF org with `huggingface-cli upload`. Confirm +visibility (`--private` for staging, omit for public release):: + + huggingface-cli upload --private FastVideo/stable-audio-open-1.0-Diffusers \ + converted_weights/stable-audio-open-1.0-Diffusers + huggingface-cli upload --private FastVideo/stable-audio-open-small-Diffusers \ + converted_weights/stable-audio-open-small-Diffusers +""" +from __future__ import annotations + +import argparse +import json +import os +import shutil +from pathlib import Path +from typing import Any + +from huggingface_hub import snapshot_download +from safetensors import safe_open +from safetensors.torch import save_file + +# Map host-pipeline prefix → output subfolder. Order matters only for +# logging. +_COMPONENT_PREFIXES: dict[str, str] = { + "model.model.": "transformer", + "pretransform.model.": "vae", + "conditioner.": "conditioner", +} + +# Subfolders we copy verbatim from the source repo. The `vae/` subfolder +# in `stabilityai/stable-audio-open-1.0` already ships in Diffusers +# format with the key naming our `OobleckVAE` expects (`decoder.block.X. +# conv_t1.weight_g` etc.) — preferred over the raw `pretransform.model.*` +# extraction below, which uses a different module nesting (`decoder. +# layers.X.layers.Y.*`) that would need a separate remap. +_PASSTHROUGH_SUBFOLDERS = ("text_encoder", "tokenizer", "scheduler", "vae") + +_MODEL_INDEX = { + "_class_name": "StableAudioPipeline", + # `ComposedPipelineBase._load_config` requires this field; the value + # is informational (we don't pin a Diffusers version). + "_diffusers_version": "0.30.0", + "_fastvideo_converted_from": None, # filled in at write time + "transformer": ["fastvideo.models.dits.stable_audio", "StableAudioDiT"], + "vae": ["fastvideo.models.vaes.oobleck", "OobleckVAE"], + "conditioner": [ + "fastvideo.models.encoders.stable_audio_conditioner", + "StableAudioMultiConditioner", + ], + "text_encoder": ["transformers", "T5EncoderModel"], + "tokenizer": ["transformers", "T5TokenizerFast"], + "scheduler": ["diffusers", "CosineDPMSolverMultistepScheduler"], +} + + +def _resolve_src(src: str) -> Path: + """Accept either a HF repo id or a local directory.""" + if os.path.isdir(src): + return Path(src) + from fastvideo.utils import resolve_hf_token + return Path(snapshot_download(repo_id=src, token=resolve_hf_token())) + + +def _rename_layernorm_keys(state: dict[str, Any]) -> dict[str, Any]: + """`nn.LayerNorm` ships `gamma`/`beta` in the official checkpoint; + torch's `nn.LayerNorm` uses `weight`/`bias`.""" + out: dict[str, Any] = {} + for k, v in state.items(): + if k.endswith(".gamma"): + out[k[: -len(".gamma")] + ".weight"] = v + elif k.endswith(".beta"): + out[k[: -len(".beta")] + ".bias"] = v + else: + out[k] = v + return out + + +def _split_state_dict(src_safetensors: Path) -> dict[str, dict[str, Any]]: + """Read `model.safetensors` and bucket its keys by the host-pipeline + prefix listed in `_COMPONENT_PREFIXES`.""" + buckets: dict[str, dict[str, Any]] = {name: {} for name in _COMPONENT_PREFIXES.values()} + skipped = 0 + with safe_open(str(src_safetensors), framework="pt") as f: + for k in f.keys(): + for prefix, subfolder in _COMPONENT_PREFIXES.items(): + if k.startswith(prefix): + buckets[subfolder][k[len(prefix):]] = f.get_tensor(k) + break + else: + skipped += 1 + if skipped: + print(f" skipped {skipped} keys (no matching prefix)") + for subfolder, state in buckets.items(): + if subfolder == "transformer": + buckets[subfolder] = _rename_layernorm_keys(state) + print(f" {subfolder}: {len(buckets[subfolder])} keys") + return buckets + + +def _detect_projection_flags( + diff_cfg: dict[str, Any], + dit_state: dict[str, Any]) -> tuple[bool, bool]: + """Decide `project_cond_tokens` / `project_global_cond` by comparing + the actual `to_cond_embed.0` / `to_global_embed.0` weight shapes + in the DiT state dict against `cond_token_dim` / `global_cond_dim`. + Upstream's factory toggles these per-variant in ways that aren't + derivable from `embed_dim` alone (SA-1.0 keeps `cond_embed_dim` = + `cond_token_dim` and uses GQA in cross-attn; SA-small projects to + `embed_dim` and uses MHA). + """ + cond_dim = diff_cfg.get("cond_token_dim") + glob_dim = diff_cfg.get("global_cond_dim") + cond_w = dit_state.get("to_cond_embed.0.weight") + glob_w = dit_state.get("to_global_embed.0.weight") + project_cond = (cond_w is not None and cond_w.shape[0] != cond_dim) + project_glob = (glob_w is not None and glob_w.shape[0] != glob_dim) + return project_cond, project_glob + + +def _component_config(model_config: dict[str, Any]) -> dict[str, dict[str, Any]]: + """Pull per-component sizing out of the official `model_config.json` + so the converted repo records the architecture authoritatively.""" + diff = model_config["model"]["diffusion"] + pre = model_config["model"]["pretransform"]["config"] + cond = model_config["model"]["conditioning"] + # Translate upstream `model_config.json` field names to FastVideo's + # `StableAudioArchConfig` fields. The standard `TransformerLoader` + # calls `dit_config.update_model_arch(json_dict)` which raises on + # unknown keys. + diff_cfg = dict(diff["config"]) + if "num_heads" in diff_cfg: + diff_cfg["num_attention_heads"] = diff_cfg.pop("num_heads") + diff_cfg.pop("transformer_type", None) # always "continuous_transformer" + # `attn_kwargs.qk_norm` (small variant) → top-level `qk_norm`. + attn_kwargs = diff_cfg.pop("attn_kwargs", None) or {} + if "qk_norm" in attn_kwargs: + diff_cfg["qk_norm"] = attn_kwargs["qk_norm"] + # `cross_attention_cond_ids` / `global_cond_ids` belong to the + # conditioner, not the DiT — keep them out of transformer/config.json. + transformer_cfg = { + "_class_name": "StableAudioDiT", + **diff_cfg, + } + # Mirror Diffusers' `AutoencoderOobleck` config field naming so + # `OobleckVAEArchConfig.update_model_arch` (called by `VAELoader`) + # accepts the keys. + enc_cfg = pre["encoder"]["config"] + dec_cfg = pre["decoder"]["config"] + vae_cfg = { + "_class_name": "AutoencoderOobleck", + "encoder_hidden_size": enc_cfg["channels"], + "downsampling_ratios": enc_cfg["strides"], + "channel_multiples": enc_cfg["c_mults"], + "decoder_channels": dec_cfg["channels"], + "decoder_input_channels": dec_cfg["latent_dim"], + "audio_channels": pre.get("io_channels", 2), + "sampling_rate": model_config.get("sample_rate", 44100), + } + conditioner_cfg = { + "_class_name": "StableAudioMultiConditioner", + "cond_dim": cond["cond_dim"], + "configs": cond["configs"], + "cross_attention_cond_ids": diff["cross_attention_cond_ids"], + "global_cond_ids": diff["global_cond_ids"], + } + return { + "transformer": transformer_cfg, + "vae": vae_cfg, + "conditioner": conditioner_cfg, + } + + +def _write_component(dst_dir: Path, name: str, state: dict[str, Any], config: dict[str, Any]) -> None: + sub = dst_dir / name + sub.mkdir(parents=True, exist_ok=True) + save_file(state, str(sub / "diffusion_pytorch_model.safetensors")) + with open(sub / "config.json", "w") as f: + json.dump(config, f, indent=2) + print(f" wrote {sub}/") + + +def _copy_passthrough(src_dir: Path, dst_dir: Path) -> list[str]: + copied: list[str] = [] + for sub in _PASSTHROUGH_SUBFOLDERS: + s = src_dir / sub + if not s.is_dir(): + continue + d = dst_dir / sub + if d.exists(): + shutil.rmtree(d) + shutil.copytree(s, d, ignore=shutil.ignore_patterns("*.bin", "*.ckpt")) + copied.append(sub) + print(f" copied {sub}/") + return copied + + +def convert(src: str, dst: str) -> None: + src_dir = _resolve_src(src) + dst_dir = Path(dst) + dst_dir.mkdir(parents=True, exist_ok=True) + print(f"src: {src_dir}") + print(f"dst: {dst_dir}") + + monolithic = src_dir / "model.safetensors" + if not monolithic.is_file(): + raise FileNotFoundError(f"Expected monolithic safetensors at {monolithic}") + cfg_path = src_dir / "model_config.json" + if not cfg_path.is_file(): + raise FileNotFoundError( + f"Expected model_config.json at {cfg_path} (Stable Audio's authoritative " + f"per-component arch config).") + with open(cfg_path) as f: + model_config = json.load(f) + + print("\n[1/4] Splitting monolithic state dict by component prefix:") + buckets = _split_state_dict(monolithic) + + print("\n[2/4] Copying passthrough subfolders (text_encoder/tokenizer/scheduler/vae):") + copied = _copy_passthrough(src_dir, dst_dir) + + print("\n[3/4] Writing per-component configs + safetensors:") + component_cfgs = _component_config(model_config) + project_cond, project_glob = _detect_projection_flags(component_cfgs["transformer"], + buckets["transformer"]) + component_cfgs["transformer"]["project_cond_tokens"] = project_cond + component_cfgs["transformer"]["project_global_cond"] = project_glob + for name in ("transformer", "vae", "conditioner"): + if name in copied: + print(f" skipped {name}/ (already copied from source in Diffusers format)") + continue + _write_component(dst_dir, name, buckets[name], component_cfgs[name]) + + print("\n[4/4] Writing model_index.json:") + index = dict(_MODEL_INDEX) + # Prefer the HF repo id over a snapshot directory (which is + # machine-local and full of hash gunk). + index["_fastvideo_converted_from"] = src if "/" in src and not os.path.isdir(src) else src_dir.name + # Drop entries for subfolders we didn't actually populate. + available = {"transformer", "vae", "conditioner", *copied} + index = {k: v for k, v in index.items() if not k or k.startswith("_") or k in available} + with open(dst_dir / "model_index.json", "w") as f: + json.dump(index, f, indent=2) + print(f" wrote {dst_dir}/model_index.json") + print(f"\nDone. Push to HF with:") + print(f" huggingface-cli upload {dst_dir}") + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__.split("\n\n")[0]) + parser.add_argument("--src", required=True, + help="HF repo id or local directory containing model.safetensors + model_config.json") + parser.add_argument("--dst", required=True, help="Output directory for the Diffusers-format repo") + args = parser.parse_args() + convert(args.src, args.dst) + + +if __name__ == "__main__": + main() diff --git a/tests/local_tests/_stable_audio_helpers.py b/tests/local_tests/_stable_audio_helpers.py new file mode 100644 index 000000000..6d71347fe --- /dev/null +++ b/tests/local_tests/_stable_audio_helpers.py @@ -0,0 +1,40 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Shared helpers for Stable Audio local parity tests. + +Centralizes the HF-token lookup, env-mirroring, and gated-repo access +probe that all parity tests in this directory need; uses +`fastvideo.utils.resolve_hf_token` as the source of truth for env-var +precedence so the tests stay in sync with the rest of the codebase. +""" +from __future__ import annotations + +import os + +from fastvideo.utils import resolve_hf_token + + +def setup_hf_env() -> None: + """Mirror the resolved token to both `HF_TOKEN` and + `HUGGINGFACE_HUB_TOKEN` so downstream libraries see it under + whichever name they happen to check. + """ + token = resolve_hf_token() + if token is None: + return + os.environ.setdefault("HF_TOKEN", token) + os.environ.setdefault("HUGGINGFACE_HUB_TOKEN", token) + + +def can_access_repo(repo_id: str, filename: str = "model_index.json") -> bool: + """True iff the resolved HF token grants access to `filename` in + `repo_id`. Cheapest possible probe (single small file fetch). + """ + token = resolve_hf_token() + if token is None: + return False + try: + from huggingface_hub import hf_hub_download + hf_hub_download(repo_id=repo_id, filename=filename, token=token) + return True + except Exception: + return False diff --git a/tests/local_tests/pipelines/test_stable_audio_a2a_parity.py b/tests/local_tests/pipelines/test_stable_audio_a2a_parity.py new file mode 100644 index 000000000..ffb1f17a5 --- /dev/null +++ b/tests/local_tests/pipelines/test_stable_audio_a2a_parity.py @@ -0,0 +1,173 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Parity test: FastVideo audio-to-audio variation vs the official +`Stability-AI/stable-audio-tools.generate_diffusion_cond(init_audio=..., +init_noise_level=...)`. + +Same harness as `test_stable_audio_pipeline_parity.py` — drives both +sides with the same prompt / seed / steps / CFG / init reference, then +compares decoded waveforms. +""" +from __future__ import annotations + +import json + +import pytest +import torch + +from fastvideo.utils import resolve_hf_token +from tests.local_tests._stable_audio_helpers import ( + can_access_repo, + setup_hf_env, +) + +_HF_REPO_ID = "stabilityai/stable-audio-open-1.0" # raw upstream weights for the official side +_FV_REPO_ID = "FastVideo/stable-audio-open-1.0-Diffusers" # converted Diffusers layout for FastVideo's loader +_MODEL_CFG = "model_config.json" +_MODEL_WEIGHTS = "model.safetensors" + + +def _stable_audio_tools_inference_available() -> bool: + try: + from stable_audio_tools.inference.generation import generate_diffusion_cond # noqa + from stable_audio_tools.models.factory import create_model_from_config # noqa + return True + except Exception: + return False + + +def _load_official(device): + from huggingface_hub import hf_hub_download + from safetensors.torch import load_file + from stable_audio_tools.models.factory import create_model_from_config + + cfg_path = hf_hub_download(repo_id=_HF_REPO_ID, filename=_MODEL_CFG, token=resolve_hf_token()) + weights_path = hf_hub_download(repo_id=_HF_REPO_ID, filename=_MODEL_WEIGHTS, + token=resolve_hf_token()) + with open(cfg_path) as f: + model_config = json.load(f) + model = create_model_from_config(model_config) + state = load_file(weights_path) + missing, unexpected = model.load_state_dict(state, strict=True) + assert not missing and not unexpected + return model.to(device).eval(), model_config + + +def _make_init_audio(seed: int, sample_rate: int, seconds: float, channels: int = 2): + """Deterministic synthetic reference: a stereo sine sweep.""" + g = torch.Generator(device="cpu").manual_seed(seed) + n = int(seconds * sample_rate) + t = torch.linspace(0, n / sample_rate, n) + freqs = torch.linspace(220.0, 880.0, n) + base = torch.sin(2 * torch.pi * freqs * t) * 0.3 + # Add a touch of noise so the encoder sees nontrivial information. + base = base + torch.randn(n, generator=g) * 0.01 + return base.unsqueeze(0).repeat(channels, 1).unsqueeze(0).contiguous() + + +@pytest.mark.skipif(not torch.cuda.is_available(), + reason="Stable Audio A2A parity requires CUDA.") +@pytest.mark.skipif(not can_access_repo(_HF_REPO_ID), + reason=f"{_HF_REPO_ID} not accessible — gated.") +@pytest.mark.skipif(not _stable_audio_tools_inference_available(), + reason="`stable_audio_tools.inference` not importable.") +def test_stable_audio_a2a_official_parity(): + setup_hf_env() + device = torch.device("cuda:0") + + prompt = "A gentle ambient pad with soft synth swells." + seed = 0 + num_inference_steps = 25 + guidance_scale = 7.0 + audio_end_in_s = 1.5 + init_noise_level = 1.0 + + off_model, off_cfg = _load_official(device) + sample_rate = int(off_cfg["sample_rate"]) + full_sample_size = int(off_cfg["sample_size"]) + + init_audio = _make_init_audio(seed=seed, sample_rate=sample_rate, + seconds=audio_end_in_s) + + cond = [{"prompt": prompt, "seconds_start": 0.0, "seconds_total": audio_end_in_s}] + + from stable_audio_tools.inference.generation import generate_diffusion_cond + audio_full = generate_diffusion_cond( + off_model, + steps=num_inference_steps, + cfg_scale=guidance_scale, + conditioning=cond, + sample_size=full_sample_size, + seed=seed, + device="cuda", + sampler_type="dpmpp-3m-sde", + sigma_min=0.3, + sigma_max=500, + init_audio=(sample_rate, init_audio.squeeze(0)), + init_noise_level=init_noise_level, + ) + end_idx = int(audio_end_in_s * sample_rate) + off_audio = audio_full.detach().float().cpu()[:, :, :end_idx] + print(f"off shape={tuple(off_audio.shape)} abs_mean={off_audio.abs().mean():.6f}") + + del off_model + import gc + gc.collect() + torch.cuda.empty_cache() + + from fastvideo import VideoGenerator + generator = VideoGenerator.from_pretrained( + _FV_REPO_ID, + num_gpus=1, + use_fsdp_inference=False, + dit_cpu_offload=False, + vae_cpu_offload=False, + text_encoder_cpu_offload=False, + ) + try: + result = generator.generate_video( + prompt=prompt, + output_path="outputs_audio/sa_a2a_parity.mp4", + save_video=False, + num_inference_steps=num_inference_steps, + guidance_scale=guidance_scale, + seed=seed, + audio_end_in_s=audio_end_in_s, + init_audio=init_audio, + init_noise_level=init_noise_level, + ) + finally: + generator.shutdown() + + fv_audio = result.get("audio") + if fv_audio is None: + ext = result.get("extra", {}) or {} + fv_audio = ext.get("decoded_audio") + assert fv_audio is not None + if not torch.is_tensor(fv_audio): + import numpy as np + fv_audio = torch.from_numpy(np.asarray(fv_audio)) + fv_audio = fv_audio.detach().float().cpu() + if fv_audio.ndim == 2 and fv_audio.shape[0] != 1: + fv_audio = fv_audio.T.unsqueeze(0) + print(f"fv shape={tuple(fv_audio.shape)} abs_mean={fv_audio.abs().mean():.6f}") + + assert fv_audio.shape == off_audio.shape + diff = (fv_audio - off_audio).abs() + diff_max = diff.max().item() + diff_mean = diff.mean().item() + diff_median = diff.median().item() + rel = abs(off_audio.abs().mean() - fv_audio.abs().mean()) / max(off_audio.abs().mean().item(), + 1e-6) + print(f"diff max={diff_max:.6f} mean={diff_mean:.6f} median={diff_median:.6f} " + f"abs_mean rel drift={rel:.4%}") + + # A2A is more sensitive than T2A — the encode-side + # `posterior.sample()` adds per-element randn noise, so even tiny + # RNG-ordering differences propagate through the SDE trajectory. + # Empirically (after hoisting the CFG batching + null_embed out of + # the sampler loop): drift ~0.5%, mean diff ~0.003. + assert rel < 0.02, f"A2A abs_mean drift {rel:.2%} > 2%" + assert diff_mean < 0.05, f"A2A mean diff {diff_mean:.4f} > 0.05" + # Both waveforms should stay in the audio range (no NaN, no blow-up). + assert torch.isfinite(fv_audio).all(), "FV A2A produced non-finite values" + assert fv_audio.abs().max().item() < 5.0, "FV A2A magnitude blew up" diff --git a/tests/local_tests/pipelines/test_stable_audio_inpaint_parity.py b/tests/local_tests/pipelines/test_stable_audio_inpaint_parity.py new file mode 100644 index 000000000..22e0e9158 --- /dev/null +++ b/tests/local_tests/pipelines/test_stable_audio_inpaint_parity.py @@ -0,0 +1,129 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Inpainting / outpainting test for FastVideo's RePaint blending path. + +Stable Audio Open 1.0 is `model_type=diffusion_cond`, **not** +`diffusion_cond_inpaint` — there is no inpaint-trained checkpoint for +this model, so an apples-to-apples parity comparison against +`generate_diffusion_cond_inpaint` would be testing the wrong thing +(that function expects a model whose conditioner accepts `inpaint_mask` +and `inpaint_masked_input`, which our model's conditioner does not). + +Instead this test validates the two invariants that any RePaint-style +blender must hold: + + 1. **Kept-region preservation** — for every sample where mask == 1, + the decoded output must round-trip back to the reference within + fp32 VAE round-trip noise (~RMS ratio close to 1, max diff + bounded). This is the core promise of mask=1: "leave this alone". + 2. **Unkept-region freshness** — for samples where mask == 0, the + output must be different from the reference (model actually + regenerated something). RMS in the unkept region should be in + a sane audio range (not silent, not blown up). + +Skips when CUDA / HF access is unavailable. +""" +from __future__ import annotations + +import pytest +import torch + +from tests.local_tests._stable_audio_helpers import ( + can_access_repo, + setup_hf_env, +) + +_HF_REPO_ID = "stabilityai/stable-audio-open-1.0" # raw upstream — only used to gate the test on access +_FV_REPO_ID = "FastVideo/stable-audio-open-1.0-Diffusers" # converted Diffusers repo FastVideo loads from + + +def _make_reference(seed: int, sample_rate: int, seconds: float): + """Stereo amplitude-modulated noise — fakes a 'beat'.""" + g = torch.Generator(device="cpu").manual_seed(seed) + n = int(seconds * sample_rate) + base = torch.randn(2, n, generator=g) * 0.1 + env = 0.5 + 0.5 * torch.sin(2 * torch.pi * 4 * torch.linspace(0, seconds, n)) + return (base * env).unsqueeze(0).contiguous() + + +@pytest.mark.skipif(not torch.cuda.is_available(), + reason="Stable Audio inpainting test requires CUDA.") +@pytest.mark.skipif(not can_access_repo(_HF_REPO_ID), + reason=f"{_HF_REPO_ID} not accessible — gated.") +def test_stable_audio_inpaint_kept_region_preserved(): + """Run inpainting with the first 1.5s kept and the next 4.5s + regenerated. Verify the kept region matches the reference and the + regenerated region is genuinely different. + """ + setup_hf_env() + sample_rate = 44100 + keep_seconds = 1.5 + total_seconds = 6.0 + keep_samples = int(keep_seconds * sample_rate) + total_samples = int(total_seconds * sample_rate) + + ref_short = _make_reference(seed=7, sample_rate=sample_rate, seconds=keep_seconds) + padded_ref = torch.zeros((1, ref_short.shape[1], total_samples), dtype=torch.float32) + padded_ref[..., :keep_samples] = ref_short + + mask = torch.zeros(total_samples, dtype=torch.float32) + mask[:keep_samples] = 1.0 + + from fastvideo import VideoGenerator + generator = VideoGenerator.from_pretrained( + _FV_REPO_ID, + num_gpus=1, + use_fsdp_inference=False, + dit_cpu_offload=False, + vae_cpu_offload=False, + text_encoder_cpu_offload=False, + ) + try: + result = generator.generate_video( + prompt="Steady lo-fi hip hop drum loop with vinyl crackle.", + output_path="outputs_audio/sa_inpaint_self.mp4", + save_video=False, + num_inference_steps=25, + guidance_scale=7.0, + seed=0, + audio_end_in_s=total_seconds, + inpaint_audio=padded_ref, + inpaint_mask=mask, + ) + finally: + generator.shutdown() + + fv_audio = result.get("audio") + if fv_audio is None: + fv_audio = (result.get("extra", {}) or {}).get("decoded_audio") + if not torch.is_tensor(fv_audio): + import numpy as np + fv_audio = torch.from_numpy(np.asarray(fv_audio)) + out = fv_audio.detach().float().cpu() + if out.ndim == 2 and out.shape[0] != 1: + out = out.T.unsqueeze(0) + print(f"out shape={tuple(out.shape)} abs_mean={out.abs().mean():.6f}") + assert out.shape[-1] >= keep_samples + + ref_kept = padded_ref[..., :keep_samples].float().cpu() + out_kept = out[..., :keep_samples] + out_unkept = out[..., keep_samples:] + + # 1. Kept region: should round-trip to within VAE round-trip noise. + # Round-trip RMS ratio for Oobleck on synthetic audio runs ~1.0–1.05 + # in the existing `test_oobleck_vae_round_trip_sanity` test, so we + # bound the per-sample-RMS ratio inside [0.5, 2.0] to confirm the + # kept region was actually preserved (and not, e.g., zeroed out). + ref_rms = ref_kept.float().pow(2).mean().sqrt().item() + out_kept_rms = out_kept.float().pow(2).mean().sqrt().item() + rms_ratio = out_kept_rms / max(ref_rms, 1e-6) + print(f"kept rms ratio={rms_ratio:.3f} ref_rms={ref_rms:.4f} out_rms={out_kept_rms:.4f}") + assert 0.5 < rms_ratio < 2.0, ( + f"kept-region RMS ratio {rms_ratio:.3f} outside [0.5, 2.0] — " + "RePaint blending did not preserve the reference") + + # 2. Unkept region: should be live audio (finite, in-range, non-silent). + assert torch.isfinite(out_unkept).all(), "unkept region produced non-finite values" + unkept_rms = out_unkept.float().pow(2).mean().sqrt().item() + print(f"unkept rms={unkept_rms:.4f}") + assert unkept_rms > 1e-3, "unkept region is silent — model produced nothing" + assert out_unkept.abs().max().item() < 5.0, "unkept region magnitude blew up" diff --git a/tests/local_tests/pipelines/test_stable_audio_pipeline_parity.py b/tests/local_tests/pipelines/test_stable_audio_pipeline_parity.py new file mode 100644 index 000000000..ff585b05d --- /dev/null +++ b/tests/local_tests/pipelines/test_stable_audio_pipeline_parity.py @@ -0,0 +1,228 @@ +# SPDX-License-Identifier: Apache-2.0 +"""End-to-end pipeline parity test for Stable Audio Open 1.0. + +Per the add-model skill (steps 5 + 13(a)), the canonical parity +reference is the **official upstream repo** (`Stability-AI/stable-audio-tools`). +This test runs `stable_audio_tools.inference.generation.generate_diffusion_cond` +on the published model + checkpoint and compares the resulting waveform +against the FastVideo `StableAudioPipeline` output for the same prompt / +seed / steps / CFG. + +All FV components are now first-class ports (REVIEW item 30 closed): + * **VAE**: bit-identical (FV's `OobleckVAE` is a 1:1 port of + upstream's `OobleckEncoder/OobleckDecoder` — verified separately + in `tests/local_tests/vaes/test_oobleck_vae_parity.py` to diff=0). + * **DiT**: native `fastvideo.models.dits.stable_audio.StableAudioDiT` + (vendored from `stable_audio_tools/models/dit.py + transformer.py`). + Forward parity vs upstream `DiffusionTransformer` is bit-identical + (diff=0) on shared random latents. + * **Conditioner**: native `StableAudioMultiConditioner` (vendored + from `stable_audio_tools/models/conditioners.py`). + * **Sampler**: `k_diffusion.sampling.sample_dpmpp_3m_sde` — pure + sampling library, the same function `generate_diffusion_cond` + calls in upstream. + * **Stage orchestration**: FV's stage split vs upstream's monolithic + `generate_diffusion_cond`. + +Empirically the end-to-end drift is ~0.015% abs_mean / max element +~0.006 on a 25-step CFG=7 run — fp32 numerical noise from non- +deterministic CUDA kernels (matmul reorder), not orchestration. + +Skips when: + * CUDA is unavailable. + * `stabilityai/stable-audio-open-1.0` is inaccessible (gated). + * `stable_audio_tools` and its deps are not installed. +""" +from __future__ import annotations + +import json + +import pytest +import torch + +from fastvideo.utils import resolve_hf_token +from tests.local_tests._stable_audio_helpers import ( + can_access_repo, + setup_hf_env, +) + +_HF_REPO_ID = "stabilityai/stable-audio-open-1.0" # raw upstream weights for the official side +_FV_REPO_ID = "FastVideo/stable-audio-open-1.0-Diffusers" # converted Diffusers layout for FastVideo's loader +_MODEL_CFG = "model_config.json" +_MODEL_WEIGHTS = "model.safetensors" + + +def _stable_audio_tools_inference_available() -> bool: + try: + from stable_audio_tools.inference.generation import ( # noqa: F401 + generate_diffusion_cond, + ) + from stable_audio_tools.models.factory import ( # noqa: F401 + create_model_from_config, + ) + return True + except Exception: + return False + + +def _load_official_diffusion_cond(device: torch.device): + """Build the official `ConditionedDiffusionModelWrapper` from + `model_config.json` and load its weights from `model.safetensors`. + """ + from huggingface_hub import hf_hub_download + from safetensors.torch import load_file + + cfg_path = hf_hub_download(repo_id=_HF_REPO_ID, filename=_MODEL_CFG, + token=resolve_hf_token()) + weights_path = hf_hub_download(repo_id=_HF_REPO_ID, filename=_MODEL_WEIGHTS, + token=resolve_hf_token()) + with open(cfg_path) as f: + model_config = json.load(f) + + from stable_audio_tools.models.factory import create_model_from_config + model = create_model_from_config(model_config) + state = load_file(weights_path) + missing, unexpected = model.load_state_dict(state, strict=True) + assert not missing and not unexpected, ( + f"official model state mismatch — missing={missing[:3]} unexpected={unexpected[:3]}" + ) + return model.to(device).eval(), model_config + + +@pytest.mark.skipif( + not torch.cuda.is_available(), + reason="Stable Audio pipeline parity requires CUDA.", +) +@pytest.mark.skipif( + not can_access_repo(_HF_REPO_ID), + reason=(f"{_HF_REPO_ID} not accessible — gated repo; set HF_TOKEN / " + f"HF_API_KEY and accept the terms on https://huggingface.co/{_HF_REPO_ID}."), +) +@pytest.mark.skipif( + not _stable_audio_tools_inference_available(), + reason=("`stable_audio_tools.inference.generation` not importable. " + "Clone https://github.com/Stability-AI/stable-audio-tools and " + "`pip install` its deps (k_diffusion, einops_exts, alias_free_torch)."), +) +def test_stable_audio_pipeline_official_parity(): + setup_hf_env() + device = torch.device("cuda:0") + + prompt = "A gentle ambient pad with soft synth swells." + negative_prompt = "low quality, distorted" + seed = 0 + num_inference_steps = 25 + guidance_scale = 7.0 + audio_end_in_s = 1.5 + + # --- Reference: official `generate_diffusion_cond` --- + off_model, off_cfg = _load_official_diffusion_cond(device) + sample_rate = int(off_cfg["sample_rate"]) + # Use the model's training sample_size (full latent length 1024). + # Generating at the full duration matches what FV/diffusers do — + # both produce the full ~47.5s clip then slice. Generating at the + # requested 1.5s would change the latent shape and so the noise + # trajectory, which would not be a meaningful comparison. + full_sample_size = int(off_cfg["sample_size"]) + + cond = [{ + "prompt": prompt, + "seconds_start": 0.0, + "seconds_total": audio_end_in_s, + }] + neg_cond = [{ + "prompt": negative_prompt, + "seconds_start": 0.0, + "seconds_total": audio_end_in_s, + }] + + from stable_audio_tools.inference.generation import generate_diffusion_cond + audio_full = generate_diffusion_cond( + off_model, + steps=num_inference_steps, + cfg_scale=guidance_scale, + conditioning=cond, + negative_conditioning=neg_cond, + sample_size=full_sample_size, + seed=seed, + device="cuda", + sampler_type="dpmpp-3m-sde", + sigma_min=0.3, + sigma_max=500, + ) + end_idx = int(audio_end_in_s * sample_rate) + off_audio = audio_full.detach().float().cpu()[:, :, :end_idx] + print( + f"off shape={tuple(off_audio.shape)} " + f"abs_mean={off_audio.abs().mean().item():.6f} " + f"range=[{off_audio.min().item():.4f}, {off_audio.max().item():.4f}]" + ) + + # Free the reference model so FV has GPU memory. + del off_model + import gc + gc.collect() + torch.cuda.empty_cache() + + # --- FastVideo path --- + from fastvideo import VideoGenerator + generator = VideoGenerator.from_pretrained( + _FV_REPO_ID, + num_gpus=1, + use_fsdp_inference=False, + dit_cpu_offload=False, + vae_cpu_offload=False, + text_encoder_cpu_offload=False, + ) + try: + result = generator.generate_video( + prompt=prompt, + negative_prompt=negative_prompt, + output_path="outputs_audio/stable_audio_parity.mp4", + save_video=False, + num_inference_steps=num_inference_steps, + guidance_scale=guidance_scale, + seed=seed, + audio_end_in_s=audio_end_in_s, + ) + finally: + generator.shutdown() + + fv_audio = result.get("audio") + if fv_audio is None: + ext = result.get("extra", {}) or {} + fv_audio = ext.get("decoded_audio") + assert fv_audio is not None, "FastVideo pipeline did not surface audio" + if not torch.is_tensor(fv_audio): + import numpy as np + fv_audio = torch.from_numpy(np.asarray(fv_audio)) + fv_audio = fv_audio.detach().float().cpu() + if fv_audio.ndim == 2 and fv_audio.shape[0] != 1: + fv_audio = fv_audio.T.unsqueeze(0) + print( + f"fv shape={tuple(fv_audio.shape)} " + f"abs_mean={fv_audio.abs().mean().item():.6f} " + f"range=[{fv_audio.min().item():.4f}, {fv_audio.max().item():.4f}]" + ) + + assert fv_audio.shape == off_audio.shape, ( + f"shape mismatch: off={off_audio.shape} fv={fv_audio.shape}" + ) + + diff = (fv_audio - off_audio).abs() + diff_max = diff.max().item() + diff_mean = diff.mean().item() + diff_median = diff.median().item() + rel = abs(off_audio.abs().mean() - fv_audio.abs().mean()) / max(off_audio.abs().mean().item(), 1e-6) + print( + f"diff max={diff_max:.6f} mean={diff_mean:.6f} median={diff_median:.6f}" + ) + print(f"abs_mean rel drift: {rel:.4%}") + + # Tight bounds (1% drift, 0.05 element-wise) since FV is now a + # first-class port of the upstream — drift is fp32 kernel noise, + # not algorithmic divergence. + assert rel < 0.01, f"abs_mean rel drift {rel:.2%} > 1% — port regression" + assert diff_max < 0.05, ( + f"max element diff {diff_max:.4f} > 0.05 — port regression" + ) diff --git a/tests/local_tests/pipelines/test_stable_audio_pipeline_smoke.py b/tests/local_tests/pipelines/test_stable_audio_pipeline_smoke.py new file mode 100644 index 000000000..06770196f --- /dev/null +++ b/tests/local_tests/pipelines/test_stable_audio_pipeline_smoke.py @@ -0,0 +1,53 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Smoke / preflight tests for the Stable Audio Open 1.0 T2A pipeline.""" +from __future__ import annotations + + +def test_stable_audio_typed_surface_preflight() -> None: + """No-GPU preflight: imports + registry + preset are wired and the + pipeline does NOT depend on diffusers model classes at import time + (REVIEW item 30 hard-ban). + """ + import fastvideo.registry # noqa: F401 + from fastvideo.api.presets import get_preset, get_presets_for_family + from fastvideo.configs.pipelines.stable_audio import StableAudioT2AConfig + from fastvideo.pipelines.basic.stable_audio.stable_audio_pipeline import ( + EntryClass, + StableAudioPipeline, + ) + from fastvideo.pipelines.basic.stable_audio.stages import ( # noqa: F401 + StableAudioConditioningStage, + StableAudioDecodingStage, + StableAudioDenoisingStage, + StableAudioLatentPreparationStage, + ) + # Native components — must import without diffusers in scope. + from fastvideo.models.dits.stable_audio import StableAudioDiT # noqa: F401 + from fastvideo.models.encoders.stable_audio_conditioner import ( # noqa: F401 + StableAudioMultiConditioner, + ) + + assert EntryClass is StableAudioPipeline + + names = {p.name for p in get_presets_for_family("stable_audio")} + assert names == {"stable_audio_open_1_0_base", "stable_audio_open_small"} + preset = get_preset("stable_audio_open_1_0_base", "stable_audio") + assert preset.defaults["num_inference_steps"] == 100 + assert preset.defaults["guidance_scale"] == 7.0 + # Audio workload: pin frame-shaped fields so the video-shaped sample + # buffer in `VideoGenerator` doesn't pre-allocate a hundreds-of-MB + # placeholder. The real output is the waveform on `result["audio"]`. + # `height`/`width` use 8 (smallest value the shared + # `InputValidationStage` accepts — rejects non-8-divisible). + assert preset.defaults["height"] == 8 + assert preset.defaults["width"] == 8 + assert preset.defaults["num_frames"] == 1 + + pc = StableAudioT2AConfig() + assert pc.sampling_rate == 44100 + assert pc.audio_channels == 2 + assert pc.guidance_scale == 7.0 + assert pc.num_inference_steps == 100 + from fastvideo.configs.models.vaes import OobleckVAEConfig + assert isinstance(pc.vae_config, OobleckVAEConfig) + assert pc.vae_config.pretrained_path == "stabilityai/stable-audio-open-1.0" diff --git a/tests/local_tests/stable-audio.md b/tests/local_tests/stable-audio.md new file mode 100644 index 000000000..7f2878258 --- /dev/null +++ b/tests/local_tests/stable-audio.md @@ -0,0 +1,131 @@ +# Local Stable Audio Tests + +End-to-end parity tests for the Stable Audio Open 1.0 pipeline + the +Oobleck VAE component. They compare FastVideo against the published +weights and the official `Stability-AI/stable-audio-tools` reference, +so they're skipped in CI and run locally on a single GPU. + +## Setup + +### 1. Hugging Face access + +Stable Audio Open 1.0 is a gated repo. Accept the terms once at +https://huggingface.co/stabilityai/stable-audio-open-1.0 and export +your token in the shell: + +```bash +export HF_TOKEN=hf_... +# any of HF_TOKEN / HUGGINGFACE_HUB_TOKEN / HF_API_KEY works +``` + +The tests skip cleanly with a helpful message if no token is set. + +### 2. Optional inference dependencies + +`k_diffusion` and friends are not part of FastVideo's base install. +One-shot: + +```bash +pip install k_diffusion einops_exts alias_free_torch torchsde +``` + +### 3. Clone the upstream reference repo + +The pipeline-level parity tests (`test_stable_audio_pipeline_parity.py`, +`test_stable_audio_a2a_parity.py`) and the VAE-component parity tests +(`test_oobleck_vae_parity.py`, `test_oobleck_vae_official_parity.py`) +import directly from `stable_audio_tools`. Clone it under the repo +root and `pip install` editable, then add it to `.gitignore`: + +```bash +cd +git clone --depth 1 https://github.com/Stability-AI/stable-audio-tools.git +pip install --no-deps -e ./stable-audio-tools +echo "/stable-audio-tools/" >> .git/info/exclude # personal ignore +``` + +The smoke test and the inpainting self-consistency test do not need the +upstream clone. + +### 4. (Optional) Pre-warm the model cache + +The first parity-test run downloads the full Stable Audio Open 1.0 +checkpoint (~3 GB). To avoid the download blocking your first test +run, fetch it ahead of time: + +```bash +python -c "from huggingface_hub import snapshot_download; \ + snapshot_download('stabilityai/stable-audio-open-1.0')" +``` + +## Running the tests + +All Stable Audio tests in one shot (~80 s on a single B200): + +```bash +pytest tests/local_tests/pipelines/test_stable_audio_pipeline_smoke.py \ + tests/local_tests/pipelines/test_stable_audio_pipeline_parity.py \ + tests/local_tests/pipelines/test_stable_audio_a2a_parity.py \ + tests/local_tests/pipelines/test_stable_audio_inpaint_parity.py \ + tests/local_tests/vaes/test_oobleck_vae_parity.py \ + tests/local_tests/vaes/test_oobleck_vae_official_parity.py \ + -v -s +``` + +Add `-s` to print the per-test diff numbers (shape / abs_mean / max +diff / drift / RMS ratios). + +### What each test covers + +| Test | Compares against | Notes | +|---|---|---| +| `test_stable_audio_pipeline_smoke.py` | — (no GPU) | Imports + registry + preset wiring. CPU-only. | +| `test_stable_audio_pipeline_parity.py` | `stable_audio_tools.inference.generate_diffusion_cond` | T2A end-to-end. Drift bound: < 1% / 0.05 element-wise. | +| `test_stable_audio_a2a_parity.py` | `generate_diffusion_cond(init_audio=...)` | Audio-to-audio variation. Drift bound: < 2% / 0.05 mean. | +| `test_stable_audio_inpaint_parity.py` | self-consistency | RePaint blending — kept-region RMS ratio in [0.5, 2.0], unkept region non-silent + bounded. SA Open 1.0 isn't inpaint-trained so there's no upstream reference to compare against. | +| `test_oobleck_vae_parity.py` | `stable_audio_tools.models.AudioAutoencoder` (real weights) | Decode + encode + round-trip; should be bit-identical (fp32, diff = 0). | +| `test_oobleck_vae_official_parity.py` | `OobleckEncoder` / `OobleckDecoder` (random init) | Architectural parity — verifies the FastVideo port matches upstream's structure. Bit-identical (diff = 0). | + +### Reproducing a single test + +Each test file is independent. Run one: + +```bash +pytest tests/local_tests/pipelines/test_stable_audio_a2a_parity.py -v -s +``` + +Expected output (approximate, on a B200, fp32, default seed = 0): + +``` +off shape=(1, 2, 66150) abs_mean=0.181810 +fv shape=(1, 2, 66150) abs_mean=0.180913 +diff max=0.030 mean=0.003 median=0.002 abs_mean rel drift=0.49% +PASSED +``` + +## Troubleshooting + +- **All tests skip with "gated repo" reason.** Check `echo $HF_TOKEN` + and that you've accepted the model terms at + https://huggingface.co/stabilityai/stable-audio-open-1.0. +- **`ModuleNotFoundError: stable_audio_tools` / `k_diffusion`.** Re-run + the setup steps; the parity tests skip cleanly if these aren't + importable but only after detecting both, so a partial install can + produce an unhelpful error mid-test. +- **A2A or inpaint test passes individually but fails in a batch + run.** The deterministic-math toggles (TF32 / cuDNN benchmark off) + are set in `StableAudioPipeline.load_modules`. If a test earlier in + the batch flipped them on, drift can spike — re-run the failing + test alone to confirm. +- **Wall-clock blew up past 5 min.** First call downloads the + checkpoint and the T5-base weights (~4 GB total) — see the + pre-warm step above. + +## Adding new parity tests for this family + +Use any of the existing `test_stable_audio_*` files as a template — the +HF-token detection (`_hf_token` / `_can_access` / `_setup_hf_env`) and +the loader for the official model (`_load_official`) are duplicated +across the parity-test files; copy them as a starting point and call +out the comparison reference (upstream class + invocation) in the file +docstring. diff --git a/tests/local_tests/vaes/test_oobleck_vae_official_parity.py b/tests/local_tests/vaes/test_oobleck_vae_official_parity.py new file mode 100644 index 000000000..d107c4551 --- /dev/null +++ b/tests/local_tests/vaes/test_oobleck_vae_official_parity.py @@ -0,0 +1,179 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Parity test: FastVideo's `OobleckVAE` port vs the official +`stable_audio_tools.models.autoencoders.OobleckEncoder/OobleckDecoder` +from Stability-AI/stable-audio-tools. + +This is the skill-compliant parity test (per `add-model` step 5: "clone +the official reference repo locally"). The companion test in this same +directory (`test_oobleck_vae_parity.py`) compares against +`diffusers.AutoencoderOobleck`, which is the canonical *runtime loader* +for the published HF weights and is already exact (diff=0). This test +adds the *structural* parity check against the official Python source. + +What it does: + 1. Build both encoders/decoders with the diffusers-shape config that + matches what Stable Audio Open 1.0 was trained for (5 downsampling + stages, latent_dim=64 / params 128, snake activations, no tanh). + 2. Walk both state_dicts in declaration order — they have the same + 183 parameters in the same order, modulo a (C,) vs (1, C, 1) + storage-shape difference for Snake alpha/beta. + 3. Copy upstream's randomly-initialized weights into FastVideo's + module via positional pairing + shape-reshape where needed. + 4. Forward identical input through both, assert close. + +Skips when: + * CUDA is unavailable. + * `stable-audio-tools` isn't importable (clone missing or its + `alias_free_torch` dep isn't installed). Install via: + git clone --depth 1 https://github.com/Stability-AI/stable-audio-tools.git + pip install --no-deps -e ./stable-audio-tools + pip install alias_free_torch +""" +from __future__ import annotations + +import os + +import pytest +import torch +from torch.testing import assert_close + + +def _stable_audio_tools_available() -> bool: + try: + from stable_audio_tools.models.autoencoders import OobleckEncoder # noqa + return True + except Exception: + return False + + +@pytest.mark.skipif( + not torch.cuda.is_available(), + reason="Oobleck official parity test requires CUDA.", +) +@pytest.mark.skipif( + not _stable_audio_tools_available(), + reason=( + "Official `stable-audio-tools` not importable. Install with:\n" + " git clone --depth 1 https://github.com/Stability-AI/stable-audio-tools.git\n" + " pip install --no-deps -e ./stable-audio-tools\n" + " pip install alias_free_torch" + ), +) +def test_oobleck_official_parity(): + from stable_audio_tools.models.autoencoders import OobleckDecoder, OobleckEncoder + from fastvideo.models.vaes.oobleck import OobleckVAE + + device = torch.device("cuda:0") + torch.manual_seed(0) + + # Build upstream with the diffusers-shape config that matches + # stabilityai/stable-audio-open-1.0/vae/config.json: + # downsampling_ratios=[2,4,4,8,8] -> upstream strides=[2,4,4,8,8] + # channel_multiples=[1,2,4,8,16] -> upstream c_mults=[1,2,4,8,16] + # decoder_input_channels=64 -> upstream encoder latent_dim=128 (mean+scale chunks) + # decoder latent_dim=64 + # audio_channels=2 (stereo) + # snake activations on, no tanh on decoder + up_enc = OobleckEncoder( + in_channels=2, channels=128, latent_dim=128, + c_mults=[1, 2, 4, 8, 16], strides=[2, 4, 4, 8, 8], + use_snake=True, antialias_activation=False, + ).to(device).eval() + up_dec = OobleckDecoder( + out_channels=2, channels=128, latent_dim=64, + c_mults=[1, 2, 4, 8, 16], strides=[2, 4, 4, 8, 8], + use_snake=True, antialias_activation=False, final_tanh=False, + ).to(device).eval() + + # FastVideo OobleckVAE — extract its encoder/decoder. + fv_vae = OobleckVAE( + encoder_hidden_size=128, + downsampling_ratios=[2, 4, 4, 8, 8], + channel_multiples=[1, 2, 4, 8, 16], + decoder_channels=128, decoder_input_channels=64, + audio_channels=2, sampling_rate=44100, + ).to(device).eval() + + # Transfer upstream weights into FastVideo via positional pairing. + _transfer_state_positional(up_enc, fv_vae.encoder, "encoder") + _transfer_state_positional(up_dec, fv_vae.decoder, "decoder") + + # Forward identical input. + sr = fv_vae.sampling_rate + n = (sr // fv_vae.hop_length) * fv_vae.hop_length # full chunks only + waveform = torch.randn((1, fv_vae.audio_channels, n), + dtype=torch.float32, device=device) + + with torch.inference_mode(): + # Encoder parity: upstream returns raw conv output (mean+scale + # concatenated along channel); FastVideo's encoder.forward returns + # the same. Diffusers wraps it in a posterior, but here we test + # the raw conv stack. + up_enc_out = up_enc(waveform).detach().float().cpu() + fv_enc_out = fv_vae.encoder(waveform).detach().float().cpu() + + print( + f"encoder: up shape={tuple(up_enc_out.shape)} abs_mean={up_enc_out.abs().mean().item():.6f}" + ) + print( + f"encoder: fv shape={tuple(fv_enc_out.shape)} abs_mean={fv_enc_out.abs().mean().item():.6f}" + ) + enc_diff = (up_enc_out - fv_enc_out).abs() + print(f"encoder diff max={enc_diff.max().item():.6e} mean={enc_diff.mean().item():.6e}") + + assert up_enc_out.shape == fv_enc_out.shape + assert_close(fv_enc_out, up_enc_out, atol=1e-5, rtol=1e-5) + + # Decoder parity on a random latent of the right shape. + torch.manual_seed(1) + latent = torch.randn((1, fv_vae.decoder_input_channels, 4), + dtype=torch.float32, device=device) + with torch.inference_mode(): + up_dec_out = up_dec(latent).detach().float().cpu() + fv_dec_out = fv_vae.decoder(latent).detach().float().cpu() + + print( + f"decoder: up shape={tuple(up_dec_out.shape)} abs_mean={up_dec_out.abs().mean().item():.6f}" + ) + print( + f"decoder: fv shape={tuple(fv_dec_out.shape)} abs_mean={fv_dec_out.abs().mean().item():.6f}" + ) + dec_diff = (up_dec_out - fv_dec_out).abs() + print(f"decoder diff max={dec_diff.max().item():.6e} mean={dec_diff.mean().item():.6e}") + + assert up_dec_out.shape == fv_dec_out.shape + assert_close(fv_dec_out, up_dec_out, atol=1e-5, rtol=1e-5) + + +def _transfer_state_positional(src: torch.nn.Module, dst: torch.nn.Module, label: str) -> None: + """Copy `src.state_dict()` into `dst.state_dict()` by walking both in + declaration order. Reshapes Snake alpha/beta from (C,) to (1, C, 1) + where needed (the only structural difference between upstream and + FastVideo's port). + + Asserts that both sides have the same number of parameters and that + each pair has the same number of elements. + """ + src_state = src.state_dict() + dst_state = dst.state_dict() + src_items = list(src_state.items()) + dst_items = list(dst_state.items()) + assert len(src_items) == len(dst_items), ( + f"{label}: param count mismatch — src={len(src_items)} dst={len(dst_items)}" + ) + new_state: dict[str, torch.Tensor] = {} + for (sk, sv), (dk, dv) in zip(src_items, dst_items): + if sv.shape == dv.shape: + new_state[dk] = sv.clone() + elif sv.numel() == dv.numel(): + new_state[dk] = sv.reshape(dv.shape).clone() + else: + raise RuntimeError( + f"{label}: numel mismatch at src={sk}{tuple(sv.shape)} " + f"vs dst={dk}{tuple(dv.shape)}" + ) + missing, unexpected = dst.load_state_dict(new_state, strict=True) + if missing or unexpected: + raise RuntimeError( + f"{label}: load_state_dict reported missing={missing[:3]} unexpected={unexpected[:3]}" + ) diff --git a/tests/local_tests/vaes/test_oobleck_vae_parity.py b/tests/local_tests/vaes/test_oobleck_vae_parity.py new file mode 100644 index 000000000..682737729 --- /dev/null +++ b/tests/local_tests/vaes/test_oobleck_vae_parity.py @@ -0,0 +1,242 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Parity test: FastVideo's first-class `OobleckVAE` port vs the +**official** `Stability-AI/stable-audio-tools` `AudioAutoencoder` loaded +from the published Stable Audio Open 1.0 checkpoint. + +(Was previously a diffusers comparison; switched to compare against the +upstream reference per the add-model skill, which calls for parity +against the official repo. The companion test +`test_oobleck_vae_official_parity.py` covers structural parity with +randomly-initialized weights; this file exercises the *real* trained +weights end-to-end through the official factory + checkpoint loader.) + +Skips when: + * CUDA is unavailable. + * `stabilityai/stable-audio-open-1.0` is inaccessible (gated repo). + * `stable_audio_tools` is not importable (clone missing or its + transitive deps not installed — `k_diffusion`, `einops_exts`, + `alias_free_torch`). + +Install official deps once via: + git clone --depth 1 https://github.com/Stability-AI/stable-audio-tools.git + pip install --no-deps -e ./stable-audio-tools + pip install k_diffusion einops_exts alias_free_torch +""" +from __future__ import annotations + +import json + +import pytest +import torch +from torch.testing import assert_close + +from fastvideo.utils import resolve_hf_token +from tests.local_tests._stable_audio_helpers import ( + can_access_repo, + setup_hf_env, +) + +_SA_AUDIO_ID = "stabilityai/stable-audio-open-1.0" +_VAE_CKPT_FILE = "vae_model.ckpt" +_VAE_CFG_FILE = "vae_model_config.json" + + +def _stable_audio_tools_available() -> bool: + try: + from stable_audio_tools.models.factory import ( # noqa: F401 + create_model_from_config, + ) + return True + except Exception: + return False + + +def _load_official_vae(device: torch.device): + """Build the official `AudioAutoencoder` from `vae_model_config.json` + and load its weights from `vae_model.ckpt`. + """ + from huggingface_hub import hf_hub_download + + cfg_path = hf_hub_download(repo_id=_SA_AUDIO_ID, filename=_VAE_CFG_FILE, + token=resolve_hf_token()) + ckpt_path = hf_hub_download(repo_id=_SA_AUDIO_ID, filename=_VAE_CKPT_FILE, + token=resolve_hf_token()) + with open(cfg_path) as f: + vcfg = json.load(f) + + from stable_audio_tools.models.factory import create_model_from_config + vae = create_model_from_config(vcfg) + + ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False) + state = ckpt.get("state_dict", ckpt) + pref = "autoencoder." + state = {k[len(pref):] if k.startswith(pref) else k: v for k, v in state.items()} + missing, unexpected = vae.load_state_dict(state, strict=True) + assert not missing and not unexpected, ( + f"official VAE state_dict mismatch — missing={missing[:3]} unexpected={unexpected[:3]}" + ) + return vae.to(device).eval() + + +_skip_no_cuda = pytest.mark.skipif( + not torch.cuda.is_available(), + reason="Oobleck VAE parity test requires CUDA.", +) +_skip_no_access = pytest.mark.skipif( + not can_access_repo(_SA_AUDIO_ID, filename=_VAE_CFG_FILE), + reason=(f"{_SA_AUDIO_ID} not accessible — gated repo; set HF_TOKEN / " + f"HF_API_KEY and accept the terms on https://huggingface.co/{_SA_AUDIO_ID}."), +) +_skip_no_official = pytest.mark.skipif( + not _stable_audio_tools_available(), + reason=("`stable_audio_tools` not importable. Clone " + "https://github.com/Stability-AI/stable-audio-tools and `pip install` " + "its deps (k_diffusion, einops_exts, alias_free_torch)."), +) + + +@_skip_no_cuda +@_skip_no_access +@_skip_no_official +def test_oobleck_vae_decode_official_parity(): + """Decode-side parity with the official AudioAutoencoder.""" + setup_hf_env() + device = torch.device("cuda:0") + + off_vae = _load_official_vae(device) + + from fastvideo.models.vaes.oobleck import OobleckVAE + fv_vae = OobleckVAE.from_pretrained( + _SA_AUDIO_ID, subfolder="vae", torch_dtype=torch.float32, + ).to(device).eval() + + torch.manual_seed(0) + latent = torch.randn( + (1, fv_vae.decoder_input_channels, 8), + dtype=torch.float32, device=device, + ) + + with torch.inference_mode(): + off_out = off_vae.decode(latent).detach().float().cpu() + fv_out = fv_vae.decode(latent).sample.detach().float().cpu() + + print( + f"off shape={tuple(off_out.shape)} " + f"abs_mean={off_out.abs().mean().item():.6f} " + f"range=[{off_out.min().item():.4f}, {off_out.max().item():.4f}]" + ) + print( + f"fv shape={tuple(fv_out.shape)} " + f"abs_mean={fv_out.abs().mean().item():.6f} " + f"range=[{fv_out.min().item():.4f}, {fv_out.max().item():.4f}]" + ) + diff = (off_out - fv_out).abs() + print(f"diff max={diff.max().item():.6e} mean={diff.mean().item():.6e}") + + assert off_out.shape == fv_out.shape + # FV is a structural rewrite of upstream's OobleckEncoder/Decoder, + # loading the same weights via `OobleckVAE.from_pretrained`. Output + # should be bit-identical at fp32. + assert_close(fv_out, off_out, atol=1e-5, rtol=1e-5) + + +@_skip_no_cuda +@_skip_no_access +@_skip_no_official +def test_oobleck_vae_encode_official_parity(): + """Encode-side parity vs the official autoencoder. + + Compares the **mean** of the diagonal Gaussian posterior, not a + sample. Official `VAEBottleneck.encode` calls `vae_sample(mean, + scale)` which adds `randn() * softplus(scale)+1e-4` — non-deterministic. + To compare apples-to-apples we ask the official side to skip the + bottleneck (`skip_bottleneck=True`), chunk to (mean, scale) + ourselves, and compare against FV's `.encode().mode()`. + """ + setup_hf_env() + device = torch.device("cuda:0") + + off_vae = _load_official_vae(device) + + from fastvideo.models.vaes.oobleck import OobleckVAE + fv_vae = OobleckVAE.from_pretrained( + _SA_AUDIO_ID, subfolder="vae", torch_dtype=torch.float32, + ).to(device).eval() + + torch.manual_seed(1) + waveform = torch.randn( + (1, fv_vae.audio_channels, 44100), + dtype=torch.float32, device=device, + ) + # Round to a multiple of hop_length so encode aligns. + n = (waveform.shape[-1] // fv_vae.hop_length) * fv_vae.hop_length + waveform = waveform[..., :n] + + with torch.inference_mode(): + # Bypass the stochastic VAEBottleneck — we want the mean. + off_pre = off_vae.encode(waveform, skip_bottleneck=True) + off_mean, _off_scale = off_pre.chunk(2, dim=1) + fv_mean = fv_vae.encode(waveform).mode() + + off_mean_cpu = off_mean.detach().float().cpu() + fv_mean_cpu = fv_mean.detach().float().cpu() + + diff = (off_mean_cpu - fv_mean_cpu).abs() + print( + f"off_mean shape={tuple(off_mean_cpu.shape)} abs_mean={off_mean_cpu.abs().mean().item():.6f}" + ) + print( + f"fv_mean shape={tuple(fv_mean_cpu.shape)} abs_mean={fv_mean_cpu.abs().mean().item():.6f}" + ) + print(f"diff max={diff.max().item():.6e} mean={diff.mean().item():.6e}") + + assert off_mean_cpu.shape == fv_mean_cpu.shape + assert_close(fv_mean_cpu, off_mean_cpu, atol=1e-5, rtol=1e-5) + + +@_skip_no_cuda +@_skip_no_access +@_skip_no_official +def test_oobleck_vae_round_trip_sanity(): + """Sanity check: encode → decode preserves shape and stays in a + sensible amplitude range. Not a parity test — just guards against + silent diverging VAE output. Uses FV alone (no comparison). + """ + setup_hf_env() + device = torch.device("cuda:0") + + from fastvideo.models.vaes.oobleck import OobleckVAE + fv_vae = OobleckVAE.from_pretrained( + _SA_AUDIO_ID, subfolder="vae", torch_dtype=torch.float32, + ).to(device).eval() + + torch.manual_seed(2) + sr = fv_vae.sampling_rate + n = (sr // fv_vae.hop_length) * fv_vae.hop_length + t = torch.linspace(0, n / sr, n, device=device) + freqs = torch.linspace(220.0, 880.0, n, device=device) + mono = torch.sin(2 * torch.pi * freqs * t) * 0.3 + waveform = mono.unsqueeze(0).repeat(2, 1).unsqueeze(0).contiguous() + + with torch.inference_mode(): + latent = fv_vae.encode(waveform).mode() + recon = fv_vae.decode(latent).sample + + assert recon.shape == waveform.shape, ( + f"round-trip shape mismatch: in={waveform.shape} out={recon.shape}" + ) + assert torch.isfinite(recon).all(), "round-trip produced non-finite values" + assert recon.abs().max().item() < 5.0, ( + f"round-trip output magnitude {recon.abs().max().item():.3f} > 5.0" + ) + + in_rms = waveform.float().pow(2).mean().sqrt().item() + out_rms = recon.float().pow(2).mean().sqrt().item() + rms_ratio = out_rms / max(in_rms, 1e-6) + print( + f"round-trip in_rms={in_rms:.4f} out_rms={out_rms:.4f} " + f"ratio={rms_ratio:.3f}" + ) + assert 0.3 < rms_ratio < 3.0, ( + f"round-trip RMS ratio {rms_ratio:.3f} outside sanity band [0.3, 3.0]" + )