audio!

app.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import json
 import tempfile
 import os
 
@@ -7,14 +8,19 @@ os.environ.setdefault("HF_HOME", r"D:\hf-cache")
 os.environ.setdefault("HUGGINGFACE_HUB_CACHE", r"D:\hf-cache\hub")
 
 import gradio as gr
+import librosa
 import numpy as np
 import torch
 import trimesh
 from diffusers import AutoencoderKL
+from diffusers.pipelines.deprecated.audio_diffusion.mel import Mel
 from huggingface_hub import hf_hub_download
+from PIL import Image
 from safetensors.torch import load_file
-from transformers import CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5TokenizerFast
+from transformers import AutoTokenizer, CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5TokenizerFast
+from transformers import ClapTextModelWithProjection
 
+from audio_dit import AudioDiT
 from pixel_dit import DiT
 from voxel_dit import VoxelDiT
 from mmdit import MMDiT
@@ -50,6 +56,27 @@ vm1_weights = hf_hub_download("bench-labs/VoxelModel-v1", "model.safetensors")
 voxel_model = VoxelDiT().to(DEV).eval()
 voxel_model.load_state_dict(load_file(vm1_weights))
 
+print("[boot] loading AudioModel v1...")
+audio_cfg = json.load(open(hf_hub_download("bench-labs/AudioModel-v1", "config.json")))
+a_dit = audio_cfg["dit"]
+audio_model = AudioDiT(x_res=a_dit["x_res"], y_res=a_dit["y_res"],
+                       text_seq_dim=a_dit["text_seq_dim"], text_pool_dim=a_dit["text_pool_dim"]).to(DEV).eval()
+audio_model.load_state_dict(load_file(hf_hub_download("bench-labs/AudioModel-v1", "model_best.safetensors")))
+audio_mel = Mel(x_res=a_dit["x_res"], y_res=a_dit["y_res"],
+                sample_rate=audio_cfg["mel"]["sample_rate"], n_fft=audio_cfg["mel"]["n_fft"],
+                hop_length=audio_cfg["mel"]["hop_length"], top_db=audio_cfg["mel"]["top_db"])
+# Fast mel->STFT pseudo-inverse (librosa's pre-0.10 behavior). librosa 0.10+ uses an
+# NNLS/L-BFGS solver here that allocates ~2GB and takes minutes on CPU for this
+# 384x256 spectrogram; the pinv gives the same Griffin-Lim output in about a second.
+audio_mel_pinv = np.linalg.pinv(
+    librosa.filters.mel(sr=audio_mel.sr, n_fft=audio_mel.n_fft, n_mels=audio_mel.y_res, dtype=np.float32)
+)
+# Text-only CLAP tower: runs the same text_model + text_projection the reference
+# sample.py uses, but skips the unused HTSAT audio tower (~150M params of dead weight).
+audio_tokenizer = AutoTokenizer.from_pretrained("laion/clap-htsat-unfused")
+audio_clap = ClapTextModelWithProjection.from_pretrained("laion/clap-htsat-unfused").to(DEV).eval()
+audio_null_seq = audio_null_pool = None
+
 print("[boot] loading PixelModel v6...")
 v6_weights = hf_hub_download("bench-labs/PixelModel-v6", "model.safetensors")
 pixel_model_v6 = MMDiT().to(DEV).eval()
@@ -129,6 +156,56 @@ def sample_image_v6(prompt: str, steps: int, cfg: float, seed: int, progress=gr.
     return (((img.clamp(-1, 1) + 1) / 2).permute(0, 2, 3, 1).numpy()[0] * 255).round().astype(np.uint8)
 
 
+AUDIO_MAX_TOKENS = 32
+
+
+@torch.no_grad()
+def encode_audio(strings):
+    t = audio_tokenizer(strings, padding="max_length", max_length=AUDIO_MAX_TOKENS,
+                        truncation=True, return_tensors="pt").to(DEV)
+    out = audio_clap(**t)
+    return out.last_hidden_state.float(), out.text_embeds.float()
+
+
+def spectrogram_to_audio(mel, img):
+    """Invert a mel spectrogram image back to audio (Griffin-Lim, 32 iters).
+
+    Same 0..255 -> dB mapping as diffusers' Mel.image_to_audio, with the fast
+    pseudo-inverse mel->STFT step (see the boot comment on audio_mel_pinv).
+    """
+    log_S = (np.frombuffer(img.tobytes(), dtype="uint8").reshape((img.height, img.width)).astype(np.float32)
+             * mel.top_db / 255 - mel.top_db)
+    power = librosa.db_to_power(log_S)
+    stft_mag = np.clip(audio_mel_pinv @ power, 0, None) ** 0.5  # power=2.0 -> magnitude
+    audio = librosa.griffinlim(stft_mag, n_iter=mel.n_iter, hop_length=mel.hop_length,
+                               n_fft=mel.n_fft, window="hann")
+    peak = np.abs(audio).max()
+    return audio if peak == 0 else audio / peak * 0.9
+
+
+@torch.no_grad()
+def sample_audio(prompt: str, steps: int, cfg: float, seed: int, progress=gr.Progress()):
+    global audio_null_seq, audio_null_pool
+    if not prompt.strip():
+        raise gr.Error("Type a prompt first.")
+    steps = int(steps)
+    g = torch.Generator(device=DEV).manual_seed(int(seed))
+    seq, pool = encode_audio([prompt])
+    if audio_null_seq is None:
+        audio_null_seq, audio_null_pool = encode_audio([""])
+    x = torch.randn(1, 1, audio_model.y_res, audio_model.x_res, device=DEV, generator=g)
+    dt = 1.0 / steps
+    for i in progress.tqdm(range(steps), desc="sampling"):
+        t = torch.full((1,), i * dt, device=DEV)
+        vc = audio_model(x, t, seq, pool)
+        vu = audio_model(x, t, audio_null_seq, audio_null_pool)
+        x = x + (vu + cfg * (vc - vu)) * dt
+    row = x[0, 0].clamp(-1, 1).float().cpu().numpy()
+    img = Image.fromarray(((row + 1) * 127.5 + 0.5).astype(np.uint8))
+    audio = spectrogram_to_audio(audio_mel, img)
+    return audio_mel.get_sample_rate(), audio
+
+
 def grid_to_glb(grid: np.ndarray) -> str:
     voxel = trimesh.voxel.VoxelGrid(encoding=grid)
     mesh = voxel.as_boxes()
@@ -141,8 +218,8 @@ def grid_to_glb(grid: np.ndarray) -> str:
 with gr.Blocks(title="BenchLabs Models") as demo:
     gr.Markdown(
         "# BenchLabs Models\n"
-        "Two tiny diffusion models, running live on CPU, no GPU behind this Space. "
-        "Both are under 45M trained parameters, so generation is slower than a hosted API "
+        "Three tiny diffusion models, running live on CPU, no GPU behind this Space. "
+        "All are under 45M trained parameters, so generation is slower than a hosted API "
         "but the whole model fits in a PNG image if you're curious — see the model pages linked below."
     )
     with gr.Tab("Text → Image (PixelModel v6)"):
@@ -188,6 +265,31 @@ with gr.Blocks(title="BenchLabs Models") as demo:
             [vox_prompt, vox_steps, vox_cfg, vox_thresh, vox_seed],
         )
 
+    with gr.Tab("Text → Audio (AudioModel v1)"):
+        gr.Markdown(
+            "17.8 seconds of sound at 22 kHz from a text prompt. The same tiny DiT + "
+            "rectified-flow recipe as the other tabs, applied to a mel spectrogram "
+            "image instead of pixels. Conditioned by the CLAP text tower; audio comes "
+            "back via Griffin-Lim (32 iterations), so expect lo-fi, slightly phasey "
+            "sound — there's no learned vocoder in v1."
+        )
+        with gr.Row():
+            with gr.Column():
+                aud_prompt = gr.Textbox(label="Prompt", placeholder="a dog barking")
+                aud_steps = gr.Slider(10, 100, value=50, step=5, label="Detail (sampling steps)")
+                aud_cfg = gr.Slider(1.0, 10.0, value=4.0, step=0.5, label="Prompt strength (CFG)")
+                aud_seed = gr.Number(value=0, precision=0, label="Seed")
+                aud_btn = gr.Button("Generate audio", variant="primary")
+            with gr.Column():
+                aud_out = gr.Audio(label="Result")
+        aud_btn.click(sample_audio, [aud_prompt, aud_steps, aud_cfg, aud_seed], aud_out)
+        gr.Examples(
+            [["a dog barking", 50, 4.0, 0],
+             ["rain falling on a roof", 50, 4.0, 0],
+             ["footsteps on gravel", 50, 4.0, 0]],
+            [aud_prompt, aud_steps, aud_cfg, aud_seed],
+        )
+
     with gr.Tab("Text → Image (PixelModel v5)"):
         gr.Markdown(
             "Good at material and light: food, landscapes, skies, interiors. "
@@ -212,7 +314,9 @@ with gr.Blocks(title="BenchLabs Models") as demo:
 
     gr.Markdown(
         "Models: [PixelModel v5](https://huggingface.co/bench-labs/PixelModel-v5) · "
-        "[VoxelModel v1](https://huggingface.co/bench-labs/VoxelModel-v1)"
+        "[PixelModel v6](https://huggingface.co/bench-labs/PixelModel-v6) · "
+        "[VoxelModel v1](https://huggingface.co/bench-labs/VoxelModel-v1) · "
+        "[AudioModel v1](https://huggingface.co/bench-labs/AudioModel-v1)"
     )
 
 if __name__ == "__main__":

audio_dit.py

@@ -0,0 +1,137 @@
+from __future__ import annotations
+
+import math
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+def modulate(x, shift, scale):
+    return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
+
+def timestep_embedding(t, dim, max_period=10000):
+    half = dim // 2
+    freqs = torch.exp(-math.log(max_period) * torch.arange(half, device=t.device) / half)
+    args = t[:, None].float() * freqs[None]
+    emb = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+    if dim % 2:
+        emb = torch.cat([emb, torch.zeros_like(emb[:, :1])], dim=-1)
+    return emb
+
+def sincos_2d(dim, grid_h, grid_w):
+    assert dim % 4 == 0
+    gy = np.arange(grid_h, dtype=np.float32)
+    gx = np.arange(grid_w, dtype=np.float32)
+    gyy, gxx = np.meshgrid(gy, gx, indexing="ij")
+    d4 = dim // 4
+    omega = 1.0 / (10000 ** (np.arange(d4, dtype=np.float32) / d4))
+    def emb1(p):
+        out = p.reshape(-1)[:, None] * omega[None]
+        return np.concatenate([np.sin(out), np.cos(out)], axis=1)
+    pe = np.concatenate([emb1(gyy), emb1(gxx)], axis=1)
+    return torch.from_numpy(pe).float()
+
+class Attention(nn.Module):
+    def __init__(self, dim, heads):
+        super().__init__()
+        self.heads = heads
+        self.q = nn.Linear(dim, dim)
+        self.kv = nn.Linear(dim, dim * 2)
+        self.proj = nn.Linear(dim, dim)
+
+    def forward(self, x, ctx=None):
+        ctx = x if ctx is None else ctx
+        B, N, C = x.shape
+        M = ctx.shape[1]
+        h = self.heads
+        q = self.q(x).reshape(B, N, h, C // h).transpose(1, 2)
+        kv = self.kv(ctx).reshape(B, M, 2, h, C // h).permute(2, 0, 3, 1, 4)
+        k, v = kv[0], kv[1]
+        o = F.scaled_dot_product_attention(q, k, v)
+        o = o.transpose(1, 2).reshape(B, N, C)
+        return self.proj(o)
+
+class Block(nn.Module):
+    def __init__(self, dim, heads, mlp_ratio=4.0):
+        super().__init__()
+        self.norm1 = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+        self.attn = Attention(dim, heads)
+        self.norm_ca = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+        self.cross = Attention(dim, heads)
+        self.norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+        hidden = int(dim * mlp_ratio)
+        self.mlp = nn.Sequential(nn.Linear(dim, hidden), nn.GELU(approximate="tanh"),
+                                 nn.Linear(hidden, dim))
+        self.ada = nn.Sequential(nn.SiLU(), nn.Linear(dim, 6 * dim))
+        self.cross_gate = nn.Parameter(torch.zeros(1))
+
+    def forward(self, x, c, text):
+        shift1, scale1, gate1, shift2, scale2, gate2 = self.ada(c).chunk(6, dim=1)
+        x = x + gate1.unsqueeze(1) * self.attn(modulate(self.norm1(x), shift1, scale1))
+        x = x + self.cross_gate * self.cross(self.norm_ca(x), text)
+        x = x + gate2.unsqueeze(1) * self.mlp(modulate(self.norm2(x), shift2, scale2))
+        return x
+
+class AudioDiT(nn.Module):
+    def __init__(self, mel_ch=1, x_res=384, y_res=256, patch=16, dim=384, depth=12,
+                 heads=6, text_seq_dim=768, text_pool_dim=512, mlp_ratio=4.0):
+        super().__init__()
+        assert x_res % patch == 0 and y_res % patch == 0
+        self.mel_ch = mel_ch
+        self.x_res = x_res
+        self.y_res = y_res
+        self.patch = patch
+        self.grid_h = y_res // patch
+        self.grid_w = x_res // patch
+        self.patch_dim = mel_ch * patch * patch
+        self.x_embed = nn.Linear(self.patch_dim, dim)
+        self.register_buffer("pos", sincos_2d(dim, self.grid_h, self.grid_w).unsqueeze(0))
+        self.t_mlp = nn.Sequential(nn.Linear(dim, dim), nn.SiLU(), nn.Linear(dim, dim))
+        self.text_proj = nn.Linear(text_seq_dim, dim)
+        self.text_pool = nn.Linear(text_pool_dim, dim)
+        self.blocks = nn.ModuleList([Block(dim, heads, mlp_ratio) for _ in range(depth)])
+        self.norm_out = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+        self.ada_out = nn.Sequential(nn.SiLU(), nn.Linear(dim, 2 * dim))
+        self.head = nn.Linear(dim, self.patch_dim)
+        self.dim = dim
+        self._init()
+
+    def _init(self):
+        for m in self.modules():
+            if isinstance(m, nn.Linear):
+                nn.init.xavier_uniform_(m.weight)
+                if m.bias is not None:
+                    nn.init.zeros_(m.bias)
+        for b in self.blocks:
+            nn.init.zeros_(b.ada[-1].weight); nn.init.zeros_(b.ada[-1].bias)
+        nn.init.zeros_(self.ada_out[-1].weight); nn.init.zeros_(self.ada_out[-1].bias)
+        nn.init.zeros_(self.head.weight); nn.init.zeros_(self.head.bias)
+
+    def patchify(self, x):
+        B, C, H, W = x.shape
+        p = self.patch
+        x = x.reshape(B, C, H // p, p, W // p, p)
+        x = x.permute(0, 2, 4, 1, 3, 5).reshape(B, (H // p) * (W // p), C * p * p)
+        return x
+
+    def unpatchify(self, x):
+        B, N, _ = x.shape
+        p = self.patch
+        gh, gw = self.grid_h, self.grid_w
+        C = self.mel_ch
+        x = x.reshape(B, gh, gw, C, p, p).permute(0, 3, 1, 4, 2, 5)
+        return x.reshape(B, C, gh * p, gw * p)
+
+    def forward(self, x, t, text_seq, text_pool):
+        x = self.x_embed(self.patchify(x)) + self.pos
+        c = self.t_mlp(timestep_embedding(t, self.dim)) + self.text_pool(text_pool)
+        text = self.text_proj(text_seq)
+        for blk in self.blocks:
+            x = blk(x, c, text)
+        shift, scale = self.ada_out(c).chunk(2, dim=1)
+        x = modulate(self.norm_out(x), shift, scale)
+        x = self.head(x)
+        return self.unpatchify(x)
+
+    def num_params(self):
+        return sum(p.numel() for p in self.parameters())

requirements.txt

@@ -11,3 +11,4 @@ scipy
 audioop-lts
 sentencepiece
 accelerate
+librosa