diff --git a/backend/lsdj/controller.py b/backend/lsdj/controller.py index a93c634..0af344c 100644 --- a/backend/lsdj/controller.py +++ b/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: diff --git a/backend/lsdj/loras.py b/backend/lsdj/loras.py new file mode 100644 index 0000000..e634456 --- /dev/null +++ b/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`` 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 (``/``) 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 diff --git a/backend/lsdj/sa3.py b/backend/lsdj/sa3.py index 32cd64f..32e35b2 100644 --- a/backend/lsdj/sa3.py +++ b/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, diff --git a/backend/tests/test_controller.py b/backend/tests/test_controller.py index 91ac380..484d54f 100644 --- a/backend/tests/test_controller.py +++ b/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 diff --git a/backend/tests/test_loras.py b/backend/tests/test_loras.py new file mode 100644 index 0000000..47a8411 --- /dev/null +++ b/backend/tests/test_loras.py @@ -0,0 +1,89 @@ +"""LoRA registry tests: name resolution at the trust boundary. + +The registry is a directory layout the Rust importer writes; here it is laid +out by hand and `resolve` — the only read the generate path performs — is +exercised against well-formed, malformed, and hostile names. +""" + +import pytest + +from lsdj import loras + +ADAPTER = b"fake safetensors bytes" + + +def install_adapter(root, base, slug, filename="adapter_model.safetensors"): + adapter_dir = root / base / slug + adapter_dir.mkdir(parents=True) + (adapter_dir / filename).write_bytes(ADAPTER) + return adapter_dir + + +class TestLorasDir: + def test_env_override_wins(self, tmp_path): + assert ( + loras.loras_dir( + env={"SA3_LORAS_HOME": str(tmp_path / "elsewhere")}, home=tmp_path + ) + == tmp_path / "elsewhere" + ) + + def test_defaults_to_the_app_support_home(self, tmp_path): + assert ( + loras.loras_dir(env={}, home=tmp_path) + == tmp_path / "Library" / "Application Support" / "LSDJai" / "sa3-loras" + ) + + +class TestResolve: + def test_resolves_an_installed_adapter(self, tmp_path): + adapter_dir = install_adapter(tmp_path, "medium", "maqam") + resolved, base = loras.resolve( + "medium/maqam", env={"SA3_LORAS_HOME": str(tmp_path)}, home=tmp_path + ) + assert resolved == adapter_dir + assert base == "medium" + + def test_resolves_a_hand_placed_safetensors_name(self, tmp_path): + # The runtime accepts any single .safetensors in the dir; so does the + # registry (a user may drop an adapter in by hand). + install_adapter(tmp_path, "small", "crackle", filename="crackle.safetensors") + _, base = loras.resolve( + "small/crackle", env={"SA3_LORAS_HOME": str(tmp_path)}, home=tmp_path + ) + assert base == "small" + + @pytest.mark.parametrize( + "name", + [ + "medium/absent", # never installed + "maqam", # no base segment + "large/maqam", # not a known base + "medium/../maqam", # traversal + "medium/.hidden", # leading dot + "medium/", # empty slug + "medium/sub/dir", # extra separator lands in the slug check + "MEDIUM/maqam", # bases are exact + ], + ) + def test_rejects_names_that_do_not_resolve(self, tmp_path, name): + install_adapter(tmp_path, "medium", "maqam") + with pytest.raises(loras.UnknownAdapter): + loras.resolve(name, env={"SA3_LORAS_HOME": str(tmp_path)}, home=tmp_path) + + def test_rejects_a_directory_without_a_safetensors(self, tmp_path): + (tmp_path / "medium" / "empty").mkdir(parents=True) + with pytest.raises(loras.UnknownAdapter): + loras.resolve( + "medium/empty", env={"SA3_LORAS_HOME": str(tmp_path)}, home=tmp_path + ) + + def test_rejects_a_directory_with_two_safetensors(self, tmp_path): + # Ambiguous contents mirror the runtime's own refusal — better to + # refuse at the boundary than let sa3_mlx fail mid-generation. + adapter_dir = install_adapter(tmp_path, "medium", "both") + (adapter_dir / "second.safetensors").write_bytes(ADAPTER) + with pytest.raises(loras.UnknownAdapter): + loras.resolve( + "medium/both", env={"SA3_LORAS_HOME": str(tmp_path)}, home=tmp_path + ) diff --git a/backend/tests/test_sa3.py b/backend/tests/test_sa3.py index 7bdb634..1afb098 100644 --- a/backend/tests/test_sa3.py +++ b/backend/tests/test_sa3.py @@ -156,6 +156,45 @@ def test_passes_the_full_generation_surface_and_init_bytes(self, checkout): ] assert (mlx_dir / ".venv" / "bin" / "init.wav").read_bytes() == init_audio + def test_passes_one_lora_group_per_adapter_with_its_strength(self, checkout): + # Issue #66 (ADR-0028): each adapter rides the argv as its own + # --lora group — the directory plus a strength=S option (the + # upstream PR #57/#65 CLI syntax). + mlx_dir = checkout(SUCCESS_STUB) + asyncio.run( + sa3.generate( + "maqam phrasing", + 120.0, + "track", + lora_dirs=["/adapters/medium/maqam", "/adapters/medium/breaks"], + lora_strengths=[0.75, 1.5], + ) + ) + argv = (mlx_dir / ".venv" / "bin" / "argv.txt").read_text().splitlines() + first = argv.index("--lora") + assert argv[first : first + 6] == [ + "--lora", + "/adapters/medium/maqam", + "strength=0.75", + "--lora", + "/adapters/medium/breaks", + "strength=1.5", + ] + + def test_lora_without_strengths_omits_the_option(self, checkout): + # No strengths → bare --lora groups; the CLI's default (1.0) applies. + mlx_dir = checkout(SUCCESS_STUB) + asyncio.run( + sa3.generate( + "vinyl spinback", 3.0, "sfx", lora_dirs=["/adapters/small/crackle"] + ) + ) + argv = (mlx_dir / ".venv" / "bin" / "argv.txt").read_text().splitlines() + lora_index = argv.index("--lora") + assert argv[lora_index + 1] == "/adapters/small/crackle" + assert not any(arg.startswith("strength=") for arg in argv) + assert "--lora-strength" not in argv + def test_tracks_run_the_medium_dit_with_its_decoder(self, checkout): # M19 (ADR-0013): tracks pair the medium DiT with SAME-L; the # pad kinds keep the small DiTs with SAME-S. diff --git a/docs/issue-66-checklist.md b/docs/issue-66-checklist.md new file mode 100644 index 0000000..dc30d7b --- /dev/null +++ b/docs/issue-66-checklist.md @@ -0,0 +1,103 @@ +# Issue #66 — SA3 LoRA adapter manager: hardware/UX checklist + +Issue #66 builds the production importer for Stable Audio 3 LoRA finetunes on +top of the spike's merge-at-load runtime (ADR-0028, `docs/spike-sa3-lora.md`): +a `sa3-loras///` registry in the app data dir owned by the Rust +shell, a `lora` field on `/api/generate` that rides `--lora`/`--lora-strength` +into the pinned CLI, adapter pickers in the generate surfaces, and a +manager section with import (HuggingFace repo id or local `.safetensors`) and +in-app delete. + +Unit tests cover name/path trust boundaries, the safetensors-header +validation (pickle refusal, convention detection, base inference), the exact +argv, the `/api/generate` contract, and the picker/manager UI. What follows +needs a real machine with the SA3 checkout warmed — the sandbox cannot run +the shell or MLX. The public PEFT adapter the spike used +(`motiftechnologies/stable-audio-3-maqam-lora`, medium base) is the reference +adapter throughout. + +## Import + +- [ ] Settings drawer → Model library shows the **LoRA adapters** section + with "No adapters installed" and an **Open folder** that reveals + `~/Library/Application Support/LSDJai/sa3-loras`. +- [ ] Enter `motiftechnologies/stable-audio-3-maqam-lora` and Install: fetch / + download / install progress shows, then the adapter lists as + `stable-audio-3-maqam-lora` — **Medium DiT (tracks)**, ~200 MB. +- [ ] Cancel works mid-download and surfaces as a clean stop, not an error. +- [ ] Import the same adapter's `adapter_model.safetensors` via **Import + file…** (download it separately first): refused as already installed + when the slug collides; imports cleanly under a different folder name. +- [ ] A pickle file (`.ckpt`/`.pt` — rename any small file) is refused by the + file picker's filter, and forcing a path at it (e.g. via the HF id of a + pickle-only repo) yields the explicit pickle refusal, not a generic + error. +- [ ] A non-LoRA `.safetensors` (e.g. a Magenta `*_state.safetensors`) is + refused with "not a recognised SA3 LoRA". + +## Generate + +- [ ] Media Explorer → Generate, engine **Track (SA3 medium)**: the LoRA + picker offers the Maqam adapter; the pad engines do NOT offer it (wrong + base), and the deck pad panels don't either. +- [ ] Compose two tracks from the same prompt + fixed conditions, adapter + None vs Maqam at ×1: audibly different in character (the spike measured + a difference as large as the signal itself). +- [ ] Strength ×0.25 vs ×1.5 audibly scales the adapter's influence. +- [ ] Backend log (the generation server's stderr) shows the CLI's + `lora: merged 168 layer(s) from 1 adapter(s)` line during a LoRA take. +- [ ] Magenta engine ignores the adapter path entirely (no `lora` in its + render request). + +## Bypass (ADR-0028's bit-exact claim) + +- [ ] Two tracks with the same prompt and `seed`, adapter **None** vs Maqam at + **×0**: byte-identical WAVs (compare SHA-256). Seed rides via the + `/api/generate` `seed` field (issue #54) — use + `scripts/verify_sa3_surface.py`-style direct calls if the UI has no seed + control. + +## Registry lifecycle + +- [ ] Quit and relaunch: the adapter is still listed (the registry is the + directory layout — nothing else to persist). +- [ ] Delete from the manager: the row disappears, the folder is gone, and an + in-flight picker choice falls back to **None** on the next generate. +- [ ] Drop a valid adapter folder in by hand (`sa3-loras/medium//` with + its `.safetensors`): the watcher lists it live, and it generates. + +## Wrong-base refusal + +- [ ] POST `/api/generate` directly with `kind: "sfx"` and the medium + adapter's name: 422 naming the base mismatch (the UI never offers the + combination; the boundary still refuses it). + +## LoRA stack (multi-adapter follow-up) + +The generate surfaces replaced the adapter/strength pickers with the +**LoRA rack**: every base-matched adapter is a toggle chip; a chip clicked +into the stack grows a trim knob (double-click parks it at ×1, ×0 dims the +slot — bit-exact bypass). `/api/generate` now takes `loras` (a list of up to +4 `{name, strength}` entries) instead of the single `lora` object. + +**Prerequisite (before anything below):** the pin is back on upstream +`Stability-AI/stable-audio-3` (our LoRA support landed as PR #57; PR #65 +added the per-adapter `--lora PATH strength=S` syntax the app now emits). +Settings drawer → Model library → the SA3 row shows **Update available** — +run the update so the installed checkout matches the pin. On the old fork +checkout, a multi-adapter generation fails in argparse +(`unrecognized arguments`); a hand-patched checkout is replaced cleanly. + +- [ ] With two medium adapters installed: chip both into the Generate tab's + rack, distinct trims (e.g. ×1 and ×0.5) — the backend log shows + `merged … from 2 adapter(s)` and the take audibly carries both. +- [ ] Trim one slot to ×0: its slot dims; same prompt + seed with that + adapter chipped out entirely is byte-identical (per-slot bypass). +- [ ] Chip a 5th adapter with 4 stacked: the chip is disabled with the + "Stack full" hint (and a direct POST with 5 entries returns 422, as + does a duplicated name). +- [ ] Deck rack takes the deck accent (A lime / B violet by default); + the Media Explorer racks take the master accent. +- [ ] Toggle a chip out and back in: it remembers its trim for the session. +- [ ] Delete a stacked adapter from the LoRA Library: the chip vanishes + from the racks and the next generate simply rides without it. diff --git a/frontend/src/App.tsx b/frontend/src/App.tsx index 326c1de..85b0b53 100644 --- a/frontend/src/App.tsx +++ b/frontend/src/App.tsx @@ -56,6 +56,7 @@ import { import { Logo } from './ui/Logo' import { Drawer } from './ui/Drawer' import { Button } from './ui/Button' +import { LoraLibrary } from './models/LoraLibrary' import { ModelManager } from './models/ModelManager' import type { StylePreset } from './presets' import { combinedRamWarning } from './ramWarning' @@ -882,6 +883,13 @@ function App() {

{t('settings.modelLibrary')}

+ {/* The LoRA library (issue #66): import / manage Stable Audio 3 LoRA + finetunes — its own top-level section so custom sound worlds don't + hide inside the weights installer. */} +
+

{t('settings.loraLibrary')}

+ +
{/* The native MCP server (ADR-0020 Phase 2): last in the list so the copy-paste connection snippets don't push the everyday controls down. */}
diff --git a/frontend/src/audio/nativeEngine.ts b/frontend/src/audio/nativeEngine.ts index 9f1dfd0..d554bff 100644 --- a/frontend/src/audio/nativeEngine.ts +++ b/frontend/src/audio/nativeEngine.ts @@ -139,7 +139,7 @@ export function subscribeLibraryChanged( // --- Model manager (issue #43) --------------------------------------------- /** A model family the manager installs (Rust `models::Family`). */ -export type ModelFamily = 'magenta' | 'sa3' +export type ModelFamily = 'magenta' | 'sa3' | 'lora' /** One installed Magenta model in `model_status` (serde camelCase). */ export type InstalledModel = { @@ -156,6 +156,23 @@ export type Sa3State = 'missing' | 'venv_missing' | 'not_warmed' | 'ready' * `models::Sa3Source`): the `sa3-pin.json` repo + commit. */ export type Sa3Source = { repo: string; commit: string } +/** The DiT family an SA3 LoRA adapter rides (issue #66): the 1024-wide small + * DiTs (the sfx/music pad kinds) or the 1536-wide medium track DiT. */ +export type LoraBase = 'small' | 'medium' + +/** One installed SA3 LoRA adapter (Rust `loras::LoraInfo`). `name` is the + * `/` id a generate request sends to the backend. */ +export type LoraAdapter = { + name: string + base: LoraBase + slug: string + sizeBytes: number + /** Import-manifest facts; null for a hand-placed adapter. */ + source: string | null + adapterType: string | null + rank: number | null +} + /** The model-manager status for both families (Rust `models::ModelStatus`). */ export type ModelStatus = { magenta: { @@ -175,6 +192,8 @@ export type ModelStatus = { /** The installed checkout differs from the pin (or is unstamped) — offer an update. */ updateAvailable: boolean } + /** The installed SA3 LoRA adapters (issue #66), sorted by name. */ + loras: LoraAdapter[] /** The in-flight install, so the manager reflects it after a close/reopen even * though the live `model://progress` events were missed while unmounted. */ installing: { family: ModelFamily; name: string } | null @@ -240,6 +259,24 @@ export function openModelFolder(family: ModelFamily): Promise { return invoke('open_model_folder', { family }) } +/** Import an SA3 LoRA adapter (issue #66) from a HuggingFace repo id or a local + * path (`.safetensors` file or PEFT adapter folder). `base` is only needed for + * adapters whose base cannot be inferred (rank-only `-xs` shapes). Resolves once + * the import has STARTED; progress arrives via `subscribeModelProgress` (family + * `lora`) and a final `models://changed`. */ +export function installLora( + source: { hfRepo: string } | { path: string }, + base?: LoraBase, +): Promise { + return invoke('install_lora', { spec: { ...source, base: base ?? null } }) +} + +/** Delete an installed adapter (small, re-downloadable — unlike the model + * families, adapters get an in-app delete). Emits `models://changed`. */ +export function deleteLora(name: string): Promise { + return invoke('delete_lora', { name }) +} + /** Subscribe to `models://changed` (the models-dir watcher / an install finishing): * re-fetch `model_status` and the deck picker's `/api/models`. */ export function subscribeModelsChanged(onChange: () => void): () => void { diff --git a/frontend/src/deck/DeckColumn.test.tsx b/frontend/src/deck/DeckColumn.test.tsx index 6e39871..3742868 100644 --- a/frontend/src/deck/DeckColumn.test.tsx +++ b/frontend/src/deck/DeckColumn.test.tsx @@ -20,8 +20,17 @@ import { type InterfaceState, } from '../audio/nativeEngine' import * as interfaceStore from '../audio/interfaceStore' +import type { LoraAdapter } from '../audio/nativeEngine' import { DeckColumn } from './DeckColumn' +// The installed-adapter list (issue #66) is stubbed per test; the default is +// none, which hides the pickers entirely. +const useLorasMock = vi.fn<() => LoraAdapter[]>(() => []) +vi.mock('../models/useLoras', async (importOriginal) => { + const original = await importOriginal() + return { ...original, useLoras: () => useLorasMock() } +}) + // The style intents are mocked flat (no rAF coalescing — that machinery has // its own tests in nativeEngine.test.ts) and wired per test to a fake pad // store below, so the projection tests exercise the full loop a real session @@ -1366,7 +1375,7 @@ describe('DeckColumn', () => { target: { value: 'vinyl spinback' }, }) fireEvent.click(screen.getByRole('button', { name: 'Generate' })) - expect(onGenerateToPad).toHaveBeenCalledWith('vinyl spinback', 'sfx', true) + expect(onGenerateToPad).toHaveBeenCalledWith('vinyl spinback', 'sfx', true, []) fireEvent.change(screen.getByLabelText('Engine'), { target: { value: 'music' }, @@ -1381,9 +1390,60 @@ describe('DeckColumn', () => { 'vinyl spinback', 'music', false, + [], ) }) + it('rides the stacked LoRA adapters and their trims into a pad generation (issue #66)', () => { + useLorasMock.mockReturnValue([ + { + name: 'small/maqam', + base: 'small', + slug: 'maqam', + sizeBytes: 200_000_000, + source: null, + adapterType: 'lora', + rank: 64, + }, + { + name: 'small/crackle', + base: 'small', + slug: 'crackle', + sizeBytes: 50_000_000, + source: null, + adapterType: 'lora', + rank: 8, + }, + // A medium adapter never reaches the pad rack (pads ride the small DiTs). + { + name: 'medium/tape', + base: 'medium', + slug: 'tape', + sizeBytes: 100_000_000, + source: null, + adapterType: 'lora', + rank: 16, + }, + ]) + const onGenerateToPad = vi.fn() + generateRow({ onGenerateToPad: onGenerateToPad as () => void }) + + expect(screen.queryByText('tape')).toBeNull() + fireEvent.click(screen.getByText('maqam')) + fireEvent.click(screen.getByText('crackle')) + fireEvent.change(screen.getByLabelText('maqam strength'), { + target: { value: '0.75' }, + }) + fireEvent.change(screen.getByLabelText('Generate prompt'), { + target: { value: 'oud phrase' }, + }) + fireEvent.click(screen.getByRole('button', { name: 'Generate' })) + expect(onGenerateToPad).toHaveBeenCalledWith('oud phrase', 'sfx', true, [ + { name: 'small/maqam', strength: 0.75 }, + { name: 'small/crackle', strength: 1 }, + ]) + }) + it('offers Magenta while the deck plays — the third engine is its own worker', () => { const onGenerateToPad = vi.fn() renderPanel( @@ -1397,7 +1457,7 @@ describe('DeckColumn', () => { target: { value: 'magenta' }, }) fireEvent.click(screen.getByRole('button', { name: 'Generate' })) - expect(onGenerateToPad).toHaveBeenCalledWith('dub chords', 'magenta', true) + expect(onGenerateToPad).toHaveBeenCalledWith('dub chords', 'magenta', true, []) }) it('caps the prompt input short of the backend limit, sparing the BPM stamp', () => { diff --git a/frontend/src/deck/DeckColumn.tsx b/frontend/src/deck/DeckColumn.tsx index 5e4de10..bb010cd 100644 --- a/frontend/src/deck/DeckColumn.tsx +++ b/frontend/src/deck/DeckColumn.tsx @@ -45,6 +45,15 @@ import { type SyncResult, type TrackState, } from './useDeck' +import { + adaptersForKind, + MAX_LORA_STACK, + stackForKind, + useLoras, + useLoraStack, + type LoraChoice, +} from '../models/useLoras' +import { LoraRack } from '../ui/LoraRack' import './deck.css' // The worker holds ~3s of lead (see backend worker pacing); the meter shows @@ -123,8 +132,13 @@ type DeckColumnProps = { onClearLoopPad: (slot: number) => void onSetLoopSeconds: (seconds: number) => void /** Generated pads (M18): fill the first empty slot from a prompt, - * with the chosen engine and one-shot/loop behaviour. */ - onGenerateToPad: (prompt: string, engine: GenerateEngine, oneShot: boolean) => void + * with the chosen engine, one-shot/loop behaviour, and LoRA stack. */ + onGenerateToPad: ( + prompt: string, + engine: GenerateEngine, + oneShot: boolean, + loras?: LoraChoice[] | null, + ) => void generateError: string | null /** Gated tempo readout (M14): null shows an honest dash. */ bpm: number | null @@ -268,6 +282,10 @@ export function DeckColumn({ } | null>(null) const [generateEngine, setGenerateEngine] = useState('sfx') const [generateOneShot, setGenerateOneShot] = useState(true) + // Per-deck LoRA stack (issue #66): both pad kinds ride the small DiTs, so + // one rack covers them; Magenta has no adapter path and hides it. + const loras = useLoras() + const padStack = useLoraStack() const [targetDraft, setTargetDraft] = useState('') // In-place prompt editing: which row is open and its draft text. const [editing, setEditing] = useState<{ text: string; draft: string } | null>( @@ -296,9 +314,14 @@ export function DeckColumn({ connected && Boolean(generateDraft.trim()) && loop.slots.some((slot) => slot.state === 'empty') + // Pads only ever ride the small DiTs; a stale slot (deleted adapter) + // drops from the request, never blocks it. + const padAdapters = adaptersForKind(loras, 'sfx') const fireGenerate = () => { if (!canGenerate) return - onGenerateToPad(generateDraft, generateEngine, generateOneShot) + const stacked = + generateEngine === 'magenta' ? [] : stackForKind(padStack.stack, loras, 'sfx') + onGenerateToPad(generateDraft, generateEngine, generateOneShot, stacked) } const statusKey = mode === 'playback' @@ -916,6 +939,19 @@ export function DeckColumn({ onChange={(value) => setGenerateOneShot(value === 'oneshot')} /> + {generateEngine !== 'magenta' && ( + ({ + name: adapter.name, + label: adapter.slug, + }))} + value={padStack.stack} + onToggle={padStack.toggle} + onStrength={padStack.setStrength} + max={MAX_LORA_STACK} + /> + )}
{ expect(result.current.loop.active).toBeNull() }) + it('rides a LoRA stack with per-adapter strengths on an SA3 pad generation (issue #66)', async () => { + const fetchMock = stubFetchOk() + const { engine } = makeFakeEngine() + const { result } = await playingDeck(engine) + + act(() => + result.current.generateToPad('oud phrase', 'sfx', true, [ + { name: 'small/maqam', strength: 0.75 }, + { name: 'small/crackle', strength: 1 }, + ]), + ) + await act(async () => {}) + expect(requestBody(fetchMock)).toEqual({ + prompt: 'oud phrase', + seconds: 4, + kind: 'sfx', + loras: [ + { name: 'small/maqam', strength: 0.75 }, + { name: 'small/crackle', strength: 1 }, + ], + }) + }) + + it('omits the loras field when the stack is empty', async () => { + const fetchMock = stubFetchOk() + const { engine } = makeFakeEngine() + const { result } = await playingDeck(engine) + + act(() => result.current.generateToPad('air horn', 'sfx', true, [])) + await act(async () => {}) + expect(requestBody(fetchMock)).toEqual({ + prompt: 'air horn', + seconds: 4, + kind: 'sfx', + }) + }) + it('auto-saves a generated pad to the samples library (raw WAV + metadata)', async () => { stubFetchOk() const { engine, channel } = makeFakeEngine() diff --git a/frontend/src/deck/useDeck.ts b/frontend/src/deck/useDeck.ts index 169032a..d6a4f97 100644 --- a/frontend/src/deck/useDeck.ts +++ b/frontend/src/deck/useDeck.ts @@ -51,6 +51,7 @@ import { type TrackLoop, } from '../audio/track' import { loadDeckSettings, updateDeckSettings } from '../persistence' +import type { LoraChoice } from '../models/useLoras' import { SAMPLE_RATE, type Beatgrid, @@ -205,8 +206,15 @@ export type DeckControls = { * dedicated render worker; first use pays its model load inside the * pending state). One-shots overlay, loops replace like captures; * music-model loops snap to whole bars while the tempo gate is - * locked and respect the quality floor. */ - generateToPad: (prompt: string, engine: GenerateEngine, oneShot: boolean) => void + * locked and respect the quality floor. An SA3 engine can ride a LoRA + * stack — adapters + per-adapter strengths (issue #66); Magenta has no + * adapter path. */ + generateToPad: ( + prompt: string, + engine: GenerateEngine, + oneShot: boolean, + loras?: LoraChoice[] | null, + ) => void /** Load a saved sample (ADR-0022) into the first empty loop slot, as a loop or * one-shot per the sample. Resolves false when every slot is full, the deck isn't * a live Realtime deck, or the body doesn't decode (surfaced via `generateError`). */ @@ -1418,7 +1426,12 @@ export function useDeck(deckId: DeckId): DeckControls { ) const generateToPad = useCallback( - (prompt: string, engine: GenerateEngine, oneShot: boolean) => { + ( + prompt: string, + engine: GenerateEngine, + oneShot: boolean, + loras?: LoraChoice[] | null, + ) => { const trimmed = prompt.trim() if (!trimmed) return const current = loopRef.current @@ -1469,6 +1482,7 @@ export function useDeck(deckId: DeckId): DeckControls { prompt: requestPrompt, seconds: requestSeconds, kind: engine, + ...(loras && loras.length > 0 ? { loras } : {}), }, ), }, diff --git a/frontend/src/i18n/en.json b/frontend/src/i18n/en.json index 88119fc..b6da6d5 100644 --- a/frontend/src/i18n/en.json +++ b/frontend/src/i18n/en.json @@ -305,6 +305,7 @@ "models": "Deck realtime models", "modelDeck": "Deck {{id}}", "modelLibrary": "Model library", + "loraLibrary": "LoRA Library", "recording": "Recording", "recordingFolder": "Recordings folder", "recordingFolderDefault": "Downloads (default)", @@ -369,7 +370,20 @@ "install": "Building and warming…", "done": "Done", "error": "Failed" - } + }, + "loras": "LoRA adapters", + "loraNone": "No adapters installed", + "loraBaseLabel": "Base", + "loraBaseAuto": "Detect base", + "loraBase": { + "small": "Small DiT (pads)", + "medium": "Medium DiT (tracks)" + }, + "loraRepo": "HuggingFace repo", + "loraRepoPlaceholder": "owner/adapter-repo or huggingface.co link", + "loraImportFile": "Import file…", + "loraFileFilter": "LoRA adapter", + "loraDelete": "Delete adapter {{name}}" }, "piano": { "heading": "MIDI keyboard", @@ -380,5 +394,13 @@ "octaveUp": "Octave up", "octaveDown": "Octave down", "key": "{{note}}" + }, + "lora": { + "rack": "LoRA", + "include": "Add {{name}} to the LoRA stack", + "exclude": "Remove {{name}} from the LoRA stack", + "strength": "{{name}} strength", + "strengthValue": "×{{value}}", + "stackFull": "Stack full — up to {{max}} adapters per generation" } } diff --git a/frontend/src/media/MediaExplorer.test.tsx b/frontend/src/media/MediaExplorer.test.tsx index 1dd9200..65ef6ef 100644 --- a/frontend/src/media/MediaExplorer.test.tsx +++ b/frontend/src/media/MediaExplorer.test.tsx @@ -18,6 +18,14 @@ vi.mock('../audio/nativeEngine', async (importOriginal) => { }) vi.mock('../audio/interfaceStore', () => ({ useInterfaceStore: vi.fn(() => null) })) +// The installed-adapter list (issue #66) is stubbed per test; the default is +// none, which hides the LoRA pickers entirely. +const useLorasMock = vi.fn<() => import('../audio/nativeEngine').LoraAdapter[]>(() => []) +vi.mock('../models/useLoras', async (importOriginal) => { + const original = await importOriginal() + return { ...original, useLoras: () => useLorasMock() } +}) + type Handlers = { onLoadPreset?: (deck: DeckId, preset: StylePreset) => void onLoadTrack?: (deck: DeckId, source: TrackSource, title: string) => Promise @@ -99,6 +107,7 @@ afterEach(() => { vi.unstubAllGlobals() vi.mocked(useInterfaceStore).mockReturnValue(null) vi.mocked(togglePianoWindow).mockClear() + useLorasMock.mockReturnValue([]) }) describe('MediaExplorer', () => { @@ -206,6 +215,90 @@ describe('MediaExplorer', () => { ).toBe(true) }) + it('rides a stacked pair of LoRA adapters with their trims on a track compose (issue #66)', async () => { + useLorasMock.mockReturnValue([ + { + name: 'medium/maqam', + base: 'medium', + slug: 'maqam', + sizeBytes: 200_000_000, + source: null, + adapterType: 'lora', + rank: 64, + }, + { + name: 'medium/breaks', + base: 'medium', + slug: 'breaks', + sizeBytes: 150_000_000, + source: null, + adapterType: 'lora', + rank: 32, + }, + // Small adapters never reach the track rack (tracks ride medium). + { + name: 'small/crackle', + base: 'small', + slug: 'crackle', + sizeBytes: 50_000_000, + source: null, + adapterType: 'lora', + rank: 8, + }, + ]) + const fetchMock = stubFetch() + renderExplorer() + fireEvent.click(screen.getByRole('tab', { name: 'Generate' })) + expect(screen.queryByText('crackle')).toBeNull() + // Chip both medium adapters into the stack; trim only maqam. + fireEvent.click(screen.getByText('maqam')) + fireEvent.click(screen.getByText('breaks')) + fireEvent.change(screen.getByLabelText('maqam strength'), { + target: { value: '1.5' }, + }) + await composeTrack('maqam study') + expect(fetchMock).toHaveBeenCalledWith( + '/api/generate', + expect.objectContaining({ + body: JSON.stringify({ + prompt: 'maqam study', + seconds: 120, + kind: 'track', + loras: [ + { name: 'medium/maqam', strength: 1.5 }, + { name: 'medium/breaks', strength: 1 }, + ], + }), + }), + ) + }) + + it('drops a chip toggled back out of the stack from the request', async () => { + useLorasMock.mockReturnValue([ + { + name: 'medium/maqam', + base: 'medium', + slug: 'maqam', + sizeBytes: 200_000_000, + source: null, + adapterType: 'lora', + rank: 64, + }, + ]) + const fetchMock = stubFetch() + renderExplorer() + fireEvent.click(screen.getByRole('tab', { name: 'Generate' })) + fireEvent.click(screen.getByText('maqam')) // in… + fireEvent.click(screen.getByText('maqam')) // …and back out + await composeTrack('clean take') + expect(fetchMock).toHaveBeenCalledWith( + '/api/generate', + expect.objectContaining({ + body: JSON.stringify({ prompt: 'clean take', seconds: 120, kind: 'track' }), + }), + ) + }) + it('previews a take in the headphones and toggles it off', async () => { stubFetch() const onPreview = vi.fn(async () => {}) diff --git a/frontend/src/media/MediaExplorer.tsx b/frontend/src/media/MediaExplorer.tsx index 7a2dda4..b2da1f1 100644 --- a/frontend/src/media/MediaExplorer.tsx +++ b/frontend/src/media/MediaExplorer.tsx @@ -18,6 +18,14 @@ import { import { useInterfaceStore } from '../audio/interfaceStore' import { useControlBus } from '../control/busContext' import { CrateBrowser } from '../crates/CrateBrowser' +import { + adaptersForKind, + MAX_LORA_STACK, + stackForKind, + useLoras, + useLoraStack, +} from '../models/useLoras' +import { LoraRack } from '../ui/LoraRack' import type { StylePreset } from '../presets' import { Button } from '../ui/Button' import { Panel } from '../ui/Panel' @@ -309,6 +317,13 @@ export function MediaExplorer({ const [engine, setEngine] = useState('track') const [seconds, setSeconds] = useState(120) const [generateError, setGenerateError] = useState(null) + // The installed SA3 LoRA adapters (issue #66) and each form's stack: the + // adapters riding the next generation, each at a trim strength. The racks + // only offer base-matched adapters; a slot that stops resolving (adapter + // deleted) silently drops from the request. + const loras = useLoras() + const trackStack = useLoraStack() + const sampleStack = useLoraStack() // Auto-save runs after a take is composed; its failure is separate from a // generation failure (the take is already playable from memory). const [saveError, setSaveError] = useState(null) @@ -653,6 +668,7 @@ export function MediaExplorer({ if (!trimmedPrompt) return const id = nextIdRef.current++ const requestEngine = sampleEngine + const requestLoras = stackForKind(sampleStack.stack, loras, requestEngine) const oneShot = sampleOneShot const clipTitle = sampleTitle.trim() || randomSongTitle() // A loop asks for the surplus tail the engine folds away on reload; a one-shot is @@ -681,6 +697,7 @@ export function MediaExplorer({ prompt: trimmedPrompt, seconds: requestSeconds, kind: requestEngine, + ...(requestLoras.length > 0 ? { loras: requestLoras } : {}), }), }) if (!response.ok) { @@ -754,6 +771,10 @@ export function MediaExplorer({ if (!trimmedPrompt) return const id = nextIdRef.current++ const requestEngine = engine + const requestLoras = + requestEngine === 'magenta' + ? [] + : stackForKind(trackStack.stack, loras, requestEngine) // The name (and on-disk filename) come from the Title field, NOT the prompt — a // blank title gets a random song title so a long/JSON prompt never becomes the // name. The row appends a session-unique #id to tell same-title siblings apart. @@ -775,7 +796,12 @@ export function MediaExplorer({ body: JSON.stringify( requestEngine === 'magenta' ? { prompt: trimmedPrompt, seconds } - : { prompt: trimmedPrompt, seconds, kind: requestEngine }, + : { + prompt: trimmedPrompt, + seconds, + kind: requestEngine, + ...(requestLoras.length > 0 ? { loras: requestLoras } : {}), + }, ), }, ) @@ -871,6 +897,10 @@ export function MediaExplorer({ } const lengths = ENGINE_LENGTHS[engine] + // The racks offer only base-matched adapters; Magenta hides the rack + // entirely (no adapter path). + const trackAdapters = engine === 'magenta' ? [] : adaptersForKind(loras, engine) + const sampleAdapters = adaptersForKind(loras, sampleEngine) function loadButtons(onLoad: (deck: DeckId) => void, name: string) { return (['a', 'b'] as const).map((deck) => ( @@ -1082,6 +1112,16 @@ export function MediaExplorer({ {t('media.generate.action')}
+ ({ + name: adapter.name, + label: adapter.slug, + }))} + value={trackStack.stack} + onToggle={trackStack.toggle} + onStrength={trackStack.setStrength} + max={MAX_LORA_STACK} + /> {tracks.length === 0 ? (

{t('media.generate.empty')}

) : ( @@ -1239,6 +1279,16 @@ export function MediaExplorer({ {t('media.generate.action')}
+ ({ + name: adapter.name, + label: adapter.slug, + }))} + value={sampleStack.stack} + onToggle={sampleStack.toggle} + onStrength={sampleStack.setStrength} + max={MAX_LORA_STACK} + /> {samples.length === 0 ? (

{t('media.samples.empty')}

) : ( diff --git a/frontend/src/models/LoraLibrary.test.tsx b/frontend/src/models/LoraLibrary.test.tsx new file mode 100644 index 0000000..d61d467 --- /dev/null +++ b/frontend/src/models/LoraLibrary.test.tsx @@ -0,0 +1,202 @@ +import { act, fireEvent, render, screen, waitFor } from '@testing-library/react' +import { beforeEach, describe, expect, it, vi } from 'vitest' + +import type { ModelProgress, ModelStatus } from '../audio/nativeEngine' + +// Capture the subscriber callbacks so tests can fire watcher / progress events. +let changedCb: (() => void) | null = null +let progressCb: ((event: ModelProgress) => void) | null = null +const modelStatus = vi.fn<() => Promise>() +const installLora = vi.fn<(source: object, base?: string) => Promise>(async () => {}) +const deleteLora = vi.fn<(name: string) => Promise>(async () => {}) +const cancelInstall = vi.fn(async () => {}) +const openModelFolder = vi.fn<(family: string) => Promise>(async () => {}) +const invoke = vi.fn<(cmd: string, args?: object) => Promise>(async () => null) + +vi.mock('../audio/nativeEngine', () => ({ + modelStatus: () => modelStatus(), + installLora: (source: object, base?: string) => installLora(source, base), + deleteLora: (name: string) => deleteLora(name), + cancelInstall: () => cancelInstall(), + openModelFolder: (family: string) => openModelFolder(family), + invoke: (cmd: string, args?: object) => invoke(cmd, args), + subscribeModelsChanged: (cb: () => void) => { + changedCb = cb + return () => {} + }, + subscribeModelProgress: (cb: (event: ModelProgress) => void) => { + progressCb = cb + return () => {} + }, +})) + +import { LoraLibrary } from './LoraLibrary' + +function status(overrides: Partial = {}): ModelStatus { + return { + magenta: { + modelsDir: '/models', + resourcesPresent: true, + installable: ['mrt2_small', 'mrt2_base'], + installed: [], + }, + sa3: { + state: 'ready', + sizeBytes: 5_000_000_000, + checkout: '/sa3', + installedSource: null, + pinnedSource: { repo: 'https://github.com/brxs/stable-audio-3', commit: 'pinned1' }, + updateAvailable: false, + }, + loras: [], + installing: null, + ...overrides, + } +} + +const maqam = { + name: 'medium/maqam', + base: 'medium' as const, + slug: 'maqam', + sizeBytes: 200_000_000, + source: 'motiftechnologies/stable-audio-3-maqam-lora', + adapterType: 'lora', + rank: 64, +} + +beforeEach(() => { + vi.clearAllMocks() + changedCb = null + progressCb = null +}) + +describe('LoraLibrary', () => { + it('lists installed adapters with base and size, and deletes one', async () => { + modelStatus.mockResolvedValue(status({ loras: [maqam] })) + render() + expect(await screen.findByText('maqam')).toBeInTheDocument() + expect(screen.getByText('Medium DiT (tracks) · 200 MB')).toBeInTheDocument() + fireEvent.click(screen.getByLabelText('Delete adapter maqam')) + expect(deleteLora).toHaveBeenCalledWith('medium/maqam') + }) + + it('reveals the registry folder for native inspection', async () => { + modelStatus.mockResolvedValue(status()) + render() + await screen.findByText('No adapters installed') + fireEvent.click(screen.getByText('Open folder')) + expect(openModelFolder).toHaveBeenCalledWith('lora') + }) + + it('imports an adapter from a HuggingFace repo id', async () => { + modelStatus.mockResolvedValue(status()) + render() + await screen.findByText('No adapters installed') + fireEvent.change(screen.getByLabelText('HuggingFace repo'), { + target: { value: 'owner/my-lora' }, + }) + fireEvent.click(screen.getByText('Install')) + expect(installLora).toHaveBeenCalledWith({ hfRepo: 'owner/my-lora' }, undefined) + }) + + it('accepts a pasted huggingface.co URL and installs the canonical repo id', async () => { + modelStatus.mockResolvedValue(status()) + render() + await screen.findByText('No adapters installed') + fireEvent.change(screen.getByLabelText('HuggingFace repo'), { + target: { + value: 'https://huggingface.co/motiftechnologies/stable-audio-3-maqam-lora', + }, + }) + fireEvent.click(screen.getByText('Install')) + expect(installLora).toHaveBeenCalledWith( + { hfRepo: 'motiftechnologies/stable-audio-3-maqam-lora' }, + undefined, + ) + }) + + it('passes an explicit base override to a repo import', async () => { + modelStatus.mockResolvedValue(status()) + render() + await screen.findByText('No adapters installed') + fireEvent.change(screen.getByLabelText('HuggingFace repo'), { + target: { value: 'owner/xs-lora' }, + }) + fireEvent.change(screen.getByLabelText('Base'), { target: { value: 'medium' } }) + fireEvent.click(screen.getByText('Install')) + expect(installLora).toHaveBeenCalledWith({ hfRepo: 'owner/xs-lora' }, 'medium') + }) + + it('imports an adapter file through the native picker', async () => { + modelStatus.mockResolvedValue(status()) + invoke.mockResolvedValue('/downloads/maqam.safetensors') + render() + await screen.findByText('No adapters installed') + fireEvent.click(screen.getByText('Import file…')) + await waitFor(() => + expect(installLora).toHaveBeenCalledWith( + { path: '/downloads/maqam.safetensors' }, + undefined, + ), + ) + expect(invoke).toHaveBeenCalledWith( + 'plugin:dialog|open', + expect.objectContaining({ options: expect.anything() }), + ) + }) + + it('shows live import progress, offers cancel, and re-fetches on change', async () => { + modelStatus.mockResolvedValue(status()) + render() + await screen.findByText('No adapters installed') + act(() => + progressCb?.({ + family: 'lora', + name: 'owner/my-lora', + stage: 'download', + message: null, + file: 'adapter_model.safetensors', + }), + ) + expect(screen.getByText(/adapter_model\.safetensors/)).toBeInTheDocument() + fireEvent.click(screen.getByText('Cancel')) + expect(cancelInstall).toHaveBeenCalledTimes(1) + act(() => + progressCb?.({ family: 'lora', name: '', stage: 'cancelled', message: null, file: null }), + ) + act(() => changedCb?.()) + await waitFor(() => expect(modelStatus).toHaveBeenCalledTimes(2)) + // A cancel is a clean stop, never an error banner. + expect(screen.queryByRole('alert')).toBeNull() + }) + + it("another family's install gates the import buttons without claiming the label", async () => { + modelStatus.mockResolvedValue(status()) + render() + await screen.findByText('No adapters installed') + fireEvent.change(screen.getByLabelText('HuggingFace repo'), { + target: { value: 'owner/my-lora' }, + }) + act(() => + progressCb?.({ family: 'sa3', name: '', stage: 'install', message: null, file: null }), + ) + // One install at a time shell-side: both actions disabled, no lora label. + expect(screen.getByText('Install')).toBeDisabled() + expect(screen.getByText('Import file…')).toBeDisabled() + expect(screen.queryByText('Cancel')).toBeNull() + }) + + it("surfaces a lora import error but not another family's", async () => { + modelStatus.mockResolvedValue(status()) + render() + await screen.findByText('No adapters installed') + act(() => + progressCb?.({ family: 'magenta', name: 'mrt2_base', stage: 'error', message: 'no weights', file: null }), + ) + expect(screen.queryByRole('alert')).toBeNull() + act(() => + progressCb?.({ family: 'lora', name: '', stage: 'error', message: 'not a recognised SA3 LoRA', file: null }), + ) + expect(screen.getByRole('alert')).toHaveTextContent('not a recognised SA3 LoRA') + }) +}) diff --git a/frontend/src/models/LoraLibrary.tsx b/frontend/src/models/LoraLibrary.tsx new file mode 100644 index 0000000..e051945 --- /dev/null +++ b/frontend/src/models/LoraLibrary.tsx @@ -0,0 +1,223 @@ +import { useCallback, useEffect, useState } from 'react' +import { useTranslation } from 'react-i18next' + +import { Button } from '../ui/Button' +import { Select } from '../ui/Select' +import { TextField } from '../ui/TextField' +import { + cancelInstall, + deleteLora, + installLora, + invoke, + modelStatus, + openModelFolder, + subscribeModelProgress, + subscribeModelsChanged, + type LoraBase, + type ModelProgress, + type ModelStatus, +} from '../audio/nativeEngine' +import { formatBytes } from './formatBytes' + +/** Accept what people actually paste as a repo: a bare `owner/name` id or a + * full huggingface.co URL (scheme/host, a /tree/… suffix, query, fragment). + * Best-effort — anything without a recognisable id inside passes through + * unchanged and the shell's validation names the problem. Mirrors the Rust + * `loras::normalize_hf_repo`. */ +function normalizeHfRepo(input: string): string { + const rest = input + .trim() + .replace(/^https?:\/\//, '') + .replace(/^(www\.)?(huggingface\.co|hf\.co)\//, '') + .split(/[?#]/)[0] + const segments = rest.split('/').filter(Boolean) + return segments.length >= 2 ? `${segments[0]}/${segments[1]}` : input.trim() +} + +/** The LoRA library (issue #66): list / import / delete Stable Audio 3 LoRA + * adapters, its own settings section beside the model manager. Same shape as + * the manager — Rust owns the lifecycle, this reads `model_status` and drives + * the import commands, and the in-flight install survives a drawer + * close/reopen via the status snapshot. */ +export function LoraLibrary() { + const { t } = useTranslation() + const [status, setStatus] = useState(null) + // The in-flight install across ALL families — the shell runs one install at + // a time, so a Magenta/SA3 download gates the import buttons here too; only + // lora events surface as this section's label/error. + const [inflight, setInflight] = useState(null) + const [error, setError] = useState(null) + // The import controls: a HuggingFace repo/link draft and the base override + // ('auto' = infer from the adapter's shapes, the normal case). + const [repo, setRepo] = useState('') + const [base, setBase] = useState<'auto' | LoraBase>('auto') + + const refresh = useCallback(() => { + modelStatus().then(setStatus).catch(() => {}) + }, []) + + useEffect(() => { + refresh() + const unsubChanged = subscribeModelsChanged(refresh) + const unsubProgress = subscribeModelProgress((event) => { + // `done` (success) and `cancelled` (user stop) both just clear the + // in-flight UI; the follow-up `models://changed` reflects the new state. + if (event.stage === 'done' || event.stage === 'cancelled') { + setInflight(null) + setError(null) + } else if (event.stage === 'error') { + setInflight(null) + // Another family's failure is the model manager's story to tell. + if (event.family === 'lora') { + setError(event.message ?? t('modelManager.stage.error')) + } + } else { + setInflight(event) + if (event.family === 'lora') setError(null) + } + }) + return () => { + unsubChanged() + unsubProgress() + } + }, [refresh, t]) + + const onInstall = useCallback( + (source: { hfRepo: string } | { path: string }, name: string) => { + setError(null) + // The display name mirrors the Rust spec's display_name (the progress + // event key), so the pending label shows before the first event lands. + setInflight({ family: 'lora', name, stage: 'fetch', message: null, file: null }) + installLora(source, base === 'auto' ? undefined : base).catch((e: unknown) => { + setInflight(null) + setError(String(e)) + }) + }, + [base], + ) + + const onImportFile = useCallback(() => { + void (async () => { + // The native file picker (dialog plugin) — WKWebView has no File System + // Access API. Only .safetensors is offered (ADR-0028's trust boundary; + // the shell refuses anything else anyway). + const path = await invoke('plugin:dialog|open', { + options: { + multiple: false, + filters: [{ name: t('modelManager.loraFileFilter'), extensions: ['safetensors'] }], + }, + }).catch(() => null) + if (!path) return // the user dismissed the picker + const name = path.replace(/\/+$/, '').split('/').pop() || path + onInstall({ path }, name) + })() + }, [onInstall, t]) + + const onDelete = useCallback((name: string) => { + setError(null) + // The follow-up `models://changed` refreshes the list. + deleteLora(name).catch((e: unknown) => setError(String(e))) + }, []) + + if (!status) { + return

{t('modelManager.loading')}

+ } + + const snapshot = status.installing + const isInstalling = inflight !== null || snapshot !== null + // The in-flight import's label — live event (detailed) or status snapshot. + const label = + inflight?.family === 'lora' + ? (() => { + const stage = t(`modelManager.stage.${inflight.stage}`, { + defaultValue: inflight.stage, + }) + // The keyed stage carries the wording; only the file (data) rides along. + return inflight.file ? `${stage} ${inflight.file}` : stage + })() + : snapshot?.family === 'lora' + ? t('modelManager.installing') + : null + const repoDraft = normalizeHfRepo(repo) + + return ( +
+ {error && ( +

+ {t('modelManager.errorPrefix', { message: error })} +

+ )} + +
+
+

{t('modelManager.loras')}

+ +
+ {status.loras.length === 0 && ( +

{t('modelManager.loraNone')}

+ )} + {status.loras.map((adapter) => ( +
+
+
{adapter.slug}
+
+ {t(`modelManager.loraBase.${adapter.base}`)} + {` · ${formatBytes(adapter.sizeBytes)}`} +
+
+
+ +
+
+ ))} +
+
+ setRepo(event.target.value)} + onKeyDown={(event) => { + if (event.key === 'Enter' && repoDraft && !isInstalling) { + onInstall({ hfRepo: repoDraft }, repoDraft) + } + }} + /> +
+