Duplicate from Localsong/Localsong

Co-authored-by: Localsong <Localsong@users.noreply.huggingface.co>

.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

LICENSE

@@ -0,0 +1,58 @@
+MIT License
+
+Copyright (c) 2026
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+
+-------------------------------------------------------------------------------
+
+Third-party code
+
+same_l_decoder.py is adapted from Stability AI's stable-audio-3
+(https://github.com/Stability-AI/stable-audio-3), used under the MIT License:
+
+    MIT License
+
+    Copyright (c) 2026 Stability AI
+
+    Permission is hereby granted, free of charge, to any person obtaining a copy
+    of this software and associated documentation files (the "Software"), to deal
+    in the Software without restriction, including without limitation the rights
+    to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+    copies of the Software, and to permit persons to whom the Software is
+    furnished to do so, subject to the following conditions:
+
+    The above copyright notice and this permission notice shall be included in all
+    copies or substantial portions of the Software.
+
+    THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+    IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+    FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+    AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+    LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+    OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+    SOFTWARE.
+
+Third-party weights
+
+The SAME-L autoencoder weights that same_l_decoder.py loads are released
+separately by Stability AI under the Stable Audio Community License
+(https://stability.ai/license).
+
+This Stability AI Model is licensed under the Stability AI Community License, Copyright ©  Stability AI Ltd. All Rights Reserved

README.md

@@ -0,0 +1,64 @@
+---
+license: mit
+---
+# Audio DiT
+
+Tag-conditioned music generation focused on instrumental game and anime songs. Choose 1-8 tags (recommend to start with 1). Generates 95s songs.
+Set a seed to reproduce a generation, or leave it at `-1` for a random seed.
+
+## Run it
+
+Needs [uv](https://docs.astral.sh/uv/). CUDA, MPS or CPU.
+
+```bash
+hf download Localsong/Localsong --local-dir Localsong
+uv run webui.py  # http://127.0.0.1:7860
+```
+
+## Architecture
+
+```
+tags ──► AudioDiT ──► 256×1024 latents ──► SAME-L decoder ──► 95 s of stereo audio
+         0.84B params            
+```
+
+**Latents.** [SAME-L](https://huggingface.co/stabilityai/SAME-L) autoencoder by Stability AI.
+
+**Denoiser.** `audio_dit.py` is a 1-D latent transformer DiT. 
+RMSNorm and QK-RMSNorm, RoPE, SwiGLU
+Timestep conditioning is adaLN-single (PixArt)
+offset. Tags condition via cross-attention.
+
+**Sparse-dense fusion** ([SPRINT](https://arxiv.org/abs/2510.21986)). The blocks
+are split into a dense 2-block encoder, a 20-block middle stack, and a dense
+2-block decoder. The encoder output is concatenated with the middle output and
+projected back down before the decoder:
+
+```
+frames ──► encoder ──┬─────────────────────────► f ──┐
+                     └─► middle (20 blocks) ──► g ──► linear([f;g]) ──► decoder ──► velocity
+```
+
+Training ran the middle stack on a random 25 % of frames and sometimes dropped
+it entirely, which makes the shallow encoder→fusion→decoder route a usable weak
+model on its own. Inference runs the middle stack on every frame
+(`AudioDiT.forward`) and keeps the shallow route as the guidance branch
+(`AudioDiT.shallow`).
+
+**Sampling.** Rectified flow: v-prediction, 50 steps Euler shift=2
+Guidance alternates between normal CFG and the shallow-path PDG.
+
+## License
+
+MIT, see [LICENSE](LICENSE).
+
+`same_l_decoder.py` is adapted from Stability AI's
+[stable-audio-3](https://github.com/Stability-AI/stable-audio-3) (MIT). The
+SAME-L weights it loads are released by Stability AI under the Stable Audio
+Community License; commercial use of them is subject to
+<https://stability.ai/license>.
+
+The model also builds on SPRINT ([2510.21986](https://arxiv.org/abs/2510.21986)),
+PixArt-α ([2310.00426](https://arxiv.org/abs/2310.00426)),
+DiT ([2212.09748](https://arxiv.org/abs/2212.09748)) and
+rectified flow ([2209.03003](https://arxiv.org/abs/2209.03003)).

audio_dit.py

@@ -0,0 +1,216 @@
+"""Tag-conditioned diffusion transformer over SAME-L music latents (0.84B params)."""
+
+from __future__ import annotations
+
+import math
+from pathlib import Path
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from safetensors.torch import load_file
+
+WEIGHTS = Path(__file__).parent / "audio_dit.safetensors"
+
+CHANNELS = 256  # SAME-L latent channels
+FRAMES = 1024  # latent frames per sample, at 10.77 frames/second
+DIM = 1536
+HEADS = 24
+HEAD_DIM = 64
+MLP_HIDDEN = 4096
+NUM_TAGS = 1470  # vocabulary rows; index 0 is padding
+MAX_TAGS = 8
+ROPE_THETA = 10000.0
+
+# SPRINT sparse-dense fusion: a dense encoder, a deep middle stack, a dense
+# decoder. Encoder and decoder blocks cross-attend to the tags, middle blocks
+# every second block. Training subsampled the middle stack; inference runs it on
+# every frame.
+ENCODER_BLOCKS = 2
+MIDDLE_BLOCKS = 20
+DECODER_BLOCKS = 2
+
+
+def rms_norm(x: torch.Tensor) -> torch.Tensor:
+    return F.rms_norm(x.float(), (x.shape[-1],), eps=1e-6).to(x.dtype)
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return rms_norm(x) * self.weight
+
+
+def rope_tables() -> tuple[torch.Tensor, torch.Tensor]:
+    freqs = torch.outer(
+        torch.arange(FRAMES).float(),
+        1.0 / (ROPE_THETA ** (torch.arange(0, HEAD_DIM, 2).float() / HEAD_DIM)),
+    )
+    return freqs.cos()[None, None], freqs.sin()[None, None]
+
+
+def apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+    x1, x2 = x.float().chunk(2, dim=-1)
+    return torch.cat([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1).to(x.dtype)
+
+
+class SelfAttention(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.qkv_x = nn.Linear(DIM, 3 * DIM, bias=False)
+        self.q_norm = RMSNorm(HEAD_DIM)
+        self.k_norm = RMSNorm(HEAD_DIM)
+        self.proj = nn.Linear(DIM, DIM, bias=False)
+
+    def forward(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+        B, S, _ = x.shape
+        q, k, v = self.qkv_x(x).view(B, S, 3, HEADS, HEAD_DIM).permute(2, 0, 3, 1, 4)
+        q = apply_rope(self.q_norm(q), cos, sin)
+        k = apply_rope(self.k_norm(k), cos, sin)
+        out = F.scaled_dot_product_attention(q, k, v)
+        return self.proj(out.transpose(1, 2).reshape(B, S, DIM))
+
+
+class CrossAttention(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.q_y = nn.Linear(DIM, DIM, bias=False)
+        self.kv_y = nn.Linear(DIM, 2 * DIM, bias=False)
+        self.q_norm = RMSNorm(HEAD_DIM)
+        self.k_norm = RMSNorm(HEAD_DIM)
+        self.proj_y = nn.Linear(DIM, DIM, bias=False)
+
+    def forward(self, x: torch.Tensor, y: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
+        B, S, _ = x.shape
+        q = self.q_norm(self.q_y(x).view(B, S, HEADS, HEAD_DIM).transpose(1, 2))
+        k, v = self.kv_y(y).view(B, MAX_TAGS, 2, HEADS, HEAD_DIM).permute(2, 0, 3, 1, 4)
+        out = F.scaled_dot_product_attention(q, self.k_norm(k), v,
+                                             attn_mask=mask[:, None, None, :])
+        return self.proj_y(out.transpose(1, 2).reshape(B, S, DIM))
+
+
+class SwiGLU(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.w12 = nn.Linear(DIM, 2 * MLP_HIDDEN, bias=False)
+        self.w3 = nn.Linear(MLP_HIDDEN, DIM, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        gate, up = self.w12(x).chunk(2, dim=-1)
+        return self.w3(F.silu(gate) * up)
+
+
+class Block(nn.Module):
+    """adaLN-single: shared timestep modulation plus a learned per-block offset."""
+
+    def __init__(self):
+        super().__init__()
+        self.attn = SelfAttention()
+        self.mlp = SwiGLU()
+        self.adaLN_offset = nn.Parameter(torch.zeros(6 * DIM))
+
+    def forward(self, x, mod, y, mask, cos, sin):
+        shift1, scale1, gate1, shift2, scale2, gate2 = (mod + self.adaLN_offset).chunk(6, -1)
+        x = x + gate1 * self.attn(rms_norm(x) * (1 + scale1) + shift1, cos, sin)
+        return x + gate2 * self.mlp(rms_norm(x) * (1 + scale2) + shift2)
+
+
+class CrossBlock(Block):
+    def __init__(self):
+        super().__init__()
+        self.cross = CrossAttention()
+        self.cross_gate = nn.Parameter(torch.zeros(DIM))
+
+    def forward(self, x, mod, y, mask, cos, sin):
+        shift1, scale1, gate1, shift2, scale2, gate2 = (mod + self.adaLN_offset).chunk(6, -1)
+        x = x + gate1 * self.attn(rms_norm(x) * (1 + scale1) + shift1, cos, sin)
+        x = x + self.cross_gate * self.cross(rms_norm(x), y, mask)
+        return x + gate2 * self.mlp(rms_norm(x) * (1 + scale2) + shift2)
+
+
+class Timestep(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.mlp = nn.Sequential(nn.Linear(256, DIM), nn.SiLU(), nn.Linear(DIM, DIM))
+
+    def forward(self, t: float, batch: int, device) -> torch.Tensor:
+        freqs = torch.exp(-math.log(10000.0) * torch.arange(128, device=device) / 128)
+        args = t * 1000.0 * freqs[None]
+        return self.mlp(torch.cat([args.cos(), args.sin()], dim=-1)).expand(batch, DIM)
+
+
+class FinalLayer(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.adaLN_modulation = nn.Linear(DIM, 2 * DIM)
+        self.linear = nn.Linear(DIM, CHANNELS)
+
+    def forward(self, x: torch.Tensor, c: torch.Tensor) -> torch.Tensor:
+        shift, scale = self.adaLN_modulation(F.silu(c)).unsqueeze(1).chunk(2, dim=-1)
+        return self.linear(rms_norm(x) * (1 + scale) + shift)
+
+
+class AudioDiT(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.x_embedder = nn.Linear(CHANNELS, DIM)
+        self.t_embedder = Timestep()
+        self.tag_embedding = nn.Embedding(NUM_TAGS, DIM, padding_idx=0)
+        self.tag_norm = RMSNorm(DIM)
+        self.null_tag = nn.Parameter(torch.zeros(DIM))
+        self.mask_token = nn.Parameter(torch.zeros(DIM))
+        self.fusion_proj = nn.Linear(2 * DIM, DIM)
+        self.adaLN_shared = nn.Linear(DIM, 6 * DIM)
+        self.encoder = nn.ModuleList(CrossBlock() for _ in range(ENCODER_BLOCKS))
+        self.middle = nn.ModuleList(
+            (CrossBlock if i % 2 == 0 else Block)() for i in range(MIDDLE_BLOCKS))
+        self.decoder = nn.ModuleList(CrossBlock() for _ in range(DECODER_BLOCKS))
+        self.final_layer = FinalLayer()
+        cos, sin = rope_tables()
+        self.register_buffer("cos", cos, persistent=False)
+        self.register_buffer("sin", sin, persistent=False)
+
+    def _condition(self, t: float, tags: torch.Tensor):
+        y = self.tag_norm(self.tag_embedding(tags))
+        valid = tags != 0
+        # An all-padding row is unconditional: it attends to the learned null tag.
+        empty = ~valid.any(dim=1)
+        y = torch.where(empty[:, None, None], self.tag_norm(self.null_tag).expand_as(y), y)
+        mask = valid.clone()
+        mask[:, 0] |= empty
+        c = self.t_embedder(t, tags.shape[0], tags.device)
+        return c, self.adaLN_shared(F.silu(c)).unsqueeze(1), y, mask
+
+    def _run(self, blocks, x, mod, y, mask):
+        for block in blocks:
+            x = block(x, mod, y, mask, self.cos, self.sin)
+        return x
+
+    def _fuse(self, f, g, c, mod, y, mask):
+        h = self.fusion_proj(torch.cat([f, g], dim=-1))
+        return self.final_layer(self._run(self.decoder, h, mod, y, mask), c).transpose(1, 2)
+
+    def forward(self, x: torch.Tensor, t: float, tags: torch.Tensor) -> torch.Tensor:
+        """Velocity at (x, t), for latents (B, 256, 1024) and tag indices (B, 8)."""
+        c, mod, y, mask = self._condition(t, tags)
+        f = self._run(self.encoder, self.x_embedder(x.transpose(1, 2)), mod, y, mask)
+        return self._fuse(f, self._run(self.middle, f, mod, y, mask), c, mod, y, mask)
+
+    def shallow(self, x: torch.Tensor, t: float, tags: torch.Tensor) -> torch.Tensor:
+        """Path-drop branch: the middle stack is replaced by [MASK] tokens."""
+        c, mod, y, mask = self._condition(t, tags)
+        f = self._run(self.encoder, self.x_embedder(x.transpose(1, 2)), mod, y, mask)
+        return self._fuse(f, self.mask_token.expand_as(f), c, mod, y, mask)
+
+
+def load_model(device: str = "cuda") -> tuple[AudioDiT, torch.Tensor, torch.Tensor]:
+    """The model plus the per-channel mean/std the training latents were scaled by."""
+    state = load_file(WEIGHTS, device=device)
+    mean = state.pop("latent_mean").view(1, CHANNELS, 1)
+    std = state.pop("latent_std").view(1, CHANNELS, 1)
+    model = AudioDiT()
+    model.load_state_dict(state)
+    return model.to(device).eval().requires_grad_(False), mean, std

audio_dit.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:68d4d5c1ccb9501905a32f9d6364460eb7de4555af7500354e9ce98414915197
+size 3365063088

pyproject.toml

@@ -0,0 +1,22 @@
+[project]
+name = "audio-dit"
+version = "0.1.0"
+description = "Instrumental music generation."
+readme = "README.md"
+requires-python = ">=3.12"
+dependencies = [
+    "gradio>=6.2.0",
+    "numpy>=2.0",
+    "safetensors>=0.7.0",
+    "torch>=2.9.1",
+]
+
+[tool.uv.sources]
+torch = [
+  { index = "pytorch-cu130", marker = "sys_platform == 'linux' or sys_platform == 'win32'" },
+]
+
+[[tool.uv.index]]
+name = "pytorch-cu130"
+url = "https://download.pytorch.org/whl/cu130"
+explicit = true

same_l_decoder.py

@@ -0,0 +1,198 @@
+# Adapted from Stability AI's stable-audio-3 (MIT License). See LICENSE.
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from safetensors.torch import load_file
+
+WEIGHTS = Path(__file__).parent / "samel" / "model.safetensors"
+
+SAMPLE_RATE = 44100
+LATENT_DIM = 256
+DIM = 1536
+DEPTH = 12
+HEADS = 24
+HEAD_DIM = 64
+ROPE_DIM = 32  
+FF_HIDDEN = 3 * DIM
+SINUSOIDAL_FROM = 5 
+PATCH = 256  
+STRIDE = 16
+GROUP = STRIDE + 1  
+WINDOW = GROUP  
+
+CHUNK = 128 
+OVERLAP = 32 
+SEQ = CHUNK * GROUP
+SAMPLES_PER_FRAME = PATCH * STRIDE
+MASK_NOISE = 0.1
+LATENT_NOISE = 1e-3
+
+
+class DynamicTanh(nn.Module):
+    def __init__(self, dim: int):
+        super().__init__()
+        self.alpha = nn.Parameter(torch.ones(1))
+        self.gamma = nn.Parameter(torch.ones(dim))
+        self.beta = nn.Parameter(torch.zeros(dim))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.gamma * F.tanh(self.alpha * x) + self.beta
+
+
+def apply_rope(t: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor:
+    rot, keep = t[..., :ROPE_DIM], t[..., ROPE_DIM:]
+    x1, x2 = rot.chunk(2, dim=-1)
+    rotated = torch.cat([-x2, x1], dim=-1)
+    return torch.cat([rot * freqs.cos() + rotated * freqs.sin(), keep], dim=-1)
+
+
+class Attention(nn.Module):
+    """Differential attention: two attention maps per head, subtracted."""
+
+    def __init__(self):
+        super().__init__()
+        self.to_qkv = nn.Linear(DIM, 5 * DIM, bias=False)
+        self.to_out = nn.Linear(DIM, DIM, bias=False)
+        self.q_norm = DynamicTanh(HEAD_DIM)
+        self.k_norm = DynamicTanh(HEAD_DIM)
+
+    def forward(self, x: torch.Tensor, freqs: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
+        B, N, _ = x.shape
+        q, k, v, q_diff, k_diff = (
+            self.to_qkv(x).view(B, N, 5, HEADS, HEAD_DIM).permute(2, 0, 3, 1, 4))
+        q, q_diff = apply_rope(self.q_norm(q).float(), freqs), apply_rope(
+            self.q_norm(q_diff).float(), freqs)
+        k, k_diff = apply_rope(self.k_norm(k).float(), freqs), apply_rope(
+            self.k_norm(k_diff).float(), freqs)
+        out = F.scaled_dot_product_attention(q.to(v.dtype), k.to(v.dtype), v, attn_mask=mask)
+        out = out - F.scaled_dot_product_attention(
+            q_diff.to(v.dtype), k_diff.to(v.dtype), v, attn_mask=mask)
+        return self.to_out(out.transpose(1, 2).reshape(B, N, DIM))
+
+
+class Sin(nn.Module):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return torch.sin(torch.pi * x)
+
+
+class GatedProjection(nn.Module):
+    def __init__(self, activation: nn.Module):
+        super().__init__()
+        self.proj = nn.Linear(DIM, 2 * FF_HIDDEN)
+        self.act = activation
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        value, gate = self.proj(x).chunk(2, dim=-1)
+        return value * self.act(gate)
+
+
+class FeedForward(nn.Module):
+    def __init__(self, activation: nn.Module):
+        super().__init__()
+
+        self.ff = nn.Sequential(
+            GatedProjection(activation), nn.Identity(), nn.Linear(FF_HIDDEN, DIM))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.ff(x)
+
+
+class Block(nn.Module):
+    def __init__(self, sinusoidal: bool):
+        super().__init__()
+        self.pre_norm = DynamicTanh(DIM)
+        self.self_attn = Attention()
+        self.ff_norm = DynamicTanh(DIM)
+        self.ff = FeedForward(Sin() if sinusoidal else nn.SiLU())
+
+    def forward(self, x: torch.Tensor, freqs: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
+        x = x + self.self_attn(self.pre_norm(x), freqs, mask)
+        return x + self.ff(self.ff_norm(x))
+
+
+class Resampler(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.new_tokens = nn.Parameter(torch.zeros(1, 1, DIM))
+        self.blocks = nn.ModuleList(Block(i >= SINUSOIDAL_FROM) for i in range(DEPTH))
+        self.mapping = nn.Conv1d(DIM, 2 * PATCH, 1)
+
+    def forward(self, x: torch.Tensor, freqs: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
+        B, T, _ = x.shape
+        slots = self.new_tokens.expand(B * T, STRIDE, DIM)
+        slots = slots + torch.randn_like(slots) * MASK_NOISE
+        x = torch.cat([x.reshape(B * T, 1, DIM), slots], dim=1).reshape(B, T * GROUP, DIM)
+        for block in self.blocks:
+            x = block(x, freqs, mask)
+        x = x.reshape(B * T, GROUP, DIM)[:, 1:]  
+        return self.mapping(x.reshape(B, T * STRIDE, DIM).transpose(1, 2))
+
+
+class SameLDecoder(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.proj_in = nn.Linear(LATENT_DIM, DIM)
+        self.resampler = Resampler()
+        self.register_buffer("running_std", torch.ones(1))
+        position = torch.arange(SEQ)
+        self.register_buffer(
+            "mask", (position[None, :] - position[:, None]).abs() <= WINDOW, persistent=False)
+        self.register_buffer(
+            "inv_freq", 1.0 / (10000.0 ** (torch.arange(0, ROPE_DIM, 2) / ROPE_DIM)),
+            persistent=False)
+
+    def _rope_freqs(self) -> torch.Tensor:
+        freqs = torch.outer(torch.arange(SEQ, device=self.inv_freq.device).float(),
+                            self.inv_freq.float())
+        return torch.cat([freqs, freqs], dim=-1)
+
+    def _decode_chunk(self, latents: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor:
+        x = latents * self.running_std
+        x = x + torch.randn_like(x) * self.running_std * LATENT_NOISE
+        x = self.resampler(self.proj_in(x.transpose(1, 2)), freqs, self.mask)
+        B, _, L = x.shape  
+        return x.view(B, 2, PATCH, L).permute(0, 1, 3, 2).reshape(B, 2, L * PATCH)
+
+    @torch.no_grad()
+    def decode(self, latents: torch.Tensor) -> torch.Tensor:
+        """(B, 256, T) latents, T >= CHUNK, to (B, 2, T * 4096) audio."""
+        frames = latents.shape[-1]
+        starts = list(range(0, frames - CHUNK + 1, CHUNK - OVERLAP))
+        if starts[-1] != frames - CHUNK:
+            starts.append(frames - CHUNK)  
+        freqs = self._rope_freqs()
+        edge = OVERLAP // 2 * SAMPLES_PER_FRAME  
+        audio = latents.new_zeros(latents.shape[0], 2, frames * SAMPLES_PER_FRAME)
+        for i, start in enumerate(starts):
+            chunk = self._decode_chunk(latents[..., start:start + CHUNK], freqs)
+            left = 0 if i == 0 else edge
+            right = chunk.shape[-1] if i == len(starts) - 1 else chunk.shape[-1] - edge
+            at = start * SAMPLES_PER_FRAME
+            audio[..., at + left:at + right] = chunk[..., left:right]
+        return audio
+
+
+def load_decoder(device: str = "cuda",
+                 dtype: torch.dtype = torch.float16) -> SameLDecoder:
+    """Load the bundled SAME-L weights, keeping only what decoding needs."""
+    raw = load_file(WEIGHTS)
+    state = {"running_std": raw["bottleneck.running_std"]}
+
+    g, v = raw["decoder.layers.3.mapping.weight_g"], raw["decoder.layers.3.mapping.weight_v"]
+    state["resampler.mapping.weight"] = g * v / v.norm(dim=(1, 2), keepdim=True)
+    for key, tensor in raw.items():
+        if key.startswith("decoder.layers.1."):
+            state[key.replace("decoder.layers.1.", "proj_in.")] = tensor
+        elif key.startswith("decoder.layers.3.") and not key.endswith(
+                ("rope.inv_freq", "mapping.weight_g", "mapping.weight_v")):
+            state[key.replace("decoder.layers.3.", "resampler.")
+                     .replace("transformers.", "blocks.")] = tensor
+
+    decoder = SameLDecoder()
+    decoder.load_state_dict(state)
+    return decoder.to(device=device, dtype=dtype).eval().requires_grad_(False)

samel/.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

samel/LICENSE.md

@@ -0,0 +1,58 @@
+STABILITY AI COMMUNITY LICENSE AGREEMENT
+	
+Last Updated: July 5, 2024
+
+1. INTRODUCTION
+
+This Agreement applies to any individual person or entity (“You”, “Your” or “Licensee”) that uses or distributes any portion or element of the Stability AI Materials  or Derivative Works thereof for any Research & Non-Commercial or Commercial purpose. Capitalized terms not otherwise defined herein are defined in Section V below.
+
+This Agreement is intended to allow research, non-commercial, and limited commercial uses of the Models free of charge. In order to ensure that certain limited commercial uses of the Models continue to be allowed, this Agreement  preserves free access to the Models for people or organizations  generating annual revenue of less than US $1,000,000 (or local currency equivalent).  
+
+By clicking “I Accept”  or by using or distributing or using any portion or element of the Stability Materials or Derivative Works, You agree that You have read, understood and are bound by the terms of this Agreement. If You are acting on behalf of a company, organization or other entity, then “You” includes you and that entity, and You agree that You: (i) are an authorized representative of such entity with the authority to bind such entity to this Agreement, and (ii) You agree to the terms of this Agreement on that entity’s behalf. 
+
+2. RESEARCH & NON-COMMERCIAL USE LICENSE
+
+Subject to the terms of this Agreement, Stability AI grants You a non-exclusive, worldwide, non-transferable, non-sublicensable, revocable and royalty-free limited license under Stability AI’s intellectual property or other rights owned by Stability AI embodied in the Stability AI Materials to use, reproduce, distribute, and create Derivative Works of, and make modifications to, the Stability AI Materials for any Research or Non-Commercial Purpose. “Research Purpose” means academic or scientific advancement, and in each case, is not primarily intended for commercial advantage or monetary compensation to You or others. “Non-Commercial Purpose” means any purpose other than a Research Purpose that is not primarily intended for commercial advantage or monetary compensation to You or others, such as personal use (i.e., hobbyist) or evaluation and testing.  
+
+3. COMMERCIAL USE LICENSE
+
+Subject to the terms of this Agreement (including the remainder of this Section III), Stability AI grants You a non-exclusive, worldwide, non-transferable, non-sublicensable, revocable and royalty-free limited license under Stability AI’s intellectual property or other rights owned by Stability AI embodied in the Stability AI Materials to use, reproduce, distribute, and create Derivative Works of, and make modifications to, the Stability AI Materials for any Commercial Purpose. “Commercial Purpose” means any purpose other than a Research Purpose or Non-Commercial Purpose that is primarily intended for commercial advantage or monetary compensation to You or others, including but not limited to, (i) creating, modifying, or distributing Your product or service, including via a hosted service or application programming interface, and (ii) for Your business’s or organization’s internal operations. 
+If You are using or distributing the Stability AI Materials for a Commercial Purpose, You must register with Stability AI at (https://stability.ai/community-license). If at any time You or Your Affiliate(s), either individually or in aggregate, generate more than USD $1,000,000 in annual revenue (or the equivalent thereof in Your local currency), regardless of whether that revenue is generated directly or indirectly from the Stability AI Materials or Derivative Works, any licenses granted to You under this Agreement shall terminate as of such date. You must request a license from Stability AI at (https://stability.ai/enterprise) , which Stability AI may grant to You in its sole discretion. If you receive Stability AI Materials, or any Derivative Works thereof, from a Licensee as part of an integrated end user product, then Section III of this Agreement will not apply to you.
+
+4. GENERAL TERMS
+
+Your Research, Non-Commercial, and Commercial License(s) under this Agreement are subject to the following terms.
+a.  Distribution & Attribution. If You distribute or make available the Stability AI Materials or a Derivative Work to a third party, or a product or service that uses any portion of them, You shall: (i) provide a copy of this Agreement to that third party, (ii) retain the following attribution notice within a "Notice" text file distributed as a part of such copies: "This Stability AI Model is licensed under the Stability AI Community License, Copyright ©  Stability AI Ltd. All Rights Reserved”, and (iii) prominently display “Powered by Stability AI” on a related website, user interface, blogpost, about page, or product documentation.  If You create a Derivative Work, You may add your own attribution notice(s) to the “Notice” text file included with that Derivative Work, provided that You clearly indicate which attributions apply to the Stability AI Materials and state in the “Notice” text file that You changed the Stability AI Materials and how it was modified.
+b.  Use Restrictions. Your use of the Stability AI Materials and Derivative Works, including any output or results of the Stability AI Materials or Derivative Works, must comply with applicable laws and regulations (including Trade Control Laws and equivalent regulations) and adhere to the Documentation and Stability AI’s AUP, which is hereby incorporated by reference. Furthermore, You will not use the Stability AI Materials or Derivative Works, or any output or results of the Stability AI Materials or Derivative Works, to create or improve any foundational generative AI model (excluding the Models or Derivative Works). 
+c.  Intellectual Property. 
+(i) Trademark License.  No trademark licenses are granted under this Agreement, and in connection with the Stability AI Materials or Derivative Works, You may not use any name or mark owned by or associated with Stability AI or any of its Affiliates, except as required under Section IV(a) herein.
+(ii)  Ownership of Derivative Works.  As between You and Stability AI, You are the owner of Derivative Works You create, subject to Stability AI’s ownership of the Stability AI Materials and any Derivative Works made by or for Stability AI.
+(iii)  Ownership of Outputs. As between You and Stability AI, You own any outputs generated from the Models or Derivative Works to the extent permitted by applicable law. 
+(iv)  Disputes.  If You or Your Affiliate(s) institute litigation or other proceedings against Stability AI (including a cross-claim or counterclaim in a lawsuit) alleging that the Stability AI Materials, Derivative Works or associated outputs or results, or any portion of any of the foregoing, constitutes infringement of intellectual property or other rights owned or licensable by You, then any licenses granted to You under this Agreement shall terminate as of the date such litigation or claim is filed or instituted. You will indemnify and hold harmless Stability AI from and against any claim by any third party arising out of or related to Your use or distribution of the Stability AI Materials or Derivative Works in violation of this Agreement.
+(v)  Feedback.  From time to time, You may provide Stability AI with verbal and/or written suggestions, comments or other feedback related to Stability AI’s existing or prospective technology, products or services (collectively, “Feedback”). You are not obligated to provide Stability AI with Feedback, but to the extent that You do, You hereby grant Stability AI a perpetual, irrevocable, royalty-free, fully-paid, sub-licensable, transferable, non-exclusive, worldwide right and license to exploit the Feedback in any manner without restriction. Your Feedback is provided “AS IS” and You make no warranties whatsoever about any Feedback.
+d.  Disclaimer Of Warranty. UNLESS REQUIRED BY APPLICABLE LAW, THE STABILITY AI MATERIALS AND ANY OUTPUT AND RESULTS THEREFROM ARE PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING, WITHOUT LIMITATION, ANY WARRANTIES OF TITLE, NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. YOU ARE SOLELY RESPONSIBLE FOR DETERMINING THE APPROPRIATENESS OR LAWFULNESS OF USING OR REDISTRIBUTING THE STABILITY AI MATERIALS, DERIVATIVE WORKS OR ANY OUTPUT OR RESULTS AND ASSUME ANY RISKS ASSOCIATED WITH YOUR USE OF THE STABILITY AI MATERIALS, DERIVATIVE WORKS AND ANY OUTPUT AND RESULTS.
+e.  Limitation Of Liability. IN NO EVENT WILL STABILITY AI OR ITS AFFILIATES BE LIABLE UNDER ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, TORT, NEGLIGENCE, PRODUCTS LIABILITY, OR OTHERWISE, ARISING OUT OF THIS AGREEMENT, FOR ANY LOST PROFITS OR ANY DIRECT, INDIRECT, SPECIAL, CONSEQUENTIAL, INCIDENTAL, EXEMPLARY OR PUNITIVE DAMAGES, EVEN IF STABILITY AI OR ITS AFFILIATES HAVE BEEN ADVISED OF THE POSSIBILITY OF ANY OF THE FOREGOING.
+f.  Term And Termination. The term of this Agreement will commence upon Your acceptance of this Agreement or access to the Stability AI Materials and will continue in full force and effect until terminated in accordance with the terms and conditions herein. Stability AI may terminate this Agreement if You are in breach of any term or condition of this Agreement. Upon termination of this Agreement, You shall delete and cease use of any Stability AI Materials or Derivative Works. Section IV(d), (e), and (g) shall survive the termination of this Agreement.
+g.  Governing Law.  This Agreement will be governed by and constructed in accordance with the laws of the United States and the State of California without regard to choice of law principles, and the UN Convention on Contracts for International Sale of Goods does not apply to this Agreement.
+
+5. DEFINITIONS 
+
+“Affiliate(s)” means any entity that directly or indirectly controls, is controlled by, or is under common control with the subject entity; for purposes of this definition, “control” means direct or indirect ownership or control of more than 50% of the voting interests of the subject entity.
+
+"Agreement" means this Stability AI Community License Agreement.
+
+“AUP” means the Stability AI Acceptable Use Policy available at (https://stability.ai/use-policy), as may be updated from time to time.
+
+"Derivative Work(s)” means (a) any derivative work of the Stability AI Materials as recognized by U.S. copyright laws and (b) any modifications to a Model, and any other model created which is based on or derived from the Model or the Model’s output, including “fine tune” and “low-rank adaptation” models derived from a Model or a Model’s output, but do not include the output of any Model.
+
+“Documentation” means any specifications, manuals, documentation, and other written information provided by Stability AI related to the Software or Models.
+
+“Model(s)" means, collectively, Stability AI’s proprietary models and algorithms, including machine-learning models, trained model weights and other elements of the foregoing listed on Stability’s Core Models Webpage available at (https://stability.ai/core-models), as may be updated from time to time.
+
+"Stability AI" or "we" means Stability AI Ltd. and its Affiliates.
+
+"Software" means Stability AI’s proprietary software made available under this Agreement now or in the future.
+
+“Stability AI Materials” means, collectively, Stability’s proprietary Models, Software and Documentation (and any portion or combination thereof) made available under this Agreement. 
+
+“Trade Control Laws” means any applicable U.S. and non-U.S. export control and trade sanctions laws and regulations.

samel/LICENSE_GEMMA.md

@@ -0,0 +1,206 @@
+# Gemma Terms of Use
+
+The terms below apply to Gemma models listed in the Appendix at bottom of this page. For Gemma 4 terms, see the [Gemma 4 license](https://ai.google.dev/gemma/apache_2).
+
+Last modified: April 1, 2026
+
+By using, reproducing, modifying, distributing, performing or displaying any
+portion or element of Gemma, Model Derivatives including via any Hosted Service,
+(each as defined below) (collectively, the "**Gemma Services**") or otherwise
+accepting the terms of this Agreement, you agree to be bound by this Agreement.
+
+## Section 1: DEFINITIONS
+
+### 1.1 Definitions
+
+(a) "**Agreement** " or "**Gemma Terms of Use**" means these terms and conditions
+that govern the use, reproduction, Distribution or modification of the Gemma
+Services and any terms and conditions incorporated by reference.
+
+(b) "**Distribution** " or "**Distribute** " means any transmission, publication,
+or other sharing of Gemma or Model Derivatives to a third party, including by
+providing or making Gemma or its functionality available as a hosted service via
+API, web access, or any other electronic or remote means ("**Hosted Service**").
+
+(c) "**Gemma** " means the set of machine learning language models, trained model
+weights and parameters identified in the [Appendix](https://ai.google.dev/gemma/terms#appendix),
+regardless of the source that you obtained it from.
+
+(d) "**Google**" means Google LLC.
+
+(e) "**Model Derivatives**" means all (i) modifications to Gemma, (ii) works based
+on Gemma, or (iii) any other machine learning model which is created by transfer
+of patterns of the weights, parameters, operations, or Output of Gemma, to that
+model in order to cause that model to perform similarly to Gemma, including
+distillation methods that use intermediate data representations or methods based
+on the generation of synthetic data Outputs by Gemma for training that model.
+For clarity, Outputs are not deemed Model Derivatives.
+
+(f) "**Output**" means the information content output of Gemma or a Model
+Derivative that results from operating or otherwise using Gemma or the Model
+Derivative, including via a Hosted Service.
+
+### 1.2
+
+As used in this Agreement, "**including** " means
+"**including without limitation**".
+
+## Section 2: ELIGIBILITY AND USAGE
+
+### 2.1 Eligibility
+
+You represent and warrant that you have the legal capacity to enter into this
+Agreement (including being of sufficient age of consent). If you are accessing
+or using any of the Gemma Services for or on behalf of a legal entity, (a) you
+are entering into this Agreement on behalf of yourself and that legal entity,
+(b) you represent and warrant that you have the authority to act on behalf of
+and bind that entity to this Agreement and (c) references to "**you** " or
+"**your**" in the remainder of this Agreement refers to both you (as an
+individual) and that entity.
+
+### 2.2 Use
+
+You may use, reproduce, modify, Distribute, perform or display any of the Gemma
+Services only in accordance with the terms of this Agreement, and must not
+violate (or encourage or permit anyone else to violate) any term of this
+Agreement.
+
+## Section 3: DISTRIBUTION AND RESTRICTIONS
+
+### 3.1 Distribution and Redistribution
+
+You may reproduce or Distribute copies of Gemma or Model Derivatives if you meet
+all of the following conditions:
+
+1. You must include the use restrictions referenced in Section 3.2 as an enforceable provision in any agreement (e.g., license agreement, terms of use, etc.) governing the use and/or distribution of Gemma or Model Derivatives and you must provide notice to subsequent users you Distribute to that Gemma or Model Derivatives are subject to the use restrictions in Section 3.2.
+2. You must provide all third party recipients of Gemma or Model Derivatives a copy of this Agreement.
+3. You must cause any modified files to carry prominent notices stating that you modified the files.
+4. All Distributions (other than through a Hosted Service) must be accompanied by a "**Notice** " text file that contains the following notice: "**Gemma is provided under and subject to the Gemma Terms of Use found at ai.google.dev/gemma/terms**".
+
+You may add your own intellectual property statement to your modifications and,
+except as set forth in this Section, may provide additional or different terms
+and conditions for use, reproduction, or Distribution of your modifications, or
+for any such Model Derivatives as a whole, provided your use, reproduction,
+modification, Distribution, performance, and display of Gemma otherwise complies
+with the terms and conditions of this Agreement. Any additional or different
+terms and conditions you impose must not conflict with the terms of this
+Agreement.
+
+### 3.2 Use Restrictions
+
+You must not use any of the Gemma Services:
+
+1. for the restricted uses set forth in the Gemma Prohibited Use Policy at [ai.google.dev/gemma/prohibited_use_policy](https://ai.google.dev/gemma/prohibited_use_policy) ("**Prohibited Use Policy**"), which is hereby incorporated by reference into this Agreement; or
+2. in violation of applicable laws and regulations.
+
+To the maximum extent permitted by law, Google reserves the right to restrict
+(remotely or otherwise) usage of any of the Gemma Services that Google
+reasonably believes are in violation of this Agreement.
+
+### 3.3 Generated Output
+
+Google claims no rights in Outputs you generate using Gemma. You and your users
+are solely responsible for Outputs and their subsequent uses.
+
+## Section 4: ADDITIONAL PROVISIONS
+
+### 4.1 Updates
+
+Google may update Gemma from time to time.
+
+### 4.2 Trademarks
+
+Nothing in this Agreement grants you any rights to use Google's trademarks,
+trade names, logos or to otherwise suggest endorsement or misrepresent the
+relationship between you and Google. Google reserves any rights not expressly
+granted herein.
+
+### 4.3 DISCLAIMER OF WARRANTY
+
+UNLESS REQUIRED BY APPLICABLE LAW, THE GEMMA SERVICES, AND OUTPUTS, ARE PROVIDED
+ON AN "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, EITHER
+EXPRESS OR IMPLIED, INCLUDING ANY WARRANTIES OR CONDITIONS OF TITLE,
+NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. YOU ARE
+SOLELY RESPONSIBLE FOR DETERMINING THE APPROPRIATENESS OF USING, REPRODUCING,
+MODIFYING, PERFORMING, DISPLAYING OR DISTRIBUTING ANY OF THE GEMMA SERVICES
+OR OUTPUTS AND ASSUME ANY AND ALL RISKS ASSOCIATED WITH YOUR USE OR DISTRIBUTION
+OF ANY OF THE GEMMA SERVICES OR OUTPUTS AND YOUR EXERCISE OF RIGHTS AND
+PERMISSIONS UNDER THIS AGREEMENT.
+
+### 4.4 LIMITATION OF LIABILITY
+
+TO THE FULLEST EXTENT PERMITTED BY APPLICABLE LAW, IN NO EVENT AND UNDER NO
+LEGAL THEORY, WHETHER IN TORT (INCLUDING NEGLIGENCE), PRODUCT LIABILITY,
+CONTRACT, OR OTHERWISE, UNLESS REQUIRED BY APPLICABLE LAW, SHALL GOOGLE OR ITS
+AFFILIATES BE LIABLE TO YOU FOR DAMAGES, INCLUDING ANY DIRECT, INDIRECT,
+SPECIAL, INCIDENTAL, EXEMPLARY, CONSEQUENTIAL, OR PUNITIVE DAMAGES, OR LOST
+PROFITS OF ANY KIND ARISING FROM THIS AGREEMENT OR RELATED TO, ANY OF THE GEMMA
+SERVICES OR OUTPUTS EVEN IF GOOGLE OR ITS AFFILIATES HAVE BEEN ADVISED OF THE
+POSSIBILITY OF SUCH DAMAGES.
+
+### 4.5 Term, Termination, and Survival
+
+The term of this Agreement will commence upon your acceptance of this Agreement
+(including acceptance by your use, modification, or Distribution, reproduction,
+performance or display of any portion or element of the Gemma Services) and will
+continue in full force and effect until terminated in accordance with the terms
+of this Agreement. Google may terminate this Agreement if you are in breach of
+any term of this Agreement. Upon termination of this Agreement, you must delete
+and cease use and Distribution of all copies of Gemma and Model Derivatives in
+your possession or control. Sections 1, 2.1, 3.3, 4.2 to 4.9 shall survive the
+termination of this Agreement.
+
+### 4.6 Governing Law and Jurisdiction
+
+This Agreement will be governed by the laws of the State of California without
+regard to choice of law principles. The UN Convention on Contracts for the
+International Sale of Goods does not apply to this Agreement. The state and
+federal courts of Santa Clara County, California shall have exclusive
+jurisdiction of any dispute arising out of this Agreement.
+
+### 4.7 Severability
+
+If any provision of this Agreement is held to be invalid, illegal or
+unenforceable, the remaining provisions shall be unaffected thereby and remain
+valid as if such provision had not been set forth herein.
+
+### 4.8 Entire Agreement
+
+This Agreement states all the terms agreed between the parties and supersedes
+all other agreements between the parties as of the date of acceptance relating
+to its subject matter.
+
+### 4.9 No Waiver
+
+Google will not be treated as having waived any rights by not exercising (or
+delaying the exercise of) any rights under this Agreement.
+
+## Appendix
+
+- [Gemma 1](https://ai.google.dev/gemma/docs/core/model_card)
+- [Gemma 1.1](https://ai.google.dev/gemma/docs/core/model_card)
+- [Gemma 2](https://ai.google.dev/gemma/docs/core/model_card_2)
+- [Gemma 3](https://ai.google.dev/gemma/docs/core/model_card_3)
+- [Gemma 3n](https://ai.google.dev/gemma/docs/3n)
+- [FunctionGemma](https://ai.google.dev/gemma/docs/functiongemma)
+- [EmbeddingGemma](https://ai.google.dev/gemma/docs/embeddinggemma)
+- [PaliGemma](https://ai.google.dev/gemma/docs/paligemma/model-card)
+- [PaliGemma 2](https://ai.google.dev/gemma/docs/paligemma/model-card-2)
+- [ShieldGemma](https://ai.google.dev/gemma/docs/shieldgemma/model_card)
+- [ShieldGemma 2](https://ai.google.dev/gemma/docs/shieldgemma/model_card_2)
+- [CodeGemma](https://ai.google.dev/gemma/docs/codegemma/model_card)
+- [CodeGemma 1.1](https://ai.google.dev/gemma/docs/codegemma/model_card)
+- [Gemma 2 JPN](https://huggingface.co/google/gemma-2-2b-jpn-it)
+- [DataGemma RIG](https://www.kaggle.com/models/google/datagemma-rig)
+- [DataGemma RAG](https://www.kaggle.com/models/google/datagemma-rag)
+- [RecurrentGemma](https://ai.google.dev/gemma/docs/recurrentgemma/model_card)
+- [Gemma Scope](https://ai.google.dev/gemma/docs/gemma_scope)
+- [Gemma-APS](https://ai.google.dev/gemma/docs/gemma-aps)
+- [T5Gemma](https://www.kaggle.com/models/google/t5gemma)
+- [VaultGemma](https://www.kaggle.com/models/google/vaultgemma)
+- [FunctionGemma](https://www.kaggle.com/models/google/functiongemma)
+- [T5Gemma 2](https://www.kaggle.com/models/google/t5gemma-2)
+- [TranslateGemma](https://www.kaggle.com/models/google/translategemma)
+
+> [!NOTE]
+> **Note:** Previous versions of these Terms are [archived here](https://ai.google.dev/gemma/terms-archive).

samel/NOTICE

@@ -0,0 +1 @@
+Gemma is provided under and subject to the Gemma Terms of Use found at ai.google.dev/gemma/terms

samel/README.md

@@ -0,0 +1,101 @@
+---
+language:
+- en
+library_name: stable-audio-3
+license: other
+license_name: stable-audio-community
+license_link: LICENSE
+tags:
+- music
+- sound-effects
+- audio
+- autoencoder
+---
+
+# SAME: A Semantically-Aligned Music Autoencoder
+
+Please note: For commercial use, please refer to [https://stability.ai/license](https://stability.ai/license)
+
+## Model Description
+Latent representations are at the heart of the majority of modern generative models. 
+In the audio domain they are typically produced by a neural-audio-codec autoencoder. 
+In this work we introduce SAME (Semantically Aligned Music autoEncoder), 
+a transformer-based autoencoder for stereo music and general audio that reaches a 4096x temporal compression ratio (roughly twice the current standard)
+while maintaining excellent reconstruction quality and strong downstream generative performance. 
+We achieve this by combining a set of semantic regularisation approaches with phase-aware reconstruction losses. 
+The architecture also delivers substantial computational cost benefits, through both its high compression ratio and its reliance on well-optimised transformer primitives. 
+Two variants (a large SAME-L and a CPU-deployable SAME-S) are released in open-weights form.
+
+## Usage
+
+This model can be used with:
+1. the [`stable-audio-3`](https://github.com/Stability-AI/stable-audio-3) inference and fine-tuning library
+2. the [`stable-audio-tools`](https://github.com/Stability-AI/stable-audio-tools) research library
+
+
+### Using with `stable-audio-3`
+import torchaudio
+from stable_audio_3 import AutoencoderModel
+
+ae = AutoencoderModel.from_pretrained("same-l")
+waveform, sr = torchaudio.load("audio.wav")
+latents = ae.encode(waveform, sr)
+audio_out = ae.decode(latents)
+
+
+### Using with `stable-audio-tools`
+
+```python
+import torch
+import torchaudio
+from einops import rearrange
+from stable_audio_tools import get_pretrained_model
+from stable_audio_tools.inference.generation import generate_diffusion_cond
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+if device == "cuda":
+  model_half = True
+
+# Download model
+model, model_config = get_pretrained_model("stabilityai/SAME-L")
+sample_rate = model_config["sample_rate"]
+sample_size = model_config["sample_size"]
+
+model = model.to(device)
+if model_half:
+  model = model.to(torch.float16)
+
+audio, sr = torchaudio.load(/path/to/audiofile)  # [channels, samples]
+if audio.shape[0] == 1:
+    audio = audio.repeat(2, 1)
+
+audio = audio.unsqueeze(0).to(device)
+if model_half:
+  audio = audio.half()
+with torch.no_grad():
+    latents = model.encode_audio(audio)  
+    reconstructed = model.decode_audio(latents)         
+reconstructed = reconstructed.squeeze(0).cpu()  
+reconstructed = reconstructed.to(torch.float32).clamp(-1, 1).mul(32767).to(torch.int16).cpu()
+
+```
+
+
+## Model Details
+* **Model type**: `SAME` is a continuous autoencoder model based on a transformer architecture.
+* **Language(s)**: English
+* **License**: [Stability AI Community License](https://stability.ai/license).
+* **Research Paper**: [https://arxiv.org/abs/2605.18613](https://arxiv.org/abs/2605.18613)
+
+
+## Training dataset
+
+### Datasets Used
+Our dataset consists of ~19,500 hours of licensed production audio from [AudioSparx](https://www.audiosparx.com/) which includes a 66/25/9% mix of music, sound effects, and instrument stems.
+
+
+
+
+
+
+

samel/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8cf93d5345c559307fc429bde896f594516fd6a1a23014c2f6055f78d8e13378
+size 3408566892

samel/model_config.json

@@ -0,0 +1,67 @@
+{
+    "model_type": "autoencoder",
+    "sample_size": 196608,
+    "sample_rate": 44100,
+    "audio_channels": 2,
+    "model": {
+        "pretransform": {
+            "type": "patched",
+            "config": {
+                "patch_size": 256,
+                "channels": 2
+            }
+        },
+        "encoder": {
+            "type": "same",
+            "requires_grad": false,
+            "config": {
+                "in_channels": 512,
+                "channels": 256,
+                "c_mults": [6],
+                "strides": [16],
+                "latent_dim": 256,
+                "transformer_depths": [12],
+                "checkpointing": true,
+                "differential": true,
+                "dyt": true,
+                "dim_heads": 64,
+                "sliding_window": [1,1],
+                "variable_stride": true,
+                "mask_noise": 0.001
+            }
+        },
+        "decoder": {
+            "type": "same",
+            "requires_grad": false,
+            "config": {
+                "out_channels": 512,
+                "channels": 256,
+                "c_mults": [6],
+                "strides": [16],
+                "latent_dim": 256,
+                "transformer_depths": [12],
+                "sinusoidal_blocks": [8],
+                "checkpointing": false,
+                "differential": true,
+                "dyt": true,
+                "dim_heads": 64,
+                "sliding_window": [1,1],
+                "variable_stride": true,
+                "mask_noise": 0.1
+            }
+        },
+        "bottleneck": {
+            "type": "softnorm",
+            "config": {
+                "dim": 256,
+                "noise_augment_dim": 0,
+                "noise_regularize": true,
+                "auto_scale": true,
+                "freeze": true
+            }
+        },
+        "latent_dim": 256,
+        "downsampling_ratio": 4096,
+        "io_channels": 2
+    }
+}

tags.json

@@ -0,0 +1,461 @@
+{
+"metadata": {
+"acoustic": 28,
+"acoustic guitar": 25,
+"afrobeat": 275,
+"aggressive": 77,
+"alternative": 286,
+"alternative metal": 150,
+"alternative pop": 267,
+"alternative r&b": 355,
+"alternative rock": 39,
+"ambient": 63,
+"ambient pop": 333,
+"americana": 271,
+"angry": 162,
+"anthem": 244,
+"anthemic": 24,
+"arena rock": 192,
+"art pop": 338,
+"art rock": 301,
+"atmospheric": 30,
+"avant-garde": 249,
+"ballad": 11,
+"bass": 97,
+"bass guitar": 137,
+"beats": 322,
+"big band": 305,
+"big beat": 396,
+"bittersweet": 65,
+"bluegrass": 314,
+"blues": 131,
+"blues rock": 207,
+"boastful": 117,
+"boom bap": 250,
+"bossa nova": 344,
+"breakbeat": 241,
+"c-pop": 22,
+"calm": 80,
+"caribbean music": 374,
+"catchy": 96,
+"cello": 142,
+"chamber pop": 413,
+"chill": 222,
+"chillout": 210,
+"chillwave": 277,
+"chiptune": 272,
+"choir": 173,
+"choral": 253,
+"cinematic": 75,
+"cinematic classical": 214,
+"classic rock": 329,
+"classical": 143,
+"classical music": 132,
+"cloud rap": 324,
+"club": 125,
+"complex": 170,
+"conscious": 199,
+"conscious hip hop": 226,
+"contemporary country": 323,
+"contemporary folk": 160,
+"country": 81,
+"country pop": 265,
+"country rock": 227,
+"cyberpunk": 299,
+"dance": 38,
+"dance-pop": 52,
+"dancehall": 318,
+"dark": 53,
+"dark ambient": 291,
+"dark pop": 348,
+"darkwave": 331,
+"death metal": 264,
+"deep": 189,
+"deep bass": 239,
+"deep house": 221,
+"disco": 161,
+"djent": 297,
+"doom metal": 390,
+"downtempo": 138,
+"dramatic": 181,
+"dream pop": 92,
+"dreamy": 95,
+"drill": 304,
+"drum and bass": 156,
+"drum machine": 43,
+"drums": 51,
+"dub": 341,
+"dubstep": 133,
+"dynamic": 255,
+"east coast hip hop": 219,
+"eclectic": 120,
+"edm": 88,
+"electric": 279,
+"electric guitar": 37,
+"electro": 99,
+"electro house": 184,
+"electro pop": 632,
+"electro swing": 307,
+"electronic": 19,
+"emo": 155,
+"emotional": 21,
+"empowering": 124,
+"empowerment": 257,
+"energetic": 18,
+"epic": 56,
+"epic music": 351,
+"ethereal": 44,
+"euphoric": 208,
+"european music": 398,
+"experimental": 106,
+"fantasy": 230,
+"fast": 223,
+"female vocal": 12,
+"film score": 157,
+"flute": 169,
+"folk": 34,
+"folk pop": 94,
+"folk rock": 177,
+"fun": 270,
+"funk": 107,
+"funky": 258,
+"future bass": 193,
+"futuristic": 119,
+"gangsta rap": 200,
+"gentle": 32,
+"glam metal": 418,
+"glam rock": 353,
+"glitch": 183,
+"gospel": 243,
+"gothic metal": 313,
+"gothic rock": 381,
+"gritty": 289,
+"groove metal": 391,
+"groovy": 202,
+"grunge": 233,
+"guitar": 62,
+"happy": 122,
+"hard rock": 70,
+"hardcore": 346,
+"hardcore hip hop": 149,
+"harmonica": 276,
+"harp": 354,
+"haunting": 260,
+"heartfelt": 136,
+"heavy": 89,
+"heavy bass": 280,
+"heavy metal": 108,
+"high-energy": 445,
+"hip hop": 47,
+"hispanic american music": 356,
+"hispanic music": 361,
+"hopeful": 66,
+"house": 87,
+"humorous": 292,
+"hypnotic": 114,
+"idm": 262,
+"indie": 201,
+"indie folk": 48,
+"indie pop": 50,
+"indie rock": 98,
+"industrial": 185,
+"inspirational": 49,
+"instrumental": 35,
+"intense": 176,
+"introspective": 71,
+"j-pop": 191,
+"j-rock": 393,
+"jamaican music": 412,
+"jazz": 84,
+"jazz fusion": 259,
+"jazz rap": 384,
+"kpop": 6,
+"latin": 325,
+"latin pop": 387,
+"lo-fi": 148,
+"longing": 121,
+"love": 69,
+"lush": 78,
+"male vocal": 20,
+"mechanical": 251,
+"medieval": 281,
+"meditative": 212,
+"melancholic": 33,
+"melancholy": 365,
+"mellow": 90,
+"melodic": 16,
+"melodic death metal": 452,
+"melodic metalcore": 438,
+"metal": 55,
+"metalcore": 196,
+"minimal": 360,
+"minimalist": 287,
+"modern classical": 290,
+"mysterious": 163,
+"nature": 206,
+"neo-psychedelia": 362,
+"neo-soul": 285,
+"neon": 218,
+"new age": 198,
+"new wave": 167,
+"nocturnal": 105,
+"nostalgic": 101,
+"nu metal": 164,
+"ominous": 171,
+"opera": 232,
+"optimistic": 153,
+"orchestral": 57,
+"outlaw country": 375,
+"party": 86,
+"passionate": 42,
+"phonk": 215,
+"piano": 14,
+"playful": 82,
+"poetic": 242,
+"pop": 10,
+"pop punk": 205,
+"pop rap": 100,
+"pop rock": 45,
+"pop soul": 236,
+"post-grunge": 273,
+"post-hardcore": 268,
+"post-industrial": 340,
+"post-punk": 245,
+"power ballad": 118,
+"power metal": 151,
+"power pop": 166,
+"powerful": 102,
+"powerful vocals": 158,
+"progressive": 127,
+"progressive electronic": 306,
+"progressive house": 220,
+"progressive metal": 203,
+"progressive rock": 152,
+"progressive trance": 377,
+"psychedelic": 109,
+"psychedelic pop": 342,
+"psychedelic rock": 225,
+"punk": 154,
+"punk rock": 188,
+"quirky": 231,
+"r&b": 67,
+"rap": 41,
+"raw": 139,
+"rebellious": 213,
+"reflective": 46,
+"reggae": 147,
+"reggaeton": 274,
+"regional music": 85,
+"repetitive": 128,
+"retro": 317,
+"reverb": 302,
+"rhythmic": 27,
+"rock": 13,
+"rock and roll": 541,
+"romantic": 17,
+"sad": 123,
+"sampling": 134,
+"saxophone": 248,
+"sensual": 180,
+"sentimental": 113,
+"shoegaze": 371,
+"singer-songwriter": 40,
+"ska": 368,
+"slow": 126,
+"slow tempo": 320,
+"smooth": 187,
+"smooth jazz": 269,
+"smooth soul": 339,
+"soft": 76,
+"soft pop": 115,
+"soft rock": 26,
+"sombre": 174,
+"soothing": 112,
+"soul": 93,
+"soulful": 178,
+"soundtrack": 7,
+"southern hip hop": 316,
+"southern rock": 453,
+"spiritual": 263,
+"spoken word": 312,
+"storytelling": 186,
+"strings": 29,
+"summer": 179,
+"surreal": 235,
+"suspenseful": 195,
+"swing": 197,
+"symphonic": 337,
+"symphonic metal": 175,
+"synth": 60,
+"synth pad": 79,
+"synth pop": 172,
+"synthesizer": 31,
+"synthetic": 1,
+"synthwave": 72,
+"tech house": 311,
+"technical": 209,
+"techno": 110,
+"thrash metal": 359,
+"trance": 129,
+"trap": 73,
+"trip hop": 211,
+"triumphant": 159,
+"trumpet": 308,
+"upbeat": 61,
+"uplifting": 15,
+"urban": 103,
+"violin": 91,
+"vocal": 294,
+"vocal jazz": 388,
+"warm": 54,
+"west coast hip hop": 392,
+"western classical music": 140
+},
+"other": {
+"A Hat in Time": 1394,
+"Ace Attorney": 922,
+"Ace Combat 7": 923,
+"Animal Crossing New Horizons": 937,
+"Aquaria": 1456,
+"Ar Tonelico": 938,
+"Arknights": 1447,
+"Assasins Creed Valhalla": 944,
+"Atelier Series": 956,
+"Baldur's Gate 3": 963,
+"Banjo Kazooie": 964,
+"Banjo Tooie": 965,
+"battle": 8,
+"Black Myth Wukong": 972,
+"Blue Reflection": 976,
+"Bravely Default": 979,
+"Castlevania Symphony of the Night": 991,
+"Cave Story": 992,
+"Celeste": 993,
+"Chrono Cross": 996,
+"Chrono Trigger": 997,
+"Clair Obscur": 999,
+"Clannad": 1000,
+"Cowboy Bebop": 1003,
+"credits": 9,
+"Crosscode": 1004,
+"Cuphead": 1005,
+"Deltarune": 1011,
+"Devil May Cry 4": 1020,
+"Devil May Cry 5": 1430,
+"Diddy Kong Racing": 1016,
+"Divinity Original Sin 2": 1457,
+"Doki Doki Literature Club": 1021,
+"Donkey Kong Country": 1024,
+"Dragon Quest IX": 1031,
+"Elder Scrolls 3 Morrowind": 1035,
+"EverQuest": 1038,
+"Fate Grand Order": 1047,
+"Final Fantasy I": 869,
+"Final Fantasy II": 870,
+"Final Fantasy III": 871,
+"Final Fantasy IV": 874,
+"Final Fantasy IX": 882,
+"Final Fantasy V": 876,
+"Final Fantasy VI": 877,
+"Final Fantasy VII": 879,
+"Final Fantasy VII Remake": 880,
+"Final Fantasy VIII": 881,
+"Final Fantasy X": 883,
+"Final Fantasy X-2": 894,
+"Final Fantasy XI": 884,
+"Final Fantasy XII": 885,
+"Final Fantasy XIII": 886,
+"Final Fantasy XIII-2": 897,
+"Final Fantasy XIII-3": 898,
+"Final Fantasy XIV": 887,
+"Final Fantasy XV": 888,
+"Fire Emblem Awakening": 1056,
+"Fire Emblem Three Houses": 1058,
+"Frieren": 1059,
+"Genshin Impact": 1065,
+"Golden Sun": 1067,
+"Granblue Fantasy": 1069,
+"Granblue Fantasy Versus": 1070,
+"Guilty Gear Strive": 1462,
+"Hades 2": 1077,
+"Haikyuu": 1078,
+"Harvest Moon": 1079,
+"Hollow Knight": 1086,
+"Hollow Knight Silksong": 1087,
+"Honkai Impact 3rd": 1105,
+"Honkai Star Rail": 1108,
+"K-ON": 1123,
+"Kingdom Hearts": 1129,
+"Kirby": 1359,
+"Lies of P": 1151,
+"Maple Story": 1157,
+"Mario & Luigi - Superstar Saga": 1163,
+"Mario Kart 8": 1444,
+"Mario Kart World": 1173,
+"Mega Man": 1181,
+"Mega Man X": 1187,
+"Metaphor Refantazio": 1429,
+"Minecraft": 1199,
+"My Hero Academia": 1206,
+"Nausicaa Valley of the Wind": 1207,
+"Neon White": 1209,
+"Nier Automata": 1216,
+"One Piece": 1222,
+"Osu!": 1389,
+"Outer Wilds": 1223,
+"Perfect Dark": 1227,
+"Persona": 1228,
+"Pokemon Black and White": 1235,
+"Pokemon Crystal": 1236,
+"Pokemon Diamond and Pearl": 1237,
+"Pokemon Fire Red and Leaf Green": 1239,
+"Pokemon Heart Gold & Soul Silver": 1240,
+"Pokemon Mystery Dungeon - Explorers of Sky": 1243,
+"Pokemon Red and Blue": 1246,
+"Pokemon Ruby and Sapphire": 1247,
+"Pokemon Scarlet and Violet": 1249,
+"Pokemon Sun and Moon": 1250,
+"Pokemon X and Y": 1252,
+"Pokemon XD - Gale of Darkness": 1253,
+"Pokopia": 1255,
+"RuneScape": 1403,
+"Sayonara Wild Hearts": 1281,
+"Sea of Stars": 1282,
+"Secret of Mana": 1283,
+"Sonic Adventure": 1302,
+"Sonic the Hedgehog 3": 1307,
+"Spirited Away": 1309,
+"Splatoon": 1467,
+"Star Ocean": 1312,
+"Stellar Blade": 1426,
+"Super Mario Galaxy": 1327,
+"Super Mario Odyssey": 1328,
+"Super Mario RPG": 1329,
+"Super Meat Boy": 1451,
+"Super Monkey Ball 2": 1331,
+"Super Smash Bros. Brawl": 1333,
+"Tales of Symphonia": 1338,
+"The Elder Scrolls IV: Oblivion": 1468,
+"The Legend of Zelda": 1144,
+"The Legend of Zelda A Link to the Past": 1145,
+"The Legend of Zelda Ocarina of Time": 1383,
+"Touhou": 1345,
+"Umamusume": 1355,
+"Umineko": 1396,
+"Undertale": 1356,
+"Utamonogatari": 1358,
+"Vampire Survivors": 1469,
+"Vocaloid": 1415,
+"Wild Arms 2": 1368,
+"Witcher 3": 1369,
+"Wuchang Fallen Feathers": 1459,
+"Wuthering Waves": 1370,
+"Xenoblade Chronicles": 1372,
+"Xenoblade Chronicles 2": 1373,
+"Xenoblade Chronicles 3": 1375,
+"Xenogears": 1376,
+"Ys": 1380,
+"Yuki Kajiura": 1401,
+"Zenless Zone Zero": 1385
+}
+}

webui.py

@@ -0,0 +1,79 @@
+from __future__ import annotations
+
+import json
+import secrets
+from itertools import pairwise
+from pathlib import Path
+
+import gradio as gr
+import torch
+
+from audio_dit import CHANNELS, FRAMES, MAX_TAGS, load_model
+from same_l_decoder import SAMPLE_RATE, load_decoder
+
+DEVICE = ("cuda" if torch.cuda.is_available() else
+          "mps" if torch.backends.mps.is_available() else "cpu")
+DTYPE = torch.float32 if DEVICE == "cpu" else torch.float16
+AUTOCAST = torch.bfloat16 if DEVICE == "cuda" else torch.float16
+
+STEPS = 50
+DEFAULT_CFG = 3.5
+TIME_SHIFT = 2.0
+
+TAG_GROUPS = json.loads((Path(__file__).parent / "tags.json").read_text(encoding="utf-8"))
+METADATA_TAGS = TAG_GROUPS["metadata"]
+OTHER_TAGS = TAG_GROUPS["other"]
+SYNTHETIC_TAG_ID = METADATA_TAGS["synthetic"]
+TAGS = METADATA_TAGS | OTHER_TAGS
+del TAGS["synthetic"]  # Source conditioning is added automatically, not shown in the UI.
+
+TIMESTEPS = [1 - i / STEPS for i in range(STEPS + 1)]
+TIMESTEPS = [TIME_SHIFT * t / (1 + (TIME_SHIFT - 1) * t) for t in TIMESTEPS]
+
+model, latent_mean, latent_std = load_model(DEVICE)
+decoder = load_decoder(device=DEVICE, dtype=DTYPE)
+
+
+@torch.no_grad()
+def generate(tag_names: list[str], cfg_scale: float, seed: int):
+    if not tag_names:
+        raise gr.Error("Pick at least one tag.")
+    seed = secrets.randbits(63) if seed == -1 else seed
+    torch.manual_seed(seed)
+    ids = [TAGS[name] for name in tag_names]
+    if all(name in METADATA_TAGS for name in tag_names):
+        if len(ids) == MAX_TAGS:
+            raise gr.Error(f"Choose at most {MAX_TAGS - 1} metadata tags.")
+        ids.append(SYNTHETIC_TAG_ID)
+    tags = torch.tensor([ids + [0] * (MAX_TAGS - len(ids))], device=DEVICE)
+    null = torch.zeros_like(tags)
+    x = torch.randn(1, CHANNELS, FRAMES, device=DEVICE)
+
+    for i, (t, t_next) in enumerate(pairwise(TIMESTEPS)):
+        # Guidance alternates between the shallow path-drop branch and a full
+        # unconditional pass as the weak model to steer away from.
+        weak = (model.shallow, model)[i % 2]
+        with torch.autocast(DEVICE, dtype=AUTOCAST, enabled=DEVICE != "cpu"):
+            cond = model(x, t, tags).float()
+            uncond = weak(x, t, null).float()
+        x = x + (t_next - t) * (uncond + cfg_scale * (cond - uncond))
+
+    audio = decoder.decode((x * latent_std + latent_mean).to(DTYPE))
+    return SAMPLE_RATE, audio[0].float().clamp(-1, 1).T.cpu().numpy()
+
+
+with gr.Blocks(title="Audio DiT") as demo:
+    gr.Markdown("# Audio DiT\nChoose 1-8 tags.")
+    tag_box = gr.Dropdown(choices=list(TAGS), value=[], multiselect=True,
+                          max_choices=MAX_TAGS, label="Tags", filterable=True)
+    cfg_scale = gr.Slider(2.0, 6.0, value=DEFAULT_CFG, step=0.5,
+                          label="CFG scale")
+    seed = gr.Number(value=-1, precision=0, minimum=-1,
+                     label="Seed (-1 = random)")
+    button = gr.Button("Generate", variant="primary")
+    audio_out = gr.Audio(label="Output", interactive=False)
+    button.click(generate, inputs=[tag_box, cfg_scale, seed], outputs=audio_out,
+                 concurrency_limit=1)
+
+if __name__ == "__main__":
+    demo.queue().launch()