Vendor stable-audio-3 for ZeroGPU

README.md

@@ -42,3 +42,11 @@ Base checkpoints are not gated, but they are intended mainly for fine-tuning and
 - `SAME-L` is GPU-first; `SAME-S` can be used for CPU autoencoder round trips.
 
 The Space is configured with `suggested_hardware: zero-a10g`.
+
+## Runtime note
+
+The upstream `stable-audio-3` Python package is vendored in this Space from
+Stability AI's public MIT-licensed repository because its package metadata pins
+Torch 2.7.1. ZeroGPU currently provides Torch 2.8.0, so installing the upstream
+package through normal dependency resolution would downgrade Torch and break the
+ZeroGPU runtime.

STABLE_AUDIO_3_LICENSE

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

requirements.txt

@@ -3,6 +3,9 @@ torchaudio==2.8.0
 gradio==6.3.0
 spaces
 hf_transfer
+numpy>=2.2.6
+packaging>=26.0
+huggingface-hub>=1.7.1
 soundfile
 einops>=0.8.2
 einops-exts>=0.0.4

stable_audio_3/__init__.py

@@ -0,0 +1,2 @@
+from stable_audio_3.model import AutoencoderModel as AutoencoderModel
+from stable_audio_3.model import StableAudioModel as StableAudioModel

stable_audio_3/cli.py

@@ -0,0 +1,292 @@
+"""
+stable-audio — command-line interface for Stable Audio 3.
+
+Basic usage::
+
+    stable-audio --model small-music -p "lo-fi hip hop beat, 90 BPM" --duration 30 -o beat.wav
+
+"""
+
+import argparse
+import os
+import torch
+import torchaudio
+
+from stable_audio_3 import StableAudioModel
+
+
+def _save_output(audio: torch.Tensor, sample_rate: int, output: str, batch_size: int):
+    """Save generated audio tensor(s) to disk."""
+    base, ext = os.path.splitext(output)
+    if not ext:
+        ext = ".wav"
+    for i in range(batch_size):
+        path = f"{base}_{i}{ext}" if batch_size > 1 else f"{base}{ext}"
+        torchaudio.save(path, audio[i].cpu(), sample_rate)
+        print(f"Saved: {path}")
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        prog="stable-audio",
+        description="Stable Audio 3 — CLI for text-to-audio, audio-to-audio, and inpainting",
+    )
+
+    # Model
+    parser.add_argument(
+        "--model",
+        default="medium",
+        choices=[
+            "medium",
+            "small-music",
+            "small-sfx",
+            "medium-base",
+            "small-music-base",
+            "small-sfx-base",
+        ],
+        help="Model to load (default: medium)",
+    )
+    parser.add_argument(
+        "--device",
+        default=None,
+        help="Device: cuda / mps / cpu (auto-detected if omitted)",
+    )
+    parser.add_argument(
+        "--no-half", action="store_true", help="Disable half-precision (fp16) on CUDA"
+    )
+
+    # Generation
+    parser.add_argument(
+        "-p",
+        "--prompt",
+        required=True,
+        nargs="+",
+        help="Text prompt(s). Pass multiple for per-batch prompts",
+    )
+    parser.add_argument(
+        "--negative-prompt", nargs="+", default=None, help="Negative prompt(s)"
+    )
+    parser.add_argument(
+        "--duration",
+        type=float,
+        nargs="+",
+        default=[120.0],
+        help="Duration in seconds (default: 120). Pass multiple for per-batch durations",
+    )
+    parser.add_argument(
+        "--steps", type=int, default=8, help="Diffusion steps (default: 8)"
+    )
+    parser.add_argument(
+        "--cfg-scale",
+        type=float,
+        default=1.0,
+        help="CFG scale (default: 1.0; try 7.0 for base models)",
+    )
+    parser.add_argument(
+        "--seed", type=int, default=-1, help="Random seed (-1 = random, default: -1)"
+    )
+    parser.add_argument(
+        "--batch-size",
+        type=int,
+        default=None,
+        help="Batch size (default: inferred from number of prompts, or 1)",
+    )
+    parser.add_argument(
+        "-o",
+        "--output",
+        default="output.wav",
+        help="Output file path (default: output.wav)",
+    )
+
+    # Audio-to-Audio
+    parser.add_argument(
+        "--init-audio",
+        default=None,
+        metavar="PATH",
+        help="Source audio file for audio-to-audio generation",
+    )
+    parser.add_argument(
+        "--init-noise-level",
+        type=float,
+        default=0.9,
+        help="Noise level for audio-to-audio (0.0–1.0, default: 0.9)",
+    )
+
+    # Inpainting / Continuation
+    parser.add_argument(
+        "--inpaint-audio",
+        default=None,
+        metavar="PATH",
+        help="Source audio file for inpainting or continuation",
+    )
+    parser.add_argument(
+        "--inpaint-start",
+        type=float,
+        action="append",
+        dest="inpaint_starts",
+        metavar="SECONDS",
+        help="Start of inpaint region in seconds. Repeat for multiple regions.",
+    )
+    parser.add_argument(
+        "--inpaint-end",
+        type=float,
+        action="append",
+        dest="inpaint_ends",
+        metavar="SECONDS",
+        help="End of inpaint region in seconds. Repeat for multiple regions.",
+    )
+
+    # Chunked decode
+    decode_group = parser.add_mutually_exclusive_group()
+    decode_group.add_argument(
+        "--chunked-decode",
+        action="store_true",
+        default=None,
+        help="Force chunked decoding on",
+    )
+    decode_group.add_argument(
+        "--no-chunked-decode",
+        action="store_true",
+        default=None,
+        help="Force chunked decoding off",
+    )
+
+    # LoRA
+    parser.add_argument(
+        "--lora-ckpt-path",
+        action="append",
+        dest="loras",
+        metavar="PATH",
+        help="LoRA checkpoint path. Repeat to stack multiple LoRAs.",
+    )
+    parser.add_argument(
+        "--lora-strength",
+        type=float,
+        default=None,
+        help="LoRA strength (applied to all LoRAs)",
+    )
+    parser.add_argument(
+        "--lora-index",
+        type=int,
+        default=None,
+        help="Target a specific LoRA index when setting strength",
+    )
+
+    args = parser.parse_args()
+
+    # --- Validate inpaint args ---
+    if (args.inpaint_starts is None) != (args.inpaint_ends is None):
+        parser.error("--inpaint-start and --inpaint-end must both be provided together")
+    if args.inpaint_starts and len(args.inpaint_starts) != len(args.inpaint_ends):
+        parser.error(
+            "--inpaint-start and --inpaint-end must be specified the same number of times"
+        )
+    if args.inpaint_starts and not args.inpaint_audio:
+        parser.error("--inpaint-start/--inpaint-end require --inpaint-audio")
+    if args.inpaint_audio and not args.inpaint_starts:
+        parser.error("--inpaint-audio requires --inpaint-start and --inpaint-end")
+
+    # --- Resolve batch size ---
+    n_prompts = len(args.prompt)
+    if args.batch_size is None:
+        batch_size = n_prompts
+    elif n_prompts > 1 and args.batch_size != n_prompts:
+        parser.error(
+            f"--batch-size {args.batch_size} does not match the number of prompts "
+            f"({n_prompts}); omit --batch-size to have it inferred automatically"
+        )
+    else:
+        batch_size = args.batch_size
+
+    # --- Validate list-flag lengths against batch size ---
+    if (
+        args.negative_prompt
+        and len(args.negative_prompt) > 1
+        and len(args.negative_prompt) != batch_size
+    ):
+        parser.error(
+            f"Got {len(args.negative_prompt)} --negative-prompt values but batch size is {batch_size}"
+        )
+    if len(args.duration) > 1 and len(args.duration) != batch_size:
+        parser.error(
+            f"Got {len(args.duration)} --duration values but batch size is {batch_size}"
+        )
+
+    # --- Build scalar / list args ---
+    prompt = args.prompt[0] if len(args.prompt) == 1 else args.prompt
+    negative_prompt = None
+    if args.negative_prompt:
+        negative_prompt = (
+            args.negative_prompt[0]
+            if len(args.negative_prompt) == 1
+            else args.negative_prompt
+        )
+    duration = args.duration[0] if len(args.duration) == 1 else args.duration
+
+    # --- chunked_decode flag ---
+    chunked_decode = None
+    if args.chunked_decode:
+        chunked_decode = True
+    elif args.no_chunked_decode:
+        chunked_decode = False
+
+    # --- Load model ---
+    print(f"Loading model '{args.model}'…")
+    model = StableAudioModel.from_pretrained(
+        args.model, device=args.device, model_half=not args.no_half
+    )
+
+    # --- LoRA ---
+    if args.loras:
+        print(f"Loading LoRA(s): {args.loras}")
+        model.load_lora(args.loras)
+    if args.lora_strength is not None:
+        model.set_lora_strength(args.lora_strength, lora_index=args.lora_index)
+
+    # --- Load audio inputs ---
+    # torchaudio.load returns (waveform, sample_rate); model.generate expects (sample_rate, waveform)
+    init_audio = None
+    if args.init_audio:
+        waveform, sr = torchaudio.load(args.init_audio)
+        init_audio = (sr, waveform)
+
+    inpaint_audio = None
+    if args.inpaint_audio:
+        waveform, sr = torchaudio.load(args.inpaint_audio)
+        inpaint_audio = (sr, waveform)
+
+    inpaint_start = None
+    inpaint_end = None
+    if args.inpaint_starts:
+        inpaint_start = (
+            args.inpaint_starts[0]
+            if len(args.inpaint_starts) == 1
+            else args.inpaint_starts
+        )
+        inpaint_end = (
+            args.inpaint_ends[0] if len(args.inpaint_ends) == 1 else args.inpaint_ends
+        )
+
+    # --- Generate ---
+    print("Generating…")
+    audio = model.generate(
+        prompt=prompt,
+        negative_prompt=negative_prompt,
+        duration=duration,
+        steps=args.steps,
+        cfg_scale=args.cfg_scale,
+        seed=args.seed,
+        batch_size=batch_size,
+        init_audio=init_audio,
+        init_noise_level=args.init_noise_level,
+        inpaint_audio=inpaint_audio,
+        inpaint_mask_start_seconds=inpaint_start,
+        inpaint_mask_end_seconds=inpaint_end,
+        chunked_decode=chunked_decode,
+    )
+
+    _save_output(audio, model.model.sample_rate, args.output, batch_size)
+
+
+if __name__ == "__main__":
+    main()

stable_audio_3/data/dataset.py

@@ -0,0 +1,510 @@
+import numpy as np
+import json
+import os
+import dill
+import random
+import time
+import torch
+import torchaudio
+
+from os import path
+from torchaudio import transforms as T
+from typing import Optional, Callable, List
+
+from .utils import Stereo, Mono, PhaseFlipper, PadCrop_Normalized_T, VolumeNorm, strip_trailing_silence
+
+AUDIO_KEYS = ("flac", "wav", "mp3", "m4a", "ogg", "opus")
+
+# fast_scandir implementation by Scott Hawley originally in https://github.com/zqevans/audio-diffusion/blob/main/dataset/dataset.py
+
+def fast_scandir(
+    dir:str,  # top-level directory at which to begin scanning
+    ext:list,  # list of allowed file extensions,
+    #max_size = 1 * 1000 * 1000 * 1000 # Only files < 1 GB
+    ):
+    "very fast `glob` alternative. from https://stackoverflow.com/a/59803793/4259243"
+    subfolders, files = [], []
+    ext = ['.'+x if x[0]!='.' else x for x in ext]  # add starting period to extensions if needed
+    try: # hope to avoid 'permission denied' by this try
+        for f in os.scandir(dir):
+            try: # 'hope to avoid too many levels of symbolic links' error
+                if f.is_dir():
+                    subfolders.append(f.path)
+                elif f.is_file():
+                    file_ext = os.path.splitext(f.name)[1].lower()
+                    is_hidden = os.path.basename(f.path).startswith(".")
+
+                    if file_ext in ext and not is_hidden:
+                        files.append(f.path)
+            except:
+                pass 
+    except:
+        pass
+
+    for dir in list(subfolders):
+        sf, f = fast_scandir(dir, ext)
+        subfolders.extend(sf)
+        files.extend(f)
+    return subfolders, files
+
+def keyword_scandir(
+    dir: str,  # top-level directory at which to begin scanning
+    ext: list,  # list of allowed file extensions
+    keywords: list,  # list of keywords to search for in the file name
+):
+    "very fast `glob` alternative. from https://stackoverflow.com/a/59803793/4259243"
+    subfolders, files = [], []
+    # make keywords case insensitive
+    keywords = [keyword.lower() for keyword in keywords]
+    # add starting period to extensions if needed
+    ext = ['.'+x if x[0] != '.' else x for x in ext]
+    banned_words = ["paxheader", "__macosx"]
+    try:  # hope to avoid 'permission denied' by this try
+        for f in os.scandir(dir):
+            try:  # 'hope to avoid too many levels of symbolic links' error
+                if f.is_dir():
+                    subfolders.append(f.path)
+                elif f.is_file():
+                    is_hidden = f.name.split("/")[-1][0] == '.'
+                    has_ext = os.path.splitext(f.name)[1].lower() in ext
+                    name_lower = f.name.lower()
+                    has_keyword = any(
+                        [keyword in name_lower for keyword in keywords])
+                    has_banned = any(
+                        [banned_word in name_lower for banned_word in banned_words])
+                    if has_ext and has_keyword and not has_banned and not is_hidden and not os.path.basename(f.path).startswith("._"):
+                        files.append(f.path)
+            except:
+                pass
+    except:
+        pass
+
+    for dir in list(subfolders):
+        sf, f = keyword_scandir(dir, ext, keywords)
+        subfolders.extend(sf)
+        files.extend(f)
+    return subfolders, files
+
+def get_audio_filenames(
+    paths: list,  # directories in which to search
+    keywords=None,
+    exts=['.wav', '.mp3', '.flac', '.ogg', '.aif', '.opus'],
+    filelist_path=None
+):
+    "recursively get a list of audio filenames"
+    filenames = []
+    if type(paths) is str:
+        paths = [paths]
+    for path in paths:               # get a list of relevant filenames
+
+        if filelist_path is None:
+            # Check for filelist.txt at the root of the directory
+            filelist_path = os.path.join(path, "filelist.txt")
+            
+        if os.path.exists(filelist_path):
+            with open(filelist_path, "r") as f:
+                files = f.readlines()
+                files = [os.path.join(path, file.strip()) for file in files]
+                filenames.extend(files)
+            continue
+
+        if keywords is not None:
+            subfolders, files = keyword_scandir(path, exts, keywords)
+        else:
+            subfolders, files = fast_scandir(path, exts)
+        filenames.extend(files)
+    return filenames
+
+def get_latent_filenames(
+    paths,  # directories in which to search
+    extension='npy',
+    filelist_path=None
+):
+    "recursively get a list of pre-encoded filenames"
+    filenames = []
+    if type(paths) is str:
+        paths = [paths]
+    for path in paths:               # get a list of relevant filenames
+
+        if filelist_path is None:
+            # Check for filelist.txt at the root of the directory
+            filelist_path = os.path.join(path, "filelist.txt")
+        
+        if os.path.exists(filelist_path):
+            with open(filelist_path, "r") as f:
+                files = f.readlines()
+                files = [os.path.join(path, file.strip()) for file in files]
+                filenames.extend(files)
+            continue
+
+        _, files = fast_scandir(path, [extension])
+        filenames.extend(files)
+
+    # Filter out silence.npy (used for silence latent padding, not a data sample)
+    filenames = [f for f in filenames if os.path.basename(f) != "silence.npy"]
+
+    # Add metadata paths
+    filenames = [(filename, filename.replace(f".{extension}", ".json")) for filename in filenames]
+
+    return filenames
+
+class LocalDatasetConfig:
+    def __init__(
+        self,
+        id: str,
+        path: str,
+        keywords: Optional[List[str]]=None,
+        custom_metadata_fn: Optional[Callable[[str], str]] = None,
+        filelist_path = None,
+        weight: float = 1.0,
+    ):
+        self.id = id
+        self.path = path
+        self.custom_metadata_fn = custom_metadata_fn
+        self.keywords = keywords
+        self.filelist_path = filelist_path
+        self.weight = weight
+
+class LatentDatasetConfig(LocalDatasetConfig):
+    def __init__(
+        self,
+        latent_extension: str = "npy",
+        filelist_path = None,
+        *args,
+        **kwargs,
+    ):
+        super().__init__(*args, **kwargs)
+        self.latent_extension = latent_extension
+        self.filelist_path = filelist_path
+        # weight is inherited from LocalDatasetConfig via **kwargs
+
+class SampleDataset(torch.utils.data.Dataset):
+    def __init__(
+        self,
+        configs,
+        sample_size=65536,
+        sample_rate=48000,
+        random_crop=True,
+        force_channels="stereo",
+        volume_norm=False,
+        volume_norm_param=(-16, 2),
+        strip_silence=False,
+        pad=True,
+    ):
+        super().__init__()
+        self.filenames = []
+        self.sample_weights = []
+
+        self.augs = torch.nn.Sequential(
+            PhaseFlipper(),
+            #nn.Identity()
+        )
+
+
+        self.root_paths = []
+
+        self.pad_crop = PadCrop_Normalized_T(sample_size, sample_rate, randomize=random_crop, pad=pad)
+        self.strip_silence = strip_silence
+
+        self.force_channels = force_channels
+
+        self.encoding = torch.nn.Sequential(
+            Stereo() if self.force_channels == "stereo" else torch.nn.Identity(),
+            Mono() if self.force_channels == "mono" else torch.nn.Identity()
+        )
+
+        self.sr = sample_rate
+
+        self.volume_norm = VolumeNorm(volume_norm_param, self.sr) if volume_norm else torch.nn.Identity()
+
+        self.custom_metadata_fns = {}
+
+        for config in configs:
+            self.root_paths.append(config.path)
+            new_files = get_audio_filenames(config.path, config.keywords, filelist_path=config.filelist_path)
+            self.filenames.extend(new_files)
+            self.sample_weights.extend([config.weight] * len(new_files))
+            if config.custom_metadata_fn is not None:
+                self.custom_metadata_fns[config.path] = dill.dumps(config.custom_metadata_fn)
+
+        print(f'Found {len(self.filenames)} files')
+
+    def load_file(self, filename):
+        ext = filename.split(".")[-1]
+
+        audio, in_sr = torchaudio.load(filename, format=ext)
+
+        if in_sr != self.sr:
+            resample_tf = T.Resample(in_sr, self.sr)
+            audio = resample_tf(audio)
+
+        return audio
+
+    def __len__(self):
+        return len(self.filenames)
+
+    def __getitem__(self, idx):
+        audio_filename = self.filenames[idx]
+        try:
+            start_time = time.time()
+            audio = self.load_file(audio_filename)
+
+            audio = self.volume_norm(audio)
+
+            if self.strip_silence:
+                audio = strip_trailing_silence(audio, self.sr)
+
+            audio, t_start, t_end, seconds_start, seconds_total, padding_mask = self.pad_crop(audio)
+
+            # Check for silence
+            if is_silence(audio):
+                return self[random.randrange(len(self))]
+
+            # Run augmentations on this sample (including random crop)
+            if self.augs is not None:
+                audio = self.augs(audio)
+
+            audio = audio.clamp(-1, 1)
+
+            # Encode the file to assist in prediction
+            if self.encoding is not None:
+                audio = self.encoding(audio)
+
+            info = {}
+
+            info["path"] = audio_filename
+
+            for root_path in self.root_paths:
+                if root_path in audio_filename:
+                    info["relpath"] = path.relpath(audio_filename, root_path)
+
+            info["timestamps"] = (t_start, t_end)
+            info["seconds_start"] = seconds_start
+            info["seconds_total"] = seconds_total
+            info["padding_mask"] = [padding_mask]
+            info["sample_rate"] = self.sr
+
+            end_time = time.time()
+
+            info["load_time"] = end_time - start_time
+
+            for custom_md_path in self.custom_metadata_fns.keys():
+                if custom_md_path in audio_filename:
+                    custom_metadata_fn = dill.loads(self.custom_metadata_fns[custom_md_path])
+                    custom_metadata = custom_metadata_fn(info, audio)
+                    info.update(custom_metadata)
+
+                if "__reject__" in info and info["__reject__"]:
+                    return self[random.randrange(len(self))]
+
+                # Provide audio inputs as their own dictionary to be merged into info, each audio element will be normalized in the same way as the main audio
+                if "__audio__" in info:
+                    for audio_key, audio_value in info["__audio__"].items():
+                        # Process the audio_value tensor, which should be a torch tensor
+                        audio_value, _, _, _, _, _ = self.pad_crop(audio_value)
+                        audio_value = audio_value.clamp(-1, 1)
+                        if self.encoding is not None:
+                            audio_value = self.encoding(audio_value)
+                        info[audio_key] = audio_value
+                
+                    del info["__audio__"]
+
+            return (audio, info)
+        except Exception as e:
+            print(f'Couldn\'t load file {audio_filename}: {e}')
+            return self[random.randrange(len(self))]
+
+
+class PreEncodedDataset(torch.utils.data.Dataset):
+    def __init__(
+        self,
+        configs: List[LatentDatasetConfig],
+        latent_crop_length=None,
+        min_length_sec=None,
+        max_length_sec=None,
+        random_crop=False,
+        tokenizers: Optional[dict] = None,
+    ):
+        super().__init__()
+        self.filenames = []
+        self.sample_weights = []
+
+        self.custom_metadata_fns = {}
+
+        self.silence_latents = {}
+
+        for config in configs:
+            new_files = get_latent_filenames(config.path, config.latent_extension, config.filelist_path)
+            self.filenames.extend(new_files)
+            self.sample_weights.extend([config.weight] * len(new_files))
+            if config.custom_metadata_fn is not None:
+                self.custom_metadata_fns[config.path] = dill.dumps(config.custom_metadata_fn)
+
+            # Load silence latent if available (for variable-length padding)
+            paths = config.path if isinstance(config.path, list) else [config.path]
+            for path in paths:
+                silence_path = os.path.join(path, "silence.npy")
+                if os.path.exists(silence_path):
+                    self.silence_latents[path] = np.load(silence_path).squeeze(0)  # [C, N]
+                    print(f'Loaded silence latent from {silence_path}')
+
+        self.latent_crop_length = latent_crop_length
+        self.random_crop = random_crop
+
+        self.min_length_sec = min_length_sec
+        self.max_length_sec = max_length_sec
+
+        # tokenizers: dict mapping metadata key -> (tokenizer, max_length)
+        # If provided, text fields will be pre-tokenized in DataLoader workers
+        self.tokenizers = tokenizers
+
+        print(f'Found {len(self.filenames)} files')
+
+    def __len__(self):
+        return len(self.filenames)
+
+    def _get_silence_for_file(self, latent_filename):
+        """Return the silence latent for the dataset that contains this file, or None."""
+        for path, silence in self.silence_latents.items():
+            if path in latent_filename:
+                return silence
+        return None
+
+    def __getitem__(self, idx):
+        latent_filename, md_filename = self.filenames[idx]
+        try:
+            latents = torch.from_numpy(np.load(latent_filename)) # [C, N]
+
+            with open(md_filename, "r") as f:
+                try:
+                    info = json.load(f)
+                except:
+                    raise Exception(f"Couldn't load metadata file {md_filename}")
+
+            info["latent_filename"] = latent_filename
+
+            if self.latent_crop_length is not None:
+                stored_length = latents.shape[1]
+
+                if stored_length > self.latent_crop_length:
+                    # Crop to latent_crop_length (existing logic)
+                    # Get the last index from the padding mask, the index of the last 1 in the sequence
+                    last_ix = len(info["padding_mask"]) - 1 - info["padding_mask"][::-1].index(1)
+
+                    if self.random_crop and last_ix > self.latent_crop_length:
+                        start = random.randint(0, last_ix - self.latent_crop_length)
+                    else:
+                        start = 0
+
+                    latents = latents[:, start:start+self.latent_crop_length]
+                    info["padding_mask"] = info["padding_mask"][start:start+self.latent_crop_length]
+                    info["latent_crop_start"] = start
+
+                elif stored_length < self.latent_crop_length:
+                    # Pad with silence latent to reach latent_crop_length
+                    pad_needed = self.latent_crop_length - stored_length
+                    silence = self._get_silence_for_file(latent_filename)
+
+                    if silence is not None:
+                        # Slice or tile silence latent to cover pad_needed frames
+                        if silence.shape[1] >= pad_needed:
+                            silence_pad = silence[:, :pad_needed]
+                        else:
+                            silence_pad = np.tile(silence, (1, (pad_needed // silence.shape[1]) + 1))[:, :pad_needed]
+                        latents = torch.cat([latents, torch.from_numpy(silence_pad)], dim=1)
+                    else:
+                        # No silence latent available — zero-pad as fallback
+                        latents = torch.nn.functional.pad(latents, (0, pad_needed))
+
+                    # Build padding_mask: valid frames from stored mask, zeros for padding
+                    info["padding_mask"] = info["padding_mask"][:stored_length] + [0] * pad_needed
+                    info["latent_crop_start"] = 0
+
+                else:
+                    # Exact match
+                    info["latent_crop_start"] = 0
+
+                info["latent_crop_length"] = self.latent_crop_length
+
+            info["padding_mask"] = [torch.tensor(info["padding_mask"])]
+
+            seconds_total = info["seconds_total"]
+
+            if self.min_length_sec is not None and seconds_total < self.min_length_sec:
+                return self[random.randrange(len(self))]
+
+            if self.max_length_sec is not None and seconds_total > self.max_length_sec:
+                return self[random.randrange(len(self))]
+
+            for custom_md_path in self.custom_metadata_fns.keys():
+                if custom_md_path in latent_filename:
+                    custom_metadata_fn = dill.loads(self.custom_metadata_fns[custom_md_path])
+                    custom_metadata = custom_metadata_fn(info, latents)
+                    info.update(custom_metadata)
+
+                if "__reject__" in info and info["__reject__"]:
+                    return self[random.randrange(len(self))]
+
+                if "__replace__" in info and info["__replace__"] is not None:
+                    # Replace the latents with the new latents if the custom metadata function returns a new set of latents
+                    latents = info["__replace__"]
+
+            info["audio"] = latents
+
+            # Pre-tokenize text fields in DataLoader workers to avoid
+            # CPU contention with the main training thread
+            if self.tokenizers is not None:
+                for key, (tokenizer, max_length) in self.tokenizers.items():
+                    if key in info and isinstance(info[key], str):
+                        # Save raw text before replacing with tokens (needed by CLAP and other text-based losses)
+                        info[f"{key}_text"] = info[key]
+                        encoded = tokenizer(
+                            info[key],
+                            truncation=True,
+                            max_length=max_length,
+                            padding="max_length",
+                            return_tensors="pt",
+                        )
+                        info[key] = {
+                            "input_ids": encoded["input_ids"].squeeze(0),
+                            "attention_mask": encoded["attention_mask"].squeeze(0),
+                        }
+
+            return (latents, info)
+        except Exception as e:
+            print(f'Couldn\'t load file {latent_filename}: {e}')
+            return self[random.randrange(len(self))]
+
+# get_dbmax and is_silence copied from https://github.com/drscotthawley/aeiou/blob/main/aeiou/core.py under Apache 2.0 License
+# License can be found in LICENSES/LICENSE_AEIOU.txt
+def get_dbmax(
+    audio,       # torch tensor of (multichannel) audio
+    ):
+    "finds the loudest value in the entire clip and puts that into dB (full scale)"
+    return 20*torch.log10(torch.flatten(audio.abs()).max()).cpu().numpy()
+
+def is_silence(
+    audio,       # torch tensor of (multichannel) audio
+    thresh=-60,  # threshold in dB below which we declare to be silence
+    ):
+    "checks if entire clip is 'silence' below some dB threshold"
+    dBmax = get_dbmax(audio)
+    return dBmax < thresh
+
+
+def collation_fn(samples):
+        batched = list(zip(*samples))
+        result = []
+        for b in batched:
+            if isinstance(b[0], (int, float)):
+                b = np.array(b)
+            elif isinstance(b[0], torch.Tensor):
+                b = torch.stack(b)
+            elif isinstance(b[0], np.ndarray):
+                b = np.array(b)
+            else:
+                b = b
+            result.append(b)
+        return result
+
+

stable_audio_3/data/utils.py

@@ -0,0 +1,257 @@
+import math
+import random
+import torch
+
+from torch import nn
+from typing import Optional, Tuple
+
+from torchaudio import transforms as T
+
+class PadCrop(nn.Module):
+    def __init__(self, n_samples, randomize=True):
+        super().__init__()
+        self.n_samples = n_samples
+        self.randomize = randomize
+
+    def __call__(self, signal):
+        n, s = signal.shape
+        start = 0 if (not self.randomize) else torch.randint(0, max(0, s - self.n_samples) + 1, []).item()
+        end = start + self.n_samples
+        output = signal.new_zeros([n, self.n_samples])
+        output[:, :min(s, self.n_samples)] = signal[:, start:end]
+        return output
+
+class PadCrop_Normalized_T(nn.Module):
+    
+    def __init__(self, n_samples: int, sample_rate: int, randomize: bool = True, pad: bool = True):
+
+        super().__init__()
+
+        self.n_samples = n_samples
+        self.sample_rate = sample_rate
+        self.randomize = randomize
+        self.pad = pad
+
+    def __call__(self, source: torch.Tensor) -> Tuple[torch.Tensor, float, float, int, int, torch.Tensor]:
+        
+        n_channels, n_samples = source.shape
+        
+        # Calculate bounds and offset
+        upper_bound = max(0, n_samples - self.n_samples)
+        offset = 0
+        if self.randomize and n_samples > self.n_samples:
+            offset = random.randint(0, upper_bound)
+
+        # Calculate normalized times
+        norm_denom = upper_bound + self.n_samples
+        t_start = offset / norm_denom
+        t_end = (offset + self.n_samples) / norm_denom
+
+        # Calculate timing info
+        seconds_start = math.floor(offset / self.sample_rate)
+        seconds_total = math.ceil(n_samples / self.sample_rate)
+
+        # Optimize for different cases
+        if n_samples >= self.n_samples:
+            # No padding needed - use view (zero-copy)
+            chunk = source[:, offset:offset + self.n_samples]
+            # Create full mask efficiently
+            padding_mask = torch.ones(self.n_samples, dtype=source.dtype, device=source.device)
+        elif not self.pad:
+            # No padding mode - return audio at natural length
+            chunk = source
+            padding_mask = torch.ones(n_samples, dtype=source.dtype, device=source.device)
+        else:
+            # Padding needed - create chunk and fill in-place
+            chunk = torch.zeros(n_channels, self.n_samples, dtype=source.dtype, device=source.device)
+            chunk[:, :n_samples] = source  # Use in-place assignment
+
+            # Create padding mask in-place
+            padding_mask = torch.zeros(self.n_samples, dtype=source.dtype, device=source.device)
+            padding_mask[:n_samples] = 1
+        
+        return (
+            chunk,
+            t_start,
+            t_end,
+            seconds_start,
+            seconds_total,
+            padding_mask
+        )
+
+def strip_trailing_silence(audio, sample_rate, threshold_db=-60, min_silence_duration=0.1):
+    """Strip silence from the end of an audio tensor.
+
+    Args:
+        audio: tensor [channels, samples]
+        sample_rate: audio sample rate
+        threshold_db: dB threshold below which audio is considered silent
+        min_silence_duration: minimum trailing silence duration in seconds to strip
+    Returns:
+        Truncated audio tensor [channels, trimmed_samples], or original if no significant trailing silence
+    """
+    n_samples = audio.shape[-1]
+    hop_length = max(1, int(sample_rate * 0.01))  # 10ms frames
+    min_silence_samples = int(sample_rate * min_silence_duration)
+    n_frames = n_samples // hop_length
+
+    if n_frames == 0:
+        return audio
+
+    # Work in float32 for precision
+    audio_f = audio.float()
+
+    # Reshape into frames and compute max absolute amplitude per frame across channels
+    trimmed = audio_f[:, :n_frames * hop_length]
+    frames = trimmed.reshape(audio_f.shape[0], n_frames, hop_length)
+    frame_peak = frames.abs().amax(dim=(0, 2))  # [n_frames] - max across channels and samples
+    frame_db = 20 * torch.log10(frame_peak + 1e-10)
+
+    # Find last frame above threshold
+    above_thresh = (frame_db > threshold_db).nonzero(as_tuple=True)[0]
+
+    if len(above_thresh) == 0:
+        # Entire audio is silent
+        return audio[:, :0]
+
+    last_active_frame = above_thresh[-1].item()
+    content_end = min((last_active_frame + 1) * hop_length, n_samples)
+
+    # Only strip if trailing silence is long enough
+    if (n_samples - content_end) < min_silence_samples:
+        return audio
+
+    return audio[:, :content_end]
+
+
+class PhaseFlipper(nn.Module):
+    "Randomly invert the phase of a signal"
+    def __init__(self, p=0.5):
+        super().__init__()
+        self.p = p
+    def __call__(self, signal):
+        return -signal if (random.random() < self.p) else signal
+        
+class Mono(nn.Module):
+  def __call__(self, signal):
+    return torch.mean(signal, dim=0, keepdims=True) if len(signal.shape) > 1 else signal
+
+class Stereo(nn.Module):
+  def __call__(self, signal):
+    signal_shape = signal.shape
+    # Check if it's mono
+    if len(signal_shape) == 1: # s -> 2, s
+        signal = signal.unsqueeze(0).repeat(2, 1)
+    elif len(signal_shape) == 2:
+        if signal_shape[0] == 1: #1, s -> 2, s
+            signal = signal.repeat(2, 1)
+        elif signal_shape[0] > 2: #?, s -> 2,s
+            signal = signal[:2, :]    
+
+    return signal
+
+class VolumeNorm(nn.Module):
+    "Volume normalization and augmentation of a signal [LUFS standard]"
+    def __init__(self, params=[-16, 2], sample_rate=16000, energy_threshold=1e-6):
+        super().__init__()
+        self.loudness = T.Loudness(sample_rate)
+        self.value = params[0]
+        self.gain_range = [-params[1], params[1]]
+        self.energy_threshold = energy_threshold
+
+    def __call__(self, signal):
+        """
+        signal: torch.Tensor [channels, time]
+        """
+        # avoid do normalisation for silence
+        energy = torch.mean(signal**2)
+        if energy < self.energy_threshold:
+            return signal
+        
+        input_loudness = self.loudness(signal)
+        # Generate a random target loudness within the specified range
+        target_loudness = self.value + (torch.rand(1).item() * (self.gain_range[1] - self.gain_range[0]) + self.gain_range[0])
+        delta_loudness = target_loudness - input_loudness
+        gain = torch.pow(10.0, delta_loudness / 20.0)
+        output = gain * signal
+
+        # Check for potentially clipped samples
+        if torch.max(torch.abs(output)) >= 1.0:
+            output = self.declip(output)
+
+        return output
+
+    def declip(self, signal):
+        """
+        Declip the signal by scaling down if any samples are clipped
+        """
+        max_val = torch.max(torch.abs(signal))
+        if max_val > 1.0:
+            signal = signal / max_val
+            signal *= 0.95
+        return signal
+
+
+def create_padding_mask_from_lengths(
+    valid_lengths: torch.Tensor,
+    total_seq_len: int,
+) -> torch.Tensor:
+    """
+    Create a boolean padding mask from per-batch valid sequence lengths.
+
+    Args:
+        valid_lengths: Tensor of shape (batch_size,) with valid length per sample
+        total_seq_len: Total sequence length of the latent
+
+    Returns:
+        Boolean tensor of shape (batch_size, total_seq_len) where True = valid, False = padding
+    """
+    device = valid_lengths.device
+    positions = torch.arange(total_seq_len, device=device).unsqueeze(0)  # (1, T)
+    padding_mask = positions < valid_lengths.unsqueeze(1)  # (B, T)
+    return padding_mask
+
+
+def compute_effective_seq_len_from_conditioning(
+    conditioning: list,
+    sample_rate: int,
+    downsampling_ratio: int = 1,
+    device: str = "cuda"
+) -> Optional[torch.Tensor]:
+    """
+    Compute effective sequence lengths from seconds_total in conditioning dicts.
+
+    Args:
+        conditioning: List of conditioning dicts, one per batch element
+        sample_rate: Audio sample rate
+        downsampling_ratio: Pretransform downsampling ratio (1 if no pretransform)
+        device: Device to place the tensor on
+
+    Returns:
+        Tensor of shape (batch_size,) with effective sequence lengths in latent space,
+        or None if seconds_total is not present in conditioning
+    """
+    if conditioning is None:
+        return None
+
+    # Check if seconds_total is present in any conditioning dict
+    if not any("seconds_total" in c for c in conditioning):
+        return None
+
+    effective_lengths = []
+    for cond_dict in conditioning:
+        if "seconds_total" in cond_dict:
+            seconds = cond_dict["seconds_total"]
+            # Convert seconds to latent sequence length
+            audio_samples = int(seconds * sample_rate)
+            latent_length = math.ceil(audio_samples / downsampling_ratio)
+            effective_lengths.append(latent_length)
+        else:
+            # If seconds_total not present for this item, use None as marker
+            effective_lengths.append(None)
+
+    # If any item is missing seconds_total, return None to fall back to full length
+    if any(l is None for l in effective_lengths):
+        return None
+
+    return torch.tensor(effective_lengths, dtype=torch.float32, device=device)

stable_audio_3/factory.py

@@ -0,0 +1,157 @@
+import typing as tp
+
+from stable_audio_3.models.diffusion import DiTWrapper, ConditionedDiffusionModelWrapper
+from stable_audio_3.models.autoencoders import (
+    AudioAutoencoder,
+    SAMEEncoder,
+    SAMEDecoder,
+)
+from stable_audio_3.models.conditioners import (
+    MultiConditioner,
+    NumberConditioner,
+    T5GemmaConditioner,
+)
+from stable_audio_3.models.bottleneck import SoftNormBottleneck
+from stable_audio_3.models.pretransforms import (
+    PatchedPretransform,
+    AutoencoderPretransform,
+)
+
+
+def create_diffusion_cond_from_config(config: tp.Dict[str, tp.Any]):
+    model_config = config["model"]
+    diffusion_config = model_config.get("diffusion", None)
+    diffusion_model_config = diffusion_config.get("config", None)
+    diffusion_objective = diffusion_config.get("diffusion_objective", "v")
+    modular_local_cond_configs = diffusion_config.get("modular_local_cond_configs", [])
+
+    diffusion_model = DiTWrapper(
+        diffusion_objective=diffusion_objective,
+        modular_local_cond_configs=modular_local_cond_configs,
+        **diffusion_model_config,
+    )
+
+    io_channels = model_config.get("io_channels", None)
+    sample_rate = config.get("sample_rate", None)
+    cross_attention_ids = diffusion_config.get("cross_attention_cond_ids", [])
+    global_cond_ids = diffusion_config.get("global_cond_ids", [])
+    input_concat_ids = diffusion_config.get("input_concat_ids", [])
+    local_add_cond_ids = diffusion_config.get("local_add_cond_ids", [])
+    modular_local_cond_ids = [c["id"] for c in modular_local_cond_configs]
+    prepend_cond_ids = diffusion_config.get("prepend_cond_ids", [])
+
+    distribution_shift_options = diffusion_config.get(
+        "distribution_shift_options", None
+    )
+    sampling_distribution_shift_options = diffusion_config.get(
+        "sampling_distribution_shift_options", None
+    )
+    mask_padding_attention = diffusion_config.get("mask_padding_attention", False)
+    use_effective_length_for_schedule = diffusion_config.get(
+        "use_effective_length_for_schedule", False
+    )
+
+    pretransform = create_pretransform_from_config(model_config, sample_rate)
+    min_input_length = pretransform.downsampling_ratio
+
+    conditioning_config = model_config.get("conditioning", None)
+
+    conditioner = create_multi_conditioner_from_conditioning_config(conditioning_config)
+
+    min_input_length *= diffusion_model.model.patch_size
+
+    return ConditionedDiffusionModelWrapper(
+        diffusion_model,
+        conditioner,
+        min_input_length=min_input_length,
+        sample_rate=sample_rate,
+        cross_attn_cond_ids=cross_attention_ids,
+        global_cond_ids=global_cond_ids,
+        input_concat_ids=input_concat_ids,
+        local_add_cond_ids=local_add_cond_ids,
+        modular_local_cond_ids=modular_local_cond_ids,
+        prepend_cond_ids=prepend_cond_ids,
+        pretransform=pretransform,
+        io_channels=io_channels,
+        distribution_shift_options=distribution_shift_options,
+        sampling_distribution_shift_options=sampling_distribution_shift_options,
+        mask_padding_attention=mask_padding_attention,
+        use_effective_length_for_schedule=use_effective_length_for_schedule,
+        diffusion_objective=diffusion_objective,
+    )
+
+
+def create_autoencoder_from_config(config, sample_rate):
+    # AE-only configs (SAME-S/SAME-L) have encoder/decoder at the top of config["model"].
+    # Full SA3 configs nest them inside config["model"]["pretransform"]["config"].
+    if "encoder" in config:
+        autoencoder_config = config
+    else:
+        autoencoder_config = config["pretransform"]["config"]
+    encoder = SAMEEncoder(**autoencoder_config["encoder"]["config"])
+    decoder = SAMEDecoder(**autoencoder_config["decoder"]["config"])
+
+    latent_dim = autoencoder_config.get("latent_dim", None)
+    downsampling_ratio = autoencoder_config.get("downsampling_ratio", None)
+    io_channels = autoencoder_config.get("io_channels", None)
+    in_channels = autoencoder_config.get("in_channels", None)
+    out_channels = autoencoder_config.get("out_channels", None)
+    soft_clip = autoencoder_config["decoder"].get("soft_clip", False)
+    pretransform = PatchedPretransform(**autoencoder_config["pretransform"]["config"])
+    bottleneck = SoftNormBottleneck(**autoencoder_config["bottleneck"]["config"])
+
+    return AudioAutoencoder(
+        encoder,
+        decoder,
+        io_channels=io_channels,
+        latent_dim=latent_dim,
+        downsampling_ratio=downsampling_ratio,
+        sample_rate=sample_rate,
+        bottleneck=bottleneck,
+        pretransform=pretransform,
+        in_channels=in_channels,
+        out_channels=out_channels,
+        soft_clip=soft_clip,
+    )
+
+
+def create_pretransform_from_config(config, sample_rate):
+    pretransform_block = config["pretransform"]
+    autoencoder = create_autoencoder_from_config(config, sample_rate)
+    chunked = pretransform_block.get("chunked", False)
+    iterate_batch = pretransform_block.get("iterate_batch", False)
+    return AutoencoderPretransform(
+        autoencoder, chunked=chunked, iterate_batch=iterate_batch
+    )
+
+
+def create_multi_conditioner_from_conditioning_config(
+    config: tp.Dict[str, tp.Any],
+) -> MultiConditioner:
+    """Create a MultiConditioner from a conditioning config dictionary."""
+    conditioners = {}
+    cond_dim = config["cond_dim"]
+
+    default_keys = config.get("default_keys", {})
+
+    pre_encoded_keys = config.get("pre_encoded_keys", [])
+
+    for conditioner_info in config["configs"]:
+        id = conditioner_info["id"]
+
+        conditioner_type = conditioner_info["type"]
+
+        conditioner_config = {"output_dim": cond_dim}
+
+        conditioner_config.update(conditioner_info["config"])
+
+        if conditioner_type == "t5gemma":
+            conditioners[id] = T5GemmaConditioner(**conditioner_config)
+        elif conditioner_type == "number":
+            conditioners[id] = NumberConditioner(**conditioner_config)
+        else:
+            raise ValueError(f"Unknown conditioner type: {conditioner_type}")
+
+    return MultiConditioner(
+        conditioners, default_keys=default_keys, pre_encoded_keys=pre_encoded_keys
+    )

stable_audio_3/inference/audio_utils.py

@@ -0,0 +1,98 @@
+import numpy as np
+import torch
+from torch import nn
+from ..data.utils import PadCrop
+
+from torchaudio import transforms as T
+
+def numpy_audio_to_tensor(audio):
+    """Convert a numpy audio array to a float32 torch tensor with shape [channels, samples]."""
+    if audio.dtype == np.float32:
+        t = torch.from_numpy(audio)
+    elif audio.dtype == np.int16:
+        t = torch.from_numpy(audio).float().div(32767)
+    elif audio.dtype == np.int32:
+        t = torch.from_numpy(audio).float().div(2147483647)
+    else:
+        raise ValueError(f"Unsupported audio dtype: {audio.dtype}")
+    if t.dim() == 1:
+        t = t.unsqueeze(0)       # [n] -> [1, n]
+    elif t.dim() == 2:
+        t = t.transpose(0, 1)   # [n, channels] -> [channels, n]
+    return t
+
+
+def preprocess_audio(sr_audio, sample_rate, sample_size, half=False):
+    """Normalize a (sr, numpy_or_tensor) tuple into (sample_rate, tensor) ready for the model."""
+    in_sr, audio = sr_audio
+
+    if isinstance(audio, np.ndarray):
+        audio = numpy_audio_to_tensor(audio)
+    elif audio.dim() == 1:
+        audio = audio.unsqueeze(0)
+    elif audio.dim() == 2:
+        audio = audio.transpose(0, 1)  # [n, 2] -> [2, n]
+
+    if half:
+        audio = audio.to(torch.float16)
+
+    if in_sr != sample_rate:
+        resample_tf = T.Resample(in_sr, sample_rate).to(audio.device).to(audio.dtype)
+        audio = resample_tf(audio)
+
+    if audio.shape[-1] > sample_size:
+        audio = audio[:, :sample_size]
+
+    return (sample_rate, audio)
+
+
+def set_audio_channels(audio, target_channels):
+    # Add channel dim if it's missing
+    if audio.dim() == 2:
+        audio = audio.unsqueeze(1)
+        
+    if target_channels == 1:
+        # Convert to mono
+        audio = audio.mean(1, keepdim=True)
+    elif target_channels == 2:
+        # Convert to stereo
+        if audio.shape[1] == 1:
+            audio = audio.repeat(1, 2, 1)
+        elif audio.shape[1] > 2:
+            audio = audio[:, :2, :]
+    return audio
+
+def prepare_audio(audio, in_sr, target_sr, target_length, target_channels, device):
+    
+    audio = audio.to(device)
+
+    if in_sr != target_sr:
+        resample_tf = T.Resample(in_sr, target_sr).to(device)
+        audio = resample_tf(audio)
+
+    audio = PadCrop(target_length, randomize=False)(audio)
+
+    # Add batch dimension
+    if audio.dim() == 1:
+        audio = audio.unsqueeze(0).unsqueeze(0)
+    elif audio.dim() == 2:
+        audio = audio.unsqueeze(0)
+
+    audio = set_audio_channels(audio, target_channels)
+
+    return audio
+
+
+class PadCrop(nn.Module):
+    def __init__(self, n_samples, randomize=True):
+        super().__init__()
+        self.n_samples = n_samples
+        self.randomize = randomize
+
+    def __call__(self, signal):
+        n, s = signal.shape
+        start = 0 if (not self.randomize) else torch.randint(0, max(0, s - self.n_samples) + 1, []).item()
+        end = start + self.n_samples
+        output = signal.new_zeros([n, self.n_samples])
+        output[:, :min(s, self.n_samples)] = signal[:, start:end]
+        return output

stable_audio_3/inference/distribution_shift.py

@@ -0,0 +1,198 @@
+import math
+import torch
+import typing as tp
+
+
+class IdentityDistributionShift:
+    """No-op distribution shift — returns timesteps unchanged."""
+    def shift(self, t: torch.Tensor, seq_len):
+        return t
+
+
+class FluxDistributionShift:
+    """Flux/SD3/Self-Flow timestep shift: t_shifted = alpha * t / (1 + (alpha-1) * t).
+
+    Convention: t=0 is data, t=1 is noise.
+    alpha > 1 shifts timesteps toward noise, appropriate for longer sequences
+    where the critical structure-from-noise transition happens at higher noise levels.
+
+    Can be used in two ways:
+    - Constant alpha: set alpha_min == alpha_max. This is how the Self-Flow paper
+      (BFL, 2025) uses it, with alpha chosen per modality/autoencoder.
+      Reference values from the paper: audio sampleshift=6.93, trainshift=1.0;
+      video sampleshift=15.0, trainshift=2.95; images sampleshift=1.78-6.93.
+    - Seq_len-dependent alpha: set different alpha_min/alpha_max. Alpha is
+      interpolated log-linearly in seq_len space (power-law), following the
+      SD3 derivation where alpha ∝ sqrt(seq_len).
+
+    Args:
+        min_length: Minimum sequence length (alpha = alpha_min here)
+        max_length: Maximum sequence length (alpha = alpha_max here)
+        alpha_min: Shift factor at min_length (1.0 = no shift)
+        alpha_max: Shift factor at max_length (1.0 = no shift)
+    """
+    def __init__(self, min_length=256, max_length=4096,
+                 alpha_min=1.0, alpha_max=1.0):
+        self.min_length = min_length
+        self.max_length = max_length
+        self.alpha_min = alpha_min
+        self.alpha_max = alpha_max
+        # Precompute for log-linear interpolation
+        self.log_alpha_min = math.log(max(alpha_min, 1e-8))
+        self.log_alpha_max = math.log(max(alpha_max, 1e-8))
+        self.log_min_seq = math.log(min_length)
+        self.log_max_seq = math.log(max_length)
+        if self.log_max_seq == self.log_min_seq:
+            self.log_max_seq += 1e-8  # prevent division by zero for constant alpha
+
+    def get_alpha(self, seq_len: tp.Union[int, torch.Tensor]):
+        """Compute alpha via log-linear interpolation in seq_len."""
+        if isinstance(seq_len, torch.Tensor):
+            seq_len = seq_len.float().clamp(self.min_length, self.max_length)
+            log_seq = torch.log(seq_len)
+            frac = (log_seq - self.log_min_seq) / (self.log_max_seq - self.log_min_seq)
+            log_alpha = self.log_alpha_min + frac * (self.log_alpha_max - self.log_alpha_min)
+            return torch.exp(log_alpha)
+        else:
+            seq_len = max(min(seq_len, self.max_length), self.min_length)
+            log_seq = math.log(seq_len)
+            frac = (log_seq - self.log_min_seq) / (self.log_max_seq - self.log_min_seq)
+            log_alpha = self.log_alpha_min + frac * (self.log_alpha_max - self.log_alpha_min)
+            return math.exp(log_alpha)
+
+    def shift(self, t: torch.Tensor, seq_len: tp.Union[int, torch.Tensor]):
+        """Shift timesteps based on sequence length.
+
+        Args:
+            t: Timesteps tensor of shape (batch_size,) or (steps,)
+            seq_len: Either a scalar int (same shift for all elements) or
+                     tensor of shape (batch_size,) for per-element shifts
+        Returns:
+            Shifted timesteps. If seq_len is a tensor and t is 1D with different size,
+            returns shape (batch_size, steps) for per-element schedules.
+        """
+        alpha = self.get_alpha(seq_len)
+
+        if isinstance(seq_len, torch.Tensor):
+            alpha = alpha.to(t.device)
+            if t.dim() == 1 and alpha.dim() == 1 and t.shape[0] != alpha.shape[0]:
+                t = t.unsqueeze(0)
+                alpha = alpha.unsqueeze(1)
+
+        return alpha * t / (1 + (alpha - 1.0) * t)
+
+
+class DistributionShift:
+    def __init__(self, base_shift=0.5, max_shift=1.15, max_length=4096, min_length=256, use_sine=False):
+        self.base_shift = base_shift
+        self.max_shift = max_shift
+        self.max_length = max_length
+        self.min_length = min_length
+        self.use_sine = use_sine
+
+    def shift(self, t: torch.Tensor, seq_len: tp.Union[int, torch.Tensor]):
+        """
+        Shift timesteps based on sequence length to adjust noise schedule.
+
+        Args:
+            t: Timesteps tensor of shape (batch_size,) or (steps,)
+            seq_len: Either a scalar int (same shift for all elements) or
+                     tensor of shape (batch_size,) for per-element shifts
+        Returns:
+            Shifted timesteps. If seq_len is a tensor and t is 1D with different size,
+            returns shape (batch_size, steps) for per-element schedules.
+        """
+        if isinstance(seq_len, torch.Tensor):
+            # Per-element sequence lengths
+            # Ensure seq_len is on the same device as t
+            seq_len = seq_len.to(t.device)
+            seq_len_clamped = seq_len.float().clamp(self.min_length, self.max_length)
+            # Handle broadcasting when t and seq_len have different sizes
+            if t.dim() == 1 and seq_len_clamped.dim() == 1 and t.shape[0] != seq_len_clamped.shape[0]:
+                # t: (steps,) -> (1, steps), seq_len: (batch,) -> (batch, 1)
+                # Result: (batch, steps)
+                t = t.unsqueeze(0)
+                seq_len_clamped = seq_len_clamped.unsqueeze(1)
+            sigma = 1.0
+            mu = - (self.base_shift + (self.max_shift - self.base_shift) * (seq_len_clamped - self.min_length) / (self.max_length - self.min_length))
+            t_out = 1 - torch.exp(mu) / (torch.exp(mu) + (1 / (1 - t) - 1) ** sigma)
+            if self.use_sine:
+                t_out = torch.sin(t_out * math.pi / 2)
+        else:
+            # Scalar path (original behavior)
+            seq_len = min(max(seq_len, self.min_length), self.max_length)
+            sigma = 1.0
+            mu = - (self.base_shift + (self.max_shift - self.base_shift) * (seq_len - self.min_length) / (self.max_length - self.min_length))
+            t_out = 1 - math.exp(mu) / (math.exp(mu) + (1 / (1 - t) - 1) ** sigma)
+
+            if self.use_sine:
+                t_out = torch.sin(t_out * math.pi / 2)
+
+        return t_out
+
+
+class LogSNRShift:
+    """Adaptive log-SNR distribution shift.
+
+    Maps t∈[0,1] to log-SNR-spaced values while preserving order (0→0, 1→1).
+    Equivalent to applying: logsnr = linspace(logsnr_end, logsnr_start, N)
+    then t = sigmoid(-logsnr), which spaces steps uniformly in log-SNR.
+
+    logsnr_start (the high-t bound) scales with sequence length following
+    the "-1 per doubling" rule:
+        logsnr_start = anchor_logsnr - rate * log₂(seq_len / anchor_length)
+
+    This captures the empirical finding that the critical log-SNR point
+    (where structure emerges from noise) drops by ~rate for each doubling
+    of sequence length. logsnr_end (the low-t bound) is fixed because
+    low-t refinement is purely local.
+    """
+
+    def __init__(self, anchor_length=2000, anchor_logsnr=-6.2,
+                 rate=1.0, logsnr_end=2.0):
+        self.anchor_length = anchor_length
+        self.anchor_logsnr = anchor_logsnr
+        self.rate = rate
+        self.logsnr_end = logsnr_end
+
+    def get_logsnr_start(self, seq_len):
+        """Compute adaptive logsnr_start: drops by `rate` per doubling of seq_len."""
+        if isinstance(seq_len, torch.Tensor):
+            log2_ratio = torch.log2(seq_len.float() / self.anchor_length)
+            return self.anchor_logsnr - self.rate * log2_ratio
+        else:
+            log2_ratio = math.log2(seq_len / self.anchor_length)
+            return self.anchor_logsnr - self.rate * log2_ratio
+
+    def shift(self, t: torch.Tensor, seq_len: tp.Union[int, torch.Tensor]):
+        """Transform t∈[0,1] to log-SNR-spaced t with adaptive bounds.
+
+        Maps through: logsnr = logsnr_end - t * (logsnr_end - logsnr_start)
+                      t_out = sigmoid(-logsnr)
+
+        Preserves order: 0→~0, 1→~1, with exact endpoint preservation.
+
+        Args:
+            t: Timesteps tensor of shape (batch_size,) or (steps,)
+            seq_len: Either a scalar int or tensor of shape (batch_size,)
+        Returns:
+            Log-SNR-spaced timesteps in [0, 1].
+        """
+        t_original = t
+        logsnr_start = self.get_logsnr_start(seq_len)
+
+        if isinstance(seq_len, torch.Tensor):
+            logsnr_start = logsnr_start.to(t.device)
+            if t.dim() == 1 and logsnr_start.dim() == 1 and t.shape[0] != logsnr_start.shape[0]:
+                t = t.unsqueeze(0)
+                logsnr_start = logsnr_start.unsqueeze(1)
+
+        # Map t through log-SNR space (monotonically: low t → high logsnr → low t_out)
+        logsnr = self.logsnr_end - t * (self.logsnr_end - logsnr_start)
+        t_out = torch.sigmoid(-logsnr)
+
+        # Preserve exact endpoints
+        t_out = torch.where(t_original <= 0, torch.zeros_like(t_out), t_out)
+        t_out = torch.where(t_original >= 1, torch.ones_like(t_out), t_out)
+
+        return t_out

stable_audio_3/inference/sampling.py

@@ -0,0 +1,523 @@
+import torch
+import typing as tp
+from tqdm import trange, tqdm
+import torch.distributions as dist
+
+from ..data.utils import create_padding_mask_from_lengths, compute_effective_seq_len_from_conditioning
+
+
+def build_schedule(
+    steps: int,
+    sigma_max: float = 1.0,
+    dist_shift = None,
+    effective_seq_len: tp.Union[int, torch.Tensor, None] = None,
+    fallback_seq_len: tp.Optional[int] = None,
+    include_endpoint: bool = True,
+    device: tp.Union[str, torch.device] = "cpu",
+) -> torch.Tensor:
+    """Build a timestep schedule for diffusion sampling.
+
+    Returns a 1D tensor of shape (N,) where N = steps+1 (if include_endpoint)
+    or steps (if not), OR a 2D tensor of shape (batch_size, N) when
+    effective_seq_len is a tensor and dist_shift produces per-element schedules.
+
+    Args:
+        steps: Number of sampling steps.
+        sigma_max: Starting noise level (1.0 for full generation, <1.0 for variations).
+        dist_shift: Optional distribution shift object (FluxDistributionShift,
+            DistributionShift, LogSNRShift, etc.). Applied to warp the linear schedule.
+        effective_seq_len: Sequence length for dist_shift. Scalar int or
+            tensor of shape (batch_size,) for per-element schedules.
+        fallback_seq_len: Fallback when effective_seq_len is None (typically x.shape[-1]).
+        include_endpoint: If True, schedule includes 0 as final value (RF samplers).
+            If False, excludes 0 (v-diffusion DDIM).
+        device: Device for the output tensor.
+    """
+    n_points = steps + 1 if include_endpoint else steps
+
+    if include_endpoint:
+        t = torch.linspace(sigma_max, 0, n_points, device=device)
+    else:
+        t = torch.linspace(sigma_max, 0, n_points + 1, device=device)[:-1]
+
+    if dist_shift is not None:
+        seq_len = effective_seq_len if effective_seq_len is not None else fallback_seq_len
+        if isinstance(seq_len, torch.Tensor):
+            # Clamp per-element sequence lengths to avoid zeros causing log/NaN issues
+            seq_len = torch.clamp(seq_len, min=1)
+        elif seq_len is not None:
+            # Clamp scalar sequence length to at least 1
+            seq_len = max(int(seq_len), 1)
+        t = dist_shift.shift(t, seq_len)
+
+        # Ensure the first timestep remains aligned with sigma_max after shifting.
+        # This keeps the schedule consistent with the initialization in sample_diffusion(),
+        # which mixes init_data using sigma_max.
+        if isinstance(t, torch.Tensor):
+            sigma_max_tensor = t.new_tensor(sigma_max)
+            if t.ndim == 1:
+                t[0] = sigma_max_tensor
+            else:
+                # For batched/per-element schedules, enforce sigma_max at the first time index.
+                t[..., 0] = sigma_max_tensor
+
+    return t
+
+
+def sample_timesteps_logsnr(batch_size, mean_logsnr=-1.2, std_logsnr=2.0):
+    """
+    Sample timesteps for diffusion training by sampling logSNR values and converting to t.
+
+    Args:
+        batch_size (int): Number of timesteps to sample
+        mean_logsnr (float): Mean of the logSNR Gaussian distribution
+        std_logsnr (float): Standard deviation of the logSNR Gaussian distribution
+
+    Returns:
+        torch.Tensor: Tensor of shape (batch_size,) containing timestep values t in [0, 1]
+    """
+    # Sample logSNR from Gaussian distribution
+    logsnr = torch.randn(batch_size) * std_logsnr + mean_logsnr
+
+    # Convert logSNR to timesteps using the logistic function
+    # Since logSNR = ln((1-t)/t), we can solve for t:
+    # t = 1 / (1 + exp(logsnr))
+    t = torch.sigmoid(-logsnr)
+
+    # Clamp values to ensure numerical stability
+    t = t.clamp(1e-4, 1 - 1e-4)
+
+    return t
+
+def sample_timesteps_logsnr_uniform(batch_size, min_logsnr=-6, max_logsnr=5.0):
+    """
+    Sample timesteps for diffusion training by sampling logSNR values and converting to t.
+
+    Args:
+        batch_size (int): Number of timesteps to sample
+        min_logsnr (float): Minimum logSNR value
+        max_logsnr (float): Maximum logSNR value
+
+    Returns:
+        torch.Tensor: Tensor of shape (batch_size,) containing timestep values t in [0, 1]
+    """
+    # Sample logSNR from uniform distribution
+    logsnr = torch.rand(batch_size) * (max_logsnr - min_logsnr) + min_logsnr
+
+    # Convert logSNR to timesteps using the logistic function
+    # Since logSNR = ln((1-t)/t), we can solve for t:
+    # t = 1 / (1 + exp(logsnr))
+    t = torch.sigmoid(-logsnr)
+
+    # Clamp values to ensure numerical stability
+    t = t.clamp(1e-4, 1 - 1e-4)
+
+    return t
+
+def truncated_logistic_normal_rescaled(shape, left_trunc=0.075, right_trunc=1):
+    """
+
+    shape: shape of the output tensor
+    left_trunc: left truncation point, fraction of probability to be discarded
+    right_trunc: right truncation boundary, should be 1 (never seen at test time)
+    """
+
+    # Step 1: Sample from the logistic normal distribution (sigmoid of normal)
+    logits = torch.randn(shape)
+
+    # Step 2: Apply the CDF transformation of the normal distribution
+    normal_dist = dist.Normal(0, 1)
+    cdf_values = normal_dist.cdf(logits)
+
+    # Step 3: Define the truncation bounds on the CDF
+    lower_bound = normal_dist.cdf(torch.logit(torch.tensor(left_trunc)))
+    upper_bound = normal_dist.cdf(torch.logit(torch.tensor(right_trunc)))
+
+    # Step 4: Rescale linear CDF values into the truncated region (between lower_bound and upper_bound)
+    truncated_cdf_values = lower_bound + (upper_bound - lower_bound) * cdf_values
+
+    # Step 5: Map back to logistic-normal space using inverse CDF
+    truncated_samples = torch.sigmoid(normal_dist.icdf(truncated_cdf_values))
+
+    # Step 6: Rescale values so that min is 0 and max is just below 1
+    rescaled_samples = (truncated_samples - left_trunc) / (right_trunc - left_trunc)
+
+    return rescaled_samples
+
+def sample_discrete_euler(model, x, sigmas, callback=None, disable_tqdm=False, **extra_args):
+    """Draws samples from a model given starting noise. Euler method
+
+    Args:
+        sigmas: Pre-computed schedule tensor. Shape (steps+1,) for global schedule
+            or (batch_size, steps+1) for per-element schedules.
+    """
+    t = sigmas
+
+    # Check if we have per-element schedules (batch_size, steps+1) or global schedule (steps+1,)
+    per_element_schedule = t.dim() == 2
+
+    t = t.to(x.device)
+    num_steps = t.shape[-1] - 1
+
+    for i in tqdm(range(num_steps), disable=disable_tqdm):
+        if per_element_schedule:
+            # Per-element schedules: t has shape (batch_size, steps+1)
+            t_curr_tensor = t[:, i].to(x.dtype)  # (batch_size,)
+            t_prev = t[:, i + 1].to(x.dtype)  # (batch_size,)
+            dt = t_prev - t_curr_tensor  # (batch_size,)
+            # Reshape for broadcasting with x: (batch_size,) -> (batch_size, 1, 1)
+            dt_broadcast = dt.view(-1, 1, 1)
+        else:
+            # Global schedule: t has shape (steps+1,)
+            t_curr = t[i]
+            t_prev = t[i + 1]
+            t_curr_tensor = t_curr * torch.ones((x.shape[0],), dtype=x.dtype, device=x.device)
+            dt = t_prev - t_curr
+            dt_broadcast = dt
+
+        v = model(x, t_curr_tensor, **extra_args)
+
+        if callback is not None:
+            denoised = x - t_curr_tensor[:, None, None] * v
+            callback({'x': x, 't': t_curr_tensor, 'sigma': t_curr_tensor, 'i': i, 'denoised': denoised})
+
+        x = x + dt_broadcast * v
+
+    # If we are on the last timestep, output the denoised data
+    return x
+
+def sample_rk4(model, x, sigmas, callback=None, disable_tqdm=False, **extra_args):
+    """Draws samples from a model given starting noise. 4th-order Runge-Kutta
+
+    Args:
+        sigmas: Pre-computed schedule tensor of shape (steps+1,).
+            Per-element schedules not supported for RK4.
+    """
+    # Make tensor of ones to broadcast the single t values
+    ts = x.new_ones([x.shape[0]])
+
+    t = sigmas
+
+    t = t.to(x.device)
+
+    for i, (t_curr, t_prev) in enumerate(tqdm(zip(t[:-1], t[1:]), disable=disable_tqdm)):
+        # Broadcast the current timestep to the correct shape
+        t_curr_tensor = t_curr * ts
+        dt = t_prev - t_curr  # we solve backwards in our formulation
+
+        k1 = model(x, t_curr_tensor, **extra_args)
+
+        if callback is not None:
+            denoised = x - t_curr * k1
+            callback({'x': x, 't': t_curr, 'sigma': t_curr, 'i': i, 'denoised': denoised})
+
+        k2 = model(x + dt / 2 * k1, (t_curr + dt / 2) * ts, **extra_args)
+        k3 = model(x + dt / 2 * k2, (t_curr + dt / 2) * ts, **extra_args)
+
+        # Clamp t_prev to avoid evaluating model at exactly t=0
+        # (models aren't trained at t=0 and may return garbage/NaN)
+        t_prev_eval = t_prev.clamp(min=1e-5)
+        k4 = model(x + dt * k3, t_prev_eval * ts, **extra_args)
+
+        x = x + dt / 6 * (k1 + 2 * k2 + 2 * k3 + k4)
+
+    # If we are on the last timestep, output the denoised data
+    return x
+
+def sample_flow_dpmpp(model, x, sigmas, callback=None, disable_tqdm=False, **extra_args):
+    """Draws samples from a model given starting noise. DPM-Solver++ for RF models
+
+    Args:
+        sigmas: Pre-computed schedule tensor. Shape (steps+1,) for global schedule
+            or (batch_size, steps+1) for per-element schedules.
+    """
+    t = sigmas
+
+    # Check if we have per-element schedules (batch_size, steps+1) or global schedule (steps+1,)
+    per_element_schedule = t.dim() == 2
+
+    t = t.to(x.device)
+    num_steps = t.shape[-1] - 1
+
+    old_denoised = None
+
+    # Clamp t to avoid numerical issues with log(0) and division by zero
+    # This prevents inf/-inf values that can cause NaN propagation
+    log_snr = lambda t: ((1-t).clamp(min=1e-10) / t.clamp(min=1e-10)).log()
+
+    for i in trange(num_steps, disable=disable_tqdm):
+        if per_element_schedule:
+            # Per-element schedules: t has shape (batch_size, steps+1)
+            t_curr = t[:, i]  # (batch_size,)
+            t_next = t[:, i + 1]  # (batch_size,)
+            t_prev = t[:, i - 1] if i > 0 else None
+            # Reshape for broadcasting with x: (batch_size,) -> (batch_size, 1, 1)
+            t_curr_broadcast = t_curr.view(-1, 1, 1)
+            t_next_broadcast = t_next.view(-1, 1, 1)
+            t_curr_tensor = t_curr  # already (batch_size,)
+        else:
+            # Global schedule: t has shape (steps+1,)
+            t_curr = t[i]
+            t_next = t[i + 1]
+            t_prev = t[i - 1] if i > 0 else None
+            t_curr_broadcast = t_curr
+            t_next_broadcast = t_next
+            t_curr_tensor = t_curr.expand(x.shape[0])
+
+        model_output = model(x, t_curr_tensor, **extra_args)
+        denoised = x - t_curr_broadcast * model_output
+
+        if callback is not None:
+            callback({'x': x, 'i': i, 't': t_curr, 'sigma': t_curr, 'denoised': denoised})
+
+        alpha_t = 1 - t_next_broadcast
+
+        # For rectified flow, compute the DPM++ coefficient directly without log_snr
+        # to avoid numerical issues at t=0 or t=1
+        # The formula is: (-h).expm1() = (t_next - t_curr) / [(1 - t_next) * t_curr]
+        # Note: t_next < t_curr, so this is negative
+        # We'll compute this directly instead of going through log_snr
+        dt = t_next_broadcast - t_curr_broadcast
+        # Clamp to avoid division by zero when t_curr or t_next are at boundaries
+        dpmpp_coeff = dt / ((1 - t_next_broadcast).clamp(min=1e-10) * t_curr_broadcast.clamp(min=1e-10))
+
+        # Check if this is the first step or the last step (t_next == 0)
+        is_first_step = old_denoised is None
+        is_last_step = (t_next_broadcast == 0).all() if per_element_schedule else (t_next == 0)
+
+        if is_first_step or is_last_step:
+            # First-order update using the directly computed coefficient
+            x = (t_next_broadcast / t_curr_broadcast.clamp(min=1e-10)) * x - alpha_t * dpmpp_coeff * denoised
+        else:
+            # Second-order update with Richardson extrapolation
+            if per_element_schedule:
+                t_prev_broadcast = t_prev.view(-1, 1, 1)
+            else:
+                t_prev_broadcast = t_prev
+            # Compute r = h_last / h in log-SNR space for second-order correction
+            # h = log_snr(t_next) - log_snr(t_curr), h_last = log_snr(t_curr) - log_snr(t_prev)
+            h = log_snr(t_next_broadcast) - log_snr(t_curr_broadcast)
+            h_last = log_snr(t_curr_broadcast) - log_snr(t_prev_broadcast)
+            r = h_last / h
+            denoised_d = (1 + 1 / (2 * r)) * denoised - (1 / (2 * r)) * old_denoised
+            x = (t_next_broadcast / t_curr_broadcast.clamp(min=1e-10)) * x - alpha_t * dpmpp_coeff * denoised_d
+
+        old_denoised = denoised
+    return x
+
+def sample_flow_pingpong(model, x, sigmas, callback=None, disable_tqdm=False, **extra_args):
+    """Draws samples from a model given starting noise. Ping-pong sampling for distilled models
+
+    Args:
+        sigmas: Pre-computed schedule tensor. Shape (steps+1,) for global schedule
+            or (batch_size, steps+1) for per-element schedules.
+    """
+    t = sigmas
+
+    # Check if we have per-element schedules (batch_size, steps+1) or global schedule (steps+1,)
+    per_element_schedule = t.dim() == 2
+
+    t = t.to(x.device)
+    num_steps = t.shape[-1] - 1
+
+    for i in trange(num_steps, disable=disable_tqdm):
+        if per_element_schedule:
+            # Per-element schedules: t has shape (batch_size, steps+1)
+            t_curr = t[:, i].to(x.dtype)  # (batch_size,)
+            t_next = t[:, i + 1].to(x.dtype)  # (batch_size,)
+            # Reshape for broadcasting with x: (batch_size,) -> (batch_size, 1, 1)
+            t_curr_broadcast = t_curr.view(-1, 1, 1)
+            t_next_broadcast = t_next.view(-1, 1, 1)
+        else:
+            # Global schedule: t has shape (steps+1,)
+            t_curr = t[i].to(x.dtype)
+            t_next = t[i + 1].to(x.dtype)
+            t_curr_broadcast = t_curr
+            t_next_broadcast = t_next
+
+        # Model forward
+        if per_element_schedule:
+            t_curr_tensor = t_curr  # already (batch_size,)
+        else:
+            t_curr_tensor = t_curr * torch.ones((x.shape[0],), dtype=x.dtype, device=x.device)
+
+        denoised = x - t_curr_broadcast * model(x, t_curr_tensor, **extra_args)
+
+        if callback is not None:
+            callback({'x': x, 'i': i, 't': t_curr, 'sigma': t_curr, 'sigma_hat': t_curr, 'denoised': denoised})
+
+        x = (1 - t_next_broadcast) * denoised + t_next_broadcast * torch.randn_like(x)
+
+    return x
+
+
+
+@torch.no_grad()
+def sample_diffusion(
+    model,
+    noise: torch.Tensor,
+    cond_inputs: dict,
+    diffusion_objective: str,
+    steps: int,
+    cfg_scale: float = 1.0,
+    # Varlen support
+    conditioning: tp.Optional[tp.List[dict]] = None,
+    sample_rate: int = 44100,
+    pretransform = None,
+    mask_padding_attention: bool = False,
+    use_effective_length_for_schedule: bool = False,
+    headroom_seconds: float = 5.0,
+    padding_mask: tp.Optional[torch.Tensor] = None,
+    # Timestep schedule
+    dist_shift = None,
+    # Sampler options
+    sampler_type: str = None,
+    batch_cfg: bool = True,
+    rescale_cfg: bool = False,
+    # CFG options
+    apg_scale: float = 1.0,
+    # Init data (variation / img2img)
+    init_data: tp.Optional[torch.Tensor] = None,
+    init_noise_level: float = 1.0,
+    # Other
+    callback = None,
+    disable_tqdm: bool = False,
+    decode: bool = True,
+    chunked_decode: tp.Optional[bool] = None,
+    **sampler_kwargs
+) -> torch.Tensor:
+    """
+    Unified sampling function for diffusion models. Handles all diffusion objectives,
+    varlen support (padding_mask + effective_seq_len), timestep scheduling, and init_data
+    for variation/img2img.
+
+    Args:
+        model: The diffusion model backbone (model.model, not the wrapper)
+        noise: Initial noise tensor of shape (B, C, T)
+        cond_inputs: Pre-processed conditioning inputs dict (merged positive + negative)
+        diffusion_objective: One of "v", "rectified_flow", "rf_denoiser"
+        steps: Number of sampling steps
+        cfg_scale: Classifier-free guidance scale
+        conditioning: List of conditioning dicts (for computing varlen from seconds_total)
+        sample_rate: Audio sample rate
+        pretransform: Optional pretransform for decoding latents and computing downsampling_ratio
+        mask_padding_attention: Whether to create padding_mask for attention
+        use_effective_length_for_schedule: Whether to use effective_seq_len for dist_shift
+        padding_mask: Optional pre-computed padding mask (B, T). If provided, skips
+            internal mask computation. Use this to ensure consistency with training masks.
+        headroom_seconds: Extra seconds beyond seconds_total for valid region
+        dist_shift: Distribution shift object for warping the timestep schedule, or None
+        sampler_type: Sampler type. For RF: "euler", "rk4", "dpmpp", "pingpong".
+            For v-diffusion: "v-ddim", "v-ddim-cfgpp", or k-diffusion types like "dpmpp-2m-sde".
+        batch_cfg: Whether to use batched CFG
+        rescale_cfg: Whether to use rescaled CFG
+        apg_scale: APG (Adaptive Projected Guidance) scale. 1.0 = full APG, 0.0 = vanilla CFG
+        init_data: Optional pre-encoded latent tensor for variation/img2img (shape: B, C, T)
+        init_noise_level: Noise level (sigma_max) when using init_data. 1.0 = full noise (no variation).
+        callback: Optional callback for progress reporting
+        disable_tqdm: Whether to disable progress bar
+        decode: Whether to decode latents using pretransform
+        **sampler_kwargs: Additional kwargs passed to sampler
+
+    Returns:
+        Generated samples (decoded audio if decode=True, else latents)
+    """
+    device = noise.device
+    batch_size = noise.shape[0]
+    latent_seq_len = noise.shape[-1]
+
+    # Compute downsampling ratio
+    downsampling_ratio = pretransform.downsampling_ratio if pretransform is not None else 1
+
+    # Default sampler_type per objective
+    if sampler_type is None:
+        sampler_type = "pingpong" if diffusion_objective == "rf_denoiser" else "euler"
+
+
+    # Compute effective_seq_len for dist_shift if enabled
+    effective_seq_len = None
+    if use_effective_length_for_schedule and conditioning is not None:
+        effective_seq_len = compute_effective_seq_len_from_conditioning(
+            conditioning, sample_rate, downsampling_ratio, device
+        )
+
+    # Create padding_mask for attention if enabled (skip if pre-computed mask provided)
+    if padding_mask is None and mask_padding_attention and conditioning is not None:
+        raw_effective_len = compute_effective_seq_len_from_conditioning(
+            conditioning, sample_rate, downsampling_ratio, device
+        )
+        if raw_effective_len is not None:
+            headroom_tokens = int(headroom_seconds * sample_rate / downsampling_ratio)
+            valid_lengths = (raw_effective_len + headroom_tokens).clamp(max=latent_seq_len).long()
+            padding_mask = create_padding_mask_from_lengths(valid_lengths, latent_seq_len)
+
+    # Determine sigma_max for schedule
+    sigma_max = init_noise_level if init_data is not None else 1.0
+
+    # Mix init_data with noise for variation/img2img
+    # For k-diffusion v-diffusion samplers, init_data is passed through to sample_k
+    # which handles mixing internally with its own sigma scaling
+    k_diff_sampler_types = {"k-heun", "k-lms", "k-dpmpp-2s-ancestral", "k-dpm-2",
+                            "k-dpm-fast", "k-dpm-adaptive", "dpmpp-2m-sde", "dpmpp-3m-sde", "dpmpp-2m"}
+
+    if init_data is not None:
+        noise = init_data * (1 - sigma_max) + noise * sigma_max
+
+    # Build common sampler kwargs (conditioning + model-level params only).
+    # disable_tqdm and callback are passed explicitly to samplers that use them,
+    # not included here, to avoid leaking into model forward() calls.
+    common_kwargs = {
+        **cond_inputs,
+        "cfg_scale": cfg_scale,
+        "batch_cfg": batch_cfg,
+        "rescale_cfg": rescale_cfg,
+        "padding_mask": padding_mask,
+        "apg_scale": apg_scale,
+        **sampler_kwargs
+    }
+
+
+    if diffusion_objective in ["rectified_flow", "rf_denoiser"]:
+        # Remove v-diffusion-specific kwargs that don't apply to RF
+        common_kwargs.pop("sigma_min", None)
+        common_kwargs.pop("sigma_max", None)
+        common_kwargs.pop("rho", None)
+
+        # Build schedule
+        sigmas = build_schedule(
+            steps=steps, sigma_max=sigma_max,
+            dist_shift=dist_shift, effective_seq_len=effective_seq_len,
+            fallback_seq_len=latent_seq_len, include_endpoint=True, device=device
+        )
+
+        # Route to sampler
+        if sampler_type == "euler":
+            sampled = sample_discrete_euler(model, noise, sigmas=sigmas, callback=callback, disable_tqdm=disable_tqdm, **common_kwargs)
+        elif sampler_type == "rk4":
+            sampled = sample_rk4(model, noise, sigmas=sigmas, callback=callback, disable_tqdm=disable_tqdm, **common_kwargs)
+        elif sampler_type == "dpmpp":
+            sampled = sample_flow_dpmpp(model, noise, sigmas=sigmas, callback=callback, disable_tqdm=disable_tqdm, **common_kwargs)
+        elif sampler_type == "pingpong":
+            sampled = sample_flow_pingpong(model, noise, sigmas=sigmas, callback=callback, disable_tqdm=disable_tqdm, **common_kwargs)
+        else:
+            raise ValueError(f"Unknown sampler_type for {diffusion_objective}: {sampler_type}")
+
+    else:
+        raise ValueError(f"Unknown diffusion_objective: {diffusion_objective}")
+
+    # Decode if requested
+    if decode and pretransform is not None:
+        sampled = sampled.to(next(pretransform.parameters()).dtype)
+        sampled = pretransform.decode(sampled, chunked=chunked_decode)
+
+        # Zero out audio beyond valid region (padding positions decode to garbage)
+        if padding_mask is not None:
+            audio_mask = padding_mask.unsqueeze(1).repeat_interleave(downsampling_ratio, dim=-1)
+            # Trim or pad to match sampled length
+            if audio_mask.shape[-1] > sampled.shape[-1]:
+                audio_mask = audio_mask[..., :sampled.shape[-1]]
+            elif audio_mask.shape[-1] < sampled.shape[-1]:
+                audio_mask = torch.nn.functional.pad(audio_mask, (0, sampled.shape[-1] - audio_mask.shape[-1]), value=False)
+            sampled = sampled * audio_mask.to(sampled.dtype)
+
+    return sampled

stable_audio_3/interface/aeiou.py

@@ -0,0 +1,303 @@
+# Modified from https://github.com/drscotthawley/aeiou/blob/main/aeiou/viz.py under Apache 2.0 License
+# License can be found in LICENSES/LICENSE_AEIOU.txt
+
+from matplotlib.backends.backend_agg import FigureCanvasAgg
+import matplotlib.cm as cm
+from matplotlib.colors import Normalize
+from matplotlib.figure import Figure
+import numpy as np
+from PIL import Image
+
+import torch
+
+import torchaudio.transforms as T
+from einops import rearrange
+
+import numpy as np
+
+def embeddings_table(tokens):
+    from wandb import Table
+    from pandas import DataFrame
+
+    "make a table of embeddings for use with wandb"
+    features, labels = [], []
+    embeddings = rearrange(tokens, 'b d n -> b n d') # each demo sample is n vectors in d-dim space
+    for i in range(embeddings.size()[0]):  # nested for's are slow but sure ;-) 
+        for j in range(embeddings.size()[1]):
+            features.append(embeddings[i,j].detach().cpu().numpy())
+            labels.append([f'demo{i}'])    # labels does the grouping / color for each point
+    features = np.array(features)
+    labels = np.concatenate(labels, axis=0)
+    cols = [f"dim_{i}" for i in range(features.shape[1])]
+    df   = DataFrame(features, columns=cols)
+    df['LABEL'] = labels
+    return Table(columns=df.columns.to_list(), data=df.values)
+
+def project_down(tokens,     # batched high-dimensional data with dims (b,d,n)
+            proj_dims=3,     # dimensions to project to
+            method='pca',    # projection method: 'pca'|'umap'
+            n_neighbors=10, # umap parameter for number of neighbors
+            min_dist=0.3,    # umap param for minimum distance
+            debug=False,     # print more info while running
+            **kwargs,        # other params to pass to umap, cf. https://umap-learn.readthedocs.io/en/latest/parameters.html 
+            ):
+    "this projects to lower dimenions, grabbing the first _`proj_dims`_ dimensions"
+    method = method.lower()
+    A = rearrange(tokens, 'b d n -> (b n) d') # put all the vectors into the same d-dim space
+    if A.shape[-1] > proj_dims: 
+        if method=='umap':
+            from umap import UMAP
+            proj_data = UMAP(n_components=proj_dims, n_neighbors=n_neighbors, min_dist=min_dist,
+                            metric='correlation', **kwargs).fit_transform(A.cpu().numpy())
+            proj_data = torch.from_numpy(proj_data).to(tokens.device)
+        else:  # pca
+            (U, S, V) = torch.pca_lowrank(A)
+            proj_data = torch.matmul(A, V[:, :proj_dims])  # this is the actual PCA projection step
+    else:
+        proj_data = A
+    if debug: print("proj_data.shape =",proj_data.shape)
+    return torch.reshape(proj_data, (tokens.size()[0], -1, proj_dims)) # put it in shape [batch, n, proj_dims]
+
+
+def proj_pca(tokens, proj_dims=3):
+    return project_down(do_proj, method='pca', proj_dims=proj_dims)
+
+def point_cloud(
+    tokens,                  # embeddings / latent vectors. shape = (b, d, n)
+    method='pca',            # projection method for 3d mapping: 'pca' | 'umap'
+    color_scheme='batch',    # 'batch': group by sample; integer n: n groups, sequentially,  otherwise color sequentially by time step
+    output_type='wandbobj',  # plotly | points | wandbobj.  NOTE: WandB can do 'plotly' directly!
+    mode='markers',    # plotly scatter mode.  'lines+markers' or 'markers'
+    size=3,            # size of the dots
+    line=dict(color='rgba(10,10,10,0.01)'),  # if mode='lines+markers', plotly line specifier. cf. https://plotly.github.io/plotly.py-docs/generated/plotly.graph_objects.scatter3d.html#plotly.graph_objects.scatter3d.Line
+    ds_preproj=1,         # EXPERIMENTAL: downsampling factor before projecting  (1=no downsampling). Could screw up colors
+    ds_preplot=1,         # EXPERIMENTAL: downsampling factor before plotting (1=no downsampling). Could screw up colors
+    debug=False,          # print more info
+    colormap=None,        # valid color map to use, None=defaults
+    darkmode=False,       # dark background, white fonts
+    layout_dict=None,      # extra plotly layout options such as camera orientation
+    rgb_float = False,   # if True, color_scheme is RGB float values
+    **kwargs,             # anything else to pass along
+    ):
+    "returns a 3D point cloud of the tokens" 
+    if ds_preproj != 1: 
+        tokens = tokens[torch.randperm(tokens.size()[0])]  # EXPERIMENTAL: to 'correct' for possible weird effects of downsampling
+        tokens = tokens[::ds_preproj]
+        if debug: print("tokens.shape =",tokens.shape)
+
+    data = project_down(tokens, method=method, debug=debug, **kwargs).cpu().numpy()
+    if debug: print("data.shape =",data.shape)
+    if data.shape[-1] < 3: # for data less than 3D, embed it in 3D 
+        data = np.pad(data, ((0,0),(0,0),(0, 3-data.shape[-1])), mode='constant', constant_values=0)
+
+    bytime = False
+    points = []         
+    if color_scheme=='batch': # all dots in same batch index same color, each batch-index unique (almost)
+        ncolors = data.shape[0]
+        cmap, norm = cm.tab20, Normalize(vmin=0, vmax=ncolors)
+    elif isinstance(color_scheme, int) or color_scheme.isnumeric():  # n groups, by batch-indices, sequentially
+        ncolors = int(color_scheme)
+        cmap, norm = cm.tab20, Normalize(vmin=0, vmax=ncolors)
+    else:                                                    # time steps match up
+        bytime, ncolors = True, data.shape[1]
+        cmap, norm = cm.viridis, Normalize(vmin=0, vmax=ncolors)
+
+    cmap = cmap if colormap is None else colormap # overwrite default cmap with user choice if given 
+
+    points = []  
+    for bi in range(data.shape[0]):  # batch index
+        if color_scheme=='batch': 
+            [r, g, b, _] = [int(255*x) for x in cmap(norm(bi+1))] 
+        elif isinstance(color_scheme, int) or color_scheme.isnumeric():
+            grouplen = data.shape[0]//(ncolors)
+            #if debug: print(f"bi, grouplen, bi//grouplen = ",bi, grouplen, bi//grouplen)
+            [r, g, b, _] = [int(255*x) for x in cmap(norm(bi//grouplen))] 
+            #if debug: print("r,g,b = ",r,g,b)
+
+        if rgb_float: [r, g, b] = [x/255 for x in [r, g, b]]
+
+        for n in range(data.shape[1]):    # across time
+            if bytime: [r, g, b, _] = [int(255*x) for x in cmap(norm(n))] 
+            points.append([data[bi,n,0], data[bi,n,1], data[bi,n,2], r, g, b])  # include dot colors with point coordinates
+
+    point_cloud = np.array(points)
+        
+    if output_type == 'points':
+        return point_cloud
+    elif output_type =='plotly':
+        import plotly.graph_objects as go
+
+        fig = go.Figure(data=[go.Scatter3d(
+            x=point_cloud[::ds_preplot,0], y=point_cloud[::ds_preplot,1], z=point_cloud[::ds_preplot,2], 
+            marker=dict(size=size,  color=point_cloud[:,3:6]),
+            mode=mode, 
+            # show batch index and time index in tooltips: 
+            text=[ f'bi: {i*ds_preplot}, ti: {j}' for i in range(data.shape[0]//ds_preplot) for j in range(data.shape[1]) ],
+            line=line,
+            )])
+        fig.update_layout(margin=dict(l=0, r=0, b=0, t=0)) # tight layout
+        if darkmode: 
+            fig.layout.template = 'plotly_dark'
+            if isinstance(darkmode, str):   # 'rgb(12,15,24)'gradio margins in dark mode
+                fig.update_layout( paper_bgcolor=darkmode) 
+        if layout_dict:
+                fig.update_layout( **layout_dict ) 
+            
+        if debug: print("point_cloud: fig made. returning")
+        return fig
+    else:
+        from wandb import Object3D
+        return Object3D(point_cloud)
+    
+def pca_point_cloud(
+    tokens,                  # embeddings / latent vectors. shape = (b, d, n)
+    color_scheme='batch',    # 'batch': group by sample, otherwise color sequentially
+    output_type='wandbobj',  # plotly | points | wandbobj.  NOTE: WandB can do 'plotly' directly!
+    mode='markers',    # plotly scatter mode.  'lines+markers' or 'markers'
+    size=3,            # size of the dots
+    line=dict(color='rgba(10,10,10,0.01)'),  # if mode='lines+markers', plotly line specifier. cf. https://plotly.github.io/plotly.py-docs/generated/plotly.graph_objects.scatter3d.html#plotly.graph_objects.scatter3d.Line
+    **kwargs,
+    ):
+    return point_cloud(tokens, method='pca', color_scheme=color_scheme, output_type=output_type,
+        mode=mode, size=size, line=line, **kwargs)
+
+def power_to_db(spec, *, amin = 1e-10):
+    magnitude = np.asarray(spec)
+
+    log_spec = 10.0 * np.log10(np.maximum(amin, magnitude))
+    log_spec -= 10.0 * np.log10(np.maximum(amin, 1))
+
+    log_spec = np.maximum(log_spec, log_spec.max() - 80)
+
+    return log_spec
+
+def mel_spectrogram(waveform, power=2.0, sample_rate=48000, db=False, n_fft=1024, n_mels=128, debug=False):
+    "calculates data array for mel spectrogram (in however many channels)"
+    win_length = None
+    hop_length = n_fft//2 # 512
+
+    mel_spectrogram_op = T.MelSpectrogram(
+        sample_rate=sample_rate, n_fft=n_fft, win_length=win_length, 
+        hop_length=hop_length, center=True, pad_mode="reflect", power=power, 
+        norm='slaney', onesided=True, n_mels=n_mels, mel_scale="htk")
+
+    melspec = mel_spectrogram_op(waveform.float())
+    if db: 
+        amp_to_db_op = T.AmplitudeToDB()
+        melspec = amp_to_db_op(melspec)
+    if debug:
+        print_stats(melspec, print=print) 
+        print(f"torch.max(melspec) = {torch.max(melspec)}")
+        print(f"melspec.shape = {melspec.shape}")
+    return melspec
+
+def spectrogram_image(
+        spec,
+        title=None,
+        ylabel='freq_bin',
+        aspect='auto',
+        xmax=None,
+        db_range=[35,120],
+        justimage=False,
+        figsize=(5, 4), # size of plot (if justimage==False)
+        context_mask=None, # 1D inpaint mask: 1 = context/kept (tinted), 0 = inpainted/generated
+        context_color=(0.8, 0.2, 0.2, 0.3), # RGBA color for context region overlay (red, visible against viridis)
+    ):
+    "Modified from PyTorch tutorial https://pytorch.org/tutorials/beginner/audio_feature_extractions_tutorial.html"
+    fig = Figure(figsize=figsize, dpi=100) if not justimage else Figure(figsize=(4.145, 4.145), dpi=100, tight_layout=True)
+    canvas = FigureCanvasAgg(fig)
+    axs = fig.add_subplot()
+    spec = spec.squeeze()
+    im = axs.imshow(power_to_db(spec), origin='lower', aspect=aspect, vmin=db_range[0], vmax=db_range[1])
+
+    # Overlay tint on context (kept) regions where mask == 1
+    if context_mask is not None:
+        context_mask = np.asarray(context_mask).flatten()
+        num_frames = spec.shape[-1]
+        # Resize mask to match spectrogram frames if needed
+        if len(context_mask) != num_frames:
+            indices = np.linspace(0, len(context_mask) - 1, num_frames).astype(int)
+            context_mask = context_mask[indices]
+        # Find contiguous context regions (mask > 0.5) and draw axvspan overlays
+        in_span = False
+        span_start = 0
+        for j in range(len(context_mask)):
+            if context_mask[j] > 0.5 and not in_span:
+                span_start = j
+                in_span = True
+            elif context_mask[j] <= 0.5 and in_span:
+                axs.axvspan(span_start - 0.5, j - 0.5, color=context_color)
+                in_span = False
+        if in_span:
+            axs.axvspan(span_start - 0.5, len(context_mask) - 0.5, color=context_color)
+
+    if xmax:
+        axs.set_xlim((0, xmax))
+    if justimage:
+        import matplotlib.pyplot as plt
+        axs.axis('off')
+        plt.tight_layout()
+    else:
+        axs.set_ylabel(ylabel)
+        axs.set_xlabel('frame')
+        axs.set_title(title or 'Spectrogram (dB)')
+        fig.colorbar(im, ax=axs)
+    canvas.draw()
+    rgba = np.asarray(canvas.buffer_rgba())
+    im = Image.fromarray(rgba)
+    if justimage: # remove tiny white border
+        b = 15 # border size
+        im = im.crop((b,b, im.size[0]-b, im.size[1]-b))
+        #print(f"im.size = {im.size}")
+    return im
+
+def audio_spectrogram_image(waveform, power=2.0, sample_rate=48000, print=print, db=False, db_range=[35,120], justimage=False, log=False, figsize=(5, 4), context_mask=None):
+    "Wrapper for calling above two routines at once, does Mel scale; Modified from PyTorch tutorial https://pytorch.org/tutorials/beginner/audio_feature_extractions_tutorial.html"
+    melspec = mel_spectrogram(waveform, power=power, db=db, sample_rate=sample_rate, debug=log)
+    melspec = melspec[0] # TODO: only left channel for now
+    return spectrogram_image(melspec, title="MelSpectrogram", ylabel='mel bins (log freq)', db_range=db_range, justimage=justimage, figsize=figsize, context_mask=context_mask)
+
+from matplotlib.ticker import AutoLocator 
+def tokens_spectrogram_image(
+        tokens,                # the embeddings themselves (in some diffusion codes these are called 'tokens')
+        aspect='auto',         # aspect ratio of plot
+        title='Embeddings',    # title to put on top
+        ylabel='index',        # label for y axis of plot
+        cmap='coolwarm',       # colormap to use. (default used to be 'viridis')
+        symmetric=True,        # make color scale symmetric about zero, i.e. +/- same extremes
+        figsize=(8, 4),       # matplotlib size of the figure
+        dpi=100,               # dpi of figure
+        mark_batches=False,     # separate batches with dividing lines
+        debug=False,           # print debugging info
+    ):
+    "for visualizing embeddings in a spectrogram-like way"
+    batch_size, dim, samples = tokens.shape
+    embeddings = rearrange(tokens, 'b d n -> (b n) d')  # expand batches in time
+    vmin, vmax = None, None
+    if symmetric:
+        vmax = torch.abs(embeddings).max()
+        vmin = -vmax
+        
+    fig = Figure(figsize=figsize, dpi=dpi)
+    canvas = FigureCanvasAgg(fig)
+    ax = fig.add_subplot()
+    if symmetric: 
+        subtitle = f'min={embeddings.min():0.4g}, max={embeddings.max():0.4g}'
+        ax.set_title(title+'\n')
+        ax.text(x=0.435, y=0.9, s=subtitle, fontsize=11, ha="center", transform=fig.transFigure)
+    else:
+        ax.set_title(title)
+    ax.set_ylabel(ylabel)
+    ax.set_xlabel('time frame (samples, in batches)')
+    if mark_batches:
+        intervals = np.arange(batch_size)*samples
+        if debug: print("intervals = ",intervals)
+        ax.vlines(intervals, -10, dim+10, color='black', linestyle='dashed', linewidth=1)
+
+    im = ax.imshow(embeddings.cpu().numpy().T, origin='lower', aspect=aspect, interpolation='none', cmap=cmap, vmin=vmin,vmax=vmax) #.T because numpy is x/y 'backwards'
+    fig.colorbar(im, ax=ax)
+    fig.tight_layout()
+    canvas.draw()
+    rgba = np.asarray(canvas.buffer_rgba())
+    return Image.fromarray(rgba)
+

stable_audio_3/interface/diffusion_cond.py

@@ -0,0 +1,644 @@
+import gc
+import numpy as np
+import gradio as gr
+import re
+import subprocess
+import torch
+import torchaudio
+import threading
+import os, time, math
+
+from einops import rearrange
+
+from stable_audio_3.interface.aeiou import audio_spectrogram_image
+from stable_audio_3.inference.distribution_shift import LogSNRShift, FluxDistributionShift, DistributionShift, IdentityDistributionShift
+from stable_audio_3.models.lora import has_lora
+from stable_audio_3.interface.reprompt import reprompt as _reprompt_fn, get_model as _reprompt_get_model, is_model_cached as _reprompt_is_model_cached
+
+stable_audio_3_model = None
+sample_size = 5324800
+sample_rate = 44100
+n_loras = 0
+_LENGTH_EXTRACT_RE = re.compile(r' Length: (\d+) seconds\.?\s*$')
+
+
+# when using a prompt in a filename
+def condense_prompt(prompt):
+    pattern = r'[\\/:*?"<>|]'
+    # Replace special characters with hyphens
+    prompt = re.sub(pattern, '-', prompt)
+    # set a character limit
+    prompt = prompt[:150]
+    # zero length prompts may lead to filenames (ie ".wav") which seem cause problems with gradio
+    if len(prompt)==0:
+        prompt = "_"
+    return prompt
+
+def generate_cond(
+        prompt,
+        negative_prompt=None,
+        seconds_total=30,
+        cfg_scale=6.0,
+        steps=250,
+        preview_every=None,
+        seed=-1,
+        sampler_type="dpmpp-3m-sde",
+        sigma_max=1000,
+        cfg_interval_min=0.0,
+        cfg_interval_max=1.0,
+        cfg_rescale=0.0,
+        cfg_norm_threshold=0.0,
+        apg_scale=1.0,
+        file_format="wav",
+        file_naming="verbose",
+        cut_to_seconds_total=False,
+        init_audio=None,
+        init_noise_level=1.0,
+        mask_maskstart=None,
+        mask_maskend=None,
+        inpaint_audio=None,
+        init_audio_type="Init audio",
+        inversion_steps=100,
+        inversion_gamma=0.3,
+        inversion_unconditional=False,
+        duration_padding_sec=6.0,
+        batch_size=1,
+        dist_shift=None,
+        *lora_args
+    ):
+
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+    gc.collect()
+
+    print(f"Prompt: {prompt}")
+
+    global preview_images
+    preview_images = []
+    if preview_every == 0:
+        preview_every = None
+
+    # Parse per-LoRA controls from trailing args
+    # Each LoRA has 5 controls: strength, interval_min, interval_max, layer_filter
+    lora_configs = None
+    if n_loras > 0 and len(lora_args) >= n_loras * 4:
+        lora_configs = []
+        for i in range(n_loras):
+            off = i * 4
+            strength = lora_args[off]
+            interval_min = lora_args[off + 1]
+            interval_max = lora_args[off + 2]
+            layer_filter = lora_args[off + 3]
+            stable_audio_3_model.set_lora_strength(strength, lora_index=i)
+            lora_configs.append({
+                "lora_index": i,
+                "interval": (interval_min, interval_max),
+                "layer_filter": layer_filter,
+            })
+
+    input_sample_size = sample_size
+
+    def progress_callback(callback_info):
+        global preview_images
+        denoised = callback_info["denoised"]
+        current_step = callback_info["i"]
+        t = callback_info["t"]
+        sigma = callback_info["sigma"]
+
+        # Extract scalar from tensor if needed (samplers pass tensors to avoid GPU sync)
+        if isinstance(t, torch.Tensor):
+            t = t[0].item() if t.dim() > 0 else t.item()
+        if isinstance(sigma, torch.Tensor):
+            sigma = sigma[0].item() if sigma.dim() > 0 else sigma.item()
+
+        log_snr = math.log(((1 - sigma) / sigma) + 1e-6)
+
+        if (current_step - 1) % preview_every == 0:
+            if stable_audio_3_model.model.pretransform is not None:
+                denoised = stable_audio_3_model.model.pretransform.decode(denoised)
+            denoised = rearrange(denoised, "b d n -> d (b n)")
+            denoised = denoised.clamp(-1, 1).mul(32767).to(torch.int16).cpu()
+            audio_spectrogram = audio_spectrogram_image(denoised, sample_rate=sample_rate)
+            preview_images.append((audio_spectrogram, f"Step {current_step} sigma={sigma:.3f} logSNR={log_snr:.3f}"))
+
+    if init_audio_type == "RF-Inversion":
+        inversion_params = {
+            "inversion_steps": inversion_steps,
+            "inversion_gamma": inversion_gamma,
+            "inversion_unconditional": inversion_unconditional,
+            "inversion_cfg_scale": 1.0,
+            "inversion_sigma_max": 1.0
+        }
+    else:
+        inversion_params = None
+
+    generate_args = {
+        "prompt": prompt,
+        "negative_prompt": negative_prompt,
+        "duration": seconds_total,
+        "steps": steps,
+        "cfg_scale": cfg_scale,
+        "cfg_interval": (cfg_interval_min, cfg_interval_max),
+        "lora_configs": lora_configs,
+        "batch_size": batch_size,
+        "sample_size": input_sample_size,
+        "seed": seed,
+        "sampler_type": sampler_type,
+        "sigma_max": sigma_max,
+        "init_audio": init_audio,
+        "init_noise_level": init_noise_level,
+        "callback": progress_callback if preview_every is not None else None,
+        "scale_phi": cfg_rescale,
+        "cfg_norm_threshold": cfg_norm_threshold,
+        "apg_scale": apg_scale,
+        "duration_padding_sec": duration_padding_sec,
+        "dist_shift": dist_shift,
+    }
+
+    # If inpainting, send mask args
+    # This will definitely change in the future
+    if inpaint_audio is not None:
+        generate_args.update({
+            "inpaint_audio": inpaint_audio,
+            "inpaint_mask_start_seconds": mask_maskstart,
+            "inpaint_mask_end_seconds": mask_maskend,
+        })
+
+    audio = stable_audio_3_model.generate(**generate_args)
+
+    # Filenaming convention
+    prompt_condensed = condense_prompt(prompt)
+    if file_naming=="verbose":
+        basename = prompt_condensed
+        if negative_prompt:
+            basename += ".neg-%s" % condense_prompt(negative_prompt)
+        basename += ".cfg%s" % (cfg_scale)
+        if sigma_max not in [1.0, 100.0]:
+            # this is a common parameter to tweak, if it's not a default value, put it in the verbose filename
+            basename += ".smx%s" % sigma_max
+        basename += ".%s" % seed
+    elif file_naming=="prompt":
+        basename = prompt_condensed
+    else:
+        # simple e.g. "output.wav"
+        basename = "output"
+
+    if file_format:
+        filename_extension = file_format.split(" ")[0].lower()
+    else:
+        filename_extension = "wav"
+    output_filename = "%s.%s" % (basename, filename_extension)
+    output_wav = "%s.wav" % basename
+
+    # Cut the extra silence off the end, if the user requested a smaller seconds_total
+    if cut_to_seconds_total:
+        audio = audio[:,:,:seconds_total*sample_rate]
+
+    # Encode the audio to WAV format
+    audio = rearrange(audio, "b d n -> d (b n)")
+    audio = audio.to(torch.float32).clamp(-1, 1).mul(32767).to(torch.int16).cpu()
+
+    # save as wav file
+    torchaudio.save(output_wav, audio, sample_rate)
+
+    # If file_format is other than wav, convert to other file format
+    cmd = ""
+    if file_format == "m4a aac_he_v2 32k":
+        # note: need to compile ffmpeg with --enable-libfdk_aac
+        cmd = f"ffmpeg -i \"{output_wav}\" -c:a libfdk_aac -profile:a aac_he_v2 -b:a 32k -y \"{output_filename}\""
+    elif file_format == "m4a aac_he_v2 64k":
+        cmd = f"ffmpeg -i \"{output_wav}\" -c:a libfdk_aac -profile:a aac_he_v2 -b:a 64k -y \"{output_filename}\""
+    elif file_format == "flac":
+        cmd = f"ffmpeg -i \"{output_wav}\" -y \"{output_filename}\""
+    elif file_format == "mp3 320k":
+        cmd = f"ffmpeg -i \"{output_wav}\" -b:a 320k -y \"{output_filename}\""
+    elif file_format == "mp3 128k":
+        cmd = f"ffmpeg -i \"{output_wav}\" -b:a 128k -y \"{output_filename}\""
+    elif file_format == "mp3 v0":
+        cmd = f"ffmpeg -i \"{output_wav}\" -q:a 0 -y \"{output_filename}\""
+    else: # wav
+        pass
+    if cmd:
+        cmd += " -loglevel error" # make output less verbose in the cmd window
+        subprocess.run(cmd, shell=True, check=True)
+
+    # Let's look at a nice spectrogram too
+    audio_spectrogram = audio_spectrogram_image(audio, sample_rate=sample_rate)
+
+    # Asynchronously delete the files after returning the output file, so as to prevent clutter in the directory
+    delete_files_async([output_wav, output_filename], 30)
+
+    return (output_filename, [audio_spectrogram, *preview_images])
+
+#  Asynchronously delete the given list of filenames after delay seconds. Sets up thread that sleeps for delay then deletes.
+def delete_files_async(filenames, delay):
+    def delete_files_after_delay(filenames, delay):
+        time.sleep(delay)  # Wait for the specified delay
+        for filename in filenames:
+            if os.path.exists(filename):
+                os.remove(filename)  # Delete the file
+    threading.Thread(target=delete_files_after_delay, args=(filenames, delay)).start()
+
+def create_sampling_ui(stable_audio_3_model, default_prompt=None):
+    global n_loras
+    diffusion_objective = stable_audio_3_model.model.diffusion_objective
+    is_rf = diffusion_objective == "rectified_flow"
+    is_rf_denoiser = diffusion_objective == "rf_denoiser" # includes ARC models
+
+    # Extract default dist_shift params from model's sampling_dist_shift
+    default_sampling_dist_shift = getattr(stable_audio_3_model.model, 'sampling_dist_shift', None)
+    default_dist_shift_type = "LogSNR"
+    default_logsnr_params = {"anchor_length": 2000, "anchor_logsnr": -6.2, "rate": 0.0, "logsnr_end": 2.0}
+    default_flux_params = {"min_length": 256, "max_length": 4096, "alpha_min": 6.93, "alpha_max": 6.93}
+    default_full_params = {"base_shift": 0.5, "max_shift": 1.15, "min_length": 256, "max_length": 4096}
+
+    if isinstance(default_sampling_dist_shift, LogSNRShift):
+        default_dist_shift_type = "LogSNR"
+        default_logsnr_params = {
+            "anchor_length": getattr(default_sampling_dist_shift, 'anchor_length', 2000),
+            "anchor_logsnr": getattr(default_sampling_dist_shift, 'anchor_logsnr', -6.2),
+            "rate": getattr(default_sampling_dist_shift, 'rate', 0.0),
+            "logsnr_end": getattr(default_sampling_dist_shift, 'logsnr_end', 2.0),
+        }
+    elif isinstance(default_sampling_dist_shift, FluxDistributionShift):
+        default_dist_shift_type = "Flux"
+        default_flux_params = {
+            "min_length": default_sampling_dist_shift.min_length,
+            "max_length": default_sampling_dist_shift.max_length,
+            "alpha_min": default_sampling_dist_shift.alpha_min,
+            "alpha_max": default_sampling_dist_shift.alpha_max,
+        }
+    elif isinstance(default_sampling_dist_shift, DistributionShift):
+        default_dist_shift_type = "Full"
+        default_full_params = {
+            "base_shift": default_sampling_dist_shift.base_shift,
+            "max_shift": default_sampling_dist_shift.max_shift,
+            "min_length": default_sampling_dist_shift.min_length,
+            "max_length": default_sampling_dist_shift.max_length,
+        }
+    elif default_sampling_dist_shift is None:
+        default_dist_shift_type = "None"
+
+    has_seconds_total = True
+
+    use_lora = has_lora(stable_audio_3_model.model)
+    lora_names = getattr(stable_audio_3_model.model, 'lora_names', [])
+    n_loras = len(lora_names)
+
+    if default_prompt is None:
+        default_prompt = ""
+
+    _reprompt_model_id = "Qwen/Qwen3.5-2B"
+    _reprompt_cached = _reprompt_is_model_cached(_reprompt_model_id)
+
+    with gr.Row():
+        with gr.Column(scale=6):
+            prompt = gr.Textbox(show_label=False, placeholder="Prompt", value=default_prompt)
+            negative_prompt = gr.Textbox(show_label=False, placeholder="Negative prompt")
+        prompt_assistant_button = gr.Button(
+            "Prompt Assistant" if _reprompt_cached else "Download Prompt Assistant (~4.2 GB)",
+            scale=1
+        )
+        generate_button = gr.Button("Generate", variant='primary', scale=1)
+
+    with gr.Row(equal_height=False):
+        with gr.Column():
+            with gr.Row(visible = True):
+                # Timing controls
+                seconds_total_slider = gr.Slider(minimum=0, maximum=sample_size//sample_rate, step=1, value=sample_size//sample_rate, label="Seconds total", visible=has_seconds_total)
+
+            with gr.Row():
+                # Steps slider
+                if is_rf:
+                    default_steps = 50
+                elif is_rf_denoiser:
+                    default_steps = 8
+
+                steps_slider = gr.Slider(minimum=1, maximum=500, step=1, value=default_steps, label="Steps")
+                # CFG scale
+                default_cfg_scale = 1.0 if is_rf_denoiser else 7.0
+                cfg_scale_slider = gr.Slider(minimum=0.0, maximum=25.0, step=0.1, value=default_cfg_scale, label="CFG scale")
+
+            # Per-LoRA controls (dynamic based on number of loaded LoRAs)
+            lora_ui_inputs = []
+            if use_lora and lora_names:
+                for i, lora_name in enumerate(lora_names):
+                    with gr.Accordion("LoRA {}: {}".format(i + 1, lora_name), open=(i == 0)):
+                        with gr.Row():
+                            strength = gr.Slider(minimum=0.0, maximum=10.0, step=0.1, value=1.0, label="strength")
+                        with gr.Row():
+                            int_min = gr.Slider(minimum=0.0, maximum=1.0, step=0.01, value=0.0, label="Interval min")
+                            int_max = gr.Slider(minimum=0.0, maximum=1.0, step=0.01, value=1.0, label="Interval max")
+                            lyr_filt = gr.Textbox(label="Layer filter", placeholder="")
+                        lora_ui_inputs.extend([strength, int_min, int_max, lyr_filt])
+
+            with gr.Accordion("Sampler params", open=False):
+                with gr.Row():
+                    # Seed
+                    seed_textbox = gr.Number(label="Seed (set to -1 for random seed)", value=-1, precision=0)
+
+                    cfg_interval_min_slider = gr.Slider(minimum=0.0, maximum=1, step=0.01, value=0.0, label="CFG interval min")
+                    cfg_interval_max_slider = gr.Slider(minimum=0.0, maximum=1, step=0.01, value=1.0, label="CFG interval max")
+
+                with gr.Row():
+                    cfg_rescale_slider = gr.Slider(minimum=0.0, maximum=1, step=0.01, value=0.0, label="CFG rescale amount")
+                    cfg_norm_threshold = gr.Slider(minimum=0.0, maximum=100, step=0.1, value=0.0, label="CFG norm threshold")
+                    apg_scale_slider = gr.Slider(minimum=0.0, maximum=1.0, step=0.1, value=1.0, label="APG scale", info="1.0=full APG, 0.0=vanilla CFG")
+
+                with gr.Row():
+                    # Sampler params
+                    if is_rf:
+                        sampler_types = ["euler", "rk4", "dpmpp"]
+                        default_sampler_type = "euler"
+                        sigma_max_max = 1.0
+                        sigma_max_default = 1.0
+                    elif is_rf_denoiser:
+                        sampler_types = ["pingpong"]
+                        default_sampler_type = "pingpong"
+                        sigma_max_max = 1.0
+                        sigma_max_default = 1.0
+                    else:
+                        sampler_types = ["dpmpp-2m-sde", "dpmpp-3m-sde", "dpmpp-2m", "k-heun", "k-lms", "k-dpmpp-2s-ancestral", "k-dpm-2", "k-dpm-adaptive", "k-dpm-fast", "v-ddim", "v-ddim-cfgpp"]
+                        default_sampler_type = "dpmpp-3m-sde"
+                        sigma_max_max = 1000.0
+                        sigma_max_default = 100.0
+
+                    sampler_type_dropdown = gr.Dropdown(sampler_types, label="Sampler type", value=default_sampler_type)
+                    sigma_max_slider = gr.Slider(minimum=0.0, maximum=sigma_max_max, step=0.1, value=sigma_max_default, label="Sigma max", visible=True)
+
+                with gr.Row():
+                    duration_padding_slider = gr.Slider(minimum=0.0, maximum=30.0, step=0.5, value=6.0, label="Duration padding (sec)")
+
+                def build_dist_shift(shift_type, p1, p2, p3, p4):
+                    """Build dist_shift from type + 4 params (meaning depends on type)."""
+                    if shift_type == "LogSNR":
+                        return LogSNRShift(anchor_length=int(p1), anchor_logsnr=p2, rate=p3, logsnr_end=p4)
+                    elif shift_type == "Flux":
+                        return FluxDistributionShift(min_length=int(p1), max_length=int(p2), alpha_min=p3, alpha_max=p4)
+                    elif shift_type == "Full":
+                        return DistributionShift(base_shift=p1, max_shift=p2, min_length=int(p3), max_length=int(p4))
+                    return IdentityDistributionShift()  # "None" = no shift
+
+                dist_shift_state = gr.State(value=default_sampling_dist_shift)
+
+                with gr.Row(visible=is_rf or is_rf_denoiser):
+                    dist_shift_type_dropdown = gr.Dropdown(
+                        ["LogSNR", "Flux", "Full", "None"],
+                        label="Sampling schedule shift",
+                        value=default_dist_shift_type,
+                        info="Distribution shift applied to sampling timesteps"
+                    )
+                with gr.Row(visible=(is_rf or is_rf_denoiser) and default_dist_shift_type == "LogSNR") as logsnr_params_row:
+                    logsnr_anchor_length_slider = gr.Slider(minimum=100, maximum=10000, step=100, value=default_logsnr_params["anchor_length"], label="Anchor length")
+                    logsnr_anchor_logsnr_slider = gr.Slider(minimum=-12.0, maximum=0.0, step=0.1, value=default_logsnr_params["anchor_logsnr"], label="Anchor log-SNR")
+                    logsnr_rate_slider = gr.Slider(minimum=-2.0, maximum=2.0, step=0.1, value=default_logsnr_params["rate"], label="Rate")
+                    logsnr_end_slider = gr.Slider(minimum=-2.0, maximum=6.0, step=0.1, value=default_logsnr_params["logsnr_end"], label="log-SNR end")
+                with gr.Row(visible=(is_rf or is_rf_denoiser) and default_dist_shift_type == "Flux") as flux_params_row:
+                    flux_min_length_slider = gr.Slider(minimum=1, maximum=10000, step=1, value=default_flux_params["min_length"], label="Min seq len")
+                    flux_max_length_slider = gr.Slider(minimum=1, maximum=10000, step=1, value=default_flux_params["max_length"], label="Max seq len")
+                    flux_alpha_min_slider = gr.Slider(minimum=0.1, maximum=20.0, step=0.1, value=default_flux_params["alpha_min"], label="Alpha min")
+                    flux_alpha_max_slider = gr.Slider(minimum=0.1, maximum=20.0, step=0.1, value=default_flux_params["alpha_max"], label="Alpha max")
+                with gr.Row(visible=(is_rf or is_rf_denoiser) and default_dist_shift_type == "Full") as full_params_row:
+                    full_base_shift_slider = gr.Slider(minimum=0.0, maximum=5.0, step=0.05, value=default_full_params["base_shift"], label="Base shift")
+                    full_max_shift_slider = gr.Slider(minimum=0.0, maximum=5.0, step=0.05, value=default_full_params["max_shift"], label="Max shift")
+                    full_min_length_slider = gr.Slider(minimum=1, maximum=10000, step=1, value=default_full_params["min_length"], label="Min length")
+                    full_max_length_slider = gr.Slider(minimum=1, maximum=10000, step=1, value=default_full_params["max_length"], label="Max length")
+
+                # Per-type slider groups for wiring to state
+                logsnr_sliders = [logsnr_anchor_length_slider, logsnr_anchor_logsnr_slider, logsnr_rate_slider, logsnr_end_slider]
+                flux_sliders = [flux_min_length_slider, flux_max_length_slider, flux_alpha_min_slider, flux_alpha_max_slider]
+                full_sliders = [full_base_shift_slider, full_max_shift_slider, full_min_length_slider, full_max_length_slider]
+                all_dist_shift_inputs = [dist_shift_type_dropdown] + logsnr_sliders + flux_sliders + full_sliders
+
+                def update_dist_shift_state(shift_type, *params):
+                    """Route the 4 relevant params to build_dist_shift based on type."""
+                    type_to_slice = {"LogSNR": params[0:4], "Flux": params[4:8], "Full": params[8:12]}
+                    p = type_to_slice.get(shift_type, (0, 0, 0, 0))
+                    return (
+                        build_dist_shift(shift_type, *p),
+                        gr.update(visible=((is_rf or is_rf_denoiser) and (shift_type == "LogSNR"))),
+                        gr.update(visible=((is_rf or is_rf_denoiser) and (shift_type == "Flux"))),
+                        gr.update(visible=((is_rf or is_rf_denoiser) and (shift_type == "Full"))),
+                    )
+
+                for component in all_dist_shift_inputs:
+                    component.change(
+                        update_dist_shift_state,
+                        inputs=all_dist_shift_inputs,
+                        outputs=[dist_shift_state, logsnr_params_row, flux_params_row, full_params_row],
+                    )
+
+            # Hidden state for batch_size (no UI control, but needed for function signature)
+            batch_size_state = gr.State(value=1)
+
+            with gr.Accordion("Output params", open=False):
+                # Output params
+                with gr.Row():
+                    file_format_dropdown = gr.Dropdown(["wav", "flac", "mp3 320k", "mp3 v0", "mp3 128k", "m4a aac_he_v2 64k", "m4a aac_he_v2 32k"], label="File format", value="wav")
+                    file_naming_dropdown = gr.Dropdown(["verbose", "prompt", "output.wav"], label="File naming", value="verbose") # ,"prompt","verbose"
+                    preview_every_slider = gr.Slider(minimum=0, maximum=100, step=1, value=0, label="Spec Preview Every")
+
+                    cut_to_seconds_total_checkbox = gr.Checkbox(label="Cut to seconds total", value=True)
+                    autoplay_checkbox = gr.Checkbox(label="Autoplay", value=False, elem_id="autoplay")
+                    infinite_radio_checkbox = gr.Checkbox(label="Infinite Radio", value=False, elem_id="infinite-radio")
+                    automatic_download_checkbox = gr.Checkbox(label="Auto Download", value=False, elem_id="automatic-download")
+
+            # Default generation tab
+            with gr.Accordion("Init audio", open=False):
+                init_audio_input = gr.Audio(label="Init audio", waveform_options=gr.WaveformOptions(show_recording_waveform=False))
+                min_noise_level = 0.01
+                max_noise_level = 1.0
+                default_noise_level = 0.9 # roughly halfway style transfer values
+                if is_rf:
+                    choices = ["Init audio","RF-Inversion"]
+                else:
+                    choices = ["Init audio"]
+
+                init_audio_type_radio = gr.Radio(label="Techniques", choices=choices, value=choices[0], visible=len(choices)>1)
+                with gr.Column(visible=True) as interface_a:
+                    init_noise_level_slider = gr.Slider(minimum=min_noise_level, maximum=max_noise_level, step=0.01, value=default_noise_level, label="Init noise level")
+                with gr.Column(visible=False) as interface_b:
+                    inversion_steps_slider = gr.Slider(minimum=1, maximum=500, step=1, value=100, label="Inversion Steps")
+                    inversion_gamma_slider = gr.Slider(minimum=0, maximum=1, step=0.1, value=0, label="Gamma", visible=True)
+                    inversion_unconditional_checkbox = gr.Checkbox(label="Unconditional", value=False)
+                    gr.HTML("<div style='opacity: 0.5; padding: 0px'>For reproduction, try empty prompt, cfg 1, gamma .3<br>\
+                        For prompt re-stylization, try cfg 1-7, gamma 0-.15, unconditional</div>")
+                def init_audio_type_switch(choice):
+                    return (
+                        gr.update(visible=(choice == "Init audio")),
+                        gr.update(visible=(choice == "RF-Inversion"))
+                    )
+                init_audio_type_radio.change(init_audio_type_switch, inputs=init_audio_type_radio, outputs=[interface_a, interface_b])
+
+            with gr.Accordion("Inpainting", open=False):
+                inpaint_audio_input = gr.Audio(label="Inpaint audio", waveform_options=gr.WaveformOptions(show_recording_waveform=False))
+                mask_maskstart_slider = gr.Slider(minimum=0.0, maximum=sample_size//sample_rate, step=0.1, value=0, label="Mask Start (sec)")
+                mask_maskend_slider = gr.Slider(minimum=0.0, maximum=sample_size//sample_rate, step=0.1, value=sample_size//sample_rate, label="Mask End (sec)")
+
+                # Update inpainting slider ranges when seconds_total changes.
+                # Only seconds_total is an input — reading the mask sliders here would cause
+                # validation errors since their values may exceed the about-to-be-reduced maximum.
+                def update_inpaint_sliders(seconds_total):
+                    max_val = max(seconds_total, 1)
+                    return (
+                        gr.update(maximum=max_val),
+                        gr.update(maximum=max_val, value=max_val),
+                    )
+                seconds_total_slider.change(update_inpaint_sliders, inputs=[seconds_total_slider], outputs=[mask_maskstart_slider, mask_maskend_slider])
+
+            inputs = [
+                prompt,
+                negative_prompt,
+                seconds_total_slider,
+                cfg_scale_slider,
+                steps_slider,
+                preview_every_slider,
+                seed_textbox,
+                sampler_type_dropdown,
+                sigma_max_slider,
+                cfg_interval_min_slider,
+                cfg_interval_max_slider,
+                cfg_rescale_slider,
+                cfg_norm_threshold,
+                apg_scale_slider,
+                file_format_dropdown,
+                file_naming_dropdown,
+                cut_to_seconds_total_checkbox,
+                init_audio_input,
+                init_noise_level_slider,
+                mask_maskstart_slider,
+                mask_maskend_slider,
+                inpaint_audio_input,
+                init_audio_type_radio,
+                inversion_steps_slider,
+                inversion_gamma_slider,
+                inversion_unconditional_checkbox,
+                duration_padding_slider,
+                batch_size_state,
+                dist_shift_state,
+            ] + lora_ui_inputs
+
+        with gr.Column():
+            audio_output = gr.Audio(label="Output audio", interactive=False,
+                    waveform_options=gr.WaveformOptions(show_recording_waveform=False))
+            audio_spectrogram_output = gr.Gallery(label="Output spectrogram", show_label=False)
+            send_to_init_button = gr.Button("Send to init audio", scale=1)
+            send_to_init_button.click(fn=lambda audio: audio, inputs=[audio_output], outputs=[init_audio_input])
+
+            send_to_inpaint_button = gr.Button("Send to inpaint audio", scale=1)
+            send_to_inpaint_button.click(fn=lambda audio: audio, inputs=[audio_output], outputs=[inpaint_audio_input])
+
+    generate_button.click(fn=generate_cond,
+        inputs=inputs,
+        outputs=[
+            audio_output,
+            audio_spectrogram_output
+        ],
+        api_name="generate")
+
+    def _prompt_assistant_or_download(text, progress=gr.Progress(track_tqdm=True)):
+        if not _reprompt_is_model_cached(_reprompt_model_id):
+            _reprompt_get_model(_reprompt_model_id)
+            return text, gr.update(), gr.update(value="Prompt Assistant")
+        _, result, _ = _reprompt_fn(text, "Auto", "", _reprompt_model_id, 128, 1.11)
+        m = _LENGTH_EXTRACT_RE.search(result)
+        if m:
+            max_seconds = sample_size // sample_rate
+            seconds = min(int(m.group(1)), max_seconds)
+            result = result[:m.start()]
+        else:
+            seconds = gr.update()
+        return result, seconds, gr.update()
+
+    prompt_assistant_button.click(
+        fn=_prompt_assistant_or_download,
+        inputs=[prompt],
+        outputs=[prompt, seconds_total_slider, prompt_assistant_button],
+        concurrency_limit=1,
+    )
+
+def create_diffusion_cond_ui(model, gradio_title="", default_prompt=None):
+    global sample_size, sample_rate, stable_audio_3_model
+
+    sample_size = model.model_config["sample_size"]
+    sample_rate = model.model_config["sample_rate"]
+    stable_audio_3_model = model
+
+
+    js ="""function run_javascript_on_page_load(){
+        const generateBtn = Array.from(document.querySelectorAll('button'))
+            .find(btn => btn.innerText.trim() === 'Generate');
+        function getAudioOutputPlayer () {
+            return [...document.querySelectorAll('label')].find(label => label.textContent.trim() === 'Output audio')?.parentElement.querySelector('audio');
+        }
+        const infiniteRadio = document.querySelector('#infinite-radio input[type="checkbox"]');
+        const autoplay = document.querySelector('#autoplay input[type="checkbox"]');
+        const automaticDownload = document.querySelector('#automatic-download input[type="checkbox"]');
+        let radioAutoStart = false;
+        let listenersSetup = false;
+        const setupListeners = () => {
+            const audioEl = getAudioOutputPlayer();
+            if (!audioEl) return;
+            audioEl.addEventListener('loadedmetadata', () => {
+                if(automaticDownload.checked){
+                    downloadAudio(audioEl);
+                }
+                if(autoplay.checked || radioAutoStart){
+                    audioEl.play();
+                    radioAutoStart = false;
+                }
+                if(infiniteRadio.checked){
+                    audioEl.addEventListener('timeupdate', function checkAudioEnd() {
+                        // Can set window.headstart (seconds) in the dev console if you want to start generating before the song is over
+                        let headstart = 1;
+                        if(window.headstart) headstart = window.headstart;
+                        if (audioEl.duration - audioEl.currentTime <= headstart) {
+                            generateBtn.click();
+                            radioAutoStart = true;
+                            audioEl.removeEventListener('timeupdate', checkAudioEnd);
+                        }
+                    });
+                }
+            });
+            listenersSetup = true;
+        };
+        generateBtn.addEventListener('click', () => {
+            if(listenersSetup) return;
+            const interval = setInterval(() => {
+                console.log("...")
+                const audioEl = document.querySelector('audio');
+                if (audioEl?.src && audioEl.src !== window.location.href) {
+                    setupListeners();
+                    clearInterval(interval);
+                }
+            }, 100);
+        });
+        // Respond to >> button on MacBookPro and on steering wheel during CarPlay
+        if ('mediaSession' in navigator) {
+            navigator.mediaSession.setActionHandler('nexttrack', () => generateBtn.click());
+            navigator.mediaSession.setActionHandler('play', () => getAudioOutputPlayer()?.play());
+            navigator.mediaSession.setActionHandler('pause', () => getAudioOutputPlayer()?.pause());
+        }
+        // Automatic Download
+        function downloadAudio(audioEl) {
+            const audioSrc = audioEl.src;
+            const link = document.createElement('a');
+            link.href = audioSrc;
+            link.download = audioSrc.substring(audioSrc.lastIndexOf('/') + 1);
+            document.body.appendChild(link);
+            link.click();
+            document.body.removeChild(link);
+        }
+    }
+    """
+
+    with gr.Blocks() as ui:
+        ui._sao_js = js
+        ui._sao_theme = gr.themes.Base()
+        if gradio_title:
+            gr.Markdown("### %s" % gradio_title)
+        with gr.Tab("Generation"):
+            create_sampling_ui(model, default_prompt=default_prompt)
+
+        # JavaScript to autoplay audio immediately after generation (if autoplay enabled)
+    return ui

stable_audio_3/interface/reprompt.py

@@ -0,0 +1,481 @@
+import re
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from huggingface_hub import scan_cache_dir
+
+# ---------------------------------------------------------------------------
+# System prompts
+# ---------------------------------------------------------------------------
+
+SYSTEM_PROMPTS = {
+    "Classifier": """You are a classification assistant.
+
+Your task is to classify the given input into ONE of the following categories:
+
+music_genre
+
+instrument
+
+sound
+
+one_shot
+
+Definitions:
+
+music_genre: A style or category of music.
+
+instrument: A musical instrument used to produce sound, or with words solo or stem.
+
+sound: A non-instrumental or general sound.
+
+one_shot: A single, isolated audio sample of an instrument or sound, usually short and intended for use in music production.
+
+Rules:
+
+Output ONLY one label: music_genre, instrument, sound, or one_shot.
+
+Do not explain your answer.
+
+If unsure, choose the closest category.
+
+EXAMPLES
+
+Music Genres:
+rock, jazz, pop, blues, reggae, techno, hip hop, classical, country, metal, folk, trance, house, dubstep, funk, soul, disco, punk, ambient, indie, grunge, gospel, bluegrass, R&B, salsa, samba, bossa nova, dance, electro, industrial, emo, ska, swing, swing jazz, hard rock, soft rock, progressive rock, psychedelic rock, acid jazz, latin, afrobeat, trap, chillhop, lo-fi, synthwave, vaporwave, EDM, drum and bass, hardcore, dub, experimental, alternative, post punk, new wave, indie pop, indie rock, glam rock, hair metal, death metal, black metal, metalcore, screamo, post rock, ambient techno, house deep, house progressive, minimal, chillout, lounge, trip hop, breakbeat, moombahton, cumbia, merengue, polka, zydeco, bhangra, k-pop, J-pop, Chinese pop, folk rock, psychedelic, surf rock, emo pop, synth pop, power pop, baroque pop, future bass, hardstyle, happy hardcore, gabber, jungle, UK garage, tropical house, moody trap, cinematic, soundtrack, orchestral
+
+Instruments:
+guitar, piano, violin, drums, bass, saxophone, flute, trumpet, cello, harp, trombone, clarinet, organ, synth, synthesizer, banjo, ukulele, sitar, marimba, accordion, tuba, oboe, bassoon, vibraphone, conga, bongo, djembe, cajon, xylophone, harmonica, dobro, mandolin, lute, koto, shamisen, erhu, bagpipes, shakuhachi, recorder, piccolo, euphonium, glockenspiel, steel drum, tabla, timpani, castanets, tambourine, triangle, dulcimer, zither, hurdy gurdy, melodica, pan flute, ocarina, didgeridoo, jaw harp, kora, santur, balalaika, sousaphone, lyre, clavinet, celesta, fife, cornet, contrabass, electric guitar, acoustic guitar, bass guitar, upright bass, electric piano, synth pad, modular synth, maracas, guiro, washboard, kalimba, udu, bodhran, tonbak, darbuka, tabla tarang, bongos, congas, cajon box, log drum, talking drum, steelpan, rainstick, dhol, ghatam, tambura, tambour, bamboo flute, reed organ, hurdy-gurdy, cimbalom, santur, psaltery, kantele, nyckelharpa, viola
+
+Sounds:
+rain, thunder, wind, footsteps, glass, explosion, ocean, beep, alarm, door, crowd, birds, fire, traffic, siren, clap, whistle, rustle, engine, bell, clock, heartbeat, dripping, typing, gunshot, gun reload, cannon, splash, pouring, chewing, sneezing, coughing, whisper, scream, giggle, laugh, growl, roar, howl, bark, meow, chirp, moo, neigh, oink, quack, frog, cricket, buzzing, hissing, airplane, helicopter, train, subway, car horn, tires, skateboard, skateboard wheels, zipper, button, rustling leaves, branches, river, waterfall, waves, snow crunch, wind chime, church bell, church organ, rainstick, wind gust, echo, footsteps gravel, footsteps wood, footsteps carpet, footsteps stairs, punch, slap, slapstick, slapstick sound, ricochet, splash water, dripping water, boiling, boiling kettle, frying, sizzling, microwave, fridge, fan, air conditioner, vacuum, typing keyboard, mouse click, door slam, gate, chain, rope, paper tearing, book closing, page flip, sewing machine, hammer, saw, drill, car engine, truck horn, airplane engine
+
+One-Shots:
+kick, snare, hi-hat, tom, crash, clap, 808, piano stab, guitar pluck, synth stab, triangle, noise burst, cymbal hit, bass hit, rimshot, sound fx, percussion hit, kick drum, snare drum, open hi-hat, closed hi-hat, ride cymbal, splash cymbal, reverse cymbal, tom hit, rim click, shaker, tambourine hit, hand clap, finger snap, cowbell, bell hit, bell toll, bell ring, church bell hit, gong hit, woodblock, claves, conga hit, bongo hit, djembe hit, tabla hit, rim tap, crash cymbal hit, metallic hit, metallic clang, glass hit, glass break, scratch vinyl, scratch record, vinyl scratch, tape stop, pitch down, pitch up, reverse hit, laser zap, laser shot, explosion one-shot, gunshot one-shot, shotgun, pistol, ricochet, swoosh, whoosh, wind gust one-shot, thunder clap, rain drop, water drip, water splash, water hit, ocean hit, fire hit, fire crackle, lava hit, lava splash, air hit, balloon pop, pop, click, snap, digital blip, 8-bit sound, synth hit, synth burst, lead stab, pad hit, chord stab, organ hit, piano chord hit, piano note hit, guitar chord hit, guitar note hit, bass stab, bass note hit, sub bass hit""",
+
+    "One-shot": """You are a music metadata expert. Given an instrument or sound, generate a descriptive prompt for a short, isolated one-shot audio sample for music production.
+
+1. Identify the instrument or sound source.
+2. Describe the playing technique or hit type (e.g., pluck, slam, tap, stab).
+3. Include details about material, timbre, or texture.
+4. Add spatial or production qualities (dry/wet, room, close-mic).
+5. Specify length: short integer in seconds (1–11 s).
+
+Examples:
+- Piano key hit with bright percussive attack and resonant wooden body. Length: 2 seconds
+- Kick drum punchy low-end hit with warm skin resonance. Length: 2 seconds
+- Snare drum rimshot accent with crisp snare wires. Length: 2 seconds
+- Acoustic guitar fingerstyle note with warm spruce tone. Length: 3 seconds
+- Bass pluck with jazzy tone and resonant wooden body. Length: 3 seconds
+- Electric guitar power chord with distortion. Length: 3 seconds
+- Metallic glitch percussion hit with sharp metallic texture. Length: 2 seconds
+- Tabla resonant tone hit with natural skin timbre. Length: 2 seconds
+- Djembe slap accent with dry wooden resonance. Length: 2 seconds
+- Synth stab with reverb tail. Length: 3 seconds
+- Violin expressive note with vibrato and rich wooden resonance. Length: 3 seconds
+- Cello legato note, cinematic, with warm resonant body. Length: 3 seconds
+- Trumpet bright accent with slightly brassy overtones. Length: 2 seconds
+- Melodic saxophone jazz riff with smooth reed timbre and a slight vibrato bend. Length: 3 seconds
+- Harp pluck with airy tone and resonant strings. Length: 2 seconds
+- Glockenspiel bell-like note with bright metallic clarity. Length: 2 seconds
+- Metallic clang sound design hit. Length: 2 seconds
+- Granular texture hit. Length: 3 seconds
+- Reversed piano hit. Length: 2 seconds
+- Synth riser effect. Length: 6 seconds
+- Percussion impact hit. Length: 2 seconds
+- Cinematic hit. Length: 2 seconds
+- Dry clap, a crisp, natural single hand clap recorded in a dead room with an extremely sharp transient and no room reflections. Length: 1 second
+- Studio hat, a classic, natural recording of 14-inch hi-hats played tightly closed, zero ring, very fast decay. Length: 1 second
+- Disco open hat, bright 14-inch open hi-hat with long, shimmering decay, perfect for disco or dance grooves. Length: 1 second
+- Pillow kick, acoustic kick drum muffled with a heavy blanket, producing a short, dry "thump" with almost zero resonance. Length: 1 second
+- Short 808, punchy 808 kick with sharp, distorted transient and fast-decaying sub-tail. Length: 1 second
+- Egg shaker, classic plastic egg shaker recorded with a small-diaphragm condenser mic, producing a light, consistent "tick" with very short sustain. Length: 1 second
+- African drums, dynamic African drums and percussion ensemble with natural acoustic textures. Length: 3 seconds
+- Latin drums, dynamic Latin drums and percussion ensemble featuring authentic rhythmic patterns. Length: 3 seconds
+- String quartet, euphoric string quartet with dynamic and emotional playing, full of expressive harmonies and movement. Length: 3 seconds
+- Piano, nostalgic, atmospheric piano piece with dynamic and emotional performance, intimate and resonant. Length: 3 seconds
+- Analogue drift pad, warm polyphonic pad with three detuned oscillators (saw + triangle), subtle pitch drift, and lush bucket-brigade chorus for wide, nostalgic stereo image. Length: 11 seconds
+- Phase distortion bass, Casio CZ-style phase-distorted sine wave warped into a jagged sawtooth for retro synth bass tone. Length: 11 seconds
+- Vibrato saxophone, bright lyrical alto sax with fast fluttery vibrato, reedy vintage tone, captured with ribbon mic for warm nostalgic sound. Length: 11 seconds
+- Lofi upright bass, upright bass recorded with ribbon mic in a wooden room, natural air with slightly boxy resonance, tape-saturated for dusty 1950s jazz feel. Length: 2 seconds""",
+
+    "Instrument": """You are a music metadata expert. Given an instrument, generate a descriptive prompt for a generative audio model.
+
+1. Identify the instrument.
+2. Add playing style or technique.
+3. Include details about material, timbre, or texture.
+4. Add musical style or mood. Specify the genre, context, or emotional character.
+5. Add spatial or production qualities.
+6. Specify BPM: Always include a BPM appropriate to the style and context.
+7. Specify length: Provide an integer in seconds (6–20 s for loops, 20–180 s for stems).
+
+Examples:
+- Synth arpeggio loop with bright detuned oscillators. BPM: 120. Length: 8 seconds
+- Chord stab loop with sharp percussive attack. BPM: 90. Length: 6 seconds
+- Guitar muted strum loop with tight rhythmic feel. BPM: 100. Length: 8 seconds
+- Pluck sequence loop with bright resonant tone. BPM: 128. Length: 10 seconds
+- Marimba and vibraphone percussive loop with resonant wooden and metallic tones. BPM: 110. Length: 12 seconds
+- Drum loop with deep muffled kick on beat one, snappy rimshot snare on beats two and four with rolling ghost note fills, and tight closed hi-hats with subtle open accents. BPM: 85. Length: 10 seconds
+- Drum groove loop with brushed snare swinging on the ride, soft feathered kick on downbeats, and light closed hi-hat taps on the upbeats. BPM: 130. Length: 12 seconds
+- Kick and hi-hat loop with four-on-the-floor punchy kick, tight closed hi-hats on every eighth note, and a sharp dry snare on beats two and four. BPM: 130. Length: 15 seconds
+- Vinyl crackle drum loop with warm low-pass filtered kick, dusty snare with tape saturation, and shuffled closed hi-hats with subtle vinyl crackle ambiance. BPM: 80. Length: 10 seconds
+- Ambient pad loop with evolving texture. BPM: 80. Length: 12 seconds
+- Melodic synth bass groove loop with pumping sidechain feel. BPM: 122. Length: 10 seconds
+- Melodic Bass slap and pop rhythm loop. BPM: 100. Length: 8 seconds
+- Acoustic bass walking line loop with natural wooden resonance. BPM: 120. Length: 12 seconds
+- String pizzicato motif loop, suspenseful, with tight string texture. BPM: 90. Length: 8 seconds
+- Brass staccato riff loop with sharp bright attack. BPM: 130. Length: 10 seconds
+- Flute airy melodic loop with wooden headjoint resonance. BPM: 100. Length: 6 seconds
+- Pan flute ambient loop with breathy timbre. BPM: 75. Length: 8 seconds
+- Clarinet riff loop with warm smooth reed tone. BPM: 120. Length: 10 seconds
+- Oboe motif loop, orchestral, with rich double reed resonance. BPM: 80. Length: 8 seconds
+- Recorder Renaissance motif loop with soft wooden timbre. BPM: 100. Length: 6 seconds
+- Electric sitar riff loop with buzzing resonant tone. BPM: 90. Length: 10 seconds
+- Koto plucked motif loop with resonant wooden strings. BPM: 90. Length: 8 seconds
+- Shamisen folk melody loop with percussive twang. BPM: 100. Length: 8 seconds
+- Banjo fingerpicking loop with metallic string resonance. BPM: 110. Length: 10 seconds
+- Mandolin tremolo loop with crisp wooden body tone. BPM: 120. Length: 10 seconds
+- Acoustic guitar chord vamp loop with natural room resonance. BPM: 110. Length: 12 seconds
+- Nylon string guitar arpeggio loop with warm, soft timbre. BPM: 90. Length: 15 seconds
+- Electric guitar riff loop with driven distorted tone. BPM: 130. Length: 10 seconds
+- Slide guitar melody loop with warm resonant glide. BPM: 100. Length: 12 seconds
+- Steel guitar slide loop with bright pedal steel tone. BPM: 95. Length: 12 seconds
+- Harpsichord arpeggio loop with crisp plucked attack. BPM: 120. Length: 10 seconds
+- Rhodes chord vamp loop with warm electric piano tone. BPM: 100. Length: 12 seconds
+- Clavinet funky rhythm loop. BPM: 105. Length: 10 seconds
+- Organ chord vamp loop with full drawbar warmth. BPM: 90. Length: 12 seconds
+- Drum loop with booming 808 kick on beat one, crisp snare on beat three, and rapid triplet hi-hat rolls with open hat accents for aggressive high-energy feel. BPM: 140. Length: 8 seconds
+- Breakbeat drum loop with chopped Amen-style snare flurries, driving kick on the one, fast sixteenth-note closed hi-hats, and syncopated open hat accents. BPM: 170. Length: 10 seconds
+- Glitch percussion loop with stuttered kick transients, randomised snare hits processed with bit-crushing, and erratic hi-hat patterns with pitch-shifted metallic ticks. BPM: 120. Length: 12 seconds
+- Metallic hits loop with distorted kick impacts, processed metal-plate snare slams, and grinding hi-hat noise bursts for aggressive mechanical texture. BPM: 120. Length: 10 seconds
+- Timpani hits loop, cinematic, with deep resonant kick-like timpani strikes on beat one, rolling snare-style timpani fills, and no hi-hats for a grand orchestral feel. BPM: 70. Length: 8 seconds
+- Snare roll loop, dramatic, with accelerating snare drum rolls building from soft to crashing, deep supporting kick pulses, and no hi-hats for maximum impact. BPM: 100. Length: 8 seconds
+- Accordion motif loop with bright reedy bellows tone. BPM: 100. Length: 10 seconds
+- Harmonica blues riff loop with expressive reed timbre. BPM: 90. Length: 10 seconds
+- Trombone riff loop with warm sliding brass tone. BPM: 120. Length: 10 seconds
+- French horn melodic loop, cinematic. BPM: 80. Length: 12 seconds
+- Soprano sax ballad loop. BPM: 70. Length: 12 seconds
+- Alto sax bebop riff loop. BPM: 200. Length: 10 seconds
+- Electric violin melodic loop with reverb. BPM: 90. Length: 10 seconds
+- String pad loop with cinematic texture. BPM: 70. Length: 15 seconds
+- Granular synth evolving texture loop. BPM: 90. Length: 15 seconds
+- Piano motif loop with soft felt hammer tone. BPM: 80. Length: 10 seconds
+- Pad and synth loop with lush detuned shimmer. BPM: 85. Length: 12 seconds
+- Synth lead loop with sidechain pumping compression. BPM: 128. Length: 10 seconds
+- Analog synth bassline loop with deep warm low-end. BPM: 122. Length: 12 seconds
+- FM synth lead motif loop with bright metallic shimmer. BPM: 110. Length: 10 seconds
+- Bass groove loop with tight rhythmic two-bar pattern. BPM: 100. Length: 16 seconds
+- Acoustic guitar fingerstyle motif loop with warm wood resonance. BPM: 90. Length: 45 seconds
+- Sombre acoustic guitar motif loop with cavernous reverb, delicate fingerpicking, and expressive melancholic tone. BPM: 70. Length: 45 seconds
+- Electric guitar rock riff motif loop. BPM: 130. Length: 40 seconds
+- Vintage electric guitar motif loop, live-recorded in a vintage studio, with expressive and dynamic solo performance. BPM: 90. Length: 40 seconds
+- Piano chord progression motif loop with rich harmonic movement. BPM: 120. Length: 60 seconds
+- String ensemble cinematic motif loop with rich wooden resonance. BPM: 80. Length: 120 seconds
+- Brass ensemble cinematic motif loop with bright metallic timbre. BPM: 90. Length: 90 seconds
+- Ethnic percussion ensemble motif loop with deep resonant djembe kick tones, slapped snare-like rim hits on congas, and layered shakers and bells providing hi-hat-like rhythmic texture with polyrhythmic patterns. BPM: 100. Length: 90 seconds
+- Synth ambient motif loop with evolving textures. BPM: 80. Length: 180 seconds
+- Motif loop with warm dusty vinyl crackle and tape saturation. BPM: 80. Length: 60 seconds
+- Synth lead and bass motif loop with bright punchy energy. BPM: 128. Length: 90 seconds
+- Funk band motif loop: bass, drums, guitar. BPM: 100. Length: 90 seconds
+- Ethnic flute motif for cinematic use. BPM: 80. Length: 30 seconds
+- Steel drum melodic motif loop with bright metallic resonance. BPM: 110. Length: 20 seconds
+- Marimba percussive motif loop with resonant wooden tone. BPM: 100. Length: 20 seconds
+- Vibraphone melodic motif loop with metallic shimmer. BPM: 90. Length: 25 seconds
+- Piano cinematic motif loop with resonant wooden tone. BPM: 80. Length: 30 seconds
+- Violin expressive cinematic motif loop with rich wooden resonance. BPM: 75. Length: 25 seconds
+- Cello expressive motif loop with deep wooden resonance. BPM: 70. Length: 30 seconds
+- Trumpet expressive motif loop with brassy overtones. BPM: 100. Length: 25 seconds
+- Sax expressive motif loop with warm reed timbre. BPM: 95. Length: 25 seconds
+- Ethnic drum ensemble motif loop with booming natural-skin bass drum kicks, sharp hand-slap snare accents on djembes and talking drums, and layered wooden and metal percussion providing rhythmic hi-hat-like patterns. BPM: 95. Length: 30 seconds
+- Ambient drone motif loop. BPM: 60. Length: 180 seconds
+- Orchestral tension motif loop. BPM: 90. Length: 150 seconds
+- Electronic track motif loop with drums, bass, synth. BPM: 128. Length: 180 seconds""",
+
+    "Music": """You are an expert musician and musicologist and prompt engineer. Transform the user's input into a detailed, vivid music prompt for a full instrumental track.
+
+1. Start with the genre or style and optional adjectives (e.g., upbeat, dreamy, aggressive).
+2. List the main instruments that define the track.
+3. Add supporting elements or layers such as pads, harmonics, effects, or field recordings.
+4. Include rhythm or percussion elements like drums, hi-hats, congas, brushes, or polyrhythms.
+5. Integrate mood and energy naturally in the sentence (e.g., "creating suspenseful tension" or "bright and uplifting").
+6. Specify the BPM.
+7. Specify the track length as an integer in seconds. Use ranges: energetic/dance 120-180s, pop/rock 180-210s, cinematic/ambient 240-300s.
+8. Combine all elements into one natural, fluid sentence. Avoid semicolons.
+
+Template:
+Genre/Style with main instruments, supporting instruments/layers, and rhythm/percussion creating mood/energy. BPM: X. Length: Y seconds
+
+Examples:
+- Jazz ballad with smooth saxophone lead, piano chords, upright bass, brushed drums, and soft strings that swing gently for a warm and cozy evening. BPM: 85. Length: 180 seconds
+- EDM festival track with pulsing synth leads, plucked arpeggios, layered pads, side-chained bass, punchy kick and snare, and hi-hat rolls creating bright, energetic, and uplifting dance energy. BPM: 128. Length: 150 seconds
+- Lo-fi hip-hop chill track with mellow electric piano, soft vinyl crackle, subtle synth pads, low-pass filtered drums, percussion loops, and soft plucked bass for a relaxed, dreamy vibe. BPM: 75. Length: 150 seconds
+- Heavy metal anthem with distorted electric guitars, bass guitar, double bass drums, and cymbal crashes with fast palm-muted riffs creating intense, aggressive energy. BPM: 160. Length: 180 seconds
+- Melancholic piano piece with soft piano lead, string pads, subtle atmospheric synths, and minimal brush percussion evoking a reflective rainy-day feeling. BPM: 60. Length: 240 seconds
+- Suspenseful electronic thriller with pulsing bass synth, arpeggiated lead synth, cinematic pads, glitchy percussion, and high string stabs creating dark and tense energy. BPM: 100. Length: 200 seconds
+- Dreamy ambient soundscape with layered pads, soft bell textures, gentle drones, and wind and water field recordings for ethereal and spacious meditation. BPM: 40. Length: 300 seconds
+- Fingerpicking acoustic guitar solo with harmonics, subtle reverb, occasional shaker and soft stomp percussion, and soft pad layers for warm intimate storytelling. BPM: 70. Length: 120 seconds
+- Synthwave 80s retro track with arpeggiated synth leads, analog pads, electric bass, punchy electronic drums, gated reverb snares, and atmospheric FX for nostalgic and vibrant energy. BPM: 110. Length: 180 seconds
+- Tribal percussion ensemble with congas, djembes, bongos, shakers, and frame drums layered with deep synthetic sub-bass in complex polyrhythms. BPM: 100. Length: 140 seconds
+- 1920s swing jazz with brass section, upright bass, piano, brushed drums, banjo, clarinet, and soft strings that swing lively for energetic dance vibes. BPM: 110. Length: 180 seconds
+- Futuristic electronic sci-fi track with pulsing bass synth, evolving lead synths, layered pads, glitch percussion, robotic FX, and sub-bass for tense cinematic energy. BPM: 125. Length: 200 seconds
+- Ambient underwater soundscape with flowing water textures, soft piano motifs, synth drones, distant bells, and underwater reverb for spacious meditative immersion. BPM: 45. Length: 300 seconds
+- Horror cinematic track with dissonant strings, eerie piano stabs, cinematic percussion including taiko and low toms, and synth FX producing suspenseful creepy tension. BPM: 90. Length: 240 seconds
+- Reggae track with offbeat guitar, warm basslines, snare, kick, congas, and horn stabs giving laid-back groovy energy. BPM: 85. Length: 150 seconds
+- Blues track with soulful electric guitar solos, walking bass, piano, and shuffle drums creating expressive and emotive storytelling. BPM: 90. Length: 180 seconds
+- Latin salsa with congas, timbales, horns, piano montunos, bass, and layered percussion for vibrant danceable energy. BPM: 120. Length: 210 seconds
+- Afrobeat track with electric guitar stabs, horns, layered percussion, congas, shakers, bass groove, and synth pads for vibrant rhythmic energy. BPM: 105. Length: 200 seconds
+- Indie rock track with electric guitar riffs, bass, live drum kit, layered synths, and subtle strings for energetic yet emotional feel. BPM: 110. Length: 180 seconds
+- Funk groove with slap bass, electric guitar chords, brass stabs, drums, congas, and rhythmic keyboards creating high-energy danceable rhythm. BPM: 105. Length: 180 seconds
+- Drum and bass track with fast breakbeat drums, deep sub-bass, sharp synth leads, pads, and atmospheric FX for high-energy club motion. BPM: 175. Length: 150 seconds
+- Dark ambient track with drones, distant bells, low rumbles, soft wind textures, and synth pads producing eerie immersive tension. BPM: 50. Length: 300 seconds
+- Tropical house track with marimba, steel drums, soft synths, smooth bass, layered percussion, and light piano riffs for sunny chill dance vibes. BPM: 110. Length: 180 seconds
+- Progressive rock track with electric guitar leads, organ, bass, drum kit, synth layers, and occasional strings for epic layered energy. BPM: 100. Length: 220 seconds
+- Music box melody with delicate metallic tones and soft resonance, lullaby style, with gentle ambient reverb. BPM: 60. Length: 20 seconds
+- Soft piano arpeggio with warm felted tone and slow attack, lullaby style, with intimate room ambience. BPM: 60. Length: 30 seconds
+- Harp gentle plucked pattern with airy resonance, lullaby style, with dreamy reverb tail. BPM: 65. Length: 25 seconds
+- Acoustic guitar fingerstyle pattern with warm nylon strings and soft dynamics, lullaby style, with subtle room resonance. BPM: 60. Length: 30 seconds
+- Ambient synth pad with smooth evolving texture and soft harmonics, lullaby style, with wide stereo ambience. BPM: 50. Length: 40 seconds
+- Early rock piano with walking left-hand bass line, shuffle rhythms, and blues scale improvisations in energetic 1950s boogie-woogie style. BPM: 160. Length: 180 seconds
+- Trip Hop track with jazzy sampled vibraphone, mid-tempo breakbeat drums, harp, Latin ethnic percussion, and sweeping cinematic strings creating airy, relaxing, soulful lounge vibes. BPM: 90. Length: 180 seconds
+- Country outlaw cinematic instrumental with blues pedal steel guitar, rustic mandolin, fiddle call-and-response, tape-driven rattly drum kit, autoharp, and soaring accordion solo for raw, emotional southern blues expression. BPM: 85. Length: 200 seconds
+- Neo Classical track with sweeping string section, elegant horns, and delicate piano creating soothing, hypnotic, modern, soft, and classic mood. BPM: 70. Length: 180 seconds
+- Art Rock desert track with desolate piano chords, western-themed rhythm guitars, unique lead guitars, rattly vintage drum kit, and supporting bass creating lonely, expansive, beautiful, and strange atmospheres. BPM: 95. Length: 180 seconds
+- Cinematic Sci-Fi score with dramatic horn section, building marcato strings, gliding bassoon, thunderous cymbals, subdued timpani, and subtle synth drones producing awe-inspiring, uplifting, epic intergalactic energy. BPM: 100. Length: 220 seconds
+- West Coast Hip Hop instrumental with cascading harp melodies, smooth Rhodes piano chops, vintage boom bap drums, and walking double bass producing raw, street, and soulful block-party vibes. BPM: 92. Length: 180 seconds
+- Synthwave futuristic track with pulsating synth bass, exciting chords, soaring leads, and reverberating drum machine patterns creating gritty, pounding, and cool energy. BPM: 110. Length: 180 seconds
+- Breakbeat track with complex percussion, intricate breakbeats, gritty synths, lush pads, and 808 bassline producing fresh, modern, futuristic, and rave-ready energy. BPM: 140. Length: 160 seconds
+- Lounge Jazz 1960s smooth track with laid-back drums, piano chords, double bass, soft electric piano, subtle flute, and unique percussion creating beautiful, atmospheric, eclectic, retro, and chill vibes. BPM: 85. Length: 180 seconds
+- Latin Jazz 1950s blissful track with laid-back Latin drums, euphoric piano chords, double bass, orchestral accompaniment, acoustic guitar, and vibraphone producing nostalgic, beautiful, atmospheric, cinematic, and chill mood. BPM: 95. Length: 180 seconds
+- Acid Jazz 1970s summertime track with smooth electric piano, trippy synth leads, laid-back vintage drum kit, fuzzy electric bass, and uplifting violin producing retro, psychedelic, jazzy, relaxing energy. BPM: 100. Length: 180 seconds
+- Progressive Soul 1970s track with feel-good piano, psychedelic organ, groovy vintage drum kit with percussion, fuzzy electric bass, and synth strings producing retro, raw, soulful, joyous atmosphere. BPM: 90. Length: 180 seconds
+- Discotheque 1970s French-inspired track with sultry piano, psychedelic guitars, groovy drum kit, fuzzy electric bass, and melancholic organ producing retro, raw, laid-back, and relaxing mood. BPM: 105. Length: 180 seconds
+- Soul Jazz 1970s track with expressive saxophone, smooth piano, groovy drum kit, rhythmic upright bass, sweeping strings, and minimal vibraphone producing retro, raw, laid-back, and epic energy. BPM: 95. Length: 180 seconds
+- Vintage R&B 1970s live studio track with subtle brass, smooth piano, sweeping strings, and minimal drums producing retro, beautiful, uplifting, nostalgic mood. BPM: 85. Length: 180 seconds
+- 50s Pop track with Latin influence, string section, bold brass, vibraphone, acoustic guitar, flute, ethnic percussion, and brushed drums creating sexy, epic, vintage, retro, melancholic, jazzy, dramatic energy. BPM: 100. Length: 180 seconds
+- A piece of calm, quiet, mellow, serene music perfect for a peaceful film score, featuring soft modulating piano, ambient sfx and foley, beautiful vibraphone, and subtle synthesizer drones. The mood is cinematic, thoughtful, serene and nostalgic. BPM: 55. Length: 300 seconds""",
+
+    "SFX": """You are a professional sound design expert. Convert the user's input into a precise, vivid sound effects description suitable for generative audio models.
+
+Describe clearly:
+- Sound source
+- Physical character (texture, timbre, material: metal, wood, glass, concrete, etc.)
+- Spatial qualities (indoor/outdoor, cave/open field/underwater, dry/reverberant, close-up/distant, echoing/muffled)
+- Temporal evolution (attack, decay, movement, transitions over time)
+- Include motion or spatial movement if applicable (passing, approaching, stereo movement)
+
+Audio length rules:
+- Very short sounds (impacts, clicks, gunshots): 1–3 seconds
+- Medium actions (footsteps, object movement, transitions): 3–6 seconds
+- Ambience / environments: 6–15 seconds
+- Always append: Length: X seconds (integer only, no decimals).
+
+Output constraints:
+- Length: 1–2 dense sentences maximum
+- Output ONLY the final rewritten prompt
+- No explanations, no formatting, no quotes
+- Use concise but dense technical language
+- Focus strictly on sound effects or ambience
+- Always append: Length: X seconds (integer only, no decimals).
+
+Quality guidelines:
+- Be specific and avoid vague terms
+- Prioritize clarity and realism
+- Combine elements into one coherent scene
+- Avoid redundancy
+
+Examples:
+- Heavy rain hitting a metal roof during a thunderstorm, distant thunder rumbles, stereo, realistic ambience. Length: 45 seconds
+- Quiet forest at dawn with birds chirping, soft wind through leaves, distant stream flowing. Length: 60 seconds
+- Busy city street at night, cars passing, muffled conversations, occasional horn, urban ambience. Length: 50 seconds
+- Ocean waves crashing against rocky cliffs, strong wind, dramatic and cinematic. Length: 70 seconds
+- Wooden door creaking open slowly in an old house, echoing interior, eerie tone. Length: 3 seconds
+- Glass bottle shattering on concrete, sharp impact, scattered fragments. Length: 2 seconds
+- Footsteps on gravel, steady walking pace, close perspective. Length: 8 seconds
+- Typing rapidly on a mechanical keyboard, crisp tactile clicks. Length: 5 seconds
+- Punch impact with deep bass hit, cinematic trailer style. Length: 2 seconds
+- Car speeding past at high velocity, doppler effect, realistic whoosh. Length: 3 seconds
+- Object falling from height and hitting ground with a heavy thud. Length: 2 seconds
+- Sword swing whooshing through air, fast motion, clean metallic tone. Length: 2 seconds
+- Futuristic laser blast, clean energy pulse, high-tech sound design. Length: 1 seconds
+- Spaceship engine humming, low frequency rumble, interior perspective. Length: 90 seconds
+- Magical spell casting, shimmering particles, rising tonal energy. Length: 8 seconds
+- Teleportation effect, glitchy digital distortion with a soft whoosh. Length: 5 seconds
+- Dark eerie drone with distant whispers, creepy, slow build tension. Length: 120 seconds
+- Sudden horror jump scare sting, sharp violin hit, cinematic. Length: 1 second
+- Metal scraping slowly in a dark tunnel, echoing and ominous. Length: 20 seconds
+- Explosion with debris scattering, deep bass, cinematic realism. Length: 4 seconds
+- Building collapsing, rumbling concrete, dust and debris falling. Length: 25 seconds
+- Fire crackling intensely, wood burning, close-up detail. Length: 80 seconds
+- Gunshot in a large empty warehouse, loud echo decay. Length: 2 seconds
+- Retro arcade coin insert sound, 8-bit style. Length: 1 second
+- Level up chime, bright, rewarding, fantasy RPG style. Length: 2 seconds
+- Error buzzer, short, digital, UI feedback. Length: 1 second
+- Menu navigation clicks, soft futuristic interface sounds. Length: 3 seconds
+- Layered soundscape: rain, thunder, footsteps, and distant sirens all blending naturally. Length: 90 seconds
+- Rapid sequence of three impacts: metal hit, glass break, wood crack, spaced evenly. Length: 4 seconds
+- Sound moving from left to right stereo field: passing motorcycle. Length: 5 seconds
+- Close vs far perspective transition: footsteps approaching then fading away. Length: 6 seconds
+- Tape stop sub drop, a massive sub-bass note that mimics a vinyl record or tape machine being turned off, the pitch and speed drop simultaneously, causing the high-end harmonics to smear and thicken as the sound grinds to a halt at a sub-sonic frequency. Length: 11 seconds
+- Gravel and leaves footsteps, the sound of a hard boot stepping onto dry leaves or gravel, crisp and natural with detailed texture. Length: 11 seconds
+- Ghostship moan, a massive, deep wooden groan with a low-frequency moan, like heavy timber under immense structural tension, swaying slowly, processed with long, dark wooden room reverb for a sense of scale. Length: 11 seconds
+- Bicycle chain, a continuous metallic whirring sound of a chain moving over sprockets, with individual teeth catching the links, processed with resonant band-pass filter to emphasize metallic singing. Length: 11 seconds
+- Warp drive, a sound that starts with a massive suck-back of ambient noise, followed by a supersonic crack and high-pitched zing that disappears into the distance, giving the sense of stretching space-time. Length: 11 seconds
+- Ice cubes, high-pitched musical clinking of hard ice hitting a thin glass, bright resonant ring with subtle liquid sloshing around the edges. Length: 11 seconds
+- Paper shuffle, the sound of a thick stack of heavy bond paper being squared up on a desk, dry papery thud with a quick fanning sound as air moves between the pages. Length: 11 seconds
+- Drawer slam, a blunt, powerful thud made by slamming a wooden desk drawer shut, pronounced low-mid body, slightly distorted for aggressive character. Length: 3 seconds""",
+}
+
+# ---------------------------------------------------------------------------
+# Track UI prefixes
+# ---------------------------------------------------------------------------
+
+TRACK_TYPE_PREFIXES = {
+    "music":      "TrackType: Music, VocalType: Instrumental, ",
+    "instrument": "TrackType: Instrument, ",
+    "sfx":        "TrackType: SFX, ",
+}
+
+# ---------------------------------------------------------------------------
+# Model cache
+# ---------------------------------------------------------------------------
+
+_cache = {"model_id": None, "tokenizer": None, "model": None}
+
+
+_WEIGHT_SUFFIXES = (".safetensors", ".bin")
+
+def is_model_cached(model_id):
+    if _cache["model_id"] == model_id:
+        return True
+    try:
+        for repo in scan_cache_dir().repos:
+            if repo.repo_id != model_id:
+                continue
+            for rev in repo.revisions:
+                if any(f.file_name.endswith(_WEIGHT_SUFFIXES) for f in rev.files):
+                    return True
+    except Exception:
+        pass
+    return False
+
+
+def get_model(model_id):
+    if model_id != _cache["model_id"]:
+        print(f"[Loading] {model_id}")
+        _cache["tokenizer"] = AutoTokenizer.from_pretrained(model_id)
+        _cache["model"] = AutoModelForCausalLM.from_pretrained(
+            model_id, torch_dtype=torch.bfloat16, device_map="auto"
+        )
+        _cache["model_id"] = model_id
+    return _cache["tokenizer"], _cache["model"]
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+def _postprocess(prompt, raw, tags_mode=True):
+    source = raw.split("</think>", 1)[1].strip() if "</think>" in raw else raw
+    if source.startswith("- "):
+        source = source[2:]
+
+    result = source.strip()
+
+    result = re.sub(r'((?:^|(?<=\.\s))\w)', lambda m: m.group().upper(), result)
+    if result:
+        result = result[0].upper() + result[1:]
+    return raw, result
+
+
+def _extract_examples(system_prompt: str) -> list:
+    return [m.group(1).strip() for m in re.finditer(r'^- (.+)', system_prompt, re.MULTILINE)]
+
+
+_ARTIFACT_RE = re.compile(r'[\[\]*#]')
+_VOCAL_RE = re.compile(
+    r'\b(vocals?|singing|singer|female|male|voice|voices|chorus|rap|rapper|chant(ing)?|lyrics?)\b',
+    re.IGNORECASE,
+)
+_LENGTH_RE = re.compile(r'\. Length: \d+ seconds\.?\s*$')
+
+
+def _has_artifacts(text: str) -> bool:
+    return (
+        bool(_ARTIFACT_RE.search(text))
+        or bool(_VOCAL_RE.search(text))
+        or not bool(_LENGTH_RE.search(text))
+        or len(text.split()) > 45
+    )
+
+
+def _run_model(tokenizer, model, system_prompt, prompt, max_new_tokens, temperature):
+    messages = [
+        {"role": "system", "content": system_prompt},
+        {"role": "user",   "content": f"Input: {prompt}\nOutput: "},
+    ]
+    text = tokenizer.apply_chat_template(
+        messages, tokenize=False, add_generation_prompt=True, enable_thinking=False,
+    )
+    inputs = tokenizer(text, return_tensors="pt").to(model.device)
+    with torch.no_grad():
+        output_ids = model.generate(
+            **inputs, max_new_tokens=int(max_new_tokens),
+            do_sample=True, temperature=float(temperature),
+        )
+    new_ids = output_ids[0][inputs["input_ids"].shape[-1]:]
+    return tokenizer.decode(new_ids, skip_special_tokens=True).strip()
+
+# ---------------------------------------------------------------------------
+# Main entry point
+# ---------------------------------------------------------------------------
+
+def reprompt(prompt, preset, system_prompt, model_id, max_new_tokens, temperature):
+    if not prompt.strip():
+        import random
+        pool = _extract_examples(SYSTEM_PROMPTS["Music"])
+        if not pool:
+            return "", "", None
+        example = random.choice(pool)
+        print(f"[Random] picked: {example}")
+        category = "music"
+        example = TRACK_TYPE_PREFIXES.get(category, "") + example
+        return example, example, category
+
+    tokenizer, model = get_model(model_id)
+
+    category = None
+    if preset == "Auto":
+        prompt_lower = prompt.lower()
+        if "tracktype: music" in prompt_lower and "vocaltype: instrumental" in prompt_lower:
+            label = "music_genre"
+        elif "tracktype: instrument" in prompt_lower:
+            label = "instrument"
+        elif "tracktype: sfx" in prompt_lower:
+            label = "sound"
+        else:
+            label = _run_model(tokenizer, model, SYSTEM_PROMPTS["Classifier"], prompt, 16, 0.1)
+            label = label.strip().lower()
+        if "music_genre" in label:
+            system_prompt = SYSTEM_PROMPTS["Music"]
+            category = "music"
+        elif "one_shot" in label:
+            system_prompt = SYSTEM_PROMPTS["One-shot"]
+            category = "one_shot"
+        elif "instrument" in label:
+            system_prompt = SYSTEM_PROMPTS["Instrument"]
+            category = "instrument"
+        else:
+            system_prompt = SYSTEM_PROMPTS["SFX"]
+            category = "sfx"
+        print(f"[Auto] classified as: {label}")
+
+    MAX_RETRIES = 5
+    for attempt in range(MAX_RETRIES):
+        raw = _run_model(tokenizer, model, system_prompt, prompt, max_new_tokens, temperature)
+        print(f"[Reprompt] attempt {attempt+1}  IN: {prompt}  OUT: {raw}")
+        if not _has_artifacts(raw):
+            break
+        if attempt < MAX_RETRIES - 1:
+            print("[Retry] artifacts detected, retrying…")
+    raw_out, result = _postprocess(prompt, raw)
+    if category:
+        result = TRACK_TYPE_PREFIXES.get(category, "") + result
+    return raw_out, result, category

stable_audio_3/loading_utils.py

@@ -0,0 +1,86 @@
+import json
+
+import torch
+from safetensors import safe_open
+from safetensors.torch import load_file
+from stable_audio_3.factory import (
+    create_autoencoder_from_config,
+    create_diffusion_cond_from_config,
+)
+
+
+def copy_state_dict(model, state_dict):
+    model_state_dict = model.state_dict()
+    state_dict = remap_state_dict_keys(state_dict, model_state_dict)
+    for key in state_dict:
+        if (
+            key in model_state_dict
+            and state_dict[key].shape == model_state_dict[key].shape
+        ):
+            model_state_dict[key] = state_dict[key]
+        else:
+            print(
+                f"Key {key} not found in target state_dict or shape mismatch. Skipping."
+            )
+
+    model.load_state_dict(model_state_dict, strict=False)
+
+
+def load_autoencoder(config_path: str, ckpt_path: str, device: str = "cpu"):
+    """Load only the autoencoder from a combined DiT+autoencoder checkpoint.
+
+    Only pretransform tensors are read from disk, directly onto the target device.
+    Standalone AE-only checkpoints (e.g. stabilityai/SAME-L / SAME-S) have no prefix.
+    """
+
+    with open(config_path) as f:
+        config = json.load(f)
+
+    autoencoder = create_autoencoder_from_config(config["model"], config["sample_rate"])
+
+    # Full DiT checkpoints nest the AE under pretransform.model.*;
+    # standalone AE-only checkpoints have no prefix.
+    nested_prefix = "pretransform.model."
+    with safe_open(ckpt_path, framework="pt", device=device) as f:
+        all_keys = list(f.keys())
+    if any(k.startswith(nested_prefix) for k in all_keys):
+        effective_prefix = nested_prefix
+    else:
+        effective_prefix = ""  # standalone AE — keys are already bare
+    with safe_open(ckpt_path, framework="pt", device=device) as f:
+        state_dict = {
+            k[len(effective_prefix) :]: f.get_tensor(k)
+            for k in all_keys
+            if k.startswith(effective_prefix)
+        }
+
+    copy_state_dict(autoencoder, state_dict)
+    return autoencoder.to(device)
+
+
+def load_diffusion_cond(
+    model_config,
+    ckpt_path: str,
+    device: str = "cuda",
+    model_half: bool = False,
+):
+    model = create_diffusion_cond_from_config(model_config)
+    copy_state_dict(model, load_file(ckpt_path))
+    model.to(device).eval().requires_grad_(False)
+    if model_half:
+        model.to(torch.float16)
+    return model
+
+
+def remap_state_dict_keys(state_dict, model_state_dict):
+    remapped = {}
+    for key, value in state_dict.items():
+        if key not in model_state_dict:
+            parts = key.split(".")
+            for i in range(1, len(parts)):
+                candidate = ".".join(parts[:i]) + "." + ".".join(parts[i + 1 :])
+                if candidate in model_state_dict:
+                    key = candidate
+                    break
+        remapped[key] = value
+    return remapped

stable_audio_3/model.py

@@ -0,0 +1,554 @@
+import json
+import numpy as np
+import torch
+import typing as tp
+from torch.nn.functional import interpolate
+
+from stable_audio_3.inference.audio_utils import prepare_audio, numpy_audio_to_tensor
+from stable_audio_3.inference.sampling import sample_diffusion
+from stable_audio_3.loading_utils import load_autoencoder, load_diffusion_cond
+from stable_audio_3.model_configs import ae_models, all_models
+from stable_audio_3.models.lora import (
+    set_lora_strength as _set_lora_strength,
+    load_and_apply_loras,
+)
+
+
+class StableAudioModel:
+    def __init__(self, model, model_config, device, model_half):
+        self.model = model
+        self.model_config = model_config
+        self.device = device
+        self.model_half = model_half
+        self.same = self.model.pretransform
+        self.dit = self.model.model
+        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
+
+    @staticmethod
+    def from_pretrained(model_name_or_path, device=None, model_half=True):
+        # Load the model and any necessary components here
+        if device is None and torch.cuda.is_available():
+            device = "cuda"
+        elif device is None and torch.backends.mps.is_available():
+            device = "mps"
+        elif device is None:
+            device = "cpu"
+
+        if not torch.cuda.is_available():
+            if model_name_or_path in ("medium", "medium-base"):
+                print(
+                    f"Warning: You are loading the {model_name_or_path} model without a GPU. This model is not designed to run on cpu"
+                )
+            model_half = False
+
+        if model_name_or_path not in all_models:
+            raise ValueError(
+                f"Unknown model '{model_name_or_path}'. Valid models: {list(all_models)}"
+            )
+
+        model_cfg = all_models[model_name_or_path]
+        local_config, local_ckpt = model_cfg.resolve()
+        with open(local_config) as f:
+            model_config = json.load(f)
+
+        model = load_diffusion_cond(
+            model_config, local_ckpt, device=device, model_half=model_half
+        )
+        model.use_lora = False
+        model.lora_names = []
+        return StableAudioModel(model, model_config, device, model_half)
+
+    def load_lora(self, lora_ckpt_paths):
+        """Load LoRA checkpoints onto the model after construction."""
+        model_type = self.model_config["model_type"]
+        svd_bases_path = self.model_config.get("svd_bases_path")
+        load_and_apply_loras(
+            self.model, lora_ckpt_paths, model_type, svd_bases_path=svd_bases_path
+        )
+
+    def set_lora_strength(self, strength: float, lora_index: int | None = None):
+        _set_lora_strength(self.model.model, strength, lora_index=lora_index)
+        _set_lora_strength(self.model.conditioner, strength, lora_index=lora_index)
+
+    @torch.inference_mode()
+    def generate(
+        self,
+        # Simple path: pass a prompt string and duration
+        prompt: str | list = None,
+        negative_prompt: str | list = None,
+        duration: float | list = 120,
+        # Generation parameters
+        steps: int = 8,
+        cfg_scale: float = 1.0,
+        batch_size: int = 1,
+        sample_size: int = 5292032,
+        truncate_output_to_duration: bool = True,
+        # Low-level path: pass pre-built conditioning dicts
+        conditioning: tp.Optional[tp.List[dict]] = None,
+        conditioning_tensors: tp.Optional[dict] = None,
+        negative_conditioning: tp.Optional[tp.List[dict]] = None,
+        negative_conditioning_tensors: tp.Optional[dict] = None,
+        seed: int = -1,
+        # Audio inputs
+        init_audio: tp.Optional[tp.Tuple[int, torch.Tensor]] = None,
+        init_noise_level: float = 1.0,
+        inpaint_audio: tp.Optional[tp.Tuple[int, torch.Tensor]] = None,
+        inpaint_mask=None,
+        inpaint_mask_start_seconds: tp.Optional[tp.Union[float, tp.List[float]]] = None,
+        inpaint_mask_end_seconds: tp.Optional[tp.Union[float, tp.List[float]]] = None,
+        # Schedule options
+        duration_padding_sec: float = 6.0,
+        apg_scale: float = 1.0,
+        dist_shift=None,
+        return_latents: bool = False,
+        chunked_decode: tp.Optional[bool] = None,
+        **sampler_kwargs,
+    ) -> torch.Tensor:
+        """
+        Generate audio.
+
+        Simple path:
+            model.generate(prompt="...", duration=30, steps=100)
+
+        Low-level path (pre-built conditioning):
+            model.generate(conditioning=[{"prompt": "...", "seconds_total": 30}], steps=100, ...)
+
+        Args:
+            prompt: The text prompt to condition on. Ignored if conditioning dicts are provided directly.
+            negative_prompt: The negative text prompt for classifier-free guidance. Ignored if negative_conditioning dicts are provided directly.
+            duration: The duration of the generated audio in seconds. Only used if conditioning dicts with "seconds_total" are not provided.
+            steps: The number of diffusion steps to use.
+            cfg_scale: Classifier-free guidance scale
+            batch_size: The batch size to use for generation.
+            sample_size: The length of the audio to generate, in samples.
+            truncate_output_to_duration: If True, truncate the output audio to the specified duration.
+            conditioning: A dictionary of conditioning parameters to use for generation.
+            conditioning_tensors: A dictionary of precomputed conditioning tensors to use for generation.
+            negative_conditioning: A dictionary of negative conditioning parameters for classifier-free guidance.
+            negative_conditioning_tensors: A dictionary of precomputed negative conditioning tensors for classifier-free guidance
+            seed: The random seed to use for generation, or -1 to use a random seed.
+            init_audio: A tuple of (sample_rate, audio) to use as the initial audio for generation.
+            init_noise_level: The noise level to use when generating from an initial audio sample.
+            inpaint_audio: A tuple of (sample_rate, audio) to use as the source audio for inpainting. The inpaint region will be determined by the inpaint_mask or inpaint_mask_start_seconds/inpaint_mask_end_seconds parameters.
+            inpaint_mask: A prebuilt mask tensor for inpainting. Shape should be [batch_size, sample_size].
+                Ignored if inpaint_mask_start_seconds/inpaint_mask_end_seconds are provided.
+            inpaint_mask_start_seconds: Start of the inpaint region in seconds. Can be a float
+                for a single region, or a list of floats for multiple non-contiguous regions.
+            inpaint_mask_end_seconds: End of the inpaint region in seconds. Can be a float
+                for a single region, or a list of floats matching inpaint_mask_start_seconds.
+            duration_padding_sec: Extra seconds to add when adapting duration (default 6.0).
+            apg_scale: APG (Adaptive Projected Guidance) scale. 1.0 = full APG, 0.0 = vanilla CFG.
+            dist_shift: Optional distribution shift override for sampling. If None, uses model.sampling_dist_shift.
+            return_latents: Whether to return the latents used for generation instead of the decoded audio.
+            chunked_decode: Whether to decode latents in overlapping chunks to reduce peak VRAM. True forces
+                chunked decoding on, False forces it off, None (default) uses the value set in the model config.
+            **sampler_kwargs: Additional keyword arguments to pass to the sampler.
+        """
+
+        device = str(self.device)
+
+        # Build conditioning from prompt string if not provided directly
+        if conditioning is None and conditioning_tensors is None:
+            assert prompt is not None, "Must provide either prompt or conditioning"
+            conditioning, negative_conditioning = self._build_conditioning_dicts(
+                prompt, negative_prompt, duration, batch_size
+            )
+
+        # Adapt sample size based on seconds_total in conditioning
+        audio_sample_size = sample_size
+        if conditioning is not None:
+            audio_sample_size = self._adapt_sample_size(
+                conditioning,
+                sample_size,
+                duration_padding_sec,
+            )
+
+        # Convert audio sample size to latent size
+        latent_sample_size = audio_sample_size
+        if self.model.pretransform is not None:
+            latent_sample_size = (
+                audio_sample_size // self.model.pretransform.downsampling_ratio
+            )
+
+        # Build inpaint mask from seconds if provided
+        if (
+            inpaint_mask_start_seconds is not None
+            and inpaint_mask_end_seconds is not None
+        ):
+            start_is_list = isinstance(inpaint_mask_start_seconds, list)
+            end_is_list = isinstance(inpaint_mask_end_seconds, list)
+            if start_is_list != end_is_list:
+                raise ValueError(
+                    "inpaint_mask_start_seconds and inpaint_mask_end_seconds must both be "
+                    "scalars or both be lists, got "
+                    f"{type(inpaint_mask_start_seconds).__name__} and "
+                    f"{type(inpaint_mask_end_seconds).__name__}."
+                )
+            starts = (
+                inpaint_mask_start_seconds
+                if start_is_list
+                else [inpaint_mask_start_seconds]
+            )
+            ends = (
+                inpaint_mask_end_seconds if end_is_list else [inpaint_mask_end_seconds]
+            )
+            if len(starts) != len(ends):
+                raise ValueError(
+                    f"inpaint_mask_start_seconds and inpaint_mask_end_seconds must have the same "
+                    f"length, got {len(starts)} and {len(ends)}."
+                )
+            inpaint_mask = torch.ones(1, audio_sample_size, device=device)
+            for start_sec, end_sec in zip(starts, ends):
+                mask_start_samples = min(
+                    int(start_sec * self.model.sample_rate),
+                    audio_sample_size,
+                )
+                mask_end_samples = min(
+                    int(end_sec * self.model.sample_rate),
+                    audio_sample_size,
+                )
+                inpaint_mask[:, mask_start_samples:mask_end_samples] = 0
+
+        # If the caller passed a prebuilt mask sized to the un-adapted sample_size (or
+        # anything longer than audio_sample_size), truncate to audio_sample_size so the
+        # downstream nearest-neighbor interpolation preserves the mask's time-domain
+        # positions instead of squashing the mask region.
+        if inpaint_mask is not None and inpaint_mask.shape[-1] > audio_sample_size:
+            inpaint_mask = inpaint_mask[:, :audio_sample_size]
+
+        # Match training: when mask_padding_attention is used, random_inpaint_mask
+        # zeroes the mask past real_sequence_length. Apply the
+        # same convention here so the mask matches the training distribution, whether
+        # it was built from seconds above or passed in by the caller.
+        if inpaint_mask is not None and conditioning is not None:
+            max_seconds = max(
+                (c.get("seconds_total", 0.0) for c in conditioning), default=0.0
+            )
+            if max_seconds > 0:
+                effective_audio_len = int(max_seconds * self.model.sample_rate)
+                mask_len = inpaint_mask.shape[-1]
+                if effective_audio_len < mask_len:
+                    inpaint_mask = inpaint_mask.clone()
+                    inpaint_mask[:, effective_audio_len:] = 0
+
+        if inpaint_mask is not None:
+            inpaint_mask = inpaint_mask.float()
+
+        # Seed and noise
+        seed = seed if seed != -1 else np.random.randint(0, 99999)
+        torch.manual_seed(seed)
+        noise = torch.randn(
+            [batch_size, self.model.io_channels, latent_sample_size], device=device
+        )
+
+        # Encode conditioning
+        if conditioning_tensors is None:
+            conditioning_tensors = self.model.conditioner(conditioning, device)
+        if (
+            negative_conditioning is not None
+            or negative_conditioning_tensors is not None
+        ):
+            if negative_conditioning_tensors is None:
+                negative_conditioning_tensors = self.model.conditioner(
+                    negative_conditioning, device
+                )
+        else:
+            negative_conditioning_tensors = {}
+
+        # Process init audio
+        if init_audio is not None:
+            init_audio, inpaint_mask = self._encode_audio_input(
+                init_audio, audio_sample_size, inpaint_mask
+            )
+            init_audio = init_audio.repeat(batch_size, 1, 1)
+
+        # Process inpaint audio
+        if inpaint_audio is not None:
+            inpaint_audio, inpaint_mask = self._encode_audio_input(
+                inpaint_audio, audio_sample_size, inpaint_mask
+            )
+            inpaint_audio = inpaint_audio.repeat(batch_size, 1, 1)
+        else:
+            if inpaint_mask is not None:
+                inpaint_mask = interpolate(
+                    inpaint_mask.unsqueeze(1), size=latent_sample_size, mode="nearest"
+                ).squeeze(1)
+
+        # Build inpaint mask tensor and masked input
+        if inpaint_mask is None:
+            mask = torch.zeros((batch_size, 1, latent_sample_size), device=device)
+        else:
+            mask = inpaint_mask.unsqueeze(1)
+        mask = mask.to(device)
+
+        inpaint_input = (
+            inpaint_audio * mask.expand_as(inpaint_audio)
+            if inpaint_audio is not None
+            else torch.zeros(
+                (batch_size, self.model.io_channels, latent_sample_size), device=device
+            )
+        )
+
+        conditioning_tensors["inpaint_mask"] = [mask]
+        conditioning_tensors["inpaint_masked_input"] = [inpaint_input]
+        conditioning_inputs = self.model.get_conditioning_inputs(conditioning_tensors)
+
+        if negative_conditioning_tensors:
+            negative_conditioning_tensors["inpaint_mask"] = [mask]
+            negative_conditioning_tensors["inpaint_masked_input"] = [inpaint_input]
+            negative_conditioning_tensors = self.model.get_conditioning_inputs(
+                negative_conditioning_tensors, negative=True
+            )
+
+        model_dtype = next(self.model.model.parameters()).dtype
+        noise = noise.type(model_dtype)
+        conditioning_inputs = {
+            k: v.type(model_dtype) if v is not None else v
+            for k, v in conditioning_inputs.items()
+        }
+
+        cond_inputs = {**conditioning_inputs, **negative_conditioning_tensors}
+
+        sampler_type = sampler_kwargs.pop("sampler_type", None)
+
+        result = sample_diffusion(
+            model=self.model.model,
+            noise=noise,
+            cond_inputs=cond_inputs,
+            diffusion_objective=self.model.diffusion_objective,
+            steps=steps,
+            cfg_scale=cfg_scale,
+            conditioning=conditioning,
+            sample_rate=self.model.sample_rate,
+            pretransform=self.model.pretransform,
+            mask_padding_attention=True,
+            use_effective_length_for_schedule=True,
+            headroom_seconds=duration_padding_sec,
+            dist_shift=dist_shift
+            if dist_shift is not None
+            else self.model.sampling_dist_shift,
+            sampler_type=sampler_type,
+            batch_cfg=True,
+            rescale_cfg=True,
+            apg_scale=apg_scale,
+            init_data=init_audio,
+            init_noise_level=init_noise_level,
+            decode=not return_latents,
+            chunked_decode=chunked_decode,
+            **sampler_kwargs,
+        )
+
+        if not return_latents:
+            result = result.to(torch.float32).clamp(-1, 1)
+
+        if not return_latents and truncate_output_to_duration:
+            if isinstance(duration, (int, float)):
+                max_length_samples = int(duration * self.model.sample_rate)
+                result = result[:, :, :max_length_samples]
+            else:
+                if torch.all(torch.tensor(duration) == duration[0]):
+                    max_length_samples = int(duration[0] * self.model.sample_rate)
+                    result = result[:, :, :max_length_samples]
+                else:
+                    # Warn that we can't truncate to a single duration if the durations are different, and return the full length output
+                    print(
+                        "Warning: Cannot truncate output to a single duration when passing a list of different durations"
+                    )
+
+        return result
+
+    # --- generate() helpers ---
+
+    @staticmethod
+    def _build_conditioning_dicts(prompt, negative_prompt, duration, batch_size):
+        """Returns (conditioning, negative_conditioning) lists of dicts."""
+
+        def _to_list(value, name):
+            """Broadcast a scalar or validate a sequence to length batch_size."""
+            if isinstance(value, (list, tuple)):
+                assert len(value) == batch_size, (
+                    f"Length of {name} ({len(value)}) must match batch_size ({batch_size})"
+                )
+                return list(value)
+            return [value] * batch_size
+
+        prompts = _to_list(prompt, "prompt")
+        durations = _to_list(duration, "duration")
+        conditioning = [
+            {"prompt": p, "seconds_total": d} for p, d in zip(prompts, durations)
+        ]
+
+        negative_conditioning = None
+        if negative_prompt is not None:
+            neg_prompts = _to_list(negative_prompt, "negative_prompt")
+            negative_conditioning = [
+                {"prompt": p, "seconds_total": d}
+                for p, d in zip(neg_prompts, durations)
+            ]
+
+        return conditioning, negative_conditioning
+
+    def _adapt_sample_size(self, conditioning, sample_size, duration_padding_sec):
+        """Returns audio_sample_size adapted from conditioning, clamped to sample_size."""
+        max_seconds = 0.0
+        for cond_dict in conditioning:
+            if "seconds_total" in cond_dict:
+                max_seconds = max(max_seconds, cond_dict["seconds_total"])
+
+        if max_seconds <= 0:
+            return sample_size
+
+        target_audio_samples = int(
+            (max_seconds + duration_padding_sec) * self.model.sample_rate
+        )
+        if self.model.pretransform is not None:
+            ds_ratio = self.model.pretransform.downsampling_ratio
+            # Round up to nearest multiple of downsampling ratio
+            target_audio_samples = (
+                (target_audio_samples + ds_ratio - 1) // ds_ratio
+            ) * ds_ratio
+            encoder_config = self.model_config["model"]["pretransform"]["config"][
+                "encoder"
+            ]["config"]
+            chunk_size = encoder_config.get("chunk_size", 32)
+            stride = encoder_config["strides"][0]  # or min(strides) if multiple
+            # For chunked attention with latent space, align to chunk size after downsampling
+            latent_align = chunk_size // stride
+            align = ds_ratio * latent_align
+            target_audio_samples = ((target_audio_samples + align - 1) // align) * align
+
+        return min(target_audio_samples, sample_size)
+
+    def _encode_audio_input(self, audio_input, audio_sample_size, inpaint_mask=None):
+        """
+        Converts a (sample_rate, audio) tuple to an encoded latent tensor.
+        If model has a pretransform, encodes to latent space and downsamples inpaint_mask to match.
+        Returns (encoded_audio, updated_inpaint_mask). encoded_audio is not yet repeated to batch size.
+        """
+        device = str(self.device)
+        in_sr, audio_data = audio_input
+        if isinstance(audio_data, np.ndarray):
+            audio_data = numpy_audio_to_tensor(audio_data)
+        io_channels = (
+            self.model.pretransform.io_channels
+            if self.model.pretransform is not None
+            else self.model.io_channels
+        )
+        audio = prepare_audio(
+            audio_data,
+            in_sr=in_sr,
+            target_sr=self.model.sample_rate,
+            target_length=audio_sample_size,
+            target_channels=io_channels,
+            device=device,
+        )
+        if self.model.pretransform is not None:
+            audio = audio.to(next(self.model.pretransform.parameters()).dtype)
+            audio = self.model.pretransform.encode(audio)
+            if inpaint_mask is not None:
+                inpaint_mask = interpolate(
+                    inpaint_mask.unsqueeze(1),
+                    size=audio.shape[-1],
+                    mode="nearest",
+                ).squeeze(1)
+        return audio, inpaint_mask
+
+
+class AutoencoderModel:
+    def __init__(self, autoencoder, sample_rate, device):
+        self.autoencoder = autoencoder
+        self.sample_rate = sample_rate
+        self.device = device
+
+    @staticmethod
+    def from_pretrained(model_name, device=None):
+        if device is None:
+            if torch.cuda.is_available():
+                device = "cuda"
+            elif torch.backends.mps.is_available():
+                device = "mps"
+            else:
+                device = "cpu"
+
+        if not torch.cuda.is_available():
+            if model_name == "same-l":
+                print(
+                    f"Warning: You are loading the {model_name} model without a GPU. This model is not designed to run on cpu"
+                )
+
+        if model_name not in ae_models:
+            raise ValueError(
+                f"Unknown autoencoder '{model_name}'. Valid models: {list(ae_models)}"
+            )
+
+        cfg = ae_models[model_name]
+        local_config, local_ckpt = cfg.resolve()
+
+        with open(local_config) as f:
+            sample_rate = json.load(f)["sample_rate"]
+
+        autoencoder = load_autoencoder(local_config, local_ckpt, device=device)
+        autoencoder.eval().requires_grad_(False)
+
+        return AutoencoderModel(autoencoder, sample_rate, device)
+
+    @torch.inference_mode()
+    def encode(self, audio, sr, chunked=False, chunk_size=128, overlap=32):
+        """Encode audio to latents.
+
+        Args:
+            audio: A single waveform tensor (C, T), a list of waveform tensors,
+                or a pre-batched tensor (B, C, T). Resampling, channel conversion,
+                and padding are handled automatically; passing sr=ae.sample_rate
+                for already-preprocessed audio skips resampling.
+            sr: Sample rate of the input audio, or a list of sample rates when
+                audio is a list.
+            chunked: If True, encode in overlapping chunks to save memory.
+            chunk_size: Chunk size in latent frames (only used when chunked=True).
+            overlap: Overlap in latent frames between chunks (only used when chunked=True).
+
+        Returns:
+            Latent tensor of shape (B, latent_dim, latent_time).
+        """
+        if isinstance(audio, list):
+            preprocessed = self.autoencoder.preprocess_audio_list_for_encoder(
+                audio, in_sr_list=sr
+            )
+        elif isinstance(audio, torch.Tensor) and audio.dim() == 3:
+            sr_list = sr if isinstance(sr, list) else [sr] * audio.shape[0]
+            preprocessed = self.autoencoder.preprocess_audio_list_for_encoder(
+                list(audio), in_sr_list=sr_list
+            )
+        else:
+            preprocessed = self.autoencoder.preprocess_audio_for_encoder(
+                audio, in_sr=sr
+            )
+        return self.autoencoder.encode_audio(
+            preprocessed.to(self.device),
+            chunked=chunked,
+            chunk_size=chunk_size,
+            overlap=overlap,
+        )
+
+    @torch.inference_mode()
+    def decode(self, latents, chunked=False, chunk_size=128, overlap=32):
+        """Decode latents to audio.
+
+        Args:
+            latents: Latent tensor of shape (B, latent_dim, latent_time).
+            chunked: If True, decode in overlapping chunks to save memory.
+            chunk_size: Chunk size in latent frames (only used when chunked=True).
+            overlap: Overlap in latent frames between chunks (only used when chunked=True).
+
+        Returns:
+            Audio tensor of shape (B, channels, samples).
+        """
+        return self.autoencoder.decode_audio(
+            latents,
+            chunked=chunked,
+            chunk_size=chunk_size,
+            overlap=overlap,
+        )

stable_audio_3/model_configs.py

@@ -0,0 +1,116 @@
+from dataclasses import dataclass
+
+from huggingface_hub import hf_hub_download, try_to_load_from_cache
+
+
+@dataclass(frozen=True)
+class ModelConfig:
+    repo_id: str
+    config_path: str
+    ckpt_path: str
+
+    def resolve(self):
+        """Download files from HuggingFace Hub and return local cached paths."""
+        local_config = hf_hub_download(repo_id=self.repo_id, filename=self.config_path)
+        local_ckpt = hf_hub_download(repo_id=self.repo_id, filename=self.ckpt_path)
+        return local_config, local_ckpt
+
+
+@dataclass(frozen=True)
+class AutoencoderModelConfig:
+    """Config for a standalone autoencoder HF repo (e.g. stabilityai/SAME-S).
+
+    resolve() first checks whether any full Stable Audio 3 checkpoint that contains the same
+    autoencoder is already cached locally.  If one is found it is used as-is
+    (load_autoencoder strips the pretransform.* prefix automatically), avoiding a
+    redundant download.  Otherwise the lightweight AE-only repo is fetched instead.
+    """
+
+    ae_repo_id: str
+    ae_config_path: str
+    ae_ckpt_path: str
+    stable_audio_3: tuple[ModelConfig, ...]
+
+    def resolve(self):
+        """Return (config_path, ckpt_path), preferring an already-cached Stable Audio 3 checkpoint."""
+        for fallback in self.stable_audio_3:
+            cached_config = try_to_load_from_cache(
+                fallback.repo_id, fallback.config_path
+            )
+            cached_ckpt = try_to_load_from_cache(fallback.repo_id, fallback.ckpt_path)
+            if isinstance(cached_config, str) and isinstance(cached_ckpt, str):
+                return cached_config, cached_ckpt
+
+        # No Stable Audio 3 checkpoint found in local cache — download the AE-only repo.
+        local_config = hf_hub_download(
+            repo_id=self.ae_repo_id, filename=self.ae_config_path
+        )
+        local_ckpt = hf_hub_download(
+            repo_id=self.ae_repo_id, filename=self.ae_ckpt_path
+        )
+        return local_config, local_ckpt
+
+
+models: dict[str, ModelConfig] = {
+    "small-music": ModelConfig(
+        "stabilityai/stable-audio-3-small-music",
+        "model_config.json",
+        "model.safetensors",
+    ),
+    "small-music-base": ModelConfig(
+        "stabilityai/stable-audio-3-small-music-base",
+        "model_config.json",
+        "model.safetensors",
+    ),
+    "small-sfx": ModelConfig(
+        "stabilityai/stable-audio-3-small-sfx",
+        "model_config.json",
+        "model.safetensors",
+    ),
+    "small-sfx-base": ModelConfig(
+        "stabilityai/stable-audio-3-small-sfx-base",
+        "model_config.json",
+        "model.safetensors",
+    ),
+    "medium": ModelConfig(
+        "stabilityai/stable-audio-3-medium",
+        "model_config.json",
+        "model.safetensors",
+    ),
+    "medium-base": ModelConfig(
+        "stabilityai/stable-audio-3-medium-base",
+        "model_config.json",
+        "model.safetensors",
+    ),
+}
+
+# Stable Audio 3 full-model configs to probe (in order) before downloading the AE-only repo.
+_small_stable_audio_3: tuple[ModelConfig, ...] = (
+    models["small-music"],
+    models["small-sfx"],
+)
+_medium_stable_audio_3: tuple[ModelConfig, ...] = (models["medium"],)
+
+ae_models: dict[str, AutoencoderModelConfig] = {
+    "same-s": AutoencoderModelConfig(
+        ae_repo_id="stabilityai/SAME-S",
+        ae_config_path="model_config.json",
+        ae_ckpt_path="model.safetensors",
+        stable_audio_3=_small_stable_audio_3,
+    ),
+    "same-l": AutoencoderModelConfig(
+        ae_repo_id="stabilityai/SAME-L",
+        ae_config_path="model_config.json",
+        ae_ckpt_path="model.safetensors",
+        stable_audio_3=_medium_stable_audio_3,
+    ),
+}
+
+base_models: dict[str, ModelConfig] = {
+    k: v for k, v in models.items() if k.endswith("-base")
+}
+
+all_models: dict[str, ModelConfig | AutoencoderModelConfig] = {
+    **models,
+    **ae_models,
+}

stable_audio_3/models/autoencoders.py

@@ -0,0 +1,639 @@
+import torch
+
+from torch import nn
+from torch.nn.utils import weight_norm
+from torchaudio import transforms as T
+from einops import rearrange
+
+from ..inference.audio_utils import prepare_audio
+from .transformer import TransformerBlock
+
+def WNConv1d(*args, **kwargs):
+    return weight_norm(nn.Conv1d(*args, **kwargs))
+
+def checkpoint(function, *args, **kwargs):
+    kwargs.setdefault("use_reentrant", False)
+    return torch.utils.checkpoint.checkpoint(function, *args, **kwargs)
+
+
+class Transpose(nn.Module):
+    def __init__(self):
+        super().__init__()
+    def forward(self, x, **kwargs):
+        return rearrange(x, '... a b -> ... b a')
+
+def _zero_pad_modulo_sequence(x, size, dim=-2):
+    input_len = x.shape[dim]
+    pad_len = (size - input_len % size) % size
+    if pad_len > 0:
+        pad_shape = list(x.shape)
+        pad_shape[dim] = pad_len
+        x = torch.cat([x, torch.zeros(pad_shape, device=x.device, dtype=x.dtype)], dim=dim)
+    return x
+
+class TransformerResamplingBlock(nn.Module):
+    def __init__(self, in_channels, out_channels, stride, sliding_window = None, chunk_size = 128, chunk_midpoint_shift = False, type = 'encoder', transformer_depth = 3, checkpointing = False,
+                 conformer = False, layer_scale = False, dim_heads = 128, differential = True, variable_stride = False, feat_scale = False,
+                 sinusoidal_blocks = 0, mask_noise = 0, ff_mult = 3, mapping_bias = True, cross_attn = False, dyt = True, conv_mapping = False, freeze_backbone = False, **kwargs):
+        super().__init__()
+        if type not in ['encoder', 'decoder']:
+            raise ValueError(f"Unknown type {type}. Must be 'encoder' or 'decoder'")
+
+        self.checkpointing = checkpointing
+
+        transformer_dim = out_channels if type == 'encoder' else in_channels
+        transformers = []
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.variable_stride = variable_stride
+        self.stride = stride
+        self.mapping = WNConv1d(in_channels, out_channels, 3 if conv_mapping else 1, padding = 'same', bias = mapping_bias) if in_channels != out_channels else nn.Identity()
+        self.chunk_size = chunk_size
+        self.chunk_midpoint_shift = chunk_midpoint_shift
+        self.type = type
+        self.mask_noise = mask_noise
+        self.sliding_window_latents = sliding_window
+
+        self.sliding_window_seq = self._get_sliding_window_size(sliding_window, stride)
+        self.input_seg_size, self.output_seg_size, self.sub_chunk_size = self._get_seg_sizes(stride)
+        self.transformer_depth = transformer_depth
+        for i in range(transformer_depth):
+            sinusoidal = True if ((transformer_depth - i) < sinusoidal_blocks) else False
+            transformers.append(TransformerBlock(transformer_dim,
+                                                 dim_heads = dim_heads,
+                                                 causal = False,
+                                                 zero_init_branch_outputs = True if not layer_scale else False,
+                                                 norm_type = 'dyt' if dyt else 'rms_norm',
+                                                 conformer = conformer,
+                                                 layer_scale = layer_scale,
+                                                 add_rope = True,
+                                                 attn_kwargs={'qk_norm': "dyt" if dyt else "rms", "qk_norm_eps": 1e-3, "differential": differential, "feat_scale": feat_scale},
+                                                 ff_kwargs={'mult': ff_mult, 'no_bias': False, "sinusoidal": sinusoidal},
+                                                 norm_kwargs = {'eps': 1e-3},
+                                                 cross_attend = cross_attn))
+
+        self.new_tokens = nn.Parameter(1e-5 * torch.randn(1, self.output_seg_size if not self.variable_stride else 1, out_channels if type == 'encoder' else in_channels))
+
+        self.transformers = nn.ModuleList(transformers)
+
+        if freeze_backbone:
+            for param in self.transformers.parameters():
+                param.requires_grad = False
+            self.new_tokens.requires_grad = False
+
+    def _get_sliding_window_size(self, window, stride, prepend_cond_length = 0):
+        if window is None:
+            return None
+        else:
+            return [(win * (stride + 1 + prepend_cond_length)) for win in window]
+
+    def _get_seg_sizes(self, stride, prepend_cond_length = 0):
+        sub_chunk_size = stride + 1 + prepend_cond_length
+        if self.sliding_window_latents is None:
+            assert (self.chunk_size % stride) == 0, f"Stride must fit evenly into chunk size:{self.chunk_size}"
+        input_seg_size = stride if self.type == 'encoder' else 1
+        output_seg_size = 1 if self.type == 'encoder' else stride
+        return input_seg_size, output_seg_size, sub_chunk_size
+
+    #@torch.compile
+    def forward(self, x, stride = None, return_features = False, override_new_tokens = None, prepend_cond = None, cross_attn_cond = None):
+        batch_size = x.shape[0]
+        input_length = x.shape[-1]
+        if return_features:
+            features = []
+
+        if stride == None:
+            input_seg_size = self.input_seg_size
+            output_seg_size = self.output_seg_size
+            sub_chunk_size = self.sub_chunk_size
+            sliding_window = self.sliding_window_seq
+        else:
+            if not self.variable_stride:
+                print("cannot override stride if variable_stride is not set")
+            prepend_cond_length = prepend_cond.shape[-2] if prepend_cond is not None else 0
+            input_seg_size, output_seg_size, sub_chunk_size = self._get_seg_sizes(stride, prepend_cond_length)
+            sliding_window = self._get_sliding_window_size(self.sliding_window_latents, stride, prepend_cond_length)
+
+        if self.type == 'encoder':
+            # Pad before mapping so silence zeros get projected through the mapping,
+            # rather than inserting raw zeros in the mapped space
+            if self.transformer_depth > 0:
+                if sliding_window is None:
+                    pad_modulo = self.chunk_size
+                else:
+                    pad_modulo = input_seg_size
+                x = _zero_pad_modulo_sequence(x, pad_modulo, dim=-1)
+            x = self.mapping(x)
+
+        if self.transformer_depth > 0:
+            x = rearrange(x, '... a b -> ... b a')
+            if return_features:
+                features.append(x)
+            if self.type != 'encoder':
+                if sliding_window is None:
+                    active_stride = stride if stride is not None else self.stride
+                    pad_modulo = self.chunk_size // active_stride
+                    x = _zero_pad_modulo_sequence(x, pad_modulo)
+                else:
+                    x = _zero_pad_modulo_sequence(x, input_seg_size)
+            x = rearrange(x, 'b (n c) d -> (b n) c d', c = input_seg_size)
+            new_token_seq_dim = -1 if not self.variable_stride else output_seg_size
+            new_tokens = self.new_tokens.expand([x.shape[0],new_token_seq_dim,-1])
+            if override_new_tokens is not None:
+                #print(f"Using override new tokens with shape {override_new_tokens.shape}, new tokens shape {new_tokens.shape}, x shape {x.shape}")
+                override_new_tokens = rearrange(override_new_tokens, 'b (n c) d -> (b n) c d', c = output_seg_size)
+                new_tokens = new_tokens + override_new_tokens
+            elif self.mask_noise > 0:
+                new_tokens = new_tokens + torch.randn_like(new_tokens) * self.mask_noise
+            x = torch.cat([x,new_tokens], dim = -2)
+            if prepend_cond is not None:
+                n = x.shape[0] // batch_size
+                cond_folded = prepend_cond.unsqueeze(1).expand(batch_size, n, prepend_cond.shape[-2], x.shape[-1]).reshape(n * batch_size, prepend_cond.shape[-2], x.shape[-1])
+                x = torch.cat([cond_folded, x], dim=-2)   
+            x = rearrange(x, '(b n) c d -> b (n c) d', b=batch_size)#.contiguous()
+
+            # Fold into contiguous chunks if no sliding window
+            if sliding_window is None:
+                prepend_cond_length = prepend_cond.shape[-2] if prepend_cond is not None else 0
+                effective_chunk_size = self.chunk_size + self.chunk_size * (1 + prepend_cond_length) // (stride if stride is not None else self.stride)
+
+            if sliding_window is None and self.chunk_midpoint_shift:
+                split = self.transformer_depth // 2
+                shift = effective_chunk_size // 2
+
+                # First half: standard chunks
+                nc = x.shape[1] // effective_chunk_size
+                x = rearrange(x, 'b (nc cc) d -> (b nc) cc d', cc=effective_chunk_size)
+                cross_attn_first = None
+                if cross_attn_cond is not None:
+                    cross_attn_first = cross_attn_cond.repeat_interleave(nc, dim=0)
+                for layer in self.transformers[:split]:
+                    if self.checkpointing:
+                        x = checkpoint(layer, x, context = cross_attn_first, self_attention_flash_sliding_window = None)
+                    else:
+                        x = layer(x, context = cross_attn_first)
+                    if return_features:
+                        features.append(rearrange(x, '(b nc) cc d -> b (nc cc) d', b=batch_size))
+                x = rearrange(x, '(b nc) cc d -> b (nc cc) d', b=batch_size)
+
+                # Second half: shifted chunks with sub-chunk repeat padding
+                # shift is always a multiple of sub_chunk_size, so slicing
+                # by shift gives whole sub-chunks in their original order
+                x = torch.cat([x[:, :shift, :], x, x[:, -shift:, :]], dim=1)
+                nc_shifted = x.shape[1] // effective_chunk_size
+                x = rearrange(x, 'b (nc cc) d -> (b nc) cc d', cc=effective_chunk_size)
+                cross_attn_second = None
+                if cross_attn_cond is not None:
+                    cross_attn_second = cross_attn_cond.repeat_interleave(nc_shifted, dim=0)
+                for layer in self.transformers[split:]:
+                    if self.checkpointing:
+                        x = checkpoint(layer, x, context = cross_attn_second, self_attention_flash_sliding_window = None)
+                    else:
+                        x = layer(x, context = cross_attn_second)
+                    if return_features:
+                        feat = rearrange(x, '(b nc) cc d -> b (nc cc) d', b=batch_size)
+                        features.append(feat[:, shift:-shift, :])
+                x = rearrange(x, '(b nc) cc d -> b (nc cc) d', b=batch_size)
+                x = x[:, shift:-shift, :]
+            else:
+                if sliding_window is None:
+                    x = rearrange(x, 'b (nc cc) d -> (b nc) cc d', cc=effective_chunk_size)
+
+                for layer in self.transformers:
+                    if self.checkpointing:
+                        x = checkpoint(layer, x, context = cross_attn_cond, self_attention_flash_sliding_window = sliding_window)
+                    else:
+                        x = layer(x, context = cross_attn_cond, self_attention_flash_sliding_window = sliding_window)
+                    if return_features:
+                        features.append(x)
+
+                # Unfold chunks back to original batch
+                if sliding_window is None:
+                    x = rearrange(x, '(b nc) cc d -> b (nc cc) d', b=batch_size)
+
+            x = rearrange(x, 'b (n c) d -> (b n) c d', c=sub_chunk_size)
+            x = x[:,-output_seg_size:,:]
+            x = rearrange(x, '(b n) c d -> b d (n c)', b=batch_size)
+        if self.type == 'decoder':
+            x = self.mapping(x)
+        if return_features:
+            return x, features
+        else:   
+            return x
+
+
+class SAMEEncoder(nn.Module):
+    def __init__(self,
+                 in_channels=2,
+                 channels=128,
+                 latent_dim=32,
+                 c_mults = [1, 2, 4, 8],
+                 strides = [2, 4, 8, 8],
+                 transformer_depths = [3,3,3,3],
+                 sliding_window = None,
+                 checkpointing = False,
+                 conformer = False,
+                 layer_scale = False,
+                 causal = False,
+                 differential = True,
+                 variable_stride = False,
+                 mask_noise = 0.0,
+                 conv_mapping = False,
+                 freeze_backbone = False,
+                 **kwargs
+        ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.strides = strides
+
+        channel_dims = [c * channels for c in c_mults]
+        channel_dims = [in_channels] + channel_dims
+
+        self.depth = len(c_mults)
+
+        layers = []
+
+        for i in range(self.depth):
+            layers += [TransformerResamplingBlock(in_channels=channel_dims[i], out_channels=channel_dims[i+1], stride=strides[i], transformer_depth = transformer_depths[i],
+                                                  sliding_window = sliding_window, checkpointing = checkpointing, conformer = conformer, layer_scale = layer_scale, causal = causal,
+                                                  differential = differential, variable_stride = variable_stride, mask_noise = mask_noise, conv_mapping = conv_mapping,
+                                                  freeze_backbone = freeze_backbone, **kwargs)]
+
+        layers += [Transpose(), nn.Linear(channel_dims[-1], latent_dim), Transpose()]
+        self.layers = nn.ModuleList(layers)
+
+        if freeze_backbone:
+            for param in self.layers[-2].parameters():
+                param.requires_grad = False
+
+    def forward(self, x, override_stride = None, return_features = False, **kwargs):
+        if override_stride != None:
+            assert isinstance(override_stride, list), "override_stride must be a list"
+            assert len(override_stride) == self.depth, "override_stride must be a list containing strides for every layer"
+        for i, layer in enumerate(self.layers):
+            if isinstance(layer, TransformerResamplingBlock):
+                if override_stride != None:
+                    stride = override_stride[i]
+                else:
+                    stride = None
+                if return_features:
+                    x, features = layer(x, stride = stride, return_features = True)
+                else:
+                    x = layer(x, stride = stride)
+            else:
+                x = layer(x)
+        if return_features:
+            return x, features
+        else:
+            return x
+
+class SAMEDecoder(nn.Module):
+    def __init__(self,
+                 out_channels=2,
+                 channels=128,
+                 latent_dim=32,
+                 c_mults = [1, 2, 4, 8],
+                 strides = [2, 4, 8, 8],
+                 transformer_depths = [3,3,3,3],
+                 sliding_window = None,
+                 checkpointing = False,
+                 conformer = False,
+                 layer_scale = False,
+                 causal = False,
+                 differential = True,
+                 variable_stride = False,
+                 sinusoidal_blocks = [0,0,0,0],
+                 mask_noise = 0.0,
+                 conv_mapping = False,
+                 freeze_backbone = False,
+                 **kwargs
+        ):
+        super().__init__()
+
+        channel_dims = [c * channels for c in c_mults]
+        channel_dims = [out_channels] + channel_dims
+
+        self.depth = len(c_mults)
+
+        layers = [Transpose(), nn.Linear(latent_dim, channel_dims[-1]), Transpose()]
+
+        for i in range(self.depth, 0, -1):
+            layers += [TransformerResamplingBlock(in_channels=channel_dims[i], out_channels=channel_dims[i-1], stride=strides[i-1], type = 'decoder', transformer_depth = transformer_depths[i-1],
+                                                  sliding_window = sliding_window, checkpointing = checkpointing, conformer = conformer, layer_scale = layer_scale, causal = causal, differential = differential,
+                                                  variable_stride = variable_stride, sinusoidal_blocks = sinusoidal_blocks[i-1], mask_noise = mask_noise, conv_mapping = conv_mapping,
+                                                  freeze_backbone = freeze_backbone, **kwargs)]
+
+        self.layers = nn.ModuleList(layers)
+
+        if freeze_backbone:
+            for param in self.layers[1].parameters():
+                param.requires_grad = False
+
+    def forward(self, x, override_stride = None, **kwargs):
+        if override_stride != None:
+            assert isinstance(override_stride, list), "override_stride must be a list"
+            assert len(override_stride) == self.depth, "override_stride must be a list containing strides for every layer"
+
+        transformer_layer_index = 0
+        for i, layer in enumerate(self.layers):
+            if isinstance(layer, TransformerResamplingBlock):
+                if override_stride != None:
+                    stride = override_stride[transformer_layer_index]
+                else:
+                    stride = None
+                x = layer(x, stride = stride)
+                transformer_layer_index += 1
+            else:
+                x = layer(x)
+        return x
+
+
+class AudioAutoencoder(nn.Module):
+    def __init__(
+        self,
+        encoder,
+        decoder,
+        latent_dim,
+        downsampling_ratio,
+        sample_rate,
+        io_channels=2,
+        bottleneck: nn.Module = None,
+        pretransform: nn.Module = None,
+        in_channels = None,
+        out_channels = None,
+        soft_clip = False,
+        freeze_pretransform = False
+    ):
+        super().__init__()  
+        self.downsampling_ratio = downsampling_ratio
+        self.sample_rate = sample_rate
+
+        self.latent_dim = latent_dim
+        self.io_channels = io_channels
+        self.in_channels = io_channels
+        self.out_channels = io_channels
+
+        self.min_length = self.downsampling_ratio
+
+        if in_channels is not None:
+            self.in_channels = in_channels
+
+        if out_channels is not None:
+            self.out_channels = out_channels
+
+        self.bottleneck = bottleneck
+
+        self.encoder = encoder
+
+        self.decoder = decoder
+
+        self.pretransform = pretransform
+
+        self.freeze_pretransform = freeze_pretransform
+        if self.pretransform is not None:
+            if self.freeze_pretransform:
+                for p in self.pretransform.parameters():
+                    p.requires_grad = False
+            else:
+                for p in self.pretransform.parameters():
+                    p.requires_grad = True
+
+        self.soft_clip = soft_clip
+
+        self.is_discrete = False
+
+    def encode(self, audio, return_info=False, skip_pretransform=False, iterate_batch=False, return_pretransform = False, **kwargs):
+
+        info = {}
+        if self.pretransform is not None and not skip_pretransform: 
+            if self.pretransform.enable_grad:
+                if iterate_batch:
+                    audios = []
+                    for i in range(audio.shape[0]):
+                        audios.append(self.pretransform.encode(audio[i:i+1]))
+                    audio = torch.cat(audios, dim=0)
+                else:
+                    audio = self.pretransform.encode(audio)
+            else:
+                with torch.no_grad():
+                    if iterate_batch:
+                        audios = []
+                        for i in range(audio.shape[0]):
+                            audios.append(self.pretransform.encode(audio[i:i+1]))
+                        audio = torch.cat(audios, dim=0)
+                    else:
+                        audio = self.pretransform.encode(audio)
+        if self.encoder is not None:
+            if iterate_batch:
+                latents = []
+                for i in range(audio.shape[0]):
+                    latents.append(self.encoder(audio[i:i+1], **kwargs))
+                latents = torch.cat(latents, dim=0)
+            else:
+                latents = self.encoder(audio, **kwargs)
+        else:
+            latents = audio
+        if self.bottleneck is not None:
+            # TODO: Add iterate batch logic, needs to merge the info dicts
+            latents, bottleneck_info = self.bottleneck.encode(latents, return_info=True, **kwargs)
+
+            info.update(bottleneck_info)
+        
+        if return_info and return_pretransform:
+            return latents, info, audio
+        elif return_info:
+            return latents, info
+        elif return_pretransform:
+            return latents, audio
+        else:
+            return latents
+
+    def decode(self, latents, iterate_batch=False, return_loss = False, **kwargs):
+        if self.bottleneck is not None:
+            if iterate_batch:
+                decoded = []
+                for i in range(latents.shape[0]):
+                    decoded.append(self.bottleneck.decode(latents[i:i+1]))
+                latents = torch.cat(decoded, dim=0)
+            else:
+                latents = self.bottleneck.decode(latents)
+        if iterate_batch:
+            decoded = []
+            for i in range(latents.shape[0]):
+                decoded.append(self.decoder(latents[i:i+1], **kwargs))
+            decoded = torch.cat(decoded, dim=0)
+        else:
+            if return_loss:
+                decoded, loss = self.decoder(latents, **kwargs)
+            else:
+                decoded = self.decoder(latents, **kwargs)
+        if self.pretransform is not None:
+            if self.pretransform.enable_grad:
+                if iterate_batch:
+                    decodeds = []
+                    for i in range(decoded.shape[0]):
+                        decodeds.append(self.pretransform.decode(decoded[i:i+1]))
+                    decoded = torch.cat(decodeds, dim=0)
+                else:
+                    decoded = self.pretransform.decode(decoded)
+            else:
+                with torch.no_grad():
+                    if iterate_batch:
+                        decodeds = []
+                        for i in range(latents.shape[0]):
+                            decodeds.append(self.pretransform.decode(decoded[i:i+1]))
+                        decoded = torch.cat(decodeds, dim=0)
+                    else:
+                        decoded = self.pretransform.decode(decoded)
+
+        if self.soft_clip:
+            decoded = torch.tanh(decoded)
+        if return_loss:
+            return decoded, loss
+        else:
+            return decoded
+          
+    
+    def preprocess_audio_for_encoder(self, audio, in_sr):
+        '''
+        Preprocess single audio tensor (Channels x Length) to be compatible with the encoder.
+        If the model is mono, stereo audio will be converted to mono.
+        Audio will be silence-padded to be a multiple of the model's downsampling ratio.
+        Audio will be resampled to the model's sample rate. 
+        The output will have batch size 1 and be shape (1 x Channels x Length)
+        '''
+        return self.preprocess_audio_list_for_encoder([audio], [in_sr])
+
+    def preprocess_audio_list_for_encoder(self, audio_list, in_sr_list):
+        '''
+        Preprocess a [list] of audio (Channels x Length) into a batch tensor to be compatable with the encoder. 
+        The audio in that list can be of different lengths and channels. 
+        in_sr can be an integer or list. If it's an integer it will be assumed it is the input sample_rate for every audio.
+        All audio will be resampled to the model's sample rate. 
+        Audio will be silence-padded to the longest length, and further padded to be a multiple of the model's downsampling ratio. 
+        If the model is mono, all audio will be converted to mono. 
+        The output will be a tensor of shape (Batch x Channels x Length)
+        '''
+        batch_size = len(audio_list)
+        if isinstance(in_sr_list, int):
+            in_sr_list = [in_sr_list]*batch_size
+        assert len(in_sr_list) == batch_size, "list of sample rates must be the same length of audio_list"
+        new_audio = []
+        max_length = 0
+        # resample & find the max length
+        for i in range(batch_size):
+            audio = audio_list[i]
+            in_sr = in_sr_list[i]
+            if len(audio.shape) == 3 and audio.shape[0] == 1:
+                # batchsize 1 was given by accident. Just squeeze it.
+                audio = audio.squeeze(0)
+            elif len(audio.shape) == 1:
+                # Mono signal, channel dimension is missing, unsqueeze it in
+                audio = audio.unsqueeze(0)
+            assert len(audio.shape)==2, "Audio should be shape (Channels x Length) with no batch dimension" 
+            # Resample audio
+            if in_sr != self.sample_rate:
+                resample_tf = T.Resample(in_sr, self.sample_rate).to(audio.device)
+                audio = resample_tf(audio)
+            new_audio.append(audio)
+            if audio.shape[-1] > max_length:
+                max_length = audio.shape[-1]
+        # Pad every audio to the same length, multiple of model's downsampling ratio
+        padded_audio_length = max_length + (self.min_length - (max_length % self.min_length)) % self.min_length
+        for i in range(batch_size):
+            # Pad it & if necessary, mixdown/duplicate stereo/mono channels to support model
+            new_audio[i] = prepare_audio(new_audio[i], in_sr=in_sr, target_sr=in_sr, target_length=padded_audio_length, 
+                target_channels=self.in_channels, device=new_audio[i].device).squeeze(0)
+        # convert to tensor 
+        return torch.stack(new_audio) 
+
+    def encode_audio(self, audio, chunked=False, overlap=32, chunk_size=128, **kwargs):
+        '''
+        Encode audios into latents. Audios should already be preprocesed by preprocess_audio_for_encoder.
+        If chunked is True, split the audio into chunks of a given maximum size chunk_size, with given overlap.
+        Overlap and chunk_size params are both measured in number of latents (not audio samples) 
+        # and therefore you likely could use the same values with decode_audio. 
+        A overlap of zero will cause discontinuity artefacts. Overlap should be => receptive field size. 
+        Every autoencoder will have a different receptive field size, and thus ideal overlap.
+        You can determine it empirically by diffing unchunked vs chunked output and looking at maximum diff.
+        The final chunk may have a longer overlap in order to keep chunk_size consistent for all chunks.
+        Smaller chunk_size uses less memory, but more compute.
+        The chunk_size vs memory tradeoff isn't linear, and possibly depends on the GPU and CUDA version
+        For example, on a A6000 chunk_size 128 is overall faster than 256 and 512 even though it has more chunks
+        '''
+        samples_per_latent = int(self.downsampling_ratio)
+        if not chunked or audio.shape[-1] < chunk_size * samples_per_latent:
+            return self.encode(audio, **kwargs)
+
+        # chunk_size/overlap are in latent units; scale to samples for slicing.
+        chunk_size_samples = chunk_size * samples_per_latent
+        hop_samples = (chunk_size - overlap) * samples_per_latent
+        total_samples = audio.shape[-1]
+
+        # Anchor the final chunk to the signal end (may overlap more than `overlap`).
+        chunk_starts = list(range(0, total_samples - chunk_size_samples + 1, hop_samples))
+        if chunk_starts[-1] != total_samples - chunk_size_samples:
+            chunk_starts.append(total_samples - chunk_size_samples)
+
+        encoded_chunks = [self.encode(audio[..., s:s + chunk_size_samples]) for s in chunk_starts]
+
+        total_latents = total_samples // samples_per_latent
+        half_overlap_latents = overlap // 2
+        output = audio.new_zeros(*encoded_chunks[0].shape[:-1], total_latents)
+        num_chunks = len(chunk_starts)
+
+        # Trim half the overlap off inner edges (edges facing a neighbour chunk).
+        for i, (start_sample, chunk) in enumerate(zip(chunk_starts, encoded_chunks)):
+            is_first = i == 0
+            is_last = i == num_chunks - 1
+            out_start = (total_latents - chunk_size) if is_last else (start_sample // samples_per_latent)
+            left  = 0 if is_first else half_overlap_latents
+            right = chunk_size if is_last else chunk_size - half_overlap_latents
+            output[..., out_start + left : out_start + right] = chunk[..., left:right]
+
+        return output
+    
+    def decode_audio(self, latents, chunked=False, overlap=32, chunk_size=128, **kwargs):
+        '''
+        Decode latents to audio. 
+        If chunked is True, split the latents into chunks of a given maximum size chunk_size, with given overlap, both of which are measured in number of latents. 
+        A overlap of zero will cause discontinuity artefacts. Overlap should be => receptive field size. 
+        Every autoencoder will have a different receptive field size, and thus ideal overlap.
+        You can determine it empirically by diffing unchunked vs chunked audio and looking at maximum diff.
+        The final chunk may have a longer overlap in order to keep chunk_size consistent for all chunks.
+        Smaller chunk_size uses less memory, but more compute.
+        The chunk_size vs memory tradeoff isn't linear, and possibly depends on the GPU and CUDA version
+        For example, on a A6000 chunk_size 128 is overall faster than 256 and 512 even though it has more chunks
+        '''
+        if not chunked or latents.shape[-1] < chunk_size:
+            return self.decode(latents, **kwargs)
+
+        # chunk_size/overlap are in latent units — same as `latents`, no scaling.
+        samples_per_latent = int(self.downsampling_ratio)
+        hop_latents = chunk_size - overlap
+        total_latents = latents.shape[-1]
+
+        # Anchor the final chunk to the signal end (may overlap more than `overlap`).
+        chunk_starts = list(range(0, total_latents - chunk_size + 1, hop_latents))
+        if chunk_starts[-1] != total_latents - chunk_size:
+            chunk_starts.append(total_latents - chunk_size)
+
+        decoded_chunks = [self.decode(latents[..., s:s + chunk_size]) for s in chunk_starts]
+
+        total_samples = total_latents * samples_per_latent
+        chunk_size_samples = chunk_size * samples_per_latent
+        half_overlap_samples = (overlap // 2) * samples_per_latent
+        output = latents.new_zeros(*decoded_chunks[0].shape[:-1], total_samples)
+        num_chunks = len(chunk_starts)
+
+        # Trim half the overlap off inner edges (edges facing a neighbour chunk).
+        for i, (start_latent, chunk) in enumerate(zip(chunk_starts, decoded_chunks)):
+            is_first = i == 0
+            is_last = i == num_chunks - 1
+            out_start = (total_samples - chunk_size_samples) if is_last else (start_latent * samples_per_latent)
+            left  = 0 if is_first else half_overlap_samples
+            right = chunk_size_samples if is_last else chunk_size_samples - half_overlap_samples
+            output[..., out_start + left : out_start + right] = chunk[..., left:right]
+
+        return output
+        

stable_audio_3/models/blocks.py

@@ -0,0 +1,84 @@
+import math
+import torch
+from torch import nn
+from torch.nn.utils import weight_norm
+
+
+def get_activation(activation, channels=None) -> nn.Module:
+    if activation == "elu":
+        return nn.ELU()
+    elif activation == "none":
+        return nn.Identity()
+    else:
+        raise ValueError(f"Unknown activation {activation}")
+
+def WNConv1d(*args, **kwargs):
+    return weight_norm(nn.Conv1d(*args, **kwargs))
+
+class ResidualUnit(nn.Module):
+    def __init__(self, in_channels, out_channels, dilation, depthwise=False, bias=True):
+        super().__init__()
+
+        self.dilation = dilation
+
+        padding = (dilation * (7-1)) // 2
+
+        self.layers = nn.Sequential(
+            get_activation("elu"),
+            WNConv1d(in_channels=in_channels, out_channels=out_channels,
+                      kernel_size=7, dilation=dilation, padding=padding, groups=1 if not depthwise else out_channels, bias=bias),
+            get_activation("elu"),
+            WNConv1d(in_channels=out_channels, out_channels=out_channels,
+                      kernel_size=1, bias=bias)
+        )
+
+    def forward(self, x):
+        res = x
+        x = self.layers(x)
+        return x + res
+
+class FourierFeatures(nn.Module):
+    def __init__(self, in_features, out_features, std=16.):
+        super().__init__()
+        assert out_features % 2 == 0
+        self.register_buffer('weight', torch.randn([out_features // 2, in_features]) * std)
+
+    def forward(self, input):
+        f = 2 * math.pi * input @ self.weight.T
+        return torch.cat([f.cos(), f.sin()], dim=-1)
+
+class ExpoFourierFeatures(nn.Module):
+    def __init__(self, dim, min_freq=0.5, max_freq=10000.0):
+        super().__init__()
+        self.dim = dim
+        self.min_freq = min_freq
+        self.max_freq = max_freq
+
+    @torch.amp.autocast("cuda",enabled=False)
+    def forward(self, t):
+        """
+        t: [B] tensor.
+        """
+        in_dtype = t.dtype
+        t = t.float()
+
+        if t.dim() == 1:
+            t = t.unsqueeze(-1)
+
+        half_dim = self.dim // 2
+
+        # Calculate frequencies (safely in FP32)
+        ramp = torch.linspace(0, 1, half_dim, device=t.device, dtype=torch.float32)
+        log_min = math.log(self.min_freq)
+        log_max = math.log(self.max_freq)
+
+        freqs = torch.exp(ramp * (log_max - log_min) + log_min)
+
+        # Calculate arguments (safely in FP32)
+        args = t * freqs * 2 * math.pi
+
+        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+
+        return embedding.to(in_dtype)
+
+

stable_audio_3/models/bottleneck.py

@@ -0,0 +1,71 @@
+import torch
+from torch import nn
+
+class SoftNormBottleneck(nn.Module):
+    def __init__(self, dim = 32, noise_augment_dim=0, noise_regularize = False, auto_scale = False, freeze = False):
+        super().__init__()
+
+        self.noise_augment_dim = noise_augment_dim
+        self.scaling_factor = nn.Parameter(torch.ones(1,dim,1))
+        self.bias = nn.Parameter(torch.zeros(1,dim,1))
+        self.noise_scaling_factor = nn.Parameter(torch.ones(1,noise_augment_dim,1))
+        self.noise_regularize = noise_regularize
+        self.freeze = freeze
+        if self.freeze:
+            self.scaling_factor.requires_grad = False
+            self.bias.requires_grad = False
+            self.noise_scaling_factor.requires_grad = False
+        if auto_scale:
+            running_std = torch.ones(1)
+            self.register_parameter("running_std", nn.Parameter(running_std, requires_grad=False))
+
+    def encode(self, x, return_info=False, **kwargs):
+        info = {}
+
+        x = x * self.scaling_factor + self.bias
+
+        if self.training and hasattr(self,"running_std") and not self.freeze:
+            # Update running std
+            self.running_std.data = (self.running_std.data * 0.999 + x.std().detach() * 0.001).clamp(min = 1e-4)
+
+        if hasattr(self, "running_std"):
+            x = x / self.running_std
+
+        if self.training and return_info:
+            var = (x.std(dim=-1) ** 2).clip(min = 1e-4)
+            logvar = torch.log(var)
+            mean = x.mean(dim=-1)
+            loss = (mean * mean + var - logvar - 1).mean()
+            var = (x.std(dim=-2) ** 2).clip(min = 1e-4)
+            logvar = torch.log(var)
+            mean = x.mean(dim=-2)
+            loss = loss + 0.4 * (mean * mean + var - logvar - 1).mean()
+            info["softnorm_loss"] = loss
+
+        if return_info:
+            return x, info
+
+        return x
+
+    def decode(self, x, **kwargs):
+        if hasattr(self, "running_std"):
+            x = x * self.running_std
+
+        if self.noise_regularize:
+            if hasattr(self, "running_std"):
+                scaling = self.running_std
+            else:
+                scaling = x.std(dim = -1).unsqueeze(-1)
+            if self.training:
+                scale = 5e-2
+            else:
+                scale = 1e-3
+            noise = torch.randn_like(x) * scaling * scale
+            x = x + noise
+
+        if self.noise_augment_dim > 0:
+            noise = self.noise_scaling_factor * torch.randn(x.shape[0], self.noise_augment_dim,
+                                x.shape[-1]).type_as(x)
+            x = torch.cat([x, noise], dim=1)
+
+        return x

stable_audio_3/models/conditioners.py

@@ -0,0 +1,318 @@
+#Heavily influenced by https://github.com/facebookresearch/audiocraft/blob/main/audiocraft/modules/conditioners.py
+
+import torch
+import logging, warnings
+import typing as tp
+from enum import Enum
+from math import pi
+from typing import List, Union
+from torch import Tensor, nn
+from einops import rearrange
+from .blocks import ExpoFourierFeatures
+from .utils import enable_torch_compile
+import os
+
+class PaddingMode(str, Enum):
+    """Enum for handling padding in text conditioner embeddings."""
+    NONE = "none"       # No padding handling (raw embeddings with pad token)
+    ZERO = "zero"       # Zero out padding positions (default)
+    LEARNED = "learned" # Use learned padding embedding
+
+class Conditioner(nn.Module):
+    def __init__(
+            self,
+            dim: int,
+            output_dim: int,
+            project_out: bool = False,
+            padding_mode: str = "zero"
+            ):
+
+        super().__init__()
+
+        self.dim = dim
+        self.output_dim = output_dim
+        self.padding_mode = padding_mode
+        self.proj_out = nn.Linear(dim, output_dim) if (dim != output_dim or project_out) else nn.Identity()
+
+        # Learned padding embedding (only created if needed)
+        if padding_mode == "learned" or padding_mode == PaddingMode.LEARNED:
+            self.padding_embedding = nn.Parameter(torch.randn(output_dim) * 0.02)
+
+    def apply_padding(self, embeddings: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
+        """
+        Apply padding handling based on padding_mode.
+
+        Args:
+            embeddings: [batch, seq_len, dim] - the embeddings to process
+            attention_mask: [batch, seq_len] bool/int, True/1 = valid token
+
+        Returns:
+            embeddings with padding handled according to mode
+        """
+        mode = self.padding_mode
+        if isinstance(mode, str):
+            mode = PaddingMode(mode)
+
+        if mode == PaddingMode.NONE:
+            return embeddings
+        elif mode == PaddingMode.ZERO:
+            return embeddings * attention_mask.unsqueeze(-1).float()
+        elif mode == PaddingMode.LEARNED:
+            mask_expanded = attention_mask.unsqueeze(-1).bool()
+            return torch.where(
+                mask_expanded,
+                embeddings,
+                self.padding_embedding.unsqueeze(0).unsqueeze(0).expand_as(embeddings)
+            )
+        else:
+            raise ValueError(f"Unknown padding mode: {mode}")
+
+    def forward(self, x: tp.Any) -> tp.Any:
+        raise NotImplementedError()
+
+class LearnedPositionalEmbedding(nn.Module):
+    def __init__(self, dim: int, std: float = 16.0):
+        super().__init__()
+        assert (dim % 2) == 0
+        half_dim = dim // 2
+        self.weights = nn.Parameter(torch.randn(half_dim) * std)
+
+    def forward(self, x: Tensor) -> Tensor:
+        x = rearrange(x, "b -> b 1")
+        freqs = x * rearrange(self.weights, "d -> 1 d") * 2 * pi
+        fouriered = torch.cat((freqs.sin(), freqs.cos()), dim=-1)
+        fouriered = torch.cat((x, fouriered), dim=-1)
+        return fouriered
+
+
+def TimePositionalEmbedding(dim: int, out_features: int) -> nn.Module:
+    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,
+        fourier_features_type: tp.Literal["learned", "expo"] = "learned"
+    ):
+        super().__init__()
+        self.features = features
+        if fourier_features_type == "expo":
+            self.embedding = nn.Sequential(ExpoFourierFeatures(dim=dim), nn.Linear(in_features=dim, out_features=features))
+        else:
+            self.embedding = TimePositionalEmbedding(dim=dim, out_features=features)
+
+    def forward(self, x: Union[List[float], Tensor]) -> Tensor:
+        if not torch.is_tensor(x):
+            device = next(self.embedding.parameters()).device
+            x = torch.tensor(x, device=device)
+        assert isinstance(x, Tensor)
+        shape = x.shape
+        x = rearrange(x, "... -> (...)")
+        embedding = self.embedding(x)
+        x = embedding.view(*shape, self.features)
+        return x  # type: ignore
+
+
+class NumberConditioner(Conditioner):
+    '''
+        Conditioner that takes a list of floats, normalizes them for a given range, and returns a list of embeddings
+    '''
+    def __init__(self, 
+                output_dim: int,
+                min_val: float=0,
+                max_val: float=1,
+                fourier_features_type : tp.Literal["learned", "expo"] = "learned"
+                ):
+        super().__init__(output_dim, output_dim)
+
+        self.min_val = min_val
+        self.max_val = max_val
+
+        self.embedder = NumberEmbedder(features=output_dim, fourier_features_type=fourier_features_type)
+
+    def forward(self, floats: tp.List[float], device=None) -> tp.Any:
+            self.embedder.to(device)
+            # Cast the inputs to floats
+            floats = [float(x) for x in floats]
+
+            floats = torch.tensor(floats).to(device)
+
+            floats = floats.clamp(self.min_val, self.max_val)
+    
+            normalized_floats = (floats - self.min_val) / (self.max_val - self.min_val)
+
+            # Cast floats to same type as embedder
+            embedder_dtype = next(self.embedder.parameters()).dtype
+            normalized_floats = normalized_floats.to(embedder_dtype)
+
+            float_embeds = self.embedder(normalized_floats).unsqueeze(1)
+    
+            return [float_embeds, torch.ones(float_embeds.shape[0], 1).to(device)]
+
+class T5GemmaConditioner(Conditioner):
+
+    T5GEMMA_MODELS = ["google/t5gemma-b-b-ul2"]
+
+    T5GEMMA_MODEL_DIMS = {
+        "google/t5gemma-b-b-ul2": 768,
+    }
+
+    def __init__(
+            self,
+            output_dim: int,
+            model_name: str = "google/t5gemma-b-b-ul2",
+            max_length: str = 128,
+            enable_grad: bool = False,
+            project_out: bool = False,
+            padding_mode: str = "zero",
+            model_path: str = None,
+            repo_id: str = None,
+            subfolder: str = None,
+    ):
+        assert model_name in self.T5GEMMA_MODELS, f"Unknown T5 model name: {model_name}"
+        super().__init__(self.T5GEMMA_MODEL_DIMS[model_name], output_dim, project_out=project_out, padding_mode=padding_mode)
+
+        load_from = model_path or repo_id or model_name
+
+        self.max_length = max_length
+        self.enable_grad = enable_grad
+
+        # Set environment variables to disable progress bars BEFORE importing transformers
+        # This is the most reliable way to suppress HuggingFace progress bars
+        prev_hf_hub = os.environ.get("HF_HUB_DISABLE_PROGRESS_BARS")
+        prev_transformers = os.environ.get("TRANSFORMERS_VERBOSITY")
+        os.environ["HF_HUB_DISABLE_PROGRESS_BARS"] = "1"
+        os.environ["TRANSFORMERS_VERBOSITY"] = "error"
+
+        # Suppress logging from transformers
+        previous_level = logging.root.manager.disable
+        logging.disable(logging.ERROR)
+
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore")
+            try:
+                from transformers import T5GemmaEncoderModel, AutoTokenizer, AutoConfig
+                logging.info(f"Loading T5Gemma tokenizer and model from: {load_from}")
+                hf_kwargs = {"subfolder": subfolder} if subfolder else {}
+                self.tokenizer = AutoTokenizer.from_pretrained(load_from, **hf_kwargs)
+                config = AutoConfig.from_pretrained(load_from, **hf_kwargs)
+                config.is_encoder_decoder = False
+                model = T5GemmaEncoderModel.from_pretrained(load_from, config=config, **hf_kwargs).train(enable_grad).requires_grad_(enable_grad)
+
+            finally:
+                logging.disable(previous_level)
+                # Restore environment variables
+                if prev_hf_hub is None:
+                    os.environ.pop("HF_HUB_DISABLE_PROGRESS_BARS", None)
+                else:
+                    os.environ["HF_HUB_DISABLE_PROGRESS_BARS"] = prev_hf_hub
+                if prev_transformers is None:
+                    os.environ.pop("TRANSFORMERS_VERBOSITY", None)
+                else:
+                    os.environ["TRANSFORMERS_VERBOSITY"] = prev_transformers
+
+        # Compile the model to reduce CPU-GPU kernel launch overhead,
+        # which is sensitive to CPU contention from DataLoader workers
+        if enable_torch_compile:
+            model = torch.compile(model)
+
+        if self.enable_grad:
+            self.model = model
+        else:
+            self.__dict__["model"] = model
+
+        self._device_initialized = False
+
+    def forward(self, inputs: tp.Union[tp.List[str], tp.List[tp.Dict[str, torch.Tensor]]], device: tp.Union[torch.device, str]) -> tp.Tuple[torch.Tensor, torch.Tensor]:
+
+        # Only move to device once (avoid overhead on every forward call)
+        if not self._device_initialized:
+            self.model.to(device)
+            self.proj_out.to(device)
+            self.model.eval()
+            self._device_initialized = True
+
+        # Handle pre-tokenized inputs (dicts with input_ids/attention_mask from DataLoader workers)
+        # or raw strings (from demo generation / inference)
+        if isinstance(inputs[0], dict):
+            input_ids = torch.stack([x["input_ids"] for x in inputs]).to(device, non_blocking=True)
+            attention_mask = torch.stack([x["attention_mask"] for x in inputs]).to(device, non_blocking=True).to(torch.bool)
+        else:
+            encoded = self.tokenizer(
+                inputs,
+                truncation=True,
+                max_length=self.max_length,
+                padding="max_length",
+                return_tensors="pt",
+            )
+            input_ids = encoded["input_ids"].to(device, non_blocking=True)
+            attention_mask = encoded["attention_mask"].to(device, non_blocking=True).to(torch.bool)
+
+        with torch.no_grad():
+            embeddings = self.model(
+                input_ids=input_ids, attention_mask=attention_mask
+            )["last_hidden_state"]
+
+        # Cast embeddings to same type as proj_out, unless proj_out is Identity
+        if not isinstance(self.proj_out, nn.Identity):
+            proj_out_dtype = next(self.proj_out.parameters()).dtype
+            embeddings = embeddings.to(proj_out_dtype)
+
+        embeddings = self.proj_out(embeddings)
+        embeddings = self.apply_padding(embeddings, attention_mask)
+
+        return embeddings, attention_mask
+
+
+class MultiConditioner(nn.Module):
+    """
+    A module that applies multiple conditioners to an input dictionary based on the keys
+
+    Args:
+        conditioners: a dictionary of conditioners with keys corresponding to the keys of the conditioning input dictionary (e.g. "prompt")
+        default_keys: a dictionary of default keys to use if the key is not in the input dictionary (e.g. {"prompt_t5": "prompt"})
+    """
+    def __init__(self, conditioners: tp.Dict[str, Conditioner], default_keys: tp.Dict[str, str] = {}, pre_encoded_keys: tp.List[str] = []):
+        super().__init__()
+
+        self.conditioners = nn.ModuleDict(conditioners)
+        self.default_keys = default_keys
+        self.pre_encoded_keys = pre_encoded_keys
+
+    def forward(self, batch_metadata: tp.List[tp.Dict[str, tp.Any]], device: tp.Union[torch.device, str]) -> tp.Dict[str, tp.Any]:
+        output = {}
+
+        for key, conditioner in self.conditioners.items():
+            condition_key = key
+
+            conditioner_inputs = []
+
+            for x in batch_metadata:
+
+                if condition_key not in x:
+                    if condition_key in self.default_keys:
+                        condition_key = self.default_keys[condition_key]
+                    else:
+                        raise ValueError(f"Conditioner key {condition_key} not found in batch metadata")
+
+                #Unwrap the condition info if it's a single-element list or tuple, this is to support collation functions that wrap everything in a list
+                if isinstance(x[condition_key], list) or isinstance(x[condition_key], tuple) and len(x[condition_key]) == 1:
+                    conditioner_input = x[condition_key][0]
+                    
+                else:
+                    conditioner_input = x[condition_key]
+
+                conditioner_inputs.append(conditioner_input)
+
+            if key in self.pre_encoded_keys:
+                output[key] = [torch.stack(conditioner_inputs, dim=0).to(device), None]
+            else:
+                output[key] = conditioner(conditioner_inputs, device)
+
+        return output
+    

stable_audio_3/models/diffusion.py

@@ -0,0 +1,253 @@
+import torch
+from torch import nn
+import typing as tp
+
+from .conditioners import MultiConditioner
+from .dit import DiffusionTransformer
+from stable_audio_3.inference.distribution_shift import FluxDistributionShift, DistributionShift, LogSNRShift, IdentityDistributionShift
+
+from time import time
+
+class Profiler:
+
+    def __init__(self):
+        self.ticks = [[time(), None]]
+
+    def tick(self, msg):
+        self.ticks.append([time(), msg])
+
+    def __repr__(self):
+        rep = 80 * "=" + "\n"
+        for i in range(1, len(self.ticks)):
+            msg = self.ticks[i][1]
+            ellapsed = self.ticks[i][0] - self.ticks[i - 1][0]
+            rep += msg + f": {ellapsed*1000:.2f}ms\n"
+        rep += 80 * "=" + "\n\n\n"
+        return rep
+
+class ConditionedDiffusionModelWrapper(nn.Module):
+    """
+    A diffusion model that takes in conditioning
+    """
+    def __init__(
+            self,
+            model: nn.Module,
+            conditioner: MultiConditioner,
+            io_channels,
+            sample_rate,
+            min_input_length: int,
+            diffusion_objective: tp.Literal["v", "rectified_flow", "rf_denoiser"] = "v",
+            distribution_shift_options = None,
+            sampling_distribution_shift_options = None,
+            mask_padding_attention: bool = False,
+            use_effective_length_for_schedule: bool = False,
+            pretransform: tp.Optional[nn.Module] = None,
+            cross_attn_cond_ids: tp.List[str] = [],
+            global_cond_ids: tp.List[str] = [],
+            input_concat_ids: tp.List[str] = [],
+            local_add_cond_ids: tp.List[str] = [],
+            modular_local_cond_ids: tp.List[str] = [],
+            prepend_cond_ids: tp.List[str] = [],
+            ):
+        super().__init__()
+
+        self.model = model
+        self.conditioner = conditioner
+        self.io_channels = io_channels
+        self.sample_rate = sample_rate
+        self.diffusion_objective = diffusion_objective
+        self.pretransform = pretransform
+        self.cross_attn_cond_ids = cross_attn_cond_ids
+        self.global_cond_ids = global_cond_ids
+        self.input_concat_ids = input_concat_ids
+        self.local_add_cond_ids = local_add_cond_ids
+        self.modular_local_cond_ids = modular_local_cond_ids
+        self.prepend_cond_ids = prepend_cond_ids
+        self.min_input_length = min_input_length
+        self.mask_padding_attention = mask_padding_attention
+        self.use_effective_length_for_schedule = use_effective_length_for_schedule
+
+        self.dist_shift = None
+        if distribution_shift_options is not None:
+            self.dist_shift = self._create_dist_shift(distribution_shift_options)
+
+        # Sampling dist_shift: separate config for inference-time schedule
+        if sampling_distribution_shift_options is not None:
+            self.sampling_dist_shift = self._create_dist_shift(sampling_distribution_shift_options)
+        else:
+            # Default: seq_len-invariant LogSNR shift matching legacy log_snr_sampling=True
+            self.sampling_dist_shift = LogSNRShift(rate=0, anchor_logsnr=-6.2, logsnr_end=2.0)
+
+    @staticmethod
+    def _create_dist_shift(options: dict):
+        """Create a distribution shift object from config options."""
+        dist_shift_type = options.get("type", "full")
+        dist_shift_kwargs = {k: v for k, v in options.items() if k != "type"}
+        if dist_shift_type == "none":
+            return IdentityDistributionShift()
+        elif dist_shift_type == "flux":
+            return FluxDistributionShift(**dist_shift_kwargs)
+        elif dist_shift_type == "full":
+            return DistributionShift(**dist_shift_kwargs)
+        elif dist_shift_type == "logsnr":
+            return LogSNRShift(**dist_shift_kwargs)
+        else:
+            raise ValueError(f"Unknown distribution shift type: {dist_shift_type}. Expected 'none', 'flux', 'full', or 'logsnr'.")     
+
+    def get_conditioning_inputs(self, conditioning_tensors: tp.Dict[str, tp.Any], negative=False):
+        cross_attention_input = None
+        cross_attention_masks = None
+        global_cond = None
+        input_concat_cond = None
+        prepend_cond = None
+        prepend_cond_mask = None
+        local_add_cond = None
+        modular_local_cond = None
+
+        if len(self.cross_attn_cond_ids) > 0:
+            # Concatenate all cross-attention inputs over the sequence dimension
+            # Assumes that the cross-attention inputs are of shape (batch, seq, channels)
+            cross_attention_input = []
+            cross_attention_masks = []
+
+            for key in self.cross_attn_cond_ids:
+                cross_attn_in, cross_attn_mask = conditioning_tensors[key]
+
+                # Add sequence dimension if it's not there
+                if len(cross_attn_in.shape) == 2:
+                    cross_attn_in = cross_attn_in.unsqueeze(1)
+                    cross_attn_mask = cross_attn_mask.unsqueeze(1)
+
+                cross_attention_input.append(cross_attn_in)
+                cross_attention_masks.append(cross_attn_mask)
+
+            cross_attention_input = torch.cat(cross_attention_input, dim=1)
+            cross_attention_masks = torch.cat(cross_attention_masks, dim=1)
+
+        if len(self.global_cond_ids) > 0:
+            # Concatenate all global conditioning inputs over the channel dimension
+            # Assumes that the global conditioning inputs are of shape (batch, channels)
+            global_conds = []
+            for key in self.global_cond_ids:
+                global_cond_input = conditioning_tensors[key][0]
+
+                global_conds.append(global_cond_input)
+
+            # Concatenate over the channel dimension
+            global_cond = torch.cat(global_conds, dim=-1)
+
+            if len(global_cond.shape) == 3:
+                global_cond = global_cond.squeeze(1)
+
+        if len(self.input_concat_ids) > 0:
+            # Concatenate all input concat conditioning inputs over the channel dimension
+            # Assumes that the input concat conditioning inputs are of shape (batch, channels, seq)
+            input_concat_cond = torch.cat([conditioning_tensors[key][0] for key in self.input_concat_ids], dim=1)
+
+        if len(self.local_add_cond_ids) > 0:
+            # Concatenate all local conditioning inputs over the channel dimension
+            # Assumes that the local conditioning inputs are of shape (batch, channels, seq)
+            local_add_cond = torch.cat([conditioning_tensors[key][0] for key in self.local_add_cond_ids], dim=1)
+
+        if len(self.modular_local_cond_ids) > 0:
+            # Keep modular local conditioning as a dict of tensors (not concatenated)
+            # Each tensor is of shape (batch, channels, seq)
+            modular_local_cond = {}
+            for key in self.modular_local_cond_ids:
+                if key in conditioning_tensors:
+                    modular_local_cond[key] = conditioning_tensors[key][0]
+            # Only set if we have any conditioning
+            if len(modular_local_cond) == 0:
+                modular_local_cond = None
+
+        if len(self.prepend_cond_ids) > 0:
+            # Concatenate all prepend conditioning inputs over the sequence dimension
+            # Assumes that the prepend conditioning inputs are of shape (batch, seq, channels)
+            prepend_conds = []
+            prepend_cond_masks = []
+
+            for key in self.prepend_cond_ids:
+                prepend_cond_input, prepend_cond_mask = conditioning_tensors[key]
+                prepend_conds.append(prepend_cond_input)
+                prepend_cond_masks.append(prepend_cond_mask)
+
+            prepend_cond = torch.cat(prepend_conds, dim=1)
+            prepend_cond_mask = torch.cat(prepend_cond_masks, dim=1)
+
+        if negative:
+            return {
+                "negative_cross_attn_cond": cross_attention_input,
+                "negative_cross_attn_mask": cross_attention_masks,
+                "negative_global_cond": global_cond,
+                "negative_input_concat_cond": input_concat_cond
+            }
+        else:
+            return {
+                "cross_attn_cond": cross_attention_input,
+                "cross_attn_mask": cross_attention_masks,
+                "global_cond": global_cond,
+                "input_concat_cond": input_concat_cond,
+                "local_add_cond": local_add_cond,
+                "modular_local_cond": modular_local_cond,
+                "prepend_cond": prepend_cond,
+                "prepend_cond_mask": prepend_cond_mask
+            }
+
+    def forward(self, x: torch.Tensor, t: torch.Tensor, cond: tp.Dict[str, tp.Any], **kwargs):
+        return self.model(x, t, **self.get_conditioning_inputs(cond), **kwargs)
+
+
+class DiTWrapper(nn.Module):
+    def __init__(
+        self,
+        diffusion_objective: str,
+        *args,
+        **kwargs
+    ):
+        super().__init__()
+
+        self.diffusion_objective = diffusion_objective
+
+        self.model = DiffusionTransformer(diffusion_objective=diffusion_objective, *args, **kwargs)
+
+    def forward(self,
+                x,
+                t,
+                cross_attn_cond=None,
+                cross_attn_mask=None,
+                negative_cross_attn_cond=None,
+                negative_cross_attn_mask=None,
+                input_concat_cond=None,
+                local_add_cond=None,
+                negative_input_concat_cond=None,
+                global_cond=None,
+                negative_global_cond=None,
+                prepend_cond=None,
+                prepend_cond_mask=None,
+                cfg_scale=1.0,
+                cfg_dropout_prob: float = 0.0,
+                batch_cfg: bool = True,
+                rescale_cfg: bool = False,
+                scale_phi: float = 0.0,
+                **kwargs):
+
+        assert batch_cfg, "batch_cfg must be True for DiTWrapper"
+        #assert negative_input_concat_cond is None, "negative_input_concat_cond is not supported for DiTWrapper"
+
+        return self.model(
+            x,
+            t,
+            cross_attn_cond=cross_attn_cond,
+            cross_attn_cond_mask=cross_attn_mask,
+            negative_cross_attn_cond=negative_cross_attn_cond,
+            negative_cross_attn_mask=negative_cross_attn_mask,
+            input_concat_cond=input_concat_cond,
+            prepend_cond=prepend_cond,
+            prepend_cond_mask=prepend_cond_mask,
+            cfg_scale=cfg_scale,
+            cfg_dropout_prob=cfg_dropout_prob,
+            scale_phi=scale_phi,
+            global_embed=global_cond,
+            local_add_cond=local_add_cond,
+            **kwargs)
+

stable_audio_3/models/dit.py

@@ -0,0 +1,643 @@
+import typing as tp
+import math
+import torch
+
+from einops import rearrange
+from torch import nn
+from torch.nn import functional as F
+
+from .blocks import FourierFeatures, ExpoFourierFeatures
+from .transformer import ContinuousTransformer        
+from .lora import LoRAParametrization, set_lora_strength, has_lora, enable_lora, disable_lora, filter_lora_layers
+
+class DiffusionTransformer(nn.Module):
+    def __init__(self,
+        io_channels=32,
+        patch_size=1,
+        embed_dim=768,
+        cond_token_dim=0,
+        project_cond_tokens=True,
+        global_cond_dim=0,
+        project_global_cond=True,
+        input_concat_dim=0,
+        prepend_cond_dim=0,
+        depth=12,
+        num_heads=8,
+        transformer_type: tp.Literal["continuous_transformer", "mm_transformer"] = "continuous_transformer",
+        global_cond_type: tp.Literal["prepend", "adaLN"] = "prepend",
+        timestep_cond_type: tp.Literal["global", "input_concat"] = "global",
+        timestep_embed_dim=None,
+        diffusion_objective: tp.Literal["v", "rectified_flow", "rf_denoiser"] = "v",
+        timestep_features_type: tp.Literal["learned", "expo"] = "learned",
+        timestep_features_dim = 256,
+        timestep_features_logsnr: bool = False,
+        modular_local_cond_configs = None,
+        **kwargs):
+
+        super().__init__()
+
+        self.cond_token_dim = cond_token_dim
+
+        # Timestep embeddings
+        self.timestep_cond_type = timestep_cond_type
+        self.timestep_features_logsnr = timestep_features_logsnr
+
+        timestep_features_dim = timestep_features_dim
+
+        if timestep_features_type == "expo":
+            self.timestep_features = ExpoFourierFeatures(timestep_features_dim, 0.5, 10000.0)
+        else:
+            self.timestep_features = FourierFeatures(1, timestep_features_dim)
+
+        if timestep_cond_type == "global":
+            timestep_embed_dim = embed_dim
+        elif timestep_cond_type == "input_concat":
+            assert timestep_embed_dim is not None, "timestep_embed_dim must be specified if timestep_cond_type is input_concat"
+            input_concat_dim += timestep_embed_dim
+
+        self.to_timestep_embed = nn.Sequential(
+            nn.Linear(timestep_features_dim, timestep_embed_dim, bias=True),
+            nn.SiLU(),
+            nn.Linear(timestep_embed_dim, timestep_embed_dim, bias=True),
+        )
+        
+        self.diffusion_objective = diffusion_objective
+
+        if cond_token_dim > 0:
+            # Conditioning tokens
+
+            cond_embed_dim = cond_token_dim if not project_cond_tokens else embed_dim
+            self.to_cond_embed = nn.Sequential(
+                nn.Linear(cond_token_dim, cond_embed_dim, bias=False),
+                nn.SiLU(),
+                nn.Linear(cond_embed_dim, cond_embed_dim, bias=False)
+            )
+        else:
+            cond_embed_dim = 0
+
+        if global_cond_dim > 0:
+            # Global conditioning
+            global_embed_dim = global_cond_dim if not project_global_cond else embed_dim
+            self.to_global_embed = nn.Sequential(
+                nn.Linear(global_cond_dim, global_embed_dim, bias=False),
+                nn.SiLU(),
+                nn.Linear(global_embed_dim, global_embed_dim, bias=False)
+            )
+
+        if prepend_cond_dim > 0:
+            # Prepend conditioning
+            self.to_prepend_embed = nn.Sequential(
+                nn.Linear(prepend_cond_dim, embed_dim, bias=False),
+                nn.SiLU(),
+                nn.Linear(embed_dim, embed_dim, bias=False)
+            )
+
+        self.input_concat_dim = input_concat_dim
+
+        dim_in = io_channels + self.input_concat_dim
+
+        self.patch_size = patch_size
+
+        # Transformer
+
+        self.transformer_type = transformer_type
+
+        self.global_cond_type = global_cond_type
+
+        transformer_dim_out = io_channels * patch_size
+
+        if self.transformer_type == "continuous_transformer":
+
+            global_dim = None
+
+            if self.global_cond_type == "adaLN":
+                # The global conditioning is projected to the embed_dim already at this point
+                global_dim = embed_dim
+
+            self.transformer = ContinuousTransformer(
+                dim=embed_dim,
+                depth=depth,
+                dim_heads=embed_dim // num_heads,
+                dim_in=dim_in * patch_size,
+                dim_out=transformer_dim_out,
+                cross_attend = cond_token_dim > 0,
+                cond_token_dim = cond_embed_dim,
+                global_cond_dim=global_dim,
+                modular_local_cond_configs=modular_local_cond_configs,
+                **kwargs
+            )
+      
+        else:
+            raise ValueError(f"Unknown transformer type: {self.transformer_type}")
+
+        self.preprocess_conv = nn.Conv1d(dim_in, dim_in, 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)
+
+    # Fixed logsnr normalization range: maps logsnr to [0, 1] preserving direction (t=0→0, t=1→1)
+    _LOGSNR_MIN = -12.0
+    _LOGSNR_MAX = 5.0
+    _LOGSNR_RANGE = _LOGSNR_MAX - _LOGSNR_MIN
+
+    def _t_to_logsnr_cond(self, t: torch.Tensor) -> torch.Tensor:
+        """Convert t to normalized logsnr in [0, 1] for timestep conditioning.
+
+        Maps t through logsnr = log((1-t)/t), clamps to fixed range,
+        then normalizes to [0, 1] preserving direction (t=0→0, t=1→1).
+        """
+        t_clamped = t.float().clamp(1e-7, 1 - 1e-7)
+        logsnr = torch.log((1 - t_clamped) / t_clamped)
+        logsnr = logsnr.clamp(self._LOGSNR_MIN, self._LOGSNR_MAX)
+        return ((self._LOGSNR_MAX - logsnr) / self._LOGSNR_RANGE).to(t.dtype)
+
+    def _call_transformer(self, x, *, prepend_inputs=None, cross_attn_cond=None,
+                         mask=None, prepend_mask=None, return_info=False,
+                         exit_layer_ix=None, local_add_cond=None,
+                         modular_local_cond=None, padding_mask=None,
+                         extra_args=None, **kwargs):
+        """Helper method to call transformer and handle early exit logic."""
+
+        output = self.transformer(x, prepend_embeds=prepend_inputs, context=cross_attn_cond,
+                                    return_info=return_info, exit_layer_ix=exit_layer_ix,
+                                    local_add_cond=local_add_cond, modular_local_cond=modular_local_cond,
+                                    padding_mask=padding_mask,
+                                    **(extra_args or {}), **kwargs)
+
+        if return_info:
+            output, info = output
+
+        # Avoid postprocessing on early exit
+        if exit_layer_ix is not None:
+            if return_info:
+                return output, info
+            else:
+                return output
+
+        return (output, info) if return_info and 'info' in locals() else output
+
+    def _forward(
+        self,
+        x,
+        t,
+        mask=None,
+        cross_attn_cond=None,
+        cross_attn_cond_mask=None,
+        input_concat_cond=None,
+        local_add_cond=None,
+        modular_local_cond=None,
+        global_embed=None,
+        prepend_cond=None,
+        prepend_cond_mask=None,
+        padding_mask=None,
+        return_info=False,
+        exit_layer_ix=None,
+        **kwargs):
+
+        if cross_attn_cond is not None:
+            cross_attn_cond = self.to_cond_embed(cross_attn_cond)
+
+        if global_embed is not None:
+            # Project the global conditioning to the embedding dimension
+            global_embed = self.to_global_embed(global_embed)
+
+        prepend_inputs = None 
+        prepend_mask = None
+        prepend_length = 0
+        if prepend_cond is not None:
+            # Project the prepend conditioning to the embedding dimension
+            prepend_cond = self.to_prepend_embed(prepend_cond)
+            
+            prepend_inputs = prepend_cond
+            if prepend_cond_mask is not None:
+                prepend_mask = prepend_cond_mask
+
+            prepend_length = prepend_cond.shape[1]
+
+        if input_concat_cond is not None:
+            # Interpolate input_concat_cond to the same length as x
+            if input_concat_cond.shape[2] != x.shape[2]:
+                input_concat_cond = F.interpolate(input_concat_cond, (x.shape[2], ), mode='nearest')
+
+            x = torch.cat([x, input_concat_cond], dim=1)
+
+        if local_add_cond is not None:
+            local_add_cond = rearrange(local_add_cond, "b c t -> b t c")
+
+        # Rearrange modular_local_cond tensors
+        if modular_local_cond is not None:
+            modular_local_cond = {
+                k: rearrange(v, "b c t -> b t c")
+                for k, v in modular_local_cond.items()
+            }
+
+        # Get the batch of timestep embeddings
+        t_cond = self._t_to_logsnr_cond(t) if self.timestep_features_logsnr else t
+        # Convert to model dtype for linear layers (t itself is kept in float32 for precision)
+        # x has already been converted to model dtype in the outer forward() method
+        t_cond = t_cond.to(x.dtype)
+        timestep_embed = self.to_timestep_embed(self.timestep_features(t_cond[:, None])) # (b, embed_dim)
+
+        # Timestep embedding is considered a global embedding. Add to the global conditioning if it exists
+
+        if self.timestep_cond_type == "global":
+            if global_embed is not None:
+                global_embed = global_embed + timestep_embed
+            else:
+                global_embed = timestep_embed
+        elif self.timestep_cond_type == "input_concat":
+            x = torch.cat([x, timestep_embed.unsqueeze(2).expand(-1, -1, x.shape[2])], dim=1)
+
+        # Add the global_embed to the prepend inputs if there is no global conditioning support in the transformer
+        if self.global_cond_type == "prepend" and global_embed is not None:
+            if prepend_inputs is None:
+                # Prepend inputs are just the global embed, and the mask is all ones
+                prepend_inputs = global_embed.unsqueeze(1)
+                prepend_mask = torch.ones((x.shape[0], 1), device=x.device, dtype=torch.bool)
+            else:
+                # Prepend inputs are the prepend conditioning + the global embed
+                prepend_inputs = torch.cat([prepend_inputs, global_embed.unsqueeze(1)], dim=1)
+                prepend_mask = torch.cat([prepend_mask, torch.ones((x.shape[0], 1), device=x.device, dtype=torch.bool)], dim=1)
+
+            prepend_length = prepend_inputs.shape[1]
+
+        x = self.preprocess_conv(x) + x
+
+        x = rearrange(x, "b c t -> b t c")
+
+        extra_args = {}
+
+        if self.global_cond_type == "adaLN":
+            extra_args["global_cond"] = global_embed
+
+        if self.patch_size > 1:
+            x = rearrange(x, "b (t p) c -> b t (c p)", p=self.patch_size)
+
+        result = self._call_transformer(
+            x,
+            prepend_inputs=prepend_inputs,
+            cross_attn_cond=cross_attn_cond,
+            mask=mask,
+            prepend_mask=prepend_mask,
+            return_info=return_info,
+            exit_layer_ix=exit_layer_ix,
+            local_add_cond=local_add_cond,
+            modular_local_cond=modular_local_cond,
+            padding_mask=padding_mask,
+            extra_args=extra_args,
+            **kwargs,
+        )
+
+        # Handle early exit (result contains both output and info)
+        if exit_layer_ix is not None:
+            return result
+
+        output = result[0] if return_info else result
+        if return_info:
+            info = result[1]
+
+        output = rearrange(output, "b t c -> b c t")[:,:,prepend_length:]       
+
+        if self.patch_size > 1:
+            output = rearrange(output, "b (c p) t -> b c (t p)", p=self.patch_size)
+
+        output = self.postprocess_conv(output) + output
+
+        if return_info:
+            return output, info
+
+        return output
+
+    def apg_project(self, v0, v1, padding_mask=None):
+        """
+        Project v0 into components parallel and orthogonal to v1.
+
+        Args:
+            v0: Tensor to project (B, C, T)
+            v1: Reference direction (B, C, T)
+            padding_mask: Optional mask (B, T) where True = valid, False = padding.
+                          If provided, only valid positions contribute to the projection.
+        """
+        dtype = v0.dtype
+        v0, v1 = v0.double(), v1.double()
+
+        if padding_mask is not None:
+            # Expand mask to match tensor shape: (B, T) -> (B, 1, T)
+            mask = padding_mask.unsqueeze(1).double()
+            # Zero out padding positions for projection computation
+            v0_masked = v0 * mask
+            v1_masked = v1 * mask
+            # Normalize only over valid positions
+            v1_norm = v1_masked.norm(dim=[-1, -2], keepdim=True).clamp(min=1e-8)
+            v1_normalized = v1_masked / v1_norm
+            # Compute projection using masked values
+            v0_parallel = (v0_masked * v1_normalized).sum(dim=[-1, -2], keepdim=True) * v1_normalized
+            # Orthogonal component: subtract parallel from original (not masked) v0
+            # but apply mask to ensure padding stays zero
+            v0_orthogonal = (v0 - (v0 * v1_normalized).sum(dim=[-1, -2], keepdim=True) * v1_normalized) * mask
+        else:
+            v1 = torch.nn.functional.normalize(v1, dim=[-1, -2])
+            v0_parallel = (v0 * v1).sum(dim=[-1, -2], keepdim=True) * v1
+            v0_orthogonal = v0 - v0_parallel
+
+        return v0_parallel.to(dtype), v0_orthogonal.to(dtype)
+
+    def forward(
+        self,
+        x,
+        t,
+        cross_attn_cond=None,
+        cross_attn_cond_mask=None,
+        negative_cross_attn_cond=None,
+        negative_cross_attn_mask=None,
+        input_concat_cond=None,
+        local_add_cond=None,
+        modular_local_cond=None,
+        global_embed=None,
+        negative_global_embed=None,
+        prepend_cond=None,
+        prepend_cond_mask=None,
+        padding_mask=None,
+        cfg_scale=1.0,
+        cfg_dropout_prob=0.0,
+        cfg_interval = (0, 1),
+        lora_interval = (0, 1),
+        lora_layer_filter = "",
+        lora_configs = None,
+        causal=False,
+        scale_phi=0.0,
+        cfg_norm_threshold=0.0,
+        apg_scale=1.0,
+        mask=None,
+        return_info=False,
+        exit_layer_ix=None,
+        **kwargs):
+
+        assert not causal, "Causal mode is not supported for DiffusionTransformer"
+
+        model_dtype = next(self.parameters()).dtype
+
+        x = x.to(model_dtype)
+
+        # Keep t in float32: the logsnr transform log((1-t)/t) amplifies bf16
+        # quantization error ~380x near t=1, causing catastrophic conditioning errors.
+        # t is a 1D batch-size tensor so float32 has zero memory impact.
+        t = t.float()
+
+        if cross_attn_cond is not None:
+            cross_attn_cond = cross_attn_cond.to(model_dtype)
+
+        if negative_cross_attn_cond is not None:
+            negative_cross_attn_cond = negative_cross_attn_cond.to(model_dtype)
+
+        if input_concat_cond is not None:
+            input_concat_cond = input_concat_cond.to(model_dtype)
+
+        if local_add_cond is not None:
+            local_add_cond = local_add_cond.to(model_dtype)
+
+        if modular_local_cond is not None:
+            modular_local_cond = {k: v.to(model_dtype) for k, v in modular_local_cond.items()}
+
+        if global_embed is not None:
+            global_embed = global_embed.to(model_dtype)
+
+        if negative_global_embed is not None:
+            negative_global_embed = negative_global_embed.to(model_dtype)
+
+        if prepend_cond is not None:
+            prepend_cond = prepend_cond.to(model_dtype)
+
+        if cross_attn_cond_mask is not None:
+            cross_attn_cond_mask = cross_attn_cond_mask.bool()
+
+            cross_attn_cond_mask = None # Temporarily disabling conditioning masks due to kernel issue for flash attention
+
+        if prepend_cond_mask is not None:
+            prepend_cond_mask = prepend_cond_mask.bool()
+
+        # Early exit bypasses CFG processing
+        if exit_layer_ix is not None:
+            assert self.transformer_type == "continuous_transformer", "exit_layer_ix is only supported for continuous_transformer"
+            return self._forward(
+                x,
+                t,
+                cross_attn_cond=cross_attn_cond,
+                cross_attn_cond_mask=cross_attn_cond_mask,
+                input_concat_cond=input_concat_cond,
+                local_add_cond=local_add_cond,
+                modular_local_cond=modular_local_cond,
+                global_embed=global_embed,
+                prepend_cond=prepend_cond,
+                prepend_cond_mask=prepend_cond_mask,
+                padding_mask=padding_mask,
+                mask=mask,
+                return_info=return_info,
+                exit_layer_ix=exit_layer_ix,
+                **kwargs
+            )
+
+        # CFG dropout
+        if cfg_dropout_prob > 0.0 and cfg_scale == 1.0:
+            if cross_attn_cond is not None:
+                null_embed = torch.zeros_like(cross_attn_cond, device=cross_attn_cond.device)
+                dropout_mask = torch.bernoulli(torch.full((cross_attn_cond.shape[0], 1, 1), cfg_dropout_prob, device=cross_attn_cond.device)).to(torch.bool)
+                cross_attn_cond = torch.where(dropout_mask, null_embed, cross_attn_cond)
+
+            if prepend_cond is not None:
+                null_embed = torch.zeros_like(prepend_cond, device=prepend_cond.device)
+                dropout_mask = torch.bernoulli(torch.full((prepend_cond.shape[0], 1, 1), cfg_dropout_prob, device=prepend_cond.device)).to(torch.bool)
+                prepend_cond = torch.where(dropout_mask, null_embed, prepend_cond)
+
+        if self.diffusion_objective == "v":
+            sigma = torch.sin(t * math.pi / 2)
+            alpha = torch.cos(t * math.pi / 2)
+        elif self.diffusion_objective in ["rectified_flow", "rf_denoiser"]:
+            sigma = t
+
+        # LoRA interval
+        if has_lora(self):
+            if lora_configs is not None:
+                # Multi-LoRA: per-LoRA interval and layer filter
+                for lora_config in lora_configs:
+                    idx = lora_config["lora_index"]
+                    interval = lora_config.get("interval", (0, 1))
+                    layer_filter = lora_config.get("layer_filter", "")
+                    if interval[0] <= sigma[0] <= interval[1]:
+                        enable_lora(self, lora_index=idx)
+                        filter_lora_layers(self, layer_filter, lora_index=idx)
+                    else:
+                        disable_lora(self, lora_index=idx)
+            else:
+                # Legacy single-LoRA path
+                if lora_interval[0] <= sigma[0] <= lora_interval[1]:
+                    enable_lora(self)
+                    filter_lora_layers(self, lora_layer_filter)
+                else:
+                    disable_lora(self)
+
+        if cfg_scale != 1.0 and (cross_attn_cond is not None or prepend_cond is not None) and (cfg_interval[0] <= sigma[0] <= cfg_interval[1]):
+
+            # Classifier-free guidance
+            # Concatenate conditioned and unconditioned inputs on the batch dimension            
+            batch_inputs = torch.cat([x, x], dim=0)
+            batch_timestep = torch.cat([t, t], dim=0)
+
+            if global_embed is not None:
+                batch_global_cond = torch.cat([global_embed, global_embed], dim=0)
+            else:
+                batch_global_cond = None
+
+            if input_concat_cond is not None:
+                batch_input_concat_cond = torch.cat([input_concat_cond, input_concat_cond], dim=0)
+            else:
+                batch_input_concat_cond = None
+
+            if local_add_cond is not None:
+                batch_local_add_cond = torch.cat([local_add_cond, local_add_cond], dim=0)
+            else:
+                batch_local_add_cond = None
+
+            if modular_local_cond is not None:
+                batch_modular_local_cond = {k: torch.cat([v, v], dim=0) for k, v in modular_local_cond.items()}
+            else:
+                batch_modular_local_cond = None
+
+            batch_cond = None
+            batch_cond_masks = None
+            
+            # Handle CFG for cross-attention conditioning
+            if cross_attn_cond is not None:
+
+                null_embed = torch.zeros_like(cross_attn_cond, device=cross_attn_cond.device)
+
+                # For negative cross-attention conditioning, replace the null embed with the negative cross-attention conditioning
+                if negative_cross_attn_cond is not None:
+
+                    # If there's a negative cross-attention mask, set the masked tokens to the null embed
+                    if negative_cross_attn_mask is not None:
+                        negative_cross_attn_mask = negative_cross_attn_mask.to(torch.bool).unsqueeze(2)
+
+                        negative_cross_attn_cond = torch.where(negative_cross_attn_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, null_embed], dim=0)
+
+                if cross_attn_cond_mask is not None:
+                    batch_cond_masks = torch.cat([cross_attn_cond_mask, cross_attn_cond_mask], dim=0)
+               
+            batch_prepend_cond = None
+            batch_prepend_cond_mask = None
+
+            if prepend_cond is not None:
+
+                null_embed = torch.zeros_like(prepend_cond, device=prepend_cond.device)
+
+                batch_prepend_cond = torch.cat([prepend_cond, null_embed], dim=0)
+                           
+                if prepend_cond_mask is not None:
+                    batch_prepend_cond_mask = torch.cat([prepend_cond_mask, prepend_cond_mask], dim=0)
+         
+
+            if mask is not None:
+                batch_masks = torch.cat([mask, mask], dim=0)
+            else:
+                batch_masks = None
+
+            if padding_mask is not None:
+                batch_padding_mask = torch.cat([padding_mask, padding_mask], dim=0)
+            else:
+                batch_padding_mask = None
+
+            batch_output = self._forward(
+                batch_inputs,
+                batch_timestep,
+                cross_attn_cond=batch_cond,
+                cross_attn_cond_mask=batch_cond_masks,
+                mask = batch_masks,
+                input_concat_cond = batch_input_concat_cond,
+                local_add_cond = batch_local_add_cond,
+                modular_local_cond=batch_modular_local_cond,
+                global_embed = batch_global_cond,
+                prepend_cond = batch_prepend_cond,
+                prepend_cond_mask = batch_prepend_cond_mask,
+                padding_mask = batch_padding_mask,
+                return_info = return_info,
+                **kwargs)
+
+            if return_info:
+                batch_output, info = batch_output
+
+            cond_output, uncond_output = torch.chunk(batch_output, 2, dim=0)
+
+            if self.diffusion_objective == "v":
+                cond_denoised = x * alpha[:, None, None] - cond_output * sigma[:, None, None]
+                uncond_denoised = x * alpha[:, None, None] - uncond_output * sigma[:, None, None]
+
+            elif self.diffusion_objective in ["rectified_flow", "rf_denoiser"]:
+                cond_denoised = x - cond_output * sigma[:, None, None]
+                uncond_denoised = x - uncond_output * sigma[:, None, None]
+
+            diff = cond_denoised - uncond_denoised
+            
+            if cfg_norm_threshold > 0:
+                if padding_mask is not None:
+                    # Only compute norm over valid positions
+                    mask = padding_mask.unsqueeze(1).float()  # (B, 1, T)
+                    diff_masked = diff * mask
+                    diff_norm = diff_masked.norm(p=2, dim=[-1, -2], keepdim=True)
+                else:
+                    diff_norm = diff.norm(p=2, dim=[-1, -2], keepdim=True)
+                scale_factor = torch.minimum(torch.ones_like(diff), cfg_norm_threshold / diff_norm)
+                diff *= scale_factor
+
+            if apg_scale == 0.0:
+                # Vanilla CFG: use full diff
+                cfg_diff = diff
+            elif apg_scale == 1.0:
+                # Full APG: use only orthogonal component
+                _, diff_orthogonal = self.apg_project(diff, cond_denoised, padding_mask=padding_mask)
+                cfg_diff = diff_orthogonal
+            else:
+                # Blended APG: interpolate between full diff and orthogonal
+                diff_parallel, diff_orthogonal = self.apg_project(diff, cond_denoised, padding_mask=padding_mask)
+                cfg_diff = apg_scale * diff_orthogonal + (1 - apg_scale) * diff
+
+            cfg_denoised = cond_denoised + (cfg_scale - 1) * cfg_diff
+                    
+            if self.diffusion_objective == "v":
+                output = (x * alpha[:, None, None] - cfg_denoised) / sigma[:, None, None]
+            elif self.diffusion_objective in ["rectified_flow", "rf_denoiser"]:
+                output = (x - cfg_denoised) / sigma[:, None, None]
+
+            # CFG Rescale
+            if scale_phi != 0.0:
+                cond_out_std = cond_output.std(dim=1, keepdim=True)
+                out_cfg_std = output.std(dim=1, keepdim=True)
+                output = scale_phi * (output * (cond_out_std/out_cfg_std)) + (1-scale_phi) * output
+           
+            if return_info:
+                info["uncond_output"] = uncond_output
+                return output, info
+
+            return output
+            
+        else:
+            return self._forward(
+                x,
+                t,
+                cross_attn_cond=cross_attn_cond,
+                cross_attn_cond_mask=cross_attn_cond_mask,
+                input_concat_cond=input_concat_cond,
+                local_add_cond=local_add_cond,
+                modular_local_cond=modular_local_cond,
+                global_embed=global_embed,
+                prepend_cond=prepend_cond,
+                prepend_cond_mask=prepend_cond_mask,
+                padding_mask=padding_mask,
+                mask=mask,
+                return_info=return_info,
+                **kwargs
+            )

stable_audio_3/models/inpainting.py

@@ -0,0 +1,242 @@
+import random
+import torch
+from enum import Enum
+from typing import List, Optional, Tuple
+
+class MaskType(Enum):
+    RANDOM_SEGMENTS = 0  # Legacy: uncontrolled ratio, overlapping segments
+    FULL_MASK = 1
+    CAUSAL_MASK = 2
+    RANDOM_SPANS = 3     # New: explicit ratio control, non-overlapping spans
+
+
+def _generate_random_spans_mask(
+    real_sequence_length: int,
+    sequence_length: int,
+    device: torch.device,
+    max_spans: int = 4,
+    mask_ratio_range: Tuple[float, float] = (0.2, 1.0),
+    span_count_weights: Optional[List[float]] = None,
+) -> torch.Tensor:
+    """
+    Generate a mask with 1-N non-overlapping contiguous spans covering a target ratio.
+
+    Samples a target mask ratio from U[mask_ratio_range], then allocates that budget
+    across a variable number of spans with Dirichlet-like random allocation.
+
+    Args:
+        real_sequence_length: Number of valid (non-padding) tokens.
+        sequence_length: Total sequence length including padding.
+        device: Torch device.
+        max_spans: Maximum number of masked spans.
+        mask_ratio_range: (lo, hi) range for uniform sampling of target mask ratio.
+        span_count_weights: Weights for sampling number of spans [1..max_spans].
+            Defaults to geometric decay [4, 2, 1, 1] (biased toward fewer spans).
+    """
+    item_mask = torch.ones((1, 1, sequence_length), device=device, dtype=torch.float32)
+
+    if real_sequence_length == 0:
+        return item_mask
+
+    # 1. Sample target mask ratio
+    lo, hi = mask_ratio_range
+    target_ratio = random.uniform(lo, hi)
+    target_masked_tokens = max(1, int(target_ratio * real_sequence_length))
+
+    # 2. Sample number of spans (biased toward fewer spans)
+    if span_count_weights is None:
+        span_count_weights = [4, 2, 1, 1]
+
+    # Truncate/pad weights to max_spans
+    weights = span_count_weights[:max_spans]
+    while len(weights) < max_spans:
+        weights.append(weights[-1])
+
+    num_spans = random.choices(range(1, max_spans + 1), weights=weights, k=1)[0]
+    num_spans = min(num_spans, real_sequence_length)
+
+    # 3. Allocate target tokens across spans (Dirichlet-like)
+    raw_weights = [random.random() for _ in range(num_spans)]
+    total_weight = sum(raw_weights)
+    span_lengths = [max(1, int(w / total_weight * target_masked_tokens)) for w in raw_weights]
+
+    # Adjust to hit exact target
+    diff = target_masked_tokens - sum(span_lengths)
+    if diff > 0:
+        span_lengths[0] += diff
+    elif diff < 0:
+        for i in range(len(span_lengths)):
+            reduce = min(-diff, span_lengths[i] - 1)
+            span_lengths[i] -= reduce
+            diff += reduce
+            if diff >= 0:
+                break
+
+    # 4. Place spans without overlap (largest first for better packing)
+    span_lengths.sort(reverse=True)
+    placed_spans = []
+
+    for length in span_lengths:
+        length = min(length, real_sequence_length)
+        if length <= 0:
+            continue
+
+        max_start = real_sequence_length - length
+        if max_start < 0:
+            continue
+
+        # Try random placement avoiding overlaps
+        placed = False
+        for _ in range(50):
+            start = random.randint(0, max_start)
+            end = start + length
+            if not any(s < end and start < e for s, e in placed_spans):
+                placed_spans.append((start, end))
+                item_mask[:, :, start:end] = 0
+                placed = True
+                break
+
+        # Fallback: scan for first valid position
+        if not placed:
+            for start in range(max_start + 1):
+                end = start + length
+                if not any(s < end and start < e for s, e in placed_spans):
+                    placed_spans.append((start, end))
+                    item_mask[:, :, start:end] = 0
+                    break
+
+    return item_mask
+
+
+def random_inpaint_mask(
+    sequence: torch.Tensor,
+    padding_masks: torch.Tensor,
+    max_mask_segments: int = 10,
+    mask_type_probabilities: Optional[List[float]] = None,
+    mask_padding: bool = False,
+    force_mask_type: Optional[MaskType] = None,
+    # RANDOM_SPANS parameters
+    max_spans: int = 4,
+    mask_ratio_range: Tuple[float, float] = (0.2, 1.0),
+    span_count_weights: Optional[List[float]] = None,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Generates random inpainting masks for a batch of latent audio sequences.
+    The output inpainting mask has 0 where data should be inpainted, and 1 where data is provided.
+
+    Supports two segment masking modes:
+        - RANDOM_SEGMENTS (legacy): Uncontrolled mask ratio, overlapping segments allowed.
+        - RANDOM_SPANS (new): Explicit target mask ratio sampled from mask_ratio_range,
+          non-overlapping contiguous spans with controllable count distribution.
+
+    Args:
+        sequence: The input sequence tensor of shape (b, c, sequence_length).
+        padding_masks: A tensor of shape (b, sequence_length)
+                       where 1 indicates real data latents and 0 indicates latents encoding silence padding.
+        max_mask_segments: (RANDOM_SEGMENTS only) Maximum number of segments.
+        mask_type_probabilities: Probabilities for each mask type. Length must match the
+                                 number of MaskType values used:
+                                 - 3 elements: [P(RANDOM_SEGMENTS), P(FULL_MASK), P(CAUSAL_MASK)]
+                                 - 4 elements: [P(RANDOM_SEGMENTS), P(FULL_MASK), P(CAUSAL_MASK), P(RANDOM_SPANS)]
+                                 If None, defaults to [0.1, 0.8, 0.1] (legacy 3-type behavior).
+        mask_padding: If True, zero out padding region in the inpaint mask for
+                      RANDOM_SEGMENTS/RANDOM_SPANS/CAUSAL_MASK. Use when training with
+                      attention masking for variable-length sequences.
+        max_spans: (RANDOM_SPANS only) Maximum number of masked spans.
+        mask_ratio_range: (RANDOM_SPANS only) (lo, hi) for uniform sampling of target mask ratio.
+        span_count_weights: (RANDOM_SPANS only) Weights for sampling number of spans [1..max_spans].
+
+    Returns:
+        A tuple containing:
+            - masked_sequence: The sequence with masks applied (original sequence where mask is 1,
+                               and usually 0 or a placeholder where mask is 0).
+            - inpaint_mask: The generated inpainting mask tensor (0 for inpaint, 1 for keep).
+    """
+    b, _, sequence_length = sequence.size()
+
+    # Skip probability validation when forcing a specific mask type
+    if force_mask_type is not None:
+        mask_types_to_sample = None
+    else:
+        if mask_type_probabilities is None:
+            mask_type_probabilities = [0.1, 0.8, 0.1]
+
+        num_probs = len(mask_type_probabilities)
+
+        # Determine which mask types are active based on probability list length
+        if num_probs == 3:
+            # Legacy: [RANDOM_SEGMENTS, FULL_MASK, CAUSAL_MASK]
+            mask_types_to_sample = [MaskType.RANDOM_SEGMENTS.value, MaskType.FULL_MASK.value, MaskType.CAUSAL_MASK.value]
+        elif num_probs == 4:
+            # Extended: [RANDOM_SEGMENTS, FULL_MASK, CAUSAL_MASK, RANDOM_SPANS]
+            mask_types_to_sample = [MaskType.RANDOM_SEGMENTS.value, MaskType.FULL_MASK.value, MaskType.CAUSAL_MASK.value, MaskType.RANDOM_SPANS.value]
+        else:
+            raise ValueError(
+                f"mask_type_probabilities must have 3 or 4 elements. Got {num_probs}."
+            )
+
+        if not torch.isclose(torch.tensor(sum(mask_type_probabilities)), torch.tensor(1.0)):
+            raise ValueError(
+                f"mask_type_probabilities must sum to 1.0. "
+                f"Current sum: {sum(mask_type_probabilities)}"
+            )
+
+    output_masks_list = []
+
+    for i in range(b):
+        padding_mask_single_item = padding_masks[i]
+        real_sequence_length = (padding_mask_single_item == 1).sum().item()
+
+        item_mask = torch.ones((1, 1, sequence_length), device=sequence.device, dtype=torch.float32)
+
+        if force_mask_type is not None:
+            current_mask_type = force_mask_type
+        else:
+            chosen_mask_value = random.choices(mask_types_to_sample, weights=mask_type_probabilities, k=1)[0]
+            current_mask_type = MaskType(chosen_mask_value)
+
+        if current_mask_type == MaskType.FULL_MASK:
+            item_mask = torch.zeros((1, 1, sequence_length), device=sequence.device, dtype=torch.float32)
+        elif real_sequence_length == 0:
+            pass  # item_mask remains all ones
+        else:
+            if current_mask_type == MaskType.RANDOM_SEGMENTS:
+                # Legacy behavior: uncontrolled ratio, overlapping segments
+                num_segments = random.randint(1, max_mask_segments)
+                max_len_per_segment_calc = max(1, real_sequence_length // num_segments)
+
+                for _ in range(num_segments):
+                    segment_length = random.randint(1, max_len_per_segment_calc)
+
+                    if real_sequence_length - segment_length < 0:
+                        continue
+                    mask_start = random.randint(0, real_sequence_length - segment_length)
+                    item_mask[:, :, mask_start : mask_start + segment_length] = 0
+
+            elif current_mask_type == MaskType.RANDOM_SPANS:
+                item_mask = _generate_random_spans_mask(
+                    real_sequence_length=real_sequence_length,
+                    sequence_length=sequence_length,
+                    device=sequence.device,
+                    max_spans=max_spans,
+                    mask_ratio_range=mask_ratio_range,
+                    span_count_weights=span_count_weights,
+                )
+
+            elif current_mask_type == MaskType.CAUSAL_MASK:
+                unmasked_prefix_len = random.randint(0, real_sequence_length)
+
+                if unmasked_prefix_len < real_sequence_length:
+                    item_mask[:, :, unmasked_prefix_len:real_sequence_length] = 0
+
+            # When using attention masking for variable-length training, zero out padding
+            # so the model doesn't see it as provided context. Without attention masking,
+            # padding is handled via mask_loss_weight and should remain as provided input.
+            if mask_padding:
+                item_mask[:, :, real_sequence_length:] = 0
+
+        output_masks_list.append(item_mask)
+
+    final_inpaint_mask = torch.cat(output_masks_list, dim=0).to(sequence.device)
+    masked_sequence = sequence * final_inpaint_mask
+    return masked_sequence, final_inpaint_mask

stable_audio_3/models/lora/__init__.py

@@ -0,0 +1,28 @@
+from .model import LoRAParametrization, add_lora, default_lora_config, merge_lora, remove_lora, remove_lora_by_index, set_lora_strength
+from .utils import (
+    apply_to_lora,
+    convert_lora_ckpt_to_safetensors,
+    disable_lora,
+    enable_lora,
+    filter_lora_layers,
+    get_lora_count,
+    has_lora,
+    get_bias_params,
+    get_lora_params,
+    get_lora_layers,
+    get_lora_state_dict,
+    load_lora_checkpoint,
+    load_multiple_lora,
+    name_is_lora,
+    remap_lora_state_dict,
+    save_lora_safetensors,
+    select_lora,
+    tie_weights,
+    untie_weights,
+    infer_global_rank,
+    cast_base_to_precision,
+    detect_dora_variant,
+    prepare_dora_state_dict,
+    resolve_adapter_type
+)
+from .loader import load_and_apply_loras

stable_audio_3/models/lora/loader.py

@@ -0,0 +1,112 @@
+"""Helpers for loading pre-trained LoRA checkpoints onto a live model."""
+import os
+from functools import partial
+
+import torch
+
+from .model import LoRAParametrization, add_lora
+from .utils import (
+    get_lora_layers,
+    infer_global_rank,
+    load_lora_checkpoint,
+    prepare_dora_state_dict,
+    remap_lora_state_dict,
+    resolve_adapter_type,
+)
+from ...verbose import vprint
+
+
+def load_and_apply_loras(model, lora_ckpt_paths, model_type, svd_bases_path=None):
+    """Load LoRA checkpoints from disk and attach them to `model`.
+
+    - Uses a two-pass approach: first resolves each LoRA's adapter type, then
+      loads SVD bases (only if any LoRA is -XS), then applies each LoRA.
+    - Handles legacy "dora" adapter_type via `resolve_adapter_type`.
+    - Only passes svd_bases to -XS LoRAs to avoid spurious "no -XS layers"
+      messages for non-XS LoRAs.
+    - Sets `model.use_lora = True` and `model.lora_names = [...]` as a
+      convenience for downstream UI code.
+
+    Args:
+        model: The loaded model. For `diffusion_cond` / `diffusion_cond_inpaint`
+            types, LoRAs are applied to `model.model` and `model.conditioner`;
+            otherwise they are applied to `model` directly.
+        lora_ckpt_paths: List of paths to `.ckpt` or `.safetensors` files.
+        model_type: String from the model config, e.g. "diffusion_cond".
+        svd_bases_path: Optional path to a precomputed SVD bases .pt file
+            (used only by -XS adapter types).
+
+    Returns:
+        List of display names (file stems) in the same order as the input paths.
+    """
+    if not lora_ckpt_paths:
+        model.use_lora = False
+        model.lora_names = []
+        return []
+
+    is_cond = model_type in ("diffusion_cond", "diffusion_cond_inpaint")
+
+    # Pass 1: load each checkpoint and resolve adapter type.
+    lora_entries = []  # (path, state_dict, config_dict, adapter_type)
+    for i, lora_path in enumerate(lora_ckpt_paths):
+        print(f"Loading LoRA {i} from {lora_path}")
+        state_dict, config_dict = load_lora_checkpoint(lora_path)
+        adapter_type_raw = config_dict.get("adapter_type", "lora")
+        adapter_type = resolve_adapter_type(adapter_type_raw, state_dict)
+        if adapter_type != adapter_type_raw:
+            vprint(f"Resolved legacy '{adapter_type_raw}' -> {adapter_type}")
+        lora_entries.append((lora_path, state_dict, config_dict, adapter_type))
+
+    # Load SVD bases only if at least one LoRA is -XS.
+    model_svd_bases = None
+    cond_svd_bases = None
+    any_xs = any(t.endswith("-xs") for _, _, _, t in lora_entries)
+    if any_xs:
+        if svd_bases_path is not None:
+            vprint(f"Loading SVD bases from {svd_bases_path}")
+            svd_bases_all = torch.load(svd_bases_path, map_location="cpu", weights_only=True)
+            model_svd_bases = {k[len("model."):]: v for k, v in svd_bases_all.items() if k.startswith("model.")}
+            cond_svd_bases = {k[len("conditioner."):]: v for k, v in svd_bases_all.items() if k.startswith("conditioner.")}
+            vprint(f"  model bases: {len(model_svd_bases)}, conditioner bases: {len(cond_svd_bases)}")
+        else:
+            print("WARNING: -XS adapter type present without svd_bases_path -- SVD will be computed on current device")
+
+    # Pass 2: apply each LoRA.
+    lora_names = []
+    for i, (lora_path, state_dict, config_dict, adapter_type) in enumerate(lora_entries):
+        rank = config_dict.get("rank", infer_global_rank(state_dict))
+        alpha = config_dict.get("alpha", rank)
+        include = config_dict.get("include", None)
+        exclude = config_dict.get("exclude", None)
+        is_xs = adapter_type.endswith("-xs")
+
+        lora_config = {
+            torch.nn.Linear: {
+                "weight": partial(LoRAParametrization.from_linear, rank=rank, lora_alpha=alpha, adapter_type=adapter_type, lora_index=i),
+            },
+            torch.nn.Conv1d: {
+                "weight": partial(LoRAParametrization.from_conv1d, rank=rank, lora_alpha=alpha, adapter_type=adapter_type, lora_index=i),
+            },
+        }
+        if is_cond:
+            add_lora(model.model, lora_config, include=include, exclude=exclude,
+                     svd_bases=model_svd_bases if is_xs else None)
+            add_lora(model.conditioner, lora_config, include=include, exclude=exclude,
+                     svd_bases=cond_svd_bases if is_xs else None)
+        else:
+            add_lora(model, lora_config, include=include, exclude=exclude,
+                     svd_bases=model_svd_bases if is_xs else None)
+
+        prepare_dora_state_dict(state_dict)
+        remapped_sd = remap_lora_state_dict(state_dict, i)
+        if is_cond:
+            model.model.load_state_dict(remapped_sd, strict=False)
+            model.conditioner.load_state_dict(remapped_sd, strict=False)
+        else:
+            model.load_state_dict(remapped_sd, strict=False)
+        lora_names.append(os.path.splitext(os.path.basename(lora_path))[0])
+
+    vprint("lora layers:", len(get_lora_layers(model)))
+    model.use_lora = True
+    model.lora_names = lora_names
+    return lora_names

stable_audio_3/models/lora/model.py

@@ -0,0 +1,471 @@
+"""
+References:
+1) the official LoRA implementation released by Microsoft:
+https://github.com/microsoft/LoRA/blob/main/loralib/layers.py
+"""
+
+import math
+import re
+from functools import partial
+from itertools import product
+
+import torch
+import torch.nn.utils.parametrize as parametrize
+from torch import nn
+
+from ...verbose import vprint
+
+
+def _canonicalize_svd_signs(U, Vh):
+    """Enforce deterministic sign convention: largest-magnitude element of each U column is positive."""
+    max_abs_idx = U.abs().argmax(dim=0)
+    signs = U[max_abs_idx, torch.arange(U.shape[1])].sign()
+    signs[signs == 0] = 1
+    return U * signs.unsqueeze(0), Vh * signs.unsqueeze(1)
+
+
+class LoRAParametrization(nn.Module):
+    def __init__(self, fan_in, fan_out, fan_in_fan_out=False, rank=4, lora_dropout_p=0.0, lora_alpha=1,
+                adapter_type="lora", W0=None, lora_index=0, svd_bases=None
+                ):
+        super().__init__()
+        dtype = torch.get_default_dtype()
+        device = W0.device if W0 is not None else None
+        # if weight is stored as (fan_out, fan_in), the memory layout of A & B follows (W + BA)x
+        # otherwise, it's x(W + AB). This allows us to tie the weights between linear layers and embeddings
+        self.swap = (lambda x: (x[1], x[0])) if fan_in_fan_out else (lambda x: x)
+        self.lora_alpha, self.rank = lora_alpha, rank
+        self.scaling = lora_alpha / rank
+        self.lora_dropout = nn.Dropout(p=lora_dropout_p) if lora_dropout_p > 0 else lambda x: x
+        self.dropout_fn = self._dropout if lora_dropout_p > 0 else lambda x: x
+        self.register_buffer("lora_dropout_mask", torch.ones(self.swap((1, fan_in)), dtype=dtype, device=device))
+        self.register_buffer("lora_strength", torch.tensor(1.0, dtype=dtype, device=device), persistent=False)
+        self.adapter_type = adapter_type
+        self.lora_index = lora_index
+
+        if self.adapter_type == "lora":
+            self.lora_A = nn.Parameter(torch.zeros(self.swap((rank, fan_in)), device=device))
+            self.lora_B = nn.Parameter(torch.zeros(self.swap((fan_out, rank)), device=device))
+            nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
+            self.forward_fn = self.lora_forward
+            self._adapter_forward_fn = self.forward_fn
+
+        elif self.adapter_type in ("dora-rows", "dora"):
+            # DoRA-rows: magnitude per row (per-output-neuron, dim=1) - paper-correct variant
+            # "dora" maps here as the default for backward compat
+            self.lora_A = nn.Parameter(torch.zeros(self.swap((rank, fan_in)), device=device))
+            self.lora_B = nn.Parameter(torch.zeros(self.swap((fan_out, rank)), device=device))
+            nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
+            self._norm_dim = 1
+            if W0 is not None:
+                W0_2d = W0.view(W0.shape[0], -1)
+                self.magnitude = nn.Parameter(W0_2d.norm(dim=self._norm_dim).to(dtype))
+            else:
+                self.magnitude = nn.Parameter(torch.ones(fan_out, dtype=dtype, device=device))
+            self.forward_fn = self.dora_forward
+            self._adapter_forward_fn = self.forward_fn
+
+        elif self.adapter_type == "dora-cols":
+            # DoRA-cols: magnitude per column (per-input-feature, dim=0) - legacy variant
+            self.lora_A = nn.Parameter(torch.zeros(self.swap((rank, fan_in)), device=device))
+            self.lora_B = nn.Parameter(torch.zeros(self.swap((fan_out, rank)), device=device))
+            nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
+            self._norm_dim = 0
+            if W0 is not None:
+                W0_2d = W0.view(W0.shape[0], -1)
+                self.magnitude = nn.Parameter(W0_2d.norm(dim=self._norm_dim).to(dtype))
+            else:
+                self.magnitude = nn.Parameter(torch.ones(fan_in, dtype=dtype, device=device))
+            self.forward_fn = self.dora_forward
+            self._adapter_forward_fn = self.forward_fn
+
+        elif self.adapter_type == "bora":
+            # BoRA: bi-dimensional DoRA - magnitude on both row and column axes
+            self.lora_A = nn.Parameter(torch.zeros(self.swap((rank, fan_in)), device=device))
+            self.lora_B = nn.Parameter(torch.zeros(self.swap((fan_out, rank)), device=device))
+            nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
+            if W0 is not None:
+                W0_2d = W0.view(W0.shape[0], -1)
+                self.magnitude_r = nn.Parameter(W0_2d.norm(dim=1).to(dtype))
+                self.magnitude_c = nn.Parameter(W0_2d.norm(dim=0).to(dtype))
+            else:
+                self.magnitude_r = nn.Parameter(torch.ones(fan_out, dtype=dtype, device=device))
+                self.magnitude_c = nn.Parameter(torch.ones(fan_in, dtype=dtype, device=device))
+            self.forward_fn = self.bora_forward
+            self._adapter_forward_fn = self.forward_fn
+
+        elif self.adapter_type == "lora-xs":
+            # frozen bases U:(fan_out,r), V:(fan_in,r); tiny trainable core M:(r,r)
+            self.register_buffer("U", torch.empty(self.swap((fan_out, self.rank)), device=device), persistent=True)
+            self.register_buffer("V", torch.empty(self.swap((fan_in,  self.rank)), device=device), persistent=True)
+            self.M_xs = nn.Parameter(torch.zeros(self.rank, self.rank, dtype=dtype, device=device))
+            if svd_bases is not None:
+                self.U.copy_(svd_bases["U"][:, :self.rank].to(dtype))
+                self.V.copy_(svd_bases["V"][:, :self.rank].to(dtype))
+            elif W0 is not None:
+                W0_2d = W0.view(W0.shape[0], -1).cpu().float()
+                U_full, S, Vh_full = torch.linalg.svd(W0_2d, full_matrices=False)
+                U_full, Vh_full = _canonicalize_svd_signs(U_full, Vh_full)
+                self.U.copy_(U_full[:, :self.rank].to(dtype))
+                self.V.copy_(Vh_full[:self.rank, :].T.to(dtype))
+            self.forward_fn = self.lora_xs_forward
+            self._adapter_forward_fn = self.forward_fn
+
+        elif self.adapter_type == "dora-rows-xs":
+            # DoRA-rows + LoRA-XS: frozen SVD bases + trainable M_xs core + row magnitude
+            self.register_buffer("U", torch.empty(self.swap((fan_out, self.rank)), device=device), persistent=True)
+            self.register_buffer("V", torch.empty(self.swap((fan_in,  self.rank)), device=device), persistent=True)
+            self.M_xs = nn.Parameter(torch.zeros(self.rank, self.rank, dtype=dtype, device=device))
+            self._norm_dim = 1
+            if svd_bases is not None:
+                self.U.copy_(svd_bases["U"][:, :self.rank].to(dtype))
+                self.V.copy_(svd_bases["V"][:, :self.rank].to(dtype))
+            if W0 is not None:
+                W0_2d = W0.view(W0.shape[0], -1)
+                if svd_bases is None:
+                    U_full, S, Vh_full = torch.linalg.svd(W0_2d.cpu().float(), full_matrices=False)
+                    U_full, Vh_full = _canonicalize_svd_signs(U_full, Vh_full)
+                    self.U.copy_(U_full[:, :self.rank].to(dtype))
+                    self.V.copy_(Vh_full[:self.rank, :].T.to(dtype))
+                self.magnitude = nn.Parameter(W0_2d.norm(dim=self._norm_dim).to(dtype))
+            else:
+                self.magnitude = nn.Parameter(torch.ones(fan_out, dtype=dtype, device=device))
+            self.forward_fn = self.dora_xs_forward
+            self._adapter_forward_fn = self.forward_fn
+
+        elif self.adapter_type == "dora-cols-xs":
+            # DoRA-cols + LoRA-XS: frozen SVD bases + trainable M_xs core + column magnitude
+            self.register_buffer("U", torch.empty(self.swap((fan_out, self.rank)), device=device), persistent=True)
+            self.register_buffer("V", torch.empty(self.swap((fan_in,  self.rank)), device=device), persistent=True)
+            self.M_xs = nn.Parameter(torch.zeros(self.rank, self.rank, dtype=dtype, device=device))
+            self._norm_dim = 0
+            if svd_bases is not None:
+                self.U.copy_(svd_bases["U"][:, :self.rank].to(dtype))
+                self.V.copy_(svd_bases["V"][:, :self.rank].to(dtype))
+            if W0 is not None:
+                W0_2d = W0.view(W0.shape[0], -1)
+                if svd_bases is None:
+                    U_full, S, Vh_full = torch.linalg.svd(W0_2d.cpu().float(), full_matrices=False)
+                    U_full, Vh_full = _canonicalize_svd_signs(U_full, Vh_full)
+                    self.U.copy_(U_full[:, :self.rank].to(dtype))
+                    self.V.copy_(Vh_full[:self.rank, :].T.to(dtype))
+                self.magnitude = nn.Parameter(W0_2d.norm(dim=self._norm_dim).to(dtype))
+            else:
+                self.magnitude = nn.Parameter(torch.ones(fan_in, dtype=dtype, device=device))
+            self.forward_fn = self.dora_xs_forward
+            self._adapter_forward_fn = self.forward_fn
+
+        elif self.adapter_type == "bora-xs":
+            # BoRA + LoRA-XS: frozen SVD bases + trainable M_xs core + row & column magnitudes
+            self.register_buffer("U", torch.empty(self.swap((fan_out, self.rank)), device=device), persistent=True)
+            self.register_buffer("V", torch.empty(self.swap((fan_in,  self.rank)), device=device), persistent=True)
+            self.M_xs = nn.Parameter(torch.zeros(self.rank, self.rank, dtype=dtype, device=device))
+            if svd_bases is not None:
+                self.U.copy_(svd_bases["U"][:, :self.rank].to(dtype))
+                self.V.copy_(svd_bases["V"][:, :self.rank].to(dtype))
+            if W0 is not None:
+                W0_2d = W0.view(W0.shape[0], -1)
+                if svd_bases is None:
+                    U_full, S, Vh_full = torch.linalg.svd(W0_2d.cpu().float(), full_matrices=False)
+                    U_full, Vh_full = _canonicalize_svd_signs(U_full, Vh_full)
+                    self.U.copy_(U_full[:, :self.rank].to(dtype))
+                    self.V.copy_(Vh_full[:self.rank, :].T.to(dtype))
+                self.magnitude_r = nn.Parameter(W0_2d.norm(dim=1).to(dtype))
+                self.magnitude_c = nn.Parameter(W0_2d.norm(dim=0).to(dtype))
+            else:
+                self.magnitude_r = nn.Parameter(torch.ones(fan_out, dtype=dtype, device=device))
+                self.magnitude_c = nn.Parameter(torch.ones(fan_in, dtype=dtype, device=device))
+            self.forward_fn = self.bora_xs_forward
+            self._adapter_forward_fn = self.forward_fn
+
+
+    def lora_forward(self, W):
+        delta = torch.matmul(*self.swap((self.lora_B, self.dropout_fn(self.lora_A)))).view(W.shape)
+        delta = self.scaling * self.lora_strength * delta
+        return (W + delta.to(W.dtype))
+    
+    def dora_forward(self, W):
+        if self.lora_strength == 0:
+            return W
+        # work in 2D so norm/magnitude ops don't broadcast badly for Conv1d/Conv2d
+        # Upcast W to LoRA param dtype for stable norm computation, cast result back.
+        orig_shape = W.shape
+        W_2d = W.view(W.shape[0], -1).to(self.lora_A.dtype)
+        # low-rank update on the *direction*
+        delta = torch.matmul(*self.swap((self.lora_B, self.dropout_fn(self.lora_A))))
+        V = W_2d + self.scaling * self.lora_strength * delta
+        # normalize along _norm_dim, then scale by per-element magnitude
+        V_hat = V / (V.norm(dim=self._norm_dim, keepdim=True) + 1e-12)
+        return (V_hat * self.magnitude.unsqueeze(self._norm_dim)).view(orig_shape).to(W.dtype)
+
+    def bora_forward(self, W):
+        if self.lora_strength == 0:
+            return W
+        orig_shape = W.shape
+        W_2d = W.view(W.shape[0], -1).to(self.lora_A.dtype)
+        delta = torch.matmul(*self.swap((self.lora_B, self.dropout_fn(self.lora_A))))
+        V = W_2d + self.scaling * self.lora_strength * delta
+        # Row-normalize, then scale by row magnitudes
+        V_r = V / (V.norm(dim=1, keepdim=True) + 1e-12)
+        intermediate = self.magnitude_r.unsqueeze(1) * V_r
+        # Column-normalize, then scale by column magnitudes
+        H_c = intermediate / (intermediate.norm(dim=0, keepdim=True) + 1e-12)
+        return (H_c * self.magnitude_c.unsqueeze(0)).view(orig_shape).to(W.dtype)
+
+    def lora_xs_forward(self, W):
+        delta = self.U @ self.M_xs.to(self.U.dtype) @ self.V.T
+        delta = delta.view(W.shape)
+        return (W + (self.scaling * self.lora_strength * delta).to(W.dtype))
+
+    def dora_xs_forward(self, W):
+        if self.lora_strength == 0:
+            return W
+        orig_shape = W.shape
+        W_2d = W.view(W.shape[0], -1).to(self.U.dtype)
+        delta = self.U @ self.M_xs.to(self.U.dtype) @ self.V.T
+        V = W_2d + self.scaling * self.lora_strength * delta
+        V_hat = V / (V.norm(dim=self._norm_dim, keepdim=True) + 1e-12)
+        return (V_hat * self.magnitude.unsqueeze(self._norm_dim)).view(orig_shape).to(W.dtype)
+
+    def bora_xs_forward(self, W):
+        if self.lora_strength == 0:
+            return W
+        orig_shape = W.shape
+        W_2d = W.view(W.shape[0], -1).to(self.U.dtype)
+        delta = self.U @ self.M_xs.to(self.U.dtype) @ self.V.T
+        V = W_2d + self.scaling * self.lora_strength * delta
+        V_r = V / (V.norm(dim=1, keepdim=True) + 1e-12)
+        intermediate = self.magnitude_r.unsqueeze(1) * V_r
+        H_c = intermediate / (intermediate.norm(dim=0, keepdim=True) + 1e-12)
+        return (H_c * self.magnitude_c.unsqueeze(0)).view(orig_shape).to(W.dtype)
+
+    def forward(self, X):
+        return self.forward_fn(X)
+
+    def _dropout(self, A):
+        # to mimic the original implementation: A @ dropout(x), we do (A * dropout(ones)) @ x
+        return A * self.lora_dropout(self.lora_dropout_mask)
+
+    def disable_lora(self):
+        self.forward_fn = lambda x: x
+
+    def enable_lora(self):
+        self.forward_fn = self._adapter_forward_fn
+
+    @classmethod
+    def _get_original_weight(cls, layer, attr_name="weight"):
+        """Get the true original weight, even when the layer is already parametrized."""
+        if hasattr(layer, 'parametrizations') and hasattr(layer.parametrizations, attr_name):
+            return getattr(layer.parametrizations, attr_name).original.detach()
+        return getattr(layer, attr_name).detach()
+
+    @classmethod
+    def from_linear(cls, layer, rank=4, lora_dropout_p=0.0, lora_alpha=1, svd_bases=None, **kwargs):
+        fan_out, fan_in = layer.weight.shape
+        W0 = cls._get_original_weight(layer)
+        return cls(
+            fan_in, fan_out, fan_in_fan_out=False, rank=rank, lora_dropout_p=lora_dropout_p, lora_alpha=lora_alpha,
+            W0=W0, svd_bases=svd_bases, **kwargs
+        )
+
+    @classmethod
+    def from_conv2d(cls, layer, rank=4, lora_dropout_p=0.0, lora_alpha=1, svd_bases=None, **kwargs):
+        fan_out, fan_in = layer.weight.view(layer.weight.shape[0], -1).shape
+        W0 = cls._get_original_weight(layer).view(layer.weight.shape[0], -1)
+        return cls(
+            fan_in, fan_out, fan_in_fan_out=False, rank=rank, lora_dropout_p=lora_dropout_p, lora_alpha=lora_alpha,
+            W0=W0, svd_bases=svd_bases, **kwargs
+        )
+
+    @classmethod
+    def from_conv1d(cls, layer, rank=4, lora_dropout_p=0.0, lora_alpha=1, svd_bases=None, **kwargs):
+        fan_out, fan_in = layer.weight.view(layer.weight.shape[0], -1).shape
+        W0 = cls._get_original_weight(layer).view(layer.weight.shape[0], -1)
+        return cls(
+            fan_in, fan_out, fan_in_fan_out=False, rank=rank, lora_dropout_p=lora_dropout_p, lora_alpha=lora_alpha,
+            W0=W0, svd_bases=svd_bases, **kwargs
+        )
+
+    @classmethod
+    def from_embedding(cls, layer, rank=4, lora_dropout_p=0.0, lora_alpha=1, svd_bases=None, **kwargs):
+        fan_in, fan_out = layer.weight.shape
+        W0 = cls._get_original_weight(layer).view(layer.weight.shape[0], -1)
+        return cls(
+            fan_in, fan_out, fan_in_fan_out=True, rank=rank, lora_dropout_p=lora_dropout_p, lora_alpha=lora_alpha,
+            W0=W0, svd_bases=svd_bases, **kwargs
+        )
+
+
+default_lora_config = {  # specify which layers to add lora to, by default only add to linear layers
+    nn.Linear: {
+        "weight": partial(LoRAParametrization.from_linear, rank=4),
+    },
+}
+
+
+# --- layer name filtering helpers ---
+
+def _expand(s):
+    """Expand bracket notation in layer filter strings, e.g. 'layers[0-5]' -> ['layers0', ..., 'layers5']."""
+    t = re.split(r'\[(\d+)-(\d+)\]', s)
+    if len(t) == 1:
+        return [s]
+    lits, starts, ends = t[0::3], t[1::3], t[2::3]
+    pools = []
+    for a, b in zip(starts, ends):
+        ai, bi = int(a), int(b)
+        step = 1 if bi >= ai else -1
+        pools.append([str(n) for n in range(ai, bi + step, step)])
+    out = []
+    for combo in product(*pools):
+        pieces = []
+        for i, lit in enumerate(lits):
+            pieces.append(lit)
+            if i < len(combo):
+                pieces.append(combo[i])
+        out.append(''.join(pieces))
+    return out
+
+
+def _matches_any(name, patterns):
+    """Check if name contains any pattern substring (with bracket expansion)."""
+    if not patterns:
+        return False
+    for pattern in patterns:
+        for expanded in _expand(pattern.strip()):
+            if expanded in name:
+                return True
+    return False
+
+
+# --- core LoRA application ---
+
+def _match_layer_type(layer, lora_config):
+    """Find the matching lora_config key for a layer, using isinstance to handle ParametrizedLinear etc."""
+    for layer_type in lora_config:
+        if isinstance(layer, layer_type):
+            return layer_type
+    return None
+
+
+def apply_lora(layer, register=True, merge=False, lora_config=default_lora_config):
+    """add lora parametrization to a layer, designed to be used with model.apply"""
+    if register:
+        matched_type = _match_layer_type(layer, lora_config)
+        if matched_type is not None:
+            for attr_name, parametrization in lora_config[matched_type].items():
+                parametrize.register_parametrization(layer, attr_name, parametrization(layer), unsafe=True)
+    else:  # this will remove all parametrizations, use with caution
+        if hasattr(layer, "parametrizations"):
+            for attr_name in layer.parametrizations.keys():
+                parametrize.remove_parametrizations(layer, attr_name, leave_parametrized=merge)
+
+
+def add_lora(model, lora_config=default_lora_config, include=None, exclude=None, svd_bases=None):
+    """Add LoRA parametrization to layers in a model.
+
+    Args:
+        model: The model to add LoRA to.
+        lora_config: Dict mapping nn.Module types to parametrization configs.
+        include: Optional list of substrings. If provided, only modules whose
+                 name contains at least one pattern get LoRA. Supports bracket
+                 expansion (e.g. "layers[0-11]").
+        exclude: Optional list of substrings. Modules matching any pattern are
+                 skipped even if they match an include pattern.
+        svd_bases: Optional dict of pre-computed SVD bases keyed by
+                   "layer_name.weight". When provided, -XS adapter types use
+                   these instead of computing SVD.
+    """
+    if svd_bases is not None or include is not None or exclude is not None:
+        # Named-module path: iterate by name so we can look up SVD bases per layer
+        n_cached = 0
+        n_recomputed = 0
+        for name, module in model.named_modules():
+            matched_type = _match_layer_type(module, lora_config)
+            if matched_type is None:
+                continue
+
+            if include is not None and not _matches_any(name, include):
+                continue
+
+            if exclude is not None and _matches_any(name, exclude):
+                continue
+
+            for attr_name, parametrization_fn in lora_config[matched_type].items():
+                layer_bases = None
+                if svd_bases is not None:
+                    bases_key = f"{name}.{attr_name}"
+                    layer_bases = svd_bases.get(bases_key)
+                param = parametrization_fn(module, svd_bases=layer_bases)
+                parametrize.register_parametrization(module, attr_name, param, unsafe=True)
+                if hasattr(param, 'U'):  # -XS layer that needs SVD bases
+                    if layer_bases is not None:
+                        n_cached += 1
+                    else:
+                        n_recomputed += 1
+        if n_cached > 0:
+            vprint(f"SVD bases: {n_cached} layers loaded from cache")
+        if n_recomputed > 0:
+            vprint(f"SVD bases: {n_recomputed} layers RECOMPUTED (missing from cache)")
+        if svd_bases is not None and n_cached == 0 and n_recomputed == 0:
+            vprint("SVD bases: provided but no -XS layers found")
+    else:
+        # Fast path: original behavior, no name inspection needed
+        model.apply(partial(apply_lora, lora_config=lora_config))
+
+def add_lora_by_name(model, target_module_names, lora_config=default_lora_config, svd_bases=None):
+    """Add LoRA to specific layers by name. Convenience wrapper around add_lora()."""
+    add_lora(model, lora_config=lora_config, include=target_module_names, svd_bases=svd_bases)
+
+
+def merge_lora(model):
+    """merge lora parametrization to all layers in a model. This will remove all parametrization"""
+    model.apply(partial(apply_lora, register=False, merge=True))
+
+def remove_lora(model):
+    """remove lora parametrization to all layers in a model. This will remove all parametrization"""
+    model.apply(partial(apply_lora, register=False, merge=False))
+
+def remove_lora_by_index(model, lora_index):
+    """Remove all LoRA parametrizations with a specific lora_index.
+
+    For modules where other LoRA entries remain, only the matching entry is deleted.
+    For modules where the removed entry was the only one, the parametrization is
+    removed entirely and the weight reverts to its original value.
+    """
+    for name, mod in model.named_modules():
+        if not hasattr(mod, 'parametrizations'):
+            continue
+        for attr in list(mod.parametrizations.keys()):
+            params = mod.parametrizations[attr]
+            to_remove = [i for i, p in enumerate(params)
+                         if isinstance(p, LoRAParametrization) and p.lora_index == lora_index]
+            if not to_remove:
+                continue
+            if len(params) == len(to_remove):
+                # All entries being removed — remove parametrization entirely
+                parametrize.remove_parametrizations(mod, attr, leave_parametrized=False)
+            else:
+                for i in reversed(to_remove):
+                    del params[i]
+
+def set_lora_strength(model, strength: float, lora_index=None):
+    """Set lora strength. If lora_index is None, sets all LoRAs. If specified, sets only that index."""
+    strength = float(strength)
+    if lora_index is None:
+        for name, buf in model.named_buffers(recurse=True):
+            if name.endswith("lora_strength"):
+                buf.fill_(strength)
+    else:
+        for p in _iter_lora_params(model):
+            if p.lora_index == lora_index:
+                p.lora_strength.fill_(strength)
+
+def _iter_lora_params(model):
+    for _, mod in model.named_modules():
+        plist = getattr(getattr(mod, "parametrizations", None), "weight", None)
+        if plist is None:
+            continue
+        for p in plist:
+            if isinstance(p, LoRAParametrization):
+                yield p

stable_audio_3/models/lora/utils.py

@@ -0,0 +1,420 @@
+import json
+
+import torch
+from safetensors import safe_open
+from safetensors.torch import save_file as _st_save_file
+from torch import nn
+
+from .model import LoRAParametrization, _expand
+from ...verbose import vprint
+
+def apply_to_lora(fn):
+    """apply a function to LoRAParametrization layers, designed to be used with model.apply"""
+
+    def apply_fn(layer):
+        if isinstance(layer, LoRAParametrization):
+            fn(layer)
+
+    return apply_fn
+
+def enable_lora(model, lora_index=None):
+    """Enable LoRA layers. If lora_index is None, enables all. If specified, enables only that index."""
+    def _enable(layer):
+        if isinstance(layer, LoRAParametrization):
+            if lora_index is None or layer.lora_index == lora_index:
+                layer.enable_lora()
+    model.apply(_enable)
+
+def disable_lora(model, lora_index=None):
+    """Disable LoRA layers. If lora_index is None, disables all. If specified, disables only that index."""
+    def _disable(layer):
+        if isinstance(layer, LoRAParametrization):
+            if lora_index is None or layer.lora_index == lora_index:
+                layer.disable_lora()
+    model.apply(_disable)
+
+def get_lora_layers(model):
+    layers = []
+    for name, m in model.named_modules():
+        plist = getattr(getattr(m, "parametrizations", None), "weight", None)
+        if plist is None:
+            continue
+        for p in plist:
+            if isinstance(p, LoRAParametrization):
+                layers.append((name, p))
+    return layers
+
+def has_lora(model):
+    """Return True if the model has at least one LoRAParametrization on any weight."""
+    return len(get_lora_layers(model))>=1
+
+def filter_lora_layers(model, lora_layer_filter, lora_index=None):
+    """Enable/disable LoRA layers by name filter. If lora_index is specified, only affects that index."""
+    lora_layer_filter = (lora_layer_filter or "").strip()
+    # lora_layer_filter: logical OR on comma separated substrings
+    # e.g. ".transformer.layers, .to_global_embed."
+    def is_filtered(layer_name):
+        filters = [x.lower().strip() for x in lora_layer_filter.split(",") if len(x)]
+        for f in filters:
+            for _f in _expand(f):
+                if _f in layer_name:
+                    return True
+        return False
+    for name, p in get_lora_layers(model):
+        if lora_index is not None and p.lora_index != lora_index:
+            continue
+        if is_filtered(name):
+            p.disable_lora()
+        else:
+            p.enable_lora()
+
+
+
+# ------------------- helper function for collecting parameters for training/saving -------------------
+
+
+def name_is_lora(name):
+    return (
+        len(name.split(".")) >= 4
+        and (name.split(".")[-4]) == "parametrizations"
+        and name.split(".")[-1] in ["lora_A", "lora_B", "M_xs", "magnitude", "magnitude_r", "magnitude_c"]
+    )
+
+
+def name_is_bias(name):
+    return name.split(".")[-1] == "bias"
+
+
+def get_params_by_name(model, print_shapes=False, name_filter=None):
+    for n, p in model.named_parameters():
+        if name_filter is None or name_filter(n):
+            if print_shapes:
+                print(n, p.shape)
+            yield p
+
+
+def get_lora_params(model, print_shapes=False):
+    return get_params_by_name(model, print_shapes=print_shapes, name_filter=name_is_lora)
+
+
+def get_bias_params(model, print_shapes=False):
+    return get_params_by_name(model, print_shapes=print_shapes, name_filter=name_is_bias)
+
+
+def get_lora_state_dict(model):
+    return {k: v for k, v in model.state_dict().items() if name_is_lora(k)}
+
+
+def cast_base_to_precision(model, precision):
+    """Cast frozen base weights to lower precision, keeping LoRA params in fp32.
+
+    Args:
+        model: Model with LoRA parametrizations applied.
+        precision: One of "bf16", "bfloat16", "fp16", "float16".
+    """
+    import torch
+
+    dtype_map = {"bf16": torch.bfloat16, "bfloat16": torch.bfloat16,
+                 "fp16": torch.float16, "float16": torch.float16}
+    target_dtype = dtype_map.get(precision)
+    if target_dtype is None:
+        print(f"Unknown base_precision={precision!r}, keeping fp32")
+        return
+    vprint(f"Casting frozen base weights to {target_dtype}")
+    model.to(target_dtype)
+    # Restore LoRA params to fp32 — optimizer needs full precision
+    n_restored = 0
+    for p in get_lora_params(model):
+        p.data = p.data.to(torch.float32)
+        n_restored += 1
+    vprint(f"  Restored {n_restored} LoRA params to fp32")
+
+
+# ------------------- helper function for loading LoRa at inference -------------------
+
+def infer_global_rank(lora_sd: dict) -> int:
+    candidates = []
+    for k, v in lora_sd.items():
+        if k.endswith(".lora_A"):
+            candidates.append(v.shape[0])
+        elif k.endswith(".lora_B"):
+            candidates.append(v.shape[1])
+        elif k.endswith(".M_xs"):
+            candidates.append(v.shape[0])
+    if not candidates:
+        raise ValueError("No LoRA/LoRA-XS tensors found to infer rank")
+    r = candidates[0]
+    # sanity check they all match (most checkpoints use one global r)
+    if any(c != r for c in candidates):
+        print(f"Multiple ranks detected in ckpt: {sorted(set(candidates))}. "
+                         "You’ll need per-layer ranks.")
+    return r
+
+
+def detect_dora_variant(state_dict):
+    """Detect DoRA variant from a state dict.
+
+    Returns "dora-cols" (dim=0, per-input-feature) or "dora-rows" (dim=1,
+    per-output-neuron) based on the shape of the first .magnitude tensor
+    found. Returns None if no magnitude tensors exist or axis is ambiguous.
+    """
+    for key, val in state_dict.items():
+        if not key.endswith(".magnitude"):
+            continue
+        if val.dim() == 2:
+            if val.shape[0] == 1:
+                return "dora-cols"   # old format: (1, fan_in)
+            elif val.shape[1] == 1:
+                return "dora-rows"   # (fan_out, 1)
+        # 1D magnitude — can't determine axis without model context
+        return None
+    return None
+
+
+def prepare_dora_state_dict(state_dict):
+    """Squeeze 2D DoRA magnitude tensors to 1D for loading. Modifies in-place."""
+    for key in list(state_dict.keys()):
+        if key.endswith(".magnitude") and state_dict[key].dim() == 2:
+            state_dict[key] = state_dict[key].squeeze()
+
+
+def resolve_adapter_type(adapter_type, state_dict=None):
+    """Resolve legacy "dora" adapter_type to the correct variant.
+
+    If adapter_type is "dora" and a checkpoint state_dict is provided,
+    attempts to detect dora-rows vs dora-cols from magnitude tensor shapes.
+    Falls back to "dora-rows" (paper-correct default).
+    All other adapter types are returned unchanged.
+    """
+    if adapter_type != "dora":
+        return adapter_type
+    if state_dict is not None:
+        detected = detect_dora_variant(state_dict)
+        if detected:
+            return detected
+    return "dora-rows"
+
+
+# ------------------- helper functions for safetensors LoRA checkpoints -------------------
+
+
+def save_lora_safetensors(state_dict, lora_config, path):
+    """Save a LoRA checkpoint in safetensors format with config as metadata.
+
+    The lora_config dict is JSON-serialized and stored under the metadata key
+    ``"lora_config"``, which safetensors stores in the file header as a string.
+
+    Args:
+        state_dict: Dict of LoRA tensors (from get_lora_state_dict).
+        lora_config: Dict with keys like rank, alpha, adapter_type, include, exclude.
+        path: Output file path (should end in .safetensors).
+    """
+    metadata = {"lora_config": json.dumps(lora_config)}
+    fp16_dict = {k: v.half() if v.is_floating_point() else v for k, v in state_dict.items()}
+    _st_save_file(fp16_dict, str(path), metadata=metadata)
+
+
+def load_lora_checkpoint(path):
+    """Load a LoRA checkpoint from either .safetensors or .ckpt format.
+
+    Returns:
+        Tuple of (state_dict, lora_config) where lora_config may be an empty
+        dict if the checkpoint has no embedded config.
+    """
+    path = str(path)
+    if path.endswith(".safetensors"):
+        with safe_open(path, framework="pt", device="cpu") as f:
+            state_dict = {key: f.get_tensor(key) for key in f.keys()}
+            metadata = f.metadata()
+        lora_config = {}
+        if metadata and "lora_config" in metadata:
+            lora_config = json.loads(metadata["lora_config"])
+        return state_dict, lora_config
+    else:
+        ckpt = torch.load(path, map_location="cpu", weights_only=False)
+        state_dict = ckpt["state_dict"]
+        lora_config = ckpt.get("lora_config", {})
+        return state_dict, lora_config
+
+
+def convert_lora_ckpt_to_safetensors(ckpt_path, output_path=None):
+    """Convert a .ckpt LoRA checkpoint to .safetensors format.
+
+    Args:
+        ckpt_path: Path to the input .ckpt file.
+        output_path: Path for the output .safetensors file.
+            If None, replaces .ckpt extension with .safetensors.
+
+    Returns:
+        The output path.
+    """
+    ckpt_path = str(ckpt_path)
+    if output_path is None:
+        if ckpt_path.endswith(".ckpt"):
+            output_path = ckpt_path[:-5] + ".safetensors"
+        else:
+            output_path = ckpt_path + ".safetensors"
+    state_dict, lora_config = load_lora_checkpoint(ckpt_path)
+    save_lora_safetensors(state_dict, lora_config, output_path)
+    return output_path
+
+
+# ------------------- helper functions for multi-LoRA stacking -------------------
+
+
+def get_lora_count(model):
+    """Return the number of stacked LoRAs (distinct lora_index values)."""
+    indices = set()
+    for _, p in get_lora_layers(model):
+        indices.add(p.lora_index)
+    return len(indices)
+
+
+def remap_lora_state_dict(state_dict, target_index):
+    """Remap LoRA state dict keys from parametrization index 0 to target_index.
+
+    Checkpoints always save keys with index 0 (e.g. *.parametrizations.weight.0.lora_A).
+    When loading the Nth stacked LoRA, we need to remap to index N.
+    """
+    if target_index == 0:
+        return state_dict
+    remapped = {}
+    for k, v in state_dict.items():
+        new_key = k.replace(
+            ".parametrizations.weight.0.",
+            ".parametrizations.weight.{}.".format(target_index)
+        )
+        remapped[new_key] = v
+    return remapped
+
+
+# ------------------- helper function for inferencing with multiple lora -------------------
+
+
+def _get_adapter_param_names(lora_layer):
+    """Return the list of trainable parameter names for this adapter type."""
+    if lora_layer.adapter_type == "lora":
+        return ["lora_A", "lora_B"]
+    elif lora_layer.adapter_type in ("dora", "dora-rows", "dora-cols"):
+        return ["lora_A", "lora_B", "magnitude"]
+    elif lora_layer.adapter_type == "bora":
+        return ["lora_A", "lora_B", "magnitude_r", "magnitude_c"]
+    elif lora_layer.adapter_type == "lora-xs":
+        return ["M_xs"]
+    elif lora_layer.adapter_type in ("dora-rows-xs", "dora-cols-xs"):
+        return ["M_xs", "magnitude"]
+    elif lora_layer.adapter_type == "bora-xs":
+        return ["M_xs", "magnitude_r", "magnitude_c"]
+    return []
+
+
+def _prepare_for_multiple_lora(lora_layer):
+    lora_layer._multi_lora_slots = []
+
+
+def _append_lora(lora_layer):
+    slot = {}
+    for name in _get_adapter_param_names(lora_layer):
+        slot[name] = nn.Parameter(getattr(lora_layer, name).clone())
+    lora_layer._multi_lora_slots.append(slot)
+
+
+def load_multiple_lora(model, lora_state_dicts):
+    model.apply(apply_to_lora(_prepare_for_multiple_lora))
+    for state_dict in lora_state_dicts:
+        _ = model.load_state_dict(state_dict, strict=False)
+        model.apply(apply_to_lora(_append_lora))
+    return model
+
+
+def _select_lora(lora_layer, index):
+    slot = lora_layer._multi_lora_slots[index]
+    for name, param in slot.items():
+        setattr(lora_layer, name, param)
+
+
+def select_lora(model, index):
+    model.apply(apply_to_lora(lambda x: _select_lora(x, index)))
+    return model
+
+# ------------------- helper function for merging multiple LoRAs -------------------
+
+def merge_loras_into_base_model(model, lora_configurations):
+    """
+    Merges multiple LoRAs with specified application weights directly into the base model's weights.
+    Works with all adapter types (LoRA, DoRA, LoRA-XS) by using the parametrization's own forward.
+    After merging, the LoRA parametrization on each layer is disabled.
+
+    Args:
+        model: The PyTorch model with LoRAParametrization already applied (via add_lora).
+        lora_configurations: A list of dicts, each specifying a LoRA to merge:
+            [
+                {
+                    'name': 'style1',
+                    'state_dict': loaded_torch_state_dict,
+                    'application_weight': 0.7
+                },
+                ...
+            ]
+    """
+    import torch
+
+    if not lora_configurations:
+        return
+
+    # Collect all parametrized layers with their LoRA instances
+    parametrized_layers = []
+    for name, module in model.named_modules():
+        if not hasattr(module, "parametrizations"):
+            continue
+        for attr_name in list(module.parametrizations.keys()):
+            plist = module.parametrizations[attr_name]
+            for p in plist:
+                if isinstance(p, LoRAParametrization):
+                    parametrized_layers.append((name, module, attr_name, p, plist.original))
+                    break
+
+    # Initialize per-layer delta accumulators
+    deltas = {id(orig): torch.zeros_like(orig.data) for _, _, _, _, orig in parametrized_layers}
+
+    for lora_cfg in lora_configurations:
+        lora_sd = lora_cfg['state_dict']
+        app_weight = lora_cfg['application_weight']
+
+        # Load this LoRA's adapter weights into the model's parametrizations
+        model.load_state_dict(lora_sd, strict=False)
+
+        # Compute each layer's delta using the parametrization's own forward
+        with torch.no_grad():
+            for name, module, attr_name, lora_p, original in parametrized_layers:
+                merged = lora_p(original)
+                deltas[id(original)] += app_weight * (merged - original)
+
+    # Apply accumulated deltas to original weights and disable LoRA
+    with torch.no_grad():
+        for name, module, attr_name, lora_p, original in parametrized_layers:
+            original.data.add_(deltas[id(original)])
+            lora_p.disable_lora()
+
+# ------------------- helper function for tying and untieing weights -------------------
+
+
+def tie_weights(linear: nn.Linear, embedding: nn.Embedding):
+    """tie the weights of the linear layer and the embedding layer both with the same lora"""
+    lora_p = embedding.parametrizations.weight[0]
+    if lora_p.adapter_type != "lora":
+        raise NotImplementedError(f"tie_weights only supports standard LoRA, got adapter_type='{lora_p.adapter_type}'")
+    # this line below is optional if the original is already tied
+    embedding.parametrizations.weight.original = linear.parametrizations.weight.original
+    embedding.parametrizations.weight[0].lora_A = linear.parametrizations.weight[0].lora_B
+    embedding.parametrizations.weight[0].lora_B = linear.parametrizations.weight[0].lora_A
+
+
+def untie_weights(linear: nn.Linear, embedding: nn.Embedding):
+    """untie the weights of the linear layer and the embedding layer"""
+    lora_p = embedding.parametrizations.weight[0]
+    if lora_p.adapter_type != "lora":
+        raise NotImplementedError(f"untie_weights only supports standard LoRA, got adapter_type='{lora_p.adapter_type}'")
+    embedding.parametrizations.weight.original = nn.Parameter(embedding.weight.original.clone())
+    embedding.parametrizations.weight[0].lora_A = nn.Parameter(embedding.parametrizations.weight[0].lora_A.clone())
+    embedding.parametrizations.weight[0].lora_B = nn.Parameter(embedding.parametrizations.weight[0].lora_B.clone())

stable_audio_3/models/pretransforms.py

@@ -0,0 +1,84 @@
+import torch
+import torch.nn as nn
+from einops import rearrange
+from torchaudio.transforms import Resample
+from .blocks import ResidualUnit, WNConv1d
+
+
+class AutoencoderPretransform(nn.Module):
+    def __init__(self, model, scale=1.0, iterate_batch=False, chunked=False):
+        super().__init__()
+        self.model = model
+        self.model.requires_grad_(False).eval()
+        self.scale = scale
+        self.downsampling_ratio = model.downsampling_ratio
+        self.io_channels = model.io_channels
+        self.enable_grad = False
+        self.iterate_batch = iterate_batch
+        self.chunked = chunked
+
+    def encode(self, x, **kwargs):
+        return self.model.encode_audio(x, chunked=self.chunked, iterate_batch=self.iterate_batch, **kwargs) / self.scale
+
+    def decode(self, z, chunked=None, **kwargs):
+        chunked = self.chunked if chunked is None else chunked
+        return self.model.decode_audio(z * self.scale, chunked=chunked, iterate_batch=self.iterate_batch, **kwargs)
+
+def fold_channels_into_batch(x):
+    x = rearrange(x, 'b c ... -> (b c) ...')
+    return x
+
+def unfold_channels_from_batch(x, channels):
+    if channels == 1:
+        return x.unsqueeze(1)
+    x = rearrange(x, '(b c) ... -> b c ...', c = channels)
+    return x
+
+
+class PatchedPretransform(nn.Module):
+    def __init__(self, channels, patch_size, oversampling = 1, postfilter_channels = 0, **kwargs):
+        super().__init__()
+        self.channels = channels
+        self.patch_size = patch_size
+        self.oversampling = oversampling
+
+        self.downsampling_ratio = patch_size
+        self.io_channels = channels
+        self.encoded_channels = channels * patch_size
+        self.enable_grad = False
+
+        if self.oversampling > 1:
+            self.input_upsampler = Resample(1, self.oversampling)
+            self.output_downsampler = Resample(self.oversampling, 1)
+
+        if postfilter_channels > 0:
+            self.postfilter = nn.Sequential(
+            WNConv1d(in_channels=channels, out_channels=postfilter_channels, kernel_size=7, padding=3, bias=True),
+            ResidualUnit(in_channels=postfilter_channels, out_channels=postfilter_channels,
+                         dilation=1, use_snake=True, bias = True),
+            ResidualUnit(in_channels=postfilter_channels, out_channels=postfilter_channels,
+                         dilation=3, use_snake=True, bias = True),
+            ResidualUnit(in_channels=postfilter_channels, out_channels=postfilter_channels,
+                         dilation=9, use_snake=True, bias = True),
+            WNConv1d(in_channels=postfilter_channels, out_channels=channels, kernel_size=7, padding=3, bias=False))
+
+    def _pad(self, x):
+        seq_len = x.shape[-1]
+        pad_len = (self.patch_size - (seq_len % self.patch_size)) % self.patch_size
+        if pad_len > 0:
+            x = torch.cat([x, torch.zeros_like(x[:, :, :pad_len])], dim=-1)
+        return x
+        
+    def encode(self, x):
+        if self.oversampling > 1:
+            x = self.input_upsampler(x)
+        x = self._pad(x)
+        x = rearrange(x, "b c (l h) -> b (c h) l", h=self.patch_size)
+        return x
+    def decode(self, x):
+        x = rearrange(x, "b (c h) l -> b c (l h)", h=self.patch_size)
+        if hasattr(self, 'postfilter'):
+            x = self.postfilter(x)
+        if self.oversampling > 1:
+            x = self.output_downsampler(x)
+        return x

stable_audio_3/models/transformer.py

@@ -0,0 +1,1272 @@
+from functools import reduce, partial
+from packaging import version
+import logging
+import math
+
+from einops import rearrange, repeat
+from einops.layers.torch import Rearrange
+import torch
+import torch.nn.functional as F
+from torch import nn, einsum
+from torch.amp import autocast
+from torch.nn.utils.parametrizations import weight_norm
+from typing import Callable, Literal, Optional
+try:
+    from torch.nn.attention.flex_attention import flex_attention, create_block_mask
+    flex_attention_available = True
+except ImportError:
+    flex_attention = None
+    create_block_mask = None
+    flex_attention_available = False
+
+try:
+    from flash_attn import flash_attn_func, flash_attn_kvpacked_func
+except ImportError as e:
+    print(e)
+    print('flash_attn not installed, disabling Flash Attention')
+    flash_attn_kvpacked_func = None
+    flash_attn_func = None
+
+try:
+    from flash_attn import flash_attn_varlen_func
+    from flash_attn.bert_padding import pad_input, unpad_input, index_first_axis
+except ImportError as e:
+    print(e)
+    print('flash_attn varlen/bert_padding not available, disabling varlen attention')
+    flash_attn_varlen_func = None
+    pad_input = None
+    unpad_input = None
+    index_first_axis = None
+
+
+def precompute_varlen_metadata(padding_mask: torch.Tensor):
+    """
+    Precompute varlen attention metadata once to avoid recomputation in every attention layer.
+
+    Args:
+        padding_mask: Boolean tensor of shape (batch, seq_len) where True = valid
+
+    Returns:
+        Dict with cu_seqlens, max_seqlen, indices, batch_size, seq_len for use in attention
+    """
+    if padding_mask is None or unpad_input is None:
+        return None
+
+    batch_size, seq_len = padding_mask.shape
+
+    # Compute cumulative sequence lengths (same for all of q, k, v)
+    seqlens = padding_mask.sum(dim=-1, dtype=torch.int32)
+    cu_seqlens = F.pad(torch.cumsum(seqlens, dim=0, dtype=torch.int32), (1, 0))
+    max_seqlen = seqlens.max().item()
+
+    # Compute indices for gathering valid tokens
+    # indices maps from packed position -> original (batch, seq) position
+    indices = torch.nonzero(padding_mask.flatten(), as_tuple=False).flatten()
+
+    return {
+        "cu_seqlens": cu_seqlens,
+        "max_seqlen": max_seqlen,
+        "indices": indices,
+        "batch_size": batch_size,
+        "seq_len": seq_len,
+    }
+
+from .utils import compile
+
+
+def _left_pad_to_match(emb, target_len):
+    """Left-pad or right-trim emb along seq dim to match target_len.
+
+    Used for local conditioning embeddings that need to align with x
+    without affecting prepended tokens (memory tokens, global cond, etc.).
+    """
+    emb_len = emb.shape[-2]
+    if emb_len < target_len:
+        return F.pad(emb, (0, 0, target_len - emb_len, 0), value=0.)
+    elif emb_len > target_len:
+        return emb[:, -target_len:, :]
+    return emb
+
+if flex_attention_available:
+    try:
+        torch._dynamo.config.cache_size_limit = 5000
+        flex_attention_compiled = torch.compile(flex_attention, dynamic=False, mode="max-autotune-no-cudagraphs")
+    except Exception as e:
+        logging.debug(f"Could not compile flex_attention, using uncompiled version: {e}")
+        flex_attention_compiled = flex_attention
+else:
+    flex_attention_compiled = None
+
+
+# Cache band block_masks for sliding-window attention fallback (flex_attention path).
+# Keyed by (seq_q, seq_k, w_left, w_right, device). create_block_mask is expensive
+# but the result is reused across all transformer layers and forward passes.
+_SLIDING_WINDOW_BLOCK_MASK_CACHE = {}
+
+def _get_sliding_window_block_mask(seq_q, seq_k, w_left, w_right, device):
+    key = (seq_q, seq_k, int(w_left), int(w_right), str(device))
+    bm = _SLIDING_WINDOW_BLOCK_MASK_CACHE.get(key)
+    if bm is None:
+        wl, wr = int(w_left), int(w_right)
+        def _band_mod(b, h, q_idx, kv_idx):
+            delta = kv_idx - q_idx
+            return (delta >= -wl) & (delta <= wr)
+        bm = create_block_mask(_band_mod, B=None, H=None, Q_LEN=seq_q, KV_LEN=seq_k, device=device)
+        _SLIDING_WINDOW_BLOCK_MASK_CACHE[key] = bm
+    return bm
+
+def _sliding_window_additive_mask(seq_q, seq_k, w_left, w_right, device, dtype):
+    """Build a (seq_q, seq_k) additive mask for masked SDPA fallback.
+    0 inside the band [i - w_left, i + w_right], -inf outside.
+    """
+    ii = torch.arange(seq_q, device=device)
+    jj = torch.arange(seq_k, device=device)
+    delta = jj[None, :] - ii[:, None]
+    in_band = (delta >= -int(w_left)) & (delta <= int(w_right))
+    mask = torch.zeros((seq_q, seq_k), dtype=dtype, device=device)
+    return mask.masked_fill(~in_band, float('-inf'))
+
+
+# Chunked-halo SDPA fallback. Math-equivalent to masked SDPA with a band
+# mask, but processes queries in non-overlapping chunks with a (w_left,
+# w_right) halo of keys/values on each side — every query stays inside its
+# chunk's softmax. Avoids materializing the O(N^2) mask.
+#
+# At realistic SAME-L decoder shapes (N=69632, W=17, packed sequence is
+# latent_length * (stride+1)): ~34x faster than full masked SDPA, and
+# ~140x less peak mask memory (~1 MB per chunk vs 9.7 GB for one N x N mask).
+# Chunk size is a tunable; 1024 is a good default at typical pretransform
+# decoder shapes. Larger chunks waste more compute on out-of-band tiles;
+# smaller chunks suffer from launch overhead.
+_SLIDING_WINDOW_CHUNK_SIZE = 1024
+
+def _sliding_window_chunked_halo_sdpa(q, k, v, w_left, w_right, chunk_size=_SLIDING_WINDOW_CHUNK_SIZE):
+    B, H, N, D = q.shape
+    outs = []
+    for q_start in range(0, N, chunk_size):
+        q_end = min(q_start + chunk_size, N)
+        k_start = max(0, q_start - int(w_left))
+        k_end = min(N, q_end + int(w_right))
+        q_c = q[..., q_start:q_end, :]
+        k_c = k[..., k_start:k_end, :]
+        v_c = v[..., k_start:k_end, :]
+        q_idx = torch.arange(q_start, q_end, device=q.device)
+        k_idx = torch.arange(k_start, k_end, device=q.device)
+        delta = k_idx[None, :] - q_idx[:, None]
+        in_band = (delta >= -int(w_left)) & (delta <= int(w_right))
+        mask = torch.zeros(delta.shape, dtype=q.dtype, device=q.device).masked_fill(~in_band, float('-inf'))
+        outs.append(F.scaled_dot_product_attention(q_c, k_c, v_c, attn_mask=mask, is_causal=False))
+    return torch.cat(outs, dim=-2)
+
+
+def checkpoint(function, *args, **kwargs):
+    kwargs.setdefault("use_reentrant", False)
+    # Preserve autocast context during recomputation to avoid dtype mismatches
+    if "context_fn" not in kwargs:
+        from torch.amp import autocast
+        import functools
+        # Get current autocast state
+        if torch.is_autocast_enabled():
+            dtype = torch.get_autocast_dtype('cuda')
+            def get_contexts():
+                return (
+                    autocast('cuda', dtype=dtype),
+                    autocast('cuda', dtype=dtype),
+                )
+            kwargs["context_fn"] = get_contexts
+    return torch.utils.checkpoint.checkpoint(function, *args, **kwargs)
+
+
+# Copied and modified from https://github.com/lucidrains/x-transformers/blob/main/x_transformers/attend.py under MIT License
+# License can be found in LICENSES/LICENSE_XTRANSFORMERS.txt
+
+def create_causal_mask(i, j, device):
+    return torch.ones((i, j), device = device, dtype = torch.bool).triu(j - i + 1)
+
+def or_reduce(masks):
+    head, *body = masks
+    for rest in body:
+        head = head | rest
+    return head
+
+# positional embeddings
+
+class AbsolutePositionalEmbedding(nn.Module):
+    def __init__(self, dim, max_seq_len):
+        super().__init__()
+        self.scale = dim ** -0.5
+        self.max_seq_len = max_seq_len
+        self.emb = nn.Embedding(max_seq_len, dim)
+
+    def forward(self, x, pos = None, seq_start_pos = None):
+        seq_len, device = x.shape[1], x.device
+        assert seq_len <= self.max_seq_len, f'you are passing in a sequence length of {seq_len} but your absolute positional embedding has a max sequence length of {self.max_seq_len}'
+
+        if pos is None:
+            pos = torch.arange(seq_len, device = device)
+
+        if seq_start_pos is not None:
+            pos = (pos - seq_start_pos[..., None]).clamp(min = 0)
+
+        pos_emb = self.emb(pos)
+        pos_emb = pos_emb * self.scale
+        return pos_emb
+
+class ScaledSinusoidalEmbedding(nn.Module):
+    def __init__(self, dim, theta = 10000):
+        super().__init__()
+        assert (dim % 2) == 0, 'dimension must be divisible by 2'
+        self.scale = nn.Parameter(torch.ones(1) * dim ** -0.5)
+
+        half_dim = dim // 2
+        freq_seq = torch.arange(half_dim).float() / half_dim
+        inv_freq = theta ** -freq_seq
+        self.register_buffer('inv_freq', inv_freq, persistent = False)
+
+    def forward(self, x, pos = None, seq_start_pos = None):
+        seq_len, device = x.shape[1], x.device
+
+        if pos is None:
+            pos = torch.arange(seq_len, device = device)
+
+        if seq_start_pos is not None:
+            pos = pos - seq_start_pos[..., None]
+
+        emb = einsum('i, j -> i j', pos, self.inv_freq)
+        emb = torch.cat((emb.sin(), emb.cos()), dim = -1)
+        return emb * self.scale
+    
+class RotaryEmbedding(nn.Module):
+    def __init__(
+        self,
+        dim,
+        use_xpos = False,
+        scale_base = 512,
+        interpolation_factor = 1.,
+        base = 10000,
+        base_rescale_factor = 1.
+    ):
+        super().__init__()
+        # proposed by reddit user bloc97, to rescale rotary embeddings to longer sequence length without fine-tuning
+        # has some connection to NTK literature
+        # https://www.reddit.com/r/LocalLLaMA/comments/14lz7j5/ntkaware_scaled_rope_allows_llama_models_to_have/
+        base *= base_rescale_factor ** (dim / (dim - 2))
+
+        inv_freq = 1. / (base ** (torch.arange(0, dim, 2).float() / dim))
+        self.register_buffer('inv_freq', inv_freq)
+
+        assert interpolation_factor >= 1.
+        self.interpolation_factor = interpolation_factor
+
+        if not use_xpos:
+            self.register_buffer('scale', None)
+            return
+
+        scale = (torch.arange(0, dim, 2) + 0.4 * dim) / (1.4 * dim)
+
+        self.scale_base = scale_base
+        self.register_buffer('scale', scale)
+
+    def forward_from_seq_len(self, seq_len):
+        device = self.inv_freq.device
+
+        t = torch.arange(seq_len, device = device)
+        return self.forward(t)
+
+    @autocast("cuda", enabled = False)
+    def forward(self, t):
+        device = self.inv_freq.device
+
+        t = t.to(torch.float32)
+
+        t = t / self.interpolation_factor
+
+        freqs = torch.einsum('i , j -> i j', t, self.inv_freq)
+        freqs = torch.cat((freqs, freqs), dim = -1)
+
+        if self.scale is None:
+            return freqs, 1.
+
+        power = (torch.arange(seq_len, device = device) - (seq_len // 2)) / self.scale_base
+        scale = self.scale ** rearrange(power, 'n -> n 1')
+        scale = torch.cat((scale, scale), dim = -1)
+
+        return freqs, scale
+
+def rotate_half(x):
+    x = rearrange(x, '... (j d) -> ... j d', j = 2)
+    x1, x2 = x.unbind(dim = -2)
+    return torch.cat((-x2, x1), dim = -1)
+
+
+@autocast("cuda", enabled = False)
+def apply_rotary_pos_emb(t, freqs, scale = 1):
+    out_dtype = t.dtype
+
+    # cast to float32 if necessary for numerical stability
+    dtype = reduce(torch.promote_types, (t.dtype, freqs.dtype, torch.float32))
+    rot_dim, seq_len = freqs.shape[-1], t.shape[-2]
+    freqs, t = freqs.to(dtype), t.to(dtype)
+    freqs = freqs[-seq_len:, :]
+
+    if t.ndim == 4 and freqs.ndim == 3:
+        freqs = rearrange(freqs, 'b n d -> b 1 n d')
+
+    # partial rotary embeddings, Wang et al. GPT-J
+    t, t_unrotated = t[..., :rot_dim], t[..., rot_dim:]
+
+    t = (t * freqs.cos() * scale ) + (rotate_half(t) * freqs.sin() * scale)
+
+    t, t_unrotated = t.to(out_dtype), t_unrotated.to(out_dtype)
+
+    return torch.cat((t, t_unrotated), dim = -1)
+
+# norms
+class DynamicTanh(nn.Module):
+    def __init__(self, dim, init_alpha=4.0, **kwargs):
+        super().__init__()
+        self.alpha = nn.Parameter(torch.ones(1) * init_alpha)
+        self.gamma = nn.Parameter(torch.ones(dim))
+        self.beta = nn.Parameter(torch.zeros(dim))
+
+    def forward(self, x):
+        x = F.tanh(self.alpha * x)
+        return self.gamma * x + self.beta
+
+class RunningInstanceNorm(nn.Module):
+    def __init__(self, dim, momentum = 0.99, eps = 1e-4, saturate = True, trainable_gain = True):
+        super().__init__()
+        self.register_buffer("running_mean", torch.zeros(1,1,dim))
+        self.register_buffer("running_std", torch.ones(1,1,dim))
+        self.saturate = saturate
+        self.eps = eps
+        self.momentum = momentum
+        self.dim = dim
+        self.trainable_gain = trainable_gain
+        if self.trainable_gain:
+            self.gain = nn.Parameter(torch.ones(1))
+    
+    def _update_stats(self, x):
+        self.running_mean = self.running_mean * self.momentum + x.detach().mean(dim = [0,1]).view(1, 1, self.dim) * (1 - self.momentum)
+        self.running_std  = (self.running_std * self.momentum + x.detach().std(dim = [0,1]).view(1, 1, self.dim) * (1 - self.momentum)).clip(min = self.eps)
+
+    def forward(self, x):
+        if self.training:
+            self._update_stats(x)
+        x = (x - self.running_mean) / self.running_std
+        if self.saturate:
+            x = torch.asinh(x)
+        if self.trainable_gain:
+            x = x * self.gain
+        return x
+        
+class LayerNorm(nn.Module):
+    def __init__(self, dim, bias=False, fix_scale=False, force_fp32=False, eps=1e-5):
+        """
+        bias-less layernorm has been shown to be more stable. most newer models have moved towards rmsnorm, also bias-less
+        """
+        super().__init__()
+
+        if fix_scale:
+            self.register_buffer("gamma", torch.ones(dim))
+        else:
+            self.gamma = nn.Parameter(torch.ones(dim))
+
+        if bias:
+            self.beta = nn.Parameter(torch.zeros(dim))
+        else:
+            self.register_buffer("beta", torch.zeros(dim))
+
+        self.eps = eps
+
+        self.force_fp32 = force_fp32
+
+    #@autocast("cuda", enabled = False)
+    def forward(self, x):
+        if not self.force_fp32:
+            return F.layer_norm(x, x.shape[-1:], weight=self.gamma, bias=self.beta, eps=self.eps)
+        else:
+            output = F.layer_norm(x.float(), x.shape[-1:], weight=self.gamma.float(), bias=self.beta.float(), eps=self.eps)
+            return output.to(x.dtype)
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim, fix_scale=False, force_fp32=False, eps=1e-5):
+        super().__init__()
+
+        if fix_scale:
+            self.register_buffer("gamma", torch.ones(dim))
+        else:
+            self.gamma = nn.Parameter(torch.ones(dim))
+
+        self.eps = eps
+
+        self.force_fp32 = force_fp32
+
+    def forward(self, x):
+        if not self.force_fp32:
+            return F.rms_norm(x, x.shape[-1:], weight=self.gamma, eps=self.eps)
+        else:
+            output = F.rms_norm(x.float(), x.shape[-1:], weight=self.gamma.float(), eps=self.eps)
+            return output.to(x.dtype)
+
+class LayerScale(nn.Module):
+    def __init__(self, dim, init_val = 1e-5):
+        super().__init__()
+        self.scale = nn.Parameter(torch.full([dim], init_val))
+    def forward(self, x):
+        return x * self.scale
+
+# feedforward
+
+class GLU(nn.Module):
+    def __init__(
+        self,
+        dim_in,
+        dim_out,
+        activation: Callable,
+        use_conv = False,
+        conv_kernel_size = 3,
+    ):
+        super().__init__()
+        self.act = activation
+        self.proj = nn.Linear(dim_in, dim_out * 2) if not use_conv else nn.Conv1d(dim_in, dim_out * 2, conv_kernel_size, padding = (conv_kernel_size // 2))
+        self.use_conv = use_conv
+
+    def forward(self, x):
+        if self.use_conv:
+            x = rearrange(x, 'b n d -> b d n')
+            x = self.proj(x)
+            x = rearrange(x, 'b d n -> b n d')
+        else:
+            x = self.proj(x)
+
+        x, gate = x.chunk(2, dim = -1)
+        return x * self.act(gate)
+
+class Sin(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x):
+        return torch.sin(3.14159265359 * x)
+
+class FeedForward(nn.Module):
+    def __init__(
+        self,
+        dim,
+        dim_out = None,
+        mult = 4,
+        no_bias = False,
+        glu = True,
+        use_conv = False,
+        conv_kernel_size = 3,
+        zero_init_output = True,
+        sinusoidal = False
+    ):
+        super().__init__()
+        inner_dim = int(dim * mult)
+
+        # Default to SwiGLU
+
+        activation = nn.SiLU() if not sinusoidal else Sin()
+
+        dim_out = dim if dim_out is None else dim_out
+
+        if glu:
+            linear_in = GLU(dim, inner_dim, activation)
+        else:
+            linear_in = nn.Sequential(
+                Rearrange('b n d -> b d n') if use_conv else nn.Identity(),
+                nn.Linear(dim, inner_dim, bias = not no_bias) if not use_conv else nn.Conv1d(dim, inner_dim, conv_kernel_size, padding = (conv_kernel_size // 2), bias = not no_bias),
+                Rearrange('b n d -> b d n') if use_conv else nn.Identity(),
+                activation
+            )
+
+        linear_out = nn.Linear(inner_dim, dim_out, bias = not no_bias) if not use_conv else nn.Conv1d(inner_dim, dim_out, conv_kernel_size, padding = (conv_kernel_size // 2), bias = not no_bias)
+
+        # init last linear layer to 0
+        if zero_init_output:
+            nn.init.zeros_(linear_out.weight)
+            if not no_bias:
+                nn.init.zeros_(linear_out.bias)
+
+
+        self.ff = nn.Sequential(
+            linear_in,
+            Rearrange('b d n -> b n d') if use_conv else nn.Identity(),
+            linear_out,
+            Rearrange('b n d -> b d n') if use_conv else nn.Identity(),
+        )
+
+    #@compile
+    def forward(self, x, varlen_metadata=None):
+        if varlen_metadata is not None and index_first_axis is not None and pad_input is not None:
+            # Pack valid tokens for efficient FFN computation (skip padding tokens)
+            # Padding positions become zeros after unpack, which is fine since FFN output
+            # is added to residual, preserving values at padding positions
+            batch_size = varlen_metadata["batch_size"]
+            seq_len = varlen_metadata["seq_len"]
+            indices = varlen_metadata["indices"]
+            dim = x.shape[-1]
+
+            # Pack to (N_valid, D)
+            x_packed = index_first_axis(x.reshape(-1, dim), indices)
+
+            # FFN on packed representation with pseudo-batch dim
+            x_packed = self.ff(x_packed.unsqueeze(0)).squeeze(0)
+
+            # Unpack back to (B, T, D)
+            return pad_input(x_packed, indices, batch_size, seq_len)
+        else:
+            return self.ff(x)
+
+class Attention(nn.Module):
+    def __init__(
+        self,
+        dim,
+        dim_heads = 64,
+        dim_context = None,
+        causal = False,
+        zero_init_output=True,
+        qk_norm_eps = 1e-6,
+        qk_norm: Literal['l2', 'ln', 'rms', 'dyt', 'none'] = 'none',
+        differential = False,
+        feat_scale = False
+    ):
+        super().__init__()
+        self.dim = dim
+        self.dim_heads = dim_heads
+
+        self.differential = differential
+
+        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:
+            if differential:
+                self.to_q = nn.Linear(dim, dim * 2, bias=False)
+                self.to_kv = nn.Linear(dim_kv, dim_kv * 3, bias=False)
+            else:
+                self.to_q = nn.Linear(dim, dim, bias=False)
+                self.to_kv = nn.Linear(dim_kv, dim_kv * 2, bias=False)
+        else:
+            if differential:
+                self.to_qkv = nn.Linear(dim, dim * 5, bias=False)
+            else:
+                self.to_qkv = nn.Linear(dim, dim * 3, bias=False)
+
+        self.to_out = nn.Linear(dim, dim, bias=False)
+
+        if zero_init_output:
+            nn.init.zeros_(self.to_out.weight)
+
+        if qk_norm not in ['l2', 'ln', 'rms', 'dyt','none']:
+            raise ValueError(f'qk_norm must be one of ["l2", "ln", "rms" ,"dyt", "none"], got {qk_norm}')
+            
+        self.qk_norm = qk_norm
+        self.qk_norm_eps = qk_norm_eps
+
+        if self.qk_norm == "ln":
+            self.q_norm = nn.LayerNorm(dim_heads, elementwise_affine=True, eps=qk_norm_eps)
+            self.k_norm = nn.LayerNorm(dim_heads, elementwise_affine=True, eps=qk_norm_eps)
+        elif self.qk_norm == "rms":
+            self.q_norm = RMSNorm(dim_heads, eps=qk_norm_eps)
+            self.k_norm = RMSNorm(dim_heads, eps=qk_norm_eps)
+        elif self.qk_norm == 'dyt':
+            self.q_norm = DynamicTanh(dim_heads)
+            self.k_norm = DynamicTanh(dim_heads)
+
+        self.feat_scale = feat_scale
+
+        if self.feat_scale:
+            self.lambda_dc = nn.Parameter(torch.zeros(dim))
+            self.lambda_hf = nn.Parameter(torch.zeros(dim))
+
+        self.causal = causal
+        
+    @compile
+    def apply_qk_layernorm(self, q, k):
+        q_type = q.dtype
+        k_type = k.dtype
+        q = self.q_norm(q).to(q_type)
+        k = self.k_norm(k).to(k_type)
+        return q, k
+
+
+    def apply_attn(self, q, k, v, causal = None, flex_attention_block_mask = None, flex_attention_score_mod = None, flash_attn_sliding_window = None, padding_mask = None, varlen_metadata = None):
+
+        if self.num_heads != self.kv_heads:
+             # Repeat interleave kv_heads to match q_heads for grouped query attention
+             heads_per_kv_head = self.num_heads // self.kv_heads
+             k, v = map(lambda t: t.repeat_interleave(heads_per_kv_head, dim = 1), (k, v))
+
+        flash_attn_available = flash_attn_func is not None
+        flash_attn_varlen_available = flash_attn_varlen_func is not None and index_first_axis is not None
+
+        if causal and (flex_attention_block_mask is not None or flex_attention_score_mod is not None):
+            flex_attention_block_mask = None
+            flex_attention_score_mod = None
+
+        if flex_attention_block_mask is not None or flex_attention_score_mod is not None:
+            # Flex attention path - use V-zeroing for padding mask
+            if padding_mask is not None:
+                mask_expanded = padding_mask.unsqueeze(1).unsqueeze(-1).to(v.dtype)
+                v = v * mask_expanded
+            out = flex_attention_compiled(q,k,v,
+                block_mask = flex_attention_block_mask,
+                score_mod = flex_attention_score_mod)
+        elif flash_attn_available and varlen_metadata is not None and flash_attn_varlen_available:
+            # Flash attention with varlen using precomputed metadata (fast path)
+            batch_size = varlen_metadata["batch_size"]
+            seq_len = varlen_metadata["seq_len"]
+            cu_seqlens = varlen_metadata["cu_seqlens"]
+            max_seqlen = varlen_metadata["max_seqlen"]
+            indices = varlen_metadata["indices"]
+
+            fa_dtype_in = q.dtype
+            # Rearrange to (B, T, H, D) for flash_attn
+            q, k, v = map(lambda t: rearrange(t, 'b h n d -> b n h d'), (q, k, v))
+
+            if fa_dtype_in != torch.float16 and fa_dtype_in != torch.bfloat16:
+                q, k, v = map(lambda t: t.to(torch.float16), (q, k, v))
+
+            # Pack q, k, v using precomputed indices (much faster than calling unpad_input 3x)
+            num_heads, head_dim = q.shape[2], q.shape[3]
+            q_unpad = index_first_axis(q.reshape(-1, num_heads, head_dim), indices)
+            k_unpad = index_first_axis(k.reshape(-1, num_heads, head_dim), indices)
+            v_unpad = index_first_axis(v.reshape(-1, num_heads, head_dim), indices)
+
+            out_unpad = flash_attn_varlen_func(
+                q_unpad, k_unpad, v_unpad,
+                cu_seqlens, cu_seqlens,
+                max_seqlen, max_seqlen,
+                causal=causal if causal is not None else False,
+                window_size=flash_attn_sliding_window if flash_attn_sliding_window is not None else (-1, -1),
+            )
+
+            # Pad output back to original shape
+            out = pad_input(out_unpad, indices, batch_size, seq_len)
+            out = rearrange(out.to(fa_dtype_in), 'b n h d -> b h n d')
+        elif flash_attn_available:
+            # Standard flash attention (no padding mask, or varlen imports not available)
+            # Apply V-zeroing fallback if padding_mask provided but we couldn't use varlen
+            if padding_mask is not None:
+                mask_expanded = padding_mask.unsqueeze(1).unsqueeze(-1).to(v.dtype)
+                v = v * mask_expanded
+            fa_dtype_in = q.dtype
+            q, k, v = map(lambda t: rearrange(t, 'b h n d -> b n h d'), (q, k, v))
+
+            if fa_dtype_in != torch.float16 and fa_dtype_in != torch.bfloat16:
+                q, k, v = map(lambda t: t.to(torch.float16), (q, k, v))
+
+            out = flash_attn_func(q, k, v, causal = causal, window_size=flash_attn_sliding_window if (flash_attn_sliding_window is not None) else [-1,-1])
+
+            out = rearrange(out.to(fa_dtype_in), 'b n h d -> b h n d')
+        else:
+            # No flash-attn available. Sliding-window fallback cascade:
+            #   Tier 2: flex_attention with band block_mask (best when torch.compile works)
+            #   Tier 3: chunked-halo masked SDPA           (math-equivalent, ~30x faster than tier 4)
+            #   Tier 4: full masked SDPA (N x N mask)      (last resort; high memory)
+            # For the no-sliding-window case, fall through to plain SDPA full attention.
+            # All apply V-zeroing for padding masks (cheap and equivalent to masking
+            # those positions out of attention output).
+            if padding_mask is not None:
+                mask_expanded = padding_mask.unsqueeze(1).unsqueeze(-1).to(v.dtype)
+                v = v * mask_expanded
+            if flash_attn_sliding_window is not None:
+                seq_q, seq_k = q.shape[2], k.shape[2]
+                wl, wr = flash_attn_sliding_window
+                handled = False
+                if flex_attention_available and flex_attention_compiled is not None:
+                    try:
+                        bm = _get_sliding_window_block_mask(seq_q, seq_k, wl, wr, q.device)
+                        out = flex_attention_compiled(q, k, v, block_mask=bm)
+                        handled = True
+                    except Exception as _flex_err:
+                        logging.debug(f"flex_attention failed, trying chunked-halo SDPA: {_flex_err}")
+                if not handled:
+                    try:
+                        out = _sliding_window_chunked_halo_sdpa(q, k, v, wl, wr)
+                        handled = True
+                    except Exception as _chunk_err:
+                        logging.debug(f"chunked-halo SDPA failed, falling back to full masked SDPA: {_chunk_err}")
+                if not handled:
+                    add_mask = _sliding_window_additive_mask(seq_q, seq_k, wl, wr, q.device, q.dtype)
+                    out = F.scaled_dot_product_attention(q, k, v, attn_mask=add_mask, is_causal=False)
+            else:
+                out = F.scaled_dot_product_attention(q, k, v, is_causal=causal if causal is not None else False)
+        return out
+
+
+    #@compile
+    def forward(
+        self,
+        x,
+        context = None,
+        rotary_pos_emb = None,
+        rotary_pos_emb_k = None,
+        causal = None,
+        flex_attention_block_mask = None,
+        flex_attention_score_mod = None,
+        flash_attn_sliding_window = None,
+        padding_mask = None,
+        varlen_metadata = None,
+    ):
+        h, kv_h, has_context = self.num_heads, self.kv_heads, context is not None
+
+        kv_input = context if has_context else x
+
+        if hasattr(self, 'to_q'):
+            # Use separate linear projections for q and k/v
+            if self.differential:
+                q, q_diff = self.to_q(x).chunk(2, dim=-1)
+                q, q_diff = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = h), (q, q_diff))
+                q = torch.stack([q, q_diff], dim = 1)
+                k, k_diff, v = self.to_kv(kv_input).chunk(3, dim=-1)
+                k, k_diff, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = kv_h), (k, k_diff, v))
+                k = torch.stack([k, k_diff], dim = 1)
+            else:
+                q = self.to_q(x)
+                q = rearrange(q, 'b n (h d) -> b h n d', h = h)
+                k, v = self.to_kv(kv_input).chunk(2, dim=-1)
+                k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = kv_h), (k, v))
+        else:
+            # Use fused linear projection
+            if self.differential:
+                q, k, v, q_diff, k_diff = self.to_qkv(x).chunk(5, dim=-1)
+                q, k, v, q_diff, k_diff  = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = h), (q, k, v, q_diff, k_diff))
+                q = torch.stack([q, q_diff], dim = 1)
+                k = torch.stack([k, k_diff], dim = 1)
+            else:
+                q, k, v = self.to_qkv(x).chunk(3, dim=-1)
+                q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = h), (q, k, v))
+
+        # Normalize q and k for cosine sim attention
+        if self.qk_norm == "l2":
+            q = F.normalize(q, dim=-1, eps=self.qk_norm_eps)
+            k = F.normalize(k, dim=-1, eps=self.qk_norm_eps)
+        elif self.qk_norm != "none":
+            q, k = self.apply_qk_layernorm(q, k)
+            # q = self.q_norm(q)
+            # k = self.k_norm(k)
+
+        if rotary_pos_emb is not None:
+            freqs, _ = rotary_pos_emb
+            q_dtype = q.dtype
+            k_dtype = k.dtype
+            q = q.to(torch.float32)
+            k = k.to(torch.float32)
+            freqs = freqs.to(torch.float32)
+        
+            q_freqs = freqs
+
+            if rotary_pos_emb_k is not None:
+                k_freqs, _ = rotary_pos_emb_k
+                k_freqs = k_freqs.to(torch.float32)
+            else:
+                k_freqs = q_freqs
+
+                if q.shape[-2] >= k.shape[-2]:
+                    ratio = q.shape[-2] / k.shape[-2]
+                    q_freqs, k_freqs = freqs, ratio * freqs
+                else:
+                    ratio = k.shape[-2] / q.shape[-2]
+                    q_freqs, k_freqs = ratio * freqs, freqs
+
+            q = apply_rotary_pos_emb(q, q_freqs)
+            k = apply_rotary_pos_emb(k, k_freqs)
+            q = q.to(v.dtype)
+            k = k.to(v.dtype)
+        
+        n, device = q.shape[-2], q.device
+
+        causal = self.causal if causal is None else causal
+
+        if n == 1 and causal:
+            causal = False
+
+        if self.differential:
+            q, q_diff = q.unbind(dim = 1)
+            k, k_diff = k.unbind(dim = 1)
+            out = self.apply_attn(q, k, v,  causal = causal, flex_attention_block_mask = flex_attention_block_mask, flex_attention_score_mod = flex_attention_score_mod, flash_attn_sliding_window = flash_attn_sliding_window, padding_mask = padding_mask, varlen_metadata = varlen_metadata)
+            out_diff = self.apply_attn(q_diff, k_diff, v, causal = causal, flex_attention_block_mask = flex_attention_block_mask, flex_attention_score_mod = flex_attention_score_mod, flash_attn_sliding_window = flash_attn_sliding_window, padding_mask = padding_mask, varlen_metadata = varlen_metadata)
+            out = out - out_diff
+        else:
+            out = self.apply_attn(q, k, v, causal = causal, flex_attention_block_mask = flex_attention_block_mask, flex_attention_score_mod = flex_attention_score_mod, flash_attn_sliding_window = flash_attn_sliding_window, padding_mask = padding_mask, varlen_metadata = varlen_metadata)
+        # merge heads
+        out = rearrange(out, ' b h n d -> b n (h d)')
+
+        # Communicate between heads
+        
+        # with autocast(enabled = False):
+        #     out_dtype = out.dtype
+        #     out = out.to(torch.float32)
+        #     out = self.to_out(out).to(out_dtype)
+        out = self.to_out(out)
+
+        if self.feat_scale:
+            if padding_mask is not None:
+                mask = padding_mask.unsqueeze(-1).to(out.dtype)  # (b, n, 1)
+                out_dc = (out * mask).sum(dim=-2, keepdim=True) / mask.sum(dim=-2, keepdim=True).clamp(min=1)
+                out_hf = out - out_dc
+                out = out + (self.lambda_dc * out_dc + self.lambda_hf * out_hf) * mask
+            else:
+                out_dc = out.mean(dim=-2, keepdim=True)
+                out_hf = out - out_dc
+                out = out + self.lambda_dc * out_dc + self.lambda_hf * out_hf
+
+        return out
+
+class ConformerModule(nn.Module):
+    def __init__(
+        self,
+        dim,
+        norm_kwargs = {},
+    ):     
+
+        super().__init__()
+
+        self.dim = dim
+        
+        self.in_norm = LayerNorm(dim, **norm_kwargs)
+        self.pointwise_conv = nn.Conv1d(dim, dim, kernel_size=1, bias=False)
+        self.glu = GLU(dim, dim, nn.SiLU())
+        self.depthwise_conv = nn.Conv1d(dim, dim, kernel_size=17, groups=dim, padding=8, bias=False)
+        self.mid_norm = LayerNorm(dim, **norm_kwargs) # This is a batch norm in the original but I don't like batch norm
+        self.swish = nn.SiLU()
+        self.pointwise_conv_2 = nn.Conv1d(dim, dim, kernel_size=1, bias=False)
+
+    #@compile
+    def forward(self, x):
+        x = self.in_norm(x)
+        x = rearrange(x, 'b n d -> b d n')
+        x = self.pointwise_conv(x)
+        x = rearrange(x, 'b d n -> b n d')
+        x = self.glu(x)
+        x = rearrange(x, 'b n d -> b d n')
+        x = self.depthwise_conv(x)
+        x = rearrange(x, 'b d n -> b n d')
+        x = self.mid_norm(x)
+        x = self.swish(x)
+        x = rearrange(x, 'b n d -> b d n')
+        x = self.pointwise_conv_2(x)
+        x = rearrange(x, 'b d n -> b n d')
+
+        return x
+
+class TransformerBlock(nn.Module):
+    def __init__(
+            self,
+            dim,
+            dim_heads = 64,
+            cross_attend = False,
+            dim_context = None,
+            global_cond_dim = None,
+            local_add_cond_dim = None,
+            modular_local_cond_configs = None,
+            causal = False,
+            zero_init_branch_outputs = True,
+            conformer = False,
+            layer_ix = -1,
+            add_rope = False,
+            layer_scale = False,
+            norm_type = 'layer_norm',
+            attn_kwargs = {},
+            ff_kwargs = {},
+            norm_kwargs = {}
+    ):
+        
+        super().__init__()
+        self.dim = dim
+        self.dim_heads = min(dim_heads,dim)
+        self.cross_attend = cross_attend
+        self.dim_context = dim_context
+        self.causal = causal
+       
+        if layer_scale and zero_init_branch_outputs:
+            print('zero_init_branch_outputs is redundant with layer_scale, setting zero_init_branch_outputs to False')
+            zero_init_branch_outputs = False
+        
+        if norm_type not in ['layer_norm', 'rms_norm', 'dyt']:
+            raise ValueError(f'norm_type must be one of ["layer_norm", "rms_norm", "dyt"], got {norm_type}')
+
+        norm_layer_map = {
+            'layer_norm': LayerNorm,
+            'rms_norm': RMSNorm,
+            'dyt': DynamicTanh
+        }
+        norm_layer = norm_layer_map[norm_type]
+
+        self.pre_norm = norm_layer(dim,**norm_kwargs)
+        self.add_rope = add_rope
+
+        self.self_attn = Attention(
+            dim,
+            dim_heads = self.dim_heads,
+            causal = causal,
+            zero_init_output=zero_init_branch_outputs,
+            **attn_kwargs
+        )
+
+        self.self_attn_scale = LayerScale(dim) if layer_scale else nn.Identity()
+
+        self.cross_attend = cross_attend
+        if cross_attend:
+            self.cross_attend_norm = norm_layer(dim, **norm_kwargs)
+            self.cross_attn = Attention(
+                dim,
+                dim_heads = self.dim_heads,
+                dim_context=dim_context,
+                causal = causal,
+                zero_init_output=zero_init_branch_outputs,
+                **attn_kwargs
+            )
+            self.cross_attn_scale = LayerScale(dim) if layer_scale else nn.Identity()
+        
+        self.ff_norm = norm_layer(dim, **norm_kwargs)
+        self.ff = FeedForward(dim, zero_init_output=zero_init_branch_outputs, **ff_kwargs)
+        self.ff_scale = LayerScale(dim) if layer_scale else nn.Identity()
+
+        self.layer_ix = layer_ix
+
+        self.conformer = None
+        if conformer:
+            self.conformer = ConformerModule(dim, norm_kwargs=norm_kwargs)
+            self.conformer_scale = LayerScale(dim) if layer_scale else nn.Identity()
+
+        self.global_cond_dim = global_cond_dim
+
+        if global_cond_dim is not None:
+            self.to_scale_shift_gate = nn.Parameter(torch.randn(6*dim)/dim**0.5)
+
+        self.local_add_cond_dim = local_add_cond_dim
+
+        if local_add_cond_dim is not None:
+            self.to_local_embed = nn.Sequential(
+                nn.Linear(local_add_cond_dim, dim),
+                nn.SiLU(),
+                nn.Linear(dim, dim)
+            )
+
+            nn.init.zeros_(self.to_local_embed[-1].weight)
+            nn.init.zeros_(self.to_local_embed[-1].bias)
+
+        else:
+            self.to_local_embed = None
+
+        # Modular local conditioning - independent projections per conditioning ID
+        self.modular_local_cond_configs = modular_local_cond_configs or []
+        self.modular_local_embeds = nn.ModuleDict()
+
+        for config in self.modular_local_cond_configs:
+            cond_id = config["id"]
+            cond_dim = config["dim"]
+            proj = nn.Sequential(
+                nn.Linear(cond_dim, dim),
+                nn.SiLU(),
+                nn.Linear(dim, dim)
+            )
+            # Zero-init output layer so new conditioning doesn't affect model initially
+            nn.init.zeros_(proj[-1].weight)
+            nn.init.zeros_(proj[-1].bias)
+            self.modular_local_embeds[cond_id] = proj
+
+        self.rope = RotaryEmbedding(self.dim_heads // 2) if add_rope else None
+
+    def _apply_local_conditioning(self, x, local_add_cond, modular_local_cond):
+        """Apply local additive and modular local conditioning to x."""
+        if local_add_cond is not None and self.to_local_embed is not None:
+            local_emb = self.to_local_embed(local_add_cond)
+            x = x + _left_pad_to_match(local_emb, x.shape[-2])
+
+        if modular_local_cond is not None and len(self.modular_local_embeds) > 0:
+            modular_sum = None
+            for cond_id, proj in self.modular_local_embeds.items():
+                if cond_id in modular_local_cond:
+                    local_emb = proj(modular_local_cond[cond_id])
+                    local_emb = _left_pad_to_match(local_emb, x.shape[-2])
+                    modular_sum = local_emb if modular_sum is None else modular_sum + local_emb
+            if modular_sum is not None:
+                x = x + modular_sum
+
+        return x
+
+    @compile
+    def forward(
+        self,
+        x,
+        context = None,
+        global_cond=None,
+        local_add_cond=None,
+        modular_local_cond=None,
+        rotary_pos_emb = None,
+        cross_attn_rotary_pos_emb = None,
+        self_attention_block_mask = None,
+        self_attention_score_mod = None,
+        cross_attention_block_mask = None,
+        cross_attention_score_mod = None,
+        self_attention_flash_sliding_window = None,
+        cross_attention_flash_sliding_window = None,
+        padding_mask = None,
+        varlen_metadata = None,
+    ):
+        if rotary_pos_emb is None and self.add_rope:
+            rotary_pos_emb = self.rope.forward_from_seq_len(x.shape[-2])
+
+        if self.global_cond_dim is not None and self.global_cond_dim > 0 and global_cond is not None:
+            
+            scale_self, shift_self, gate_self, scale_ff, shift_ff, gate_ff = (self.to_scale_shift_gate + global_cond).unsqueeze(1).chunk(6, dim=-1)
+
+            # self-attention with adaLN
+            residual = x
+            x = self.pre_norm(x)
+            x = x * (1 + scale_self) + shift_self
+            x = self.self_attn(x, rotary_pos_emb = rotary_pos_emb, flex_attention_block_mask = self_attention_block_mask, flex_attention_score_mod = self_attention_score_mod, flash_attn_sliding_window = self_attention_flash_sliding_window, padding_mask = padding_mask, varlen_metadata = varlen_metadata)
+            x = x * torch.sigmoid(1 - gate_self)
+            x = self.self_attn_scale(x)
+            x = x + residual
+
+            if context is not None and self.cross_attend:
+                if cross_attn_rotary_pos_emb is not None:
+                    x = x + self.cross_attn_scale(self.cross_attn(self.cross_attend_norm(x), rotary_pos_emb = rotary_pos_emb, rotary_pos_emb_k = cross_attn_rotary_pos_emb, context = context, flex_attention_block_mask = cross_attention_block_mask, flex_attention_score_mod = cross_attention_score_mod, flash_attn_sliding_window = cross_attention_flash_sliding_window))
+                else:
+                    x = x + self.cross_attn_scale(self.cross_attn(self.cross_attend_norm(x), context = context, flex_attention_block_mask = cross_attention_block_mask, flex_attention_score_mod = cross_attention_score_mod, flash_attn_sliding_window = cross_attention_flash_sliding_window))
+
+            if self.conformer is not None:
+                x = x + self.conformer_scale(self.conformer(x))
+
+            x = self._apply_local_conditioning(x, local_add_cond, modular_local_cond)
+
+            # feedforward with adaLN
+            residual = x
+            x = self.ff_norm(x)
+            x = x * (1 + scale_ff) + shift_ff
+            x = self.ff(x, varlen_metadata=varlen_metadata)
+            x = x * torch.sigmoid(1 - gate_ff)
+            x = self.ff_scale(x)
+            x = x + residual
+
+        else:
+            x = x + self.self_attn_scale(self.self_attn(self.pre_norm(x), rotary_pos_emb = rotary_pos_emb, flex_attention_block_mask = self_attention_block_mask, flex_attention_score_mod = self_attention_score_mod, flash_attn_sliding_window = self_attention_flash_sliding_window, padding_mask = padding_mask, varlen_metadata = varlen_metadata))
+
+            if context is not None and self.cross_attend:
+                if cross_attn_rotary_pos_emb is not None:
+                    x = x + self.cross_attn_scale(self.cross_attn(self.cross_attend_norm(x), rotary_pos_emb = rotary_pos_emb, rotary_pos_emb_k = cross_attn_rotary_pos_emb, context = context, flex_attention_block_mask = cross_attention_block_mask, flex_attention_score_mod = cross_attention_score_mod, flash_attn_sliding_window = cross_attention_flash_sliding_window))
+                else:
+                    x = x + self.cross_attn_scale(self.cross_attn(self.cross_attend_norm(x), context = context, flex_attention_block_mask = cross_attention_block_mask, flex_attention_score_mod = cross_attention_score_mod, flash_attn_sliding_window = cross_attention_flash_sliding_window))
+                    
+            if self.conformer is not None:
+                x = x + self.conformer_scale(self.conformer(x))
+
+            x = self._apply_local_conditioning(x, local_add_cond, modular_local_cond)
+
+            x = x + self.ff_scale(self.ff(self.ff_norm(x), varlen_metadata=varlen_metadata))
+            
+
+        return x
+        
+class ContinuousTransformer(nn.Module):
+    def __init__(
+        self,
+        dim,
+        depth,
+        *,
+        dim_in = None,
+        dim_out = None,
+        dim_heads = 64,
+        cross_attend=False,
+        cond_token_dim=None,
+        final_cross_attn_ix=-1,
+        global_cond_dim=None,
+        local_add_cond_dim=None,
+        modular_local_cond_configs=None,
+        causal=False,
+        rotary_pos_emb=True,
+        cross_attn_rotary_pos_emb=False,
+        zero_init_branch_outputs=True,
+        conformer=False,
+        use_sinusoidal_emb=False,
+        use_abs_pos_emb=False,
+        abs_pos_emb_max_length=10000,
+        num_memory_tokens=0,
+        sliding_window=None,
+        **kwargs
+        ):
+
+        super().__init__()
+
+        self.dim = dim
+        self.depth = depth
+        self.causal = causal
+        self.layers = nn.ModuleList([])
+
+        self.project_in = nn.Linear(dim_in, dim, bias=False) if dim_in is not None else nn.Identity()
+        self.project_out = nn.Linear(dim, dim_out, bias=False) if dim_out is not None else nn.Identity()
+
+        if rotary_pos_emb:
+            self.rotary_pos_emb = RotaryEmbedding(max(dim_heads // 2, 32))
+        else:
+            self.rotary_pos_emb = None
+
+        if cross_attn_rotary_pos_emb:
+            self.cross_attn_rotary_pos_emb = RotaryEmbedding(max(dim_heads // 2, 32))
+        else:
+            self.cross_attn_rotary_pos_emb = None
+
+        self.num_memory_tokens = num_memory_tokens
+        if num_memory_tokens > 0:
+            self.memory_tokens = nn.Parameter(torch.randn(num_memory_tokens, dim))
+
+        self.use_sinusoidal_emb = use_sinusoidal_emb
+        if use_sinusoidal_emb:
+            self.pos_emb = ScaledSinusoidalEmbedding(dim)
+
+        self.use_abs_pos_emb = use_abs_pos_emb
+        if use_abs_pos_emb:
+            self.pos_emb = AbsolutePositionalEmbedding(dim, abs_pos_emb_max_length + self.num_memory_tokens)
+
+        self.global_cond_embedder = None
+        if global_cond_dim is not None:
+            self.global_cond_embedder = nn.Sequential(
+                nn.Linear(global_cond_dim, dim),
+                nn.SiLU(),
+                nn.Linear(dim, dim * 6)
+            )
+
+        self.final_cross_attn_ix = final_cross_attn_ix
+
+        self.sliding_window = sliding_window
+
+        for i in range(depth):
+            should_cross_attend = cross_attend and (self.final_cross_attn_ix == -1 or i <= (self.final_cross_attn_ix))
+            self.layers.append(
+                TransformerBlock(
+                    dim,
+                    dim_heads = dim_heads,
+                    cross_attend = should_cross_attend,
+                    dim_context = cond_token_dim,
+                    global_cond_dim = global_cond_dim,
+                    local_add_cond_dim = local_add_cond_dim,
+                    modular_local_cond_configs = modular_local_cond_configs,
+                    causal = causal,
+                    zero_init_branch_outputs = zero_init_branch_outputs,
+                    conformer=conformer,
+                    layer_ix=i,
+                    **kwargs
+                )
+            )
+        
+    def forward(
+        self,
+        x,
+        context = None,
+        prepend_embeds = None,
+        global_cond = None,
+        local_add_cond = None,
+        modular_local_cond = None,
+        return_info = False,
+        use_checkpointing = True,
+        exit_layer_ix = None,
+        padding_mask: Optional[torch.Tensor] = None,
+        **kwargs
+    ):
+        batch, seq, device = *x.shape[:2], x.device
+
+        model_dtype = next(self.parameters()).dtype
+        x = x.to(model_dtype)
+
+        info = {
+            "hidden_states": [],
+        }
+
+        x = self.project_in(x)
+
+        if prepend_embeds is not None:
+            prepend_length, prepend_dim = prepend_embeds.shape[1:]
+
+            assert prepend_dim == x.shape[-1], 'prepend dimension must match sequence dimension'
+
+            x = torch.cat((prepend_embeds, x), dim = -2)
+
+        if self.num_memory_tokens > 0:
+            memory_tokens = self.memory_tokens.expand(batch, -1, -1)
+            x = torch.cat((memory_tokens, x), dim=1)
+
+        if self.rotary_pos_emb is not None:
+            rotary_pos_emb = self.rotary_pos_emb.forward_from_seq_len(x.shape[1])
+        else:
+            rotary_pos_emb = None
+
+        if self.cross_attn_rotary_pos_emb is not None:
+            cross_attn_rotary_pos_emb = self.cross_attn_rotary_pos_emb.forward_from_seq_len(context.shape[-1])
+        else:
+            cross_attn_rotary_pos_emb = None
+
+        if self.use_sinusoidal_emb or self.use_abs_pos_emb:
+            x = x + self.pos_emb(x)
+
+        if global_cond is not None and self.global_cond_embedder is not None:
+            global_cond = self.global_cond_embedder(global_cond)
+
+        # Extend padding mask for prepended tokens if provided
+        extended_padding_mask = None
+        varlen_metadata = None
+        if padding_mask is not None:
+            # Compute total prepend length (memory tokens + prepend_embeds)
+            prepend_length = self.num_memory_tokens
+            if prepend_embeds is not None:
+                prepend_length += prepend_embeds.shape[1]
+
+            # Prepend tokens are always valid for attention
+            if prepend_length > 0:
+                prepend_valid = torch.ones(batch, prepend_length, device=device, dtype=torch.bool)
+                extended_padding_mask = torch.cat([prepend_valid, padding_mask], dim=-1)
+            else:
+                extended_padding_mask = padding_mask
+
+            # Precompute varlen metadata once for all layers (major performance optimization)
+            # Only compute if varlen attention is actually available
+            if flash_attn_varlen_func is not None and index_first_axis is not None:
+                varlen_metadata = precompute_varlen_metadata(extended_padding_mask)
+
+        # Iterate over the transformer layers
+        for layer_ix, layer in enumerate(self.layers):
+
+            layer_kwargs = {
+                "context": context,
+                "rotary_pos_emb": rotary_pos_emb,
+                "cross_attn_rotary_pos_emb": cross_attn_rotary_pos_emb,
+                "global_cond": global_cond,
+                "local_add_cond": local_add_cond,
+                "modular_local_cond": modular_local_cond,
+                "self_attention_flash_sliding_window": self.sliding_window,
+                "padding_mask": extended_padding_mask,
+                "varlen_metadata": varlen_metadata
+            }
+
+            if use_checkpointing:
+                x = checkpoint(layer, x, **layer_kwargs, **kwargs)
+            else:
+                x = layer(x, **layer_kwargs, **kwargs)
+
+            if return_info:
+                info["hidden_states"].append(x)
+
+            if exit_layer_ix is not None and layer_ix == exit_layer_ix:
+                x = x[:, self.num_memory_tokens:, :]
+
+                if return_info:
+                    return x, info
+                
+                return x
+
+        x = x[:, self.num_memory_tokens:, :]
+
+        x = self.project_out(x)
+
+        if return_info:
+            return x, info
+        
+        return x

stable_audio_3/models/utils.py

@@ -0,0 +1,19 @@
+import torch
+import os
+
+torch._dynamo.config.cache_size_limit = max(64, torch._dynamo.config.cache_size_limit)
+torch._dynamo.config.suppress_errors = True
+
+# Get torch.compile flag from environment variable ENABLE_TORCH_COMPILE
+
+enable_torch_compile = os.environ.get("ENABLE_TORCH_COMPILE", "0") == "1"
+
+def compile(function, *args, **kwargs):
+    
+    if enable_torch_compile:
+        try:
+            return torch.compile(function, *args, **kwargs)
+        except RuntimeError:
+            return function
+
+    return function

stable_audio_3/training/diffusion.py

@@ -0,0 +1,1143 @@
+import math
+import pytorch_lightning as pl
+from pytorch_lightning.loggers import WandbLogger, CometLogger
+import os
+import torch
+import gc
+import typing as tp
+import torchaudio
+
+from einops import rearrange
+from safetensors.torch import save_file
+from functools import partial
+from torch.nn import functional as F
+
+from ..interface.aeiou import audio_spectrogram_image
+from ..inference.sampling import truncated_logistic_normal_rescaled, sample_timesteps_logsnr, sample_timesteps_logsnr_uniform, sample_diffusion
+from ..models.diffusion import ConditionedDiffusionModelWrapper
+from ..models.inpainting import random_inpaint_mask, MaskType
+from ..models.lora import add_lora, get_lora_params, get_lora_state_dict, LoRAParametrization, get_lora_layers, save_lora_safetensors, resolve_adapter_type, prepare_dora_state_dict, cast_base_to_precision
+from .utils import create_optimizer_from_config, create_scheduler_from_config, log_audio, log_image, log_metric, get_rank, create_augmented_padding_mask, compute_masked_loss, compute_normalized_mse, resize_padding_mask, StaggeredLogger, compute_per_elem_trim, trim_and_concat
+from time import time
+
+class Profiler:
+
+    def __init__(self):
+        self.ticks = [[time(), None]]
+
+    def tick(self, msg):
+        self.ticks.append([time(), msg])
+
+    def __repr__(self):
+        rep = 80 * "=" + "\n"
+        for i in range(1, len(self.ticks)):
+            msg = self.ticks[i][1]
+            ellapsed = self.ticks[i][0] - self.ticks[i - 1][0]
+            rep += msg + f": {ellapsed*1000:.2f}ms\n"
+        rep += 80 * "=" + "\n\n\n"
+        return rep
+
+class DiffusionCondTrainingWrapper(pl.LightningModule):
+    '''
+    Wrapper for training a conditional audio diffusion model.
+    '''
+    def __init__(
+            self,
+            model: ConditionedDiffusionModelWrapper,
+            lr: float = None,
+            mask_loss_weight: float = 0.0,
+            mask_padding_attention: bool = False,
+            silence_extension_scale_seconds: float = 0.0,
+            use_ema: bool = True,
+            log_loss_info: bool = False,
+            optimizer_configs: dict = None,
+            pre_encoded: bool = False,
+            cfg_dropout_prob = 0.1,
+            timestep_sampler: tp.Literal["uniform", "logit_normal", "trunc_logit_normal", "log_snr", "log_snr_uniform"] = "uniform",
+            timestep_sampler_options: tp.Optional[tp.Dict[str, tp.Any]] = None,
+            validation_timesteps = [0.1, 0.3, 0.5, 0.7, 0.9],
+            p_one_shot: float = 0.0,
+            inpainting_config: dict = None,
+            use_effective_length_for_schedule: bool = False,
+            sample_rate: int = 44100,
+            sample_size: int = None,
+            loss_normalization: tp.Literal["none", "timestep", "sample", "sample_channel"] = "none",
+            loss_norm_eps: float = 1e-6,
+            lora_config: tp.Optional[tp.Dict[str, tp.Any]] = None,
+            lora_state_dict: tp.Optional[tp.Dict[str, tp.Any]] = None,
+            svd_bases_path: tp.Optional[str] = None,
+            log_every_n_steps: int = 10,
+            ot_coupling: bool = False,
+            base_precision: tp.Optional[str] = None,
+    ):
+        super().__init__()
+
+        self.ot_coupling = ot_coupling
+
+        self.diffusion = model
+
+        self.lora_config = lora_config
+        if self.lora_config is not None:
+            # Don't use EMA with LoRA
+            use_ema = False
+            # Freeze the pre-trained model weights
+            self.diffusion.model.eval().requires_grad_(False)
+            self.diffusion.conditioner.eval().requires_grad_(False)
+            rank = self.lora_config.get("rank", 8)
+            lora_alpha = self.lora_config.get("alpha", rank)
+            adapter_type = self.lora_config.get("adapter_type", "lora")
+            include = self.lora_config.get("include", None)
+            exclude = self.lora_config.get("exclude", None)
+            # Resolve legacy "dora" to rows/cols variant
+            adapter_type = resolve_adapter_type(adapter_type, lora_state_dict)
+            print(f"LoRA config: rank={rank}, alpha={lora_alpha}, adapter_type={adapter_type}")
+            if include:
+                print(f"  include: {include}")
+            if exclude:
+                print(f"  exclude: {exclude}")
+            # Load pre-computed SVD bases for -XS adapter types
+            svd_bases = None
+            if svd_bases_path is not None:
+                print(f"Loading SVD bases from {svd_bases_path}")
+                svd_bases = torch.load(svd_bases_path, map_location="cpu", weights_only=True)
+            elif adapter_type.endswith("-xs"):
+                print("WARNING: -XS adapter without svd_bases_path — SVD will be computed per layer")
+            lora_config = {
+                torch.nn.Linear: {
+                    "weight": partial(LoRAParametrization.from_linear, rank=rank, lora_alpha=lora_alpha, adapter_type=adapter_type),
+                },
+                torch.nn.Conv1d: {
+                    "weight": partial(LoRAParametrization.from_conv1d, rank=rank, lora_alpha=lora_alpha, adapter_type=adapter_type),
+                }
+            }
+            # Add LoRA to the model
+            add_lora(self.diffusion.model, lora_config, include=include, exclude=exclude, svd_bases=svd_bases)
+            # Add LoRA to the conditioner
+            add_lora(self.diffusion.conditioner, lora_config, include=include, exclude=exclude, svd_bases=svd_bases)
+            print("lora layers:", len(get_lora_layers(self.diffusion)))
+
+            if lora_state_dict is not None:
+                # Old DoRA checkpoints saved magnitude as 2D (1,fan_in) or (fan_out,1);
+                # current code expects 1D. Squeeze so old checkpoints still load.
+                prepare_dora_state_dict(lora_state_dict)
+                self.diffusion.model.load_state_dict(lora_state_dict, strict=False)
+                self.diffusion.conditioner.load_state_dict(lora_state_dict, strict=False)
+
+            # Cast frozen base weights to lower precision if requested
+            if base_precision:
+                cast_base_to_precision(self.diffusion.model, base_precision)
+                cast_base_to_precision(self.diffusion.conditioner, base_precision)
+                if self.diffusion.pretransform is not None:
+                    self.diffusion.pretransform.to(
+                        torch.bfloat16 if base_precision in ("bf16", "bfloat16") else torch.float16
+                    )
+
+        self.diffusion_ema = None
+        self.mask_loss_weight = mask_loss_weight
+
+        # Attention masking for padded tokens
+        # Backward compat: if passed from training config, propagate to model
+        if mask_padding_attention and not self.diffusion.mask_padding_attention:
+            import warnings
+            warnings.warn("mask_padding_attention in training config is deprecated. Move to model.diffusion config.", FutureWarning)
+            self.diffusion.mask_padding_attention = mask_padding_attention
+        self.mask_padding_attention = self.diffusion.mask_padding_attention
+        self.silence_extension_scale_seconds = silence_extension_scale_seconds
+
+        self.cfg_dropout_prob = cfg_dropout_prob
+
+        self.rng = torch.quasirandom.SobolEngine(1, scramble=True)
+
+        self.timestep_sampler = timestep_sampler     
+
+        self.timestep_sampler_options = {} if timestep_sampler_options is None else timestep_sampler_options
+
+        if self.timestep_sampler == "log_snr":
+            self.mean_logsnr = self.timestep_sampler_options.get("mean_logsnr", -1.2)
+            self.std_logsnr = self.timestep_sampler_options.get("std_logsnr", 2.0)
+        elif self.timestep_sampler == "log_snr_uniform":
+            self.min_logsnr = self.timestep_sampler_options.get("min_logsnr", -6.0)
+            self.max_logsnr = self.timestep_sampler_options.get("max_logsnr", 5.0)
+
+        self.p_one_shot = p_one_shot
+
+        self.diffusion_objective = model.diffusion_objective
+
+        self.log_loss_info = log_loss_info
+
+        self._staggered_logger = StaggeredLogger(every_n_steps=log_every_n_steps)
+
+        assert lr is not None or optimizer_configs is not None, "Must specify either lr or optimizer_configs in training config"
+
+        if optimizer_configs is None:
+            optimizer_configs = {
+                "diffusion": {
+                    "optimizer": {
+                        "type": "Adam",
+                        "config": {
+                            "lr": lr
+                        }
+                    }
+                }
+            }
+        else:
+            if lr is not None:
+                print(f"WARNING: learning_rate and optimizer_configs both specified in config. Ignoring learning_rate and using optimizer_configs.")
+
+        self.optimizer_configs = optimizer_configs
+
+        self.pre_encoded = pre_encoded
+
+        # Loss normalization by target magnitude
+        # Options: "none", "timestep", "sample", "sample_channel"
+        self.loss_normalization = loss_normalization
+        self.loss_norm_eps = loss_norm_eps
+
+        # Inpainting
+        self.inpainting_config = inpainting_config
+        
+        if self.inpainting_config is not None:
+            self.inpaint_mask_kwargs = self.inpainting_config.get("mask_kwargs", {})
+
+        # Per-element schedule shift based on effective (unpadded) sequence length
+        # Backward compat: if passed from training config, propagate to model
+        if use_effective_length_for_schedule and not self.diffusion.use_effective_length_for_schedule:
+            import warnings
+            warnings.warn("use_effective_length_for_schedule in training config is deprecated. Move to model.diffusion config.", DeprecationWarning)
+            self.diffusion.use_effective_length_for_schedule = use_effective_length_for_schedule
+        self.use_effective_length_for_schedule = self.diffusion.use_effective_length_for_schedule
+        self.sample_rate = sample_rate
+        self.sample_size = sample_size
+
+        # FSDP
+        self.use_fsdp = False
+
+        # Validation
+        self.validation_timesteps = validation_timesteps
+
+        self.validation_step_outputs = {}
+
+        for validation_timestep in self.validation_timesteps:
+            self.validation_step_outputs[f'val/loss_{validation_timestep:.1f}'] = []
+
+    def configure_optimizers(self):
+        diffusion_opt_config = self.optimizer_configs['diffusion']
+
+        if self.lora_config is not None:
+            opt_params = [*get_lora_params(self.diffusion.model), *get_lora_params(self.diffusion.conditioner)]
+        elif diffusion_opt_config['optimizer'].get('type') == 'MuonAdamW':
+            # Pass (name, param) tuples so MuonAdamW can match fused layer patterns
+            opt_params = [(n, p) for n, p in self.diffusion.named_parameters() if p.requires_grad]
+        else:
+            # Only include parameters that require gradients (excludes frozen pretransform, conditioner, etc.)
+            opt_params = [p for p in self.diffusion.parameters() if p.requires_grad]
+
+        opt_diff = create_optimizer_from_config(diffusion_opt_config['optimizer'], opt_params)
+
+        if "scheduler" in diffusion_opt_config:
+            sched_diff = create_scheduler_from_config(diffusion_opt_config['scheduler'], opt_diff)
+            sched_diff_config = {
+                "scheduler": sched_diff,
+                "interval": "step"
+            }
+            return [opt_diff], [sched_diff_config]
+
+        return [opt_diff]
+
+    def training_step(self, batch, batch_idx):
+        reals, metadata = batch
+
+        p = Profiler()
+
+        if reals.ndim == 4 and reals.shape[0] == 1:
+            reals = reals[0]
+
+        diffusion_input = reals
+
+        p.tick("setup")
+
+        #with torch.amp.autocast(device_type="cuda"):
+        conditioning = self.diffusion.conditioner(metadata, self.device)
+
+        # Create batch tensor of padding masks from the metadata
+        # If padding_mask not provided, assume all positions are valid (no padding)
+        if all("padding_mask" in md for md in metadata):
+            padding_masks = torch.stack([md["padding_mask"][0] for md in metadata], dim=0).to(self.device)  # Shape (batch_size, sequence_length)
+        else:
+            # All-True mask: everything is signal, no padding
+            padding_masks = torch.ones(diffusion_input.shape[0], diffusion_input.shape[-1], dtype=torch.bool, device=self.device)
+
+        p.tick("conditioning")
+
+        if self.diffusion.pretransform is not None:
+            self.diffusion.pretransform.to(self.device)
+
+            if not self.pre_encoded:
+                with torch.cuda.amp.autocast(), torch.set_grad_enabled(self.diffusion.pretransform.enable_grad):
+                    self.diffusion.pretransform.train(self.diffusion.pretransform.enable_grad)
+                    diffusion_input = self.diffusion.pretransform.encode(diffusion_input)
+                    p.tick("pretransform")
+                    padding_masks = resize_padding_mask(padding_masks, diffusion_input.shape[-1])
+            else:
+                # Apply scale to pre-encoded latents if needed, as the pretransform encode function will not be run
+                if hasattr(self.diffusion.pretransform, "scale") and self.diffusion.pretransform.scale != 1.0:
+                    diffusion_input = diffusion_input / self.diffusion.pretransform.scale
+
+
+                if padding_masks.shape[-1] != diffusion_input.shape[-1]:
+                    padding_masks = resize_padding_mask(padding_masks, diffusion_input.shape[-1])
+
+        if self.timestep_sampler == "uniform":
+            # Draw uniformly distributed continuous timesteps
+            t = self.rng.draw(reals.shape[0])[:, 0].to(self.device)
+        elif self.timestep_sampler == "logit_normal":
+            t = torch.sigmoid(torch.randn(reals.shape[0], device=self.device))
+        elif self.timestep_sampler == "trunc_logit_normal":
+            # Draw from logistic truncated normal distribution
+            t = truncated_logistic_normal_rescaled(reals.shape[0]).to(self.device)
+
+            # Flip the distribution
+            t = 1 - t
+        elif self.timestep_sampler == "log_snr":
+            t = sample_timesteps_logsnr(reals.shape[0], mean_logsnr=self.mean_logsnr, std_logsnr=self.std_logsnr).to(self.device)
+        elif self.timestep_sampler == "log_snr_uniform":
+            t = sample_timesteps_logsnr_uniform(reals.shape[0], min_logsnr=self.min_logsnr, max_logsnr=self.max_logsnr).to(self.device)
+        else:
+            raise ValueError(f"Invalid timestep_sampler: {self.timestep_sampler}")
+
+        if self.diffusion.dist_shift is not None:
+            # Compute sequence length for schedule shift
+            if self.use_effective_length_for_schedule:
+                # Use per-element effective lengths derived from seconds_total (rounded up)
+                # This matches inference which computes effective length from seconds_total conditioning
+                # Fall back to padding_masks.sum() if seconds_total is not available
+                if all("seconds_total" in md for md in metadata):
+                    downsampling_ratio = self.diffusion.pretransform.downsampling_ratio if self.diffusion.pretransform is not None else 1
+                    effective_seq_len = torch.tensor(
+                        [int(math.ceil(int(md["seconds_total"] * self.sample_rate) / downsampling_ratio)) for md in metadata],
+                        device=self.device
+                    )
+                else:
+                    # Fallback: use padding mask sum
+                    effective_seq_len = padding_masks.sum(dim=-1)
+            else:
+                # Use total sequence length (original behavior)
+                effective_seq_len = diffusion_input.shape[2]
+            
+            # Shift the distribution
+            t = self.diffusion.dist_shift.shift(t, effective_seq_len)
+
+        if self.p_one_shot > 0:
+            # Set t to 1 with probability p_one_shot
+            t = torch.where(torch.rand_like(t) < self.p_one_shot, torch.ones_like(t), t)
+
+        # Calculate the noise schedule parameters for those timesteps
+        if self.diffusion_objective in ["rectified_flow", "rf_denoiser"]:
+            alphas, sigmas = 1-t, t
+
+        # Combine the ground truth data and the noise
+        alphas = alphas[:, None, None]
+        sigmas = sigmas[:, None, None]
+        noise = torch.randn_like(diffusion_input)
+
+        # Minibatch OT coupling: find optimal noise permutation for straighter transport paths
+        # Based on MelodyFlow (arXiv:2407.03648v2) Section 2.5.2
+        # Uses GPU-only Sinkhorn approximation to avoid CPU sync
+        if self.ot_coupling and diffusion_input.shape[0] > 1:
+            with torch.no_grad():
+                # Flatten to [batch, features] for distance computation
+                data_flat = diffusion_input.reshape(diffusion_input.shape[0], -1)
+                noise_flat = noise.reshape(noise.shape[0], -1)
+                # Squared L2 cost via matmul (faster than cdist, same optimal assignment)
+                aa = (data_flat * data_flat).sum(dim=1, keepdim=True)
+                bb = (noise_flat * noise_flat).sum(dim=1, keepdim=True)
+                cost_matrix = aa + bb.T - 2.0 * (data_flat @ noise_flat.T)
+                # Sinkhorn assignment (GPU-only, no CPU sync)
+                log_P = -cost_matrix / cost_matrix.detach().mean() # normalize for numerical stability
+                for _ in range(20):
+                    log_P = log_P - torch.logsumexp(log_P, dim=1, keepdim=True)
+                    log_P = log_P - torch.logsumexp(log_P, dim=0, keepdim=True)
+                # Sequential assignment from soft permutation matrix (guarantees valid permutation)
+                P = log_P.exp()
+                B = P.shape[0]
+                col_indices = torch.empty(B, dtype=torch.long, device=P.device)
+                used = torch.zeros(B, dtype=torch.bool, device=P.device)
+                for i in range(B):
+                    P[i, used] = -1
+                    col_indices[i] = P[i].argmax()
+                    used[col_indices[i]] = True
+                noise = noise[col_indices]
+
+        noised_inputs = diffusion_input * alphas + noise * sigmas
+
+        if self.diffusion_objective == "v":
+            targets = noise * alphas - diffusion_input * sigmas
+        elif self.diffusion_objective in ["rectified_flow", "rf_denoiser"]:
+            targets = noise - diffusion_input
+
+        p.tick("noise")
+
+        extra_args = {}
+
+        # Compute downsampling ratio for attention mask creation
+        downsampling_ratio = self.diffusion.pretransform.downsampling_ratio if self.diffusion.pretransform is not None else 1
+
+        # Create augmented padding mask with random silence extension
+        if self.mask_padding_attention and self.silence_extension_scale_seconds > 0:
+            augmented_padding_mask = create_augmented_padding_mask(
+                padding_masks,
+                silence_extension_scale_seconds=self.silence_extension_scale_seconds,
+                sample_rate=self.sample_rate,
+                downsampling_ratio=downsampling_ratio,
+            )
+        else:
+            augmented_padding_mask = padding_masks
+
+        # Loss mask defines signal vs padding regions for loss computation
+        # - mask_loss_weight controls padding contribution (0 = signal only)
+        # - When mask_padding_attention=True: only compute loss on signal (padding saw no attention)
+        loss_mask = augmented_padding_mask.to(torch.bool)
+
+        # Pass padding mask for attention masking - model handles prepend extension
+        if self.mask_padding_attention:
+            extra_args["padding_mask"] = augmented_padding_mask
+
+        if self.inpainting_config is not None:
+
+            # Max mask size is the full sequence length
+            max_mask_length = diffusion_input.shape[2]
+
+            # Create a mask of random length for a random slice of the input
+            inpaint_masked_input, inpaint_mask = random_inpaint_mask(diffusion_input, padding_masks=augmented_padding_mask, mask_padding=self.mask_padding_attention, **self.inpaint_mask_kwargs)
+
+            conditioning['inpaint_mask'] = [inpaint_mask]
+            conditioning['inpaint_masked_input'] = [inpaint_masked_input]
+
+            # Only compute loss on inpainted region (where model is generating)
+            loss_mask = loss_mask & ~inpaint_mask.squeeze(1).to(torch.bool)
+
+        output = self.diffusion(noised_inputs, t, cond=conditioning, cfg_dropout_prob = self.cfg_dropout_prob, **extra_args)
+        p.tick("diffusion")
+
+        if self.log_loss_info:
+            # Loss debugging logs
+            num_loss_buckets = 10
+            bucket_size = 1 / num_loss_buckets
+            loss_all = F.mse_loss(output, targets, reduction="none")
+
+            sigmas = rearrange(self.all_gather(sigmas), "w b c n -> (w b) c n").squeeze()
+
+            # gather loss_all across all GPUs
+            loss_all = rearrange(self.all_gather(loss_all), "w b c n -> (w b) c n")
+
+            # Bucket loss values based on corresponding sigma values, bucketing sigma values by bucket_size
+            loss_all = torch.stack([loss_all[(sigmas >= i) & (sigmas < i + bucket_size)].mean() for i in torch.arange(0, 1, bucket_size).to(self.device)])
+
+            # Log bucketed losses with corresponding sigma bucket values, if it's not NaN
+            debug_log_dict = {
+                f"model/loss_all_{i/num_loss_buckets:.1f}": loss_all[i].detach() for i in range(num_loss_buckets) if not torch.isnan(loss_all[i])
+            }
+
+            self.log_dict(debug_log_dict)
+
+        p.tick("loss_debug")
+
+        # Compute std only over non-padded positions when masking is active
+        if loss_mask is not None and self.mask_padding_attention:
+            mask_expanded = loss_mask.unsqueeze(1)  # [B, 1, T]
+            std_data = diffusion_input[mask_expanded.expand_as(diffusion_input)].std()
+            std_targets = targets[mask_expanded.expand_as(targets)].std().detach()
+        else:
+            std_data = diffusion_input.std()
+            std_targets = targets.std().detach()
+
+        log_dict = {
+            'train/std_data': std_data,
+            'train/std_targets': std_targets,
+            'train/lr': self.trainer.optimizers[0].param_groups[0]['lr']
+        }
+
+        p.tick("std_compute")
+
+        # Compute normalized MSE (normalization only affects non-"none" modes)
+        mse_loss_full = compute_normalized_mse(output, targets, loss_mask, self.loss_normalization, self.loss_norm_eps)
+
+        p.tick("mse_loss")
+
+        # Compute loss with signal/padding separation (returns already-detached metrics)
+        loss, signal_mean, padding_mean = compute_masked_loss(
+            mse_loss_full, loss_mask, self.mask_padding_attention, self.mask_loss_weight
+        )
+        mse_loss = loss
+
+        p.tick("masked_loss")
+
+        # When attention masking is on, compute_masked_loss excludes everything outside
+        # loss_mask (which now excludes inpaint context). Add context reconstruction loss
+        # so the model learns to preserve context regions during inpainting.
+        # (When mask_padding_attention=False, context is already included via mask_loss_weight.)
+        context_loss_mean = torch.tensor(0.0, device=loss.device)
+        if (self.inpainting_config is not None
+                and self.mask_padding_attention
+                and self.mask_loss_weight > 0):
+            # Context = inpaint_mask=1 (keep) AND padding_mask=1 (real audio, not padding)
+            inpaint_context = inpaint_mask.squeeze(1).to(torch.bool) & augmented_padding_mask.to(torch.bool)
+            n_ctx = inpaint_context.sum(dim=1) * mse_loss_full.shape[1]  # per-sample count
+            if n_ctx.sum() > 0:
+                context_vals = torch.where(inpaint_context.unsqueeze(1), mse_loss_full, 0.0)
+                context_loss_mean = (context_vals.sum(dim=(1, 2)) / (n_ctx + 1e-8)).mean()
+                loss = loss + context_loss_mean * self.mask_loss_weight
+
+        # Log separate signal/padding/context losses for monitoring
+        log_dict["train/mse_signal"] = signal_mean
+        log_dict["train/mse_masked_loss"] = padding_mean
+        log_dict["train/mse_context_loss"] = context_loss_mean.detach()
+
+        log_dict["train/mse_loss"] = mse_loss.detach()
+        log_dict["train/loss"] = loss.detach()
+
+        # Stash for external callbacks (e.g. loss-by-timestep logging)
+        self._last_t = t.detach()
+        self._last_per_elem_loss = mse_loss_full.detach().mean(dim=(1, 2))
+
+        self._staggered_logger.log(log_dict, self)
+
+        #p.tick("log_dict")
+        #print(f"Profiler: {p}")
+        return loss
+
+    def validation_step(self, batch, batch_idx):
+
+        reals, metadata = batch
+
+        if reals.ndim == 4 and reals.shape[0] == 1:
+            reals = reals[0]
+
+        diffusion_input = reals
+
+        with torch.amp.autocast("cuda"), torch.no_grad():
+            conditioning = self.diffusion.conditioner(metadata, self.device)
+
+        # Create batch tensor of padding masks from the metadata
+        if all("padding_mask" in md for md in metadata):
+            padding_masks = torch.stack([md["padding_mask"][0] for md in metadata], dim=0).to(self.device)
+        else:
+            padding_masks = torch.ones(diffusion_input.shape[0], diffusion_input.shape[-1], dtype=torch.bool, device=self.device)
+
+        if self.diffusion.pretransform is not None:
+            self.diffusion.pretransform.to(self.device)
+
+            if not self.pre_encoded:
+                with torch.amp.autocast("cuda"), torch.no_grad():
+                    self.diffusion.pretransform.train(self.diffusion.pretransform.enable_grad)
+                    diffusion_input = self.diffusion.pretransform.encode(diffusion_input)
+                    padding_masks = resize_padding_mask(padding_masks, diffusion_input.shape[-1])
+            else:
+                # Apply scale to pre-encoded latents if needed, as the pretransform encode function will not be run
+                if hasattr(self.diffusion.pretransform, "scale") and self.diffusion.pretransform.scale != 1.0:
+                    diffusion_input = diffusion_input / self.diffusion.pretransform.scale
+
+                if padding_masks.shape[-1] != diffusion_input.shape[-1]:
+                    padding_masks = resize_padding_mask(padding_masks, diffusion_input.shape[-1])
+
+        # Use padding mask directly for validation (no silence extension augmentation)
+        loss_mask = padding_masks.to(torch.bool)
+
+        extra_args = {}
+        if self.mask_padding_attention:
+            extra_args["padding_mask"] = padding_masks
+
+        # Set up inpainting conditioning for validation (FULL_MASK: all zeros)
+        if self.inpainting_config is not None:
+            inpaint_mask = torch.zeros(diffusion_input.shape[0], 1, diffusion_input.shape[2], device=self.device)
+            inpaint_masked_input = torch.zeros_like(diffusion_input)
+            conditioning['inpaint_mask'] = [inpaint_mask]
+            conditioning['inpaint_masked_input'] = [inpaint_masked_input]
+
+        for validation_timestep in self.validation_timesteps:
+
+            t = torch.full((reals.shape[0],), validation_timestep, device=self.device)
+
+            # Calculate the noise schedule parameters for those timesteps
+            if self.diffusion_objective in ["v"]:
+                alphas, sigmas = get_alphas_sigmas(t)
+            elif self.diffusion_objective in ["rectified_flow", "rf_denoiser"]:
+                alphas, sigmas = 1-t, t
+
+            # Combine the ground truth data and the noise
+            alphas = alphas[:, None, None]
+            sigmas = sigmas[:, None, None]
+            noise = torch.randn_like(diffusion_input)
+            noised_inputs = diffusion_input * alphas + noise * sigmas
+
+            if self.diffusion_objective == "v":
+                targets = noise * alphas - diffusion_input * sigmas
+            elif self.diffusion_objective in ["rectified_flow", "rf_denoiser"]:
+                targets = noise - diffusion_input
+
+            with torch.amp.autocast("cuda"), torch.no_grad():
+                output = self.diffusion(noised_inputs, t, cond=conditioning, cfg_dropout_prob = 0, **extra_args)
+
+                mse_loss_full = compute_normalized_mse(output, targets, loss_mask, self.loss_normalization, self.loss_norm_eps)
+                val_loss, _, _ = compute_masked_loss(
+                    mse_loss_full, loss_mask, self.mask_padding_attention, self.mask_loss_weight
+                )
+
+                self.validation_step_outputs[f'val/loss_{validation_timestep:.1f}'].append(val_loss.item())
+
+    def on_validation_epoch_end(self):
+        log_dict = {}
+        for validation_timestep in self.validation_timesteps:
+            outputs_key = f'val/loss_{validation_timestep:.1f}'
+            val_loss = sum(self.validation_step_outputs[outputs_key]) / len(self.validation_step_outputs[outputs_key])
+
+            # Gather losses across all GPUs
+            val_loss = self.all_gather(val_loss).mean().item()
+
+            log_metric(self.logger, outputs_key, val_loss, step=self.global_step)
+
+        # Get average over all timesteps
+        val_loss = torch.tensor([val for val in self.validation_step_outputs.values()]).mean()
+
+        # Gather losses across all GPUs
+        val_loss = self.all_gather(val_loss).mean().item()
+
+        log_metric(self.logger, 'val/avg_loss', val_loss, step=self.global_step)
+
+        # Reset validation losses
+        for validation_timestep in self.validation_timesteps:
+            self.validation_step_outputs[f'val/loss_{validation_timestep:.1f}'] = []
+
+
+    def export_model(self, path, use_safetensors=False):
+        if self.diffusion_ema is not None:
+            self.diffusion.model = self.diffusion_ema.ema_model
+
+        if use_safetensors:
+            save_file(self.diffusion.state_dict(), path)
+        else:
+            torch.save({"state_dict": self.diffusion.state_dict()}, path)
+
+    def export_lora_safetensors(self, path):
+        """Export LoRA weights as a safetensors file with embedded config."""
+        if self.lora_config is None:
+            raise ValueError("No LoRA config -- this wrapper is not in LoRA mode")
+        state_dict = {
+            **get_lora_state_dict(self.diffusion.model),
+            **get_lora_state_dict(self.diffusion.conditioner)
+        }
+        save_lora_safetensors(state_dict, self.lora_config, path)
+
+    def on_save_checkpoint(self, checkpoint):
+        if self.lora_config is not None:
+            checkpoint.clear()
+            checkpoint['state_dict'] = {
+                **get_lora_state_dict(self.diffusion.model),
+                **get_lora_state_dict(self.diffusion.conditioner)
+            }
+            checkpoint['lora_config'] = self.lora_config
+
+class DiffusionCondInpaintDemoCallback(pl.Callback):
+    def __init__(
+        self,
+        demo_every=2000,
+        demo_steps=250,
+        sample_size=65536,
+        sample_rate=48000,
+        demo_cfg_scales: tp.Optional[tp.List[int]] = [3, 5, 7],
+        demo_conditioning: tp.Optional[tp.List[tp.Dict[str, tp.Any]]] = None,
+        inpaint_demo_config: tp.Optional[tp.Dict[str, int]] = None,
+        num_demos: int = 0,
+        demo_dl=None,
+    ):
+        super().__init__()
+        self.demo_every = demo_every
+        self.demo_steps = demo_steps
+        self.demo_samples = sample_size
+        self.sample_rate = sample_rate
+        self.demo_cfg_scales = demo_cfg_scales
+        self.demo_conditioning = demo_conditioning or []
+        self.last_demo_step = -1
+
+        # Map config keys to MaskType enum
+        self._mask_type_map = {
+            "num_random_segments": MaskType.RANDOM_SEGMENTS,
+            "num_full_mask": MaskType.FULL_MASK,
+            "num_causal": MaskType.CAUSAL_MASK,
+            "num_random_spans": MaskType.RANDOM_SPANS,
+        }
+
+        # Legacy fallback: if no inpaint_demo_config but num_demos is set,
+        # use num_demos items with random mask sampling (old behavior)
+        if inpaint_demo_config is not None:
+            self.inpaint_demo_config = inpaint_demo_config
+            self.legacy_inpaint_demos = False
+        elif num_demos > 0:
+            self.inpaint_demo_config = {}
+            self.legacy_inpaint_demos = True
+            self.legacy_num_demos = num_demos
+        else:
+            self.inpaint_demo_config = {}
+            self.legacy_inpaint_demos = False
+
+        # Total inpainting demos needed from batch
+        if self.legacy_inpaint_demos:
+            self.num_inpaint_demos = self.legacy_num_demos
+        else:
+            self.num_inpaint_demos = sum(
+                self.inpaint_demo_config.get(k, 0) for k in self._mask_type_map
+            )
+
+        if demo_dl is not None:
+            self.demo_dl = iter(demo_dl)
+        else:
+            self.demo_dl = None
+
+        self._teacher_demo_done = False
+
+    def _generate_prompt_demos(self, module, trainer, is_rank_zero=True):
+        """Generate full t2m demos from specified prompts (FULL_MASK)."""
+        if not self.demo_conditioning:
+            return [], []
+
+        demo_cond = self.demo_conditioning
+        num_demos = len(demo_cond)
+
+        demo_samples = self.demo_samples
+        if module.diffusion.pretransform is not None:
+            demo_samples = demo_samples // module.diffusion.pretransform.downsampling_ratio
+
+        # Conditioning from prompts
+        conditioning = module.diffusion.conditioner(demo_cond, module.device)
+
+        # FULL_MASK: all-zero inpaint conditioning
+        io_channels = module.diffusion.io_channels
+        inpaint_mask = torch.zeros(num_demos, 1, demo_samples, device=module.device)
+        inpaint_masked_input = torch.zeros(num_demos, io_channels, demo_samples, device=module.device)
+        conditioning['inpaint_mask'] = [inpaint_mask]
+        conditioning['inpaint_masked_input'] = [inpaint_masked_input]
+
+        cond_inputs = module.diffusion.get_conditioning_inputs(conditioning)
+
+        noise = torch.randn(num_demos, io_channels, demo_samples, device=module.device)
+        model_dtype = next(module.diffusion.parameters()).dtype
+        noise = noise.to(model_dtype)
+
+        per_elem_trim = compute_per_elem_trim(demo_cond, self.sample_rate, margin_seconds=2)
+
+        model = module.diffusion_ema.ema_model if module.diffusion_ema is not None else module.diffusion.model
+
+        all_audio = []
+        all_context_masks = []
+
+        for cfg_scale in self.demo_cfg_scales:
+            if is_rank_zero:
+                print(f"Generating prompt demos for cfg scale {cfg_scale}")
+
+            with torch.amp.autocast("cuda"):
+                fakes = sample_diffusion(
+                    model=model,
+                    noise=noise,
+                    cond_inputs=cond_inputs,
+                    diffusion_objective=module.diffusion_objective,
+                    steps=self.demo_steps,
+                    cfg_scale=cfg_scale,
+                    conditioning=demo_cond,
+                    sample_rate=self.sample_rate,
+                    pretransform=module.diffusion.pretransform,
+                    mask_padding_attention=module.diffusion.mask_padding_attention,
+                    use_effective_length_for_schedule=module.diffusion.use_effective_length_for_schedule,
+                    headroom_seconds=5.0,
+                    dist_shift=module.diffusion.sampling_dist_shift,
+                    batch_cfg=True,
+                    disable_tqdm=not is_rank_zero,
+                    decode=True
+                )
+
+            fakes = trim_and_concat(fakes, per_elem_trim)
+
+            all_audio.append(fakes)
+
+        # Latent-resolution all-zeros mask (no context for prompt demos),
+        # trimmed to match the per-element audio durations
+        ds_ratio = module.diffusion.pretransform.downsampling_ratio if module.diffusion.pretransform is not None else 1
+        latent_trim = [t // ds_ratio if t is not None else None for t in per_elem_trim] if per_elem_trim is not None else None
+        latent_mask = torch.zeros(num_demos, 1, demo_samples)
+        context_mask = trim_and_concat(latent_mask, latent_trim).squeeze(0).cpu()
+        all_context_masks = [context_mask] * len(self.demo_cfg_scales)
+
+        del noise, conditioning, cond_inputs, inpaint_mask, inpaint_masked_input
+        torch.cuda.empty_cache()
+
+        return all_audio, all_context_masks
+
+    def _generate_inpaint_demos(self, module, trainer, is_rank_zero=True):
+        """Generate inpainting demos from batch data with forced mask types."""
+        if self.num_inpaint_demos == 0 or self.demo_dl is None:
+            return [], []
+
+        demo_reals, metadata = next(self.demo_dl)
+
+        if demo_reals.ndim == 4 and demo_reals.shape[0] == 1:
+            demo_reals = demo_reals[0]
+
+        demo_reals = demo_reals[:self.num_inpaint_demos]
+        metadata = metadata[:self.num_inpaint_demos]
+        model_dtype = next(module.diffusion.parameters()).dtype
+        demo_reals = demo_reals.to(module.device, dtype=model_dtype)
+
+        if not module.pre_encoded:
+            if module.diffusion.pretransform is not None:
+                module.diffusion.pretransform.to(module.device)
+                demo_reals = module.diffusion.pretransform.encode(demo_reals)
+        else:
+            if hasattr(module.diffusion.pretransform, "scale") and module.diffusion.pretransform.scale != 1.0:
+                demo_reals = demo_reals / module.diffusion.pretransform.scale
+
+        padding_masks = torch.stack([md["padding_mask"][0] for md in metadata], dim=0).to(module.device)
+        if padding_masks.shape[-1] != demo_reals.shape[-1]:
+            padding_masks = resize_padding_mask(padding_masks, demo_reals.shape[-1])
+        mask_padding = module.diffusion.mask_padding_attention
+
+        if self.legacy_inpaint_demos:
+            # Legacy: random mask type sampling (old behavior)
+            masked_input, mask = random_inpaint_mask(
+                demo_reals, padding_masks=padding_masks,
+                mask_padding=mask_padding,
+                **module.inpaint_mask_kwargs
+            )
+        else:
+            # New: forced mask types per config
+            all_masks = []
+            all_masked_inputs = []
+            idx = 0
+            for config_key, mask_type in self._mask_type_map.items():
+                count = self.inpaint_demo_config.get(config_key, 0)
+                if count == 0:
+                    continue
+                subset_reals = demo_reals[idx:idx+count]
+                subset_padding = padding_masks[idx:idx+count]
+                mi, m = random_inpaint_mask(
+                    subset_reals, padding_masks=subset_padding,
+                    mask_padding=mask_padding, force_mask_type=mask_type,
+                    **module.inpaint_mask_kwargs
+                )
+                all_masks.append(m)
+                all_masked_inputs.append(mi)
+                idx += count
+
+            mask = torch.cat(all_masks, dim=0)
+            masked_input = torch.cat(all_masked_inputs, dim=0)
+
+        conditioning = module.diffusion.conditioner(metadata, module.device)
+        conditioning['inpaint_mask'] = [mask]
+        conditioning['inpaint_masked_input'] = [masked_input]
+
+        cond_inputs = module.diffusion.get_conditioning_inputs(conditioning)
+
+        demo_samples = demo_reals.shape[2]
+        noise = torch.randn(demo_reals.shape[0], module.diffusion.io_channels, demo_samples, device=module.device)
+        model_dtype = next(module.diffusion.parameters()).dtype
+        noise = noise.to(model_dtype)
+
+        per_elem_trim = compute_per_elem_trim(metadata, self.sample_rate, margin_seconds=2)
+
+        # Trim and concatenate context mask at latent resolution,
+        # using same trimming basis as audio (per_elem_trim // ds_ratio)
+        ds_ratio = module.diffusion.pretransform.downsampling_ratio if module.diffusion.pretransform is not None else 1
+        latent_trim = [t // ds_ratio if t is not None else None for t in per_elem_trim] if per_elem_trim is not None else None
+
+        # Zero out padding region in mask for display — the mask is initialized to 1,
+        # so without mask_padding the padding frames show as false context in the overlay
+        display_mask = mask * padding_masks.unsqueeze(1)
+
+        context_mask = trim_and_concat(display_mask, latent_trim).squeeze(0).cpu()
+
+        model = module.diffusion_ema.ema_model if module.diffusion_ema is not None else module.diffusion.model
+
+        all_audio = []
+        all_context_masks = []
+
+        for cfg_scale in self.demo_cfg_scales:
+            if is_rank_zero:
+                print(f"Generating inpaint demos for cfg scale {cfg_scale}")
+
+            with torch.amp.autocast("cuda"):
+                fakes = sample_diffusion(
+                    model=model,
+                    noise=noise,
+                    cond_inputs=cond_inputs,
+                    diffusion_objective=module.diffusion_objective,
+                    steps=self.demo_steps,
+                    cfg_scale=cfg_scale,
+                    conditioning=metadata,
+                    sample_rate=self.sample_rate,
+                    pretransform=module.diffusion.pretransform,
+                    mask_padding_attention=module.diffusion.mask_padding_attention,
+                    use_effective_length_for_schedule=module.diffusion.use_effective_length_for_schedule,
+                    headroom_seconds=5.0,
+                    dist_shift=module.diffusion.sampling_dist_shift,
+                    batch_cfg=True,
+                    disable_tqdm=not is_rank_zero,
+                    decode=True
+                )
+
+            fakes = trim_and_concat(fakes, per_elem_trim)
+
+            all_audio.append(fakes)
+            all_context_masks.append(context_mask)
+
+        del noise, conditioning, cond_inputs, mask, masked_input, padding_masks, demo_reals
+        torch.cuda.empty_cache()
+
+        return all_audio, all_context_masks
+
+    @torch.no_grad()
+    def on_train_batch_end(self, trainer, module: DiffusionCondTrainingWrapper, outputs, batch, batch_idx):
+        if (trainer.global_step - 1) % self.demo_every != 0 or self.last_demo_step == trainer.global_step:
+            return
+
+        is_rank_zero = get_rank() == 0
+
+        module.eval()
+
+        self.last_demo_step = trainer.global_step
+
+        try:
+            # Generate both types of demos, freeing intermediates between phases
+            prompt_audio, prompt_masks = self._generate_prompt_demos(module, trainer, is_rank_zero)
+            torch.cuda.empty_cache()
+
+            inpaint_audio, inpaint_masks = self._generate_inpaint_demos(module, trainer, is_rank_zero)
+            torch.cuda.empty_cache()
+
+            # Combine per cfg scale (prompt_audio and inpaint_audio have one entry per cfg scale)
+            if is_rank_zero:
+                for i, cfg_scale in enumerate(self.demo_cfg_scales):
+                    parts = []
+                    mask_parts = []
+
+                    if i < len(prompt_audio):
+                        parts.append(prompt_audio[i])
+                        mask_parts.append(prompt_masks[i])
+
+                    if i < len(inpaint_audio):
+                        parts.append(inpaint_audio[i])
+                        mask_parts.append(inpaint_masks[i])
+
+                    if not parts:
+                        continue
+
+                    combined_audio = torch.cat(parts, dim=-1)
+                    combined_mask = torch.cat(mask_parts, dim=-1) if mask_parts else None
+
+                    filename = f'demo_cfg_{cfg_scale}_{trainer.global_step:08}.wav'
+                    combined_audio = combined_audio.to(torch.float32).div(torch.max(torch.abs(combined_audio))).mul(32767).to(torch.int16).cpu()
+                    torchaudio.save(filename, combined_audio, self.sample_rate)
+
+                    log_audio(trainer.logger, f'demo_cfg_{cfg_scale}', filename, self.sample_rate)
+                    log_image(trainer.logger, f'demo_melspec_left_cfg_{cfg_scale}', audio_spectrogram_image(combined_audio, context_mask=combined_mask))
+                    if isinstance(trainer.logger, (WandbLogger, CometLogger)):
+                        os.remove(filename)
+
+            # Teacher ODE warmup diagnostic: mirror the exact ODE warmup sample_diffusion call
+            # and decode the target to verify teacher output quality.
+            # Only runs on the first demo.
+            # Generates both prompt and inpaint demos, consistent with the main callback.
+            teacher_ref = getattr(module, '_teacher', None) or getattr(module, 'teacher_model', None)
+            if not self._teacher_demo_done and teacher_ref is not None:
+                self._teacher_demo_done = True
+                if is_rank_zero:
+                    print("Generating teacher ODE warmup diagnostic")
+                try:
+                    pretransform = module.diffusion.pretransform  # Shared pretransform (not on teacher)
+                    io_channels = teacher_ref.io_channels
+                    ode_warmup_config = getattr(module, 'ode_warmup_config', {})
+                    teacher_cfg = getattr(module, 'ode_warmup_cfg', self.demo_cfg_scales[0])
+                    ode_steps = getattr(module, 'ode_n_sampling_steps', 20)
+                    mask_padding = module.diffusion.mask_padding_attention
+                    ds_ratio = pretransform.downsampling_ratio if pretransform is not None else 1
+
+                    # --- Teacher prompt demos (FULL_MASK, same as _generate_prompt_demos) ---
+                    prompt_target = None
+                    prompt_per_elem_trim = None
+                    prompt_context_mask = None
+
+                    demo_cond = self.demo_conditioning
+                    if demo_cond:
+                        num_demos = len(demo_cond)
+                        demo_samples = self.demo_samples
+                        if pretransform is not None:
+                            demo_samples = demo_samples // ds_ratio
+
+                        with torch.no_grad():
+                            teacher_conditioning = teacher_ref.conditioner(demo_cond, module.device)
+                        inpaint_mask = torch.zeros(num_demos, 1, demo_samples, device=module.device)
+                        inpaint_masked_input = torch.zeros(num_demos, io_channels, demo_samples, device=module.device)
+                        teacher_conditioning['inpaint_mask'] = [inpaint_mask]
+                        teacher_conditioning['inpaint_masked_input'] = [inpaint_masked_input]
+                        with torch.no_grad():
+                            teacher_cond_inputs = teacher_ref.get_conditioning_inputs(teacher_conditioning)
+
+                        noise = torch.randn(num_demos, io_channels, demo_samples, device=module.device)
+                        noise = noise.to(next(teacher_ref.parameters()).dtype)
+                        prompt_per_elem_trim = compute_per_elem_trim(demo_cond, self.sample_rate, margin_seconds=2)
+
+                        prompt_target = sample_diffusion(
+                            model=teacher_ref.model,
+                            noise=noise,
+                            cond_inputs=teacher_cond_inputs,
+                            diffusion_objective=teacher_ref.diffusion_objective,
+                            steps=ode_steps,
+                            cfg_scale=teacher_cfg,
+                            conditioning=demo_cond,
+                            sample_rate=teacher_ref.sample_rate,
+                            pretransform=pretransform,
+                            mask_padding_attention=mask_padding,
+                            use_effective_length_for_schedule=module.diffusion.use_effective_length_for_schedule,
+                            padding_mask=None,
+                            dist_shift=teacher_ref.sampling_dist_shift,
+                            sampler_type=ode_warmup_config.get('sampler', 'dpmpp'),
+                            batch_cfg=True,
+                            disable_tqdm=not is_rank_zero,
+                            decode=False,
+                        )
+
+                        prompt_latent_trim = [t // ds_ratio if t is not None else None for t in prompt_per_elem_trim] if prompt_per_elem_trim is not None else None
+                        prompt_context_mask = trim_and_concat(
+                            torch.zeros(num_demos, 1, demo_samples), prompt_latent_trim
+                        ).squeeze(0).cpu()
+
+                    # --- Teacher inpaint demos (same mask logic as _generate_inpaint_demos) ---
+                    inpaint_target = None
+                    inpaint_per_elem_trim = None
+                    inpaint_context_mask = None
+
+                    if self.num_inpaint_demos > 0 and self.demo_dl is not None:
+                        try:
+                            inpaint_reals, inpaint_metadata = next(self.demo_dl)
+                            if inpaint_reals.ndim == 4 and inpaint_reals.shape[0] == 1:
+                                inpaint_reals = inpaint_reals[0]
+                            inpaint_reals = inpaint_reals[:self.num_inpaint_demos]
+                            inpaint_metadata = inpaint_metadata[:self.num_inpaint_demos]
+                            inpaint_reals = inpaint_reals.to(module.device)
+
+                            if not module.pre_encoded:
+                                if pretransform is not None:
+                                    inpaint_reals = pretransform.encode(inpaint_reals)
+                            else:
+                                if hasattr(pretransform, "scale") and pretransform.scale != 1.0:
+                                    inpaint_reals = inpaint_reals / pretransform.scale
+
+                            inpaint_padding_masks = torch.stack(
+                                [md["padding_mask"][0] for md in inpaint_metadata], dim=0
+                            ).to(module.device)
+
+                            if self.legacy_inpaint_demos:
+                                masked_input, mask = random_inpaint_mask(
+                                    inpaint_reals, padding_masks=inpaint_padding_masks,
+                                    mask_padding=mask_padding, **module.inpaint_mask_kwargs
+                                )
+                            else:
+                                all_masks = []
+                                all_masked_inputs = []
+                                idx = 0
+                                for config_key, mask_type in self._mask_type_map.items():
+                                    count = self.inpaint_demo_config.get(config_key, 0)
+                                    if count == 0:
+                                        continue
+                                    mi, m = random_inpaint_mask(
+                                        inpaint_reals[idx:idx+count],
+                                        padding_masks=inpaint_padding_masks[idx:idx+count],
+                                        mask_padding=mask_padding, force_mask_type=mask_type,
+                                        **module.inpaint_mask_kwargs
+                                    )
+                                    all_masks.append(m)
+                                    all_masked_inputs.append(mi)
+                                    idx += count
+                                mask = torch.cat(all_masks, dim=0)
+                                masked_input = torch.cat(all_masked_inputs, dim=0)
+
+                            with torch.no_grad():
+                                inpaint_teacher_cond = teacher_ref.conditioner(inpaint_metadata, module.device)
+                            inpaint_teacher_cond['inpaint_mask'] = [mask]
+                            inpaint_teacher_cond['inpaint_masked_input'] = [masked_input]
+                            with torch.no_grad():
+                                inpaint_cond_inputs = teacher_ref.get_conditioning_inputs(inpaint_teacher_cond)
+
+                            inpaint_samples = inpaint_reals.shape[2]
+                            inpaint_noise = torch.randn(
+                                inpaint_reals.shape[0], io_channels, inpaint_samples, device=module.device
+                            ).to(next(teacher_ref.parameters()).dtype)
+                            inpaint_per_elem_trim = compute_per_elem_trim(inpaint_metadata, self.sample_rate, margin_seconds=2)
+
+                            inpaint_target = sample_diffusion(
+                                model=teacher_ref.model,
+                                noise=inpaint_noise,
+                                cond_inputs=inpaint_cond_inputs,
+                                diffusion_objective=teacher_ref.diffusion_objective,
+                                steps=ode_steps,
+                                cfg_scale=teacher_cfg,
+                                conditioning=inpaint_metadata,
+                                sample_rate=teacher_ref.sample_rate,
+                                pretransform=pretransform,
+                                mask_padding_attention=mask_padding,
+                                use_effective_length_for_schedule=module.diffusion.use_effective_length_for_schedule,
+                                padding_mask=None,
+                                dist_shift=teacher_ref.sampling_dist_shift,
+                                sampler_type=ode_warmup_config.get('sampler', 'dpmpp'),
+                                batch_cfg=True,
+                                disable_tqdm=not is_rank_zero,
+                                decode=False,
+                            )
+
+                            # Context mask for overlay (same as _generate_inpaint_demos)
+                            display_mask = mask * inpaint_padding_masks.unsqueeze(1)
+                            inpaint_latent_trim = [t // ds_ratio if t is not None else None for t in inpaint_per_elem_trim] if inpaint_per_elem_trim is not None else None
+                            inpaint_context_mask = trim_and_concat(display_mask, inpaint_latent_trim).squeeze(0).cpu()
+                        except StopIteration:
+                            if is_rank_zero:
+                                print("Teacher diagnostic: no inpaint batch available from demo_dl")
+
+                    # --- Combine and log (same pattern as main callback) ---
+                    if is_rank_zero:
+                        parts = []
+                        mask_parts = []
+
+                        if prompt_target is not None:
+                            decoded_prompt = pretransform.decode(prompt_target.float())
+                            decoded_prompt = trim_and_concat(decoded_prompt, prompt_per_elem_trim)
+                            parts.append(decoded_prompt)
+                            mask_parts.append(prompt_context_mask)
+
+                        if inpaint_target is not None:
+                            decoded_inpaint = pretransform.decode(inpaint_target.float())
+                            decoded_inpaint = trim_and_concat(decoded_inpaint, inpaint_per_elem_trim)
+                            parts.append(decoded_inpaint)
+                            mask_parts.append(inpaint_context_mask)
+
+                        if parts:
+                            combined_audio = torch.cat(parts, dim=-1)
+                            combined_mask = torch.cat(mask_parts, dim=-1) if mask_parts else None
+                            filename = f'demo_teacher_target_{trainer.global_step:08}.wav'
+                            combined_audio = combined_audio.to(torch.float32).div(torch.max(torch.abs(combined_audio))).mul(32767).to(torch.int16).cpu()
+                            torchaudio.save(filename, combined_audio, self.sample_rate)
+                            log_audio(trainer.logger, f'demo_teacher_target', filename, self.sample_rate)
+                            log_image(trainer.logger, f'demo_teacher_target_melspec', audio_spectrogram_image(combined_audio, context_mask=combined_mask))
+                            os.remove(filename)
+
+                    del prompt_target, inpaint_target
+                except Exception as e:
+                    if is_rank_zero:
+                        print(f"Teacher ODE warmup diagnostic failed: {e}")
+                        import traceback
+                        traceback.print_exc()
+
+        except Exception as e:
+            if is_rank_zero:
+                print(f'{type(e).__name__}: {e}')
+            raise e
+        finally:
+            gc.collect()
+            torch.cuda.empty_cache()
+            module.train()            

stable_audio_3/training/utils.py

@@ -0,0 +1,504 @@
+from pytorch_lightning.loggers import WandbLogger, CometLogger
+from ..interface.aeiou import pca_point_cloud
+
+import math
+import wandb
+import torch
+import torch.nn.functional as F
+import os
+import typing as tp
+
+def get_rank():
+    """Get rank of current process."""
+
+    if "SLURM_PROCID" in os.environ:
+        return int(os.environ["SLURM_PROCID"])
+
+    if not torch.distributed.is_available() or not torch.distributed.is_initialized():
+        return 0
+
+    return torch.distributed.get_rank()
+
+class InverseLR(torch.optim.lr_scheduler._LRScheduler):
+    """Implements an inverse decay learning rate schedule with an optional exponential
+    warmup. When last_epoch=-1, sets initial lr as lr.
+    inv_gamma is the number of steps/epochs required for the learning rate to decay to
+    (1 / 2)**power of its original value.
+    Args:
+        optimizer (Optimizer): Wrapped optimizer.
+        inv_gamma (float): Inverse multiplicative factor of learning rate decay. Default: 1.
+        power (float): Exponential factor of learning rate decay. Default: 1.
+        warmup (float): Exponential warmup factor (0 <= warmup < 1, 0 to disable)
+            Default: 0.
+        final_lr (float): The final learning rate. Default: 0.
+        last_epoch (int): The index of last epoch. Default: -1.
+    """
+
+    def __init__(self, optimizer, inv_gamma=1., power=1., warmup=0., final_lr=0.,
+                 last_epoch=-1):
+        self.inv_gamma = inv_gamma
+        self.power = power
+        if not 0. <= warmup < 1:
+            raise ValueError('Invalid value for warmup')
+        self.warmup = warmup
+        self.final_lr = final_lr
+        super().__init__(optimizer, last_epoch)
+
+    def get_lr(self):
+        if not self._get_lr_called_within_step:
+            import warnings
+            warnings.warn("To get the last learning rate computed by the scheduler, "
+                          "please use `get_last_lr()`.")
+
+        return self._get_closed_form_lr()
+
+    def _get_closed_form_lr(self):
+        warmup = 1 - self.warmup ** (self.last_epoch + 1)
+        lr_mult = (1 + self.last_epoch / self.inv_gamma) ** -self.power
+        return [warmup * max(self.final_lr, base_lr * lr_mult)
+                for base_lr in self.base_lrs]
+
+def create_optimizer_from_config(optimizer_config, parameters):
+    """Create optimizer from config.
+
+    Args:
+        parameters (iterable): parameters to optimize.
+        optimizer_config (dict): optimizer config.
+
+    Returns:
+        torch.optim.Optimizer: optimizer.
+    """
+
+    optimizer_type = optimizer_config["type"]
+
+    if optimizer_type == "FusedAdam":
+        from deepspeed.ops.adam import FusedAdam
+        optimizer = FusedAdam(parameters, **optimizer_config["config"])
+    elif optimizer_type == "CAdamW":
+        from stable_audio_tools.training.optims import CAdamW
+        optimizer = CAdamW(parameters, **optimizer_config["config"])
+    elif optimizer_type == "CLion":
+        from stable_audio_tools.training.optims import CLion
+        optimizer = CLion(parameters, **optimizer_config["config"])
+    elif optimizer_type == "AdamW8bit":
+        from bitsandbytes.optim import AdamW8bit
+        optimizer = AdamW8bit(parameters, **optimizer_config["config"])
+    elif optimizer_type == "MuonAdamW":
+        from stable_audio_tools.training.optims import MuonAdamW
+        optimizer = MuonAdamW(parameters, **optimizer_config["config"])
+    else:
+        optimizer_fn = getattr(torch.optim, optimizer_type)
+        optimizer = optimizer_fn(parameters, **optimizer_config["config"])
+    return optimizer
+
+def create_scheduler_from_config(scheduler_config, optimizer):
+    """Create scheduler from config.
+
+    Args:
+        scheduler_config (dict): scheduler config.
+        optimizer (torch.optim.Optimizer): optimizer.
+
+    Returns:
+        torch.optim.lr_scheduler._LRScheduler: scheduler.
+    """
+    if scheduler_config["type"] == "InverseLR":
+        scheduler_fn = InverseLR
+    else:
+        scheduler_fn = getattr(torch.optim.lr_scheduler, scheduler_config["type"])
+    scheduler = scheduler_fn(optimizer, **scheduler_config["config"])
+    return scheduler
+
+def logger_project_name(logger) -> str:
+    if isinstance(logger, WandbLogger):
+        return logger.experiment.project
+    elif isinstance(logger, CometLogger):
+        return logger.name
+
+def log_metric(logger, key, value, step=None):
+    from pytorch_lightning.loggers import WandbLogger, CometLogger
+    if isinstance(logger, WandbLogger):
+        logger.experiment.log({key: value})
+    elif isinstance(logger, CometLogger):
+        logger.experiment.log_metrics({key: value}, step=step)
+
+def log_audio(logger, key, audio_path, sample_rate, caption=None, step=None):
+    if isinstance(logger, WandbLogger):
+        logger.experiment.log({key: wandb.Audio(audio_path, sample_rate=sample_rate, caption=caption)})
+    elif isinstance(logger, CometLogger):
+        logger.experiment.log_audio(audio_path, file_name=key, sample_rate=sample_rate, step=step)
+
+def log_image(logger, key, img_data, step=None):
+    if isinstance(logger, WandbLogger):
+        logger.experiment.log({key: wandb.Image(img_data)})
+    elif isinstance(logger, CometLogger):
+        logger.experiment.log_image(img_data, name=key, step=step)
+
+def log_point_cloud(logger, key, tokens, caption=None):
+    if isinstance(logger, WandbLogger):
+        point_cloud = pca_point_cloud(tokens)
+        logger.experiment.log({key: point_cloud})
+    elif isinstance(logger, CometLogger):
+        point_cloud = pca_point_cloud(tokens, rgb_float=True, output_type="points")
+        #logger.experiment.log_points_3d(scene_name=key, points=point_cloud)
+
+
+def compute_per_elem_trim(conditioning, sample_rate, margin_seconds=5.0):
+    """Compute per-element trim lengths from seconds_total in conditioning dicts.
+
+    Returns a list of trim lengths (in audio samples) or None if no seconds_total found.
+    """
+    if not any("seconds_total" in c for c in conditioning):
+        return None
+    margin_samples = int(margin_seconds * sample_rate)
+    per_elem_trim = []
+    for c in conditioning:
+        if "seconds_total" in c:
+            per_elem_trim.append(int(c["seconds_total"] * sample_rate) + margin_samples)
+        else:
+            per_elem_trim.append(None)
+    return per_elem_trim
+
+def trim_and_concat(x, per_elem_trim):
+    """Per-element trim and concatenate along time axis.
+
+    Trims each batch element to its own length (from seconds_total + margin),
+    removing trailing padding silence before concatenation.
+
+    Args:
+        x: (b, d, n) tensor or list of (d, n) tensors
+        per_elem_trim: list of trim lengths per element, or None for no trimming
+    """
+    items = [x[i] for i in range(x.shape[0])] if isinstance(x, torch.Tensor) and x.dim() == 3 else x
+    if per_elem_trim is None:
+        return torch.cat(items, dim=-1)
+    parts = []
+    for i, elem in enumerate(items):
+        if per_elem_trim[i] is not None:
+            parts.append(elem[..., :min(per_elem_trim[i], elem.shape[-1])])
+        else:
+            parts.append(elem)
+    return torch.cat(parts, dim=-1)
+
+
+class StaggeredLogger:
+    """Accumulates log values over N steps and flushes averaged metrics.
+
+    Avoids per-step CUDA syncs (.item()) and prevents PL from caching
+    stale values between flush intervals by logging directly to the
+    experiment logger.
+
+    Usage in a PL module::
+
+        # In __init__:
+        self._staggered_logger = StaggeredLogger(every_n_steps=10)
+
+        # At end of training_step (pass detached tensors or Python scalars):
+        self._staggered_logger.log(log_dict, self)
+    """
+
+    def __init__(self, every_n_steps: int = 10):
+        self.every_n_steps = every_n_steps
+        self._accum: dict = {}
+        self._counts: dict = {}
+        self._steps: int = 0
+
+    def log(self, log_dict: dict, module) -> None:
+        """Accumulate *log_dict* and flush every *every_n_steps* calls.
+
+        NaN/inf values are handled at flush time: if an accumulated sum
+        is non-finite (due to a NaN step poisoning the window), that
+        metric is simply omitted from the flush.  This avoids the
+        expensive per-step CUDA syncs that an eager ``isfinite`` check
+        would require.
+
+        Args:
+            log_dict: Mapping of metric names to detached GPU tensors or
+                Python scalars.  Tensors are kept on-device until flush.
+            module: The ``pl.LightningModule`` instance (used for
+                ``log_dict``, ``logger``, and ``global_step``).
+        """
+        for k, v in log_dict.items():
+            val = v.detach() if torch.is_tensor(v) else v
+            if k not in self._accum:
+                self._accum[k] = val
+                self._counts[k] = 1
+            else:
+                self._accum[k] = self._accum[k] + val
+                self._counts[k] += 1
+        self._steps += 1
+
+        if self._steps >= self.every_n_steps:
+            flushed = {}
+            for k, v in self._accum.items():
+                c = self._counts[k]
+                if c > 0:
+                    val = v.item() / c if torch.is_tensor(v) else v / c
+                    # Skip non-finite values (NaN from a bad step poisoned this window)
+                    if math.isfinite(val):
+                        flushed[k] = val
+            # Update progress bar only (logger=False avoids PL caching stale values)
+            if flushed:
+                module.log_dict(flushed, prog_bar=True, on_step=True, logger=False)
+                # Log directly to experiment (bypasses PL metric caching)
+                if module.logger is not None:
+                    module.logger.log_metrics(flushed, step=module.global_step)
+            self._accum = {}
+            self._counts = {}
+            self._steps = 0
+
+
+def resize_padding_mask(padding_mask: torch.Tensor, target_length: int) -> torch.Tensor:
+    """Resize a padding mask to target_length using ceiling-based length scaling.
+
+    Unlike F.interpolate(mode="nearest"), this ensures any target position
+    that partially overlaps valid audio is marked valid (rounds up).
+
+    Args:
+        padding_mask: (B, T) or (T,) tensor where 1/True = valid, 0/False = padding
+        target_length: desired output length
+    Returns:
+        Resized boolean mask with same ndim as input
+    """
+    if padding_mask.ndim == 1:
+        valid_length = padding_mask.sum()
+        source_length = padding_mask.shape[0]
+        valid_target_length = torch.ceil(
+            valid_length.float() * target_length / source_length
+        ).long().clamp(max=target_length)
+        positions = torch.arange(target_length, device=padding_mask.device)
+        return positions < valid_target_length
+    else:
+        valid_lengths = padding_mask.sum(dim=-1)  # (B,)
+        source_length = padding_mask.shape[-1]
+        valid_target_lengths = torch.ceil(
+            valid_lengths.float() * target_length / source_length
+        ).long().clamp(max=target_length)
+        positions = torch.arange(target_length, device=padding_mask.device).unsqueeze(0)
+        return positions < valid_target_lengths.unsqueeze(1)
+
+
+def create_augmented_padding_mask(
+    padding_mask: torch.Tensor,
+    silence_extension_scale_seconds: float = 0.0,
+    sample_rate: int = 44100,
+    downsampling_ratio: int = 1,
+) -> torch.Tensor:
+    """
+    Augment padding mask by randomly extending the valid region to include silence.
+
+    This helps the model learn to handle silence at the end of audio by randomly
+    including some of the padded (silent) region as valid.
+
+    Uses an exponential distribution so most extensions are small (within ~scale seconds)
+    but can occasionally extend to the full sequence length.
+
+    Args:
+        padding_mask: Boolean tensor of shape (batch, seq_len) where True = valid audio
+        silence_extension_scale_seconds: Scale (mean) of exponential distribution in seconds.
+            ~63% of samples within 1x scale, ~86% within 2x, ~95% within 3x.
+        sample_rate: Audio sample rate (used to convert seconds to tokens)
+        downsampling_ratio: Downsampling ratio from pretransform (to convert samples to latent tokens)
+
+    Returns:
+        Augmented padding mask of shape (batch, seq_len) where True = valid region (audio + extended silence)
+    """
+    if silence_extension_scale_seconds <= 0:
+        return padding_mask
+
+    batch_size, seq_len = padding_mask.shape
+    device = padding_mask.device
+
+    # Find the last valid position for each batch element
+    valid_lengths = padding_mask.sum(dim=-1)  # (batch,)
+
+    # Calculate scale in tokens for exponential distribution
+    scale_tokens = silence_extension_scale_seconds * sample_rate / downsampling_ratio
+
+    # Sample from exponential distribution (biased toward smaller values)
+    # Exponential: ~63% within 1x scale, ~86% within 2x, ~95% within 3x
+    random_extensions = torch.empty(batch_size, device=device).exponential_(lambd=1.0/scale_tokens)
+    random_extensions = random_extensions.long()
+
+    # Extend valid lengths, clamping to sequence length (allows extension to full padding)
+    augmented_lengths = torch.clamp(valid_lengths + random_extensions, max=seq_len)
+
+    # Create augmented mask: True where position < augmented_length
+    positions = torch.arange(seq_len, device=device).unsqueeze(0)  # (1, seq_len)
+    augmented_mask = positions < augmented_lengths.unsqueeze(1)  # (batch, seq_len)
+
+    return augmented_mask
+
+
+def compute_normalized_mse(
+    pred: torch.Tensor,
+    target: torch.Tensor,
+    loss_mask: torch.Tensor,
+    loss_normalization: str = "none",
+    loss_norm_eps: float = 1e-6,
+) -> torch.Tensor:
+    """
+    Compute MSE normalized by detached target magnitude.
+
+    Prevents high-magnitude latents from dominating the loss,
+    ensuring quiet passages receive equal learning signal.
+
+    Normalization is only applied to valid (non-padded) regions.
+    Padded regions retain raw MSE for compatibility with mask_loss_weight.
+
+    Args:
+        pred: [B, C, T] predicted tensor
+        target: [B, C, T] target tensor
+        loss_mask: [B, T] boolean, True = valid (non-padded), or None
+        loss_normalization: One of "none", "timestep", "sample", "sample_channel"
+        loss_norm_eps: Epsilon for numerical stability
+
+    Returns:
+        Per-element MSE [B, C, T], normalized for signal, raw for padding
+    """
+    mse = (pred - target) ** 2  # [B, C, T]
+
+    if loss_normalization == "none":
+        return mse
+
+    if loss_mask is None:
+        mask_expanded = torch.ones(pred.shape[0], 1, pred.shape[2], device=pred.device, dtype=torch.bool)
+    else:
+        mask_expanded = loss_mask.unsqueeze(1)  # [B, 1, T]
+
+    with torch.no_grad():
+        if loss_normalization == "timestep":
+            # Var across channels at each timestep [B, 1, T]
+            mag_sq = torch.mean((target - torch.mean(target, dim=1, keepdim=True)) ** 2, dim=1, keepdim=True) + loss_norm_eps
+        else:
+            # For sample/sample_channel modes, exclude padding via NaN masking
+            masked_targets = torch.where(mask_expanded, target, float('nan'))
+
+            if loss_normalization == "sample":
+                sample_mean = torch.nanmean(masked_targets, dim=(1, 2), keepdim=True)
+                mag_sq = torch.nanmean((masked_targets - sample_mean) ** 2, dim=(1, 2), keepdim=True) + loss_norm_eps
+            elif loss_normalization == "sample_channel":
+                channel_mean = torch.nanmean(masked_targets, dim=2, keepdim=True)
+                mag_sq = torch.nanmean((masked_targets - channel_mean) ** 2, dim=2, keepdim=True) + loss_norm_eps
+            else:
+                raise ValueError(f"Unknown loss normalization mode: {loss_normalization}")
+
+            # Handle edge case where entire sample is padding (all NaN -> NaN variance)
+            mag_sq = torch.where(torch.isnan(mag_sq), torch.ones_like(mag_sq), mag_sq)
+
+    # Only normalize signal regions; keep padding MSE raw for mask_loss_weight compatibility
+    normalized_mse = mse / mag_sq
+    return torch.where(mask_expanded, normalized_mse, mse)
+
+
+def compute_masked_loss(
+    loss_full: torch.Tensor,
+    loss_mask: torch.Tensor,
+    mask_padding_attention: bool,
+    mask_loss_weight: float = 0.0,
+) -> tuple:
+    """
+    Compute loss with separate signal and padding contributions.
+
+    Args:
+        loss_full: Full loss tensor of shape (B, C, T)
+        loss_mask: Boolean mask of shape (B, T) where True = signal, False = padding
+        mask_padding_attention: If True, only compute loss on signal region
+        mask_loss_weight: Weight for padding loss when mask_padding_attention is False
+
+    Returns:
+        Tuple of (loss, signal_mean, padding_mean) where:
+            - loss: The final scalar loss
+            - signal_mean: Mean loss on signal region (for logging)
+            - padding_mean: Mean loss on padding region (for logging)
+    """
+    # Compute separate signal and padding losses
+    signal = torch.where(loss_mask.unsqueeze(1), loss_full, 0.0)
+    signal_sum = signal.sum(dim=(1, 2))
+    n_channels = loss_full.shape[1]
+    signal_count = loss_mask.sum(dim=1) * n_channels
+
+    padding = torch.where(~loss_mask.unsqueeze(1), loss_full, 0.0)
+    padding_sum = padding.sum(dim=(1, 2))
+    padding_count = (~loss_mask).sum(dim=1) * n_channels
+
+    if mask_padding_attention:
+        # When attention masking is on, only compute loss on signal region
+        # (padding saw no attention, so loss there is meaningless)
+        per_sample_loss = signal_sum / (signal_count + 1e-8)
+        loss = per_sample_loss.mean()
+    else:
+        # Weighted loss: signal + mask_loss_weight * padding
+        w = mask_loss_weight
+        denom = signal_count + w * padding_count + 1e-8
+        per_sample_loss = (signal_sum + w * padding_sum) / denom
+        loss = per_sample_loss.mean()
+
+    # Compute separate means for logging
+    signal_mean = signal.sum() / (signal_count.sum() + 1e-8)
+    padding_mean = padding.sum() / (padding_count.sum() + 1e-8)
+
+    return loss, signal_mean.detach(), padding_mean.detach()
+
+
+def masked_mean(
+    values: torch.Tensor,
+    mask: torch.Tensor,
+) -> torch.Tensor:
+    """
+    Compute mean over masked (valid) positions only.
+
+    Args:
+        values: Tensor of shape (B, C, T) or (B, T)
+        mask: Boolean mask of shape (B, T) where True = valid, False = padding
+
+    Returns:
+        Scalar mean over valid positions only
+    """
+    if values.ndim == 2:
+        # (B, T) -> add channel dim
+        values = values.unsqueeze(1)
+
+    # Expand mask to match values shape: (B, T) -> (B, 1, T)
+    mask_expanded = mask.unsqueeze(1)
+
+    # Zero out padding positions
+    masked_values = torch.where(mask_expanded, values, torch.zeros_like(values))
+
+    # Count valid positions (per sample, times channels)
+    n_channels = values.shape[1]
+    valid_count = mask.sum(dim=1) * n_channels  # (B,)
+
+    # Sum over all dimensions and divide by valid count
+    per_sample_sum = masked_values.sum(dim=(1, 2))  # (B,)
+    per_sample_mean = per_sample_sum / (valid_count + 1e-8)
+
+    return per_sample_mean.mean()
+
+
+def masked_sum(
+    values: torch.Tensor,
+    mask: torch.Tensor,
+    dim: tp.List[int],
+) -> torch.Tensor:
+    """
+    Compute sum over masked (valid) positions only.
+
+    Args:
+        values: Tensor of shape (B, C, T) or (B, T)
+        mask: Boolean mask of shape (B, T) where True = valid, False = padding
+        dim: Dimensions to sum over (should include the time dimension)
+
+    Returns:
+        Sum over valid positions, reduced over specified dims
+    """
+    if values.ndim == 2:
+        # (B, T) -> add channel dim for consistency
+        mask_expanded = mask
+    else:
+        # (B, C, T) -> expand mask
+        mask_expanded = mask.unsqueeze(1)
+
+    # Zero out padding positions
+    masked_values = torch.where(mask_expanded, values, torch.zeros_like(values))
+
+    return masked_values.sum(dim=tuple(dim))

stable_audio_3/verbose.py

@@ -0,0 +1,20 @@
+"""Shared verbose flag for user-facing prints.
+Defaults off. Toggle via `set_verbose(True)` (e.g. from run_gradio.py --verbose)
+or by setting the environment variable `SAT_VERBOSE=1`.
+Warnings and errors should continue to use plain `print()` — `vprint()` is only
+for informational/progress messages that are noisy at default verbosity.
+"""
+
+import os
+
+VERBOSE = os.environ.get("SAT_VERBOSE", "0") == "1"
+
+
+def set_verbose(v: bool):
+    global VERBOSE
+    VERBOSE = bool(v)
+
+
+def vprint(*args, **kwargs):
+    if VERBOSE:
+        print(*args, **kwargs)