bug fix

app.py

@@ -2,16 +2,20 @@ from __future__ import annotations
 
 import io
 import os
-import random
 import sys
+import random
 from pathlib import Path
-from typing import Dict, List
+from typing import Dict, List, Optional, Tuple
 
 import gradio as gr
+import matplotlib
+matplotlib.use("Agg")
 import matplotlib.pyplot as plt
+import matplotlib.patches as mpatches
 import numpy as np
 import pretty_midi
 import torch
+import torch.nn as nn
 import torch.nn.functional as F
 from huggingface_hub import hf_hub_download
 from PIL import Image
@@ -24,262 +28,431 @@ if str(SRC_DIR) not in sys.path:
 from compound import AXIS_SIZES, N_AXES, SENTINELS, STEP_BOS, STEP_EOS, decode_compound
 from compound_model import CompoundGPT, CompoundGPTConfig, default_compound_config
 
-HF_REPO_ID = os.getenv("HF_REPO_ID", "prajanya23/bachgpt-midi")
-HF_TOKEN = os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACE_HUB_TOKEN")
-TMP_MIDI_PATH = "/tmp/output.mid"
+# ── Config ────────────────────────────────────────────────────────────────────
+HF_REPO_ID  = os.getenv("HF_REPO_ID",  "Prajanya23/Coda")
+HF_TOKEN    = os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACE_HUB_TOKEN")
+TMP_MIDI    = "/tmp/coda_output.mid"
+N_PREFIX    = 8
+CLAP_DIM    = 256
+
+# Checkpoint paths inside the HF repo (set via Space secrets to override)
+GPT_FILE    = os.getenv("GPT_CKPT",    "checkpoints/compound_best.pt")
+CLAP_FILE   = os.getenv("CLAP_CKPT",   "checkpoints/clap_compound_best.pt")
+PREFIX_FILE = os.getenv("PREFIX_CKPT", "checkpoints/prefix_projector_best.pt")
+
+EXAMPLES = [
+    "a slow melancholic piano piece in a minor key with sparse flowing notes",
+    "an upbeat jazz trio with piano bass and drums syncopated and energetic",
+    "ambient electronic music with synthesizer pads slow and atmospheric",
+    "fast energetic rock band with electric guitar and drums",
+    "a gentle classical piece for piano and strings moderate tempo",
+    "a funky groove with bass guitar and brass instruments",
+    "a soft acoustic guitar piece fingerpicked quiet and introspective",
+    "an orchestral piece with strings and brass building to a climax",
+]
+
+VOICE_COLORS = ["#534AB7", "#0F6E56", "#BA7517", "#993C1D",
+                "#185FA5", "#639922", "#A32D2D", "#D4537E"]
+
+# ── Model cache ───────────────────────────────────────────────────────────────
+_CACHE: Dict[str, object] = {}
+
+
+def _dl(filename: str) -> str:
+    return hf_hub_download(repo_id=HF_REPO_ID, filename=filename, token=HF_TOKEN)
+
+
+# ── GPT loader ────────────────────────────────────────────────────────────────
+
+def _load_gpt() -> CompoundGPT:
+    if "gpt" in _CACHE:
+        return _CACHE["gpt"]                                            # type: ignore[return-value]
+    ckpt = torch.load(_dl(GPT_FILE), map_location="cpu", weights_only=True)
+    cfg  = default_compound_config()
+    raw  = ckpt.get("config") if isinstance(ckpt, dict) else None
+    if isinstance(raw, dict):
+        for k, v in raw.items():
+            if hasattr(cfg, k):
+                setattr(cfg, k, v)
+    model = CompoundGPT(cfg)
+    model.load_state_dict(ckpt.get("model_state_dict", ckpt), strict=False)
+    model.eval()
+    _CACHE["gpt"] = model
+    return model
 
-# If your uploaded filenames differ, set these in Space variables.
-V1_FILENAME = os.getenv("V1_CKPT_FILENAME", "")
-V2_FILENAME = os.getenv("V2_CKPT_FILENAME", "")
 
-CHECKPOINT_CANDIDATES: Dict[str, List[str]] = {
-    "v1": [V1_FILENAME, "compound_v1_3m.pt", "compound_3m.pt", "v1_3m.pt"],
-    "v2": [V2_FILENAME, "compound_best.pt", "compound_v2_25m.pt", "compound_25m.pt"],
-}
+# ── CLAP text encoder (lightweight — no CompoundGPT inside) ──────────────────
+
+class _TextEncoder(nn.Module):
+    """Sentence-transformer + CLAP text projection, reconstructed from checkpoint."""
+
+    def __init__(self, clap_state: dict):
+        super().__init__()
+        from sentence_transformers import SentenceTransformer
+        self._st = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
+
+        # Extract text_projection weights from CLAP state dict
+        proj = {k[len("text_projection."):]: v
+                for k, v in clap_state.items()
+                if k.startswith("text_projection.")}
+
+        weight_keys = sorted(k for k in proj if k.endswith(".weight"))
+        layers: List[nn.Module] = []
+        for i, wk in enumerate(weight_keys):
+            w  = proj[wk]
+            bk = wk.replace(".weight", ".bias")
+            lin = nn.Linear(w.shape[1], w.shape[0], bias=(bk in proj))
+            lin.weight.data.copy_(w)
+            if bk in proj:
+                lin.bias.data.copy_(proj[bk])
+            layers.append(lin)
+            if i < len(weight_keys) - 1:
+                layers.append(nn.ReLU())
+        self._proj = nn.Sequential(*layers)
+
+    @torch.no_grad()
+    def encode(self, text: str) -> torch.Tensor:
+        raw  = self._st.encode([text], convert_to_tensor=True, show_progress_bar=False)
+        emb  = self._proj(raw.float())
+        return F.normalize(emb, dim=-1)          # (1, 256)
+
+
+# ── Prefix projector (reconstructed from checkpoint) ─────────────────────────
+
+class _PrefixProjector(nn.Module):
+    def __init__(self, state: dict, n_prefix: int = N_PREFIX):
+        super().__init__()
+        self.n_prefix = n_prefix
+        weight_keys   = sorted(k for k in state if k.endswith(".weight"))
+        self._gpt_dim = state[weight_keys[-1]].shape[0] // n_prefix
+
+        layers: List[nn.Module] = []
+        for i, wk in enumerate(weight_keys):
+            w  = state[wk]
+            bk = wk.replace(".weight", ".bias")
+            lin = nn.Linear(w.shape[1], w.shape[0], bias=(bk in state))
+            lin.weight.data.copy_(w)
+            if bk in state:
+                lin.bias.data.copy_(state[bk])
+            layers.append(lin)
+            if i < len(weight_keys) - 1:
+                layers.append(nn.GELU())
+        self._net = nn.Sequential(*layers)
+
+    @torch.no_grad()
+    def forward(self, text_emb: torch.Tensor) -> torch.Tensor:   # (1,256) → (1,8,768)
+        return self._net(text_emb).view(text_emb.shape[0], self.n_prefix, self._gpt_dim)
+
+
+def _load_clap() -> Tuple[Optional[_TextEncoder], Optional[_PrefixProjector]]:
+    if "text_enc" in _CACHE:
+        return _CACHE.get("text_enc"), _CACHE.get("prefix_proj")    # type: ignore
 
+    try:
+        # weights_only=False needed because args is argparse.Namespace
+        clap_ckpt   = torch.load(_dl(CLAP_FILE),   map_location="cpu", weights_only=False)
+        prefix_ckpt = torch.load(_dl(PREFIX_FILE), map_location="cpu", weights_only=False)
 
-def _download_checkpoint(model_key: str) -> str:
-    last_error: Exception | None = None
-    for filename in CHECKPOINT_CANDIDATES[model_key]:
-        if not filename:
-            continue
-        try:
-            return hf_hub_download(
-                repo_id=HF_REPO_ID,
-                filename=filename,
-                token=HF_TOKEN,
-            )
-        except Exception as exc:  # pragma: no cover - defensive for remote failures
-            last_error = exc
-    raise RuntimeError(
-        f"Unable to download {model_key} checkpoint from repo '{HF_REPO_ID}'. "
-        f"Tried: {[f for f in CHECKPOINT_CANDIDATES[model_key] if f]}. "
-        "If this is a private/gated repo, set HF_TOKEN (or HUGGINGFACE_HUB_TOKEN) "
-        "with read access and optionally set HF_REPO_ID to the correct model repo. "
-        f"Last error: {last_error}"
-    )
+        clap_state   = clap_ckpt.get("model_state_dict",   clap_ckpt)
+        prefix_state = prefix_ckpt.get("model_state_dict", prefix_ckpt)
 
+        text_enc     = _TextEncoder(_LightStateDict(clap_state))
+        prefix_proj  = _PrefixProjector(_LightStateDict(prefix_state))
+        text_enc.eval(); prefix_proj.eval()
 
-def _load_compound_model(ckpt_path: str) -> CompoundGPT:
-    # CPU-only per requirement.
-    ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=True)
-    cfg = default_compound_config()
-    raw_cfg = ckpt.get("config") if isinstance(ckpt, dict) else None
-    if isinstance(raw_cfg, dict):
-        for field in CompoundGPTConfig.__dataclass_fields__.keys():
-            if field in raw_cfg:
-                setattr(cfg, field, raw_cfg[field])
-    model = CompoundGPT(cfg).to("cpu")
-    state = ckpt.get("model_state_dict", ckpt)
-    model.load_state_dict(state, strict=False)
-    model.eval()
-    return model
+        _CACHE["text_enc"]    = text_enc
+        _CACHE["prefix_proj"] = prefix_proj
+        print("[CODA] CLAP text conditioning loaded ✓")
+        return text_enc, prefix_proj
 
+    except Exception as exc:
+        print(f"[CODA] CLAP not available ({exc}). Falling back to unconditioned generation.")
+        _CACHE["text_enc"] = _CACHE["prefix_proj"] = None
+        return None, None
 
-def _sample_axis(logits: torch.Tensor, temperature: float, top_k: int) -> int:
-    scaled = logits / temperature
-    if top_k > 0 and top_k < scaled.numel():
-        values, _ = torch.topk(scaled, top_k)
-        cutoff = values[-1]
-        scaled = torch.where(
-            scaled < cutoff,
-            torch.tensor(float("-inf"), device=scaled.device),
-            scaled,
-        )
-    probs = F.softmax(scaled, dim=-1)
-    return int(torch.multinomial(probs, num_samples=1).item())
 
+class _LightStateDict(dict):
+    """Passthrough so _TextEncoder / _PrefixProjector constructors work with raw state dicts."""
+    pass
 
-def _random_seed_step() -> List[int]:
-    step = [random.randrange(size) for size in AXIS_SIZES]
-    if step[0] == STEP_EOS:
-        step[0] = STEP_BOS
-    return step
 
+# ── Compound step embeddings (needed to prepend prefix) ───────────────────────
 
-def _parse_seed_text(seed_text: str) -> List[List[int]]:
-    # Format: one or more steps, separated by ';'
-    # Each step: "a,b,c,d,e,f,g" (7 axes)
-    text = seed_text.strip()
-    if not text:
-        return [_random_seed_step()]
+def _compound_embeds(model: CompoundGPT, step_ids: torch.Tensor) -> torch.Tensor:
+    """
+    Sum per-axis embeddings to get (B, T, n_embd) float tensor.
+    Searches common attribute names used in CompoundGPT implementations.
+    """
+    B, T, _ = step_ids.shape
+    d = model.config.n_embd
+    emb = torch.zeros(B, T, d)
 
-    steps: List[List[int]] = []
-    chunks = [chunk.strip() for chunk in text.split(";") if chunk.strip()]
-    for chunk in chunks:
-        parts = [p.strip() for p in chunk.split(",")]
-        if len(parts) != N_AXES:
-            continue
-        try:
-            step = [int(p) for p in parts]
-        except ValueError:
-            continue
-        bounded = [
-            max(0, min(value, AXIS_SIZES[idx] - 1))
-            for idx, value in enumerate(step)
-        ]
-        steps.append(bounded)
+    axis_list = None
+    if hasattr(model, "input_embeds"):
+        axis_list = model.input_embeds
+    else:
+        for attr in ["axis_embeds", "axis_embed", "embed_axes", "wtes"]:
+            if hasattr(model, attr):
+                axis_list = getattr(model, attr); break
+            if hasattr(model, "transformer") and hasattr(model.transformer, attr):
+                axis_list = getattr(model.transformer, attr); break
+
+    if axis_list is None:
+        raise AttributeError(
+            "Cannot find axis embedding layers in CompoundGPT (tried 'input_embeds' "
+            "and common fallbacks). Check compound_model.py."
+        )
+
+    for i, layer in enumerate(axis_list):
+        emb = emb + layer(step_ids[:, :, i])
+    return emb
 
-    return steps if steps else [_random_seed_step()]
 
+# ── Sampling ──────────────────────────────────────────────────────────────────
+
+def _sample(logits: torch.Tensor, temp: float, top_k: int, top_p: float) -> int:
+    scaled = logits / max(temp, 1e-6)
+    # top-k
+    if 0 < top_k < scaled.numel():
+        v, _ = torch.topk(scaled, top_k)
+        scaled = scaled.masked_fill(scaled < v[-1], float("-inf"))
+    probs = F.softmax(scaled, dim=-1)
+    # top-p (nucleus)
+    if 0.0 < top_p < 1.0:
+        sp, si = torch.sort(probs, descending=True)
+        cum = torch.cumsum(sp, dim=-1)
+        sp[cum - sp > top_p] = 0.0
+        probs = torch.zeros_like(scaled).scatter_(0, si, sp)
+        probs = probs / probs.sum().clamp(min=1e-8)
+    return int(torch.multinomial(probs, 1).item())
+
+
+# ── Generation ────────────────────────────────────────────────────────────────
 
 @torch.no_grad()
-def _generate_steps(
-    model: CompoundGPT,
-    seed_steps: List[List[int]],
-    temperature: float,
-    top_k: int,
-    max_new_steps: int,
+def _generate(
+    model:        CompoundGPT,
+    prefix_embs:  Optional[torch.Tensor],   # (1, N_PREFIX, n_embd) or None
+    temperature:  float,
+    top_k:        int,
+    top_p:        float,
+    max_steps:    int,
 ) -> List[List[int]]:
-    generated = [list(s) for s in seed_steps]
-    if not generated:
-        bos = list(SENTINELS)
-        bos[0] = STEP_BOS
-        generated.append(bos)
-
-    for _ in range(max_new_steps):
-        step_ids = torch.tensor([generated], dtype=torch.long, device="cpu")
-        if step_ids.size(1) > model.config.block_size:
+
+    bos = list(SENTINELS); bos[0] = STEP_BOS
+    generated: List[List[int]] = [bos]
+    conditioned = prefix_embs is not None
+
+    for _ in range(max_steps):
+        step_ids = torch.tensor([generated], dtype=torch.long)   # (1, T, 7)
+        T = step_ids.shape[1]
+        n_pre = prefix_embs.shape[1] if conditioned else 0
+
+        if T + n_pre > model.config.block_size:
             break
-        position_ids = torch.arange(step_ids.size(1), dtype=torch.long).unsqueeze(0)
-        logits_per_axis = model(idx=step_ids, position_ids=position_ids)
-
-        next_step: List[int] = []
-        for axis_logits in logits_per_axis:
-            sampled = _sample_axis(
-                logits=axis_logits[0, -1, :],
-                temperature=temperature,
-                top_k=top_k,
-            )
-            next_step.append(sampled)
+
+        if conditioned:
+            try:
+                step_e    = _compound_embeds(model, step_ids)        # (1, T, d)
+                full_e    = torch.cat([prefix_embs, step_e], dim=1)  # (1, n_pre+T, d)
+                pos_ids   = torch.arange(full_e.shape[1]).unsqueeze(0)
+                logits    = model(inputs_embeds=full_e, position_ids=pos_ids)
+            except (AttributeError, TypeError):
+                # Embedding layer name mismatch — fall back silently
+                conditioned = False
+                prefix_embs = None
+
+        if not conditioned:
+            pos_ids = torch.arange(T).unsqueeze(0)
+            logits  = model(idx=step_ids, position_ids=pos_ids)
+
+        next_step = [
+            _sample(ax[0, -1, :], temperature, top_k, top_p)
+            for ax in logits
+        ]
 
         if next_step[0] == STEP_EOS:
-            next_step = [STEP_EOS] + SENTINELS[1:]
-            generated.append(next_step)
             break
-
         generated.append(next_step)
 
     return generated
 
 
-def _render_pianoroll(pm: pretty_midi.PrettyMIDI) -> Image.Image:
-    notes = []
-    for inst in pm.instruments:
-        for note in inst.notes:
-            notes.append((note.start, note.end, note.pitch, note.velocity))
+# ── Piano roll ────────────────────────────────────────────────────────────────
+
+def _piano_roll(pm: pretty_midi.PrettyMIDI) -> Image.Image:
+    notes = [
+        (n.start, n.end, n.pitch, n.velocity, i)
+        for i, inst in enumerate(pm.instruments)
+        for n in inst.notes
+    ]
+
+    fig, ax = plt.subplots(figsize=(12, 3.2), dpi=140)
+    fig.patch.set_facecolor("#F9F8F5")
+    ax.set_facecolor("#F9F8F5")
 
-    fig, ax = plt.subplots(figsize=(12, 4), dpi=120)
     if notes:
-        for start, end, pitch, velocity in notes:
-            width = max(0.01, end - start)
-            color = plt.cm.magma(velocity / 127.0)
-            ax.broken_barh([(start, width)], (pitch - 0.45, 0.9), facecolors=color)
-        max_end = max(n[1] for n in notes)
-        min_pitch = min(n[2] for n in notes) - 2
-        max_pitch = max(n[2] for n in notes) + 2
-        ax.set_xlim(0.0, max_end + 0.25)
-        ax.set_ylim(min_pitch, max_pitch)
-    else:
-        ax.text(
-            0.5,
-            0.5,
-            "No notes generated",
-            ha="center",
-            va="center",
-            transform=ax.transAxes,
-        )
+        min_p = max(0,   min(n[2] for n in notes) - 3)
+        max_p = min(127, max(n[2] for n in notes) + 3)
+        max_t = max(n[1] for n in notes)
+
+        for start, end, pitch, vel, vi in notes:
+            ax.broken_barh(
+                [(start, max(0.02, end - start))],
+                (pitch - 0.42, 0.84),
+                facecolors=VOICE_COLORS[vi % len(VOICE_COLORS)],
+                alpha=0.45 + 0.55 * vel / 127,
+                linewidth=0,
+            )
 
-    ax.set_title("Generated Piano Roll")
-    ax.set_xlabel("Time (seconds)")
-    ax.set_ylabel("Pitch (MIDI)")
-    ax.grid(alpha=0.2)
+        ax.set_xlim(0, max_t + 0.2)
+        ax.set_ylim(min_p, max_p)
+
+        n_inst = len(pm.instruments)
+        if n_inst > 1:
+            patches = [
+                mpatches.Patch(color=VOICE_COLORS[i % len(VOICE_COLORS)],
+                               label=pm.instruments[i].name or f"voice {i+1}")
+                for i in range(n_inst)
+            ]
+            ax.legend(handles=patches, fontsize=7, loc="upper right",
+                      framealpha=0.8, edgecolor="none", facecolor="#F9F8F5")
+    else:
+        ax.text(0.5, 0.5, "No notes generated", ha="center", va="center",
+                transform=ax.transAxes, color="#888780", fontsize=11)
+
+    ax.set_xlabel("Time (s)", fontsize=9, color="#5F5E5A")
+    ax.set_ylabel("Pitch",    fontsize=9, color="#5F5E5A")
+    ax.tick_params(labelsize=8, colors="#5F5E5A")
+    for sp in ax.spines.values():
+        sp.set_visible(False)
+    ax.spines["bottom"].set_visible(True)
+    ax.spines["left"].set_visible(True)
+    ax.spines["bottom"].set_color("#D3D1C7"); ax.spines["bottom"].set_linewidth(0.5)
+    ax.spines["left"].set_color("#D3D1C7");   ax.spines["left"].set_linewidth(0.5)
+    ax.grid(axis="y", alpha=0.12, linewidth=0.5, color="#888780")
 
     buf = io.BytesIO()
-    fig.tight_layout()
-    fig.savefig(buf, format="png")
+    fig.tight_layout(pad=0.6)
+    fig.savefig(buf, format="png", facecolor="#F9F8F5")
     plt.close(fig)
     buf.seek(0)
     return Image.open(buf).convert("RGB")
 
 
-MODELS: Dict[str, CompoundGPT] = {}
-
-
-def _get_model(model_key: str) -> CompoundGPT:
-    if model_key in MODELS:
-        return MODELS[model_key]
-    ckpt_path = _download_checkpoint(model_key)
-    model = _load_compound_model(ckpt_path)
-    MODELS[model_key] = model
-    return model
-
+# ── Main callable ─────────────────────────────────────────────────────────────
 
 def generate(
-    model_choice: str,
+    prompt:      str,
     temperature: float,
-    top_k: int,
-    max_new_tokens: int,
-    seed_text: str,
+    top_k:       int,
+    top_p:       float,
+    max_steps:   int,
 ):
-    model_key = "v1" if model_choice.startswith("v1") else "v2"
+    prompt = (prompt or "").strip()
+    if not prompt:
+        raise gr.Error("Please enter a text prompt describing the music you want.")
+
     try:
-        model = _get_model(model_key)
-    except RuntimeError as exc:
-        raise gr.Error(str(exc)) from exc
-
-    seed_steps = _parse_seed_text(seed_text)
-    steps = _generate_steps(
-        model=model,
-        seed_steps=seed_steps,
-        temperature=float(temperature),
-        top_k=int(top_k),
-        max_new_steps=int(max_new_tokens),
+        model = _load_gpt()
+    except Exception as exc:
+        raise gr.Error(f"Failed to load model: {exc}") from exc
+
+    text_enc, prefix_proj = _load_clap()
+    prefix_embs = None
+    mode        = "unconditioned"
+
+    if text_enc is not None and prefix_proj is not None:
+        try:
+            text_emb    = text_enc.encode(prompt)          # (1, 256)
+            prefix_embs = prefix_proj(text_emb)            # (1, 8, 768)
+            mode        = "CLAP-conditioned"
+        except Exception as exc:
+            print(f"[CODA] Text conditioning failed: {exc}")
+
+    steps = _generate(
+        model=model, prefix_embs=prefix_embs,
+        temperature=float(temperature), top_k=int(top_k),
+        top_p=float(top_p), max_steps=int(max_steps),
     )
-    steps = [s for s in steps if int(s[0]) != 9]  # drop STEP_PB for decode safety
+    steps = [s for s in steps if int(s[0]) != 9]   # drop STEP_PB
 
     pm = decode_compound(steps)
-    pm.write(TMP_MIDI_PATH)
-    image = _render_pianoroll(pm)
-    return image, TMP_MIDI_PATH
+    pm.write(TMP_MIDI)
+    image = _piano_roll(pm)
+
+    all_notes = [n for inst in pm.instruments for n in inst.notes]
+    n_notes   = len(all_notes)
+    n_voices  = len(pm.instruments)
+    duration  = round(max((n.end for n in all_notes), default=0.0), 1)
+    pitches   = [n.pitch for n in all_notes]
+    p_std     = round(float(np.std(pitches)), 1) if pitches else 0.0
+
+    info = (
+        f"**Mode:** {mode} &nbsp;|&nbsp; "
+        f"**Notes:** {n_notes} &nbsp;|&nbsp; "
+        f"**Voices:** {n_voices} &nbsp;|&nbsp; "
+        f"**Duration:** {duration}s &nbsp;|&nbsp; "
+        f"**Pitch σ:** {p_std}"
+    )
 
+    return image, TMP_MIDI, info
 
-with gr.Blocks(title="CODA MIDI Generator") as demo:
-    gr.Markdown("# CODA MIDI Generator")
-    gr.Markdown("CPU inference: generation usually takes **30–60s**.")
 
-    with gr.Row():
-        model_choice = gr.Radio(
-            choices=["v1 (3M)", "v2 (25M)"],
-            value="v2 (25M)",
-            label="Model",
-        )
-        temperature = gr.Slider(0.5, 1.5, value=1.0, step=0.01, label="Temperature")
-        top_k = gr.Slider(1, 100, value=50, step=1, label="Top-k")
-
-    max_new_tokens = gr.Slider(
-        128,
-        1024,
-        value=512,
-        step=1,
-        label="Max new tokens",
+# ── UI ────────────────────────────────────────────────────────────────────────
+
+CSS = """
+.prompt-box textarea       { font-size: 15px !important; line-height: 1.6 !important; }
+.generate-btn              { font-size: 15px !important; }
+.info-md p                 { font-size: 13px; color: var(--body-text-color-subdued); }
+.gr-examples table td      { font-size: 12px; }
+footer                     { display: none !important; }
+"""
+
+with gr.Blocks(title="CODA", css=CSS) as demo:
+
+    gr.Markdown("## CODA")
+    gr.Markdown(
+        "Text-conditioned symbolic MIDI generation · compound tokenization + CLAP alignment  \n"
+        "<small>CPU inference — generation takes ~30–60 s</small>"
     )
-    seed_text = gr.Textbox(
-        label="Optional seed text",
-        placeholder="Optional compound seed. Format: a,b,c,d,e,f,g; ...  (leave empty for random vocab seed)",
+
+    prompt = gr.Textbox(
+        placeholder="Describe the music you want to generate…",
         lines=2,
+        show_label=False,
+        elem_classes=["prompt-box"],
     )
 
-    run_btn = gr.Button("Generate", variant="primary")
+    gr.Examples(
+        examples=[[p] for p in EXAMPLES],
+        inputs=[prompt],
+        label="Example prompts",
+        examples_per_page=8,
+    )
 
     with gr.Row():
-        roll_out = gr.Image(type="pil", label="Piano Roll")
-        midi_out = gr.File(label="MIDI")
+        temperature = gr.Slider(0.5, 1.5, value=0.9,  step=0.05, label="Temperature")
+        top_k       = gr.Slider(1,   100, value=40,   step=1,    label="Top-k")
+
+    with gr.Accordion("Advanced", open=False):
+        with gr.Row():
+            top_p     = gr.Slider(0.5, 1.0, value=0.92, step=0.01, label="Top-p (nucleus)")
+            max_steps = gr.Slider(64,  512, value=256,  step=64,   label="Max steps")
+
+    gen_btn = gr.Button("Generate", variant="primary", elem_classes=["generate-btn"])
+
+    roll_out = gr.Image(type="pil", label="Piano roll", show_download_button=True)
+
+    with gr.Row():
+        midi_out = gr.File(label="Download MIDI")
+
+    info_out = gr.Markdown("", elem_classes=["info-md"])
 
-    run_btn.click(
+    gen_btn.click(
         fn=generate,
-        inputs=[model_choice, temperature, top_k, max_new_tokens, seed_text],
-        outputs=[roll_out, midi_out],
+        inputs=[prompt, temperature, top_k, top_p, max_steps],
+        outputs=[roll_out, midi_out, info_out],
     )
 
 

requirements.txt

@@ -1,9 +1,8 @@
---extra-index-url https://download.pytorch.org/whl/cpu
-torch==2.11.0
-gradio>=4.0.0
-huggingface_hub>=0.20.0
-pretty_midi>=0.2.11
-matplotlib>=3.7
-numpy>=1.24
-Pillow>=10.0.0
-mido>=1.3.0
+torch
+gradio
+pretty_midi
+matplotlib
+numpy
+Pillow
+huggingface_hub
+sentence-transformers