Add MiniCPM5-1B prompt enrichment for ZeroGPU Spaces (Stage 2)

A Space has no Ollama daemon, so enrichment fell back to the plain non-LLM path (bland titles). Add a MiniCPM5-1B backend (standard LlamaForCausalLM, no trust_remote_code, thinking mode off, robust JSON parse) loaded on cuda at startup and run INSIDE the @spaces.GPU call alongside MusicGen - one GPU acquisition per vend. enrich_prompt now dispatches by environment: MiniCPM on ZeroGPU, Ollama locally, plain fallback on any failure (a bad enricher can't crash the app - startup load is guarded). No new dependencies (transformers already covers it). Local mps/cpu/stub paths and progress are unchanged.

Co-Authored-By: Claude Opus 4.8 <[email protected]>

app.py

@@ -3,17 +3,19 @@
 Gradio Server backend: serves the Three.js frontend and exposes the
 generation API.
 
-Pipeline: user vibe -> Ollama (small LLM) enriches it into a MusicGen
-prompt + cassette title + ambience pick -> MusicGen renders the music ->
-ambience.py loops a background bed (waves, crackle, rain…) underneath.
-MusicGen ignores texture words in prompts, hence the separate bed.
+Pipeline: user vibe -> a small LLM enriches it into a MusicGen prompt +
+cassette title + ambience pick -> MusicGen renders the music -> ambience.py
+loops a background bed (waves, crackle, rain…) underneath. MusicGen ignores
+texture words in prompts, hence the separate bed. The enrichment LLM is
+MiniCPM (on cuda) on a ZeroGPU Space, or a local Ollama daemon in dev.
 
 Env knobs:
   LOFINITY_ENGINE   musicgen (default) | stub
   LOFINITY_DURATION clip length in seconds (default 30, the single-shot max)
   LOFINITY_DEVICE   cuda | mps | cpu (default: cuda on ZeroGPU, else mps if available)
-  OLLAMA_URL        default http://localhost:11434
-  OLLAMA_MODEL      default llama3.2:3b
+  LOFINITY_ENRICHER MiniCPM model id for ZeroGPU enrichment (default MiniCPM5-1B)
+  OLLAMA_URL        default http://localhost:11434  (local enrichment)
+  OLLAMA_MODEL      default llama3.2:3b              (local enrichment)
 """
 
 import base64
@@ -87,7 +89,7 @@ app = Server(title="LoFinity")
 # frontend polls /api/progress to fill its brewing bar.
 _PROGRESS = {"done": 0, "total": 1}
 
-# --- prompt enrichment (Ollama) ----------------------------------------------
+# --- prompt enrichment --------------------------------------------------------
 
 ENRICH_SYSTEM = """\
 You are the creative brain of LoFinity, a magical vending machine that sells
@@ -116,47 +118,138 @@ user: studying at midnight
 {"music_prompt": "lofi chill, rhodes piano, muted guitar, soft bass, focused and calm, slow tempo, 75 bpm, instrumental", "title": "Midnight Study Session", "ambience": "vinyl_crackle"}"""
 
 
-def enrich_prompt(prompt: str) -> tuple[str, str, str]:
-    """Vibe -> (music_prompt, cassette title, ambience slug), with a plain
-    fallback if the local LLM is unreachable or returns junk."""
-    import ambience
+# MiniCPM enrichment LLM (ZeroGPU only — a Space has no Ollama daemon).
+# MiniCPM5-1B is a standard LlamaForCausalLM (no trust_remote_code, fast
+# tokenizer) with a switchable <think> mode we keep OFF so the reply is direct
+# JSON. Needs transformers>=5.6 (the Space's latest satisfies it); no extra deps.
+ENRICHER_MODEL = os.getenv("LOFINITY_ENRICHER", "openbmb/MiniCPM5-1B")
+_enricher = None
+_enricher_lock = threading.Lock()
+_enricher_disabled = False  # set if the model can't load; forces the fallback
+
+
+def load_enricher():
+    """Lazy-load the MiniCPM enrichment LLM on cuda (ZeroGPU). Like MusicGen it is
+    placed on cuda at module level; standard Llama arch, so no remote code."""
+    global _enricher
+    with _enricher_lock:
+        if _enricher is None:
+            import torch  # noqa: F401 — needed so the .to('cuda') below resolves
+            from transformers import AutoModelForCausalLM, AutoTokenizer
+
+            print(f"[lofinity] loading enricher {ENRICHER_MODEL} on cuda…")
+            tok = AutoTokenizer.from_pretrained(ENRICHER_MODEL)
+            model = AutoModelForCausalLM.from_pretrained(ENRICHER_MODEL, torch_dtype="auto")
+            model.to("cuda")
+            model.eval()
+            _enricher = (tok, model)
+            print("[lofinity] enricher ready")
+    return _enricher
+
 
+def _parse_enrich_json(text: str) -> dict:
+    """Pull the first {...} object out of an LLM reply (it may wrap the JSON in
+    prose or ```json fences, or leak a <think> block); {} if nothing parses."""
+    import re
+
+    if "</think>" in text:  # belt-and-suspenders if thinking ever leaks through
+        text = text.rsplit("</think>", 1)[1]
+    m = re.search(r"\{.*\}", text, re.DOTALL)
+    if not m:
+        return {}
     try:
-        r = httpx.post(
-            f"{OLLAMA_URL}/api/chat",
-            json={
-                "model": OLLAMA_MODEL,
-                "messages": [
-                    {"role": "system", "content": ENRICH_SYSTEM},
-                    {"role": "user", "content": prompt},
-                ],
-                "format": "json",
-                "stream": False,
-                "options": {"temperature": 0.8, "num_predict": 220},
-            },
-            timeout=45,
+        return json.loads(m.group(0))
+    except Exception:  # noqa: BLE001
+        return {}
+
+
+def _finalize_enrichment(data: dict):
+    """Shared post-processing for any backend: validate, force the genre to lead,
+    snap the ambience to a renderable bed. Returns a tuple, or None if unusable."""
+    import ambience
+
+    music_prompt = str(data.get("music_prompt") or "").strip()
+    title = str(data.get("title") or "").strip()[:48]
+    if not (music_prompt and title):
+        return None
+    # belt and suspenders: the genre must lead even if the LLM drifts
+    if "lofi" not in music_prompt.lower():
+        music_prompt = f"lofi chill, {music_prompt}"
+    # whatever the LLM picked, snap it to a bed we can actually render
+    return music_prompt, title, ambience.normalize_slug(data.get("ambience"))
+
+
+def _enrich_minicpm(prompt: str):
+    """Enrich via MiniCPM on cuda. MUST run inside @spaces.GPU. Returns a tuple or
+    None (caller falls back). Thinking mode off so the reply is direct JSON."""
+    if _enricher_disabled:
+        return None
+    import torch
+
+    tok, model = load_enricher()
+    messages = [
+        {"role": "system", "content": ENRICH_SYSTEM},
+        {"role": "user", "content": prompt},
+    ]
+    inputs = tok.apply_chat_template(
+        messages, tokenize=True, add_generation_prompt=True,
+        enable_thinking=False, return_dict=True, return_tensors="pt",
+    ).to(model.device)
+    with torch.no_grad():
+        out = model.generate(
+            **inputs, max_new_tokens=220, do_sample=True, temperature=0.7, top_p=0.95
         )
-        r.raise_for_status()
-        data = json.loads(r.json()["message"]["content"])
-        music_prompt = str(data.get("music_prompt") or "").strip()
-        title = str(data.get("title") or "").strip()[:48]
-        if music_prompt and title:
-            # belt and suspenders: the genre must lead even if the LLM drifts
-            if "lofi" not in music_prompt.lower():
-                music_prompt = f"lofi chill, {music_prompt}"
-            # whatever the LLM picked, snap it to a bed we can actually render
-            return music_prompt, title, ambience.normalize_slug(data.get("ambience"))
-    except Exception as e:  # noqa: BLE001 — any failure means "use fallback"
-        print(f"[lofinity] ollama enrichment failed ({e!r}), using fallback")
-    fallback_title = f"{prompt[:28].title()} Tape" if prompt.strip() else "Untitled Tape"
+    reply = tok.decode(out[0][inputs["input_ids"].shape[-1]:], skip_special_tokens=True)
+    return _finalize_enrichment(_parse_enrich_json(reply))
+
+
+def _enrich_ollama(prompt: str):
+    """Enrich via a local Ollama daemon. Returns a tuple or None on failure."""
+    r = httpx.post(
+        f"{OLLAMA_URL}/api/chat",
+        json={
+            "model": OLLAMA_MODEL,
+            "messages": [
+                {"role": "system", "content": ENRICH_SYSTEM},
+                {"role": "user", "content": prompt},
+            ],
+            "format": "json",
+            "stream": False,
+            "options": {"temperature": 0.8, "num_predict": 220},
+        },
+        timeout=45,
+    )
+    r.raise_for_status()
+    return _finalize_enrichment(json.loads(r.json()["message"]["content"]))
+
+
+def _enrich_fallback(prompt: str) -> tuple[str, str, str]:
+    """Plain, LLM-free enrichment — used whenever the chosen backend fails."""
+    import ambience
+
+    title = f"{prompt[:28].title()} Tape" if prompt.strip() else "Untitled Tape"
     return (
-        f"lofi chill, {prompt}, mellow and warm, soft drums, "
-        "slow tempo, instrumental",
-        fallback_title,
+        f"lofi chill, {prompt}, mellow and warm, soft drums, slow tempo, instrumental",
+        title,
         ambience.DEFAULT,
     )
 
 
+def enrich_prompt(prompt: str) -> tuple[str, str, str]:
+    """Vibe -> (music_prompt, cassette title, ambience slug). Backend is chosen by
+    environment: MiniCPM on ZeroGPU, Ollama locally; a plain fallback covers any
+    failure. On ZeroGPU this MUST be called inside @spaces.GPU (MiniCPM is cuda)."""
+    backend = _enrich_minicpm if IS_ZEROGPU else _enrich_ollama
+    try:
+        result = backend(prompt)
+        if result:
+            return result
+        print("[lofinity] enrichment returned junk, using fallback")
+    except Exception as e:  # noqa: BLE001 — any failure means "use fallback"
+        print(f"[lofinity] enrichment failed ({e!r}), using fallback")
+    return _enrich_fallback(prompt)
+
+
 # --- audio engines ------------------------------------------------------------
 
 _musicgen = None
@@ -199,6 +292,11 @@ def load_musicgen():
 # placements done at startup are far more efficient than per-call transfers.
 if IS_ZEROGPU and ENGINE != "stub":
     load_musicgen()
+    try:
+        load_enricher()
+    except Exception as e:  # noqa: BLE001 — a bad enricher must not kill the app
+        _enricher_disabled = True
+        print(f"[lofinity] enricher load failed ({e!r}); vends use the plain fallback")
 
 
 def encode_wav(samples, rate: int) -> str:
@@ -308,22 +406,25 @@ def musicgen_engine(music_prompt: str, seconds: int = CHUNK_S, progress_cb=None)
     return samples, rate
 
 
-def _gpu_budget(music_prompt: str, seconds: int = CHUNK_S) -> int:
-    """GPU seconds to request from ZeroGPU for a brew of this length: per-chunk
-    render time plus headroom. Tighter budgets earn better queue priority, and
-    the signature must mirror gpu_musicgen so ZeroGPU can pass it the same args."""
+def _gpu_budget(prompt: str, seconds: int = CHUNK_S) -> int:
+    """GPU seconds to request from ZeroGPU for a brew of this length: MiniCPM
+    enrichment + per-chunk MusicGen render plus headroom. Tighter budgets earn
+    better queue priority; the signature must mirror gpu_brew so ZeroGPU can pass
+    it the same args."""
     chunks = max(1, round(int(seconds) / CHUNK_S))
-    return 20 + 25 * chunks  # 30s->45, 60s->70, 90s->95
+    return 25 + 25 * chunks  # enrichment + 30s->50, 60s->75, 90s->100
 
 
 @spaces.GPU(duration=_gpu_budget)
-def gpu_musicgen(music_prompt: str, seconds: int = CHUNK_S) -> tuple:
-    """ZeroGPU entry point — runs MusicGen on the real GPU and returns the audio
-    to the web process. No progress_cb on purpose: this body executes in a
-    separate GPU worker, so _PROGRESS updates can't reach /api/progress yet (the
-    brewing garden is wired up in a later stage); on the Space the brew jumps to
-    done. Locally @spaces.GPU is a no-op, but this path is only taken on Spaces."""
-    return musicgen_engine(music_prompt, seconds)
+def gpu_brew(prompt: str, seconds: int = CHUNK_S) -> tuple:
+    """ZeroGPU entry point — enrichment (MiniCPM) AND MusicGen on the real GPU in
+    a single acquisition. Takes the raw vibe and returns
+    (music_prompt, title, bed, samples, rate). No progress_cb: this body runs in
+    a separate GPU worker, so _PROGRESS can't reach /api/progress yet (Stage 3) —
+    the brewing garden jumps to done on the Space. This path is Space-only."""
+    music_prompt, title, bed = enrich_prompt(prompt)
+    samples, rate = musicgen_engine(music_prompt, seconds)
+    return music_prompt, title, bed, samples, rate
 
 
 def stub_engine(_music_prompt: str, seconds: int = CHUNK_S, progress_cb=None) -> tuple:
@@ -356,26 +457,29 @@ def generate_song(prompt: str, seconds: int = DEFAULT_SECONDS) -> dict:
 
     # snap whatever the slider sends to a length we can actually build
     seconds = min(ALLOWED_SECONDS, key=lambda s: abs(s - int(seconds)))
-    # reset progress up front, BEFORE the (sometimes slow) Ollama enrich step, so
-    # a poll arriving early sees this brew at 0% rather than the last one at 100%
+    # reset progress up front, BEFORE the (sometimes slow) enrich step, so a poll
+    # arriving early sees this brew at 0% rather than the last one at 100%
     chunks = max(1, round(seconds / CHUNK_S))
     _PROGRESS.update(done=0, total=chunks)
-    music_prompt, title, bed = enrich_prompt(prompt)
-    print(f"[lofinity] brewing {title!r} ({seconds}s) :: {music_prompt} [+ {bed}]")
-    if ENGINE == "stub":
-        samples, rate = stub_engine(
-            music_prompt, seconds,
-            progress_cb=lambda d, t: _PROGRESS.update(done=d, total=t),
-        )
-    elif IS_ZEROGPU:
-        # the GPU body runs in a separate worker process, so progress can't
-        # stream back here yet (Stage 3); the brewing bar jumps 0->100% on Space
-        samples, rate = gpu_musicgen(music_prompt, seconds)
+
+    if IS_ZEROGPU and ENGINE != "stub":
+        # On ZeroGPU both enrichment (MiniCPM) and MusicGen need the real GPU, so
+        # they share ONE @spaces.GPU acquisition. That worker runs in a separate
+        # process, so progress can't stream back yet (Stage 3): the bar jumps.
+        print(f"[lofinity] brewing on GPU :: {prompt!r} ({seconds}s)")
+        music_prompt, title, bed, samples, rate = gpu_brew(prompt, seconds)
+        print(f"[lofinity] brewed {title!r} :: {music_prompt} [+ {bed}]")
     else:
-        samples, rate = musicgen_engine(
+        # Local / stub: enrich in-process (Ollama or fallback), then render with
+        # live per-chunk progress for the brewing garden.
+        music_prompt, title, bed = enrich_prompt(prompt)
+        print(f"[lofinity] brewing {title!r} ({seconds}s) :: {music_prompt} [+ {bed}]")
+        engine = stub_engine if ENGINE == "stub" else musicgen_engine
+        samples, rate = engine(
             music_prompt, seconds,
             progress_cb=lambda d, t: _PROGRESS.update(done=d, total=t),
         )
+
     _PROGRESS.update(done=chunks, total=chunks)
     try:
         samples = ambience.mix(samples, rate, bed)

requirements.txt

@@ -3,8 +3,10 @@
 # a no-op when it's absent, so it's only actually required on the Space.
 spaces
 torch  # ZeroGPU requires torch >=2.8; the Space runtime supplies the CUDA build
-transformers
-# Ollama must be running locally with the model below pulled:
+transformers  # MusicGen + (on ZeroGPU) the MiniCPM5-1B enricher; needs >=5.6 for MiniCPM5
+# Enrichment LLM: on a ZeroGPU Space it's openbmb/MiniCPM5-1B, pulled from the Hub
+# at startup (no extra deps — standard Llama arch + fast tokenizer). For LOCAL dev
+# instead, run Ollama with the model below pulled:
 #   ollama pull llama3.2:3b
 # The app runtime needs nothing more (ambience.py reads beds via stdlib `wave`).
 # To (re)populate the sampled ambience beds in assets/ambience/, pick one: