ace-step: patch upstream infer-api + missing runtime deps + cache mount

Three upstream gaps surfaced once /generate was actually exercised:

  1. infer-api.py builds an 18-arg positional tuple but the pipeline
     expects 24 — first missing arg is `format`, so audio_duration
     shifts into format's slot and the pipeline calls len() on an
     int. Ship a patched copy of infer-api.py and COPY over upstream's
     in the Dockerfile. Also handle empty lora_name_or_path -> "none"
     (empty string trips HF Hub's repo-id validator).
  2. torchcodec + ffmpeg are required by the WAV save path but neither
     is in upstream requirements.txt. Without them every /generate
     runs to completion and then 500s at write-time.
  3. ACE-Step caches checkpoints at /root/.cache/ace-step/checkpoints
     (HARDCODED, not honored by HF_HOME). Mount our persistent dir
     there so the ~7 GB model survives container recreates.

Bench on A6000 (cached model, lo-fi hip hop, 60-step euler/apg):
  10s @ 27 steps -> 9.4s  (0.94x)
  30s @ 60 steps -> 11.2s (0.37x, ~2.7x realtime)
  60s @ 60 steps -> 14.8s (0.24x, ~4x realtime)

playbooks/deploy-ace-step.yaml

@@ -46,6 +46,12 @@ steps:
       dest: "{{ compose_dir }}/Dockerfile"
       mode: "0644"
 
+  - name: Upload patched infer-api.py (fixes upstream 18-arg tuple bug)
+    upload:
+      src: stacks/ace-step/infer-api.py
+      dest: "{{ compose_dir }}/infer-api.py"
+      mode: "0644"
+
   - name: Seed .env from template (only if absent)
     upload:
       src: stacks/ace-step/.env.example

stacks/ace-step/Dockerfile

@@ -35,9 +35,13 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
         git \
         curl \
         ca-certificates \
+        ffmpeg \
     && apt-get clean \
     && rm -rf /var/lib/apt/lists/* \
     && ln -sf /usr/bin/python3 /usr/bin/python
+# ffmpeg is required by torchcodec at runtime — torchcodec dlopens
+# libavcodec/libavformat. Without it the WAV save step fails with
+# "Could not load libtorchcodec".
 
 RUN python -m venv /opt/venv
 ENV PATH="/opt/venv/bin:$PATH"
@@ -60,7 +64,17 @@ RUN pip install --no-cache-dir --upgrade pip \
         --index-url https://download.pytorch.org/whl/cu126 \
     && pip install --no-cache-dir hf_transfer peft \
     && pip install --no-cache-dir -r requirements.txt \
-    && pip install --no-cache-dir .
+    && pip install --no-cache-dir . \
+    && pip install --no-cache-dir torchcodec
+# torchcodec — required by torchaudio's save_with_torchcodec (the save
+# path ACE-Step uses on output). Not pulled in by upstream's
+# requirements.txt; without it, /generate runs to completion and then
+# 500s at the WAV write step.
+
+# Replace upstream's infer-api.py with our patched copy. Upstream's
+# version builds an 18-arg positional tuple but the pipeline expects
+# 24 — see infer-api.py header comment for the fix.
+COPY infer-api.py /app/infer-api.py
 
 EXPOSE 8000
 CMD ["python3", "infer-api.py"]

stacks/ace-step/compose.yaml

@@ -42,7 +42,11 @@ services:
       # ~5-10 GB checkpoint download survives container recreates.
       - HF_HOME=/app/hf_cache
     volumes:
-      - ${ACE_STEP_CHECKPOINTS_DIR}:/app/checkpoints
+      # ACE-Step has a HARDCODED checkpoint cache at
+      # /root/.cache/ace-step/checkpoints — not honored by HF_HOME.
+      # Mount our persistent dir there so model re-pull doesn't
+      # happen on every container recreate.
+      - ${ACE_STEP_CHECKPOINTS_DIR}:/root/.cache/ace-step/checkpoints
       - ${ACE_STEP_OUTPUTS_DIR}:/app/outputs
       - ${ACE_STEP_LOGS_DIR}:/app/logs
       - ${ACE_STEP_CACHE_DIR}:/app/hf_cache

stacks/ace-step/infer-api.py

@@ -0,0 +1,140 @@
+# Patched ACE-Step REST API.
+#
+# Upstream's infer-api.py builds an 18-element positional `params`
+# tuple and splats it into the pipeline. The pipeline's __call__
+# actually expects 24 positional args — upstream's tuple is missing
+# `format` at the head and the audio2audio/lora block at the tail.
+# The result of the missing `format` head: every subsequent arg shifts
+# left, the int `audio_duration` lands in `format`'s slot, and the
+# pipeline blows up with "object of type 'int' has no len()" when it
+# tries to validate the format string.
+#
+# This file copies upstream's behavior verbatim except for the params
+# tuple, which now matches the canonical 24-arg pipeline signature
+# inferred from the Gradio handler in acestep/ui/components.py.
+# (Defaults for the new fields mirror Gradio's defaults: format=wav,
+# audio2audio_enable=False, ref_audio_strength=0.5, no ref audio, no
+# LoRA.)
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+from typing import List, Optional
+import os
+from acestep.pipeline_ace_step import ACEStepPipeline
+import uuid
+
+app = FastAPI(title="ACEStep Pipeline API (patched)")
+
+
+class ACEStepInput(BaseModel):
+    checkpoint_path: str
+    bf16: bool = True
+    torch_compile: bool = False
+    device_id: int = 0
+    output_path: Optional[str] = None
+    audio_duration: float
+    prompt: str
+    lyrics: str
+    infer_step: int
+    guidance_scale: float
+    scheduler_type: str
+    cfg_type: str
+    omega_scale: float
+    actual_seeds: List[int]
+    guidance_interval: float
+    guidance_interval_decay: float
+    min_guidance_scale: float
+    use_erg_tag: bool
+    use_erg_lyric: bool
+    use_erg_diffusion: bool
+    oss_steps: List[int]
+    guidance_scale_text: float = 0.0
+    guidance_scale_lyric: float = 0.0
+    # Optional new fields exposed for completeness — Gradio defaults.
+    # lora_name_or_path: empty string trips HF Hub's repo-id validator
+    # ("Repo id ... cannot ... be empty"), so default to None and only
+    # forward it as a string when actually set.
+    audio_format: str = "wav"
+    audio2audio_enable: bool = False
+    ref_audio_strength: float = 0.5
+    ref_audio_input: Optional[str] = None
+    lora_name_or_path: Optional[str] = None
+    lora_weight: float = 1.0
+
+
+class ACEStepOutput(BaseModel):
+    status: str
+    output_path: Optional[str]
+    message: str
+
+
+def initialize_pipeline(
+    checkpoint_path: str, bf16: bool, torch_compile: bool, device_id: int
+) -> ACEStepPipeline:
+    os.environ["CUDA_VISIBLE_DEVICES"] = str(device_id)
+    return ACEStepPipeline(
+        checkpoint_dir=checkpoint_path,
+        dtype="bfloat16" if bf16 else "float32",
+        torch_compile=torch_compile,
+    )
+
+
+@app.post("/generate", response_model=ACEStepOutput)
+async def generate_audio(input_data: ACEStepInput):
+    try:
+        model_demo = initialize_pipeline(
+            input_data.checkpoint_path,
+            input_data.bf16,
+            input_data.torch_compile,
+            input_data.device_id,
+        )
+
+        params = (
+            input_data.audio_format,
+            input_data.audio_duration,
+            input_data.prompt,
+            input_data.lyrics,
+            input_data.infer_step,
+            input_data.guidance_scale,
+            input_data.scheduler_type,
+            input_data.cfg_type,
+            input_data.omega_scale,
+            ", ".join(map(str, input_data.actual_seeds)),
+            input_data.guidance_interval,
+            input_data.guidance_interval_decay,
+            input_data.min_guidance_scale,
+            input_data.use_erg_tag,
+            input_data.use_erg_lyric,
+            input_data.use_erg_diffusion,
+            ", ".join(map(str, input_data.oss_steps)),
+            input_data.guidance_scale_text,
+            input_data.guidance_scale_lyric,
+            input_data.audio2audio_enable,
+            input_data.ref_audio_strength,
+            input_data.ref_audio_input,
+            input_data.lora_name_or_path or "none",
+            input_data.lora_weight,
+        )
+
+        output_path = input_data.output_path or f"/app/outputs/output_{uuid.uuid4().hex}.wav"
+
+        model_demo(*params, save_path=output_path)
+
+        return ACEStepOutput(
+            status="success",
+            output_path=output_path,
+            message="Audio generated successfully",
+        )
+
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Error generating audio: {str(e)}")
+
+
+@app.get("/health")
+async def health_check():
+    return {"status": "healthy"}
+
+
+if __name__ == "__main__":
+    import uvicorn
+
+    uvicorn.run(app, host="0.0.0.0", port=8000)