From 3b27ed6fba310e9e75e0193c1982d8f90c844382 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Wed, 4 Feb 2026 20:04:51 +0100 Subject: [PATCH] Correctly handle model dir Signed-off-by: Ettore Di Giacinto --- Makefile | 1 + backend/python/ace-step/backend.py | 23 +++++++++++++---------- 2 files changed, 14 insertions(+), 10 deletions(-) diff --git a/Makefile b/Makefile index 92ef60b59..1c5ce3347 100644 --- a/Makefile +++ b/Makefile @@ -473,6 +473,7 @@ BACKEND_QWEN_TTS = qwen-tts|python|.|false|true BACKEND_QWEN_ASR = qwen-asr|python|.|false|true BACKEND_VOXCPM = voxcpm|python|.|false|true BACKEND_WHISPERX = whisperx|python|.|false|true +BACKEND_ACE_STEP = ace-step|python|.|false|true # Helper function to build docker image for a backend # Usage: $(call docker-build-backend,BACKEND_NAME,DOCKERFILE_TYPE,BUILD_CONTEXT,PROGRESS_FLAG,NEEDS_BACKEND_ARG) diff --git a/backend/python/ace-step/backend.py b/backend/python/ace-step/backend.py index a7d635a52..6b61b30cc 100644 --- a/backend/python/ace-step/backend.py +++ b/backend/python/ace-step/backend.py @@ -370,6 +370,7 @@ def _generate_audio_sync(servicer, payload, dst_path): class BackendServicer(backend_pb2_grpc.BackendServicer): def __init__(self): self.model_path = None + self.model_dir = None self.checkpoint_dir = None self.project_root = None self.options = {} @@ -386,20 +387,20 @@ class BackendServicer(backend_pb2_grpc.BackendServicer): model_name = (request.Model or "").strip() model_file = (getattr(request, "ModelFile", None) or "").strip() + # Model dir: where we store checkpoints (always under LocalAI models path, never backend dir) if model_path and model_name: - self.checkpoint_dir = model_path - self.project_root = os.path.dirname(model_path) - self.model_path = os.path.join(model_path, model_name) + model_dir = os.path.join(model_path, model_name) elif model_file: - self.model_path = model_file - self.checkpoint_dir = os.path.dirname(model_file) - self.project_root = os.path.dirname(self.checkpoint_dir) + model_dir = model_file else: - self.model_path = model_name or "." - self.checkpoint_dir = os.path.dirname(self.model_path) if self.model_path else "." - self.project_root = os.path.dirname(self.checkpoint_dir) if self.checkpoint_dir else "." + model_dir = os.path.abspath(model_name or ".") + self.model_dir = model_dir + self.checkpoint_dir = os.path.join(model_dir, "checkpoints") + self.project_root = model_dir + self.model_path = os.path.join(self.checkpoint_dir, model_name or os.path.basename(model_dir.rstrip("/\\"))) - config_path = model_name or os.path.basename(self.model_path.rstrip("/\\")) + config_path = model_name or os.path.basename(model_dir.rstrip("/\\")) + os.makedirs(self.checkpoint_dir, exist_ok=True) # Auto-download DiT model and VAE if missing (same as upstream) if config_path: @@ -413,6 +414,8 @@ class BackendServicer(backend_pb2_grpc.BackendServicer): print(f"[ace-step] Warning: VAE download failed: {e}", file=sys.stderr) self.dit_handler = AceStepHandler() + # Patch handler so it uses our model dir instead of site-packages/checkpoints + self.dit_handler._get_project_root = lambda: self.project_root device = self.options.get("device", "auto") use_flash = self.options.get("use_flash_attention", True) if isinstance(use_flash, str):