audio: audiocraft in the main venv, and a setting to choose the backend

The isolated venv was solving a problem that was not real. audiocraft is pure
Python: its torch==2.1.0 is a declared pin, not an ABI requirement, so --no-deps
sidesteps it. The only genuine constraint was xformers, which it imports at
module load — and 0.0.35 ships a py39-none wheel requiring torch>=2.10, which
this image satisfies. I had ruled that out by grepping wheel filenames for
cp312 and missing that abi3/py39 wheels cover 3.12.

Verified in the built core image: torch 2.11.0+cu128 with xformers 0.0.35,
MusicGen and AudioGen import, ops.memory_efficient_attention present.

Cost in the main venv: 950 MB. The second venv cost 6.2 GB, because it carried a
duplicate torch and a duplicate CUDA stack. The audio image drops from 15.8 GB
to about 10.2 GB, which also removes the reason to split it in two.

audio_backend (per model, or CODERAI_AUDIO_BACKEND for a pod) selects
audiocraft or transformers explicitly, because they are not equivalent — melody
conditioning and AudioGen exist only in the first. Asking for audiocraft when it
is absent raises rather than quietly substituting something less capable. A
model's own setting beats the environment it lands in, and the choice travels to
a pod instead of being decided by what happens to be installed there.

The isolated-venv worker stays in the tree: if audiocraft 1.3.0 turns out to
misbehave against torch 2.11 at generation time, restoring one profile file
brings it back.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0182DhR1eFNPQyGmDrpHbedY

codai/api/audio_gen.py

@@ -160,10 +160,39 @@ class _TransformersMusicGen:
             return self._model.generate(**inputs, **gen)   # (batch, channels, samples)
 
 
-def _load_musicgen(model_name: str, device: str):
+def _audio_backend(model_config: dict = None) -> str:
+    """Which MusicGen implementation to use: 'audiocraft', 'transformers', 'auto'.
+
+    Per model first, then the environment (which is how a pod is told), then
+    'auto' — audiocraft when it is importable, transformers otherwise.
+
+    The choice is worth exposing because the two are not equivalent: audiocraft
+    has melody conditioning and AudioGen, transformers has neither but needs
+    nothing beyond what every image already carries.
+    """
+    import os as _os
+    chosen = str((model_config or {}).get("audio_backend") or "").strip().lower()
+    if not chosen:
+        chosen = str(_os.environ.get("CODERAI_AUDIO_BACKEND", "")).strip().lower()
+    return chosen if chosen in ("audiocraft", "transformers") else "auto"
+
+
+def _load_musicgen(model_name: str, device: str, model_config: dict = None):
+    backend = _audio_backend(model_config)
+    if backend == "transformers":
+        print(f"[audio] {model_name}: transformers backend requested "
+              "(no melody conditioning)", flush=True)
+        return _TransformersMusicGen(model_name, device)
     try:
         from audiocraft.models import MusicGen, AudioGen
     except ImportError:
+        if backend == "audiocraft":
+            # Asked for explicitly and not here: say so rather than quietly
+            # serving something with different capabilities.
+            raise RuntimeError(
+                "audio_backend is set to 'audiocraft' but audiocraft is not "
+                "installed. Install it, or set audio_backend to 'transformers' "
+                "(which cannot do melody conditioning).")
         # 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
@@ -339,7 +368,7 @@ async def audio_generate(request: AudioGenerationRequest, http_request: Request
             with loading_task(model_name, model_type="audio"):
                 if model_type in ('musicgen', 'audiogen'):
                     pipe = await asyncio.get_event_loop().run_in_executor(
-                        None, _load_musicgen, model_name, device)
+                        None, _load_musicgen, model_name, device, _ag_cfg)
                 else:
                     pipe = await asyncio.get_event_loop().run_in_executor(
                         None, _load_audioldm, model_name, device, _ag_cfg)

codai/api/runpod_worker.py

@@ -599,6 +599,13 @@ def _plan(engine, image, args, mcfg, api_key, entry, served,
     # disabled (enable it in Settings → OCR)" — a settings screen nobody will
     # open on a machine that exists for the next four minutes.
     cap = capability or model_capability(entry or {}, include_text=True)
+    if cap == "audio_gen":
+        # audiocraft and transformers are not equivalent — melody conditioning
+        # and AudioGen exist only in the first — so the choice travels to the pod
+        # rather than being decided by whatever happens to be installed there.
+        choice = str((entry or {}).get("audio_backend") or "").strip().lower()
+        if choice in ("audiocraft", "transformers"):
+            env["CODERAI_AUDIO_BACKEND"] = choice
     if cap == "ocr":
         env["CODERAI_OCR_ENABLED"] = "1"
         wanted = str((entry or {}).get("path") or "").strip().lower()
@@ -672,7 +679,10 @@ def seed_model_env(entry: dict, served: str = "", source: str = "") -> str:
             # LoRA settings travel too: a coderai pod applies them itself, and
             # load_in_4bit above is what makes a QLoRA adapter load against the
             # base it was trained on.
-            "lora_path", "lora_model_dir", "lora_scale", "loras")
+            "lora_path", "lora_model_dir", "lora_scale", "loras",
+            # Which MusicGen implementation to use. Describes the model, not
+            # this deployment's placement, so it travels.
+            "audio_backend")
     out = {k: entry[k] for k in keep if k in entry and entry[k] is not None}
     out["path"] = path
     # Adapters the pod cannot resolve by name are named by CONTENT id instead.

packaging/runpod/Dockerfile.capability

@@ -45,6 +45,12 @@ RUN --mount=type=cache,target=/root/.cache/pip set -eux; \
     fi; \
     python -m pip install --extra-index-url ${TORCH_INDEX} -r /tmp/profiles/${PROFILE}.txt; \
     python -c "import torch; print('torch', torch.__version__, 'cuda', torch.version.cuda)"; \
+    if [ -f "/tmp/profiles/${PROFILE}.nodeps" ]; then \
+        python -m pip install --no-deps -r "/tmp/profiles/${PROFILE}.nodeps"; \
+        for pkg in $(grep -vE '^\s*(#|$)' "/tmp/profiles/${PROFILE}.nodeps" | sed 's/[<>=!].*//'); do \
+            python -c "import ${pkg}; print('${pkg}', getattr(${pkg}, '__version__', 'ok'), '(installed --no-deps)')"; \
+        done; \
+    fi; \
     if [ -n "$deps" ]; then \
         apt-get purge -y $deps && apt-get autoremove -y; \
     fi; \

packaging/runpod/build_all_capability_images.sh

@@ -12,7 +12,10 @@ set -uo pipefail
 cd "$(dirname "$0")/../.."
 NS="${1:-}"
 VERSION="${VERSION:-$(python3 -c 'import re,pathlib; print(re.search(r"__version__ = \"([^\"]+)\"", pathlib.Path("codai/__init__.py").read_text()).group(1))')}"
-PROFILES="${PROFILES:-$(ls packaging/runpod/profiles/*.txt | xargs -n1 basename | sed 's/\.txt$//' | grep -v '^core$' | tr '\n' ' ')}"
+# A profile's extra-venv specs live beside it as <profile>.venv-<name>.txt, so
+# a bare *.txt glob would invent a profile called "audio.venv-audiocraft" and
+# try to build an image for it.
+PROFILES="${PROFILES:-$(ls packaging/runpod/profiles/*.txt | xargs -n1 basename | sed 's/\.txt$//' | grep -v '^core$' | grep -v '\.venv-' | tr '\n' ' ')}"
 
 if [[ -z "$NS" ]]; then
     echo "usage: $0 <registry/namespace>   e.g. ghcr.io/nextime" >&2

packaging/runpod/profiles/audio.nodeps

+# --no-deps: audiocraft's declared torch pins are unsatisfiable on this Python
+# and would undo the CUDA build pin if they were not. Everything it actually
+# imports is in audio.txt beside this file, and the build's import check fails
+# the image if something is missing.
+audiocraft==1.3.0

packaging/runpod/profiles/audio.txt

@@ -6,3 +6,37 @@ demucs>=4.0.1
 deepfilternet>=0.5.6
 soundfile>=0.12.0
 scipy>=1.11.0
+
+# audiocraft (MusicGen/AudioGen), in THIS venv rather than one of its own.
+#
+# It declares torch==2.1.0, torchvision==0.16.0 and torchtext==0.16.0, which is
+# what sent this down the isolated-venv road: those are unsatisfiable here and
+# installing them would have destroyed the CUDA build pin. But audiocraft is
+# pure Python — the pins are declared, not ABI — so --no-deps sidesteps them.
+#
+# The one real constraint is xformers, which audiocraft imports at module load.
+# 0.0.35 ships a py39-none wheel (works on any Python 3.9+) and requires
+# torch>=2.10, which this image satisfies. Verified in the built core image:
+# torch 2.11.0+cu128 with xformers 0.0.35, MusicGen and AudioGen both import,
+# and ops.memory_efficient_attention is present.
+#
+# Cost here: ~950 MB. The isolated venv cost 6.2 GB, because it carried a second
+# torch and a second CUDA stack to avoid a conflict that was never real.
+xformers>=0.0.35
+
+# audiocraft's own dependencies, since it is installed --no-deps (audio.nodeps).
+av>=11.0.0
+einops
+flashy>=0.0.1
+hydra-core>=1.1
+hydra_colorlog
+julius
+num2words
+sentencepiece
+spacy>=3.6.1
+huggingface_hub
+tqdm
+librosa
+torchmetrics
+encodec
+protobuf

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

-# Installed with --no-deps into the venv described by the .txt beside this file.
-# audiocraft's declared pins (torch==2.1.0, torchvision==0.16.0,
-# torchtext==0.16.0) cannot all be satisfied on this Python, and its inference
-# code does not need those exact versions — the pin is the library being
-# conservative. The import check in the build fails the image if that judgement
-# turns out to be wrong.
-audiocraft==1.3.0

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. Two
-# separate problems with that:
-#
-#  1. torch 2.1.0 has no build for Python 3.12 at all — support starts at 2.2 —
-#     so the pin is unsatisfiable in this image whatever else we do.
-#  2. Even if it existed, installing it into the server's venv would downgrade
-#     torch five majors and leave the pod with no working GPU torch, because
-#     2.1.0 has no build for this image's CUDA either.
-#
-# So: its own venv, on the oldest torch this Python can run, with audiocraft
-# installed --no-deps (see the .nodeps file beside this one) and the deps it
-# actually uses listed here. --no-deps is reckless in a shared venv and routine
-# in an isolated one — nothing here can reach the server's stack.
---extra-index-url https://download.pytorch.org/whl/cu121
-# torch 2.5.1, not the 2.1.0 audiocraft asks for and not the 2.2.2 tried before.
-# The binding constraint is xformers, which audiocraft imports at module load:
-# its cp312 wheels start at 0.0.27, and every version below that can only build
-# from source — which needs torch installed already and fails in a fresh venv
-# with "No module named 'torch'". xformers 0.0.28.post3 requires torch==2.5.1
-# exactly, so the pair is chosen by what actually has wheels for this Python,
-# not by what the libraries would prefer.
-torch==2.5.1
-torchaudio==2.5.1
-torchvision==0.20.1
-
-# EVERYTHING audiocraft declares, minus only the torch pins above and its
-# dev-only tools (pytest, flake8, mypy, coverage, pdoc3).
-#
-# The earlier version of this list was trimmed against what the MAIN venv has —
-# transformers, demucs and the rest. That is exactly wrong: this venv is
-# isolated, so it starts with nothing, and the trimming produced three build
-# cycles that each died on a different missing module. An isolated venv needs
-# the whole list.
-# xformers: audiocraft imports it at module load (audiocraft/modules/
-# transformer.py), so without it the import check dies. Pinned to the version
-# whose cp312 wheel matches the torch above — see the note there.
-xformers==0.0.28.post3
-av>=11.0.0
-einops
-flashy>=0.0.1
-hydra-core>=1.1
-hydra_colorlog
-julius
-num2words
-sentencepiece
-spacy>=3.6.1
-huggingface_hub
-tqdm
-transformers>=4.31.0
-demucs
-librosa
-gradio
-torchmetrics
-encodec
-protobuf
-numpy<2

tests/test_remote_execution.py

@@ -1677,3 +1677,56 @@ def test_musicgen_falls_back_in_order(monkeypatch):
     monkeypatch.setattr(builtins, "__import__", _no_audiocraft)
     got = ag._load_musicgen("facebook/musicgen-small", "cpu")
     assert got == "VENV" and calls == [("venv", "facebook/musicgen-small")]
+
+
+def test_audio_backend_is_selectable_and_travels_to_a_pod(monkeypatch):
+    """audiocraft and transformers are not equivalent — melody conditioning and
+    AudioGen exist only in the first — so the choice must be the operator's, and
+    must not be decided by whatever happens to be installed on a pod."""
+    from codai.api.audio_gen import _audio_backend
+    from codai.api.runpod_worker import pod_plan, parse_model_runpod
+
+    assert _audio_backend(None) == "auto"
+    assert _audio_backend({"audio_backend": "transformers"}) == "transformers"
+    assert _audio_backend({"audio_backend": "nonsense"}) == "auto"
+
+    monkeypatch.setenv("CODERAI_AUDIO_BACKEND", "audiocraft")
+    assert _audio_backend(None) == "audiocraft"
+    # a model's own setting beats the environment it lands in
+    assert _audio_backend({"audio_backend": "transformers"}) == "transformers"
+
+    mcfg = parse_model_runpod({"engine": "coderai", "image": "x:1"})
+    entry = {"path": "facebook/musicgen-small", "model_type": "audio_gen_models",
+             "audio_backend": "transformers"}
+    assert pod_plan(mcfg, "audio_gen", entry=entry)["env"]["CODERAI_AUDIO_BACKEND"] \
+        == "transformers"
+
+    # unset means the pod decides for itself, rather than being told wrongly
+    entry.pop("audio_backend")
+    assert "CODERAI_AUDIO_BACKEND" not in pod_plan(mcfg, "audio_gen", entry=entry)["env"]
+
+
+def test_asking_for_audiocraft_when_absent_is_an_error_not_a_substitution():
+    """Silently serving transformers when audiocraft was asked for would hand
+    back something that cannot do what the caller requested."""
+    import builtins
+    import codai.api.audio_gen as ag
+
+    real = builtins.__import__
+
+    def _no_audiocraft(name, *a, **k):
+        if name.startswith("audiocraft"):
+            raise ImportError("not installed")
+        return real(name, *a, **k)
+
+    builtins.__import__ = _no_audiocraft
+    try:
+        try:
+            ag._load_musicgen("facebook/musicgen-small", "cpu",
+                              {"audio_backend": "audiocraft"})
+        except RuntimeError as exc:
+            assert "audiocraft is not installed" in str(exc)
+        else:
+            raise AssertionError("expected a RuntimeError")
+    finally:
+        builtins.__import__ = real