audio: support audiocraft properly, in a venv of its own

The transformers fallback keeps MusicGen working on a pod, but it is not the
same thing: audiocraft has melody conditioning and AudioGen, and transformers
has neither. Supporting it is not a matter of adding a requirement — it pins
torch==2.1.0, torchvision==0.16.0 and torchtext==0.16.0, and torch 2.1.0 has no
build for the CUDA these images run, so installing it beside the server would
leave the pod with no working GPU torch at all.

So it gets its own venv and is driven as a subprocess, the same answer
parler-tts and the OCR engines already use here. AudiocraftModel deliberately
wears audiocraft's own interface (set_generation_params / generate /
generate_with_chroma / sample_rate) so the generation path cannot tell whether
the library is in this process, in a venv next door, or absent.

The venv is baked into the audio image rather than built on first use: a pod is
rented by the second, and pip-installing a library during a cold start costs
more than carrying the layer.

Order of preference: in-process audiocraft, then the isolated venv, then
transformers — which needs no venv but cannot do melody.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0182DhR1eFNPQyGmDrpHbedY

codai/api/audio_gen.py

@@ -164,8 +164,20 @@ def _load_musicgen(model_name: str, device: str):
     try:
         from audiocraft.models import MusicGen, AudioGen
     except ImportError:
+        # Next best: audiocraft in its own venv. It pins torch==2.1.0, which has
+        # no build for the CUDA this server runs, so it cannot share this venv —
+        # but it can live next door and be driven as a subprocess, which is what
+        # keeps melody conditioning available at all.
+        try:
+            from codai.api import audiocraft_worker
+            if audiocraft_worker.available():
+                print(f"[audio] using audiocraft from {audiocraft_worker._VENV}",
+                      flush=True)
+                return audiocraft_worker.AudiocraftModel(model_name)
+        except Exception as exc:
+            print(f"[audio] isolated audiocraft venv unusable ({exc})", flush=True)
         print(f"[audio] audiocraft not installed — serving {model_name} through "
-              "transformers instead", flush=True)
+              "transformers instead (no melody conditioning)", flush=True)
         return _TransformersMusicGen(model_name, device)
     name_lower = model_name.lower()
     if 'audiogen' in name_lower:

codai/api/audiocraft_worker.py

+# CoderAI - OpenAI-compatible API server
+# Copyright (C) 2026 Stefy Lanza <stefy@nexlab.net>
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+
+"""Drive audiocraft in its own venv, and look like a MusicGen model doing it.
+
+audiocraft pins torch==2.1.0. The server runs torch 2.11 built for a specific
+CUDA, and torch 2.1.0 has no build for that CUDA — so installing audiocraft
+beside the server does not merely risk a conflict, it leaves the GPU unusable.
+The same problem parler-tts and the OCR engines have, solved the same way: its
+own venv, driven as a subprocess.
+
+What this exposes is deliberately the shape audiocraft's own MusicGen has
+(``set_generation_params``, ``generate``, ``sample_rate``), because the audio
+generation path already speaks that. The caller cannot tell whether audiocraft
+is in this process, in a venv next door, or absent — in which case the path
+falls back to transformers' own MusicGen, which needs no venv but has no melody
+conditioning.
+"""
+
+import base64
+import io
+import json
+import os
+import subprocess
+import threading
+from pathlib import Path
+
+_REPO_ROOT = Path(__file__).resolve().parents[2]
+_SERVICE = _REPO_ROOT / "tools" / "audiocraft_service.py"
+
+#: Where the venv lives. The pod images bake it at /opt/coderai/venvs/audiocraft
+#: — a pod is disposable and building it on first use would pay minutes of GPU
+#: rental to install a library. A local install defaults under ~/.coderai.
+_VENV = Path(os.environ.get("CODERAI_AUDIOCRAFT_VENV")
+             or ("/opt/coderai/venvs/audiocraft"
+                 if Path("/opt/coderai/venvs/audiocraft").exists()
+                 else os.path.expanduser("~/.coderai/audiocraft_venv")))
+
+_lock = threading.RLock()
+
+
+def venv_python() -> Path:
+    return _VENV / ("Scripts" if os.name == "nt" else "bin") / (
+        "python.exe" if os.name == "nt" else "python")
+
+
+def available() -> bool:
+    """True when there is a venv with audiocraft in it to talk to."""
+    py = venv_python()
+    if not py.is_file() or not _SERVICE.is_file():
+        return False
+    try:
+        return subprocess.run([str(py), "-c", "import audiocraft"],
+                              capture_output=True, timeout=120).returncode == 0
+    except Exception:
+        return False
+
+
+class AudiocraftModel:
+    """An audiocraft model living in another interpreter."""
+
+    def __init__(self, model_name: str):
+        self._proc = subprocess.Popen(
+            [str(venv_python()), str(_SERVICE)],
+            stdin=subprocess.PIPE, stdout=subprocess.PIPE,
+            text=True, bufsize=1,
+            env={**os.environ, "PYTHONUNBUFFERED": "1"})
+        self._params: dict = {"duration": 30.0}
+        reply = self._call({"op": "load", "model": model_name})
+        if not reply.get("ok"):
+            self.close()
+            raise RuntimeError(f"audiocraft could not load {model_name!r}: "
+                               f"{reply.get('error')}")
+        self._sample_rate = int(reply.get("sample_rate") or 32000)
+
+    # -- the audiocraft-shaped surface ------------------------------------ #
+    @property
+    def sample_rate(self) -> int:
+        return self._sample_rate
+
+    def set_generation_params(self, **kwargs):
+        self._params.update({k: v for k, v in kwargs.items() if v is not None})
+
+    def generate(self, prompts):
+        return self._audio({"op": "generate", "prompt": list(prompts)[0],
+                            **self._params})
+
+    def generate_with_chroma(self, prompts, melody_wav, sr):
+        """Melody conditioning — the capability transformers does not have."""
+        import torchaudio
+        buf = io.BytesIO()
+        wav = melody_wav[0] if melody_wav.dim() == 3 else melody_wav
+        torchaudio.save(buf, wav.cpu(), int(sr), format="wav")
+        return self._audio({"op": "generate", "prompt": list(prompts)[0],
+                            "melody_wav": base64.b64encode(buf.getvalue()).decode(),
+                            **self._params})
+
+    # -- transport --------------------------------------------------------- #
+    def _audio(self, request: dict):
+        reply = self._call(request)
+        if not reply.get("ok"):
+            raise RuntimeError(f"audiocraft generation failed: {reply.get('error')}")
+        import torch, torchaudio
+        wav, _sr = torchaudio.load(io.BytesIO(base64.b64decode(reply["b64_wav"])))
+        # The caller indexes [0, 0] like audiocraft's own (batch, channels, samples).
+        return wav.unsqueeze(0) if wav.dim() == 2 else wav
+
+    def _call(self, request: dict) -> dict:
+        with _lock:
+            if self._proc.poll() is not None:
+                return {"ok": False, "error": "the audiocraft worker has exited"}
+            self._proc.stdin.write(json.dumps(request) + "\n")
+            self._proc.stdin.flush()
+            line = self._proc.stdout.readline()
+        if not line:
+            return {"ok": False, "error": "the audiocraft worker returned nothing"}
+        try:
+            return json.loads(line)
+        except Exception as exc:
+            return {"ok": False, "error": f"unparseable worker reply: {exc}"}
+
+    def close(self) -> None:
+        try:
+            if self._proc.poll() is None:
+                self._proc.terminate()
+                self._proc.wait(timeout=10)
+        except Exception:
+            pass
+
+    def __del__(self):
+        self.close()

packaging/runpod/Dockerfile.capability

@@ -50,6 +50,21 @@ RUN --mount=type=cache,target=/root/.cache/pip set -eux; \
     fi; \
     rm -rf /tmp/profiles /var/lib/apt/lists/* /root/.cargo /root/.rustup
 
+# A profile may need a second venv for a library that cannot share the first —
+# audiocraft pins torch==2.1.0, which has no build for this image's CUDA. Built
+# here rather than on first use: a pod is rented by the second, and installing a
+# library during a cold start costs more than carrying the layer.
+ARG PROFILE
+COPY packaging/runpod/profiles/ /tmp/venv-profiles/
+RUN set -eux; \
+    extra="/tmp/venv-profiles/${PROFILE}.venv-audiocraft.txt"; \
+    if [ -f "$extra" ]; then \
+        python -m venv /opt/coderai/venvs/audiocraft; \
+        /opt/coderai/venvs/audiocraft/bin/pip install --no-cache-dir -r "$extra"; \
+        /opt/coderai/venvs/audiocraft/bin/python -c "import audiocraft; print('audiocraft', audiocraft.__version__)"; \
+    fi; \
+    rm -rf /tmp/venv-profiles
+
 COPY . /opt/coderai/app
 
 # Fail the BUILD, not the pod, if the profile is missing something the app needs

packaging/runpod/profiles/audio.venv-audiocraft.txt

+# A SECOND venv inside the audio image, for audiocraft alone.
+#
+# audiocraft pins torch==2.1.0, torchvision==0.16.0 and torchtext==0.16.0. The
+# server runs torch 2.11 built for a specific CUDA, and torch 2.1.0 has no build
+# for that CUDA at all — so installing audiocraft beside the server does not
+# merely risk a conflict, it leaves the pod with no working GPU torch. Its own
+# venv is the only way both can be right.
+#
+# Baked at build time, not on first use: a pod is disposable and rented by the
+# second, so paying minutes of GPU rental to pip-install a library on every cold
+# start would cost more than the image layer does.
+--extra-index-url https://download.pytorch.org/whl/cu121
+torch==2.1.0
+torchaudio==2.1.0
+audiocraft==1.3.0

tests/test_remote_execution.py

@@ -1635,3 +1635,45 @@ def test_the_gateway_finds_the_ocr_engine_field():
 
     # JSON bodies honour the field too.
     assert _model_from_body(b'{"engine":"doctr"}', "application/json", "engine") == "doctr"
+
+
+def test_audiocraft_runs_in_its_own_venv_and_looks_like_musicgen(monkeypatch):
+    """audiocraft pins torch==2.1.0, which has no build for the CUDA this
+    server runs — installing it beside the server leaves the GPU unusable. It
+    gets its own venv, like parler-tts and the OCR engines, and the generation
+    path must not be able to tell the difference."""
+    from codai.api import audiocraft_worker as aw
+
+    # The surface the audio path actually calls.
+    for name in ("sample_rate", "set_generation_params", "generate",
+                 "generate_with_chroma"):
+        assert hasattr(aw.AudiocraftModel, name), name
+
+    # available() is honest when there is no venv to talk to.
+    monkeypatch.setattr(aw, "_VENV", Path("/nonexistent/venv"))
+    assert aw.available() is False
+
+
+def test_musicgen_falls_back_in_order(monkeypatch):
+    """In-process audiocraft, then the isolated venv, then transformers. The
+    last one has no melody conditioning, so it must be last, not first."""
+    import codai.api.audio_gen as ag
+    from codai.api import audiocraft_worker as aw
+
+    calls = []
+    monkeypatch.setattr(aw, "available", lambda: True)
+    monkeypatch.setattr(aw, "AudiocraftModel",
+                        lambda name: calls.append(("venv", name)) or "VENV")
+
+    # No in-process audiocraft: the venv wins over transformers.
+    import builtins
+    real_import = builtins.__import__
+
+    def _no_audiocraft(name, *a, **k):
+        if name == "audiocraft.models" or name == "audiocraft":
+            raise ImportError("no audiocraft here")
+        return real_import(name, *a, **k)
+
+    monkeypatch.setattr(builtins, "__import__", _no_audiocraft)
+    got = ag._load_musicgen("facebook/musicgen-small", "cpu")
+    assert got == "VENV" and calls == [("venv", "facebook/musicgen-small")]

tools/audiocraft_service.py

+#!/usr/bin/env python3
+# CoderAI - OpenAI-compatible API server
+# Copyright (C) 2026 Stefy Lanza <stefy@nexlab.net>
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+
+"""audiocraft in its own venv, spoken to over stdin/stdout.
+
+audiocraft hard-pins torch==2.1.0, torchvision==0.16.0 and torchtext==0.16.0.
+The server runs torch 2.11 built for a specific CUDA, and there is no version of
+this argument that both sides win: installing audiocraft into the server's venv
+downgrades torch five major versions and leaves the GPU unusable, because torch
+2.1.0 has no build for that CUDA at all.
+
+So it lives in a venv of its own and is driven as a subprocess, exactly like the
+parler-tts and OCR stacks that have the same problem. This script is the far
+side: one JSON object per line in, one per line out.
+
+  {"op": "load", "model": "facebook/musicgen-small"}   -> {"ok": true, "sample_rate": 32000}
+  {"op": "generate", "prompt": "...", "duration": 2.0} -> {"ok": true, "b64_wav": "..."}
+  {"op": "ping"}                                        -> {"ok": true}
+
+Audio comes back as base64 WAV rather than a path: the caller may be a pod with
+no shared filesystem, and a few seconds of audio is small enough that a file
+would only add a thing to clean up.
+"""
+
+import base64
+import io
+import json
+import sys
+
+
+_state = {"model": None, "name": ""}
+
+
+def _load(name: str) -> dict:
+    from audiocraft.models import MusicGen, AudioGen
+    if "audiogen" in name.lower():
+        model = AudioGen.get_pretrained(name)
+    else:
+        model = MusicGen.get_pretrained(name)
+    _state["model"], _state["name"] = model, name
+    return {"ok": True, "sample_rate": int(model.sample_rate)}
+
+
+def _generate(req: dict) -> dict:
+    model = _state["model"]
+    if model is None:
+        return {"ok": False, "error": "no model loaded"}
+    params = {"duration": float(req.get("duration") or 10.0)}
+    for key in ("top_k", "top_p", "temperature", "cfg_coef"):
+        if req.get(key) is not None:
+            params[key] = req[key]
+    model.set_generation_params(**params)
+
+    melody_b64 = req.get("melody_wav")
+    if melody_b64:
+        # Melody conditioning is the reason to keep audiocraft at all —
+        # transformers' MusicGen has no equivalent.
+        import torch, torchaudio
+        raw = base64.b64decode(melody_b64)
+        wav, sr = torchaudio.load(io.BytesIO(raw))
+        out = model.generate_with_chroma([req.get("prompt") or ""],
+                                         wav.unsqueeze(0), sr)
+    else:
+        out = model.generate([req.get("prompt") or ""])
+
+    import torchaudio
+    buf = io.BytesIO()
+    torchaudio.save(buf, out[0].cpu(), int(model.sample_rate), format="wav")
+    return {"ok": True, "b64_wav": base64.b64encode(buf.getvalue()).decode(),
+            "sample_rate": int(model.sample_rate)}
+
+
+def main() -> int:
+    for line in sys.stdin:
+        line = line.strip()
+        if not line:
+            continue
+        try:
+            req = json.loads(line)
+            op = req.get("op")
+            if op == "ping":
+                out = {"ok": True}
+            elif op == "load":
+                out = _load(str(req.get("model") or ""))
+            elif op == "generate":
+                out = _generate(req)
+            else:
+                out = {"ok": False, "error": f"unknown op {op!r}"}
+        except Exception as exc:
+            # Never die on one bad request: the caller keeps the process for the
+            # next one, and a crash here costs a model load.
+            out = {"ok": False, "error": f"{type(exc).__name__}: {exc}"}
+        sys.stdout.write(json.dumps(out) + "\n")
+        sys.stdout.flush()
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())