app.py

from __future__ import annotations

import io
import os
import base64
import sys
import random
from pathlib import Path
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

ROOT = Path(__file__).resolve().parent
SRC_DIR = ROOT / "src"
if str(SRC_DIR) not in sys.path:
    sys.path.insert(0, str(SRC_DIR))

from compound import AXIS_SIZES, N_AXES, SENTINELS, STEP_BOS, STEP_EOS, decode_compound
from compound_model import CompoundGPT, CompoundGPTConfig, default_compound_config

# ── 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
CTX_WINDOW  = 64   # sliding window: keeps cost constant per step

# 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


# ── 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.")}

        # Only 2D tensors = Linear weights (excludes LayerNorm/BN gamma which are 1D)
        weight_keys = sorted(k for k in proj
                             if k.endswith(".weight") and proj[k].dim() == 2)
        print(f"[CODA] text_projection keys found: {weight_keys}")
        if not weight_keys:
            print(f"[CODA] WARNING: no 2D weights in text_projection. Available: {sorted(proj.keys())[:10]}")
            self._proj = nn.Identity()
        else:
            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") and state[k].dim() == 2)
        print(f"[CODA] prefix_projector keys found: {weight_keys}")
        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)

        # out_ln matches PrefixProjector.out_ln in prefix_projector.py
        self._out_ln = nn.LayerNorm(self._gpt_dim)
        ln_w = state.get("out_ln.weight")
        ln_b = state.get("out_ln.bias")
        if ln_w is not None:
            self._out_ln.weight.data.copy_(ln_w)
        if ln_b is not None:
            self._out_ln.bias.data.copy_(ln_b)

    @torch.no_grad()
    def forward(self, text_emb: torch.Tensor) -> torch.Tensor:   # (1,256) → (1,8,768)
        out = self._net(text_emb).view(text_emb.shape[0], self.n_prefix, self._gpt_dim)
        return self._out_ln(out)


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)

        clap_state   = clap_ckpt.get("model_state_dict", clap_ckpt)

        # Prefix projector may be saved under different keys depending on training script
        prefix_state = None
        for key in ["model_state_dict", "projector_state_dict", "prefix_state_dict", "state_dict"]:
            if key in prefix_ckpt:
                prefix_state = prefix_ckpt[key]
                print(f"[CODA] prefix ckpt key used: '{key}'")
                break
        if prefix_state is None:
            # Top-level might already be the state dict — check if values are tensors
            if any(hasattr(v, "dim") for v in prefix_ckpt.values()):
                prefix_state = prefix_ckpt
                print("[CODA] prefix ckpt: using top-level as state dict")
            else:
                print(f"[CODA] prefix ckpt keys: {list(prefix_ckpt.keys())}")
                raise KeyError("Cannot find projector weights in prefix checkpoint")

        text_enc     = _TextEncoder(_LightStateDict(clap_state))
        prefix_proj  = _PrefixProjector(_LightStateDict(prefix_state))
        text_enc.eval(); prefix_proj.eval()

        _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


class _LightStateDict(dict):
    """Passthrough so _TextEncoder / _PrefixProjector constructors work with raw state dicts."""
    pass


# ── Compound step embeddings (needed to prepend prefix) ───────────────────────

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.d_model
    emb = torch.zeros(B, T, d)

    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


# ── 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(
    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]]:

    bos = list(SENTINELS); bos[0] = STEP_BOS
    generated: List[List[int]] = [bos]
    conditioned = prefix_embs is not None

    for _ in range(max_steps):
        # Sliding window: only feed last CTX_WINDOW steps to keep cost constant
        ctx      = generated[-CTX_WINDOW:] if len(generated) > CTX_WINDOW else generated
        step_ids = torch.tensor([ctx], dtype=torch.long)   # (1, ≤CTX_WINDOW, 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

        if conditioned:
            try:
                step_e  = _compound_embeds(model, step_ids)
                full_e  = torch.cat([prefix_embs, step_e], dim=1)  # (1, 8+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):
                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:
            break
        generated.append(next_step)

    return generated


# ── 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")

    if notes:
        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_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(pad=0.6)
    fig.savefig(buf, format="png", facecolor="#F9F8F5")
    plt.close(fig)
    buf.seek(0)
    return Image.open(buf).convert("RGB")




# ── MIDI player HTML ──────────────────────────────────────────────────────────

def _midi_player_html(midi_path: str) -> str:
    """
    Inline MIDI player using @tonejs/midi + Web Audio oscillators.
    Converts MIDI to audio in the browser. Loads the parser from cdn.jsdelivr.net
    (matches Hugging Face Spaces allow_list — unpkg is typically blocked).
    """
    with open(midi_path, "rb") as f:
        b64 = base64.b64encode(f.read()).decode()
    return f"""
<div id="coda-player" style="font-family:sans-serif;padding:8px 0">
  <div style="display:flex;align-items:center;gap:10px;flex-wrap:wrap">
    <button id="coda-play"
      style="padding:7px 20px;background:#534AB7;color:#fff;border:none;
             border-radius:6px;cursor:pointer;font-size:14px;font-weight:500">
      ▶ Play
    </button>
    <button id="coda-stop"
      style="padding:7px 16px;background:#F1EFE8;color:#2C2C2A;border:1px solid #D3D1C7;
             border-radius:6px;cursor:pointer;font-size:14px">
      ■ Stop
    </button>
    <span id="coda-status"
      style="font-size:12px;color:#888780">
      Loading player…
    </span>
  </div>
  <div id="coda-progress-wrap"
    style="margin-top:8px;height:4px;background:#E8E6DF;border-radius:2px;overflow:hidden">
    <div id="coda-progress"
      style="height:100%;width:0%;background:#534AB7;transition:width 0.25s linear;border-radius:2px">
    </div>
  </div>
</div>

<script>
(function() {{
  const B64 = "{b64}";
  const STATUS = document.getElementById("coda-status");
  const PLAY   = document.getElementById("coda-play");
  const STOP   = document.getElementById("coda-stop");
  const PROG   = document.getElementById("coda-progress");

  // Decode base64 MIDI → ArrayBuffer
  function b64ToBuffer(b64) {{
    const bin = atob(b64);
    const buf = new Uint8Array(bin.length);
    for (let i = 0; i < bin.length; i++) buf[i] = bin.charCodeAt(i);
    return buf.buffer;
  }}

  // Simple MIDI player using Web MIDI API if available, else AudioContext beeps
  let audioCtx = null;
  let sources  = [];
  let playing  = false;
  let startTime = 0;
  let duration  = 0;
  let rafId     = null;

  function stopAll() {{
    sources.forEach(s => {{ try {{ s.stop(); }} catch(e) {{}} }});
    sources = [];
    playing = false;
    if (rafId) cancelAnimationFrame(rafId);
    PROG.style.width = "0%";
    PLAY.textContent = "▶ Play";
  }}

  async function loadAndPlay() {{
    if (playing) {{ stopAll(); return; }}

    PLAY.textContent = "Loading…";
    STATUS.textContent = "Parsing MIDI…";

    // Load @tonejs/midi from jsDelivr (on HF Spaces, allow_list usually permits this host only)
    if (typeof Midi === "undefined") {{
      await new Promise((resolve, reject) => {{
        const s = document.createElement("script");
        s.src = "https://cdn.jsdelivr.net/npm/@tonejs/midi@2.0.28/build/Midi.js";
        s.crossOrigin = "anonymous";
        s.onload = resolve;
        s.onerror = () => reject(new Error("Could not load MIDI parser (CDN blocked?)"));
        document.head.appendChild(s);
      }});
    }}

    let midi;
    try {{
      midi = new Midi(b64ToBuffer(B64));
    }} catch(e) {{
      STATUS.textContent = "Failed to parse MIDI: " + e.message;
      PLAY.textContent = "▶ Play";
      return;
    }}

    audioCtx = audioCtx || new (window.AudioContext || window.webkitAudioContext)();
    if (audioCtx.state === "suspended") await audioCtx.resume();

    // Schedule all notes using OscillatorNode (simple synth — no soundfont)
    const scheduleStart = audioCtx.currentTime + 0.05;
    sources = [];
    duration = typeof midi.duration === "number" && midi.duration > 0 ? midi.duration : 0;

    midi.tracks.forEach((track, ti) => {{
      const waveforms = ["sine","triangle","square","sawtooth"];
      const wave = waveforms[ti % waveforms.length];
      track.notes.forEach(note => {{
        const t0 = scheduleStart + note.time;
        const t1 = t0 + Math.max(note.duration, 0.02);
        const vel = Math.max(Number(note.velocity) || 0, 0.02);
        const peak = Math.min(0.35, vel * 0.22);
        const freq = 440 * Math.pow(2, (note.midi - 69) / 12);
        const osc  = audioCtx.createOscillator();
        const gain = audioCtx.createGain();
        osc.type = wave;
        osc.frequency.value = freq;
        gain.gain.setValueAtTime(0.0001, t0);
        gain.gain.exponentialRampToValueAtTime(peak, t0 + 0.02);
        gain.gain.exponentialRampToValueAtTime(0.0001, t1 + 0.04);
        osc.connect(gain);
        gain.connect(audioCtx.destination);
        osc.start(t0);
        osc.stop(t1 + 0.06);
        sources.push(osc);
        duration = Math.max(duration, note.time + note.duration);
      }});
    }});

    if (duration <= 0) {{
      STATUS.textContent = "No notes in MIDI";
      PLAY.textContent = "▶ Play";
      return;
    }}

    playing   = true;
    startTime = scheduleStart;
    PLAY.textContent  = "⏹ Stop";
    STATUS.textContent = `Playing · ${{Math.round(duration)}}s`;

    function tick() {{
      if (!playing) return;
      const elapsed = audioCtx.currentTime - startTime;
      const pct = duration > 0 ? Math.min(100, (elapsed / duration) * 100) : 100;
      PROG.style.width = pct + "%";
      if (elapsed < duration + 0.15) {{
        rafId = requestAnimationFrame(tick);
      }} else {{
        stopAll();
        STATUS.textContent = "Done";
      }}
    }}
    tick();
  }}

  PLAY.addEventListener("click", loadAndPlay);
  STOP.addEventListener("click", () => {{ stopAll(); STATUS.textContent = "Stopped"; }});
  STATUS.textContent = "Ready";
}})();
</script>
"""

# ── Main callable ─────────────────────────────────────────────────────────────

def generate(
    prompt:      str,
    temperature: float,
    top_k:       int,
    top_p:       float,
    max_steps:   int,
):
    prompt = (prompt or "").strip()
    if not prompt:
        raise gr.Error("Please enter a text prompt describing the music you want.")

    try:
        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

    pm = decode_compound(steps)
    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}"
    )

    player_html = _midi_player_html(TMP_MIDI)
    return image, TMP_MIDI, player_html, info


# ── 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") 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>"
    )

    prompt = gr.Textbox(
        placeholder="Describe the music you want to generate…",
        lines=2,
        show_label=False,
        elem_classes=["prompt-box"],
    )

    gr.Examples(
        examples=[[p] for p in EXAMPLES],
        inputs=[prompt],
        label="Example prompts",
        examples_per_page=8,
    )

    with gr.Row():
        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=150,  step=50,   label="Max steps")

    gen_btn = gr.Button("Generate", variant="primary", elem_classes=["generate-btn"])

    roll_out = gr.Image(type="pil", label="Piano roll")

    with gr.Row():
        midi_out = gr.File(label="Download MIDI")

    player_out = gr.HTML("")
    info_out = gr.Markdown("", elem_classes=["info-md"])

    gen_btn.click(
        fn=generate,
        inputs=[prompt, temperature, top_k, top_p, max_steps],
        outputs=[roll_out, midi_out, player_out, info_out],
    )


if __name__ == "__main__":
    demo.launch(css=CSS)