diff --git a/playbooks/deploy-ace-step.yaml b/playbooks/deploy-ace-step.yaml index d6dd410..2af9bec 100644 --- a/playbooks/deploy-ace-step.yaml +++ b/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 diff --git a/stacks/ace-step/Dockerfile b/stacks/ace-step/Dockerfile index 250ffc0..44ebb3c 100644 --- a/stacks/ace-step/Dockerfile +++ b/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"] diff --git a/stacks/ace-step/compose.yaml b/stacks/ace-step/compose.yaml index 8de4475..79d2340 100644 --- a/stacks/ace-step/compose.yaml +++ b/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 diff --git a/stacks/ace-step/infer-api.py b/stacks/ace-step/infer-api.py new file mode 100644 index 0000000..d53624d --- /dev/null +++ b/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)