sa3 LoRA Manager

backend/lsdj/controller.py

@@ -27,7 +27,7 @@
 from starlette.datastructures import UploadFile
 from starlette.exceptions import HTTPException as StarletteHTTPException
 
-from . import engine, sa3
+from . import engine, loras, sa3
 from .worker import run_deck_worker
 
 logger = logging.getLogger(__name__)
@@ -445,6 +445,66 @@ def _validate_generate_request(
             )
         options["seed"] = seed
 
+    if "loras" in parsed:
+        stack = parsed["loras"]
+        if not isinstance(stack, list):
+            raise HTTPException(status_code=422, detail="'loras' must be an array")
+        if len(stack) > loras.MAX_LORA_STACK:
+            raise HTTPException(
+                status_code=422,
+                detail=f"'loras' holds at most {loras.MAX_LORA_STACK} adapters",
+            )
+        lora_dirs: list[str] = []
+        lora_strengths: list[float] = []
+        seen: set[str] = set()
+        for entry in stack:
+            if not isinstance(entry, dict):
+                raise HTTPException(
+                    status_code=422, detail="'loras' entries must be objects"
+                )
+            name = entry.get("name")
+            if not isinstance(name, str):
+                raise HTTPException(
+                    status_code=422, detail="'loras[].name' must be a string"
+                )
+            if name in seen:
+                # A repeated adapter would double its delta silently; refuse.
+                raise HTTPException(
+                    status_code=422, detail=f"duplicate adapter {name!r}"
+                )
+            seen.add(name)
+            try:
+                adapter_dir, base = loras.resolve(name)
+            except loras.UnknownAdapter as error:
+                raise HTTPException(status_code=422, detail=str(error)) from None
+            if loras.KIND_BASES[kind] != base:
+                raise HTTPException(
+                    status_code=422,
+                    detail=(
+                        f"adapter '{name}' rides the {base} DiT and cannot apply "
+                        f"to kind '{kind}'"
+                    ),
+                )
+            strength = entry.get("strength", 1.0)
+            if (
+                isinstance(strength, bool)
+                or not isinstance(strength, (int, float))
+                or not math.isfinite(strength)
+                or not loras.MIN_LORA_STRENGTH <= strength <= loras.MAX_LORA_STRENGTH
+            ):
+                raise HTTPException(
+                    status_code=422,
+                    detail=(
+                        "'loras[].strength' must be "
+                        f"{loras.MIN_LORA_STRENGTH:g}-{loras.MAX_LORA_STRENGTH:g}"
+                    ),
+                )
+            lora_dirs.append(str(adapter_dir))
+            lora_strengths.append(float(strength))
+        if lora_dirs:
+            options["lora_dirs"] = lora_dirs
+            options["lora_strengths"] = lora_strengths
+
     if "inpaint_range" in parsed:
         inpaint_range = parsed["inpaint_range"]
         if not isinstance(inpaint_range, list) or len(inpaint_range) != 2:

backend/lsdj/loras.py

@@ -0,0 +1,88 @@
+"""Stable Audio 3 LoRA adapter registry — the read side (issue #66, ADR-0028).
+
+Adapters live on disk under the app-owned data dir, one directory per
+adapter, organised by the DiT family they ride:
+
+    ~/Library/Application Support/LSDJai/sa3-loras/<base>/<slug>/
+
+``base`` is ``small`` (the 1024-wide sm-sfx / sm-music DiTs) or ``medium``
+(the 1536-wide track DiT). An adapter directory holds its ``.safetensors``
+(plus the sibling ``adapter_config.json`` when the PEFT convention applies)
+and the importer's ``lora.json`` manifest. The Rust shell owns the lifecycle
+(import / validate / delete, mirroring the model manager, issue #43); this
+module only reads the registry: the generate path resolves a client-supplied
+adapter name to the directory handed to ``sa3_mlx.py`` as ``--lora``.
+"""
+
+import os
+import pathlib
+import re
+
+# The two DiT families an adapter can ride, and which generation kind uses
+# which. sm-sfx and sm-music share one architecture, so a "small" adapter
+# applies to both kinds; the medium DiT is the track engine (sa3.KINDS).
+BASES = ("small", "medium")
+KIND_BASES = {"sfx": "small", "music": "small", "track": "medium"}
+
+# Trust-boundary bounds for the `--lora-strength` knob (mirrored by
+# `controller.generate_audio`). 0 is the bit-exact bypass (ADR-0028); the
+# spike measured 2.0 as already strong, so 4 is a guard rail, not a UX limit.
+MIN_LORA_STRENGTH = 0.0
+MAX_LORA_STRENGTH = 4.0
+
+# Adapters per generation. The merge stacks linearly (ADR-0028), but an
+# unbounded list is unbounded argv and load time — and past a few adapters
+# the summed deltas swamp the base anyway. Mirrored by the LoraRack UI.
+MAX_LORA_STACK = 4
+
+# One path segment of an adapter name: no separators, no leading dot — the
+# name a client sends can only ever address a directory INSIDE the registry.
+_SLUG = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]*$")
+
+
+class UnknownAdapter(Exception):
+    """An adapter name that does not resolve to an installed adapter."""
+
+
+def loras_dir(
+    env: dict | None = None, home: pathlib.Path | None = None
+) -> pathlib.Path:
+    """The registry root. $SA3_LORAS_HOME wins (tests, dev overrides);
+    otherwise the app-owned data dir, beside the SA3 checkout. Mirrors the
+    Rust `loras::loras_dir`."""
+    env = os.environ if env is None else env
+    home = pathlib.Path.home() if home is None else home
+    override = env.get("SA3_LORAS_HOME", "")
+    if override:
+        return pathlib.Path(override).expanduser()
+    return home / "Library" / "Application Support" / "LSDJai" / "sa3-loras"
+
+
+def _adapter_file(adapter_dir: pathlib.Path) -> pathlib.Path | None:
+    """The adapter's .safetensors inside its directory, or None. The importer
+    writes exactly one; tolerate a hand-placed dir the same way the runtime's
+    `_resolve_path` does (one .safetensors, any name)."""
+    if not adapter_dir.is_dir():
+        return None
+    hits = sorted(
+        entry
+        for entry in adapter_dir.iterdir()
+        if entry.is_file() and entry.suffix == ".safetensors"
+    )
+    return hits[0] if len(hits) == 1 else None
+
+
+def resolve(
+    name: str, env: dict | None = None, home: pathlib.Path | None = None
+) -> tuple[pathlib.Path, str]:
+    """Resolve a client-supplied adapter name (``<base>/<slug>``) to its
+    directory. Raises UnknownAdapter for anything that is not a well-formed
+    name of an installed adapter — malformed names never touch the
+    filesystem, so a name cannot escape the registry root."""
+    base, _, slug = name.partition("/")
+    if base not in BASES or not _SLUG.match(slug):
+        raise UnknownAdapter(f"unknown adapter {name!r}")
+    adapter_dir = loras_dir(env, home) / base / slug
+    if _adapter_file(adapter_dir) is None:
+        raise UnknownAdapter(f"unknown adapter {name!r}")
+    return adapter_dir, base

backend/lsdj/sa3.py

@@ -12,8 +12,9 @@
 import os
 import pathlib
 import tempfile
+from collections.abc import Sequence
 
-# CLI vocabulary of scripts/sa3_mlx.py at the pinned commit (bccf5b7).
+# CLI vocabulary of scripts/sa3_mlx.py at the pinned commit (sa3-pin.json).
 # Pads use the small DiTs with the SAME-S decoder; tracks (M19, ADR-0013)
 # the medium DiT, which pairs with SAME-L.
 KINDS = {"sfx": "sm-sfx", "music": "sm-music", "track": "medium"}
@@ -174,6 +175,8 @@ async def generate(
     cfg: float | None = None,
     apg: float | None = None,
     seed: int | None = None,
+    lora_dirs: Sequence[str] | None = None,
+    lora_strengths: Sequence[float] | None = None,
 ) -> bytes:
     """Run one generation and return the WAV bytes.
 
@@ -220,6 +223,14 @@ async def generate(
                 argv.extend(("--apg", f"{apg:g}"))
             if seed is not None:
                 argv.extend(("--seed", str(seed)))
+            if lora_dirs:
+                # One --lora group per adapter (upstream's PR #57/#65 syntax):
+                # the directory plus its strength=S option. The CLI resolves
+                # the .safetensors inside and merges all deltas at DiT load.
+                for index, lora_dir in enumerate(lora_dirs):
+                    argv.extend(("--lora", lora_dir))
+                    if lora_strengths is not None:
+                        argv.append(f"strength={lora_strengths[index]:g}")
             process = await asyncio.create_subprocess_exec(
                 *argv,
                 cwd=mlx_dir,

backend/tests/test_controller.py

@@ -142,6 +142,136 @@ async def fake_generate(prompt, seconds, kind, **options):
     ]
 
 
+@pytest.fixture
+def lora_registry(tmp_path, monkeypatch):
+    """A registry pointed at via the env override: two small adapters, one
+    medium, plus five smalls (stack0-4) for the over-the-cap case."""
+    slugs = [("small", "crackle"), ("small", "hiss"), ("medium", "maqam")]
+    slugs += [("small", f"stack{n}") for n in range(5)]
+    for base, slug in slugs:
+        adapter_dir = tmp_path / base / slug
+        adapter_dir.mkdir(parents=True)
+        (adapter_dir / "adapter_model.safetensors").write_bytes(b"fake")
+    monkeypatch.setenv("SA3_LORAS_HOME", str(tmp_path))
+    return tmp_path
+
+
+def test_generate_forwards_the_lora_stack_with_aligned_strengths(
+    client, monkeypatch, lora_registry
+):
+    # Two adapters ride together; a missing strength defaults to 1.0 so the
+    # dirs and strengths lists stay aligned position by position.
+    calls = []
+
+    async def fake_generate(prompt, seconds, kind, **options):
+        calls.append(options)
+        return b"RIFFwav"
+
+    monkeypatch.setattr(controller.sa3, "generate", fake_generate)
+    response = client.post(
+        "/api/generate",
+        json=generate_request(
+            loras=[
+                {"name": "small/crackle", "strength": 1.5},
+                {"name": "small/hiss"},
+            ]
+        ),
+    )
+    assert response.status_code == 200
+    assert calls == [
+        {
+            "lora_dirs": [
+                str(lora_registry / "small" / "crackle"),
+                str(lora_registry / "small" / "hiss"),
+            ],
+            "lora_strengths": [1.5, 1.0],
+        }
+    ]
+
+
+@pytest.mark.parametrize("strength", [0.0, 4.0])
+def test_generate_accepts_the_lora_strength_boundaries(
+    client, monkeypatch, lora_registry, strength
+):
+    # 0 is the bit-exact bypass (ADR-0028), 4 the guard-rail ceiling.
+    calls = []
+
+    async def fake_generate(prompt, seconds, kind, **options):
+        calls.append(options)
+        return b"RIFFwav"
+
+    monkeypatch.setattr(controller.sa3, "generate", fake_generate)
+    response = client.post(
+        "/api/generate",
+        json=generate_request(loras=[{"name": "small/crackle", "strength": strength}]),
+    )
+    assert response.status_code == 200
+    assert calls[0]["lora_strengths"] == [strength]
+
+
+def test_generate_treats_an_empty_lora_stack_as_no_adapters(
+    client, monkeypatch, lora_registry
+):
+    calls = []
+
+    async def fake_generate(prompt, seconds, kind, **options):
+        calls.append(options)
+        return b"RIFFwav"
+
+    monkeypatch.setattr(controller.sa3, "generate", fake_generate)
+    response = client.post("/api/generate", json=generate_request(loras=[]))
+    assert response.status_code == 200
+    assert calls == [{}]
+
+
+@pytest.mark.parametrize(
+    "loras",
+    [
+        {"name": "small/crackle"},  # not an array
+        ["small/crackle"],  # entry not an object
+        [{"name": None}],
+        [{"name": "small/absent"}],  # not installed
+        [{"name": "small/../crackle"}],  # traversal
+        [{"name": "medium/maqam"}],  # medium adapter cannot ride an sfx (small) DiT
+        [{"name": "small/crackle", "strength": -0.1}],
+        [{"name": "small/crackle", "strength": 4.1}],
+        [{"name": "small/crackle", "strength": True}],
+        [{"name": "small/crackle", "strength": "1"}],
+        # A repeated adapter would double its delta silently.
+        [{"name": "small/crackle"}, {"name": "small/crackle"}],
+        # One over MAX_LORA_STACK, and every name resolves — only the cap trips.
+        [{"name": f"small/stack{n}"} for n in range(5)],
+    ],
+)
+def test_generate_rejects_invalid_lora_stacks(
+    client, monkeypatch, lora_registry, loras
+):
+    async def fake_generate(prompt, seconds, kind, **options):  # pragma: no cover
+        raise AssertionError("invalid input must not reach generation")
+
+    monkeypatch.setattr(controller.sa3, "generate", fake_generate)
+    response = client.post("/api/generate", json=generate_request(loras=loras))
+    assert response.status_code == 422
+
+
+def test_generate_rejects_a_nan_lora_strength(client, monkeypatch, lora_registry):
+    # httpx's json= encoder refuses NaN, but Python's json.loads parses it —
+    # so it can reach the server, and the boundary must catch it.
+    async def fake_generate(prompt, seconds, kind, **options):  # pragma: no cover
+        raise AssertionError("invalid input must not reach generation")
+
+    monkeypatch.setattr(controller.sa3, "generate", fake_generate)
+    response = client.post(
+        "/api/generate",
+        content=(
+            '{"prompt": "x", "seconds": 3, "kind": "sfx", '
+            '"loras": [{"name": "small/crackle", "strength": NaN}]}'
+        ),
+        headers={"content-type": "application/json"},
+    )
+    assert response.status_code == 422
+
+
 @pytest.mark.parametrize("channels", [1, 2])
 def test_generate_forwards_multipart_init_audio_and_inpaint(
     client, monkeypatch, channels