kiapi 0.1.0__py3-none-any.whl

kiapi/capabilities/acestep/_helpers/register.py

@@ -0,0 +1,96 @@
+"""Register acestep's models + OpenAPI description in the global registries."""
+
+import os
+
+from kiapi.core.capability import CapabilitySpec, capability_spec_registry
+from kiapi.core.model import ModelSpec, model_registry
+from kiapi.core.setup import HfSnapshotResource, PythonVenvResource
+
+from .._constants.description import DESCRIPTION
+from .._models import acestep
+from .._operations.resolve_ace_step_paths import resolve_ace_step_paths
+from .._settings import settings_manager
+
+_ACESTEP_PACKAGE = "ace-step @ git+https://github.com/ace-step/ACE-Step-1.5.git"
+
+
+def register() -> None:
+    paths = resolve_ace_step_paths(settings_manager.get_settings())
+    capability_spec_registry.register(
+        CapabilitySpec(
+            name="acestep",
+            domain="audio",
+            title="kiapi ACE-Step API",
+            summary="Generate songs, covers, repaints, and stems with ACE-Step.",
+            description=DESCRIPTION,
+            openapi_path="/v1/audio/acestep/openapi.json",
+            docs_path="/v1/audio/acestep/docs",
+            redoc_path="/v1/audio/acestep/redoc",
+            path_prefixes=("/v1/audio/acestep",),
+        )
+    )
+
+    # framework="rss" -> kiapi-RSS-based measurement (child RSS not captured, so the
+    # estimate below is what the budget uses). resident=True -> held until evicted/TTL.
+    model_registry.register(
+        ModelSpec(
+            name="xl-base",
+            family="acestep",
+            domain="audio",
+            repo="acestep-v15-xl-base",
+            module=acestep,
+            weight_gb=27.0,  # est. turbo 12.2 GB + (19-4.5) DiT delta
+            peak_headroom_gb=6.0,
+            framework="rss",
+            priority=0,
+            default=True,
+            setup_resources=(
+                PythonVenvResource(
+                    path=paths.venv_path,
+                    python="3.12",
+                    packages=(os.environ.get("KIAPI_ACESTEP_SPEC", _ACESTEP_PACKAGE),),
+                    import_name="acestep",
+                    label_name="acestep-venv",
+                    disk_gb=1.3,
+                ),
+                HfSnapshotResource(
+                    repo="ACE-Step/Ace-Step1.5",
+                    disk_gb=9.4,
+                    local_dir=paths.checkpoint_dir,
+                ),
+                HfSnapshotResource(
+                    repo="ACE-Step/acestep-v15-xl-base",
+                    disk_gb=19.0,
+                    local_dir=f"{paths.checkpoint_dir}/acestep-v15-xl-base",
+                ),
+            ),
+        )
+    )
+    model_registry.register(
+        ModelSpec(
+            name="turbo",
+            family="acestep",
+            domain="audio",
+            repo="acestep-v15-turbo",
+            module=acestep,
+            weight_gb=12.0,  # measured child RSS ~12.2 GB (LLM + turbo DiT)
+            peak_headroom_gb=4.0,
+            framework="rss",
+            priority=0,
+            setup_resources=(
+                PythonVenvResource(
+                    path=paths.venv_path,
+                    python="3.12",
+                    packages=(os.environ.get("KIAPI_ACESTEP_SPEC", _ACESTEP_PACKAGE),),
+                    import_name="acestep",
+                    label_name="acestep-venv",
+                    disk_gb=1.3,
+                ),
+                HfSnapshotResource(
+                    repo="ACE-Step/Ace-Step1.5",
+                    disk_gb=9.4,
+                    local_dir=paths.checkpoint_dir,
+                ),
+            ),
+        )
+    )

kiapi/capabilities/acestep/_models/acestep.py

@@ -0,0 +1,105 @@
+"""Handler for ace-step presets (xl-base / turbo) via an isolated subprocess.
+
+Each preset is a resident spec whose ``payload`` is an :class:`AceStepEngine` - a
+long-lived ace-step worker subprocess holding the shared LLM + that preset's DiT.
+``load`` spawns it (paying the model-load cost once), ``release`` terminates it
+(freeing memory on eviction / TTL / shutdown), and ``generate_track`` runs one
+generation over the protocol and stores the produced WAV in the global file
+store. The memory manager treats the subprocess's footprint via the registry
+``weight_gb`` estimate (a child process's RSS isn't visible in kiapi's own).
+"""
+
+import logging
+import time
+from pathlib import Path
+from typing import Any
+
+import psutil  # type: ignore
+
+from kiapi.core.file import FileStore
+from kiapi.core.job import ProgressReporter
+from kiapi.core.model import ModelSpec
+from kiapi.core.workdir import create_work_dir
+
+from .._operations.resolve_ace_step_paths import resolve_ace_step_paths
+from .._services.ace_step_engine import AceStepEngine
+from .._settings import settings_manager
+from .._views.track_params import TrackParams
+
+FEATURES = {"text", "audio"}
+logger = logging.getLogger(__name__)
+
+
+def load(spec: ModelSpec) -> AceStepEngine:
+    s = settings_manager.get_settings()
+    paths = resolve_ace_step_paths(s)
+    Path(paths.project_root).mkdir(parents=True, exist_ok=True)
+    engine = AceStepEngine.spawn(
+        python_path=paths.python_path,
+        preset_name=spec.name,
+        project_root=paths.project_root,
+        checkpoint_dir=paths.checkpoint_dir,
+        llm_model=s.llm_model,
+        ready_timeout_s=s.ready_timeout_s,
+    )
+    # The child's memory isn't in kiapi's RSS; log it so the registry estimate
+    # (weight_gb) can be tuned to the real footprint.
+    try:
+        rss_gb = psutil.Process(engine.proc.pid).memory_info().rss / (1024**3)
+        logger.info(
+            "%s subprocess ready - child RSS ~%.1f GB (registry estimate %.1f)",
+            spec.name,
+            rss_gb,
+            spec.weight_gb,
+        )
+    except Exception:
+        pass
+    return engine
+
+
+def release(payload: AceStepEngine) -> None:
+    payload.release()
+
+
+def warmup(payload: AceStepEngine) -> None:
+    """No-op: spawning already loaded the LLM + DiT weights. A real generation
+    would cost ~25 s, which warmup shouldn't pay."""
+
+
+def generate_track(
+    payload: AceStepEngine,
+    params: TrackParams,
+    files: FileStore,
+) -> dict[str, Any]:
+    """Run one generation over the subprocess, store the WAV, return metadata."""
+    s = settings_manager.get_settings()
+    save_dir = create_work_dir("audio/acestep")
+    t0 = time.time()
+    engine_params = params.engine_params()
+    # generate_track runs on the single worker thread, so the job's ambient
+    # ProgressReporter is bound here; ace-step's native milestones (forwarded over
+    # the subprocess protocol) push straight onto the job (no-op when unbound).
+    reporter = ProgressReporter.current()
+    audio_path = payload.generate(
+        params.task,
+        engine_params,
+        str(save_dir),
+        timeout_s=s.job_timeout_s,
+        on_progress=reporter.update,
+    )
+    total_s = round(time.time() - t0, 2)
+
+    meta = {
+        "task": params.task,
+        "params": engine_params,
+        "timings": {"total_s": total_s},
+        **params.meta_extra(),
+    }
+    rec = files.put_path(
+        audio_path,
+        filename=Path(audio_path).name,
+        content_type="audio/wav",
+        meta=meta,
+        move=True,
+    )
+    return {"file_id": rec.file_id, "audio_bytes": rec.size, **meta}

kiapi/capabilities/acestep/_operations/handle_cover.py

@@ -0,0 +1,23 @@
+"""ACE-Step cover service entry (worker-thread thunk body)."""
+
+from kiapi.capabilities import resolve_file_ref
+from kiapi.core.app import AppContext
+from kiapi.core.file import FileID
+from kiapi.core.job import JobResult
+
+from .._views.cover_request import CoverRequest
+from .resolve_cover_params import resolve_cover_params
+from .resolve_engine import resolve_engine
+
+
+def handle_cover(ctx: AppContext, req: CoverRequest) -> tuple[JobResult, list[FileID]]:
+    spec, module, engine = resolve_engine(ctx, req.model)
+    source = resolve_file_ref(ctx.file_store, req.source, kind="source")
+    params = resolve_cover_params(
+        req,
+        variant=spec.name,
+        source_file_id=source.file_id,
+        src_audio=source.path,
+    )
+    result = module.generate_track(engine, params, ctx.file_store)
+    return result, [result["file_id"]]

kiapi/capabilities/acestep/_operations/handle_extract.py

@@ -0,0 +1,38 @@
+"""ACE-Step extract service entry (worker-thread thunk body)."""
+
+from kiapi.capabilities import resolve_file_ref
+from kiapi.core.app import AppContext
+from kiapi.core.file import FileID
+from kiapi.core.job import JobResult
+
+from .._views.extract_request import ExtractRequest
+from .resolve_engine import resolve_engine
+from .resolve_extract_params import resolve_extract_params
+
+
+def handle_extract(
+    ctx: AppContext, req: ExtractRequest
+) -> tuple[JobResult, list[FileID]]:
+    spec, module, engine = resolve_engine(ctx, req.model)
+    source = resolve_file_ref(ctx.file_store, req.source, kind="source")
+
+    stems = []
+    artifacts = []
+    for target in req.targets:
+        params = resolve_extract_params(
+            req,
+            variant=spec.name,
+            source_file_id=source.file_id,
+            src_audio=source.path,
+            target=target,
+        )
+        res = module.generate_track(engine, params, ctx.file_store)
+        stems.append(res)
+        artifacts.append(res["file_id"])
+
+    return {
+        "task": "extract",
+        "source_file_id": source.file_id,
+        "targets": req.targets,
+        "stems": stems,
+    }, artifacts

kiapi/capabilities/acestep/_operations/handle_generate.py

@@ -0,0 +1,18 @@
+"""ACE-Step generate service entry (worker-thread thunk body)."""
+
+from kiapi.core.app import AppContext
+from kiapi.core.file import FileID
+from kiapi.core.job import JobResult
+
+from .._views.generate_request import GenerateRequest
+from .resolve_engine import resolve_engine
+from .resolve_generate_params import resolve_generate_params
+
+
+def handle_generate(
+    ctx: AppContext, req: GenerateRequest
+) -> tuple[JobResult, list[FileID]]:
+    spec, module, engine = resolve_engine(ctx, req.model)
+    params = resolve_generate_params(req, variant=spec.name)
+    result = module.generate_track(engine, params, ctx.file_store)
+    return result, [result["file_id"]]

kiapi/capabilities/acestep/_operations/handle_repaint.py

@@ -0,0 +1,25 @@
+"""ACE-Step repaint service entry (worker-thread thunk body)."""
+
+from kiapi.capabilities import resolve_file_ref
+from kiapi.core.app import AppContext
+from kiapi.core.file import FileID
+from kiapi.core.job import JobResult
+
+from .._views.repaint_request import RepaintRequest
+from .resolve_engine import resolve_engine
+from .resolve_repaint_params import resolve_repaint_params
+
+
+def handle_repaint(
+    ctx: AppContext, req: RepaintRequest
+) -> tuple[JobResult, list[FileID]]:
+    spec, module, engine = resolve_engine(ctx, req.model)
+    source = resolve_file_ref(ctx.file_store, req.source, kind="source")
+    params = resolve_repaint_params(
+        req,
+        variant=spec.name,
+        source_file_id=source.file_id,
+        src_audio=source.path,
+    )
+    result = module.generate_track(engine, params, ctx.file_store)
+    return result, [result["file_id"]]

kiapi/capabilities/acestep/_operations/resolve_ace_step_paths.py

@@ -0,0 +1,30 @@
+from pathlib import Path
+
+from kiapi.core.app import get_user_data_dir
+
+from .._schemas.ace_step_paths import AceStepPaths
+from .._settings import AceStepSettings
+
+
+def resolve_ace_step_paths(settings: AceStepSettings) -> AceStepPaths:
+    root = Path(get_user_data_dir()).expanduser() / "acestep"
+    python_path = (
+        Path(settings.python_path).expanduser()
+        if settings.python_path is not None
+        else root / ".venv" / "bin" / "python"
+    )
+    project_root = (
+        Path(settings.project_root).expanduser()
+        if settings.project_root is not None
+        else root / "project"
+    )
+    checkpoint_dir = (
+        Path(settings.checkpoint_dir).expanduser()
+        if settings.checkpoint_dir is not None
+        else root / "checkpoints"
+    )
+    return AceStepPaths.from_paths(
+        python_path=python_path,
+        project_root=project_root,
+        checkpoint_dir=checkpoint_dir,
+    )

kiapi/capabilities/acestep/_operations/resolve_cover_params.py

@@ -0,0 +1,21 @@
+"""Merge an ACE-Step cover request into the model-layer contract."""
+
+from .._views.cover_params import CoverParams
+from .._views.cover_request import CoverRequest
+
+
+def resolve_cover_params(
+    req: CoverRequest, *, variant: str, source_file_id: str, src_audio: str
+) -> CoverParams:
+    return CoverParams(
+        model=variant,
+        src=source_file_id,
+        src_audio=src_audio,
+        prompt=req.prompt,
+        strength=req.strength,
+        duration=req.duration,
+        seed=req.seed,
+        inference_steps=req.inference_steps,
+        guidance_scale=req.guidance_scale,
+        shift=req.shift,
+    )

kiapi/capabilities/acestep/_operations/resolve_engine.py

@@ -0,0 +1,10 @@
+from typing import Any
+
+from kiapi.core.app import AppContext
+from kiapi.core.model import ModelSpec, model_registry
+
+
+def resolve_engine(ctx: AppContext, model: str) -> tuple[ModelSpec, Any, object]:
+    spec = model_registry.resolve("acestep", model)
+    ctx.ensure_model_ready(spec)
+    return spec, spec.module, ctx.memory_manager.acquire(spec)

kiapi/capabilities/acestep/_operations/resolve_extract_params.py

@@ -0,0 +1,24 @@
+"""Merge an ACE-Step extract request into one stem's model-layer contract."""
+
+from .._views.extract_params import ExtractParams
+from .._views.extract_request import ExtractRequest
+
+
+def resolve_extract_params(
+    req: ExtractRequest,
+    *,
+    variant: str,
+    source_file_id: str,
+    src_audio: str,
+    target: str,
+) -> ExtractParams:
+    return ExtractParams(
+        model=variant,
+        src=source_file_id,
+        src_audio=src_audio,
+        target=target,
+        seed=req.seed,
+        inference_steps=req.inference_steps,
+        guidance_scale=req.guidance_scale,
+        shift=req.shift,
+    )

kiapi/capabilities/acestep/_operations/resolve_generate_params.py

@@ -0,0 +1,18 @@
+"""Merge an ACE-Step generate request into the model-layer contract."""
+
+from .._views.generate_params import GenerateParams
+from .._views.generate_request import GenerateRequest
+
+
+def resolve_generate_params(req: GenerateRequest, *, variant: str) -> GenerateParams:
+    return GenerateParams(
+        model=variant,
+        prompt=req.prompt,
+        lyrics=req.lyrics,
+        duration=req.duration,
+        lang=req.lang,
+        seed=req.seed,
+        inference_steps=req.inference_steps,
+        guidance_scale=req.guidance_scale,
+        shift=req.shift,
+    )

kiapi/capabilities/acestep/_operations/resolve_repaint_params.py

@@ -0,0 +1,22 @@
+"""Merge an ACE-Step repaint request into the model-layer contract."""
+
+from .._views.repaint_params import RepaintParams
+from .._views.repaint_request import RepaintRequest
+
+
+def resolve_repaint_params(
+    req: RepaintRequest, *, variant: str, source_file_id: str, src_audio: str
+) -> RepaintParams:
+    return RepaintParams(
+        model=variant,
+        src=source_file_id,
+        src_audio=src_audio,
+        prompt=req.prompt,
+        start=req.start,
+        end=req.end,
+        strength=req.strength,
+        seed=req.seed,
+        inference_steps=req.inference_steps,
+        guidance_scale=req.guidance_scale,
+        shift=req.shift,
+    )

kiapi/capabilities/acestep/_schemas/ace_step_paths.py

@@ -0,0 +1,29 @@
+from pathlib import Path
+
+from pydantic import BaseModel, ConfigDict
+
+
+class AceStepPaths(BaseModel):
+    """Resolved local paths used by the ACE-Step subprocess and setup resources."""
+
+    model_config = ConfigDict(frozen=True)
+
+    python_path: str
+    venv_path: str
+    project_root: str
+    checkpoint_dir: str
+
+    @classmethod
+    def from_paths(
+        cls,
+        *,
+        python_path: Path,
+        project_root: Path,
+        checkpoint_dir: Path,
+    ) -> "AceStepPaths":
+        return cls(
+            python_path=str(python_path),
+            venv_path=str(python_path.parent.parent),
+            project_root=str(project_root),
+            checkpoint_dir=str(checkpoint_dir),
+        )

kiapi/capabilities/acestep/_services/ace_step_engine.py

@@ -0,0 +1,143 @@
+"""kiapi-side proxy for the ace-step worker subprocess.
+
+An :class:`AceStepEngine` owns one long-lived subprocess (launched from the
+ace-step venv) holding the LLM + one preset's DiT resident. It is the ``payload``
+the memory manager keeps for an ace-step spec; ``release()`` (called on eviction /
+TTL / shutdown) terminates the subprocess and frees its memory.
+
+Only ever used from the single worker thread (single-flight), so the simple
+request/response over pipes needs no concurrency beyond a guard lock.
+"""
+
+import json
+import subprocess
+import threading
+import time
+from collections.abc import Callable
+from pathlib import Path
+
+from .._exceptions.ace_step_engine_error import AceStepEngineError
+
+SENTINEL = "@@KIAPI@@"
+_WORKER = str(Path(__file__).parents[1] / "worker_subprocess.py")
+
+
+class AceStepEngine:
+    def __init__(self, proc: subprocess.Popen) -> None:
+        self.proc = proc
+        self._lock = threading.Lock()
+        self._counter = 0
+
+    # -- lifecycle ------------------------------------------------------------
+
+    @classmethod
+    def spawn(
+        cls,
+        *,
+        python_path: str,
+        preset_name: str,
+        project_root: str,
+        checkpoint_dir: str,
+        llm_model: str,
+        ready_timeout_s: float,
+    ) -> "AceStepEngine":
+        proc = subprocess.Popen(
+            [
+                python_path,
+                _WORKER,
+                "--preset-name",
+                preset_name,
+                "--project-root",
+                project_root,
+                "--checkpoint-dir",
+                checkpoint_dir,
+                "--llm-model",
+                llm_model,
+            ],
+            stdin=subprocess.PIPE,
+            stdout=subprocess.PIPE,
+            stderr=None,  # inherit: ace-step logs stream to kiapi's stderr
+            text=True,
+            bufsize=1,
+        )
+        engine = cls(proc)
+        engine._await_event("ready", timeout_s=ready_timeout_s)
+        return engine
+
+    def release(self) -> None:
+        proc = self.proc
+        if proc.poll() is None:
+            try:
+                proc.terminate()
+                proc.wait(timeout=10)
+            except Exception:
+                proc.kill()
+        for stream in (proc.stdin, proc.stdout):
+            try:
+                if stream is not None:
+                    stream.close()
+            except Exception:
+                pass
+
+    # -- protocol -------------------------------------------------------------
+
+    def _read_proto(self, timeout_s: float | None = None) -> dict:
+        """Read the next ``@@KIAPI@@`` protocol line, skipping any other noise."""
+        deadline = (time.monotonic() + timeout_s) if timeout_s else None
+        assert self.proc.stdout is not None
+        for line in self.proc.stdout:  # blocks; ace-step noise on fd1 is skipped
+            line = line.rstrip("\n")
+            if line.startswith(SENTINEL):
+                return json.loads(line[len(SENTINEL) :])  # type: ignore
+            if deadline is not None and time.monotonic() > deadline:
+                break
+        if self.proc.poll() is not None:
+            raise AceStepEngineError(
+                f"ace-step worker exited (code {self.proc.returncode}) before replying"
+            )
+        raise AceStepEngineError("ace-step worker produced no protocol output")
+
+    def _await_event(self, event: str, *, timeout_s: float) -> dict:
+        msg = self._read_proto(timeout_s=timeout_s)
+        if msg.get("event") == "error":
+            raise AceStepEngineError(msg.get("error", "worker init error"))
+        if msg.get("event") != event:
+            raise AceStepEngineError(f"expected event {event!r}, got {msg!r}")
+        return msg
+
+    def generate(
+        self,
+        task: str,
+        params: dict,
+        save_dir: str,
+        *,
+        timeout_s: float,
+        on_progress: Callable[[float, str | None], None] | None = None,
+    ) -> str:
+        with self._lock:
+            self._counter += 1
+            rid = str(self._counter)
+            req = {"id": rid, "task": task, "params": params, "save_dir": save_dir}
+            if self.proc.stdin is None or self.proc.poll() is not None:
+                raise AceStepEngineError("ace-step worker is not running")
+            self.proc.stdin.write(json.dumps(req, ensure_ascii=False) + "\n")
+            self.proc.stdin.flush()
+
+            deadline = time.monotonic() + timeout_s
+            while True:
+                remaining = deadline - time.monotonic()
+                if remaining <= 0:
+                    raise AceStepEngineError("ace-step worker timed out")
+                msg = self._read_proto(timeout_s=remaining)
+
+                if msg.get("event") == "progress":
+                    # Forward only our own progress; ignore any stray/stale line.
+                    if on_progress is not None and msg.get("id") == rid:
+                        on_progress(float(msg.get("fraction", 0.0)), msg.get("label"))
+                    continue
+
+                if msg.get("id") != rid:
+                    raise AceStepEngineError(f"reply id mismatch: {msg!r}")
+                if not msg.get("ok"):
+                    raise AceStepEngineError(msg.get("error", "generation failed"))
+                return msg["path"]  # type: ignore