feat(assets): SA3 音效补齐 radio 11+1 与 dungeon 13源+15pan（token 走 hf-mirror 实测可用）

- 模型 17/17 文件经镜像全量缓存（gated+token，直连本机不通；sfx.md/setup.md/asset-pipeline.md 镜像结论同步更新）
- 修复迁移遗留：stable_audio_3/models/ 源码子包被 .gitignore models/ 规则误伤丢失，从上游恢复并加白名单
- radio：arg_muffled/brake_screech 等叙事件落地（接线原有 L111/L154 零改动）+ phone_ring SYN→SA3 升级；dungeon 13 源 A/B 对齐（RMS 偏差 ≤0.3dB）全量同名替换，SYN 原版备份 tools/asset_fix_backup/sa3_ab_20260902/
- 关键玩法语义达标：echo 距离档 0.83<1.43<2.18<2.98s 严格递增；one-shot 真峰值全 ≤-3dBFS（解码回读法）
- check_all_assets 8 游戏全绿（radio 254/254、dungeon 303/303）；playtest radio 221/dungeon 1610 全绿

tools/stable-audio-3/.gitignore

@@ -174,3 +174,8 @@ optimized/tflite/models/tflite/lora_cache/
 # Rendered audio from the CLI / gradio backends. *.wav above already covers most of it,
 # but the --file-format flag also writes .flac and .mp3.
 optimized/*/output/
+
+# 误伤修正（2026-09-02）：上面的 models/ 规则会忽略源码子包 stable_audio_3/models/
+# （factory.py/loading_utils.py 的依赖，迁移机器时曾因未入库而丢失，从上游 GitHub 重取）。
+!stable_audio_3/models/
+!stable_audio_3/models/**

tools/stable-audio-3/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
+        

tools/stable-audio-3/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)
+
+

tools/stable-audio-3/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

tools/stable-audio-3/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
+    

tools/stable-audio-3/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)
+

tools/stable-audio-3/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.float(), v1.float()
+
+        if padding_mask is not None:
+            # Expand mask to match tensor shape: (B, T) -> (B, 1, T)
+            mask = padding_mask.unsqueeze(1).float()
+            # 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
+            )

tools/stable-audio-3/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

tools/stable-audio-3/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

tools/stable-audio-3/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

tools/stable-audio-3/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 list(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

tools/stable-audio-3/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())

tools/stable-audio-3/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

tools/stable-audio-3/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