write songs with a model the server cannot hold in its own interpreter

ACE-Step generates a whole track -- style prompt, sung lyrics, a length in
seconds -- and it ships a package pinning transformers 4.50 and spacy 3.8.4.
spacy 3.8.4 publishes no cp313 wheel, so this is not a stack the server would
rather not hold: it is one that cannot be installed on the server's
interpreter. It gets a python 3.12 venv of its own and an HTTP service in it,
driven the way CrisperWhisper, Canary and pyannote are. Built on the first
request to such a model, so a model that is configured and never called costs
nothing; `uv venv` covers the hosts that have python3.12 without ensurepip.

Three upstream details that fail silently rather than loudly:

* the pipeline resolves its own dtype as `bfloat16 if dtype == "bfloat16" else
  float32`, so asking it for float16 quietly doubles the VRAM. ACE_PIPELINE_DTYPE
  is read after that switch, so precision goes through there instead.
* constructing the pipeline loads NOTHING -- it reads the checkpoints on the
  first call. coderai measures a model's footprint as the free-VRAM delta
  across the load, and that window closes when the worker reports healthy, so
  a lazy load would measure ~0 GB and eviction would treat a resident 3.5B
  model as free. The service loads before answering /health.
* `quantized=True` does not quantize: it loads a separate int4wo export from
  the same directory layout. Missing, that is a FileNotFoundError minutes into
  a load, so the worker names the file it wants first.

The handle cached in the manager exposes cleanup(), which is what eviction
calls, and here it stops the worker -- the only thing that frees its VRAM.

`duration` and `cfg_coef` are shared with MusicGen and default to 10s/3.0,
which would make every ACE-Step call a 10-second clip at a guidance it is not
tuned for. The branch reads which fields the client actually SENT, so an
omitted one gets ACE-Step's default and an explicit `duration: 10` still means
10. `lyrics` and `steps` become declared fields: both branches already read
`steps` off the request, where only extra="allow" was letting it through.

The venv is ~6.4 GB, so the default image does not carry it and --with-acestep
bakes it in -- copying the dev venv when there is one, building it when there
is not. Ubuntu 22.04's system python is 3.10, so a baked 3.12 venv also needs
a relocatable interpreter bundled beside it, like the 3.10 venvs.

Verified on a 3090: healthy after ~3.5 min, 8.5 GB VRAM, a real 48 kHz stereo
WAV from a 10s/12-step request. Suite at its 15-failure baseline, 5292 passed.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_014T3rZnLhydX9QqUcfBbqLu

codai/api/audio_gen.py

@@ -314,6 +314,30 @@ def _generate_audio(pipe, model_name: str, request: AudioGenerationRequest, task
         audio_np = wav[0, 0].cpu().numpy()
         sr = pipe.sample_rate
 
+    elif model_type == 'ace_step':
+        # The worker returns a finished wav: generation happened in another
+        # interpreter, so there is no array to resample or clip here, and the
+        # pipeline exposes no per-step callback to report progress from.
+        _aud_progress_reset(0, unit="s")
+        # Which fields the CLIENT actually sent, not which have a value: the
+        # shared request defaults are MusicGen's (10s, cfg 3.0) and would make
+        # every ACE-Step call a 10-second clip at a guidance it is not tuned
+        # for. A field left out falls back to ACE-Step's own default; a field
+        # explicitly set to 10 still means 10.
+        _sent = getattr(request, 'model_fields_set', set())
+        wav = pipe.generate(
+            prompt=request.prompt,
+            lyrics=getattr(request, 'lyrics', None) or "",
+            duration=(float(request.duration) if 'duration' in _sent
+                      and request.duration else 60.0),
+            steps=int(getattr(request, 'steps', None) or 60),
+            guidance_scale=(float(request.cfg_coef) if 'cfg_coef' in _sent
+                            and request.cfg_coef is not None else 15.0),
+            seed=request.seed,
+        )
+        task_registry.raise_if_cancelled(task_id)
+        return wav, 'wav'
+
     elif model_type == 'stable_audio':
         num_steps = int(getattr(request, 'steps', None) or 100)
         _aud_progress_reset(num_steps, unit="it")
@@ -439,16 +463,14 @@ async def audio_generate(request: AudioGenerationRequest, http_request: Request
                     pipe = await asyncio.get_event_loop().run_in_executor(
                         None, _load_stable_audio, model_name, device, _ag_cfg)
                 elif model_type == 'ace_step':
-                    # ACE-Step ships its own package, which pins a stack this
-                    # venv cannot hold. It belongs in an isolated-venv worker
-                    # like CrisperWhisper and Canary; until that exists, say
-                    # so plainly rather than loading it as a MusicGen and
-                    # failing with a shape error three minutes later.
-                    raise RuntimeError(
-                        f"{model_name} is an ACE-Step model, which needs the "
-                        "`acestep` package in its own venv; coderai has no "
-                        "ACE-Step worker yet. Use a Stable Audio or MusicGen "
-                        "model for now.")
+                    # ACE-Step ships its own package pinning a stack this venv
+                    # cannot hold, so it runs in an isolated venv behind a
+                    # local HTTP service, like CrisperWhisper and Canary. The
+                    # handle cached here exposes cleanup(), so a model swap
+                    # evicts it and frees its VRAM like any other model.
+                    from codai.api.acestep_worker import AceStep
+                    pipe = await asyncio.get_event_loop().run_in_executor(
+                        None, AceStep, model_name, _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

@@ -1454,7 +1454,8 @@ def _plan(engine, image, args, mcfg, api_key, entry, served,
     for _var, _name in (("CODERAI_PYANNOTE_VENV", "pyannote"),
                         ("CODERAI_NEMO_VENV", "nemo"),
                         ("CODERAI_CRISPERWHISPER_VENV", "transformers"),
-                        ("CODERAI_XTTS_VENV", "TTS")):
+                        ("CODERAI_XTTS_VENV", "TTS"),
+                        ("CODERAI_ACESTEP_VENV", "acestep")):
         env.setdefault(_var, f"/opt/coderai/venvs/{_name}")
 
     if cap in ("tts", "voice"):

codai/pydantic/audiogenrequest.py

@@ -29,6 +29,15 @@ class AudioGenerationRequest(BaseModel):
     temperature: Optional[float] = 1.0
     cfg_coef: Optional[float] = 3.0        # classifier-free guidance coefficient
     seed: Optional[int] = None
+    # Diffusion models (Stable Audio, ACE-Step) count steps; the autoregressive
+    # ones (MusicGen, AudioGen) ignore this. Both branches already read it off
+    # the request -- undeclared it only arrived because `extra="allow"` let it
+    # through, so it was absent from the schema and from the docs.
+    steps: Optional[int] = None
+    # Sung lyrics, for the models that take them separately from the style
+    # prompt (ACE-Step). Supports its [verse]/[chorus] tags and the special
+    # values the model documents, e.g. "[inst]" for an instrumental.
+    lyrics: Optional[str] = None
     # Reference audio for melody conditioning (MusicGen Melody)
     melody: Optional[str] = None           # base64/URL
     # Voice profile for singing/speech conditioning

docs/README.md

@@ -31,6 +31,7 @@ reference; these are the deep dives.
 | [expressive-tts.md](expressive-tts.md) | Expressive text-to-speech |
 | [image-prompts.md](image-prompts.md) | Image prompts: the `(tag:1.3)` / `[tag]` / `BREAK` grammar, prompts past CLIP's 77 tokens, `clip_skip`, negative prompts, textual-inversion embeddings, size validation |
 | [zimage-lora-training.md](zimage-lora-training.md) | LoRA training for image models |
+| [text-to-music.md](text-to-music.md) | `POST /v1/audio/generate` with ACE-Step: style prompt + sung lyrics → a 48 kHz stereo track, the python-3.12 worker it needs, and the two image variants (`--with-acestep`) |
 | [dtype-auto-selection.md](dtype-auto-selection.md) | How model dtype is chosen automatically |
 | [model-defaults-and-sources.md](model-defaults-and-sources.md) | A model's own `defaults` (steps/guidance/sampler/modes/licence), the sampler names and the flow-matching boundary, and where weights come from: repo, single file, `org/repo:file` or `civitai:<versionId>` |
 

docs/text-to-music.md

+# Text to music: ACE-Step
+
+`POST /v1/audio/generate` with an ACE-Step model writes a song — a style prompt,
+optional sung lyrics, a length in seconds. It joins MusicGen, AudioGen,
+AudioLDM2 and Stable Audio on the same endpoint; the model's name selects the
+branch, so nothing about the request shape is ACE-Step-specific.
+
+```bash
+curl -s http://localhost:8000/v1/audio/generate \
+  -H "Authorization: Bearer $KEY" -H 'Content-Type: application/json' \
+  -d '{"model":"ACE-Step",
+       "prompt":"slow synthwave, warm analog pads, 90 bpm",
+       "lyrics":"[verse]\nthe tide goes out\n[chorus]\nand takes the light",
+       "duration":60, "steps":60, "seed":42}'
+```
+
+| field | ACE-Step default | notes |
+|---|---|---|
+| `prompt` | — | the *style*, not the words: genre, instruments, tempo, mood |
+| `lyrics` | `""` | supports its `[verse]` / `[chorus]` tags; `[inst]` for an instrumental |
+| `duration` | 60 s | the whole track is generated at once, not stitched |
+| `steps` | 60 | diffusion steps |
+| `cfg_coef` | 15.0 | guidance; ACE-Step wants far more of it than MusicGen |
+| `seed` | random | reproducible with everything else held equal |
+
+Output is 48 kHz stereo WAV.
+
+`duration` and `cfg_coef` are shared across every model on this endpoint and
+their field defaults are MusicGen's (10 s, 3.0). An ACE-Step call that leaves
+them out gets **ACE-Step's** defaults instead, because a 10-second song at a
+guidance it was not tuned for is not a useful reading of "unspecified". A value
+you actually send is always used, including `duration: 10`.
+
+## Why it runs in its own interpreter
+
+ACE-Step ships its own `acestep` package, which pins `transformers==4.50` and
+`spacy==3.8.4`. The server runs transformers 5.x, and spacy 3.8.4 publishes no
+cp313 wheel at all — so this is not merely a dependency the server would rather
+not hold, it is one that **cannot be installed on the server's interpreter**.
+
+So it runs the way CrisperWhisper, Canary, pyannote and MeloTTS do: a python
+**3.12** venv of its own, with `tools/acestep_service.py` inside it speaking
+HTTP on a local port, driven by `codai/api/acestep_worker.py`.
+
+- The venv is built on the **first request** to an ACE-Step model and reused
+  afterwards. A model that is configured but never called costs nothing.
+- `python3.12 -m venv` is tried first; on hosts that have `/usr/bin/python3.12`
+  without the `python3.12-venv` package (common on Debian) it falls back to
+  `uv venv --python 3.12`, which will also fetch a 3.12 if the host has none.
+- `CODERAI_ACESTEP_VENV` overrides the location; `service_url` (or
+  `CODERAI_ACESTEP_SERVICE_URL`) points at a worker running elsewhere instead,
+  and then nothing is built, spawned or downloaded locally.
+
+The weights come from whatever the model's `path` resolves to. An HF snapshot of
+`ACE-Step/ACE-Step-v1-3.5B` already has the four subdirectories the pipeline
+expects (`ace_step_transformer`, `music_dcae_f8c8`, `music_vocoder`,
+`umt5-base`), so the worker points the pipeline straight at the snapshot
+coderai already has rather than letting it fetch a second copy.
+
+### It evicts like any other model
+
+The handle cached in the model manager exposes `cleanup()`, which is what the
+manager calls when it evicts a model to make room for the next one — and here
+that stops the worker process, which is the only thing that frees its VRAM.
+
+The service also loads the checkpoints **before** it answers its health check.
+That is not tidiness: coderai measures a model's footprint as the device-wide
+free-VRAM delta across the load, and that window closes when the worker reports
+healthy. Constructing `ACEStepPipeline` loads nothing — upstream reads the
+checkpoints on the first `__call__` — so a lazy load would measure ~0 GB and
+eviction would treat a resident 3.5B model as free. Measured on a 3090:
+**~8.5 GB** at bf16.
+
+### Precision and quantization
+
+`precision` is honoured exactly. Upstream resolves its own `dtype` argument as
+`torch.bfloat16 if dtype == "bfloat16" else torch.float32`, so asking it for
+float16 silently yields float32 — twice the VRAM, no error. The worker sets
+`ACE_PIPELINE_DTYPE`, which is read after that switch, instead.
+
+The quantization switches (`load_in_8bit` / `load_in_4bit`) ask upstream for
+`quantized=True`, which does **not** quantize anything: it loads ACE-Step's
+separate int4wo export (`*_int4wo.bin`) from the same directory layout. If those
+files are not in the checkpoint directory the worker says so at load time,
+naming the missing file, rather than failing with `FileNotFoundError` minutes
+in. Point the model's path at the quantized checkpoint, or leave the switches
+off.
+
+## The two image variants
+
+The venv is ~6.4 GB, so the **default image does not carry it** — it builds on
+first use, like every other isolated worker. A second variant bakes it in, for
+installs that would rather pay the size than the first-call wait:
+
+```bash
+packaging/linux/build_oci_image.sh                  # default: no ACE-Step venv
+packaging/linux/build_oci_image.sh --with-acestep   # venv pre-built, baked in
+```
+
+`--with-acestep` **copies the development venv** when there is one
+(`$CODERAI_ACESTEP_VENV`, else `~/.coderai/acestep_venv`) — it is already built
+and already the thing that was tested. Only when there is none does it build one
+with `uv`. `--acestep-venv PATH` bundles a specific venv and implies
+`--with-acestep`.
+
+Because the base image is Ubuntu 22.04, whose system python is 3.10, a baked
+3.12 venv also needs a **relocatable 3.12 interpreter** bundled beside it; the
+build fetches one with `uv python install 3.12` if the host has none, and the
+image re-points the venv's `pyvenv.cfg` at `/opt/coderai/py312` — the same
+arrangement the 3.10 lip-sync and LongCat venvs use.
+
+Both variants resolve the venv the same way at runtime, so neither needs an
+environment variable: the worker uses `/opt/coderai/acestep_venv` (local image)
+or `/opt/coderai/venvs/acestep` (pod images) when one is there, and
+`~/.coderai/acestep_venv` otherwise.

packaging/linux/Dockerfile.oci-venv

@@ -58,7 +58,11 @@ RUN set -eux; \
       mkdir -p /opt/coderai/py310; \
       rsync -a /tmp/local-bundle/py310/ /opt/coderai/py310/; \
     fi; \
-    for d in lipsync_venv longcat_venv Wav2Lip SadTalker LongCat-Video ds4 colibri kimi-k3-in-c; do \
+    if [ -d /tmp/local-bundle/py312 ]; then \
+      mkdir -p /opt/coderai/py312; \
+      rsync -a /tmp/local-bundle/py312/ /opt/coderai/py312/; \
+    fi; \
+    for d in lipsync_venv longcat_venv acestep_venv Wav2Lip SadTalker LongCat-Video ds4 colibri kimi-k3-in-c; do \
       if [ -d "/tmp/local-bundle/$d" ]; then \
         mkdir -p "/opt/coderai/$d"; \
         rsync -a "/tmp/local-bundle/$d/" "/opt/coderai/$d/"; \
@@ -71,6 +75,11 @@ RUN set -eux; \
         for p in python python3 python3.10; do ln -sf /opt/coderai/py310/bin/python3.10 "/opt/coderai/$v/bin/$p"; done; \
       fi; \
     done; \
+    cfg=/opt/coderai/acestep_venv/pyvenv.cfg; \
+    if [ -f "$cfg" ]; then \
+      sed -i 's|^home *=.*|home = /opt/coderai/py312/bin|; s|^command *=.*|command = /opt/coderai/py312/bin/python3.12|' "$cfg"; \
+      for p in python python3 python3.12; do ln -sf /opt/coderai/py312/bin/python3.12 "/opt/coderai/acestep_venv/bin/$p"; done; \
+    fi; \
     if [ -d /tmp/local-bundle/local-bin ]; then \
       mkdir -p /opt/coderai/staged-local-bin; \
       rsync -a /tmp/local-bundle/local-bin/ /opt/coderai/staged-local-bin/; \

packaging/linux/build_oci_image.sh

@@ -36,6 +36,14 @@ LONGCAT_VENV="${CODERAI_LONGCAT_VENV:-$HOME/.coderai/longcat_venv}"
 # …and its CODE: LongCatVideoPipeline lives in the repo's own `longcat_video` package,
 # not on PyPI, so the source is bundled like Wav2Lip/SadTalker. Weights are NOT.
 LONGCAT_SRC="${CODERAI_LONGCAT_SRC:-$HOME/.coderai/LongCat-Video}"
+# ACE-Step (text-to-music) runs in a PYTHON 3.12 venv of its own -- it pins
+# spacy==3.8.4, which publishes no cp313 wheel, so it cannot share the server's
+# interpreter at all. The venv is ~6.4 GB, which is why it is OFF by default:
+# the slim image builds it on the first request to an ACE-Step model, exactly
+# as every other isolated worker does, and the --with-acestep variant bakes it
+# in for installs that would rather pay the size than the first-call wait.
+ACESTEP_VENV="${CODERAI_ACESTEP_VENV:-$HOME/.coderai/acestep_venv}"
+INCLUDE_ACESTEP=0
 WAV2LIP_DIR="${CODERAI_WAV2LIP_SRC:-$HOME/.coderai/Wav2Lip}"
 SADTALKER_DIR="${CODERAI_SADTALKER_SRC:-$HOME/.coderai/SadTalker}"
 DS4_DIR="${CODERAI_DS4_DIR:-$HOME/.coderai/ds4}"
@@ -80,6 +88,12 @@ Options:
                           $CODERAI_PARLER_VENV or ~/.coderai/parler_venv if present).
   --no-parler             Do not bundle the Parler-TTS venv overlay.
   --no-tools              Do not bundle the lip-sync (wav2lip/sadtalker) venvs or ds4.
+  --with-acestep          Bake the ACE-Step python-3.12 venv into the image (~6.4 GB).
+                          Off by default: without it the worker builds the venv on the
+                          first request to an ACE-Step model. Copies the existing dev
+                          venv ($CODERAI_ACESTEP_VENV or ~/.coderai/acestep_venv) when
+                          there is one, and builds it with uv when there is not.
+  --acestep-venv PATH     Bundle this ACE-Step venv (implies --with-acestep).
   -t, --tag TAG           Image tag to create (default: coderai:local or OCI_IMAGE from versions.env).
   --versioned             Auto-derive the tag from the app version + build mode + GPU:
                             local-venv build -> coderai:full_<gpu>_<version>
@@ -154,6 +168,21 @@ while [[ $# -gt 0 ]]; do
       INCLUDE_TOOLS=0
       shift
       ;;
+    --with-acestep)
+      BUILD_MODE="venv"
+      INCLUDE_ACESTEP=1
+      shift
+      ;;
+    --acestep-venv)
+      BUILD_MODE="venv"
+      if [[ $# -lt 2 ]]; then
+        echo "Error: --acestep-venv requires a path" >&2
+        exit 2
+      fi
+      ACESTEP_VENV="$2"
+      INCLUDE_ACESTEP=1
+      shift 2
+      ;;
     --include-local-bin)
       BUILD_MODE="venv"
       if [[ $# -lt 2 ]]; then
@@ -413,6 +442,70 @@ prepare_venv_bundle() {
     fi
   fi
 
+  # ACE-Step's python-3.12 venv, and the interpreter it has to run on. The base
+  # image is Ubuntu 22.04, whose system python is 3.10, so a 3.12 venv needs its
+  # own standalone interpreter bundled alongside it and re-pointed during the
+  # image build -- the same arrangement as the py3.10 venvs below.
+  if [[ "$INCLUDE_ACESTEP" == "1" ]]; then
+    # The dev install is the preferred source: it is already built and already
+    # the thing being tested. Only build one when there is none.
+    if [[ ! -d "$ACESTEP_VENV" ]]; then
+      echo "No ACE-Step venv at $ACESTEP_VENV — building it (needs uv; several GB)…"
+      local _uv
+      _uv="$(command -v uv || true)"
+      [[ -z "$_uv" && -x "$HOME/.local/bin/uv" ]] && _uv="$HOME/.local/bin/uv"
+      if [[ -z "$_uv" ]]; then
+        echo "Error: --with-acestep needs either an existing venv at $ACESTEP_VENV or uv to build one" >&2
+        exit 2
+      fi
+      # unsafe-best-match is required, not a tuning knob: uv's default is
+      # first-index-wins, so with the PyTorch index in that requirements file it
+      # resolves unrelated packages off it and pins e.g. requests 2.28.1.
+      "$_uv" venv --python 3.12 "$ACESTEP_VENV"
+      "$_uv" pip install --python "$ACESTEP_VENV/bin/python" \
+             --index-strategy unsafe-best-match \
+             -r "$ROOT_DIR/requirements-acestep.txt"
+    fi
+    if [[ -d "$ACESTEP_VENV" ]]; then
+      local _ace_excl=(--exclude '__pycache__' --exclude '*.pyc' --exclude '*.dist-info/RECORD')
+      rsync -a "${_ace_excl[@]}" "$ACESTEP_VENV/" "$bundle/acestep_venv/"
+      echo "Bundled ACE-Step venv from: $ACESTEP_VENV"
+      # Find a RELOCATABLE 3.12. The dev venv's pyvenv.cfg usually says
+      # `home = /usr/bin`, which is the host's own python and cannot be copied
+      # into the image, so that case falls through to a uv-managed build.
+      local py312_dir=""
+      local _home
+      _home="$(sed -n 's/^home *= *//p' "$ACESTEP_VENV/pyvenv.cfg" 2>/dev/null | head -1)"
+      if [[ -n "$_home" && "$_home" != "/usr/bin" && "$_home" != "/usr/local/bin" ]]; then
+        py312_dir="$(dirname "$_home")"
+      fi
+      if [[ -z "$py312_dir" ]]; then
+        py312_dir="$(find "$HOME/.local/share/uv/python" -maxdepth 1 -type d \
+                     -name 'cpython-3.12*-linux-*' 2>/dev/null | sort | tail -1)"
+      fi
+      if [[ -z "$py312_dir" ]]; then
+        local _uv2
+        _uv2="$(command -v uv || true)"
+        [[ -z "$_uv2" && -x "$HOME/.local/bin/uv" ]] && _uv2="$HOME/.local/bin/uv"
+        if [[ -n "$_uv2" ]]; then
+          echo "Fetching a relocatable Python 3.12 for the image…"
+          "$_uv2" python install 3.12 >/dev/null 2>&1 || true
+          py312_dir="$(find "$HOME/.local/share/uv/python" -maxdepth 1 -type d \
+                       -name 'cpython-3.12*-linux-*' 2>/dev/null | sort | tail -1)"
+        fi
+      fi
+      if [[ -n "$py312_dir" && -d "$py312_dir" ]]; then
+        mkdir -p "$bundle/py312"
+        rsync -a "$py312_dir/" "$bundle/py312/"
+        echo "Bundled standalone Python 3.12 from: $py312_dir"
+      else
+        echo "Error: bundled the ACE-Step venv but found no relocatable Python 3.12 to run it on." >&2
+        echo "       Install one with 'uv python install 3.12' and rebuild, or drop --with-acestep." >&2
+        exit 2
+      fi
+    fi
+  fi
+
   # Isolated lip-sync tools, LongCat-Video and the ds4 engine. The py3.10 venvs share
   # one standalone Python 3.10 (read from whichever venv's pyvenv.cfg `home` is present);
   # it's bundled once and every venv is re-pointed at it during the image build.

requirements-acestep.txt

@@ -5,23 +5,47 @@
 # server cannot hold -- the server runs transformers 5.x and a much newer
 # diffusers -- so it runs in its own venv and coderai talks to it over HTTP,
 # the same shape as CrisperWhisper and NeMo-Canary
-# (codai/api/acestep_worker.py + tools/acestep_service.py, both still to be
-# written). Do NOT merge these into requirements.txt.
+# (codai/api/acestep_worker.py + tools/acestep_service.py). Do NOT merge these
+# into requirements.txt.
 #
-# NOT YET INSTALLED. The `acestep` package is not on PyPI, so it installs from
-# the upstream git repository, and nothing here has been run or verified. The
-# audio_gen loader therefore refuses an ACE-Step model with an explicit
-# message rather than loading it as a MusicGen and failing with a shape error
-# minutes later. Install this venv, then introspect the package before
-# writing the service -- the API below the pin is not guessed at here.
+# The `acestep` package is not on PyPI, so it installs from the upstream git
+# repository at a pinned commit.
+#
+# BUILD THIS VENV WITH PYTHON 3.12, not the server's 3.13:
+#
+#     python3.12 -m venv ~/.coderai/acestep_venv
+#
+# ...except that a Debian host often has /usr/bin/python3.12 WITHOUT the
+# python3.12-venv package, so ensurepip is missing and that command fails. The
+# worker falls back to `uv venv --python 3.12`, which also fetches a 3.12 if
+# the host has none; with uv, `--index-strategy unsafe-best-match` is
+# mandatory (its default is first-index-wins, which resolves every shared name
+# off the PyTorch index below and pins e.g. requests 2.28.1).
+#
+# ACE-Step pins spacy==3.8.4, which has no cp313 wheel (3.8.7+ do, 3.8.4 does
+# not), so the install cannot be satisfied on 3.13 at all. Upstream's pin is
+# exact and respecting it is cheaper than fighting it. codai/api/lipsync's
+# venv already sets the precedent of a per-worker interpreter (3.10).
 #
 # CUDA build for the 3090; cu124 wheels run on this box's newer driver.
+#
+# 2.6.0 and not 2.5.1: there is no torchaudio==2.5.1+cu124 wheel for the
+# interpreters this can run on -- torch 2.5.1 has one, torchaudio does not, so
+# the pair cannot be satisfied. ACE-Step itself pins neither (its
+# requirements.txt asks for bare torch/torchaudio/torchvision), so the pair
+# here only has to be internally consistent.
 
 --extra-index-url https://download.pytorch.org/whl/cu124
-torch==2.5.1+cu124
-torchaudio==2.5.1+cu124
+torch==2.6.0+cu124
+torchaudio==2.6.0+cu124
+torchvision==0.21.0+cu124
 soundfile>=0.12.0
 
-# ACE-Step itself, from the upstream repository. Pin the commit you install so
-# a rebuild is reproducible:
-#   acestep @ git+https://github.com/ace-step/ACE-Step.git@<commit>
+# ACE-Step itself, from the upstream repository. Pinned to a commit, not a
+# branch, so a rebuild gets the same code: `main` would silently change what
+# this venv contains between builds.
+# The DISTRIBUTION is ace-step (setup.py name="ace_step"); the importable
+# module is `acestep`. Naming it `acestep` here makes pip discard the wheel it
+# just built: "has inconsistent name: expected 'acestep', but metadata has
+# 'ace-step'".
+ace-step @ git+https://github.com/ace-step/ACE-Step.git@1bee4c9f5b43e30995f8d4d33b3919197ce1bd68

tests/test_acestep_worker.py

+"""ACE-Step: an isolated-venv worker that evicts like any other model.
+
+ACE-Step pins spacy==3.8.4, which publishes no cp313 wheel, so it runs on its
+own python 3.12 behind a local HTTP service. The things worth pinning here are
+the ones that fail silently rather than loudly:
+
+* a precision that upstream's two-way dtype switch would turn into float32,
+* request defaults that belong to MusicGen leaking into an ACE-Step call,
+* a handle that does not free its VRAM when the manager evicts it.
+"""
+import importlib.util
+import os
+import sys
+import types
+from pathlib import Path
+
+import pytest
+
+from codai.api import acestep_worker
+from codai.api.audio_gen import _detect_audio_gen_type, _generate_audio
+from codai.pydantic.audiogenrequest import AudioGenerationRequest
+
+_SERVICE = Path(__file__).resolve().parents[1] / "tools" / "acestep_service.py"
+
+
+def _load_service():
+    """Import the service module the way nothing else can: it normally runs in
+    the 3.12 venv, but its top level is stdlib-only, so it imports here."""
+    spec = importlib.util.spec_from_file_location("_acestep_service", _SERVICE)
+    mod = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(mod)
+    return mod
+
+
+# --------------------------------------------------------------------------- #
+# routing
+# --------------------------------------------------------------------------- #
+
+@pytest.mark.parametrize("name", [
+    "ACE-Step/ACE-Step-v1-3.5B",
+    "/AI/models/acestep-v1",
+    "ace-step-v1-3.5b",
+])
+def test_an_ace_step_model_routes_to_the_ace_step_branch(name):
+    assert _detect_audio_gen_type(name) == "ace_step"
+
+
+def test_other_audio_models_are_untouched():
+    assert _detect_audio_gen_type("facebook/musicgen-small") == "musicgen"
+    assert _detect_audio_gen_type("stabilityai/stable-audio-open-1.0") == "stable_audio"
+
+
+# --------------------------------------------------------------------------- #
+# precision
+# --------------------------------------------------------------------------- #
+
+@pytest.mark.parametrize("precision,expect", [
+    (None, "bfloat16"),
+    ("bf16", "bfloat16"),
+    ("f16", "float16"),
+    ("fp16", "float16"),
+    ("f32", "float32"),
+    ("fp32", "float32"),
+    ("nonsense", "bfloat16"),
+])
+def test_precision_becomes_a_name_torch_actually_has(precision, expect):
+    cfg = {} if precision is None else {"precision": precision}
+    assert acestep_worker._dtype_name(cfg) == expect
+
+
+def test_the_service_forces_the_dtype_past_the_pipelines_two_way_switch():
+    """Upstream resolves `torch.bfloat16 if dtype == "bfloat16" else
+    torch.float32`, so asking it for float16 gets float32 -- twice the VRAM,
+    no error. ACE_PIPELINE_DTYPE is read after that switch, so the service
+    must set it, or `precision: fp16` is a lie."""
+    src = _SERVICE.read_text()
+    assert 'os.environ["ACE_PIPELINE_DTYPE"] = a.dtype' in src
+
+
+# --------------------------------------------------------------------------- #
+# venv resolution
+# --------------------------------------------------------------------------- #
+
+def test_an_explicit_venv_wins(monkeypatch):
+    monkeypatch.setenv("CODERAI_ACESTEP_VENV", "/somewhere/else")
+    assert acestep_worker.resolve_venv_dir() == Path("/somewhere/else")
+
+
+def test_without_a_baked_venv_it_falls_back_under_the_home_dir(monkeypatch):
+    monkeypatch.delenv("CODERAI_ACESTEP_VENV", raising=False)
+    monkeypatch.setattr(acestep_worker, "_BAKED_VENVS",
+                        (Path("/nonexistent/a"), Path("/nonexistent/b")))
+    assert acestep_worker.resolve_venv_dir() == Path(
+        os.path.expanduser("~/.coderai/acestep_venv"))
+
+
+def test_both_baked_layouts_are_accepted(monkeypatch, tmp_path):
+    """The local image bakes /opt/coderai/<name>_venv, the pod images
+    /opt/coderai/venvs/<name>. A worker that knows only one of them rebuilds a
+    6 GB venv on a machine rented by the second."""
+    monkeypatch.delenv("CODERAI_ACESTEP_VENV", raising=False)
+    first, second = tmp_path / "a", tmp_path / "b"
+    (second / "bin").mkdir(parents=True)
+    monkeypatch.setattr(acestep_worker, "_BAKED_VENVS", (first, second))
+    assert acestep_worker.resolve_venv_dir() == second
+
+
+# --------------------------------------------------------------------------- #
+# generation: the defaults are ACE-Step's, not MusicGen's
+# --------------------------------------------------------------------------- #
+
+class _FakePipe:
+    def __init__(self):
+        self.calls = []
+
+    def generate(self, **kw):
+        self.calls.append(kw)
+        return b"RIFF....fake wav"
+
+
+def _run(**fields):
+    pipe = _FakePipe()
+    req = AudioGenerationRequest(model="ACE-Step/ACE-Step-v1-3.5B",
+                                 prompt="a slow piano ballad", **fields)
+    out, ext = _generate_audio(pipe, "ACE-Step/ACE-Step-v1-3.5B", req)
+    return pipe.calls[0], out, ext
+
+
+def test_the_wav_from_the_worker_is_returned_as_it_arrives():
+    """It is already a finished wav. Re-encoding it here would mean decoding
+    audio this process has no reason to touch."""
+    _kw, out, ext = _run()
+    assert (out, ext) == (b"RIFF....fake wav", "wav")
+
+
+def test_an_unsent_duration_means_ace_steps_default_not_musicgens():
+    kw, _o, _e = _run()
+    assert kw["duration"] == 60.0
+    assert kw["guidance_scale"] == 15.0
+    assert kw["steps"] == 60
+
+
+def test_a_duration_the_client_really_asked_for_is_honoured_even_at_the_shared_default():
+    """10.0 is the shared field default AND a legitimate request. Guessing
+    'unset' from the value would silently stretch it to 60 seconds."""
+    kw, _o, _e = _run(duration=10.0, cfg_coef=3.0)
+    assert kw["duration"] == 10.0
+    assert kw["guidance_scale"] == 3.0
+
+
+def test_lyrics_and_steps_reach_the_worker():
+    kw, _o, _e = _run(lyrics="[verse]\nthe tide goes out", steps=27, seed=7)
+    assert kw["lyrics"].startswith("[verse]")
+    assert kw["steps"] == 27
+    assert kw["seed"] == 7
+
+
+def test_lyrics_and_steps_are_declared_fields_not_just_extras():
+    """They were readable only because the request model allows extras, so
+    they were missing from the schema and from the generated docs."""
+    assert {"lyrics", "steps"} <= set(AudioGenerationRequest.model_fields)
+
+
+# --------------------------------------------------------------------------- #
+# eviction
+# --------------------------------------------------------------------------- #
+
+def test_the_handle_frees_its_vram_the_way_the_manager_asks(monkeypatch):
+    """The manager's unload dispatch looks for cleanup() first. For a worker
+    the VRAM is in another process, so cleanup() must stop it -- nothing else
+    frees it."""
+    stopped = []
+    monkeypatch.setattr(acestep_worker, "ensure_service",
+                        lambda name, cfg=None: "http://127.0.0.1:1")
+    monkeypatch.setattr(acestep_worker, "stop_service", stopped.append)
+    handle = acestep_worker.AceStep("ACE-Step/ACE-Step-v1-3.5B", {})
+    assert callable(getattr(handle, "cleanup", None))
+    handle.cleanup()
+    assert stopped == ["ACE-Step/ACE-Step-v1-3.5B"]
+
+
+def test_the_weights_load_before_the_service_reports_healthy():
+    """coderai measures a model's VRAM as the free-memory delta across the
+    load, and that window closes when the worker goes healthy. Constructing
+    ACEStepPipeline loads nothing -- it reads the checkpoints on first call --
+    so the service has to load them itself, or the model measures ~0 GB and
+    eviction treats a resident 3.5B model as free."""
+    src = _SERVICE.read_text()
+    load_at = min(src.index("load_checkpoint"), src.index("ENGINE._ensure()"))
+    assert load_at < src.index("serve_forever()")
+
+
+def test_quantized_without_the_int4_weights_says_what_is_missing():
+    """`quantized=True` does not quantize: it loads ACE-Step's separate int4wo
+    export. Against an ordinary checkpoint that is a FileNotFoundError several
+    minutes into a load."""
+    mod = _load_service()
+    eng = mod.Engine(checkpoint_dir="/nonexistent/ckpt", quantized=True)
+    with pytest.raises(RuntimeError) as err:
+        eng._ensure()
+    msg = str(err.value)
+    assert "int4" in msg and "/nonexistent/ckpt" in msg

tests/test_ltx2_audio_and_engines.py

@@ -63,17 +63,25 @@ def test_stable_audio_is_not_routed_to_audioldm():
     assert _detect_audio_gen_type("stabilityai/stable-audio-open-1.0") != "audioldm"
 
 
-def test_stable_audio_has_a_loader_and_ace_step_says_why_it_has_none():
+def test_stable_audio_and_ace_step_each_have_their_own_loader():
+    """ACE-Step used to raise here, naming the venv it needed; it now has that
+    venv and a worker in it (see tests/test_acestep_worker.py). What must not
+    come back is the old failure mode: loading it as a MusicGen and failing
+    with a shape error minutes later."""
     import inspect
 
     from codai.api import audio_gen
     assert hasattr(audio_gen, "_load_stable_audio")
     src = inspect.getsource(audio_gen)
-    i = src.index("elif model_type == 'ace_step':")
-    window = src[i:i + 700]
-    # Named plainly rather than loaded as a MusicGen and failing with a shape
-    # error minutes later.
-    assert "acestep" in window and "own venv" in window
+    # Two branches now carry that condition -- one to load, one to generate --
+    # so take the LAST, which is the loader.
+    load = src[src.rindex("elif model_type == 'ace_step':"):][:900]
+    assert "acestep_worker" in load and "AceStep" in load
+    assert "_load_musicgen" not in load
+    # ...and the generate branch returns the worker's wav rather than falling
+    # through to the array-to-int16 writer, which would see no audio_np at all.
+    gen = src[src.index("elif model_type == 'ace_step':"):][:1400]
+    assert "return wav, 'wav'" in gen
 
 
 def test_a_stereo_waveform_is_written_as_stereo():

tools/acestep_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.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program. If not, see <https://www.gnu.org/licenses/>.
+
+"""ACE-Step text-to-music, run in its own interpreter.
+
+    POST /generate  -> audio/wav bytes
+                       body: {"prompt", "lyrics", "duration", "steps",
+                              "guidance_scale", "seed", "scheduler_type"}
+    GET  /health    -> {"ok": true}
+
+This file runs inside ``~/.coderai/acestep_venv`` (or the image's baked
+/opt/coderai/acestep_venv), NEVER in the server's venv, and it must not import
+anything from ``codai``. Two reasons it cannot share:
+
+* ACE-Step pins ``transformers==4.50.0`` and ``spacy==3.8.4``; the server runs
+  transformers 5.x. Installing ACE-Step into the server venv would downgrade
+  it, exactly as parler-tts does.
+* ``spacy==3.8.4`` has no cp313 wheel, so this venv is PYTHON 3.12 while the
+  server is 3.13. They cannot be the same interpreter, let alone the same venv.
+
+The pipeline writes .wav files and returns their paths; this reads the first
+one back and returns the bytes, because that is the shape coderai's audio_gen
+expects from every other backend.
+"""
+
+import argparse
+import json
+import os
+import sys
+import tempfile
+import threading
+import traceback
+from http.server import BaseHTTPRequestHandler, HTTPServer
+
+ENGINE = None
+_LOCK = threading.Lock()
+
+
+class Engine:
+    """Holds the loaded ACEStepPipeline. Loaded once, on first use."""
+
+    def __init__(self, checkpoint_dir=None, device_id=0, dtype="bfloat16",
+                 cpu_offload=False, quantized=False, overlapped_decode=False):
+        self.kw = dict(checkpoint_dir=checkpoint_dir or None,
+                       device_id=int(device_id), dtype=dtype,
+                       cpu_offload=bool(cpu_offload),
+                       quantized=bool(quantized),
+                       overlapped_decode=bool(overlapped_decode))
+        self.pipe = None
+
+    def _ensure(self):
+        if self.pipe is not None:
+            return
+        # `quantized` does not quantize anything: it loads a DIFFERENT set of
+        # weights (ACE-Step's int4wo export) from the same directory layout. If
+        # those files are not there the failure is a bare FileNotFoundError
+        # minutes in, so say what is missing first -- before the import, which
+        # pulls in torch and several seconds of it.
+        ckpt = self.kw.get("checkpoint_dir")
+        if self.kw.get("quantized") and ckpt:
+            need = os.path.join(ckpt, "ace_step_transformer",
+                                "diffusion_pytorch_model_int4wo.bin")
+            if not os.path.isfile(need):
+                raise RuntimeError(
+                    "this model is configured quantized, but ACE-Step's "
+                    f"int4 weights are not in {ckpt} (expected {need}). "
+                    "Quantized ACE-Step is a separate upstream checkpoint, "
+                    "not a conversion of this one: point the model's path at "
+                    "it, or turn the quantization switch off.")
+        from acestep.pipeline_ace_step import ACEStepPipeline
+        sys.stderr.write(f"acestep: loading pipeline {self.kw}\n")
+        sys.stderr.flush()
+        self.pipe = ACEStepPipeline(**self.kw)
+        # Constructing the pipeline loads NOTHING -- it sets `loaded = False`
+        # and reads the checkpoints on the first __call__. Do it here, so the
+        # VRAM this model needs is allocated while the parent is still
+        # measuring the load (see main()).
+        if self.kw.get("quantized"):
+            self.pipe.load_quantized_checkpoint(self.pipe.checkpoint_dir)
+        else:
+            self.pipe.load_checkpoint(self.pipe.checkpoint_dir)
+        sys.stderr.write("acestep: checkpoints loaded\n")
+        sys.stderr.flush()
+
+    def generate(self, body: dict) -> bytes:
+        self._ensure()
+        # manual_seeds is a LIST in this API; a bare int is silently ignored,
+        # which would make `seed` look honoured while every call differed.
+        seed = body.get("seed")
+        seeds = None
+        if seed is not None:
+            seeds = [int(seed)] if not isinstance(seed, (list, tuple)) else list(seed)
+
+        with tempfile.TemporaryDirectory() as d:
+            out = os.path.join(d, "out.wav")
+            call = dict(
+                format="wav",
+                prompt=body.get("prompt") or "",
+                # Empty string, not None: the pipeline treats lyrics as text and
+                # None reaches the tokenizer.
+                lyrics=body.get("lyrics") or "",
+                audio_duration=float(body.get("duration") or 60.0),
+                infer_step=int(body.get("steps") or 60),
+                guidance_scale=float(body.get("guidance_scale") or 15.0),
+                scheduler_type=str(body.get("scheduler_type") or "euler"),
+                manual_seeds=seeds,
+                save_path=out,
+            )
+            paths = self.pipe(**call)
+            wavs = [p for p in (paths or [])
+                    if isinstance(p, str) and p.lower().endswith(".wav")
+                    and os.path.isfile(p)]
+            if not wavs:
+                raise RuntimeError(
+                    f"ACE-Step produced no .wav (returned {paths!r})")
+            with open(wavs[0], "rb") as f:
+                return f.read()
+
+
+class Handler(BaseHTTPRequestHandler):
+    def _send(self, code, body=b"", ctype="application/json"):
+        self.send_response(code)
+        self.send_header("Content-Type", ctype)
+        self.send_header("Content-Length", str(len(body)))
+        self.end_headers()
+        if body:
+            self.wfile.write(body)
+
+    def log_message(self, fmt, *args):      # quieter: the parent pumps our stderr
+        sys.stderr.write("acestep: " + (fmt % args) + "\n")
+
+    def do_GET(self):
+        if self.path.rstrip("/") == "/health":
+            self._send(200, json.dumps({"ok": True}).encode())
+        else:
+            self._send(404, b'{"error":"not found"}')
+
+    def do_POST(self):
+        if self.path.rstrip("/") != "/generate":
+            self._send(404, b'{"error":"not found"}')
+            return
+        try:
+            n = int(self.headers.get("Content-Length") or 0)
+            body = json.loads(self.rfile.read(n) or b"{}")
+        except Exception as e:
+            self._send(400, json.dumps({"error": f"bad request: {e}"}).encode())
+            return
+        # One generation at a time: a single pipeline on one GPU, and two
+        # concurrent calls would interleave on the same modules.
+        try:
+            with _LOCK:
+                wav = ENGINE.generate(body)
+        except Exception as e:
+            traceback.print_exc(file=sys.stderr)
+            self._send(500, json.dumps({"error": f"{type(e).__name__}: {e}"}).encode())
+            return
+        self._send(200, wav, ctype="audio/wav")
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--host", default="127.0.0.1")
+    ap.add_argument("--port", type=int, required=True)
+    ap.add_argument("--checkpoint-dir", default=None)
+    ap.add_argument("--device-id", type=int, default=0)
+    ap.add_argument("--dtype", default="bfloat16")
+    ap.add_argument("--cpu-offload", action="store_true")
+    ap.add_argument("--quantized", action="store_true")
+    ap.add_argument("--overlapped-decode", action="store_true")
+    a = ap.parse_args()
+
+    # Honour the configured precision EXACTLY. The pipeline's own `dtype`
+    # argument is a two-way switch -- `torch.bfloat16 if dtype == "bfloat16"
+    # else torch.float32` -- so asking it for float16 quietly gets float32,
+    # twice the VRAM on a 3.5B model. ACE_PIPELINE_DTYPE is read after that
+    # switch and takes any torch dtype name, so it is the one that actually
+    # works.
+    import torch
+    if not isinstance(getattr(torch, a.dtype, None), torch.dtype):
+        sys.stderr.write(f"acestep: unknown dtype {a.dtype!r}\n")
+        return 2
+    os.environ["ACE_PIPELINE_DTYPE"] = a.dtype
+
+    global ENGINE
+    ENGINE = Engine(checkpoint_dir=a.checkpoint_dir, device_id=a.device_id,
+                    dtype=a.dtype, cpu_offload=a.cpu_offload,
+                    quantized=a.quantized,
+                    overlapped_decode=a.overlapped_decode)
+    # Load the weights BEFORE answering /health, as canary_service does. Not a
+    # detail: coderai measures a model's VRAM as the device-wide free-memory
+    # delta across the load, and that window closes when the parent sees this
+    # process go healthy. Loading lazily on the first /generate would measure
+    # ~0 GB, and eviction would then treat a resident 3.5B model as free and
+    # over-commit the GPU. The parent watches for the process dying, so taking
+    # minutes here is safe.
+    ENGINE._ensure()
+    srv = HTTPServer((a.host, a.port), Handler)
+    sys.stderr.write(f"acestep: ready, listening on {a.host}:{a.port}\n")
+    sys.stderr.flush()
+    srv.serve_forever()
+
+
+if __name__ == "__main__":
+    sys.exit(main() or 0)