From a3fc31b811d1e11aef3e9d7471a451e51b2825ca Mon Sep 17 00:00:00 2001 From: Jake Hartnell Date: Fri, 17 Jul 2026 19:40:28 +0000 Subject: [PATCH 1/9] Thread SA3 LoRA adapters through the generation server (issue #66) The registry read side: lsdj/loras.py resolves a client-supplied / adapter name to its on-disk directory (Rust owns the write side), /api/generate accepts a validated {name, strength} lora object at the trust boundary, and sa3.generate rides it as --lora/--lora-strength on the pinned fork's CLI (ADR-0028). Co-Authored-By: Claude Fable 5 --- backend/lsdj/controller.py | 39 ++++++++++- backend/lsdj/loras.py | 83 +++++++++++++++++++++++ backend/lsdj/sa3.py | 8 +++ backend/tests/test_controller.py | 112 +++++++++++++++++++++++++++++++ backend/tests/test_loras.py | 89 ++++++++++++++++++++++++ backend/tests/test_sa3.py | 29 ++++++++ 6 files changed, 359 insertions(+), 1 deletion(-) create mode 100644 backend/lsdj/loras.py create mode 100644 backend/tests/test_loras.py diff --git a/backend/lsdj/controller.py b/backend/lsdj/controller.py index a93c634..e1b1c11 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,43 @@ def _validate_generate_request( ) options["seed"] = seed + if "lora" in parsed: + lora = parsed["lora"] + if not isinstance(lora, dict): + raise HTTPException(status_code=422, detail="'lora' must be an object") + name = lora.get("name") + if not isinstance(name, str): + raise HTTPException(status_code=422, detail="'lora.name' must be a string") + 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}'" + ), + ) + options["lora_dir"] = str(adapter_dir) + if "strength" in lora: + strength = lora["strength"] + 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=( + "'lora.strength' must be " + f"{loras.MIN_LORA_STRENGTH:g}-{loras.MAX_LORA_STRENGTH:g}" + ), + ) + options["lora_strength"] = float(strength) + 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..7e850e3 --- /dev/null +++ b/backend/lsdj/loras.py @@ -0,0 +1,83 @@ +"""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 + +# 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..6177032 100644 --- a/backend/lsdj/sa3.py +++ b/backend/lsdj/sa3.py @@ -174,6 +174,8 @@ async def generate( cfg: float | None = None, apg: float | None = None, seed: int | None = None, + lora_dir: str | None = None, + lora_strength: float | None = None, ) -> bytes: """Run one generation and return the WAV bytes. @@ -220,6 +222,12 @@ async def generate( argv.extend(("--apg", f"{apg:g}")) if seed is not None: argv.extend(("--seed", str(seed))) + if lora_dir is not None: + # The adapter directory (ADR-0028): the CLI resolves the + # .safetensors inside it, and the merge happens at DiT load. + argv.extend(("--lora", lora_dir)) + if lora_strength is not None: + argv.extend(("--lora-strength", f"{lora_strength: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..eb66c06 100644 --- a/backend/tests/test_controller.py +++ b/backend/tests/test_controller.py @@ -142,6 +142,118 @@ async def fake_generate(prompt, seconds, kind, **options): ] +@pytest.fixture +def lora_registry(tmp_path, monkeypatch): + """A registry with one adapter per base, pointed at via the env override.""" + for base, slug in (("small", "crackle"), ("medium", "maqam")): + 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_adapter_and_strength( + 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(lora={"name": "small/crackle", "strength": 1.5}), + ) + assert response.status_code == 200 + assert calls == [ + { + "lora_dir": str(lora_registry / "small" / "crackle"), + "lora_strength": 1.5, + } + ] + + +@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(lora={"name": "small/crackle", "strength": strength}), + ) + assert response.status_code == 200 + assert calls[0]["lora_strength"] == strength + + +def test_generate_lora_strength_is_optional(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(lora={"name": "small/crackle"}) + ) + assert response.status_code == 200 + assert calls == [{"lora_dir": str(lora_registry / "small" / "crackle")}] + + +@pytest.mark.parametrize( + "lora", + [ + "small/crackle", # 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"}, + ], +) +def test_generate_rejects_invalid_lora_requests( + client, monkeypatch, lora_registry, lora +): + 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(lora=lora)) + 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", ' + '"lora": {"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..0e715e1 100644 --- a/backend/tests/test_sa3.py +++ b/backend/tests/test_sa3.py @@ -156,6 +156,35 @@ 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_the_lora_adapter_dir_and_strength(self, checkout): + # Issue #66 (ADR-0028): the adapter directory rides the argv as + # --lora; the strength knob only appears when set. + mlx_dir = checkout(SUCCESS_STUB) + asyncio.run( + sa3.generate( + "maqam phrasing", + 120.0, + "track", + lora_dir="/adapters/medium/maqam", + lora_strength=0.75, + ) + ) + argv = (mlx_dir / ".venv" / "bin" / "argv.txt").read_text().splitlines() + lora_index = argv.index("--lora") + assert argv[lora_index + 1] == "/adapters/medium/maqam" + assert argv[argv.index("--lora-strength") + 1] == "0.75" + + def test_lora_without_strength_omits_the_flag(self, checkout): + mlx_dir = checkout(SUCCESS_STUB) + asyncio.run( + sa3.generate( + "vinyl spinback", 3.0, "sfx", lora_dir="/adapters/small/crackle" + ) + ) + argv = (mlx_dir / ".venv" / "bin" / "argv.txt").read_text().splitlines() + assert argv[argv.index("--lora") + 1] == "/adapters/small/crackle" + 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. From 1ae7a46e4f7f0a6553567dbae1e013f2ab74b899 Mon Sep 17 00:00:00 2001 From: Jake Hartnell Date: Fri, 17 Jul 2026 19:49:41 +0000 Subject: [PATCH 2/9] Own the SA3 LoRA adapter lifecycle in the Rust shell (issue #66) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit loras.rs imports adapters from a HuggingFace repo id or a local path on the shared install thread (model://progress family "lora"), validates them at the ADR-0028 trust boundary — pickle refused before any read, safetensors JSON header only, base inferred from the disjoint DiT layer widths with an explicit choice for rank-only -xs shapes — and places them under sa3-loras/// with an import manifest. Adapters ride model_status, get an in-app delete (small, re-downloadable, never iCloud-managed), and a registry watcher mirrors watch_models. Co-Authored-By: Claude Fable 5 --- src-tauri/src/lib.rs | 6 + src-tauri/src/loras.rs | 1074 ++++++++++++++++++++++++++++++++++++++ src-tauri/src/models.rs | 93 +++- src-tauri/src/watcher.rs | 40 ++ 4 files changed, 1185 insertions(+), 28 deletions(-) create mode 100644 src-tauri/src/loras.rs diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs index b3ed56f..fa50551 100644 --- a/src-tauri/src/lib.rs +++ b/src-tauri/src/lib.rs @@ -43,6 +43,7 @@ mod commands; mod decode; mod generation; mod library; +mod loras; mod mcp; mod midi; mod models; @@ -582,6 +583,9 @@ pub fn run() { // when a model folder appears/disappears (`models://changed`). app.manage(models::InstallManager::new()); watcher::watch_models(app.handle().clone(), models::magenta_models_dir()); + // The LoRA adapter registry (issue #66) reflects native adds/deletes + // live, like the models dir. + watcher::watch_loras(app.handle().clone(), loras::loras_dir()); app.manage(host); // The shell-level interface-state store (ADR-0020): the single source of // truth for the semantic / audio-param interface state the webview @@ -826,6 +830,8 @@ pub fn run() { models::update_model, models::cancel_install, models::open_model_folder, + loras::install_lora, + loras::delete_lora, ]) .build(tauri::generate_context!()) .expect("error while building tauri application") diff --git a/src-tauri/src/loras.rs b/src-tauri/src/loras.rs new file mode 100644 index 0000000..d6aa525 --- /dev/null +++ b/src-tauri/src/loras.rs @@ -0,0 +1,1074 @@ +//! SA3 LoRA adapter manager (issue #66, ADR-0028): import, list, and delete +//! Stable Audio 3 LoRA finetunes, mirroring the model manager's split — Rust +//! owns the lifecycle, `backend/lsdj/loras.py` owns the generate-time read. +//! +//! The registry is a directory layout, discovered from the filesystem like the +//! models (no central index file): +//! +//! ```text +//! ~/Library/Application Support/LSDJai/sa3-loras/// +//! adapter_model.safetensors (the adapter — any single *.safetensors) +//! adapter_config.json (PEFT convention only) +//! lora.json (import manifest: source / type / rank) +//! ``` +//! +//! `` is `small` (the 1024-wide sm-sfx / sm-music DiTs) or `medium` +//! (1536-wide). The trust boundary is ADR-0028's: **only `.safetensors` is +//! accepted** — pickle-backed files (`.ckpt`/`.pt`/`.pth`/`.bin`) are refused +//! before any read, and validation parses only the safetensors JSON header +//! (shapes + metadata), never tensor data. The base is inferred from the +//! adapter's own layer widths (a `lora_A` runs `rank × fan_in`, and the two +//! DiTs share no width), with the PEFT config's `base_model_name_or_path` and +//! an explicit user choice as fallbacks for adapters whose shapes are +//! rank-only (`-xs`). + +use std::collections::BTreeMap; +use std::io::Read; +use std::path::{Path, PathBuf}; +use std::process::Command; + +use serde::{Deserialize, Serialize}; + +use crate::models::{cancelled, dir_size, stream_child, InstallShared, Progress}; + +/// The two DiT families an adapter can ride (`loras.BASES` in Python). +pub const BASES: &[&str] = &["small", "medium"]; + +// Layer widths that identify a base from adapter shapes alone — measured from +// the pinned fork's `dit_mlx.py` (embed 1024, ff-inner 4096, cond 768) and +// `dit_mlx_medium.py` (embed 1536, ff-inner 6144). The sets are disjoint, so +// one matching fan-in decides. +const SMALL_WIDTHS: &[u64] = &[1024, 4096, 768]; +const MEDIUM_WIDTHS: &[u64] = &[1536, 6144]; + +// Pickle-backed extensions ADR-0028 refuses outright. +const PICKLE_EXTS: &[&str] = &["ckpt", "pt", "pth", "bin"]; + +// A safetensors JSON header beyond this is not a plausible adapter — bail +// before allocating attacker-controlled sizes. +const MAX_HEADER_BYTES: u64 = 64 * 1024 * 1024; + +const MANIFEST: &str = "lora.json"; + +/// The registry root. `$SA3_LORAS_HOME` wins (dev/test override); otherwise the +/// app-owned data dir, beside the SA3 checkout. Mirrors `loras.loras_dir`. +pub fn loras_dir() -> PathBuf { + if let Some(override_home) = std::env::var_os("SA3_LORAS_HOME") { + if !override_home.is_empty() { + return PathBuf::from(override_home); + } + } + crate::models::app_support_base().join("sa3-loras") +} + +// --- Registry discovery ---------------------------------------------------- + +/// One installed adapter, for the manager UI and the generate pickers. +#[derive(Serialize, Clone, Debug, PartialEq)] +#[serde(rename_all = "camelCase")] +pub struct LoraInfo { + /// `/` — the name the generate request sends to the backend. + name: String, + base: String, + slug: String, + size_bytes: u64, + /// Import manifest facts, when present (absent for hand-placed adapters). + source: Option, + adapter_type: Option, + rank: Option, +} + +/// The import manifest written beside the adapter (`lora.json`). +#[derive(Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +struct LoraManifest { + source: String, + convention: String, + adapter_type: String, + rank: Option, +} + +/// A well-formed adapter directory: exactly one `*.safetensors` inside (the +/// same rule as `loras._adapter_file` in Python and the runtime's resolver). +fn adapter_file(dir: &Path) -> Option { + let entries = std::fs::read_dir(dir).ok()?; + let mut hits: Vec = entries + .flatten() + .map(|entry| entry.path()) + .filter(|path| { + path.is_file() && path.extension().is_some_and(|ext| ext == "safetensors") + }) + .collect(); + if hits.len() == 1 { + hits.pop() + } else { + None + } +} + +/// One path segment of an adapter name: no separators, no leading dot — a name +/// can only ever address a directory INSIDE the registry (mirrors `loras._SLUG`). +fn valid_slug(slug: &str) -> bool { + let mut chars = slug.chars(); + let Some(first) = chars.next() else { + return false; + }; + first.is_ascii_alphanumeric() + && chars.all(|c| c.is_ascii_alphanumeric() || matches!(c, '.' | '_' | '-')) +} + +/// Parse a client-supplied `/` name; anything malformed is refused +/// before it can touch the filesystem. +fn parse_name(name: &str) -> Result<(&str, &str), String> { + let (base, slug) = name.split_once('/').unwrap_or((name, "")); + if !BASES.contains(&base) || !valid_slug(slug) { + return Err(format!("unknown adapter '{name}'")); + } + Ok((base, slug)) +} + +/// Every installed adapter under `root`, sorted by name. Best-effort like the +/// model discovery: unreadable entries and malformed directories are skipped. +pub fn discover(root: &Path) -> Vec { + let mut adapters = Vec::new(); + for base in BASES { + let Ok(entries) = std::fs::read_dir(root.join(base)) else { + continue; + }; + for entry in entries.flatten() { + let dir = entry.path(); + let slug = entry.file_name().to_string_lossy().into_owned(); + if !dir.is_dir() || !valid_slug(&slug) || adapter_file(&dir).is_none() { + continue; + } + let manifest: Option = std::fs::read_to_string(dir.join(MANIFEST)) + .ok() + .and_then(|data| serde_json::from_str(&data).ok()); + adapters.push(LoraInfo { + name: format!("{base}/{slug}"), + base: (*base).to_string(), + slug, + size_bytes: dir_size(&dir), + source: manifest.as_ref().map(|m| m.source.clone()), + adapter_type: manifest.as_ref().map(|m| m.adapter_type.clone()), + rank: manifest.as_ref().and_then(|m| m.rank), + }); + } + } + adapters.sort_by(|a, b| a.name.cmp(&b.name)); + adapters +} + +/// The installed-adapter names, for the watcher's changed-set comparison. +pub fn discover_names(root: &Path) -> std::collections::BTreeSet { + discover(root).into_iter().map(|info| info.name).collect() +} + +// --- Validation (the ADR-0028 trust boundary) ------------------------------ + +/// What the safetensors header + config told us about an adapter. +#[derive(Debug, PartialEq)] +pub struct AdapterFacts { + convention: Convention, + adapter_type: String, + rank: Option, + /// The base inferred from layer widths / the PEFT config; `None` when the + /// adapter is shape-anonymous (rank-only `-xs` tensors, no config hint). + inferred_base: Option<&'static str>, +} + +#[derive(Debug, PartialEq, Clone, Copy)] +enum Convention { + /// SA3-native `train_lora.py` output — config in the safetensors metadata. + Native, + /// HuggingFace `peft` — config in a sibling `adapter_config.json`. + Peft, +} + +impl Convention { + fn as_str(self) -> &'static str { + match self { + Convention::Native => "native", + Convention::Peft => "peft", + } + } +} + +fn is_pickle(path: &Path) -> bool { + path.extension() + .and_then(|ext| ext.to_str()) + .is_some_and(|ext| { + let lower = ext.to_ascii_lowercase(); + PICKLE_EXTS.contains(&lower.as_str()) + }) +} + +/// The parsed safetensors header: tensor name → shape, plus `__metadata__`. +/// Reads ONLY the 8-byte length + JSON header — never tensor data, never a +/// pickle (`mx.load`-equivalent trust posture without loading anything). +fn read_safetensors_header( + path: &Path, +) -> Result<(BTreeMap>, BTreeMap), String> { + let file_name = path.file_name().unwrap_or_default().to_string_lossy(); + if is_pickle(path) { + return Err(format!( + "refusing the pickle-format adapter '{file_name}' — only .safetensors \ + adapters are accepted (a .ckpt/.pt is unpickled by torch.load and can \ + execute arbitrary code)" + )); + } + if !path.extension().is_some_and(|ext| ext == "safetensors") { + return Err(format!("'{file_name}' is not a .safetensors adapter")); + } + let mut file = std::fs::File::open(path) + .map_err(|e| format!("cannot open '{file_name}': {e}"))?; + let mut len_bytes = [0u8; 8]; + file.read_exact(&mut len_bytes) + .map_err(|_| format!("'{file_name}' is not a safetensors file"))?; + let header_len = u64::from_le_bytes(len_bytes); + if header_len == 0 || header_len > MAX_HEADER_BYTES { + return Err(format!("'{file_name}' is not a safetensors file")); + } + let mut header = vec![0u8; header_len as usize]; + file.read_exact(&mut header) + .map_err(|_| format!("'{file_name}' is not a safetensors file"))?; + let parsed: serde_json::Map = serde_json::from_slice(&header) + .map_err(|_| format!("'{file_name}' is not a safetensors file"))?; + + let mut tensors = BTreeMap::new(); + let mut metadata = BTreeMap::new(); + for (key, value) in parsed { + if key == "__metadata__" { + if let Some(map) = value.as_object() { + for (meta_key, meta_value) in map { + if let Some(text) = meta_value.as_str() { + metadata.insert(meta_key.clone(), text.to_string()); + } + } + } + continue; + } + let shape = value + .get("shape") + .and_then(|s| s.as_array()) + .map(|dims| dims.iter().filter_map(|d| d.as_u64()).collect::>()) + .unwrap_or_default(); + tensors.insert(key, shape); + } + Ok((tensors, metadata)) +} + +/// Tensor keys of the SA3-native convention +/// (`.parametrizations.weight.0.`). +const NATIVE_MARKER: &str = ".parametrizations.weight.0."; + +/// Infer the base from the fan-in widths of 2-D `lora_A` tensors (shape +/// `rank × fan_in`; the DiT widths are disjoint between the bases). `Err` on a +/// contradiction, `Ok(None)` when nothing matched (rank-only shapes). +fn base_from_widths(tensors: &BTreeMap>) -> Result, String> { + let mut small = false; + let mut medium = false; + for (key, shape) in tensors { + let is_lora_a = key.ends_with(".lora_A.weight") || key.ends_with(".lora_A"); + if !is_lora_a || shape.len() != 2 { + continue; + } + let fan_in = shape[1]; + small |= SMALL_WIDTHS.contains(&fan_in); + medium |= MEDIUM_WIDTHS.contains(&fan_in); + } + match (small, medium) { + (true, true) => Err( + "the adapter mixes small-DiT and medium-DiT layer widths — not a valid \ + SA3 LoRA" + .into(), + ), + (true, false) => Ok(Some("small")), + (false, true) => Ok(Some("medium")), + (false, false) => Ok(None), + } +} + +/// The PEFT `adapter_config.json` fields the importer reads. +#[derive(Deserialize, Default)] +struct PeftConfig { + r: Option, + use_dora: Option, + base_model_name_or_path: Option, +} + +/// The SA3-native `lora_config` metadata fields the importer reads. +#[derive(Deserialize, Default)] +struct NativeConfig { + adapter_type: Option, + rank: Option, +} + +fn base_from_model_name(model_name: &str) -> Option<&'static str> { + let lower = model_name.to_ascii_lowercase(); + if lower.contains("medium") { + Some("medium") + } else if lower.contains("small") || lower.contains("sm-") { + Some("small") + } else { + None + } +} + +/// Infer a rank from the tensors when no config declares it: `lora_A` is +/// `rank × fan_in`, `M_xs` is `rank × rank`. +fn rank_from_tensors(tensors: &BTreeMap>) -> Option { + for (key, shape) in tensors { + let is_rank_first = key.ends_with(".lora_A.weight") + || key.ends_with(".lora_A") + || key.ends_with(".M_xs"); + if is_rank_first { + if let Some(&rank) = shape.first() { + return u32::try_from(rank).ok(); + } + } + } + None +} + +/// Validate one adapter file structurally (ADR-0028): recognised convention, +/// inferable type/rank, and a base read from the shapes. `peft_config_path` is +/// the sibling `adapter_config.json` when the file follows that convention. +pub fn validate_adapter(path: &Path) -> Result { + let (tensors, metadata) = read_safetensors_header(path)?; + + if tensors.keys().any(|key| key.contains(NATIVE_MARKER)) { + let config: NativeConfig = metadata + .get("lora_config") + .and_then(|raw| serde_json::from_str(raw).ok()) + .unwrap_or_default(); + // Legacy 'dora' means 'dora-rows' (the paper-correct default) — + // mirrors the runtime's resolve_adapter_type. + let adapter_type = match config.adapter_type.as_deref() { + None => "lora".to_string(), + Some("dora") => "dora-rows".to_string(), + Some(other) => other.to_string(), + }; + return Ok(AdapterFacts { + convention: Convention::Native, + adapter_type, + rank: config.rank.or_else(|| rank_from_tensors(&tensors)), + inferred_base: base_from_widths(&tensors)?, + }); + } + + if tensors.keys().any(|key| key.ends_with(".lora_A.weight")) { + let config_path = path.with_file_name("adapter_config.json"); + if !config_path.is_file() { + return Err( + "the PEFT adapter is missing its adapter_config.json sibling".into(), + ); + } + let config: PeftConfig = std::fs::read_to_string(&config_path) + .ok() + .and_then(|data| serde_json::from_str(&data).ok()) + .ok_or("the adapter's adapter_config.json is unreadable")?; + let inferred_base = match base_from_widths(&tensors)? { + Some(base) => Some(base), + None => config + .base_model_name_or_path + .as_deref() + .and_then(base_from_model_name), + }; + return Ok(AdapterFacts { + convention: Convention::Peft, + adapter_type: if config.use_dora.unwrap_or(false) { + "dora-rows".to_string() + } else { + "lora".to_string() + }, + rank: config.r.or_else(|| rank_from_tensors(&tensors)), + inferred_base, + }); + } + + Err( + "not a recognised SA3 LoRA (no SA3-native parametrization keys and no PEFT \ + lora_A/lora_B keys)" + .into(), + ) +} + +// --- Import ---------------------------------------------------------------- + +/// What to import: exactly one of a HuggingFace repo id or a local path +/// (a `.safetensors` file or a PEFT adapter directory), plus an optional +/// explicit base for shape-anonymous adapters. +#[derive(Deserialize, Clone)] +#[serde(rename_all = "camelCase")] +pub struct ImportSpec { + pub hf_repo: Option, + pub path: Option, + pub base: Option, +} + +impl ImportSpec { + /// The display name for progress events; doubles as the early shape check + /// (exactly one source, well-formed repo id, known base). + pub fn display_name(&self) -> Result { + if let Some(base) = self.base.as_deref() { + if !BASES.contains(&base) { + return Err(format!("unknown base '{base}'")); + } + } + match (self.hf_repo.as_deref(), self.path.as_deref()) { + (Some(repo), None) => { + if !valid_hf_repo(repo) { + return Err(format!("'{repo}' is not a HuggingFace repo id")); + } + Ok(repo.to_string()) + } + (None, Some(path)) => Ok(Path::new(path) + .file_name() + .unwrap_or_default() + .to_string_lossy() + .into_owned()), + _ => Err("exactly one of a HuggingFace repo or a local path is required".into()), + } + } +} + +/// `owner/name`, each segment alphanumeric-headed with `. _ -` allowed — enough +/// to build the API/resolve URLs without the id smuggling path or query parts. +fn valid_hf_repo(repo: &str) -> bool { + match repo.split_once('/') { + Some((owner, name)) => valid_slug(owner) && valid_slug(name), + None => false, + } +} + +/// A registry slug derived from a repo/file name: invalid characters become +/// `-`, a leading non-alphanumeric is stripped. +fn slugify(name: &str) -> Result { + let cleaned: String = name + .chars() + .map(|c| { + if c.is_ascii_alphanumeric() || matches!(c, '.' | '_' | '-') { + c + } else { + '-' + } + }) + .collect(); + let slug = cleaned + .trim_start_matches(|c: char| !c.is_ascii_alphanumeric()) + .to_string(); + if valid_slug(&slug) { + Ok(slug) + } else { + Err(format!("cannot derive an adapter name from '{name}'")) + } +} + +/// Pick the adapter files to download from a HF repo's file list: the +/// conventional `adapter_model.safetensors`, or the single `*.safetensors` — +/// plus the sibling `adapter_config.json` when present. Pickle-only repos get +/// the trust-boundary refusal, not a generic "nothing found". The file list is +/// untrusted (it comes back from the HF API): only root-level names made of +/// slug characters are considered, so a name can neither escape the download +/// dir nor smuggle URL path/query parts. +fn choose_hf_files(filenames: &[String]) -> Result, String> { + let config = filenames + .iter() + .find(|name| name.as_str() == "adapter_config.json"); + let safetensors: Vec<&String> = filenames + .iter() + .filter(|name| name.ends_with(".safetensors") && valid_slug(name)) + .collect(); + let adapter = if safetensors + .iter() + .any(|name| name.as_str() == "adapter_model.safetensors") + { + "adapter_model.safetensors".to_string() + } else { + match safetensors.as_slice() { + [single] => (*single).clone(), + [] => { + let has_pickle = filenames.iter().any(|name| { + is_pickle(Path::new(name.as_str())) + }); + return Err(if has_pickle { + "the repo only ships pickle-format weights (.ckpt/.pt/.bin), \ + which are refused — only .safetensors adapters are accepted" + .into() + } else { + "no .safetensors adapter in the repo".into() + }); + } + _ => return Err("the repo holds more than one .safetensors — not a single adapter".into()), + } + }; + let mut files = vec![adapter]; + if let Some(config) = config { + files.push(config.clone()); + } + Ok(files) +} + +/// The HF model-info response — only the file list is read. +#[derive(Deserialize)] +struct HfModelInfo { + siblings: Vec, +} + +#[derive(Deserialize)] +struct HfSibling { + rfilename: String, +} + +/// Run one adapter import to completion (on the install thread): fetch or copy +/// the files, validate them (ADR-0028), resolve the base, and place the adapter +/// directory into the registry. +pub(crate) fn install( + progress: &Progress, + shared: &InstallShared, + spec: &ImportSpec, +) -> Result<(), String> { + let staged = match (spec.hf_repo.as_deref(), spec.path.as_deref()) { + (Some(repo), None) => fetch_hf_adapter(progress, shared, repo)?, + (None, Some(path)) => stage_local_adapter(Path::new(path))?, + _ => return Err("exactly one of a HuggingFace repo or a local path is required".into()), + }; + let result = (|| { + cancelled(shared)?; + progress("install", None, None); + let facts = validate_adapter(&staged.adapter)?; + let base = resolve_base(&facts, spec.base.as_deref())?; + let slug = slugify(&staged.slug_seed)?; + place_adapter(&loras_dir(), base, &slug, &staged, &facts) + })(); + // The HF staging dir is cleaned up on every exit; a local import's source + // files are the user's and stay put. + if let Some(temp) = &staged.temp { + let _ = std::fs::remove_dir_all(temp); + } + result +} + +/// A staged (temp or user-supplied) adapter before validation: the +/// `.safetensors` path, its optional PEFT config sibling, and what to derive +/// the registry slug from. +struct StagedAdapter { + adapter: PathBuf, + config: Option, + slug_seed: String, + /// The temp dir to clean up after placing (None for a local import, whose + /// source files are the user's and must stay put). + temp: Option, +} + +/// Reconcile the inferred base with an explicit choice. An explicit base wins +/// only when the shapes are silent; a contradiction is refused with the +/// reasoning (the issue's "incompatible adapters refused with clear reasoning"). +fn resolve_base( + facts: &AdapterFacts, + explicit: Option<&str>, +) -> Result<&'static str, String> { + match (facts.inferred_base, explicit) { + (Some(inferred), Some(chosen)) if inferred != chosen => Err(format!( + "the adapter's layer widths identify the {inferred} DiT — it cannot ride \ + the {chosen} base" + )), + (Some(inferred), _) => Ok(inferred), + (None, Some(chosen)) => BASES + .iter() + .find(|base| **base == chosen) + .copied() + .ok_or_else(|| format!("unknown base '{chosen}'")), + (None, None) => Err( + "cannot tell which DiT this adapter rides (its tensors are rank-only) — \ + pick the base (small or medium) explicitly" + .into(), + ), + } +} + +/// Fetch an adapter from HuggingFace into a temp dir: the model-info file list +/// first, then each chosen file via the resolve endpoint. Uses `curl` like the +/// SA3 checkout fetch (no HTTP stack in the shell). +fn fetch_hf_adapter( + progress: &Progress, + shared: &InstallShared, + repo: &str, +) -> Result { + if !valid_hf_repo(repo) { + return Err(format!("'{repo}' is not a HuggingFace repo id")); + } + let temp = std::env::temp_dir().join(format!("lsdj-lora-{}", std::process::id())); + let _ = std::fs::remove_dir_all(&temp); + std::fs::create_dir_all(&temp).map_err(|e| format!("cannot create temp dir: {e}"))?; + + progress("fetch", None, None); + let info_path = temp.join("model-info.json"); + let mut curl = Command::new("curl"); + curl.args(["-fLsS", "-o"]) + .arg(&info_path) + .arg(format!("https://huggingface.co/api/models/{repo}")); + stream_child(shared, "hf-info", curl, |_| {}) + .map_err(|e| format!("cannot reach the HuggingFace repo '{repo}': {e}"))?; + cancelled(shared)?; + let info: HfModelInfo = std::fs::read_to_string(&info_path) + .ok() + .and_then(|data| serde_json::from_str(&data).ok()) + .ok_or_else(|| format!("unexpected HuggingFace response for '{repo}'"))?; + let filenames: Vec = info + .siblings + .into_iter() + .map(|sibling| sibling.rfilename) + .collect(); + let files = choose_hf_files(&filenames)?; + + let mut adapter = None; + let mut config = None; + for file in &files { + cancelled(shared)?; + progress("download", None, Some(file.clone())); + let dest = temp.join(file); + let mut curl = Command::new("curl"); + curl.args(["-fLsS", "-o"]) + .arg(&dest) + .arg(format!("https://huggingface.co/{repo}/resolve/main/{file}")); + stream_child(shared, "hf-download", curl, |_| {}) + .map_err(|e| format!("download of '{file}' failed: {e}"))?; + if file.ends_with(".safetensors") { + adapter = Some(dest); + } else { + config = Some(dest); + } + } + Ok(StagedAdapter { + adapter: adapter.ok_or("the repo download produced no adapter")?, + config, + // The repo's own name seeds the slug (`owner/name` → `name`). + slug_seed: repo.split('/').next_back().unwrap_or(repo).to_string(), + temp: Some(temp), + }) +} + +/// Stage a local import: a `.safetensors` file (native, or PEFT with its +/// sibling config next to it) or a PEFT adapter directory. +fn stage_local_adapter(path: &Path) -> Result { + let file_name = path.file_name().unwrap_or_default().to_string_lossy(); + if is_pickle(path) { + return Err(format!( + "refusing the pickle-format adapter '{file_name}' — only .safetensors \ + adapters are accepted (a .ckpt/.pt is unpickled by torch.load and can \ + execute arbitrary code)" + )); + } + let adapter = if path.is_dir() { + adapter_file(path).ok_or_else(|| { + format!("expected one .safetensors adapter in '{file_name}'") + })? + } else { + path.to_path_buf() + }; + let config_path = adapter.with_file_name("adapter_config.json"); + let slug_seed = adapter + .file_stem() + .unwrap_or_default() + .to_string_lossy() + .into_owned(); + // A generic HF artifact name would collide across adapters; the folder + // name is the distinctive one. + let slug_seed = if slug_seed == "adapter_model" { + adapter + .parent() + .and_then(|parent| parent.file_name()) + .map(|name| name.to_string_lossy().into_owned()) + .unwrap_or(slug_seed) + } else { + slug_seed + }; + Ok(StagedAdapter { + adapter, + config: config_path.is_file().then_some(config_path), + slug_seed, + temp: None, + }) +} + +/// Copy the staged files into `///` (built in a sibling temp +/// dir, renamed into place so a failure never leaves a half-adapter), write the +/// manifest, and clean the staging dir up. +fn place_adapter( + root: &Path, + base: &str, + slug: &str, + staged: &StagedAdapter, + facts: &AdapterFacts, +) -> Result<(), String> { + let dest = root.join(base).join(slug); + if dest.exists() { + return Err(format!( + "an adapter named '{base}/{slug}' is already installed — delete it first" + )); + } + std::fs::create_dir_all(root.join(base)) + .map_err(|e| format!("cannot create the adapter registry: {e}"))?; + let staging = root.join(base).join(format!(".{slug}.importing")); + let _ = std::fs::remove_dir_all(&staging); + std::fs::create_dir_all(&staging) + .map_err(|e| format!("cannot stage the adapter: {e}"))?; + + let place = (|| -> Result<(), String> { + let adapter_name = staged + .adapter + .file_name() + .ok_or("the adapter path has no file name")?; + std::fs::copy(&staged.adapter, staging.join(adapter_name)) + .map_err(|e| format!("cannot copy the adapter: {e}"))?; + if let Some(config) = &staged.config { + std::fs::copy(config, staging.join("adapter_config.json")) + .map_err(|e| format!("cannot copy adapter_config.json: {e}"))?; + } + let manifest = LoraManifest { + source: staged.slug_seed.clone(), + convention: facts.convention.as_str().to_string(), + adapter_type: facts.adapter_type.clone(), + rank: facts.rank, + }; + let json = serde_json::to_string_pretty(&manifest) + .map_err(|e| format!("cannot write the manifest: {e}"))?; + std::fs::write(staging.join(MANIFEST), json) + .map_err(|e| format!("cannot write the manifest: {e}"))?; + std::fs::rename(&staging, &dest) + .map_err(|e| format!("cannot place the adapter: {e}")) + })(); + if place.is_err() { + let _ = std::fs::remove_dir_all(&staging); + } + place +} + +// --- Tauri commands -------------------------------------------------------- + +/// Import an adapter from a HuggingFace repo id or a local path (issue #66). +/// Runs on the shared install thread; progress arrives as `model://progress` +/// with family `lora`, completion as `models://changed`. +#[tauri::command] +pub fn install_lora( + installer: tauri::State<'_, crate::models::InstallManager>, + app: tauri::AppHandle, + spec: ImportSpec, +) -> Result<(), String> { + installer.install_lora(app, spec) +} + +/// Delete an installed adapter. In-app deletion is fine here, unlike the model +/// families: an adapter is a small re-downloadable artifact in the app-owned +/// dir (never iCloud-managed), and the issue asks for the full +/// install/list/delete lifecycle. +#[tauri::command] +pub fn delete_lora(app: tauri::AppHandle, name: String) -> Result<(), String> { + use tauri::Emitter; + let (base, slug) = parse_name(&name)?; + let dir = loras_dir().join(base).join(slug); + if adapter_file(&dir).is_none() { + return Err(format!("unknown adapter '{name}'")); + } + std::fs::remove_dir_all(&dir).map_err(|e| format!("cannot delete '{name}': {e}"))?; + let _ = app.emit("models://changed", ()); + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn temp_root(tag: &str) -> PathBuf { + let tmp = std::env::temp_dir().join(format!("lsdj-lora-{tag}-{}", std::process::id())); + let _ = std::fs::remove_dir_all(&tmp); + std::fs::create_dir_all(&tmp).unwrap(); + tmp + } + + /// Build a minimal safetensors file: the 8-byte length + JSON header, with + /// zeroed tensor data appended (sizes don't matter — only the header is read). + fn write_safetensors(path: &Path, tensors: &[(&str, &[u64])], metadata: &[(&str, &str)]) { + let mut header = serde_json::Map::new(); + if !metadata.is_empty() { + let meta: serde_json::Map = metadata + .iter() + .map(|(k, v)| (k.to_string(), serde_json::Value::String(v.to_string()))) + .collect(); + header.insert("__metadata__".into(), meta.into()); + } + for (name, shape) in tensors { + header.insert( + name.to_string(), + serde_json::json!({"dtype": "F16", "shape": shape, "data_offsets": [0, 0]}), + ); + } + let json = serde_json::to_vec(&serde_json::Value::Object(header)).unwrap(); + let mut bytes = (json.len() as u64).to_le_bytes().to_vec(); + bytes.extend_from_slice(&json); + std::fs::write(path, bytes).unwrap(); + } + + #[test] + fn a_pickle_extension_is_refused_before_any_read() { + let tmp = temp_root("pickle"); + let path = tmp.join("adapter.ckpt"); + std::fs::write(&path, b"not even opened").unwrap(); + let error = validate_adapter(&path).unwrap_err(); + assert!(error.contains("pickle"), "unexpected error: {error}"); + let _ = std::fs::remove_dir_all(&tmp); + } + + #[test] + fn a_non_safetensors_file_is_refused() { + let tmp = temp_root("garbage"); + let path = tmp.join("adapter.safetensors"); + std::fs::write(&path, b"RIFF not a safetensors").unwrap(); + assert!(validate_adapter(&path).is_err()); + let _ = std::fs::remove_dir_all(&tmp); + } + + #[test] + fn a_peft_adapter_yields_its_facts_and_medium_base() { + let tmp = temp_root("peft"); + let path = tmp.join("adapter_model.safetensors"); + write_safetensors( + &path, + &[ + ( + "base_model.model.transformer.layers.0.self_attn.to_qkv.lora_A.weight", + &[64, 1536], + ), + ( + "base_model.model.transformer.layers.0.self_attn.to_qkv.lora_B.weight", + &[7680, 64], + ), + ], + &[], + ); + std::fs::write( + tmp.join("adapter_config.json"), + r#"{"r": 64, "lora_alpha": 128, "use_dora": false, + "base_model_name_or_path": "stabilityai/stable-audio-3-medium"}"#, + ) + .unwrap(); + let facts = validate_adapter(&path).unwrap(); + assert_eq!(facts.convention, Convention::Peft); + assert_eq!(facts.adapter_type, "lora"); + assert_eq!(facts.rank, Some(64)); + assert_eq!(facts.inferred_base, Some("medium")); + let _ = std::fs::remove_dir_all(&tmp); + } + + #[test] + fn a_peft_adapter_without_its_config_is_refused() { + let tmp = temp_root("peft-noconfig"); + let path = tmp.join("adapter_model.safetensors"); + write_safetensors( + &path, + &[("x.lora_A.weight", &[8, 1024]), ("x.lora_B.weight", &[1024, 8])], + &[], + ); + let error = validate_adapter(&path).unwrap_err(); + assert!(error.contains("adapter_config.json"), "unexpected error: {error}"); + let _ = std::fs::remove_dir_all(&tmp); + } + + #[test] + fn a_native_adapter_reads_its_metadata_config_and_small_base() { + let tmp = temp_root("native"); + let path = tmp.join("finetune.safetensors"); + write_safetensors( + &path, + &[ + ( + "transformer.layers.0.self_attn.to_qkv.parametrizations.weight.0.lora_A", + &[16, 1024], + ), + ( + "transformer.layers.0.self_attn.to_qkv.parametrizations.weight.0.lora_B", + &[3072, 16], + ), + ], + &[("lora_config", r#"{"adapter_type": "dora", "rank": 16, "alpha": 32}"#)], + ); + let facts = validate_adapter(&path).unwrap(); + assert_eq!(facts.convention, Convention::Native); + assert_eq!(facts.adapter_type, "dora-rows"); + assert_eq!(facts.rank, Some(16)); + assert_eq!(facts.inferred_base, Some("small")); + let _ = std::fs::remove_dir_all(&tmp); + } + + #[test] + fn an_xs_adapter_is_shape_anonymous_and_needs_an_explicit_base() { + let tmp = temp_root("xs"); + let path = tmp.join("finetune-xs.safetensors"); + write_safetensors( + &path, + &[( + "transformer.layers.0.self_attn.to_qkv.parametrizations.weight.0.M_xs", + &[32, 32], + )], + &[("lora_config", r#"{"adapter_type": "lora-xs", "rank": 32}"#)], + ); + let facts = validate_adapter(&path).unwrap(); + assert_eq!(facts.inferred_base, None); + assert_eq!(facts.rank, Some(32)); + // Unresolvable without a choice; resolvable with one; a contradiction + // elsewhere is refused. + assert!(resolve_base(&facts, None).is_err()); + assert_eq!(resolve_base(&facts, Some("medium")), Ok("medium")); + let shaped = AdapterFacts { + convention: Convention::Native, + adapter_type: "lora".into(), + rank: Some(16), + inferred_base: Some("medium"), + }; + assert!(resolve_base(&shaped, Some("small")).is_err()); + let _ = std::fs::remove_dir_all(&tmp); + } + + #[test] + fn a_random_safetensors_is_not_a_lora() { + let tmp = temp_root("notlora"); + let path = tmp.join("weights.safetensors"); + write_safetensors(&path, &[("model.embed.weight", &[512, 1024])], &[]); + let error = validate_adapter(&path).unwrap_err(); + assert!(error.contains("not a recognised"), "unexpected error: {error}"); + let _ = std::fs::remove_dir_all(&tmp); + } + + #[test] + fn hf_file_choice_prefers_the_convention_and_refuses_pickle_only_repos() { + let names = |list: &[&str]| list.iter().map(|s| s.to_string()).collect::>(); + // The conventional pair. + assert_eq!( + choose_hf_files(&names(&[ + "README.md", + "adapter_config.json", + "adapter_model.safetensors" + ])) + .unwrap(), + vec!["adapter_model.safetensors", "adapter_config.json"] + ); + // A single differently-named safetensors, no config (SA3-native). + assert_eq!( + choose_hf_files(&names(&["README.md", "maqam.safetensors"])).unwrap(), + vec!["maqam.safetensors"] + ); + // Pickle-only → the trust-boundary refusal, by name. + let error = choose_hf_files(&names(&["adapter.ckpt", "README.md"])).unwrap_err(); + assert!(error.contains("pickle"), "unexpected error: {error}"); + // Nothing usable at all. + assert!(choose_hf_files(&names(&["README.md"])).is_err()); + // Ambiguous. + assert!(choose_hf_files(&names(&["a.safetensors", "b.safetensors"])).is_err()); + } + + #[test] + fn adapter_names_parse_and_hostile_ones_are_refused() { + assert_eq!(parse_name("medium/maqam").unwrap(), ("medium", "maqam")); + for hostile in [ + "maqam", + "large/maqam", + "medium/../maqam", + "medium/.hidden", + "medium/", + "medium/sub/dir", + "MEDIUM/maqam", + ] { + assert!(parse_name(hostile).is_err(), "accepted {hostile:?}"); + } + } + + #[test] + fn hf_repo_ids_validate() { + assert!(valid_hf_repo("motiftechnologies/stable-audio-3-maqam-lora")); + for hostile in [ + "no-slash", + "a/b/c", + "../x/y", + "a/..", + "a/b?x=1", + "https://huggingface.co/a/b", + ] { + assert!(!valid_hf_repo(hostile), "accepted {hostile:?}"); + } + } + + #[test] + fn a_local_import_lands_in_the_registry_and_discovery_sees_it() { + let root = temp_root("import"); + let source_dir = root.join("source"); + std::fs::create_dir_all(&source_dir).unwrap(); + let source = source_dir.join("maqam.safetensors"); + write_safetensors( + &source, + &[( + "transformer.layers.0.self_attn.to_out.parametrizations.weight.0.lora_A", + &[8, 1536], + )], + &[("lora_config", r#"{"adapter_type": "lora", "rank": 8}"#)], + ); + + let registry = root.join("registry"); + let staged = StagedAdapter { + adapter: source.clone(), + config: None, + slug_seed: "maqam".into(), + temp: None, + }; + let facts = validate_adapter(&source).unwrap(); + let base = resolve_base(&facts, None).unwrap(); + place_adapter(®istry, base, "maqam", &staged, &facts).unwrap(); + + // The user's source file stays put; the registry holds the copy. + assert!(source.is_file()); + let installed = discover(®istry); + assert_eq!(installed.len(), 1); + assert_eq!(installed[0].name, "medium/maqam"); + assert_eq!(installed[0].adapter_type.as_deref(), Some("lora")); + assert_eq!(installed[0].rank, Some(8)); + + // A second import under the same name is refused, not overwritten. + let error = place_adapter(®istry, base, "maqam", &staged, &facts).unwrap_err(); + assert!(error.contains("already installed"), "unexpected error: {error}"); + let _ = std::fs::remove_dir_all(&root); + } + + #[test] + fn discovery_skips_malformed_directories() { + let root = temp_root("discover"); + // Well-formed. + let good = root.join("small").join("crackle"); + std::fs::create_dir_all(&good).unwrap(); + write_safetensors(&good.join("crackle.safetensors"), &[("x.lora_A", &[4, 1024])], &[]); + // No safetensors. + std::fs::create_dir_all(root.join("small").join("empty")).unwrap(); + // Two safetensors — ambiguous, skipped (matches the Python resolver). + let two = root.join("medium").join("two"); + std::fs::create_dir_all(&two).unwrap(); + write_safetensors(&two.join("a.safetensors"), &[("x.lora_A", &[4, 1536])], &[]); + write_safetensors(&two.join("b.safetensors"), &[("x.lora_A", &[4, 1536])], &[]); + // A dot-dir never becomes a name. + let hidden = root.join("medium").join(".importing"); + std::fs::create_dir_all(&hidden).unwrap(); + write_safetensors(&hidden.join("x.safetensors"), &[("x.lora_A", &[4, 1536])], &[]); + + let names: Vec = discover(&root).into_iter().map(|info| info.name).collect(); + assert_eq!(names, vec!["small/crackle".to_string()]); + let _ = std::fs::remove_dir_all(&root); + } + + #[test] + fn slugs_are_derived_and_sanitised() { + assert_eq!(slugify("stable-audio-3-maqam-lora").unwrap(), "stable-audio-3-maqam-lora"); + assert_eq!(slugify("My Adapter (v2)").unwrap(), "My-Adapter--v2-"); + assert_eq!(slugify("..sneaky").unwrap(), "sneaky"); + assert!(slugify("...").is_err()); + assert!(slugify("").is_err()); + } +} diff --git a/src-tauri/src/models.rs b/src-tauri/src/models.rs index 08466ba..209e531 100644 --- a/src-tauri/src/models.rs +++ b/src-tauri/src/models.rs @@ -71,7 +71,7 @@ pub fn magenta_models_dir() -> PathBuf { /// The app-owned data root for model weights — `~/Library/Application Support/ /// LSDJai`. Kept out of `~/Documents` (which users may sync to iCloud, where /// multi-GB weights don't belong and Finder ops on offloaded files fail, -8013). -fn app_support_base() -> PathBuf { +pub(crate) fn app_support_base() -> PathBuf { home_dir() .join("Library") .join("Application Support") @@ -233,7 +233,7 @@ fn sa3_update_available(installed: Option<&Sa3Source>, pinned: &Sa3Source, prese /// into the shared cache; the target size is the meaningful "how big is this"). /// Best-effort: unreadable entries are skipped, and a symlinked directory is not /// traversed (so it cannot loop). -fn dir_size(path: &Path) -> u64 { +pub(crate) fn dir_size(path: &Path) -> u64 { let mut total = 0u64; let Ok(entries) = std::fs::read_dir(path) else { return 0; @@ -356,6 +356,9 @@ pub struct ActiveInstall { pub struct ModelStatus { magenta: MagentaStatus, sa3: Sa3Status, + /// The installed SA3 LoRA adapters (issue #66), listed with the models so + /// the manager drawer and the generate pickers share one snapshot. + loras: Vec, installing: Option, } @@ -394,6 +397,7 @@ fn status(active: Option<(Family, String)>) -> ModelStatus { pinned_source: pinned, update_available, }, + loras: crate::loras::discover(&crate::loras::loras_dir()), installing: active.map(|(family, name)| ActiveInstall { family, name, @@ -410,6 +414,9 @@ fn status(active: Option<(Family, String)>) -> ModelStatus { pub enum Family { Magenta, Sa3, + /// SA3 LoRA adapters (issue #66) — same progress/changed channels, its own + /// import commands (an adapter needs a source + optional base, not a name). + Lora, } /// The `model://progress` payload the webview renders as a live install bar. @@ -461,7 +468,7 @@ fn sa3_install_script() -> PathBuf { /// `active` names the in-flight job so `model_status` can report it — the manager /// reflects an install even after the drawer was closed and reopened (the live /// `model://progress` events are missed while it's unmounted). -struct InstallShared { +pub(crate) struct InstallShared { busy: AtomicBool, cancelled: AtomicBool, current_child: Mutex>, @@ -504,37 +511,64 @@ impl InstallManager { name: Option, update: bool, ) -> Result<(), String> { - if family == Family::Magenta { - let name = name.as_deref().ok_or("a model name is required")?; - if !INSTALLABLE_MODELS.contains(&name) { - return Err(format!("unknown model '{name}'")); + match family { + Family::Magenta => { + let name = name.ok_or("a model name is required")?; + if !INSTALLABLE_MODELS.contains(&name.as_str()) { + return Err(format!("unknown model '{name}'")); + } + let model = name.clone(); + self.start(app, family, name, move |progress, shared| { + install_magenta(progress, shared, &model) + }) } + // `model://progress` carries the model name for Magenta, "" for SA3. + Family::Sa3 => self.start(app, family, String::new(), move |progress, shared| { + install_sa3(progress, shared, update) + }), + Family::Lora => Err("adapters are installed via install_lora".into()), } + } + + /// Import an SA3 LoRA adapter (issue #66) on the same install thread and + /// event channels; `spec` names the source (HuggingFace repo or local path) + /// and an optional explicit base. + pub fn install_lora( + &self, + app: AppHandle, + spec: crate::loras::ImportSpec, + ) -> Result<(), String> { + let name = spec.display_name()?; + self.start(app, Family::Lora, name, move |progress, shared| { + crate::loras::install(progress, shared, &spec) + }) + } + + /// The shared install-thread dance: claim the single install slot, run `job` + /// with a progress sink wired to `model://progress` (as `family`/`name`), + /// then emit the terminal event and `models://changed`. + fn start( + &self, + app: AppHandle, + family: Family, + name: String, + job: impl FnOnce(&Progress, &InstallShared) -> Result<(), String> + Send + 'static, + ) -> Result<(), String> { if self.shared.busy.swap(true, Ordering::AcqRel) { return Err("an install is already running".into()); } self.shared.cancelled.store(false, Ordering::Release); *self.shared.active.lock().unwrap_or_else(|p| p.into_inner()) = - Some((family, name.clone().unwrap_or_default())); + Some((family, name.clone())); let shared = self.shared.clone(); std::thread::Builder::new() .name("lsdj-model-install".into()) .spawn(move || { - let name = name.unwrap_or_default(); - // `model://progress` carries the model name for Magenta, "" for SA3. - let event_name = if family == Family::Magenta { - name.clone() - } else { - String::new() - }; let progress_app = app.clone(); let progress = move |stage: &str, message: Option, file: Option| { - emit(&progress_app, family, &event_name, stage, message, file); - }; - let result = match family { - Family::Magenta => install_magenta(&progress, &shared, &name), - Family::Sa3 => install_sa3(&progress, &shared, update), + emit(&progress_app, family, &name, stage, message, file); }; + let result = job(&progress, &shared); *shared.current_child.lock().unwrap_or_else(|p| p.into_inner()) = None; match result { Ok(()) => emit(&app, family, "", "done", None, None), @@ -618,7 +652,7 @@ fn kill_group(child: &mut Child) { } } -fn cancelled(shared: &InstallShared) -> Result<(), String> { +pub(crate) fn cancelled(shared: &InstallShared) -> Result<(), String> { if shared.cancelled.load(Ordering::Acquire) { Err("cancelled".into()) } else { @@ -630,7 +664,7 @@ fn cancelled(shared: &InstallShared) -> Result<(), String> { /// stderr to the app log (so the pipe cannot fill and deadlock). Parks the child /// in `shared` so cancel/shutdown can kill it. Returns an error on a non-zero /// exit, a cancel, or a spawn/wait failure. -fn stream_child( +pub(crate) fn stream_child( shared: &InstallShared, label: &str, mut cmd: Command, @@ -685,7 +719,7 @@ struct SidecarLine { /// A progress sink: `(stage, message, file)`. Injected so the install driver is /// decoupled from `AppHandle` — production wires it to a `model://progress` /// emit; tests record the events while the install actually runs. -type Progress = dyn Fn(&str, Option, Option); +pub(crate) type Progress = dyn Fn(&str, Option, Option); fn install_magenta(progress: &Progress, shared: &InstallShared, name: &str) -> Result<(), String> { progress("download", None, None); @@ -861,17 +895,20 @@ pub fn cancel_install(installer: tauri::State<'_, InstallManager>) { } /// Reveal a family's folder in the OS file manager so the user can inspect or -/// remove models natively (in-app deletion is intentionally absent — moving -/// multi-GB weights to the Trash fails on iCloud-managed / dataless files, and -/// the watcher reflects a native delete live anyway). Magenta opens its models -/// dir; SA3 opens its checkout (or the app-owned SA3 home if not installed yet). -/// Creates the folder if it does not exist. +/// remove models natively (in-app deletion is intentionally absent for the two +/// model families — moving multi-GB weights to the Trash fails on +/// iCloud-managed / dataless files; adapters are small and DO get an in-app +/// delete — and the watcher reflects a native delete live anyway). Magenta +/// opens its models dir; SA3 opens its checkout (or the app-owned SA3 home if +/// not installed yet); LoRA opens the adapter registry. Creates the folder if +/// it does not exist. #[tauri::command] pub fn open_model_folder(app: AppHandle, family: Family) -> Result<(), String> { use tauri_plugin_opener::OpenerExt; let dir = match family { Family::Magenta => magenta_models_dir(), Family::Sa3 => sa3_status().1.unwrap_or_else(sa3_app_home), + Family::Lora => crate::loras::loras_dir(), }; std::fs::create_dir_all(&dir).map_err(|e| format!("cannot create folder: {e}"))?; app.opener() diff --git a/src-tauri/src/watcher.rs b/src-tauri/src/watcher.rs index ce37881..d5485e5 100644 --- a/src-tauri/src/watcher.rs +++ b/src-tauri/src/watcher.rs @@ -143,6 +143,46 @@ pub fn watch_models(app: AppHandle, models_dir: PathBuf) { .expect("failed to spawn lsdj models-watch thread"); } +/// Watch the LoRA adapter registry (issue #66) and emit `models://changed` when +/// the set of complete adapters changes — an import finishing, a hand-drop, or a +/// delete (in-app or native). Same shape as [`watch_models`]: recursive watch, +/// settled-burst re-scan, emit only on a real set change so the importer's +/// staging dir never fires. Best-effort. +pub fn watch_loras(app: AppHandle, loras_dir: PathBuf) { + std::thread::Builder::new() + .name("lsdj-loras-watch".into()) + .spawn(move || { + let _ = std::fs::create_dir_all(&loras_dir); + let (tx, rx) = mpsc::channel::>(); + let mut watcher = match notify::recommended_watcher(tx) { + Ok(w) => w, + Err(e) => { + eprintln!("lsdj-app: loras watcher unavailable ({e}); manager re-lists on open"); + return; + } + }; + if let Err(e) = watcher.watch(&loras_dir, RecursiveMode::Recursive) { + eprintln!("lsdj-app: cannot watch {} ({e})", loras_dir.display()); + return; + } + + let mut last = crate::loras::discover_names(&loras_dir); + loop { + if rx.recv().is_err() { + return; // the watcher (sole sender) was dropped — app shutdown. + } + // Coalesce the burst, then re-scan once it settles. + while rx.recv_timeout(DEBOUNCE).is_ok() {} + let now = crate::loras::discover_names(&loras_dir); + if now != last { + last = now; + let _ = app.emit("models://changed", ()); + } + } + }) + .expect("failed to spawn lsdj loras-watch thread"); +} + /// Classify one FS event: `(touched_songs, touched_samples)`. Only an audio-file /// change in one of the watched folders counts; `registry.json` and any other path /// are ignored. `songs`/`samples` are the candidate dir prefixes (canonical + as From 9f6a64796c6c3f0ff8980aedf50ad202fe297c35 Mon Sep 17 00:00:00 2001 From: Jake Hartnell Date: Fri, 17 Jul 2026 20:01:41 +0000 Subject: [PATCH 3/9] Surface LoRA adapters in the manager and generate UIs (issue #66) A third model-manager section lists adapters by base with import (HF repo id or the native file picker, with an explicit-base override for shape-anonymous -xs adapters) and in-app delete. The media explorer's generate/samples forms and each deck's pad panel gain a base-matched adapter picker + strength control that ride the /api/generate lora field; a stale choice (deleted adapter, base switch) honestly falls back to none. Co-Authored-By: Claude Fable 5 --- frontend/src/audio/nativeEngine.ts | 39 ++++++- frontend/src/deck/DeckColumn.test.tsx | 57 +++++++++- frontend/src/deck/DeckColumn.tsx | 55 +++++++++- frontend/src/deck/useDeck.test.tsx | 20 ++++ frontend/src/deck/useDeck.ts | 19 +++- frontend/src/i18n/en.json | 24 ++++- frontend/src/media/MediaExplorer.test.tsx | 56 ++++++++++ frontend/src/media/MediaExplorer.tsx | 98 ++++++++++++++++- frontend/src/models/ModelManager.test.tsx | 80 +++++++++++++- frontend/src/models/ModelManager.tsx | 126 ++++++++++++++++++++++ frontend/src/models/useLoras.ts | 50 +++++++++ frontend/src/ui/ui.css | 12 +++ 12 files changed, 622 insertions(+), 14 deletions(-) create mode 100644 frontend/src/models/useLoras.ts 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..bf4ebdd 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, null) fireEvent.change(screen.getByLabelText('Engine'), { target: { value: 'music' }, @@ -1381,9 +1390,53 @@ describe('DeckColumn', () => { 'vinyl spinback', 'music', false, + null, ) }) + it('rides the chosen LoRA adapter + strength 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, + }, + // A medium adapter never reaches the pad picker (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 }) + + const picker = screen.getByLabelText('LoRA') + expect( + Array.from(picker.querySelectorAll('option')).map((option) => option.value), + ).toEqual(['none', 'small/maqam']) + fireEvent.change(picker, { target: { value: 'small/maqam' } }) + fireEvent.change(screen.getByLabelText('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, + }) + }) + it('offers Magenta while the deck plays — the third engine is its own worker', () => { const onGenerateToPad = vi.fn() renderPanel( @@ -1397,7 +1450,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, null) }) 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..ecf7bd3 100644 --- a/frontend/src/deck/DeckColumn.tsx +++ b/frontend/src/deck/DeckColumn.tsx @@ -45,6 +45,12 @@ import { type SyncResult, type TrackState, } from './useDeck' +import { + adaptersForKind, + LORA_STRENGTHS, + useLoras, + type LoraChoice, +} from '../models/useLoras' import './deck.css' // The worker holds ~3s of lead (see backend worker pacing); the meter shows @@ -124,7 +130,12 @@ type DeckColumnProps = { 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 + onGenerateToPad: ( + prompt: string, + engine: GenerateEngine, + oneShot: boolean, + lora?: LoraChoice | null, + ) => void generateError: string | null /** Gated tempo readout (M14): null shows an honest dash. */ bpm: number | null @@ -268,6 +279,11 @@ export function DeckColumn({ } | null>(null) const [generateEngine, setGenerateEngine] = useState('sfx') const [generateOneShot, setGenerateOneShot] = useState(true) + // Per-deck LoRA choice (issue #66): both pad kinds ride the small DiTs, so + // one picker covers them; Magenta has no adapter path and hides it. + const loras = useLoras() + const [generateLora, setGenerateLora] = useState('none') + const [generateLoraStrength, setGenerateLoraStrength] = useState(1) 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 +312,19 @@ export function DeckColumn({ connected && Boolean(generateDraft.trim()) && loop.slots.some((slot) => slot.state === 'empty') + // Pads only ever ride the small DiTs; a stale choice (deleted adapter) + // displays — and requests — as none. + const padAdapters = adaptersForKind(loras, 'sfx') + const padLora = padAdapters.some((adapter) => adapter.name === generateLora) + ? generateLora + : 'none' const fireGenerate = () => { if (!canGenerate) return - onGenerateToPad(generateDraft, generateEngine, generateOneShot) + const lora = + generateEngine === 'magenta' || padLora === 'none' + ? null + : { name: padLora, strength: generateLoraStrength } + onGenerateToPad(generateDraft, generateEngine, generateOneShot, lora) } const statusKey = mode === 'playback' @@ -915,6 +941,31 @@ export function DeckColumn({ ]} onChange={(value) => setGenerateOneShot(value === 'oneshot')} /> + {generateEngine !== 'magenta' && padAdapters.length > 0 && ( + ({ + value: String(value), + label: t('media.generate.loraStrengthValue', { value }), + }))} + onChange={(value) => setGenerateLoraStrength(Number(value))} + /> + )}
diff --git a/frontend/src/deck/useDeck.test.tsx b/frontend/src/deck/useDeck.test.tsx index ed1288d..65e837e 100644 --- a/frontend/src/deck/useDeck.test.tsx +++ b/frontend/src/deck/useDeck.test.tsx @@ -976,6 +976,26 @@ describe('useDeck generated pads', () => { expect(result.current.loop.active).toBeNull() }) + it('rides a LoRA adapter + strength 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, + }), + ) + await act(async () => {}) + expect(requestBody(fetchMock)).toEqual({ + prompt: 'oud phrase', + seconds: 4, + kind: 'sfx', + lora: { name: 'small/maqam', strength: 0.75 }, + }) + }) + 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..6a992c3 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,14 @@ 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 + * adapter + strength (issue #66); Magenta has no adapter path. */ + generateToPad: ( + prompt: string, + engine: GenerateEngine, + oneShot: boolean, + lora?: 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 +1425,12 @@ export function useDeck(deckId: DeckId): DeckControls { ) const generateToPad = useCallback( - (prompt: string, engine: GenerateEngine, oneShot: boolean) => { + ( + prompt: string, + engine: GenerateEngine, + oneShot: boolean, + lora?: LoraChoice | null, + ) => { const trimmed = prompt.trim() if (!trimmed) return const current = loopRef.current @@ -1469,6 +1481,7 @@ export function useDeck(deckId: DeckId): DeckControls { prompt: requestPrompt, seconds: requestSeconds, kind: engine, + ...(lora ? { lora } : {}), }, ), }, diff --git a/frontend/src/i18n/en.json b/frontend/src/i18n/en.json index 88119fc..e59b013 100644 --- a/frontend/src/i18n/en.json +++ b/frontend/src/i18n/en.json @@ -72,7 +72,11 @@ "saveFailed": "Couldn't save {{title}} to the songs folder: {{message}}", "openFolderFailed": "Couldn't open the songs folder: {{message}}", "deleteFailed": "Couldn't delete {{title}}: {{message}}", - "empty": "No tracks yet — compose one from a prompt." + "empty": "No tracks yet — compose one from a prompt.", + "lora": "LoRA", + "loraNone": "None", + "loraStrength": "Strength", + "loraStrengthValue": "×{{value}}" }, "samples": { "prompt": "Loop prompt", @@ -232,7 +236,8 @@ "kindOneShot": "One-shot", "kindLoop": "Loop", "action": "Generate", - "failed": "Generation failed: {{message}}" + "failed": "Generation failed: {{message}}", + "lora": "LoRA" }, "play": "Play", "stop": "Stop", @@ -369,7 +374,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", + "loraImportFile": "Import file…", + "loraFileFilter": "LoRA adapter", + "loraDelete": "Delete adapter {{name}}" }, "piano": { "heading": "MIDI keyboard", diff --git a/frontend/src/media/MediaExplorer.test.tsx b/frontend/src/media/MediaExplorer.test.tsx index 1dd9200..ae8a2c4 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,53 @@ describe('MediaExplorer', () => { ).toBe(true) }) + it('rides the chosen LoRA adapter + strength 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, + }, + // Small adapters never reach the track picker (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' })) + const picker = screen.getByLabelText('LoRA') + expect( + Array.from(picker.querySelectorAll('option')).map((option) => option.value), + ).toEqual(['none', 'medium/maqam']) + fireEvent.change(picker, { target: { value: 'medium/maqam' } }) + fireEvent.change(screen.getByLabelText('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', + lora: { name: 'medium/maqam', strength: 1.5 }, + }), + }), + ) + }) + 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..919bb2a 100644 --- a/frontend/src/media/MediaExplorer.tsx +++ b/frontend/src/media/MediaExplorer.tsx @@ -18,6 +18,13 @@ import { import { useInterfaceStore } from '../audio/interfaceStore' import { useControlBus } from '../control/busContext' import { CrateBrowser } from '../crates/CrateBrowser' +import { + adaptersForKind, + LORA_STRENGTHS, + useLoras, + type LoraChoice, +} from '../models/useLoras' +import type { LoraAdapter } from '../audio/nativeEngine' import type { StylePreset } from '../presets' import { Button } from '../ui/Button' import { Panel } from '../ui/Panel' @@ -131,6 +138,20 @@ const ENGINES = Object.keys(ENGINE_LENGTHS) as TrackEngine[] const TRACK_ENGINES: TrackEngine[] = ['track', 'magenta'] const SAMPLE_ENGINES: SampleEngine[] = ['sfx', 'music'] +/** A form's LoRA choice as the `/api/generate` `lora` field, or null when it + * no longer resolves (adapter deleted, or the engine switched to the other + * base) — the request then simply rides without an adapter. */ +function loraChoiceFor( + loras: LoraAdapter[], + kind: 'sfx' | 'music' | 'track', + name: string, + strength: number, +): LoraChoice | null { + if (name === 'none') return null + const matches = adaptersForKind(loras, kind).some((adapter) => adapter.name === name) + return matches ? { name, strength } : null +} + function formatLength(seconds: number): string { return `${Math.floor(seconds / 60)}:${String(seconds % 60).padStart(2, '0')}` } @@ -309,6 +330,15 @@ 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 choice: an + // adapter name or 'none', plus the merge strength. The pickers only offer + // base-matched adapters; a choice that stops resolving (deleted, or the + // engine switched base) falls back to none at request time. + const loras = useLoras() + const [loraName, setLoraName] = useState('none') + const [loraStrength, setLoraStrength] = useState(1) + const [sampleLoraName, setSampleLoraName] = useState('none') + const [sampleLoraStrength, setSampleLoraStrength] = useState(1) // 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 +683,12 @@ export function MediaExplorer({ if (!trimmedPrompt) return const id = nextIdRef.current++ const requestEngine = sampleEngine + const requestLora = loraChoiceFor( + loras, + requestEngine, + sampleLoraName, + sampleLoraStrength, + ) 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 +717,7 @@ export function MediaExplorer({ prompt: trimmedPrompt, seconds: requestSeconds, kind: requestEngine, + ...(requestLora ? { lora: requestLora } : {}), }), }) if (!response.ok) { @@ -754,6 +791,10 @@ export function MediaExplorer({ if (!trimmedPrompt) return const id = nextIdRef.current++ const requestEngine = engine + const requestLora = + requestEngine === 'magenta' + ? null + : loraChoiceFor(loras, requestEngine, loraName, loraStrength) // 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 +816,12 @@ export function MediaExplorer({ body: JSON.stringify( requestEngine === 'magenta' ? { prompt: trimmedPrompt, seconds } - : { prompt: trimmedPrompt, seconds, kind: requestEngine }, + : { + prompt: trimmedPrompt, + seconds, + kind: requestEngine, + ...(requestLora ? { lora: requestLora } : {}), + }, ), }, ) @@ -871,6 +917,24 @@ export function MediaExplorer({ } const lengths = ENGINE_LENGTHS[engine] + // The pickers offer only base-matched adapters; a stale choice (deleted, or + // the engine switched base) displays — and requests — as none. + const trackAdapters = engine === 'magenta' ? [] : adaptersForKind(loras, engine) + const trackLora = trackAdapters.some((adapter) => adapter.name === loraName) + ? loraName + : 'none' + const sampleAdapters = adaptersForKind(loras, sampleEngine) + const sampleLora = sampleAdapters.some((adapter) => adapter.name === sampleLoraName) + ? sampleLoraName + : 'none' + const loraOptions = (adapters: LoraAdapter[]) => [ + { value: 'none', label: t('media.generate.loraNone') }, + ...adapters.map((adapter) => ({ value: adapter.name, label: adapter.slug })), + ] + const strengthOptions = LORA_STRENGTHS.map((value) => ({ + value: String(value), + label: t('media.generate.loraStrengthValue', { value }), + })) function loadButtons(onLoad: (deck: DeckId) => void, name: string) { return (['a', 'b'] as const).map((deck) => ( @@ -1069,6 +1133,22 @@ export function MediaExplorer({ } }} /> + {trackAdapters.length > 0 && ( + setLoraStrength(Number(value))} + /> + )} + )} + {sampleLora !== 'none' && ( + Promise>(asy const updateModel = vi.fn<(family: string) => Promise>(async () => {}) const cancelInstall = vi.fn(async () => {}) const openModelFolder = vi.fn<(family: string) => Promise>(async () => {}) +const installLora = vi.fn<(source: object, base?: string) => Promise>(async () => {}) +const deleteLora = vi.fn<(name: string) => Promise>(async () => {}) +const invoke = vi.fn<(cmd: string, args?: object) => Promise>(async () => null) vi.mock('../audio/nativeEngine', () => ({ modelStatus: () => modelStatus(), @@ -18,6 +21,9 @@ vi.mock('../audio/nativeEngine', () => ({ updateModel: (family: string) => updateModel(family), cancelInstall: () => cancelInstall(), openModelFolder: (family: string) => openModelFolder(family), + installLora: (source: object, base?: string) => installLora(source, base), + deleteLora: (name: string) => deleteLora(name), + invoke: (cmd: string, args?: object) => invoke(cmd, args), subscribeModelsChanged: (cb: () => void) => { changedCb = cb return () => {} @@ -46,6 +52,7 @@ function status(overrides: Partial = {}): ModelStatus { pinnedSource: { repo: 'https://github.com/brxs/stable-audio-3', commit: 'pinned1' }, updateAvailable: false, }, + loras: [], installing: null, ...overrides, } @@ -131,12 +138,14 @@ describe('ModelManager', () => { ) render() await screen.findByText('mrt2_small') - const buttons = screen.getAllByText('Open folder') // Magenta + SA3 (present) - expect(buttons).toHaveLength(2) + const buttons = screen.getAllByText('Open folder') // Magenta + SA3 (present) + LoRA + expect(buttons).toHaveLength(3) fireEvent.click(buttons[0]) expect(openModelFolder).toHaveBeenCalledWith('magenta') fireEvent.click(buttons[1]) expect(openModelFolder).toHaveBeenCalledWith('sa3') + fireEvent.click(buttons[2]) + expect(openModelFolder).toHaveBeenCalledWith('lora') }) it('treats a cancel as a clean stop, not an error', async () => { @@ -201,4 +210,71 @@ describe('ModelManager', () => { ) expect(screen.getByRole('alert')).toHaveTextContent('no weights') }) + + it('lists installed LoRA adapters with base and size, and deletes one', async () => { + modelStatus.mockResolvedValue( + status({ + loras: [ + { + name: 'medium/maqam', + base: 'medium', + slug: 'maqam', + sizeBytes: 200_000_000, + source: 'motiftechnologies/stable-audio-3-maqam-lora', + adapterType: 'lora', + rank: 64, + }, + ], + }), + ) + 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('imports an adapter from a HuggingFace repo id', async () => { + modelStatus.mockResolvedValue(status()) + render() + await screen.findByText('mrt2_base') + fireEvent.change(screen.getByLabelText('HuggingFace repo'), { + target: { value: 'owner/my-lora' }, + }) + // The lora Install is the last one in the DOM (after the model rows). + const installs = screen.getAllByText('Install') + fireEvent.click(installs[installs.length - 1]) + expect(installLora).toHaveBeenCalledWith({ hfRepo: 'owner/my-lora' }, undefined) + }) + + it('passes an explicit base override to a repo import', async () => { + modelStatus.mockResolvedValue(status()) + render() + await screen.findByText('mrt2_base') + fireEvent.change(screen.getByLabelText('HuggingFace repo'), { + target: { value: 'owner/xs-lora' }, + }) + fireEvent.change(screen.getByLabelText('Base'), { target: { value: 'medium' } }) + const installs = screen.getAllByText('Install') + fireEvent.click(installs[installs.length - 1]) + 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('mrt2_base') + 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() }), + ) + }) }) diff --git a/frontend/src/models/ModelManager.tsx b/frontend/src/models/ModelManager.tsx index 5d401d3..39d6e57 100644 --- a/frontend/src/models/ModelManager.tsx +++ b/frontend/src/models/ModelManager.tsx @@ -3,14 +3,20 @@ import type { ReactNode } 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, installModel, + invoke, modelStatus, openModelFolder, subscribeModelProgress, subscribeModelsChanged, updateModel, + type LoraBase, type ModelFamily, type ModelProgress, type ModelStatus, @@ -34,6 +40,10 @@ export function ModelManager() { const [status, setStatus] = useState(null) const [progress, setProgress] = useState(null) const [error, setError] = useState(null) + // The LoRA import controls (issue #66): a HuggingFace repo id draft and the + // base override ('auto' = infer from the adapter's shapes, the normal case). + const [loraRepo, setLoraRepo] = useState('') + const [loraBase, setLoraBase] = useState<'auto' | LoraBase>('auto') const refresh = useCallback(() => { modelStatus().then(setStatus).catch(() => {}) @@ -82,6 +92,43 @@ export function ModelManager() { }) }, []) + const onInstallLora = 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. + setProgress({ family: 'lora', name, stage: 'fetch', message: null, file: null }) + installLora(source, loraBase === 'auto' ? undefined : loraBase).catch((e: unknown) => { + setProgress(null) + setError(String(e)) + }) + }, + [loraBase], + ) + + const onImportLoraFile = 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 + onInstallLora({ path }, name) + })() + }, [onInstallLora, t]) + + const onDeleteLora = 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')}

} @@ -109,6 +156,15 @@ export function ModelManager() { const sa3Present = sa3.state !== 'missing' const sa3Ready = sa3.state === 'ready' const sa3Label = installLabel('sa3', '') + // The in-flight LoRA import's label, keyed by the display name the spec + // derives (the repo id or file name) — live event or status snapshot. + const loraLabel = + progress?.family === 'lora' + ? progressLabel(progress) + : snapshot?.family === 'lora' + ? t('modelManager.installing') + : null + const loraRepoDraft = loraRepo.trim() // Cancel while this row's install is in flight (label set), else its primary // action (Install / Repair, or nothing). @@ -215,6 +271,76 @@ export function ModelManager() {
+ +
+
+

{t('modelManager.loras')}

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

{t('modelManager.loraNone')}

+ )} + {status.loras.map((adapter) => ( +
+
+
{adapter.slug}
+
+ {t(`modelManager.loraBase.${adapter.base}`)} + {` · ${formatBytes(adapter.sizeBytes)}`} +
+
+
+ +
+
+ ))} +
+
+ setLoraRepo(event.target.value)} + onKeyDown={(event) => { + if (event.key === 'Enter' && loraRepoDraft && !isInstalling) { + onInstallLora({ hfRepo: loraRepoDraft }, loraRepoDraft) + } + }} + /> +
+ setBase(value as 'auto' | LoraBase)} + /> + {label ? ( + + ) : ( + <> + + + + )} +
+ {label &&
{label}
} +
+ + ) +} diff --git a/frontend/src/models/ModelManager.test.tsx b/frontend/src/models/ModelManager.test.tsx index a2c4d65..1631058 100644 --- a/frontend/src/models/ModelManager.test.tsx +++ b/frontend/src/models/ModelManager.test.tsx @@ -11,9 +11,6 @@ const installModel = vi.fn<(family: string, name?: string) => Promise>(asy const updateModel = vi.fn<(family: string) => Promise>(async () => {}) const cancelInstall = vi.fn(async () => {}) const openModelFolder = vi.fn<(family: string) => Promise>(async () => {}) -const installLora = vi.fn<(source: object, base?: string) => Promise>(async () => {}) -const deleteLora = vi.fn<(name: string) => Promise>(async () => {}) -const invoke = vi.fn<(cmd: string, args?: object) => Promise>(async () => null) vi.mock('../audio/nativeEngine', () => ({ modelStatus: () => modelStatus(), @@ -21,9 +18,6 @@ vi.mock('../audio/nativeEngine', () => ({ updateModel: (family: string) => updateModel(family), cancelInstall: () => cancelInstall(), openModelFolder: (family: string) => openModelFolder(family), - installLora: (source: object, base?: string) => installLora(source, base), - deleteLora: (name: string) => deleteLora(name), - invoke: (cmd: string, args?: object) => invoke(cmd, args), subscribeModelsChanged: (cb: () => void) => { changedCb = cb return () => {} @@ -138,14 +132,12 @@ describe('ModelManager', () => { ) render() await screen.findByText('mrt2_small') - const buttons = screen.getAllByText('Open folder') // Magenta + SA3 (present) + LoRA - expect(buttons).toHaveLength(3) + const buttons = screen.getAllByText('Open folder') // Magenta + SA3 (present) + expect(buttons).toHaveLength(2) fireEvent.click(buttons[0]) expect(openModelFolder).toHaveBeenCalledWith('magenta') fireEvent.click(buttons[1]) expect(openModelFolder).toHaveBeenCalledWith('sa3') - fireEvent.click(buttons[2]) - expect(openModelFolder).toHaveBeenCalledWith('lora') }) it('treats a cancel as a clean stop, not an error', async () => { @@ -211,87 +203,4 @@ describe('ModelManager', () => { expect(screen.getByRole('alert')).toHaveTextContent('no weights') }) - it('lists installed LoRA adapters with base and size, and deletes one', async () => { - modelStatus.mockResolvedValue( - status({ - loras: [ - { - name: 'medium/maqam', - base: 'medium', - slug: 'maqam', - sizeBytes: 200_000_000, - source: 'motiftechnologies/stable-audio-3-maqam-lora', - adapterType: 'lora', - rank: 64, - }, - ], - }), - ) - 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('imports an adapter from a HuggingFace repo id', async () => { - modelStatus.mockResolvedValue(status()) - render() - await screen.findByText('mrt2_base') - fireEvent.change(screen.getByLabelText('HuggingFace repo'), { - target: { value: 'owner/my-lora' }, - }) - // The lora Install is the last one in the DOM (after the model rows). - const installs = screen.getAllByText('Install') - fireEvent.click(installs[installs.length - 1]) - 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('mrt2_base') - fireEvent.change(screen.getByLabelText('HuggingFace repo'), { - target: { - value: 'https://huggingface.co/motiftechnologies/stable-audio-3-maqam-lora', - }, - }) - const installs = screen.getAllByText('Install') - fireEvent.click(installs[installs.length - 1]) - 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('mrt2_base') - fireEvent.change(screen.getByLabelText('HuggingFace repo'), { - target: { value: 'owner/xs-lora' }, - }) - fireEvent.change(screen.getByLabelText('Base'), { target: { value: 'medium' } }) - const installs = screen.getAllByText('Install') - fireEvent.click(installs[installs.length - 1]) - 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('mrt2_base') - 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() }), - ) - }) }) diff --git a/frontend/src/models/ModelManager.tsx b/frontend/src/models/ModelManager.tsx index 633db36..f759009 100644 --- a/frontend/src/models/ModelManager.tsx +++ b/frontend/src/models/ModelManager.tsx @@ -3,46 +3,19 @@ import type { ReactNode } 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, installModel, - invoke, modelStatus, openModelFolder, subscribeModelProgress, subscribeModelsChanged, updateModel, - type LoraBase, type ModelFamily, type ModelProgress, type ModelStatus, } from '../audio/nativeEngine' - -/** 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() -} - -function formatBytes(bytes: number): string { - if (bytes <= 0) return '0 MB' - const gb = bytes / 1e9 - if (gb >= 1) return `${gb.toFixed(1)} GB` - return `${Math.max(1, Math.round(bytes / 1e6))} MB` -} +import { formatBytes } from './formatBytes' /** The model manager (issue #43): install / on-disk size for both model * families, with live install progress and cancel. Rust owns the lifecycle; this @@ -55,10 +28,6 @@ export function ModelManager() { const [status, setStatus] = useState(null) const [progress, setProgress] = useState(null) const [error, setError] = useState(null) - // The LoRA import controls (issue #66): a HuggingFace repo id draft and the - // base override ('auto' = infer from the adapter's shapes, the normal case). - const [loraRepo, setLoraRepo] = useState('') - const [loraBase, setLoraBase] = useState<'auto' | LoraBase>('auto') const refresh = useCallback(() => { modelStatus().then(setStatus).catch(() => {}) @@ -107,43 +76,6 @@ export function ModelManager() { }) }, []) - const onInstallLora = 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. - setProgress({ family: 'lora', name, stage: 'fetch', message: null, file: null }) - installLora(source, loraBase === 'auto' ? undefined : loraBase).catch((e: unknown) => { - setProgress(null) - setError(String(e)) - }) - }, - [loraBase], - ) - - const onImportLoraFile = 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 - onInstallLora({ path }, name) - })() - }, [onInstallLora, t]) - - const onDeleteLora = 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')}

} @@ -171,15 +103,6 @@ export function ModelManager() { const sa3Present = sa3.state !== 'missing' const sa3Ready = sa3.state === 'ready' const sa3Label = installLabel('sa3', '') - // The in-flight LoRA import's label, keyed by the display name the spec - // derives (the repo id or file name) — live event or status snapshot. - const loraLabel = - progress?.family === 'lora' - ? progressLabel(progress) - : snapshot?.family === 'lora' - ? t('modelManager.installing') - : null - const loraRepoDraft = normalizeHfRepo(loraRepo) // Cancel while this row's install is in flight (label set), else its primary // action (Install / Repair, or nothing). @@ -286,76 +209,6 @@ export function ModelManager() { - -
-
-

{t('modelManager.loras')}

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

{t('modelManager.loraNone')}

- )} - {status.loras.map((adapter) => ( -
-
-
{adapter.slug}
-
- {t(`modelManager.loraBase.${adapter.base}`)} - {` · ${formatBytes(adapter.sizeBytes)}`} -
-
-
- -
-
- ))} -
-
- setLoraRepo(event.target.value)} - onKeyDown={(event) => { - if (event.key === 'Enter' && loraRepoDraft && !isInstalling) { - onInstallLora({ hfRepo: loraRepoDraft }, loraRepoDraft) - } - }} - /> -
- ({ - value: adapter.name, - label: adapter.slug, - })), - ]} - onChange={setGenerateLora} - /> - )} - {generateEngine !== 'magenta' && padLora !== 'none' && ( - - )} - {trackLora !== 'none' && ( -
+ {trackAdapters.length > 0 && ( + ({ + 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')}

) : ( @@ -1298,22 +1260,6 @@ export function MediaExplorer({ } }} /> - {sampleAdapters.length > 0 && ( - setSampleLoraStrength(Number(value))} - /> - )}