From f7ffed87dca1e098f3d8dcb3c621fc2422212c53 Mon Sep 17 00:00:00 2001 From: RyanOnTheInside <7623207+ryanontheinside@users.noreply.github.com> Date: Wed, 22 Jul 2026 17:36:43 -0400 Subject: [PATCH 1/6] feat(lora): format sniffing, family compatibility, stem-collision handling (SA3 LoRA phase 0) - lora_metadata: header-only safetensors sniff -> lora_family (ace/sa3), SA3 rank/adapter_type/base_model synthesis from the embedded lora_config, composite cache key (weights + sidecar + trigger mtimes), canonical_sa3_lineage for trainer/HF/runtime checkpoint spellings - engine/lora: _compute_deltas raises on zero PEFT-pattern matches so a wrong-family enable fails loudly instead of silently no-opping; register_lora accepts an explicit lora_id - ace_backend: lora_compatible gains the family axis (sa3 = hard no, unknown stays permissive) - paths: assign_lora_ids disambiguates duplicate-stem discoveries with parent-dir suffixes; register_library and /api/loras share it so the HTTP and engine catalogs agree on ids --- acestep/engine/lora.py | 37 ++- acestep/lora_metadata.py | 261 +++++++++++++++++++--- acestep/paths.py | 43 ++++ acestep/streaming/ace_backend.py | 16 +- demos/realtime_motion_graph_web/server.py | 17 +- tests/unit/test_lora_metadata.py | 236 +++++++++++++++++++ tests/unit/test_lora_resolution.py | 212 ++++++++++++++++++ 7 files changed, 773 insertions(+), 49 deletions(-) diff --git a/acestep/engine/lora.py b/acestep/engine/lora.py index f1cac724..a9b0d786 100644 --- a/acestep/engine/lora.py +++ b/acestep/engine/lora.py @@ -156,6 +156,7 @@ def _make_id(path: str) -> str: def register_lora( self, path: str, name: Optional[str] = None, + lora_id: Optional[str] = None, ) -> str: """Add a LoRA to the catalog without materializing deltas. @@ -163,8 +164,14 @@ def register_lora( the existing id and leaves any in-flight prewarm / enabled state alone. The existing entry's name is NOT overwritten on re-register; pass an explicit ``name`` only on first registration. + + ``lora_id`` overrides the default filename-stem id; the library + scan passes pre-disambiguated ids (see + :func:`acestep.paths.assign_lora_ids`) so same-stem files across + roots don't silently shadow each other. """ - lora_id = self._make_id(path) + if lora_id is None: + lora_id = self._make_id(path) if lora_id in self._loras: existing = self._loras[lora_id] if existing.path != str(path): @@ -197,15 +204,15 @@ def register_library( first, then each extra dir; each root sorted by filename). Missing directories are silently skipped. """ - from acestep.paths import discover_all_loras, discover_loras + from acestep.paths import assign_lora_ids, discover_all_loras, discover_loras if directory is None: files = discover_all_loras() else: files = discover_loras(directory) ids: List[str] = [] - for p in files: + for p, lora_id in assign_lora_ids(files): try: - ids.append(self.register_lora(str(p))) + ids.append(self.register_lora(str(p), lora_id=lora_id)) except Exception as e: logger.warning("Failed to register {}: {}", p, e) if files: @@ -301,6 +308,28 @@ def _compute_deltas( param_name = parts.replace(".lora_B.weight", ".weight") pairs.setdefault(param_name, {})["B"] = tensor + # Zero PEFT-pattern matches means this is not an ACE-Step LoRA + # at all. Without this check the pair map is empty, the + # shape-mismatch guard below never trips (it only fires on + # matched pairs), and enable_lora "succeeds" having applied + # nothing — a silent no-op. Wrong-family files must fail loudly. + if not pairs: + sa3_like = any( + ".parametrizations.weight." in k for k in raw + ) + hint = ( + " The key layout matches the SA3 (stable-audio-3) " + "parametrization format; SA3-family LoRAs cannot load " + "on an ACE-Step engine." + if sa3_like else "" + ) + raise RuntimeError( + f"LoRA {Path(lora_path).name} is not an ACE-Step LoRA: " + f"none of its {len(raw)} tensors match the PEFT " + f"lora_A.weight/lora_B.weight key pattern, so there is " + f"nothing to apply.{hint}" + ) + device = self._delta_compute_device() deltas: Dict[str, torch.Tensor] = {} total_bytes = 0 diff --git a/acestep/lora_metadata.py b/acestep/lora_metadata.py index e5ffa859..a9771030 100644 --- a/acestep/lora_metadata.py +++ b/acestep/lora_metadata.py @@ -29,8 +29,18 @@ warning and falls back to the ``.trigger.txt`` path, so a broken sidecar never takes the WS catalog broadcast down with it. -3. A small ``(path, mtime_ns)`` memoization layer so a catalog refresh - over ~30 LoRAs costs one ``stat`` per entry, not a JSON parse. +3. A small memoization layer keyed on the mtimes of every file that + feeds a record (weights header, ``metadata.json``, ``.trigger.txt``) + so a catalog refresh over ~30 LoRAs costs a few ``stat`` calls per + entry, not a JSON parse or a header read. + +It also owns the *format* axis of LoRA identity: a header-only sniff of +the ``.safetensors`` key layout classifies each file into a +``lora_family`` ("ace" for PEFT-style ``lora_A.weight``/``lora_B.weight`` +pairs, "sa3" for stable-audio-3 ``.parametrizations.weight..`` +entries), and for SA3 files the embedded ``lora_config`` header (written +by the trainer) supplies rank / adapter_type / base_model when no +sidecar documents them. No tensor data is ever read here. """ from __future__ import annotations @@ -83,6 +93,17 @@ class LoraMetadata: # undocumented LoRAs. base_model: Optional[str] = None base_model_scale: Optional[str] = None + # Weight-format family sniffed from the safetensors header key + # layout: "ace" (PEFT ``lora_A.weight``/``lora_B.weight`` pairs), + # "sa3" (torch-parametrize ``.parametrizations.weight..`` + # entries), or ``None`` (unknown / unreadable header). Code-derived + # from the weights file itself — a sidecar cannot override it. + lora_family: Optional[str] = None + # SA3-only fields synthesized from the embedded ``lora_config`` + # header (the trainer writes rank/alpha/adapter_type/base_model + # there). ``None`` on ACE files and on SA3 files without a config. + adapter_type: Optional[str] = None + rank: Optional[int] = None # True iff a valid metadata.json was loaded for this record. Lets # callers distinguish "rich metadata" from "synthesized fallback" # without inspecting individual field nullity. @@ -114,35 +135,183 @@ def lora_scale_compatible( return lora_scale == checkpoint_scale -# (sidecar_path_str, mtime_ns) -> parsed record -_cache: dict[tuple[str, int], LoraMetadata] = {} +# Safetensors headers are capped at 100 MB by the format spec; anything +# claiming more is corrupt (or not a safetensors file at all). +_MAX_HEADER_BYTES = 100_000_000 + +_ACE_KEY_MARKERS = (".lora_A.weight", ".lora_B.weight") +_SA3_KEY_MARKER = ".parametrizations.weight." + + +def read_safetensors_header(path: Path | str) -> Optional[dict[str, Any]]: + """Read just the JSON header of a ``.safetensors`` file (8-byte + little-endian length prefix + JSON blob). No tensor data is read, + so this is safe on multi-GB files and on every catalog refresh. + + Returns ``None`` on any IO/parse problem — sniffing must never take + the catalog broadcast down with it. + """ + try: + with open(path, "rb") as f: + prefix = f.read(8) + if len(prefix) < 8: + return None + n = int.from_bytes(prefix, "little") + if n <= 0 or n > _MAX_HEADER_BYTES: + return None + raw = f.read(n) + header = json.loads(raw) + except (OSError, ValueError, UnicodeDecodeError): + return None + return header if isinstance(header, dict) else None + + +def _family_from_keys(keys: list[str]) -> Optional[str]: + """Classify a LoRA's weight-format family from its tensor key names. + + SA3 keys look like + ``model.…​.parametrizations.weight.0.lora_A`` (torch parametrize, + no ``.weight`` suffix on the adapter tensors); ACE/PEFT keys look + like ``base_model.model.….lora_A.weight``. The two markers cannot + co-occur on one key, and a file mixing both families doesn't exist + in the wild — SA3 wins the tie because its marker is the more + specific one. + """ + saw_ace = False + for k in keys: + if _SA3_KEY_MARKER in k: + return "sa3" + if not saw_ace and any(m in k for m in _ACE_KEY_MARKERS): + saw_ace = True + return "ace" if saw_ace else None + + +def sniff_lora_file(lora_path: Path | str) -> tuple[Optional[str], dict[str, Any]]: + """Header-only sniff of a LoRA ``.safetensors``. + + Returns ``(family, config)`` where ``family`` is "ace" / "sa3" / + ``None`` (see :func:`_family_from_keys`) and ``config`` is the + parsed ``lora_config`` JSON embedded in the header's + ``__metadata__`` block (SA3 trainer convention: rank, alpha, + adapter_type, include/exclude, base_model, …). ``config`` is ``{}`` + when absent or malformed. For SA3 files without an explicit rank, + the rank is inferred from the first adapter tensor's shape (same + rule as upstream's ``infer_global_rank``). + """ + header = read_safetensors_header(lora_path) + if not header: + return None, {} + keys = [k for k in header if k != "__metadata__"] + family = _family_from_keys(keys) + + cfg: dict[str, Any] = {} + meta = header.get("__metadata__") + if isinstance(meta, dict): + raw_cfg = meta.get("lora_config") + if isinstance(raw_cfg, str) and raw_cfg: + try: + parsed = json.loads(raw_cfg) + except json.JSONDecodeError as exc: + logger.warning( + "LoRA %s has a malformed embedded lora_config (%s); " + "ignoring it", + lora_path, + exc, + ) + else: + if isinstance(parsed, dict): + cfg = parsed + + if family == "sa3" and "rank" not in cfg: + for k in keys: + spec = header.get(k) + shape = spec.get("shape") if isinstance(spec, dict) else None + if not shape: + continue + if k.endswith(".lora_A") or k.endswith(".M_xs"): + cfg["rank"] = int(shape[0]) + break + if k.endswith(".lora_B") and len(shape) > 1: + cfg["rank"] = int(shape[1]) + break + return family, cfg + + +# Runtime SA3 checkpoint ids as the loaders know them (SA3Context +# model_id). Training happens against the "-base" variants; both sides +# canonicalize to these. +_SA3_RUNTIME_LINEAGES = frozenset({"medium", "small-music", "small-sfx"}) + + +def canonical_sa3_lineage(value: Optional[str]) -> Optional[str]: + """Map an SA3 base-model identifier to its runtime checkpoint id + ("medium" / "small-music" / "small-sfx"). + + Accepts the trainer spellings ("medium-base"), the runtime ids + themselves, and full HF repo ids (with or without the org prefix: + "stabilityai/stable-audio-3-medium-base"). Case-insensitive. + Returns ``None`` for anything unrecognized — callers treat ``None`` + as "unknown lineage, don't filter" per the permissive seam + contract, so a new checkpoint spelling degrades to permissive + rather than hiding LoRAs. + """ + if not value: + return None + s = str(value).strip().lower() + s = s.rsplit("/", 1)[-1] + if s.startswith("stable-audio-3-"): + s = s[len("stable-audio-3-"):] + if s.endswith("-base"): + s = s[: -len("-base")] + return s if s in _SA3_RUNTIME_LINEAGES else None + + +# (lora_path_str, weights_mtime_ns, sidecar_mtime_ns, trigger_mtime_ns) +# -> parsed record. -1 stands for "file absent" so a sidecar/trigger +# appearing (or the weights file being replaced) changes the key. +_cache: dict[tuple[str, int, int, int], LoraMetadata] = {} + + +def _mtime_ns_or_absent(p: Path) -> int: + try: + return p.stat().st_mtime_ns + except OSError: + return -1 def load_lora_metadata(lora_path: Path | str) -> LoraMetadata: """Load metadata for a LoRA ``.safetensors`` at ``lora_path``. Returns a normalized :class:`LoraMetadata` covering the three input - states documented at the module level. Never raises on malformed - sidecars — falls back through ``metadata.json`` → ``.trigger.txt`` - → bare in that order, logging warnings as it goes. + states documented at the module level, plus the header-derived + format fields (``lora_family``, and for SA3 files + ``adapter_type``/``rank``/synthesized ``base_model``). Never raises + on malformed sidecars — falls back through ``metadata.json`` → + ``.trigger.txt`` → bare in that order, logging warnings as it goes. + A present sidecar wins for the display/trigger/genre fields; the + weight-format family always comes from the header sniff. """ p = Path(lora_path) stem = p.stem sidecar = _metadata_sidecar(p) - # Cache by mtime_ns when the sidecar exists. We don't cache the - # "no metadata.json" path because there's nothing to invalidate - # against (a file appearing wouldn't bust a cache keyed on nothing). - try: - st = sidecar.stat() - cache_key: Optional[tuple[str, int]] = (str(sidecar), st.st_mtime_ns) - except OSError: - cache_key = None + weights_mtime = _mtime_ns_or_absent(p) + sidecar_mtime = _mtime_ns_or_absent(sidecar) + trigger_mtime = _mtime_ns_or_absent(lora_sidecar(p, ".trigger.txt")) - if cache_key is not None and cache_key in _cache: - return _cache[cache_key] + cache_key: Optional[tuple[str, int, int, int]] = None + if weights_mtime != -1 or sidecar_mtime != -1 or trigger_mtime != -1: + cache_key = (str(p), weights_mtime, sidecar_mtime, trigger_mtime) + if cache_key in _cache: + return _cache[cache_key] - if cache_key is not None: + family: Optional[str] = None + header_cfg: dict[str, Any] = {} + if weights_mtime != -1: + family, header_cfg = sniff_lora_file(p) + + md: Optional[LoraMetadata] = None + if sidecar_mtime != -1: try: raw = json.loads(sidecar.read_text(encoding="utf-8")) except (OSError, UnicodeDecodeError, json.JSONDecodeError) as exc: @@ -153,21 +322,49 @@ def load_lora_metadata(lora_path: Path | str) -> LoraMetadata: ) else: md = _from_schema(raw, stem) - _cache[cache_key] = md - return md - - # Legacy .trigger.txt fallback. lora_trigger() already returns "" - # on miss/IO error so this is a single string check. - legacy = lora_trigger(p) - if legacy: - return LoraMetadata( - id=stem, - name=stem, - primary_trigger_word=legacy, - trigger_words=[legacy], - ) - return LoraMetadata(id=stem, name=stem) + if md is None: + # Legacy .trigger.txt fallback. lora_trigger() already returns "" + # on miss/IO error so this is a single string check. + legacy = lora_trigger(p) + if legacy: + md = LoraMetadata( + id=stem, + name=stem, + primary_trigger_word=legacy, + trigger_words=[legacy], + ) + else: + md = LoraMetadata(id=stem, name=stem) + + _apply_header_fields(md, family, header_cfg) + + if cache_key is not None: + _cache[cache_key] = md + return md + + +def _apply_header_fields( + md: LoraMetadata, family: Optional[str], cfg: dict[str, Any], +) -> None: + """Overlay header-derived fields onto a record. + + ``lora_family`` is always header-truth (a sidecar cannot claim a + different weight format than the file actually has). The SA3 + ``lora_config`` fields only fill gaps the sidecar left: base_model + from a sidecar wins over the trainer's embedded value. + """ + md.lora_family = family + if family != "sa3" or not cfg: + return + at = cfg.get("adapter_type") + if isinstance(at, str) and at: + md.adapter_type = at + md.rank = _optional_int(cfg.get("rank")) + if md.base_model is None: + bm = cfg.get("base_model") + if isinstance(bm, str) and bm: + md.base_model = bm def _metadata_sidecar(lora_path: Path) -> Path: diff --git a/acestep/paths.py b/acestep/paths.py index 3869e2c4..59fc3551 100644 --- a/acestep/paths.py +++ b/acestep/paths.py @@ -241,6 +241,49 @@ def discover_all_loras() -> list[Path]: return out +def assign_lora_ids(paths: list[Path]) -> list[tuple[Path, str]]: + """Assign a unique catalog id to each discovered LoRA path. + + Ids are filename stems (the wire-stable convention the whole LoRA + surface is keyed on). When two files in one scan share a stem, the + FIRST occurrence in scan order keeps the bare stem — preserving the + id the first-wins registrar historically kept, so operator configs + referencing it don't break — and every later occurrence is + disambiguated with its parent directory name (``--``), + falling back to a numeric suffix if that still collides. Collisions + are logged loudly; a collision must never silently shadow a + loadable file. + + Shared by the engine's ``register_library()`` scan and the demo + server's ``/api/loras`` listing so the HTTP catalog and the + engine-side catalog agree on ids. + """ + used: set[str] = set() + out: list[tuple[Path, str]] = [] + collided: list[tuple[Path, str]] = [] + for p in paths: + stem = p.stem + if stem not in used: + lora_id = stem + else: + candidate = f"{stem}--{p.parent.name}" if p.parent.name else stem + lora_id = candidate + n = 2 + while lora_id in used or lora_id == stem: + lora_id = f"{candidate}-{n}" + n += 1 + collided.append((p, lora_id)) + used.add(lora_id) + out.append((p, lora_id)) + if collided: + print( + "[paths] LoRA filename-stem collisions detected; later files " + "disambiguated so none is shadowed: " + + ", ".join(f"{lid!r} <- {p}" for p, lid in collided) + ) + return out + + def lora_sidecar(lora_path: Path | str, suffix: str) -> Path: """Resolve a LoRA weight file to a sibling sidecar with ``suffix``. diff --git a/acestep/streaming/ace_backend.py b/acestep/streaming/ace_backend.py index e03e9cc6..c1756646 100644 --- a/acestep/streaming/ace_backend.py +++ b/acestep/streaming/ace_backend.py @@ -347,9 +347,19 @@ def capabilities(self) -> Capabilities: ) def lora_compatible(self, metadata: dict) -> bool: - """ACE's only compatibility axis today: a 2B-trained LoRA - cannot refit onto the XL (5B) DiT and vice versa. Unknown - scale on either side is compatible per the seam contract.""" + """ACE's compatibility axes: weight-format family and + base-model scale. + + Family: a file whose sniffed ``lora_family`` names a different + family (e.g. "sa3") can never load on an ACE engine — hard no. + Unknown family stays permissive per the seam contract. + + Scale: a 2B-trained LoRA cannot refit onto the XL (5B) DiT and + vice versa. Unknown scale on either side is compatible. + """ + family = metadata.get("lora_family") + if family and family != "ace": + return False return lora_scale_compatible( metadata.get("base_model_scale"), self.checkpoint_scale, ) diff --git a/demos/realtime_motion_graph_web/server.py b/demos/realtime_motion_graph_web/server.py index c53271c8..302d1dba 100644 --- a/demos/realtime_motion_graph_web/server.py +++ b/demos/realtime_motion_graph_web/server.py @@ -286,6 +286,7 @@ def _process_request(connection, request): if path_only == "/api/loras": from acestep.lora_metadata import load_lora_metadata, lora_scale_compatible from acestep.paths import ( + assign_lora_ids, checkpoint_scale, discover_all_loras, extra_lora_dirs, @@ -297,19 +298,15 @@ def _process_request(connection, request): # register_library() scan so the HTTP catalog and the # engine-side catalog stay in lockstep. entries = [] - seen_ids: set[str] = set() scale = checkpoint_scale(_CHECKPOINT) - for p in discover_all_loras(): - # Same-stem dedup mirrors LoRAManager.register_lora's - # first-wins behavior so the UI can't see a phantom id - # the engine refused to register. - if p.stem in seen_ids: - continue - seen_ids.add(p.stem) + for p, lora_id in assign_lora_ids(discover_all_loras()): + # Same stem-collision disambiguation as the engine's + # register_library() scan so the UI sees exactly the + # ids the engine registered. md = load_lora_metadata(p).to_wire() entries.append({ - "id": p.stem, - "name": md.get("name") or p.stem, + "id": lora_id, + "name": md.get("name") or lora_id, "path": str(p), "state": "registered", "strength": 0.0, diff --git a/tests/unit/test_lora_metadata.py b/tests/unit/test_lora_metadata.py index 668b70a1..37068950 100644 --- a/tests/unit/test_lora_metadata.py +++ b/tests/unit/test_lora_metadata.py @@ -20,8 +20,10 @@ from acestep.lora_metadata import ( CURRENT_SCHEMA_VERSION, LoraMetadata, + canonical_sa3_lineage, clear_cache, load_lora_metadata, + sniff_lora_file, ) @@ -311,3 +313,237 @@ def test_empty_strings_filtered_from_trigger_lists(tmp_path): assert md.trigger_words == ["good"] assert md.tags == ["a"] + + +# --------------------------------------------------------------------------- +# Header sniffing: lora_family + SA3 lora_config synthesis +# --------------------------------------------------------------------------- +# +# These build REAL (tiny) safetensors files because the sniff reads the +# actual on-disk header; the zero-byte _weights stand-in above stays +# valid for the sidecar-only tests (an unreadable header degrades to +# family=None, never raises). + +import os +import time + +import torch +from safetensors.torch import save_file + + +def _bump_mtime(p: Path) -> None: + """Force a distinct mtime_ns so back-to-back writes can't land in + the same filesystem timestamp granule.""" + t = time.time() + 1.0 + os.utime(p, (t, t)) + + +def _write_sa3_lora( + tmp_path: Path, + stem: str = "sa3lora", + config: dict | None = None, + omit_config: bool = False, +) -> Path: + """Tiny SA3-format LoRA: parametrize-style keys at index 0, fp16, + lora_config JSON embedded in the safetensors header — exactly the + layout the vendored stable-audio-3 trainer saves.""" + prefix = "model.transformer.layers.0.self_attn.to_qkv.parametrizations.weight.0" + tensors = { + f"{prefix}.lora_A": torch.zeros(4, 16, dtype=torch.float16), + f"{prefix}.lora_B": torch.zeros(16, 4, dtype=torch.float16), + } + if config is None: + config = { + "rank": 4, + "alpha": 4, + "adapter_type": "lora", + "base_model": "medium-base", + } + metadata = None if omit_config else {"lora_config": json.dumps(config)} + p = tmp_path / f"{stem}.safetensors" + save_file(tensors, str(p), metadata=metadata) + return p + + +def _write_ace_lora(tmp_path: Path, stem: str = "acelora") -> Path: + """Tiny ACE/PEFT-format LoRA: lora_A.weight / lora_B.weight pairs.""" + tensors = { + "base_model.model.q.lora_A.weight": torch.zeros(4, 16), + "base_model.model.q.lora_B.weight": torch.zeros(8, 4), + } + p = tmp_path / f"{stem}.safetensors" + save_file(tensors, str(p)) + return p + + +def test_sniff_sa3_family_and_config(tmp_path): + p = _write_sa3_lora(tmp_path) + family, cfg = sniff_lora_file(p) + assert family == "sa3" + assert cfg["rank"] == 4 + assert cfg["adapter_type"] == "lora" + assert cfg["base_model"] == "medium-base" + + md = load_lora_metadata(p) + assert md.lora_family == "sa3" + assert md.adapter_type == "lora" + assert md.rank == 4 + assert md.base_model == "medium-base" # synthesized: no sidecar + assert md.has_metadata is False + + +def test_sniff_sa3_rank_inferred_without_config(tmp_path): + """No embedded lora_config: family still detected from key shapes + and rank inferred from the first adapter tensor (upstream's + infer_global_rank rule: lora_A shape[0]).""" + p = _write_sa3_lora(tmp_path, omit_config=True) + family, cfg = sniff_lora_file(p) + assert family == "sa3" + assert cfg == {"rank": 4} + + md = load_lora_metadata(p) + assert md.lora_family == "sa3" + assert md.rank == 4 + assert md.adapter_type is None + assert md.base_model is None + + +def test_sniff_ace_family(tmp_path): + p = _write_ace_lora(tmp_path) + family, cfg = sniff_lora_file(p) + assert family == "ace" + assert cfg == {} + + md = load_lora_metadata(p) + assert md.lora_family == "ace" + assert md.adapter_type is None + assert md.rank is None + + +def test_sniff_non_lora_safetensors_is_family_none(tmp_path): + p = tmp_path / "weights.safetensors" + save_file({"encoder.weight": torch.zeros(2, 2)}, str(p)) + family, cfg = sniff_lora_file(p) + assert family is None + assert cfg == {} + assert load_lora_metadata(p).lora_family is None + + +def test_sniff_unreadable_header_degrades_to_none(tmp_path): + """Zero-byte and garbage files must never raise — the sniff runs + on every catalog refresh.""" + empty = tmp_path / "empty.safetensors" + empty.write_bytes(b"") + garbage = tmp_path / "garbage.safetensors" + garbage.write_bytes(b"\xff" * 64) + + assert sniff_lora_file(empty) == (None, {}) + assert sniff_lora_file(garbage) == (None, {}) + assert load_lora_metadata(empty).lora_family is None + assert load_lora_metadata(garbage).lora_family is None + + +def test_sidecar_wins_display_fields_family_stays_header_truth(tmp_path): + """A sidecar owns display/trigger/base_model fields, but the + weight-format family always comes from the header — a sidecar + cannot claim a different format than the file actually has.""" + p = _write_sa3_lora(tmp_path, stem="styled") + raw = _full_metadata("styled") + raw["model"]["base_model"] = "sidecar-says-this" + (tmp_path / "styled.metadata.json").write_text( + json.dumps(raw), encoding="utf-8" + ) + + md = load_lora_metadata(p) + + assert md.has_metadata is True + assert md.name == "Industrial" # sidecar display name wins + assert md.base_model == "sidecar-says-this" # sidecar wins over header + assert md.lora_family == "sa3" # header truth + assert md.adapter_type == "lora" + assert md.rank == 4 + + +def test_cache_invalidates_on_weights_replacement(tmp_path): + """Replacing the weights file (same path, new content + mtime) must + re-sniff: header-derived fields would otherwise go stale. This is + exactly what the old sidecar-only cache key got wrong.""" + p = _write_ace_lora(tmp_path, stem="swapme") + assert load_lora_metadata(p).lora_family == "ace" + + _write_sa3_lora(tmp_path, stem="swapme") + _bump_mtime(p) + + assert load_lora_metadata(p).lora_family == "sa3" + + +def test_cache_invalidates_when_sidecar_appears(tmp_path): + """Bare records are now cached (keyed on the weights mtime), so a + sidecar appearing later must change the key and be picked up.""" + p = _write_ace_lora(tmp_path, stem="latecomer") + md1 = load_lora_metadata(p) + assert md1.has_metadata is False + + sidecar = tmp_path / "latecomer.metadata.json" + sidecar.write_text(json.dumps(_full_metadata("latecomer")), encoding="utf-8") + _bump_mtime(sidecar) + + md2 = load_lora_metadata(p) + assert md2.has_metadata is True + assert md2.name == "Industrial" + assert md2.lora_family == "ace" + + +def test_cache_invalidates_when_trigger_txt_appears(tmp_path): + """Same for the legacy .trigger.txt sidecar: the old code never + cached bare records so a trigger appearing was picked up instantly; + the composite key must preserve that behavior.""" + p = _write_ace_lora(tmp_path, stem="trig") + assert load_lora_metadata(p).primary_trigger_word is None + + trigger = tmp_path / "trig.trigger.txt" + trigger.write_text("late-word", encoding="utf-8") + _bump_mtime(trigger) + + assert load_lora_metadata(p).primary_trigger_word == "late-word" + + +def test_to_wire_carries_family_fields(tmp_path): + wire = load_lora_metadata(_write_sa3_lora(tmp_path)).to_wire() + assert wire["lora_family"] == "sa3" + assert wire["adapter_type"] == "lora" + assert wire["rank"] == 4 + + +# --------------------------------------------------------------------------- +# canonical_sa3_lineage +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "spelling,expected", + [ + # runtime ids pass through + ("medium", "medium"), + ("small-music", "small-music"), + ("small-sfx", "small-sfx"), + # trainer "-base" spellings + ("medium-base", "medium"), + ("small-music-base", "small-music"), + ("small-sfx-base", "small-sfx"), + # full HF repo ids, with and without the org prefix + ("stabilityai/stable-audio-3-medium", "medium"), + ("stabilityai/stable-audio-3-medium-base", "medium"), + ("stable-audio-3-small-music-base", "small-music"), + # case-insensitive, whitespace-tolerant + ("Medium-Base", "medium"), + (" small-music ", "small-music"), + # unrecognized -> None (permissive downstream) + ("acestep-v15-turbo", None), + ("large", None), + ("", None), + (None, None), + ], +) +def test_canonical_sa3_lineage(spelling, expected): + assert canonical_sa3_lineage(spelling) == expected diff --git a/tests/unit/test_lora_resolution.py b/tests/unit/test_lora_resolution.py index e42c3aad..29a4fb42 100644 --- a/tests/unit/test_lora_resolution.py +++ b/tests/unit/test_lora_resolution.py @@ -191,3 +191,215 @@ def test_resolve_lora_id_without_engine_is_identity(self, tmp_path): ss = _stub_session(tmp_path) ss.engine_obj = None assert ss._resolve_lora_id("Ambient") == "Ambient" + + +# --------------------------------------------------------------------------- +# Family axis: ACE backend predicate + mixed-family catalog annotation +# --------------------------------------------------------------------------- + +import torch +import torch.nn as nn +from safetensors.torch import save_file + +from acestep.engine.lora import EagerLoRAManager +from acestep.streaming.ace_backend import ACEStepBackend + + +def _write_sa3_file(tmp_path: Path, stem: str = "sa3style") -> Path: + """Tiny SA3-format LoRA file (parametrize-style keys), real bytes so + the metadata sniff classifies it.""" + prefix = ( + "model.transformer.layers.0.self_attn.to_qkv" + ".parametrizations.weight.0" + ) + p = tmp_path / f"{stem}.safetensors" + save_file( + { + f"{prefix}.lora_A": torch.zeros(4, 16, dtype=torch.float16), + f"{prefix}.lora_B": torch.zeros(16, 4, dtype=torch.float16), + }, + str(p), + metadata={ + "lora_config": json.dumps( + {"rank": 4, "alpha": 4, "adapter_type": "lora", + "base_model": "medium-base"} + ) + }, + ) + return p + + +class _AceStub: + """Bare-attribute stand-in so ACEStepBackend.lora_compatible can run + unbound without constructing the full backend.""" + + def __init__(self, scale="2B"): + self.checkpoint_scale = scale + + +class TestAceFamilyCompatible: + def test_sa3_family_is_hard_no(self): + assert not ACEStepBackend.lora_compatible( + _AceStub(), {"lora_family": "sa3"}, + ) + + def test_sa3_family_is_hard_no_even_with_matching_scale(self): + # The family axis is checked before scale: an SA3 file whose + # sidecar happens to claim a matching scale still can't load. + assert not ACEStepBackend.lora_compatible( + _AceStub("2B"), + {"lora_family": "sa3", "base_model_scale": "2B"}, + ) + + def test_ace_family_falls_through_to_scale_axis(self): + assert ACEStepBackend.lora_compatible( + _AceStub("2B"), {"lora_family": "ace", "base_model_scale": "2B"}, + ) + assert not ACEStepBackend.lora_compatible( + _AceStub("2B"), {"lora_family": "ace", "base_model_scale": "5B"}, + ) + + def test_unknown_family_stays_permissive(self): + assert ACEStepBackend.lora_compatible( + _AceStub("2B"), {"lora_family": None, "base_model_scale": None}, + ) + assert ACEStepBackend.lora_compatible(_AceStub("2B"), {}) + + +class _FamilyBackend: + """Stub backend running the REAL ACE predicate (family + scale).""" + + checkpoint_scale = "2B" + lora_compatible = ACEStepBackend.lora_compatible + + +class TestMixedFamilyCatalog: + def test_catalog_annotates_sa3_entry_incompatible(self, tmp_path): + clear_cache() + ss = _StubSession([ + _Desc(_write_lora(tmp_path, "ambient-v1", "Ambient", "2B")), + _Desc(_write_sa3_file(tmp_path, "sa3style")), + ]) + ss.backend = _FamilyBackend() + verdicts = { + e["id"]: e["compatible"] for e in ss.lora_catalog_payload() + } + assert verdicts == {"ambient-v1": True, "sa3style": False} + + def test_alias_of_sa3_only_name_misses(self, tmp_path): + """Alias resolution is restricted to the compatible subset, so a + display-name reference to an SA3-only entry misses; the exact id + still passes through (and fails loudly at the enable boundary).""" + clear_cache() + ss = _StubSession([_Desc(_write_sa3_file(tmp_path, "sa3style"))]) + ss.backend = _FamilyBackend() + assert ss._resolve_lora_id("sa3style") == "sa3style" # exact id + assert ss._resolve_lora_id("SA3STYLE") == "SA3STYLE" # alias: miss + + +# --------------------------------------------------------------------------- +# Hard format validation at the enable boundary (ACE manager) +# --------------------------------------------------------------------------- + + +def _tiny_eager_manager() -> EagerLoRAManager: + decoder = nn.Module() + decoder.q = nn.Linear(16, 8, bias=False) + return EagerLoRAManager(decoder=decoder, device=torch.device("cpu")) + + +class TestAceEnableRejectsWrongFamily: + def test_sa3_file_raises_loudly(self, tmp_path): + """The exact-id bypass means an SA3 file CAN reach the ACE + enable path; it must raise, not silently no-op (the pre-fix + behavior: empty pair map -> empty deltas -> 'enabled' with + nothing applied).""" + p = _write_sa3_file(tmp_path) + mgr = _tiny_eager_manager() + mgr.register_lora(str(p)) + import pytest + with pytest.raises(RuntimeError, match="not an ACE-Step LoRA"): + mgr.enable_lora("sa3style", strength=1.0) + + def test_sa3_file_error_names_the_family(self, tmp_path): + p = _write_sa3_file(tmp_path) + mgr = _tiny_eager_manager() + mgr.register_lora(str(p)) + import pytest + with pytest.raises(RuntimeError, match="SA3"): + mgr.enable_lora("sa3style", strength=1.0) + + def test_non_lora_file_raises_without_sa3_hint(self, tmp_path): + p = tmp_path / "notalora.safetensors" + save_file({"encoder.weight": torch.zeros(2, 2)}, str(p)) + mgr = _tiny_eager_manager() + mgr.register_lora(str(p)) + import pytest + with pytest.raises(RuntimeError, match="not an ACE-Step LoRA") as ei: + mgr.enable_lora("notalora", strength=1.0) + assert "SA3" not in str(ei.value) + + def test_real_ace_file_still_loads(self, tmp_path): + """Regression guard for the new raise: a well-formed PEFT file + through the REAL _compute_deltas still enables.""" + p = tmp_path / "real.safetensors" + save_file( + { + "base_model.model.q.lora_A.weight": torch.ones(4, 16) * 0.5, + "base_model.model.q.lora_B.weight": torch.ones(8, 4) * 0.5, + }, + str(p), + ) + mgr = _tiny_eager_manager() + mgr.register_lora(str(p)) + mgr.enable_lora("real", strength=1.0) + assert mgr.get_lora("real").state == "enabled" + + +# --------------------------------------------------------------------------- +# Duplicate-stem collision handling in discovery +# --------------------------------------------------------------------------- + +from acestep.paths import assign_lora_ids + + +class TestAssignLoraIds: + def test_unique_stems_keep_bare_ids(self, tmp_path): + paths = [tmp_path / "a.safetensors", tmp_path / "b.safetensors"] + assert assign_lora_ids(paths) == [ + (paths[0], "a"), (paths[1], "b"), + ] + + def test_collision_first_keeps_stem_later_gets_parent_suffix(self, tmp_path): + p1 = tmp_path / "lib" / "vibes.safetensors" + p2 = tmp_path / "training_out" / "vibes.safetensors" + assert assign_lora_ids([p1, p2]) == [ + (p1, "vibes"), (p2, "vibes--training_out"), + ] + + def test_triple_collision_same_parent_name_gets_numeric_suffix(self, tmp_path): + p1 = tmp_path / "r1" / "out" / "x.safetensors" + p2 = tmp_path / "r2" / "out" / "x.safetensors" + p3 = tmp_path / "r3" / "out" / "x.safetensors" + ids = [lid for _, lid in assign_lora_ids([p1, p2, p3])] + assert ids[0] == "x" + assert ids[1] == "x--out" + assert ids[2] == "x--out-2" + assert len(set(ids)) == 3 + + def test_register_library_registers_all_collided_files(self, tmp_path, monkeypatch): + """End-to-end through the eager manager: a same-stem pair in two + subdirectories must yield TWO catalog entries (the old + first-wins registrar silently dropped the second).""" + (tmp_path / "lib").mkdir() + (tmp_path / "extra").mkdir() + (tmp_path / "lib" / "dupe.safetensors").write_bytes(b"") + (tmp_path / "extra" / "dupe.safetensors").write_bytes(b"") + + mgr = _tiny_eager_manager() + ids = mgr.register_library(tmp_path) + + assert sorted(ids) == ["dupe", "dupe--extra"] or sorted(ids) == [ + "dupe", "dupe--lib", + ] + assert len(mgr.list_loras()) == 2 From 00cd32cd188846b14907e6eff231bda4dcad771b Mon Sep 17 00:00:00 2001 From: RyanOnTheInside <7623207+ryanontheinside@users.noreply.github.com> Date: Wed, 22 Jul 2026 17:46:38 -0400 Subject: [PATCH 2/6] feat(sa3-lora): phase 0.5 de-risk harness, all five checks pass scripts/sa3/lora_derisk_phase05.py runs churn / rollback / svd / hygiene / bench against the real vendored stable_audio_3 package and the medium checkpoint. Measured on the 5090: - slot churn + strength targeting bit-exact through middle-slot removal (validates the manager-owned direct-copy install design) - mid-enable failure rollback restores the model bit-identically - -xs SVD cost: 70.3s CPU full enable on medium; cache would be ~3.84 GB fp16 full-basis (confirms phase-1 rejection of -xs) - cached-context teardown leaves pristine weights and bit-identical forwards (the SA3Backend.close() contract) - eager DiT step: 56ms base, +49%/1 LoRA, +114%/3 LoRA naive; parametrize.cached() collapses 3-LoRA overhead to +15%, so the per-tick cached() wrap is promoted to a phase 1 requirement --- scripts/sa3/lora_derisk_phase05.py | 562 +++++++++++++++++++++++++++++ 1 file changed, 562 insertions(+) create mode 100644 scripts/sa3/lora_derisk_phase05.py diff --git a/scripts/sa3/lora_derisk_phase05.py b/scripts/sa3/lora_derisk_phase05.py new file mode 100644 index 00000000..a2df68ca --- /dev/null +++ b/scripts/sa3/lora_derisk_phase05.py @@ -0,0 +1,562 @@ +"""SA3 LoRA Phase 0.5 de-risk prototype (throwaway harness). + +Gates the design decisions D4 (manager-owned slots, transactional +enable) and D6a (eager fallback viability) of notes/SA3_LORA_PLAN.md +BEFORE any production wiring. Five checks against the real vendored +``stable_audio_3`` package + the medium checkpoint: + +1. churn — enable A,B,C; disable B (middle slot); enable D; verify + every adapter's contribution and per-id strength + targeting survive the physical ParametrizationList shift + that ``remove_lora_by_index`` causes. +2. rollback — simulate a mid-enable failure (partial weight install), + roll back via ``remove_lora_by_index``, verify the model + is bit-identical to its pre-enable state. +3. svd — measure the real CPU SVD wall time an ``-xs`` adapter + would pay at apply time on medium (sizes the disk-cache + work item, or confirms the Phase-1 rejection). +4. hygiene — enable adapters (DiT + conditioner), tear down the way + ``SA3Backend.close()`` will, verify the process-cached + model is pristine: bitwise params/buffers, no leftover + parametrizations, and a bit-identical DiT forward. +5. bench — eager medium DiT per-tick latency with 0/1/3 rank-16 + LoRAs at batch (pipeline depth) 1 and 4, on real + conditioning. Prices D1's parametrize overhead and D6a's + eager fallback. + +Synthetic rank-16 adapters (deterministic randn fills) stand in for +trained files: every check here is about mechanics and cost, not +audio quality, and synthetic weights exercise the identical code path +(``add_lora`` + direct tensor install + ``set_lora_strength``). + +Run (idle GPU recommended for stable bench numbers): + .venv/Scripts/python.exe scripts/sa3/lora_derisk_phase05.py + .venv/Scripts/python.exe scripts/sa3/lora_derisk_phase05.py --checks churn,rollback + .venv/Scripts/python.exe scripts/sa3/lora_derisk_phase05.py --checks svd --svd-limit 40 +""" + +from __future__ import annotations + +import argparse +import sys +import time +from functools import partial +from pathlib import Path + +# --- sys.path: repo root FIRST (so `acestep` is ours, not a sibling shadow). +_HERE = Path(__file__).resolve().parent +_REPO_ROOT = next(p for p in (_HERE, *_HERE.parents) if (p / "pyproject.toml").exists()) +for _p in (str(_REPO_ROOT),): + while _p in sys.path: + sys.path.remove(_p) + sys.path.insert(0, _p) + +import torch # noqa: E402 + + +# --------------------------------------------------------------------------- +# Adapter plumbing (the exact mechanics SA3LoRAManager will use) +# --------------------------------------------------------------------------- + + +def _lora_modules(): + """Vendored package imports, resolved after ensure_sa3_paths.""" + from stable_audio_3.models.lora.model import ( + LoRAParametrization, + add_lora, + remove_lora_by_index, + set_lora_strength, + ) + from stable_audio_3.models.lora.utils import get_lora_layers, has_lora + + return ( + LoRAParametrization, add_lora, remove_lora_by_index, + set_lora_strength, get_lora_layers, has_lora, + ) + + +def _roots(sam): + """The exact application roots the vendored loader uses for + diffusion_cond models: model.model (DiT) + model.conditioner.""" + return [sam.model.model, sam.model.conditioner] + + +def _iter_params_of_index(sam, lora_index: int): + """All LoRAParametrization objects belonging to one adapter, + selected by their ``lora_index`` attribute (NOT physical position — + this is the D4 install mechanism that is immune to prior + removals).""" + LoRAParametrization = _lora_modules()[0] + for root in _roots(sam): + for _name, mod in root.named_modules(): + plist = getattr(getattr(mod, "parametrizations", None), "weight", None) + if plist is None: + continue + for p in plist: + if isinstance(p, LoRAParametrization) and p.lora_index == lora_index: + yield p + + +def add_adapter(sam, lora_index: int, *, rank: int = 16, seed: int, strength: float = 1.0): + """Register a plain-LoRA adapter on both roots and install synthetic + weights by direct copy into the parametrization objects. + + Mirrors load_and_apply_loras' registration exactly (from_linear / + from_conv1d partials, per-index), then installs weights the D4 way: + write into the objects selected by lora_index, bypassing the + index-remapped load_state_dict path entirely. + """ + import torch.nn as nn + (LoRAParametrization, add_lora, _rm, set_lora_strength, _gl, _hl) = _lora_modules() + + cfg = { + nn.Linear: {"weight": partial( + LoRAParametrization.from_linear, rank=rank, lora_alpha=rank, + adapter_type="lora", lora_index=lora_index, + )}, + nn.Conv1d: {"weight": partial( + LoRAParametrization.from_conv1d, rank=rank, lora_alpha=rank, + adapter_type="lora", lora_index=lora_index, + )}, + } + for root in _roots(sam): + add_lora(root, cfg) + + g = torch.Generator(device="cpu").manual_seed(seed) + n = 0 + with torch.no_grad(): + for p in _iter_params_of_index(sam, lora_index): + # Deterministic non-zero A and B so the delta is non-trivial. + p.lora_A.copy_(torch.randn( + p.lora_A.shape, generator=g, dtype=torch.float32, + ).mul_(0.02).to(p.lora_A.device)) + p.lora_B.copy_(torch.randn( + p.lora_B.shape, generator=g, dtype=torch.float32, + ).mul_(0.02).to(p.lora_B.device)) + n += 1 + for root in _roots(sam): + set_lora_strength(root, strength, lora_index=lora_index) + return n + + +def remove_adapter(sam, lora_index: int): + (_L, _a, remove_lora_by_index, _s, _gl, _hl) = _lora_modules() + for root in _roots(sam): + remove_lora_by_index(root, lora_index) + + +def set_strength(sam, lora_index: int, strength: float): + (_L, _a, _r, set_lora_strength, _gl, _hl) = _lora_modules() + for root in _roots(sam): + set_lora_strength(root, strength, lora_index=lora_index) + + +# --------------------------------------------------------------------------- +# Probes and snapshots +# --------------------------------------------------------------------------- + + +def probe_module(sam, fqn: str): + """Resolve a module by FQN; a ``cond:`` prefix roots the lookup at + the conditioner instead of the DiT.""" + if fqn.startswith("cond:"): + mod, fqn = sam.model.conditioner, fqn[len("cond:"):] + else: + mod = sam.model.model + for part in fqn.split("."): + mod = getattr(mod, part) + return mod + + +PROBE_FQNS = [ + "model.transformer.layers.0.self_attn.to_qkv", + "model.transformer.layers.11.ff.ff.0.proj", + "model.transformer.layers.23.self_attn.to_out", + # The single conditioner-side target on medium (the seconds_total + # embedder Linear) — proves conditioner application + teardown. + "cond:conditioners.seconds_total.embedder.embedding.1", +] + + +def expected_weight(mod) -> torch.Tensor: + """Recompute the effective weight by replaying the parametrization + chain in physical order with the same op order/dtype casts as the + vendored lora_forward, so a correct install is BIT-identical.""" + plist = mod.parametrizations.weight + w = plist.original.detach().clone() + for p in plist: + delta = (p.lora_B.detach() @ p.lora_A.detach()).view(w.shape) + delta = p.scaling * p.lora_strength * delta + w = w + delta.to(w.dtype) + return w + + +def check_probes(sam, label: str) -> bool: + ok = True + for fqn in PROBE_FQNS: + mod = probe_module(sam, fqn) + if not hasattr(mod, "parametrizations"): + continue + got = mod.weight.detach() + want = expected_weight(mod) + if not torch.equal(got, want): + max_err = (got.float() - want.float()).abs().max().item() + print(f" PROBE MISMATCH [{label}] {fqn}: max_err={max_err:.3e}") + ok = False + return ok + + +def snapshot_state(sam) -> dict[str, torch.Tensor]: + """CPU clones of every param + buffer under both application roots, + keyed by root-qualified name. The pristine reference for rollback + and hygiene bit-identity checks.""" + snap: dict[str, torch.Tensor] = {} + for ri, root in enumerate(_roots(sam)): + for name, p in root.named_parameters(): + snap[f"{ri}.{name}"] = p.detach().cpu().clone() + for name, b in root.named_buffers(): + snap[f"{ri}.{name}"] = b.detach().cpu().clone() + return snap + + +def compare_state(sam, snap: dict[str, torch.Tensor], label: str) -> bool: + now = snapshot_state(sam) + ok = True + missing = set(snap) - set(now) + extra = set(now) - set(snap) + if missing: + print(f" STATE [{label}]: {len(missing)} tensors MISSING, e.g. {sorted(missing)[:3]}") + ok = False + if extra: + print(f" STATE [{label}]: {len(extra)} EXTRA tensors, e.g. {sorted(extra)[:3]}") + ok = False + n_diff = 0 + for k in snap.keys() & now.keys(): + if not torch.equal(snap[k], now[k]): + if n_diff < 3: + print(f" STATE [{label}]: tensor differs: {k}") + n_diff += 1 + if n_diff: + print(f" STATE [{label}]: {n_diff} tensors differ bitwise") + ok = False + return ok + + +def no_parametrizations_left(sam) -> bool: + (_L, _a, _r, _s, _gl, has_lora) = _lora_modules() + for root in _roots(sam): + if has_lora(root): + return False + for _name, mod in root.named_modules(): + if getattr(mod, "parametrizations", None): + return False + return True + + +# --------------------------------------------------------------------------- +# Checks +# --------------------------------------------------------------------------- + + +def check_churn(sam) -> bool: + print("[churn] enable A(0), B(1), C(2)...") + for idx, seed in ((0, 100), (1, 200), (2, 300)): + n = add_adapter(sam, idx, seed=seed, strength=1.0) + print(f" adapter idx={idx}: {n} parametrizations installed") + ok = check_probes(sam, "A+B+C") + + print("[churn] disable B (middle slot, physical positions shift)...") + remove_adapter(sam, 1) + ok &= check_probes(sam, "A+C after removing B") + + print("[churn] enable D(3) after the shift (direct-copy install)...") + add_adapter(sam, 3, seed=400, strength=1.0) + ok &= check_probes(sam, "A+C+D") + + print("[churn] strength targeting by id after churn...") + set_strength(sam, 2, 0.5) # C + ok &= check_probes(sam, "C@0.5") + set_strength(sam, 3, 0.25) # D + ok &= check_probes(sam, "D@0.25") + set_strength(sam, 0, 0.0) # A muted + ok &= check_probes(sam, "A@0") + + # Verify the STRENGTH landed on the right adapters (attribute-level). + for idx, want in ((0, 0.0), (2, 0.5), (3, 0.25)): + for p in _iter_params_of_index(sam, idx): + got = float(p.lora_strength) + if got != want: + print(f" STRENGTH MISMATCH idx={idx}: got {got}, want {want}") + ok = False + break + + # Cleanup for the next check. + for idx in (0, 2, 3): + remove_adapter(sam, idx) + if not no_parametrizations_left(sam): + print(" CLEANUP FAILED: parametrizations left after removing all") + ok = False + print(f"[churn] {'PASS' if ok else 'FAIL'}") + return ok + + +def check_rollback(sam) -> bool: + print("[rollback] pre-state: adapter A(0) enabled...") + add_adapter(sam, 0, seed=100, strength=1.0) + pre = snapshot_state(sam) + + print("[rollback] attempt enable of X(1), fail mid-install...") + try: + import torch.nn as nn + (LoRAParametrization, add_lora, _r, _s, _gl, _hl) = _lora_modules() + cfg = { + nn.Linear: {"weight": partial( + LoRAParametrization.from_linear, rank=16, lora_alpha=16, + adapter_type="lora", lora_index=1, + )}, + nn.Conv1d: {"weight": partial( + LoRAParametrization.from_conv1d, rank=16, lora_alpha=16, + adapter_type="lora", lora_index=1, + )}, + } + for root in _roots(sam): + add_lora(root, cfg) + # Partial install: fill only the first few parametrizations, + # then hit the simulated truncated-state-dict error. + g = torch.Generator(device="cpu").manual_seed(999) + with torch.no_grad(): + for i, p in enumerate(_iter_params_of_index(sam, 1)): + if i >= 5: + raise RuntimeError("simulated truncated state dict") + p.lora_B.copy_(torch.randn( + p.lora_B.shape, generator=g, dtype=torch.float32, + ).to(p.lora_B.device)) + except RuntimeError as exc: + print(f" enable failed as intended: {exc}") + remove_adapter(sam, 1) # the transactional rollback + + ok = compare_state(sam, pre, "post-rollback vs pre-enable") + ok &= check_probes(sam, "A after rollback") + + remove_adapter(sam, 0) + if not no_parametrizations_left(sam): + print(" CLEANUP FAILED after rollback check") + ok = False + print(f"[rollback] {'PASS' if ok else 'FAIL'}") + return ok + + +def check_svd(sam, limit: int | None) -> bool: + """Time the CPU SVDs an -xs enable would run (upstream model.py: + W0.view(out,-1).cpu().float() -> torch.linalg.svd(full_matrices=False)).""" + import torch.nn as nn + mods = [] + for root_name, root in (("model", sam.model.model), ("conditioner", sam.model.conditioner)): + for name, mod in root.named_modules(): + if isinstance(mod, (nn.Linear, nn.Conv1d)): + mods.append((root_name, name, mod)) + total = len(mods) + if limit: + mods = mods[:limit] + print(f"[svd] timing CPU SVD over {len(mods)}/{total} Linear/Conv1d weights...") + t_total = 0.0 + worst = (0.0, "") + cache_bytes_fp16 = 0 + for i, (root_name, name, mod) in enumerate(mods): + W = mod.weight.detach() + W2 = W.view(W.shape[0], -1).cpu().float() + t0 = time.perf_counter() + torch.linalg.svd(W2, full_matrices=False) + dt = time.perf_counter() - t0 + t_total += dt + if dt > worst[0]: + worst = (dt, f"{root_name}.{name} {tuple(W2.shape)}") + r = min(W2.shape) + cache_bytes_fp16 += (W2.shape[0] * r + W2.shape[1] * r) * 2 + if (i + 1) % 50 == 0: + print(f" {i + 1}/{len(mods)}: cumulative {t_total:.1f}s") + scale = total / len(mods) if mods else 1.0 + print( + f"[svd] measured {len(mods)} layers: {t_total:.1f}s total " + f"(worst {worst[0] * 1000:.0f}ms @ {worst[1]}); " + f"extrapolated full enable: ~{t_total * scale:.1f}s; " + f"full-basis cache ~{cache_bytes_fp16 * scale / 1e9:.2f} GB fp16" + ) + print("[svd] PASS (measurement only)") + return True + + +def _bench_cond(ctx, duration: float, steps: int): + cond = ctx.prepare_cond(prompt="warm analog house groove, 124 bpm", + duration=duration, steps=steps) + return cond + + +def _dit_call_ms(ctx, cond, batch: int, iters: int = 50, warmup: int = 10) -> float: + from acestep.engine.sa3_stream_helpers import stack_sa3_cond_bundles + stacked = stack_sa3_cond_bundles([cond.cond_bundle] * batch) + x = torch.randn( + batch, ctx.latent_channels, cond.latent_frames, + device=ctx.device, dtype=ctx.dtype, + ) + t = torch.full((batch,), 0.5, device=ctx.device, dtype=ctx.dtype) + dit = ctx.dit + with torch.no_grad(): + for _ in range(warmup): + dit(x, t, **stacked) + torch.cuda.synchronize() + times = [] + for _ in range(iters): + t0 = time.perf_counter() + dit(x, t, **stacked) + torch.cuda.synchronize() + times.append((time.perf_counter() - t0) * 1000) + times.sort() + return times[len(times) // 2] + + +def check_bench(sam, ctx) -> bool: + print("[bench] preparing real conditioning (duration=30s, steps=8)...") + cond = _bench_cond(ctx, duration=30.0, steps=8) + print(f" latent_frames={cond.latent_frames}") + + results: dict[str, dict[int, float]] = {} + configs = [("0 LoRA", []), ("1 LoRA", [0]), ("3 LoRA", [0, 1, 2])] + installed: list[int] = [] + for label, idxs in configs: + for idx in idxs: + if idx not in installed: + add_adapter(sam, idx, seed=100 + idx, strength=1.0) + installed.append(idx) + results[label] = {} + for batch in (1, 4): + ms = _dit_call_ms(ctx, cond, batch) + results[label][batch] = ms + print(f" {label:7s} B={batch}: {ms:7.2f} ms/step (median)") + + # parametrize.cached() experiment (D1 candidate mitigation). + import torch.nn.utils.parametrize as parametrize + with parametrize.cached(): + ms_cached = _dit_call_ms(ctx, cond, 1, iters=30, warmup=5) + print(f" 3 LoRA B=1 under parametrize.cached(): {ms_cached:7.2f} ms/step") + + for idx in installed: + remove_adapter(sam, idx) + torch.cuda.empty_cache() + + base = results["0 LoRA"] + print("[bench] overhead vs 0 LoRA:") + for label in ("1 LoRA", "3 LoRA"): + for batch in (1, 4): + ovh = (results[label][batch] / base[batch] - 1.0) * 100 + print(f" {label} B={batch}: +{ovh:.1f}%") + print("[bench] PASS (measurement only)") + return True + + +def check_hygiene(sam, ctx) -> bool: + print("[hygiene] baseline DiT forward on the pristine model...") + cond = _bench_cond(ctx, duration=10.0, steps=8) + from acestep.engine.sa3_stream_helpers import stack_sa3_cond_bundles + stacked = stack_sa3_cond_bundles([cond.cond_bundle]) + x = torch.randn( + 1, ctx.latent_channels, cond.latent_frames, + device=ctx.device, dtype=ctx.dtype, + generator=torch.Generator(device=ctx.device.type).manual_seed(1528), + ) + t = torch.full((1,), 0.5, device=ctx.device, dtype=ctx.dtype) + with torch.no_grad(): + ref_out = ctx.dit(x, t, **stacked).clone() + pre = snapshot_state(sam) + pre_use_lora = getattr(sam.model, "use_lora", None) + pre_lora_names = getattr(sam.model, "lora_names", None) + + print("[hygiene] session lifetime: enable 2 adapters + loader flags...") + add_adapter(sam, 0, seed=100, strength=1.0) + add_adapter(sam, 1, seed=200, strength=0.7) + sam.model.use_lora = True + sam.model.lora_names = ["synthetic-a", "synthetic-b"] + with torch.no_grad(): + lora_out = ctx.dit(x, t, **stacked).clone() + if torch.equal(lora_out, ref_out): + print(" WARNING: adapters had NO effect on the forward (unexpected)") + + print("[hygiene] teardown (the SA3Backend.close() contract)...") + remove_adapter(sam, 0) + remove_adapter(sam, 1) + if pre_use_lora is None: + if hasattr(sam.model, "use_lora"): + del sam.model.use_lora + else: + sam.model.use_lora = pre_use_lora + if pre_lora_names is None: + if hasattr(sam.model, "lora_names"): + del sam.model.lora_names + else: + sam.model.lora_names = pre_lora_names + torch.cuda.empty_cache() + + ok = no_parametrizations_left(sam) + if not ok: + print(" FAIL: parametrizations remain after teardown") + ok &= compare_state(sam, pre, "post-teardown vs pristine") + + print("[hygiene] next-session forward on the same cached context...") + with torch.no_grad(): + post_out = ctx.dit(x, t, **stacked) + if not torch.equal(post_out, ref_out): + max_err = (post_out.float() - ref_out.float()).abs().max().item() + print(f" FAIL: post-teardown forward differs (max_err={max_err:.3e})") + ok = False + print(f"[hygiene] {'PASS' if ok else 'FAIL'}") + return ok + + +# --------------------------------------------------------------------------- + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--model", default="medium") + ap.add_argument( + "--checks", default="churn,rollback,svd,hygiene,bench", + help="comma list from: churn,rollback,svd,hygiene,bench", + ) + ap.add_argument( + "--svd-limit", type=int, default=None, + help="only time the first N SVDs and extrapolate", + ) + args = ap.parse_args() + wanted = [c.strip() for c in args.checks.split(",") if c.strip()] + + from acestep.engine.sa3_context import SA3Context + print(f"loading SA3Context(model_id={args.model!r})...") + t0 = time.perf_counter() + ctx = SA3Context(model_id=args.model) + sam = ctx.sam + print(f"loaded in {time.perf_counter() - t0:.1f}s; dtype={ctx.dtype}") + + results: dict[str, bool] = {} + for name in wanted: + if name == "churn": + results[name] = check_churn(sam) + elif name == "rollback": + results[name] = check_rollback(sam) + elif name == "svd": + results[name] = check_svd(sam, args.svd_limit) + elif name == "hygiene": + results[name] = check_hygiene(sam, ctx) + elif name == "bench": + results[name] = check_bench(sam, ctx) + else: + print(f"unknown check: {name}") + results[name] = False + + print("\n=== Phase 0.5 summary ===") + for name, ok in results.items(): + print(f" {name:9s} {'PASS' if ok else 'FAIL'}") + return 0 if all(results.values()) else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) From 9813e9a747e4898cc0ad35cd02d5ad810062d91a Mon Sep 17 00:00:00 2001 From: RyanOnTheInside <7623207+ryanontheinside@users.noreply.github.com> Date: Wed, 22 Jul 2026 18:09:46 -0400 Subject: [PATCH 3/6] feat(sa3-lora): eager SA3 LoRA end-to-end through the backend facade (phase 1) - engine/lora: extract LoRACatalogBase (catalog dict, state machine, prewarm executor, inspection) from the delta-merge LoRAManagerBase; ACE behavior byte-identical, existing tests untouched and green - engine/sa3_lora (new): SA3LoRAManager over the vendored stable_audio_3 parametrization engine. Manager-owned id->lora_index slots (monotonic, never reused) with direct-copy weight install, transactional enable with bit-identical rollback, hard format/ lineage validation at materialize (-xs rejected with a clear error, ACE files fail loudly), conditioner-target detection, and a close() that returns the process-cached model to pristine state. gc.collect before empty_cache is load-bearing for VRAM reclaim (parametrize's injected-class cycles defer frees to the cycle collector) - generator_backend/diffusion_backend: LoRA facade on the protocol with engine_obj-delegating defaults (ACE unchanged) - session: every engine_obj LoRA touchpoint now routes through the facade (_apply_lora_pending, _enabled_lora_ids, _resolve_lora_id, lora_catalog_payload) - sa3_backend: lora capability bit, lora_str_ knob expansion, per-tick strength reads with the 0.02 gate, has_pending_refit over the pending queues, conditioner-mutation serialization + cond-bundle rebuild through the prompt-swap path, interim TRT->eager DiT swap while LoRAs are active, per-tick parametrize.cached() wrap (phase 0.5 benchmark: +114% naive -> +15% cached at 3 LoRAs), close() - sa3_session: manager construction against the cached context, library registration, startup alias resolution/prewarm, knob seeding - families: manager threaded through backend_init; sa3 knob universe gains the lora_str placeholder for the homonym guard - tests: test_lora_sa3_manager (real vendored parametrize on a tiny tree, CPU), test_lora_facade_session (pending-drain via stub backend), SA3 backend LoRA wiring tests; scripts/sa3/lora_smoke.py GPU smoke passes twice back-to-back on the cached medium context --- acestep/engine/lora.py | 260 ++++++----- acestep/engine/sa3_lora.py | 593 +++++++++++++++++++++++++ acestep/streaming/diffusion_backend.py | 44 ++ acestep/streaming/families.py | 13 +- acestep/streaming/generator_backend.py | 62 ++- acestep/streaming/sa3_backend.py | 286 +++++++++++- acestep/streaming/sa3_session.py | 100 ++++- acestep/streaming/session.py | 17 +- scripts/sa3/lora_smoke.py | 208 +++++++++ tests/unit/test_lora_facade_session.py | 168 +++++++ tests/unit/test_lora_resolution.py | 27 +- tests/unit/test_lora_sa3_manager.py | 458 +++++++++++++++++++ tests/unit/test_sa3_backend.py | 213 +++++++++ 13 files changed, 2295 insertions(+), 154 deletions(-) create mode 100644 acestep/engine/sa3_lora.py create mode 100644 scripts/sa3/lora_smoke.py create mode 100644 tests/unit/test_lora_facade_session.py create mode 100644 tests/unit/test_lora_sa3_manager.py diff --git a/acestep/engine/lora.py b/acestep/engine/lora.py index a9b0d786..8168da44 100644 --- a/acestep/engine/lora.py +++ b/acestep/engine/lora.py @@ -68,80 +68,32 @@ class _LoRAEntry: materialized_bytes: int = 0 -class LoRAManagerBase(abc.ABC): - """Catalog + lifecycle + delta math; subclasses do the engine writeback. - - Subclasses must: - - - Populate ``self._base_weights`` (dict: param_name -> CPU/GPU tensor in - the live weight's native dtype) and ``self._param_dtype`` (dict: - param_name -> torch.dtype) during __init__, then call - ``super().__init__()`` to set up the lifecycle bookkeeping. - - Implement :meth:`_apply_to_engine`, which writes - ``base + sum_over_enabled(strength * delta)`` to the live weights - for the supplied ``param_names``. - - Param naming is decoder-relative (matches ``decoder.named_parameters()``). - The TRT subclass adds the ``decoder.`` prefix only at the refitter - boundary. +class LoRACatalogBase(abc.ABC): + """The family-agnostic slice of a LoRA manager: catalog dict, + REGISTERED → MATERIALIZING → MATERIALIZED (→ ENABLED) state + machine, background prewarm executor, and read-only inspection. + + Two manager families share this layer without sharing weight + semantics: the delta-merge managers (:class:`LoRAManagerBase` — ACE + eager + TRT refit, ``base + Σ strength·delta``) and the SA3 + parametrization manager (:mod:`acestep.engine.sa3_lora`), whose + materialize/enable transaction is a different shape (staged state + dicts + torch parametrize registration). Enable/disable/strength + live on the subclasses — their transaction and rollback semantics + differ by design. + + Subclasses implement :meth:`_materialize_worker` (worker-thread + body: fill the entry's staged payload without touching the live + weights; on failure restore REGISTERED state and re-raise). """ def __init__(self) -> None: - # Subclasses set _base_weights, _param_dtype, and _param_numel - # before super().__init__. _param_numel is the element count of - # the base weight (orientation-independent) so _compute_deltas - # can sanity-check that a LoRA was trained for THIS base model - # rather than silently producing wrong-shape deltas (e.g. a 2B - # LoRA materialized against an XL engine). - if not hasattr(self, "_base_weights"): - self._base_weights: Dict[str, torch.Tensor] = {} - if not hasattr(self, "_param_dtype"): - self._param_dtype: Dict[str, torch.dtype] = {} - if not hasattr(self, "_param_numel"): - self._param_numel: Dict[str, int] = {} - # Library + lifecycle state. Insertion order is preserved so # ``remove_lora(-1)`` can pop the most-recently-registered entry, # matching the legacy stack-style API. self._loras: Dict[str, _LoRAEntry] = {} - self._ever_dirty: Set[str] = set() self._executor: Optional[concurrent.futures.ThreadPoolExecutor] = None - # ------------------------------------------------------------------ - # Subclass hook - # ------------------------------------------------------------------ - - @abc.abstractmethod - def _apply_to_engine(self, param_names: Set[str]) -> None: - """Write ``base + Σ strength·delta`` to the live weights. - - ``param_names`` is the set of decoder-relative param names that - need updating. Subclasses iterate, accumulate enabled-LoRA - contributions in a buffer, and push to the live weights. - """ - - # ------------------------------------------------------------------ - # Lifecycle transition hooks (default: no-ops) - # - # The base lifecycle (enable_lora / disable_lora) calls these around - # state transitions so subclasses can attach side effects without - # overriding the lifecycle proper. The motivating case is the eager - # backend's "promote materialized deltas to GPU on enable, drop the - # GPU mirror on disable" pattern: deltas live cheaply in CPU RAM - # while merely MATERIALIZED, but the active set lives in VRAM so - # slider-driven refits stay zero-copy on the device. - # ------------------------------------------------------------------ - - def _on_enabled(self, entry: "_LoRAEntry") -> None: - """Hook: ``entry`` just transitioned to ENABLED. Called BEFORE - the contributing-strength refit fires so subclasses can stage - runtime data the refit will read.""" - - def _on_disabled(self, entry: "_LoRAEntry") -> None: - """Hook: ``entry`` just transitioned away from ENABLED. Called - AFTER ``entry.deltas`` is cleared but BEFORE the rollback refit - fires (so subclasses can drop their mirrors first).""" - # ------------------------------------------------------------------ # Library: catalog without RAM cost # ------------------------------------------------------------------ @@ -261,6 +213,138 @@ def prewarm_lora(self, lora_id: str) -> concurrent.futures.Future: logger.info("Prewarming LoRA: {}", lora_id) return entry.future + @abc.abstractmethod + def _materialize_worker(self, entry: _LoRAEntry) -> None: + """Worker-thread body: load + stage the entry's payload without + touching the live weights, then flip it to MATERIALIZED. On + failure restore REGISTERED (when still MATERIALIZING) and + re-raise. If the entry was concurrently disabled (state changed + away from MATERIALIZING), the result must be dropped rather + than resurrected.""" + + # ------------------------------------------------------------------ + # Inspection + # ------------------------------------------------------------------ + + def list_loras(self) -> List[LoRADescriptor]: + return [ + LoRADescriptor( + id=e.lora_id, path=e.path, name=e.name, + state=e.state.value, strength=e.strength, + materialized_bytes=e.materialized_bytes, + ) + for e in self._loras.values() + ] + + def get_lora(self, lora_id: str) -> LoRADescriptor: + e = self._require_entry(lora_id) + return LoRADescriptor( + id=e.lora_id, path=e.path, name=e.name, + state=e.state.value, strength=e.strength, + materialized_bytes=e.materialized_bytes, + ) + + def _require_entry(self, lora_id: str) -> _LoRAEntry: + if lora_id not in self._loras: + raise ValueError(f"LoRA {lora_id!r} not registered") + return self._loras[lora_id] + + @property + def has_active_loras(self) -> bool: + return any(e.state == LoRAState.ENABLED for e in self._loras.values()) + + @property + def active_lora_count(self) -> int: + return sum( + 1 for e in self._loras.values() if e.state == LoRAState.ENABLED + ) + + @property + def active_lora_ids(self) -> List[str]: + return [ + e.lora_id for e in self._loras.values() + if e.state == LoRAState.ENABLED + ] + + @property + def total_materialized_bytes(self) -> int: + return sum(e.materialized_bytes for e in self._loras.values()) + + +class LoRAManagerBase(LoRACatalogBase): + """Catalog + lifecycle + delta math; subclasses do the engine writeback. + + The delta-merge manager family (ACE lineage): materialization + computes full-rank ``B @ A`` deltas and enable/disable/strength + recompose ``base + Σ strength·delta`` into the live weights. + + Subclasses must: + + - Populate ``self._base_weights`` (dict: param_name -> CPU/GPU tensor in + the live weight's native dtype) and ``self._param_dtype`` (dict: + param_name -> torch.dtype) during __init__, then call + ``super().__init__()`` to set up the lifecycle bookkeeping. + - Implement :meth:`_apply_to_engine`, which writes + ``base + sum_over_enabled(strength * delta)`` to the live weights + for the supplied ``param_names``. + + Param naming is decoder-relative (matches ``decoder.named_parameters()``). + The TRT subclass adds the ``decoder.`` prefix only at the refitter + boundary. + """ + + def __init__(self) -> None: + # Subclasses set _base_weights, _param_dtype, and _param_numel + # before super().__init__. _param_numel is the element count of + # the base weight (orientation-independent) so _compute_deltas + # can sanity-check that a LoRA was trained for THIS base model + # rather than silently producing wrong-shape deltas (e.g. a 2B + # LoRA materialized against an XL engine). + if not hasattr(self, "_base_weights"): + self._base_weights: Dict[str, torch.Tensor] = {} + if not hasattr(self, "_param_dtype"): + self._param_dtype: Dict[str, torch.dtype] = {} + if not hasattr(self, "_param_numel"): + self._param_numel: Dict[str, int] = {} + + super().__init__() + self._ever_dirty: Set[str] = set() + + # ------------------------------------------------------------------ + # Subclass hook + # ------------------------------------------------------------------ + + @abc.abstractmethod + def _apply_to_engine(self, param_names: Set[str]) -> None: + """Write ``base + Σ strength·delta`` to the live weights. + + ``param_names`` is the set of decoder-relative param names that + need updating. Subclasses iterate, accumulate enabled-LoRA + contributions in a buffer, and push to the live weights. + """ + + # ------------------------------------------------------------------ + # Lifecycle transition hooks (default: no-ops) + # + # The base lifecycle (enable_lora / disable_lora) calls these around + # state transitions so subclasses can attach side effects without + # overriding the lifecycle proper. The motivating case is the eager + # backend's "promote materialized deltas to GPU on enable, drop the + # GPU mirror on disable" pattern: deltas live cheaply in CPU RAM + # while merely MATERIALIZED, but the active set lives in VRAM so + # slider-driven refits stay zero-copy on the device. + # ------------------------------------------------------------------ + + def _on_enabled(self, entry: "_LoRAEntry") -> None: + """Hook: ``entry`` just transitioned to ENABLED. Called BEFORE + the contributing-strength refit fires so subclasses can stage + runtime data the refit will read.""" + + def _on_disabled(self, entry: "_LoRAEntry") -> None: + """Hook: ``entry`` just transitioned away from ENABLED. Called + AFTER ``entry.deltas`` is cleared but BEFORE the rollback refit + fires (so subclasses can drop their mirrors first).""" + def _materialize_worker(self, entry: _LoRAEntry) -> None: """Worker-thread body. Loads safetensors, computes deltas, writes them to the entry. Engine state untouched. @@ -670,53 +754,9 @@ def apply_lora(self, lora_path: str, strength: float = 1.0) -> str: return lora_id # ------------------------------------------------------------------ - # Inspection + # Inspection (delta-specific; catalog inspection lives on the base) # ------------------------------------------------------------------ - def list_loras(self) -> List[LoRADescriptor]: - return [ - LoRADescriptor( - id=e.lora_id, path=e.path, name=e.name, - state=e.state.value, strength=e.strength, - materialized_bytes=e.materialized_bytes, - ) - for e in self._loras.values() - ] - - def get_lora(self, lora_id: str) -> LoRADescriptor: - e = self._require_entry(lora_id) - return LoRADescriptor( - id=e.lora_id, path=e.path, name=e.name, - state=e.state.value, strength=e.strength, - materialized_bytes=e.materialized_bytes, - ) - - def _require_entry(self, lora_id: str) -> _LoRAEntry: - if lora_id not in self._loras: - raise ValueError(f"LoRA {lora_id!r} not registered") - return self._loras[lora_id] - - @property - def has_active_loras(self) -> bool: - return any(e.state == LoRAState.ENABLED for e in self._loras.values()) - - @property - def active_lora_count(self) -> int: - return sum( - 1 for e in self._loras.values() if e.state == LoRAState.ENABLED - ) - - @property - def active_lora_ids(self) -> List[str]: - return [ - e.lora_id for e in self._loras.values() - if e.state == LoRAState.ENABLED - ] - - @property - def total_materialized_bytes(self) -> int: - return sum(e.materialized_bytes for e in self._loras.values()) - @property def refittable_param_count(self) -> int: return len(self._param_dtype) diff --git a/acestep/engine/sa3_lora.py b/acestep/engine/sa3_lora.py new file mode 100644 index 00000000..7e695714 --- /dev/null +++ b/acestep/engine/sa3_lora.py @@ -0,0 +1,593 @@ +"""SA3LoRAManager: DEMON's LoRA lifecycle over the vendored +``stable_audio_3`` parametrization engine. + +Design (notes/SA3_LORA_PLAN.md, D1/D4/D5): the manager conforms to the +public manager surface the backend facade and session code consume +(register / prewarm / enable / disable / set_strength / list / +descriptors) by subclassing the shared catalog layer +(:class:`~acestep.engine.lora.LoRACatalogBase`), but its weight +semantics are NOT the ACE delta merge — application is the vendored +package's own math (``LoRAParametrization`` registered via +``torch.nn.utils.parametrize``, live ``lora_strength`` buffer, +``alpha/rank`` scaling, DoRA/BoRA magnitude renormalization), so every +adapter variant renders exactly as the trainer defined it. + +Key mechanics, each validated by the Phase 0.5 GPU prototype +(``scripts/sa3/lora_derisk_phase05.py``): + +* **Manager-owned slots.** Upstream's ``remove_lora_by_index`` + physically deletes ``ParametrizationList`` entries (shifting later + positions) while its state-dict loading addresses physical indices — + so after removing a middle adapter, index-remapped ``load_state_dict`` + can hit the wrong slot. This manager therefore owns the authoritative + id → ``lora_index`` mapping (a monotonic counter, never reused) and + installs weights by **copying tensors directly into the + parametrization objects it just created**, selected by their + ``lora_index`` attribute — immune to physical shifts. +* **Registration is per-module, driven by the file's own key set**, + using the exact vendored constructors (``from_linear`` / + ``from_conv1d`` + ``parametrize.register_parametrization``, the same + calls upstream ``apply_lora`` makes). This is equivalent to + upstream's ``add_lora(include/exclude)`` + ``load_state_dict`` + pathway because the trainer saves adapter tensors for exactly the + modules its filters parametrized — and it never leaves zero-weight + parametrizations behind to tax the forward pass. +* **Transactional enable.** ENABLED commits only after registration, + weight install, and the strength-buffer write all succeed; any + failure rolls the just-allocated index back out of both roots before + re-raising (bit-identical restore, prototype check 2). +* **Teardown.** The SA3 torch model is process-cached across sessions + (``sa3_session.get_sa3_context``), so :meth:`close` must hand the + next session a pristine model: it removes every index this manager + ever allocated (coping with a partially-failed prior enable) and + shuts the prewarm executor down. +* **`-xs` variants are rejected at materialize** with a clear error: + their apply-time CPU SVD measured 70.3 s on medium (prototype check + 3), far beyond the stall pre-coverage envelope; support waits on the + checkpoint-versioned SVD disk cache work item. + +Application roots are upstream's exact ones for ``diffusion_cond`` +models: ``sam.model.model`` (the DiT tree) and ``sam.model.conditioner`` +— never ``sam.model`` wholesale (which would also sweep the +pretransform and double-traverse the conditioner). + +Threading contract (same as ACE's managers): register / enable / +disable / set_strength run on the inference-owning thread (the +session's pending-drain rendezvous); prewarm materializes on a +background executor and touches only staged CPU state, never the model. +""" + +from __future__ import annotations + +import time +from dataclasses import dataclass, field +from pathlib import Path +from typing import Dict, Optional, Tuple + +from loguru import logger +import torch + +from acestep.engine.lora import ( + LoRACatalogBase, + LoRAState, + _LoRAEntry, +) + +# Adapter tensor names the vendored trainer saves, by adapter family. +_SA3_PARAM_NAMES = { + "lora_A", "lora_B", "magnitude", "magnitude_r", "magnitude_c", "M_xs", +} +_SA3_KEY_MARKER = ".parametrizations.weight." +_ACE_KEY_MARKERS = (".lora_A.weight", ".lora_B.weight") + + +def _vendored_lora(): + """Import the vendored LoRA modules through the managed sys.path + setup SA3Context already uses — never a pip dependency.""" + from acestep.engine.sa3_helpers import ensure_sa3_paths + + ensure_sa3_paths() + from stable_audio_3.models.lora import model as lora_model # noqa: PLC0415 + from stable_audio_3.models.lora import utils as lora_utils # noqa: PLC0415 + + return lora_model, lora_utils + + +@dataclass +class _StagedSA3File: + """CPU-staged, validated payload for one materialized file.""" + + adapter_type: str + rank: int + alpha: float + # (root_index, module_fqn) -> {param_name: cpu tensor}. root_index + # 0 = DiT root, 1 = conditioner root. + weights: Dict[Tuple[int, str], Dict[str, torch.Tensor]] = field( + default_factory=dict, + ) + touches_conditioner: bool = False + total_bytes: int = 0 + + +class SA3LoRAManager(LoRACatalogBase): + """SA3 LoRA lifecycle over the vendored parametrization engine. + + ``model_root`` / ``conditioner_root`` are the live application + roots (``sam.model.model`` / ``sam.model.conditioner``); + ``checkpoint_id`` is the runtime model id ("medium" / + "small-music" / "small-sfx") the lineage check compares against. + """ + + def __init__( + self, + *, + model_root: torch.nn.Module, + conditioner_root: Optional[torch.nn.Module] = None, + checkpoint_id: str = "", + ) -> None: + if model_root is None: + raise ValueError("SA3LoRAManager requires a model_root module") + self._model_root = model_root + self._conditioner_root = conditioner_root + self._checkpoint_id = str(checkpoint_id) + + # Authoritative id -> lora_index mapping for ENABLED adapters. + # Indices come from a monotonic counter and are never reused, so + # a stale index can never alias a live adapter (D4). + self._index_by_id: Dict[str, int] = {} + self._index_counter = 0 + # Every index ever allocated — the close() sweep target, so a + # partially-failed enable whose rollback also failed still gets + # stripped from the process-cached model at teardown. + self._allocated_indices: set[int] = set() + # Staged payloads for MATERIALIZED/ENABLED entries, keyed by id. + # Written by the prewarm worker (CPU only), read on the + # inference thread — same single-worker + GIL discipline as the + # ACE managers' entry.deltas. + self._staged: Dict[str, _StagedSA3File] = {} + + super().__init__() + logger.info( + "SA3 LoRA manager ready: checkpoint_id={} conditioner={}", + self._checkpoint_id, conditioner_root is not None, + ) + + # ------------------------------------------------------------------ + # Root / module helpers + # ------------------------------------------------------------------ + + def _roots(self) -> list: + roots = [self._model_root] + if self._conditioner_root is not None: + roots.append(self._conditioner_root) + return roots + + def _resolve_module(self, fqn: str): + """Locate ``fqn`` under the application roots. Returns + ``(root_index, module)`` or ``(None, None)`` on a miss.""" + for ri, root in enumerate(self._roots()): + mod = root + found = True + for part in fqn.split("."): + if not hasattr(mod, part): + found = False + break + mod = getattr(mod, part) + if found and isinstance(mod, torch.nn.Module): + return ri, mod + return None, None + + # ------------------------------------------------------------------ + # Materialize (background executor; GPU-free) + # ------------------------------------------------------------------ + + def _materialize_worker(self, entry: _LoRAEntry) -> None: + t0 = time.perf_counter() + try: + staged = self._stage_file(entry.path) + except Exception: + if entry.state == LoRAState.MATERIALIZING: + entry.state = LoRAState.REGISTERED + entry.future = None + raise + if entry.state == LoRAState.MATERIALIZING: + self._staged[entry.lora_id] = staged + entry.materialized_bytes = staged.total_bytes + entry.state = LoRAState.MATERIALIZED + elapsed = (time.perf_counter() - t0) * 1000 + logger.info( + "Materialized SA3 LoRA {} ({} modules, {:.1f} MB, " + "adapter={}, rank={}) in {:.1f}ms", + entry.lora_id, len(staged.weights), + staged.total_bytes / 1e6, staged.adapter_type, + staged.rank, elapsed, + ) + + def _stage_file(self, lora_path: str) -> _StagedSA3File: + """Read + validate one file into a CPU-staged payload. + + Hard validation at this boundary (D3): wrong-family key layouts, + ``-xs`` variants, lineage mismatches, and + no-module-matches-the-model all raise typed errors — a bad file + must never become a silent no-op enable. + """ + _, lora_utils = _vendored_lora() + name = Path(lora_path).name + + if not str(lora_path).endswith(".safetensors"): + # Upstream's .ckpt path is torch.load with pickles — an RCE + # vector we never expose at runtime. + raise RuntimeError( + f"SA3 LoRA {name}: only .safetensors files are accepted " + f"(.ckpt is a pickle and is never loaded at runtime; " + f"convert with stable_audio_3's " + f"convert_lora_ckpt_to_safetensors)" + ) + + sd, config = lora_utils.load_lora_checkpoint(lora_path) + + # Family validation from the actual key layout. + if not any(_SA3_KEY_MARKER in k for k in sd): + ace_like = any( + m in k for k in sd for m in _ACE_KEY_MARKERS + ) + hint = ( + " The key layout matches the ACE-Step (PEFT) format; " + "ACE LoRAs cannot load on an SA3 model." + if ace_like else "" + ) + raise RuntimeError( + f"LoRA {name} is not an SA3 LoRA: none of its " + f"{len(sd)} tensors use the parametrization key layout " + f"(…​.parametrizations.weight..).{hint}" + ) + + adapter_type = lora_utils.resolve_adapter_type( + config.get("adapter_type", "lora"), sd, + ) + if adapter_type.endswith("-xs"): + raise RuntimeError( + f"SA3 LoRA {name}: adapter type {adapter_type!r} is not " + f"yet supported — -xs variants need SVD bases of the " + f"base weights at apply time (~70s CPU on medium, " + f"measured), which requires the planned per-checkpoint " + f"SVD disk cache. Use a non-xs export of this adapter." + ) + + rank = int(config.get("rank") or lora_utils.infer_global_rank(sd)) + alpha = float(config.get("alpha") or rank) + + # Lineage check (the SA3 analog of ACE's "2B LoRA on XL engine" + # message). Permissive when the file doesn't say or the spelling + # is unrecognized, per the seam contract. + from acestep.lora_metadata import canonical_sa3_lineage + + lineage = canonical_sa3_lineage(config.get("base_model")) + if lineage is not None and self._checkpoint_id and ( + lineage != self._checkpoint_id + ): + raise RuntimeError( + f"SA3 LoRA {name} was trained against the " + f"{config.get('base_model')!r} base model " + f"(runtime lineage {lineage!r}) but the loaded " + f"checkpoint is {self._checkpoint_id!r}. Load it on a " + f"matching session." + ) + + # DoRA magnitude tensors may be saved 2D; squeeze to the 1D the + # live parametrization params use (upstream loader parity). + lora_utils.prepare_dora_state_dict(sd) + + # Group tensors by module FQN and resolve against the live tree. + by_fqn: Dict[str, Dict[str, torch.Tensor]] = {} + for key, tensor in sd.items(): + marker_at = key.find(_SA3_KEY_MARKER) + if marker_at < 0: + continue + fqn = key[:marker_at] + tail = key[marker_at + len(_SA3_KEY_MARKER):] + parts = tail.split(".", 1) + if len(parts) != 2 or parts[1] not in _SA3_PARAM_NAMES: + continue + by_fqn.setdefault(fqn, {})[parts[1]] = tensor + + pin = torch.cuda.is_available() + staged = _StagedSA3File( + adapter_type=adapter_type, rank=rank, alpha=alpha, + ) + unmatched = 0 + shape_mismatch = 0 + first_mismatch: Optional[tuple] = None + for fqn, params in by_fqn.items(): + ri, mod = self._resolve_module(fqn) + if ri is None or not isinstance( + mod, (torch.nn.Linear, torch.nn.Conv1d), + ): + unmatched += 1 + continue + weight = getattr(mod, "weight", None) + if weight is None: + unmatched += 1 + continue + w2 = weight.view(weight.shape[0], -1) + fan_out, fan_in = int(w2.shape[0]), int(w2.shape[1]) + a = params.get("lora_A") + b = params.get("lora_B") + bad = ( + (a is not None and tuple(a.shape) != (rank, fan_in)) + or (b is not None and tuple(b.shape) != (fan_out, rank)) + ) + if bad: + shape_mismatch += 1 + if first_mismatch is None: + got = tuple(a.shape) if a is not None else tuple(b.shape) + first_mismatch = (fqn, got, (fan_out, fan_in)) + continue + cpu_params = {} + for pname, t in params.items(): + t = t.detach().contiguous() + if pin: + t = t.pin_memory() + cpu_params[pname] = t + staged.total_bytes += t.numel() * t.element_size() + staged.weights[(ri, fqn)] = cpu_params + if ri == 1: + staged.touches_conditioner = True + + if not staged.weights: + detail = "" + if first_mismatch is not None: + fqn, got, fans = first_mismatch + detail = ( + f" E.g. {fqn!r}: adapter shape {got} does not fit " + f"the base weight (fan_out, fan_in)={fans}." + ) + raise RuntimeError( + f"SA3 LoRA {name} does not fit the loaded " + f"{self._checkpoint_id or 'model'}: " + f"{unmatched} module(s) missing from the model tree, " + f"{shape_mismatch} with mismatched shapes.{detail} " + f"Common cause: a LoRA trained for a different SA3 " + f"checkpoint size." + ) + if unmatched or shape_mismatch: + logger.warning( + "SA3 LoRA {}: {} modules staged, {} unmatched, {} " + "shape-mismatched (partial overlap; the rest applies)", + name, len(staged.weights), unmatched, shape_mismatch, + ) + return staged + + # ------------------------------------------------------------------ + # Enable / disable / strength (inference thread) + # ------------------------------------------------------------------ + + def touches_conditioner(self, lora_id: str) -> bool: + """Whether ``lora_id``'s staged payload targets the conditioner + (drives the backend's cond-bundle rebuild, D5). False for + entries that aren't materialized.""" + staged = self._staged.get(lora_id) + return bool(staged is not None and staged.touches_conditioner) + + def enable_lora( + self, lora_id: str, strength: Optional[float] = None, + ) -> None: + """Promote a LoRA to ENABLED — transactional. + + Registers parametrizations on the staged modules, installs the + staged weights by direct copy, writes the strength buffer, and + only then commits ENABLED state + the id→index mapping. On any + failure the just-allocated index is removed from both roots + (bit-identical restore, Phase 0.5 check 2) before re-raising — + the loud-failure path, never a half-applied adapter. + + Atomic enable-at-strength holds structurally: everything here + runs on the inference thread inside the pending-drain + rendezvous, so no forward pass can observe the intermediate + state, and the strength buffer is written before commit. + """ + entry = self._require_entry(lora_id) + if strength is not None: + entry.strength = float(strength) + if entry.state == LoRAState.ENABLED: + return + if entry.state == LoRAState.MATERIALIZING: + assert entry.future is not None + entry.future.result() + if entry.state == LoRAState.REGISTERED: + t0 = time.perf_counter() + staged = self._stage_file(entry.path) + self._staged[lora_id] = staged + entry.materialized_bytes = staged.total_bytes + entry.state = LoRAState.MATERIALIZED + logger.info( + "Materialized SA3 LoRA {} inline in {:.1f}ms", + lora_id, (time.perf_counter() - t0) * 1000, + ) + staged = self._staged[lora_id] + + lora_model, _ = _vendored_lora() + from functools import partial + + import torch.nn.utils.parametrize as parametrize + + idx = self._index_counter + self._index_counter += 1 + self._allocated_indices.add(idx) + + t0 = time.perf_counter() + try: + # 1. Register parametrizations on exactly the staged + # modules, via the vendored constructors, keeping a + # direct handle on each created object so the install + # step below cannot depend on physical list positions. + created: list = [] + for (ri, fqn), params in staged.weights.items(): + _ri, mod = self._resolve_module(fqn) + if mod is None: + raise RuntimeError( + f"SA3 LoRA {lora_id}: staged module {fqn!r} " + f"vanished from the model tree" + ) + if isinstance(mod, torch.nn.Conv1d): + ctor = lora_model.LoRAParametrization.from_conv1d + else: + ctor = lora_model.LoRAParametrization.from_linear + param_fn = partial( + ctor, + rank=staged.rank, lora_alpha=staged.alpha, + adapter_type=staged.adapter_type, lora_index=idx, + ) + p_obj = param_fn(mod) + parametrize.register_parametrization( + mod, "weight", p_obj, unsafe=True, + ) + created.append((p_obj, params)) + # 2. Direct-copy install into the objects just created (D4: + # bypasses index-remapped load_state_dict entirely, so + # enable is immune to prior removals' physical shifts). + with torch.no_grad(): + for p_obj, params in created: + for pname, tensor in params.items(): + target = getattr(p_obj, pname, None) + if target is None: + logger.warning( + "SA3 LoRA {}: adapter param {} missing " + "on live parametrization; skipped", + lora_id, pname, + ) + continue + target.copy_(tensor.to( + device=target.device, dtype=target.dtype, + non_blocking=True, + )) + # 3. Strength buffer — before commit, so the first forward + # after the rendezvous already sees the target strength. + for root in self._roots(): + lora_model.set_lora_strength( + root, float(entry.strength), lora_index=idx, + ) + except Exception: + self._remove_index(idx) + raise + + # 4. Commit. + self._index_by_id[lora_id] = idx + entry.state = LoRAState.ENABLED + entry.future = None + logger.info( + "Enabled SA3 LoRA {} (index={}, {} modules, adapter={}, " + "strength={:.2f}, conditioner={}) in {:.1f}ms", + lora_id, idx, len(staged.weights), staged.adapter_type, + entry.strength, staged.touches_conditioner, + (time.perf_counter() - t0) * 1000, + ) + + def _remove_index(self, idx: int) -> None: + """Strip every parametrization carrying ``idx`` from both roots + (upstream removal; restores the original weight when a module's + last adapter goes).""" + lora_model, _ = _vendored_lora() + for root in self._roots(): + lora_model.remove_lora_by_index(root, idx) + + def disable_lora(self, lora_id: str) -> None: + """Remove the adapter from the live model, drop staged CPU + tensors, and return the freed VRAM to the driver. Strength is + preserved on the entry so re-enable returns to the same slider + position (ACE convention).""" + entry = self._require_entry(lora_id) + if entry.state == LoRAState.REGISTERED: + return + if entry.state == LoRAState.MATERIALIZING and entry.future is not None: + try: + entry.future.result() + except Exception: + pass + + was_enabled = entry.state == LoRAState.ENABLED + idx = self._index_by_id.pop(lora_id, None) + entry.state = LoRAState.REGISTERED + entry.materialized_bytes = 0 + entry.future = None + self._staged.pop(lora_id, None) + + if was_enabled and idx is not None: + self._remove_index(idx) + # Same fragmentation rationale as the ACE managers: hand the + # freed adapter tensors back to the driver at the actual + # free event. The gc.collect() first is load-bearing here: + # torch's parametrize injects a per-module class whose type + # object sits in reference cycles, so the removed + # ParametrizationList (and the adapter parameters it holds) + # waits on the cycle collector — measured on medium, the + # ~38 MB rank-8 footprint stays "allocated" until then. + if torch.cuda.is_available(): + import gc + + gc.collect() + torch.cuda.empty_cache() + logger.info( + "Disabled SA3 LoRA {} (was_enabled={})", lora_id, was_enabled, + ) + + def set_lora_strength(self, lora_id: str, strength: float) -> None: + """Live strength update: a buffer write on every parametrization + of this adapter — no recompute, no refit (upstream's realtime + mechanism). Raises for non-enabled ids, matching ACE.""" + entry = self._require_entry(lora_id) + if entry.state != LoRAState.ENABLED: + raise ValueError( + f"LoRA {lora_id!r} is not enabled (state={entry.state.value}). " + "Call enable_lora() first." + ) + if entry.strength == strength: + return + old = entry.strength + entry.strength = float(strength) + idx = self._index_by_id[lora_id] + lora_model, _ = _vendored_lora() + for root in self._roots(): + lora_model.set_lora_strength( + root, float(strength), lora_index=idx, + ) + logger.info( + "SA3 LoRA {} strength: {:.3f} -> {:.3f}", lora_id, old, strength, + ) + + # ------------------------------------------------------------------ + # Teardown + # ------------------------------------------------------------------ + + def close(self) -> None: + """Return the process-cached model to pristine state: strip + every index this manager ever allocated (not just the currently + enabled ones — a partially-failed enable whose rollback also + failed still gets swept), drop staged payloads, and shut down + the prewarm executor. Idempotent. Phase 0.5 check 4 validated + that this restore is bitwise-complete.""" + for idx in sorted(self._allocated_indices): + try: + self._remove_index(idx) + except Exception as e: + logger.warning( + "SA3 LoRA teardown: removing index {} raised: {}", + idx, e, + ) + self._allocated_indices.clear() + self._index_by_id.clear() + self._staged.clear() + for entry in self._loras.values(): + if entry.state != LoRAState.REGISTERED: + entry.state = LoRAState.REGISTERED + entry.materialized_bytes = 0 + entry.future = None + self._loras.clear() + if self._executor is not None: + self._executor.shutdown(wait=False, cancel_futures=True) + self._executor = None + if torch.cuda.is_available(): + import gc + + gc.collect() # break the injected-class cycles (see disable) + torch.cuda.empty_cache() diff --git a/acestep/streaming/diffusion_backend.py b/acestep/streaming/diffusion_backend.py index 3efe178e..6a7f31ac 100644 --- a/acestep/streaming/diffusion_backend.py +++ b/acestep/streaming/diffusion_backend.py @@ -86,6 +86,50 @@ def lora_compatible(self, metadata: dict) -> bool: # Families with a real axis override (ACE: base-model scale). return True + # ---- LoRA facade (D2): engine_obj-delegating defaults ------------------- + # + # ACE's LoRA managers live inside its DiffusionEngine (``engine_obj`` + # on the concrete backend); these defaults route the session's LoRA + # plumbing there so ACE behavior is unchanged by the facade. SA3 + # overrides the whole block with its own manager. A family with + # neither (engine_obj None / absent) reports unavailable and fails + # loudly on mutation. + + def _lora_engine(self): + return getattr(self, "engine_obj", None) + + def _require_lora_engine(self): + eng = self._lora_engine() + if eng is None: + raise RuntimeError( + f"backend {self.name!r} has no LoRA engine; the session's " + "capability gate should have rejected this command" + ) + return eng + + def lora_available(self) -> bool: + eng = self._lora_engine() + return bool(eng is not None and getattr(eng, "lora_available", False)) + + def register_lora(self, path: str) -> str: + return self._require_lora_engine().register_lora(path) + + def prewarm_lora(self, lora_id: str): + return self._require_lora_engine().prewarm_lora(lora_id) + + def enable_lora(self, lora_id: str, strength=None) -> None: + self._require_lora_engine().enable_lora(lora_id, strength=strength) + + def disable_lora(self, lora_id: str) -> None: + self._require_lora_engine().disable_lora(lora_id) + + def set_lora_strength(self, lora_id: str, strength: float) -> None: + self._require_lora_engine().set_lora_strength(lora_id, strength) + + def list_loras(self) -> list: + eng = self._lora_engine() + return eng.list_loras() if eng is not None else [] + def rebuild_imminent(self, knobs: dict) -> bool: return False diff --git a/acestep/streaming/families.py b/acestep/streaming/families.py index 5c46daf5..1235a49a 100644 --- a/acestep/streaming/families.py +++ b/acestep/streaming/families.py @@ -86,6 +86,12 @@ def _make_sa3(ss): steps=int(ss.config.steps), depth=int(ss.state.current_depth), vae_window_s=float(ss.vae_window), + # SA3 LoRA (plan Phase 1): the create path constructs the + # family manager against the process-cached model and stashes + # it here; .get so an in-process payload predating the LoRA + # surface stays LoRA-less rather than failing assembly. + lora_manager=init.get("lora_manager"), + use_lora=bool(ss.use_lora), ) @@ -191,7 +197,12 @@ def _acestep_knob_universe(): def _sa3_knob_universe(): from acestep.streaming.sa3_backend import sa3_knob_specs - return sa3_knob_specs() + # One representative LoRA-strength knob (the per-id specs all come + # from the shared lora_strength_spec factory, so one placeholder id + # covers the pattern — same convention as the ACE universe). The + # name is shared with ACE's universe deliberately: the homonym + # guard proves the spec shapes are identical across families. + return sa3_knob_specs(loras=[""]) FAMILY_KNOB_UNIVERSES = { diff --git a/acestep/streaming/generator_backend.py b/acestep/streaming/generator_backend.py index b68d44ed..ce19ad1f 100644 --- a/acestep/streaming/generator_backend.py +++ b/acestep/streaming/generator_backend.py @@ -243,18 +243,62 @@ def lora_compatible(self, metadata: dict) -> bool: catalog entry. Backend-owned so the compatibility axes can grow per family: - today ACE checks the trained base-model scale against its - checkpoint; an SA3 LoRA lineage would add its own checks here - without the session or the wire changing. The session projects - this verdict as the ``compatible`` bool on every catalog entry - (clients decide whether to hide or grey incompatible rows) and - restricts enable-time alias resolution to the compatible - subset. Implementations must be permissive on unknowns — - a LoRA or engine the predicate can't classify is compatible, - never hidden. + ACE checks the weight-format family + trained base-model scale + against its checkpoint; SA3 checks family + training lineage. + The session projects this verdict as the ``compatible`` bool on + every catalog entry (clients decide whether to hide or grey + incompatible rows) and restricts enable-time alias resolution + to the compatible subset. Implementations must be permissive on + unknowns — a LoRA or engine the predicate can't classify is + compatible, never hidden. """ ... + # ---- LoRA facade ----------------------------------------------------- + # + # The session's LoRA plumbing (pending-drain, catalog payload, alias + # resolution, knob-manifest rebuilds, startup enablement) speaks to + # the backend through these instead of reaching into the ACE + # engine_obj — the seam that lets SA3 (a parametrization manager + # over a process-cached torch model) and ACE (a delta/refit manager + # inside the DiffusionEngine) serve the same wire surface. + # DiffusionBackend provides engine_obj-delegating defaults, so ACE + # backends inherit their historical behavior unchanged. + + def lora_available(self) -> bool: + """Whether this backend has a live LoRA manager to drive. The + capability bit (``Capabilities.lora``) gates the wire commands; + this gates the session-side plumbing (catalog, pending drain).""" + ... + + def register_lora(self, path: str) -> str: + """Add a file to the backend's LoRA catalog (startup + ``config.lora_paths``); returns the catalog id.""" + ... + + def prewarm_lora(self, lora_id: str): + """Kick off background materialization; returns a Future.""" + ... + + def enable_lora(self, lora_id: str, strength: Optional[float] = None) -> None: + """Enable at ``strength`` in one shot (atomic-strength + contract). Runs on the runner thread inside the pending-drain + rendezvous. Failures raise — never a silent no-op.""" + ... + + def disable_lora(self, lora_id: str) -> None: + """Disable and release the adapter's GPU footprint.""" + ... + + def set_lora_strength(self, lora_id: str, strength: float) -> None: + """Live strength update for an ENABLED LoRA.""" + ... + + def list_loras(self) -> list: + """Catalog descriptors (id/path/name/state/strength/bytes); + empty when no manager exists.""" + ... + # ---- hot loop -------------------------------------------------------- def sync_source(self, ctx: TickContext) -> None: diff --git a/acestep/streaming/sa3_backend.py b/acestep/streaming/sa3_backend.py index 6452e12f..84e59737 100644 --- a/acestep/streaming/sa3_backend.py +++ b/acestep/streaming/sa3_backend.py @@ -69,7 +69,11 @@ Capabilities, TickContext, ) -from acestep.streaming.knobs import KnobSpec, knob_specs as registry_knob_specs +from acestep.streaming.knobs import ( + KnobSpec, + knob_specs as registry_knob_specs, + lora_strength_spec, +) # Delivery rate (v1): SA3's native 44.1 kHz is resampled at the decode # boundary so everything downstream of the backend stays at the engine @@ -86,7 +90,38 @@ ) -def sa3_knob_specs() -> list: +def sa3_lora_compatible(metadata: dict, model_id: str) -> bool: + """The SA3 compatibility predicate: weight-format family + + training lineage. + + Family: a file whose sniffed ``lora_family`` names a different + family (e.g. "ace") can never load on the SA3 parametrization + engine — hard no. Unknown family stays permissive. + + Lineage: the file's ``base_model`` (sidecar wins, else the + embedded trainer config) canonicalized to a runtime checkpoint id + must match the loaded model. Unknown/unrecognized lineage stays + permissive per the seam contract. + + Module-level (not a method) because the create path + (:mod:`acestep.streaming.sa3_session`) needs the same verdict for + startup alias resolution before the backend exists — mirroring how + the ACE create path applies its scale axis directly. + """ + from acestep.lora_metadata import canonical_sa3_lineage + + family = metadata.get("lora_family") + if family and family != "sa3": + return False + lineage = canonical_sa3_lineage(metadata.get("base_model")) + if lineage is None or not model_id: + # Unknown on either side is compatible (the same permissive + # stance as the scale axis). + return True + return lineage == model_id + + +def sa3_knob_specs(loras: tuple | list = ()) -> list: """The SA3 family knob manifest (backend-owned, plan §3.3). ``seed``, ``steps_override``, ``x0_target``, ``feedback`` and @@ -97,6 +132,11 @@ def sa3_knob_specs() -> list: impossible instead). ``sa3_denoise`` / ``sa3_shift`` are family-prefixed because ACE's ``denoise`` and ``shift`` mean different things. + + ``loras`` is the session's enabled-LoRA id list; each id expands to + the shared ``lora_str_`` strength spec — the same registry + factory ACE uses, so the knob's shape (and therefore its wire + semantics) cannot fork across families. """ shared = {s.name: s for s in registry_knob_specs(False)} return [ @@ -124,7 +164,7 @@ def sa3_knob_specs() -> list: shared["feedback_depth"], shared["seed"], shared["steps_override"], - ] + ] + [lora_strength_spec(lid) for lid in loras] class SA3Backend(DiffusionBackend): @@ -174,6 +214,23 @@ def __init__( # SA3Context); None on directly-constructed test backends, where # handle_swap_source fails loudly instead. source_encoder: Optional[Callable] = None, + # SA3 LoRA (notes/SA3_LORA_PLAN.md Phase 1): the family's + # parametrization manager (acestep.engine.sa3_lora), constructed + # by the create path against the process-cached model. None = + # no LoRA surface (directly-constructed test backends). + lora_manager=None, + # Session-level LoRA toggle (config.lora); gates the capability + # bit, the per-tick strength reads, and has_pending_refit. + use_lora: bool = False, + # The always-resident eager DiT module, for the interim + # TRT→eager swap (D6a): when a LoRA enables on a TRT-DiT + # session the adapter's dit swaps to this; when the last + # disables it swaps back. None/identical when the session is + # already eager (swap becomes a no-op). + eager_dit=None, + # Runtime checkpoint id ("medium"/"small-music"/…) for the + # lineage axis of lora_compatible. + model_id: str = "", ): super().__init__(adapter=adapter, codec=codec) self._cond = cond @@ -241,6 +298,24 @@ def __init__( # only ever holds this to snapshot, never across pipeline work. self._control_lock = threading.Lock() + # ---- LoRA (plan D4/D5/D6a) ------------------------------------ + self._lora_mgr = lora_manager + self._use_lora = bool(use_lora) + self._model_id = str(model_id) + # D5 serialization point: conditioner EXECUTION (prompt-swap + # T5Gemma captures on the command thread) and conditioner + # MUTATION (LoRA parametrization changes at the runner + # rendezvous) are mutually exclusive under this lock. Mutating + # parametrizations under a concurrently executing conditioner is + # not safe; the create path's captures predate the backend, so + # they need no lock. + self._conditioner_lock = threading.Lock() + # D6a: the accelerated (possibly TRT) DiT the session was built + # with, and the always-resident eager module. _sync_dit_for_lora + # flips adapter.dit between them by LoRA activity. + self._dit_accel = adapter.dit if adapter is not None else None + self._dit_eager = eager_dit if eager_dit is not None else self._dit_accel + # Rendered-audio cache: one full decode+resample per fresh # latent (SAME-S decodes the whole window in ~11 ms); window # renders slice it, so gap-fill re-renders are bit-stable. @@ -367,6 +442,11 @@ def _source_encoder(waveform, sample_rate, sample_size): prompt_rebuilder=_prompt_rebuilder, prompt_tags=prompt, source_encoder=_source_encoder, + # D6a: the eager DiT module stays resident on the context + # even when make_dit returned a TRT wrapper — the interim + # LoRA fallback swaps to it. + eager_dit=context.dit, + model_id=str(context.model_id), **kwargs, ) @@ -400,9 +480,13 @@ def capabilities(self) -> Capabilities: # ACE prepare_source body, so duration/conditioning stay fixed. # render_anchor_queue: batch pad prewarm, drained by the shared # pipeline_runner exactly like the scalar stationary anchor. + # lora: on when the session asked for it (config.lora) AND the + # create path built the family manager — the same gate shape as + # ACE's use_lora bit. return Capabilities( refines_audio=True, loop_band=True, swap=True, render_anchor_queue=True, + lora=bool(self._use_lora and self._lora_mgr is not None), ) def geometry(self) -> AudioGeometry: @@ -414,7 +498,142 @@ def geometry(self) -> AudioGeometry: ) def knob_specs(self, lora_ids=()) -> list: - return sa3_knob_specs() + return sa3_knob_specs(loras=list(lora_ids or [])) + + def lora_compatible(self, metadata: dict) -> bool: + return sa3_lora_compatible(metadata, self._model_id) + + def has_pending_refit(self) -> bool: + """True when ``before_tick`` is about to apply LoRA commands — + the same pending-queue read as ACE's implementation, so the + runner pre-covers the enable stall (which can synchronously + materialize + register 200+ parametrizations).""" + if not self._use_lora or self.state is None: + return False + try: + with self.state._lock: + return bool( + self.state.pending_enable or self.state.pending_disable + ) + except AttributeError: + return False + + # ---- LoRA facade (D2 overrides; see notes/SA3_LORA_PLAN.md) --------- + + def _require_lora_manager(self): + if self._lora_mgr is None: + raise RuntimeError( + "SA3 backend has no LoRA manager; the session's " + "capability gate should have rejected this command" + ) + return self._lora_mgr + + def lora_available(self) -> bool: + return self._lora_mgr is not None + + def register_lora(self, path: str) -> str: + return self._require_lora_manager().register_lora(path) + + def prewarm_lora(self, lora_id: str): + return self._require_lora_manager().prewarm_lora(lora_id) + + def list_loras(self) -> list: + return self._lora_mgr.list_loras() if self._lora_mgr else [] + + def enable_lora(self, lora_id: str, strength=None) -> None: + """Transactional manager enable + the backend-side effects: + the D6a DiT swap and, for conditioner-targeting files, the D5 + cond-bundle rebuild. Runs on the runner thread inside the + pending-drain rendezvous (session._apply_lora_pending).""" + mgr = self._require_lora_manager() + # D5: conditioner mutation is mutually exclusive with + # conditioner execution (prompt-swap captures on the command + # thread). DiT-only files pay this lock for microseconds. + with self._conditioner_lock: + mgr.enable_lora(lora_id, strength=strength) + self._sync_dit_for_lora() + if mgr.touches_conditioner(lora_id): + self._rebuild_conditioning_after_lora("enable_lora", lora_id) + + def disable_lora(self, lora_id: str) -> None: + mgr = self._require_lora_manager() + # Read the conditioner flag BEFORE disable drops the staged + # payload. + touched = mgr.touches_conditioner(lora_id) + with self._conditioner_lock: + mgr.disable_lora(lora_id) + self._sync_dit_for_lora() + if touched: + self._rebuild_conditioning_after_lora("disable_lora", lora_id) + + def set_lora_strength(self, lora_id: str, strength: float) -> None: + mgr = self._require_lora_manager() + with self._conditioner_lock: + mgr.set_lora_strength(lora_id, strength) + # A strength change on a conditioner-targeting LoRA changes the + # conditioner's output too — the cached cond bundle must follow. + # (Per-tick slider sweeps on such files pay a T5Gemma capture + # per applied step; the common DiT-only case skips this + # entirely.) + if mgr.touches_conditioner(lora_id): + self._rebuild_conditioning_after_lora("set_lora_strength", lora_id) + + def _sync_dit_for_lora(self) -> None: + """D6a interim: LoRA weights live on the eager torch modules, so + a TRT-DiT session swaps to the eager DiT while any LoRA is + active and back when the last disables. Loud log both ways; + no-op for sessions already on the eager DiT.""" + if self.adapter is None or self._dit_eager is None: + return + if self._dit_accel is self._dit_eager: + return # eager session; nothing to swap + active = bool(self._lora_mgr is not None and self._lora_mgr.has_active_loras) + if active and self.adapter.dit is not self._dit_eager: + self.adapter.dit = self._dit_eager + logger.warning( + "sa3_dit_swap reason=lora_active dit=eager (TRT engine " + "cannot serve LoRA weights until the refit path lands)" + ) + elif not active and self.adapter.dit is not self._dit_accel: + self.adapter.dit = self._dit_accel + logger.info("sa3_dit_swap reason=lora_inactive dit=accelerated") + + def _rebuild_conditioning_after_lora(self, op: str, lora_id: str) -> None: + """D5: re-run the conditioner (through the exact prompt-swap + path, which takes the conditioner lock itself) after a + conditioner-mutating LoRA change commits, so the cached + cond_bundle reflects the mutated weights.""" + if self._prompt_rebuilder is None: + logger.warning( + "sa3_lora_cond_rebuild_skipped op={} id={} " + "reason=no_prompt_rebuilder", op, lora_id, + ) + return + tags = ( + getattr(self.state, "prompt_text", None) + if self.state is not None else None + ) + tags_b = ( + getattr(self.state, "prompt_text_b", None) + if self.state is not None else None + ) + if not tags: + # Directly-constructed backends without session state: fall + # back to the newest tags the cond history recorded. + with self._control_lock: + tags = self._cond_history[-1][2] + if not tags: + logger.warning( + "sa3_lora_cond_rebuild_skipped op={} id={} reason=no_tags", + op, lora_id, + ) + return + t0 = time.perf_counter() + self.handle_set_prompt(tags, tags_b=tags_b) + logger.info( + "sa3_lora_cond_rebuilt op={} id={} rebuild_ms={:.1f}", + op, lora_id, (time.perf_counter() - t0) * 1000, + ) def playable_duration_s(self): return self._cond.audio_sample_size / SA3_SAMPLE_RATE @@ -445,7 +664,11 @@ def handle_set_prompt(self, tags: str, *, tags_b: Optional[str] = None) -> None: "handle_set_prompt requires the from_context assembly" ) t0 = time.perf_counter() - cond, sched_factory = self._prompt_rebuilder(tags, self._steps) + # Conditioner EXECUTION under the D5 lock: a concurrent LoRA + # mutation of the conditioner's parametrizations (runner + # rendezvous) must not interleave with this capture. + with self._conditioner_lock: + cond, sched_factory = self._prompt_rebuilder(tags, self._steps) rebuild_ms = (time.perf_counter() - t0) * 1000 if int(cond.latent_frames) != int(self._cond.latent_frames): # Duration is fixed for the session lifetime, so the latent @@ -457,7 +680,8 @@ def handle_set_prompt(self, tags: str, *, tags_b: Optional[str] = None) -> None: f"duration is fixed per session" ) if tags_b and tags_b != tags: - cond_b, _ = self._prompt_rebuilder(tags_b, self._steps) + with self._conditioner_lock: + cond_b, _ = self._prompt_rebuilder(tags_b, self._steps) if int(cond_b.latent_frames) != int(cond.latent_frames): raise ValueError( f"sa3 prompt-B swap changed latent_frames " @@ -604,6 +828,25 @@ def _prepare_tick(self, knobs: dict, ctx: TickContext) -> dict: self.adapter.shift_alpha = shift self.pipeline.invalidate_schedule_cache() + # Per-LoRA live strength (the ACE per-tick convention): iterate + # the catalog so the active set can change at runtime; strength + # only flows for ENABLED entries, gated by the shared 0.02 + # slider-delta threshold. Under the parametrization engine this + # is a buffer write (no recompute); the gate also pre-guards the + # future TRT refit cost. + if self._use_lora and self._lora_mgr is not None: + for desc in self._lora_mgr.list_loras(): + if desc.state != "enabled": + continue + try: + lora_str = float( + knobs.get(f"lora_str_{desc.id}", desc.strength), + ) + except (TypeError, ValueError): + continue + if abs(lora_str - desc.strength) > 0.02: + self.set_lora_strength(desc.id, lora_str) + # Source-lock strength rides the shared override so a strength # bump engages the blend on in-flight slots submitted while it # was 0 — the ACE runner's exact per-tick convention. @@ -690,7 +933,21 @@ def _generate(self, prep: dict): # realizations (incoherent audio). See SlotRequest. sde_noise_seeded=True, )) - latent = self.pipeline.tick() # engine-layout [1, T, 256] | None + # With LoRAs active, wrap the tick in parametrize.cached(): + # without it, every weight ACCESS recomputes W + scaled delta, + # and the measured overhead is +49% per step at 1 LoRA / +114% + # at 3; under cached() each parametrized weight is computed once + # per tick and the 3-LoRA overhead collapses to ~+15% (Phase 0.5 + # entry benchmark — this wrap is a requirement, not a tweak). + # The cache is scoped to this context and strength changes only + # land between ticks, so it can never serve a stale weight. + if self._lora_mgr is not None and self._lora_mgr.has_active_loras: + import torch.nn.utils.parametrize as parametrize + + with parametrize.cached(): + latent = self.pipeline.tick() + else: + latent = self.pipeline.tick() # engine-layout [1, T, 256] | None if latent is not None: # The request this latent was generated from (valid only # right after a finishing tick — see StreamPipeline). @@ -842,6 +1099,21 @@ def render_full(self): pcm=self._rendered_audio(self._current_result), start_sample=0, ) + # ---- teardown ---------------------------------------------------------------- + + def close(self) -> None: + """Session teardown (called by StreamingSession.close). The SA3 + torch model is process-cached and shared with every later + session, so the LoRA manager must strip every parametrization + this session installed — the next session gets a pristine model + (plan D4 teardown; bitwise-validated by Phase 0.5 check 4).""" + if self._lora_mgr is not None: + try: + with self._conditioner_lock: + self._lora_mgr.close() + except Exception as exc: + logger.warning("sa3_lora_teardown_raised error={}", exc) + # ---- bookkeeping ------------------------------------------------------------- def on_fresh_generation(self, knobs: dict) -> None: diff --git a/acestep/streaming/sa3_session.py b/acestep/streaming/sa3_session.py index 55473de4..efc19099 100644 --- a/acestep/streaming/sa3_session.py +++ b/acestep/streaming/sa3_session.py @@ -172,6 +172,80 @@ def create_sa3_session( (SAMPLE_RATE, waveform), cond.audio_sample_size, ) + # ---- LoRA (notes/SA3_LORA_PLAN.md Phase 1) ------------------------ + # The family manager is constructed here — against the + # process-cached context's model, BEFORE StreamingSession.__init__ + # builds the backend — which settles the D2 construction-order item: + # startup catalog registration, alias resolution, and prewarm all + # run with a live manager, exactly like the ACE create path runs + # them against its engine_obj. The manager rides ``backend_init`` + # into families._make_sa3, so SA3Backend wraps the same instance. + from pathlib import Path + + from acestep.engine.sa3_lora import SA3LoRAManager + from acestep.lora_metadata import load_lora_metadata + from acestep.streaming.sa3_backend import sa3_lora_compatible + + use_lora = bool(config.lora) + lora_mgr = SA3LoRAManager( + model_root=context.sam.model.model, + conditioner_root=context.sam.model.conditioner, + checkpoint_id=model_id, + ) + lora_mgr.register_library() + initial_enable_ids: list[str] = [] + lora_strengths_init: dict[str, float] = dict(config.lora_strengths) + if use_lora: + from acestep.streaming.session import resolve_lora_reference + + # Same reference resolution as the runtime enable path, with + # the SA3 compatibility predicate applied directly (the backend + # doesn't exist until __init__) — mirroring how the ACE create + # path applies its scale axis. + entries = [] + for d in lora_mgr.list_loras(): + md = load_lora_metadata(d.path) + entries.append(( + d.id, + md.name or d.name, + sa3_lora_compatible(md.to_wire(), model_id), + )) + for lid in list(config.enabled_loras): + resolved = resolve_lora_reference(lid, entries) + if resolved is None: + logger.warning("lora_id_not_in_catalog id={}", lid) + continue + if resolved != lid: + logger.info( + "lora_alias_resolved requested={} id={}", lid, resolved, + ) + if lid in lora_strengths_init: + lora_strengths_init.setdefault( + resolved, lora_strengths_init[lid], + ) + if resolved not in initial_enable_ids: + initial_enable_ids.append(resolved) + for p in list(config.lora_paths): + pp = Path(p) + if not pp.exists(): + logger.warning("lora_path_missing path={}", p) + continue + try: + lid = lora_mgr.register_lora(str(pp)) + if lid not in initial_enable_ids: + initial_enable_ids.append(lid) + except Exception as e: + logger.exception( + "lora_register_failed path={} error={}", p, e, + ) + for lid in initial_enable_ids: + try: + lora_mgr.prewarm_lora(lid) + except Exception as e: + logger.exception("lora_prewarm_failed id={} error={}", lid, e) + if not initial_enable_ids: + logger.info("lora_startup_empty reason=catalog_only") + # Playable geometry comes from the conditioning capture (duration + # SA3's padding, what the DiT actually generates), not the request. playable_s = cond.audio_sample_size / context.sample_rate @@ -190,7 +264,9 @@ def create_sa3_session( src_np = src_np[:n_48k] n_channels = src_np.shape[1] if src_np.ndim > 1 else 1 - virtual_knobs = KnobState(sa3_knob_specs()) + virtual_knobs = KnobState(sa3_knob_specs( + loras=initial_enable_ids if use_lora else [], + )) state = SessionState( source=None, # no PreparedSource: swap/timbre/structure are capability-gated off @@ -218,6 +294,14 @@ def create_sa3_session( # tensors with no close() — GC reclaims them, same as the ACE # path's conditioning. with ExitStack() as cleanup: + # The manager holds no GPU state before an enable, but its + # prewarm executor + staged CPU tensors must not outlive a + # failed create. + cleanup.callback( + _cleanup_create_resource, + "sa3_lora_manager", + lora_mgr.close, + ) audio_eng = AudioEngine(src_np, SAMPLE_RATE) cleanup.callback( _cleanup_create_resource, @@ -242,9 +326,14 @@ def create_sa3_session( initial_upload_stems=None, initial_stem_error=None, initial_stem_source_mode=None, - initial_enable_ids=[], - lora_strengths_init={}, - lora_available=False, + initial_enable_ids=initial_enable_ids, + lora_strengths_init=lora_strengths_init, + # The manager exists whenever the context loaded (the eager + # torch model is always mutable), so the catalog surface is + # live even on sessions created with config.lora off — + # mirroring ACE, where lora_available reflects the engine + # and use_lora gates the commands. + lora_available=True, max_pipeline_depth=SA3_MAX_PIPELINE_DEPTH, max_seconds=playable_s, walk_window=False, @@ -252,7 +341,7 @@ def create_sa3_session( vae_window=SA3_VAE_WINDOW_S, crop_seconds=0.0, use_sde=False, - use_lora=False, + use_lora=use_lora, k1_name="sa3_denoise", # Construction payload for families._make_sa3 (the registry # factory assembles SA3Backend.from_context from this + the @@ -265,6 +354,7 @@ def create_sa3_session( "duration_s": duration_s, "dit_backend": dit_backend, "codec_backend": codec_backend, + "lora_manager": lora_mgr, }, ) cleanup.pop_all() diff --git a/acestep/streaming/session.py b/acestep/streaming/session.py index ca00dd16..7be3b42c 100644 --- a/acestep/streaming/session.py +++ b/acestep/streaming/session.py @@ -629,9 +629,12 @@ def _require_capability(self, capability: str, command: str) -> None: ) def _enabled_lora_ids(self) -> list: - if not (self.use_lora and self.engine_obj is not None): + # Backend facade (D2): the catalog lives behind the backend — + # ACE delegates to its engine_obj, SA3 to its parametrization + # manager. list_loras() is [] when no manager exists. + if not self.use_lora: return [] - return [d.id for d in self.engine_obj.list_loras() if d.state == "enabled"] + return [d.id for d in self.backend.list_loras() if d.state == "enabled"] # ---- Snapshot / catalog helpers ------------------------------------- @@ -690,7 +693,7 @@ def lora_catalog_payload(self) -> list: if not self.lora_available: return [] out = [] - for d in self.engine_obj.list_loras(): + for d in self.backend.list_loras(): # ``metadata`` is the full normalized record from the # LoRA's ``.metadata.json`` sidecar (falling back to # a synthesized record from ``.trigger.txt``, or a sparse @@ -719,10 +722,10 @@ def _resolve_lora_id(self, requested: str) -> str: string on a miss so downstream behavior (engine-side enable/disable failure logs) is unchanged for ids that never existed.""" - if not (self.lora_available and self.engine_obj is not None): + if not self.lora_available: return requested entries = [] - for d in self.engine_obj.list_loras(): + for d in self.backend.list_loras(): metadata = load_lora_metadata(d.path).to_wire() entries.append(( d.id, @@ -958,14 +961,14 @@ def _apply_lora_pending(self) -> None: return for lid in local_disable: try: - self.engine_obj.disable_lora(lid) + self.backend.disable_lora(lid) self.virtual_knobs.remove_knob(lora_strength_spec(lid).name) logger.info("lora_disabled id={}", lid) except Exception as e: logger.exception("lora_disable_failed id={} error={}", lid, e) for lid, strength in local_enable: try: - self.engine_obj.enable_lora(lid, strength=strength) + self.backend.enable_lora(lid, strength=strength) logger.info( "lora_enabled id={} strength={}", lid, strength, diff --git a/scripts/sa3/lora_smoke.py b/scripts/sa3/lora_smoke.py new file mode 100644 index 00000000..d55e2fbc --- /dev/null +++ b/scripts/sa3/lora_smoke.py @@ -0,0 +1,208 @@ +"""SA3 LoRA GPU smoke (notes/SA3_LORA_PLAN.md Phase 1). + +Drives the PRODUCTION manager (:class:`acestep.engine.sa3_lora.SA3LoRAManager`) +against a loaded medium SA3Context: register → prewarm → enable at +strength → DiT forward effect check → strength sweep 0→2 → disable → +VRAM accounting → teardown + re-session hygiene (the process-cached +model must come back bitwise pristine, twice). + +Pass ``--lora `` to smoke a real trained file (underfit output). +Without it, a synthetic rank-8 adapter is written against the live +module tree (DiT attention/FF + the seconds_total conditioner Linear) — +every mechanic is exercised identically; only musical quality needs a +trained file. + +Run: + .venv/Scripts/python.exe scripts/sa3/lora_smoke.py + .venv/Scripts/python.exe scripts/sa3/lora_smoke.py --lora path/to/file.safetensors +""" + +from __future__ import annotations + +import argparse +import sys +import time +from pathlib import Path + +_HERE = Path(__file__).resolve().parent +_REPO_ROOT = next(p for p in (_HERE, *_HERE.parents) if (p / "pyproject.toml").exists()) +for _p in (str(_REPO_ROOT),): + while _p in sys.path: + sys.path.remove(_p) + sys.path.insert(0, _p) + +import torch # noqa: E402 + + +def _vram_mb() -> float: + return torch.cuda.memory_allocated() / 1e6 if torch.cuda.is_available() else 0.0 + + +def _synthesize_lora(ctx, out_dir: Path, rank: int = 8) -> Path: + """Write a trainer-faithful synthetic file against the LIVE tree: + every self_attn/cross_attn/ff Linear in the DiT plus the + seconds_total conditioner Linear, fp16, lora_config in the header.""" + import torch.nn as nn + + from stable_audio_3.models.lora.utils import save_lora_safetensors + + g = torch.Generator().manual_seed(1528) + sd = {} + n = 0 + for name, mod in ctx.sam.model.model.named_modules(): + if not isinstance(mod, nn.Linear): + continue + if not any(k in name for k in ("self_attn", "cross_attn", ".ff.")): + continue + w2 = mod.weight.view(mod.weight.shape[0], -1) + prefix = f"{name}.parametrizations.weight.0" + sd[f"{prefix}.lora_A"] = torch.randn(rank, w2.shape[1], generator=g) * 0.05 + sd[f"{prefix}.lora_B"] = torch.randn(w2.shape[0], rank, generator=g) * 0.05 + n += 1 + for name, mod in ctx.sam.model.conditioner.named_modules(): + if isinstance(mod, nn.Linear): + w2 = mod.weight.view(mod.weight.shape[0], -1) + prefix = f"{name}.parametrizations.weight.0" + sd[f"{prefix}.lora_A"] = torch.randn(rank, w2.shape[1], generator=g) * 0.05 + sd[f"{prefix}.lora_B"] = torch.randn(w2.shape[0], rank, generator=g) * 0.05 + n += 1 + path = out_dir / "smoke_synthetic_lora.safetensors" + save_lora_safetensors( + sd, + {"rank": rank, "alpha": rank, "adapter_type": "lora", + "base_model": f"{ctx.model_id}-base"}, + path, + ) + print(f"synthesized {n}-module rank-{rank} adapter -> {path}") + return path + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--model", default="medium") + ap.add_argument("--lora", default=None, help="path to a trained SA3 LoRA") + args = ap.parse_args() + + from acestep.engine.sa3_context import SA3Context + from acestep.engine.sa3_lora import SA3LoRAManager + + print(f"loading SA3Context({args.model!r})...") + ctx = SA3Context(model_id=args.model) + sam = ctx.sam + + # Snapshot / probe helpers from the phase-0.5 harness (scripts/sa3 + # is on sys.path via ensure_sa3_paths, which SA3Context ran). + import lora_derisk_phase05 as h + + pristine = h.snapshot_state(sam) + cond = ctx.prepare_cond( + prompt="warm analog house groove, 124 bpm", duration=20.0, steps=8, + ) + from acestep.engine.sa3_stream_helpers import stack_sa3_cond_bundles + stacked = stack_sa3_cond_bundles([cond.cond_bundle]) + x = torch.randn( + 1, ctx.latent_channels, cond.latent_frames, + device=ctx.device, dtype=ctx.dtype, + generator=torch.Generator(device=ctx.device.type).manual_seed(7), + ) + t = torch.full((1,), 0.5, device=ctx.device, dtype=ctx.dtype) + + def forward(): + import torch.nn.utils.parametrize as parametrize + with torch.no_grad(), parametrize.cached(): + return ctx.dit(x, t, **stacked).clone() + + baseline = forward() + ok = True + + if args.lora: + lora_path = Path(args.lora) + else: + # Outside the repo AND outside every LoRA scan root (loras_dir + + # lora_extra_dirs), so the synthetic file can never show up in a + # real catalog. + from acestep.paths import models_dir + + out_dir = models_dir() / "sa3" / "smoke" + out_dir.mkdir(parents=True, exist_ok=True) + lora_path = _synthesize_lora(ctx, out_dir) + + def run_session(label: str) -> bool: + nonlocal ok + session_ok = True + mgr = SA3LoRAManager( + model_root=sam.model.model, + conditioner_root=sam.model.conditioner, + checkpoint_id=ctx.model_id, + ) + v0 = _vram_mb() + lid = mgr.register_lora(str(lora_path)) + mgr.prewarm_lora(lid).result(timeout=120) + d = mgr.get_lora(lid) + print(f"[{label}] materialized: {d.materialized_bytes / 1e6:.1f} MB, " + f"conditioner={mgr.touches_conditioner(lid)}") + + t0 = time.perf_counter() + mgr.enable_lora(lid, strength=1.0) + print(f"[{label}] enable_ms={(time.perf_counter() - t0) * 1000:.1f} " + f"vram_delta_mb={_vram_mb() - v0:.1f}") + + out_1 = forward() + if torch.equal(out_1, baseline): + print(f"[{label}] FAIL: enabled adapter had no forward effect") + session_ok = False + + # Strength sweep 0 -> 2 (the knob range): monotone engagement and + # exact identity at 0. + sweep_ms = [] + prev = None + for s in (0.0, 0.5, 1.0, 1.5, 2.0): + t0 = time.perf_counter() + mgr.set_lora_strength(lid, s) + sweep_ms.append((time.perf_counter() - t0) * 1000) + out_s = forward() + if s == 0.0 and not torch.equal(out_s, baseline): + print(f"[{label}] FAIL: strength 0 is not bit-identical to base") + session_ok = False + if prev is not None and torch.equal(out_s, prev): + print(f"[{label}] FAIL: strength {s} output identical to previous") + session_ok = False + prev = out_s + print(f"[{label}] strength sweep applied, max set_ms={max(sweep_ms):.2f}") + + v_enabled = _vram_mb() + mgr.disable_lora(lid) + out_off = forward() + if not torch.equal(out_off, baseline): + print(f"[{label}] FAIL: post-disable forward differs from baseline") + session_ok = False + print(f"[{label}] disable reclaimed {v_enabled - _vram_mb():.1f} MB") + + # Enable again, then close WITHOUT disabling — the session-death + # path. The process-cached model must come back pristine. + mgr.enable_lora(lid, strength=0.8) + mgr.close() + if not h.no_parametrizations_left(sam): + print(f"[{label}] FAIL: parametrizations left after close()") + session_ok = False + if not h.compare_state(sam, pristine, f"{label} post-close"): + session_ok = False + out_post = forward() + if not torch.equal(out_post, baseline): + print(f"[{label}] FAIL: post-close forward differs from baseline") + session_ok = False + print(f"[{label}] {'PASS' if session_ok else 'FAIL'}") + ok &= session_ok + return session_ok + + # Two back-to-back "sessions" on the same cached context — the + # cross-session hygiene contract. + run_session("session-1") + run_session("session-2") + + print(f"\n=== lora_smoke: {'PASS' if ok else 'FAIL'} ===") + return 0 if ok else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/unit/test_lora_facade_session.py b/tests/unit/test_lora_facade_session.py new file mode 100644 index 00000000..b7fd5343 --- /dev/null +++ b/tests/unit/test_lora_facade_session.py @@ -0,0 +1,168 @@ +"""Session-level LoRA plumbing through the backend facade (D2). + +The session's pending-drain (``_apply_lora_pending``), knob-manifest +rebuild, and enabled-id tracking must speak only to +``self.backend.*`` — never to the ACE ``engine_obj`` — so the SA3 +family (and any future one) rides the identical plumbing. Exercised +with a recording stub backend and unbound session methods: no GPU, no +model load, same pattern as test_lora_resolution. +""" + +from __future__ import annotations + +import sys +import threading +import types +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).parent.parent.parent)) + +from acestep.streaming.knobs import KnobState, lora_strength_spec +from acestep.streaming.session import StreamingSession + + +class _Desc: + def __init__(self, lora_id, state="registered", strength=0.0): + self.id = lora_id + self.name = lora_id + self.path = f"/nonexistent/{lora_id}.safetensors" + self.state = state + self.strength = strength + self.materialized_bytes = 0 + + +class _RecordingBackend: + """Facade-shaped stub: catalog + mutation recording. ``fail_ids`` + simulate engine-side enable failures (wrong family, bad file).""" + + def __init__(self, ids=(), fail_ids=()): + self._descs = {i: _Desc(i) for i in ids} + self.fail_ids = set(fail_ids) + self.calls: list = [] + + def list_loras(self): + return list(self._descs.values()) + + def enable_lora(self, lora_id, strength=None): + self.calls.append(("enable", lora_id, strength)) + if lora_id in self.fail_ids: + raise RuntimeError(f"stub enable failure for {lora_id}") + d = self._descs.setdefault(lora_id, _Desc(lora_id)) + d.state = "enabled" + if strength is not None: + d.strength = float(strength) + + def disable_lora(self, lora_id): + self.calls.append(("disable", lora_id)) + if lora_id in self._descs: + self._descs[lora_id].state = "registered" + + def knob_specs(self, lora_ids=()): + return [lora_strength_spec(lid) for lid in lora_ids] + + def lora_compatible(self, metadata): + return True + + +class _Bus: + def __init__(self): + self.events: list = [] + + def publish(self, event): + self.events.append(event) + + +class _DrainSession: + """Just enough StreamingSession surface for _apply_lora_pending.""" + + lora_available = True + use_lora = True + + _apply_lora_pending = StreamingSession._apply_lora_pending + _enabled_lora_ids = StreamingSession._enabled_lora_ids + _rebuild_knob_specs = StreamingSession._rebuild_knob_specs + lora_catalog_payload = StreamingSession.lora_catalog_payload + + def __init__(self, backend): + self.backend = backend + self.state = types.SimpleNamespace( + _lock=threading.Lock(), + pending_enable=[], + pending_disable=[], + ) + self.virtual_knobs = KnobState([]) + self.bus = _Bus() + self._knob_specs_by_name = {} + + +def test_drain_enables_through_facade_and_allocates_knob(): + be = _RecordingBackend(ids=["ambient"]) + ss = _DrainSession(be) + ss.state.pending_enable.append(("ambient", 0.7)) + + ss._apply_lora_pending() + + assert ("enable", "ambient", 0.7) in be.calls + # Knob slot allocated with the enable strength as its default. + values = ss.virtual_knobs.get_all_values() + assert values["lora_str_ambient"] == 0.7 + # The cached validation spec map was rebuilt from the backend's + # manifest for the new enabled set. + assert "lora_str_ambient" in ss._knob_specs_by_name + # Catalog broadcast published. + assert len(ss.bus.events) == 1 + catalog = ss.bus.events[0].catalog + assert [e["id"] for e in catalog] == ["ambient"] + assert catalog[0]["state"] == "enabled" + + +def test_drain_disable_removes_knob_and_rebuilds(): + be = _RecordingBackend(ids=["ambient"]) + ss = _DrainSession(be) + ss.state.pending_enable.append(("ambient", 1.0)) + ss._apply_lora_pending() + assert "lora_str_ambient" in ss.virtual_knobs.get_all_values() + + ss.state.pending_disable.append("ambient") + ss._apply_lora_pending() + + assert ("disable", "ambient") in be.calls + assert "lora_str_ambient" not in ss.virtual_knobs.get_all_values() + assert "lora_str_ambient" not in ss._knob_specs_by_name + + +def test_drain_enable_failure_is_contained(): + """A backend-side enable failure (wrong family, bad file) is logged, + allocates no knob, and doesn't break the drain for other entries.""" + be = _RecordingBackend(ids=["good", "bad"], fail_ids=["bad"]) + ss = _DrainSession(be) + ss.state.pending_enable.append(("bad", 1.0)) + ss.state.pending_enable.append(("good", 0.5)) + + ss._apply_lora_pending() # must not raise + + assert ("enable", "bad", 1.0) in be.calls + assert ("enable", "good", 0.5) in be.calls + values = ss.virtual_knobs.get_all_values() + assert "lora_str_bad" not in values + assert values["lora_str_good"] == 0.5 + assert ss._enabled_lora_ids() == ["good"] + + +def test_drain_noop_without_pending(): + be = _RecordingBackend(ids=["ambient"]) + ss = _DrainSession(be) + ss._apply_lora_pending() + assert be.calls == [] + assert ss.bus.events == [] + + +def test_enabled_lora_ids_reads_backend_catalog(): + be = _RecordingBackend(ids=["a", "b"]) + ss = _DrainSession(be) + assert ss._enabled_lora_ids() == [] + be.enable_lora("b", strength=1.0) + assert ss._enabled_lora_ids() == ["b"] + ss2 = _DrainSession(be) + ss2.use_lora = False + assert ss2._enabled_lora_ids() == [] diff --git a/tests/unit/test_lora_resolution.py b/tests/unit/test_lora_resolution.py index 29a4fb42..d66c7bc9 100644 --- a/tests/unit/test_lora_resolution.py +++ b/tests/unit/test_lora_resolution.py @@ -108,18 +108,16 @@ def __init__(self, path: Path): self.materialized_bytes = 0 -class _StubEngine: +class _ScaleBackend: + """Backend stub carrying the catalog (the D2 facade's list_loras) + and the ACE scale-axis predicate for a 2B checkpoint.""" + def __init__(self, descs): self._descs = descs def list_loras(self): return self._descs - -class _ScaleBackend: - """Backend stub whose predicate is the ACE scale axis for a 2B - checkpoint.""" - def lora_compatible(self, metadata: dict) -> bool: return lora_scale_compatible(metadata.get("base_model_scale"), "2B") @@ -128,8 +126,7 @@ class _StubSession: lora_available = True def __init__(self, descs): - self.engine_obj = _StubEngine(descs) - self.backend = _ScaleBackend() + self.backend = _ScaleBackend(descs) lora_catalog_payload = StreamingSession.lora_catalog_payload _resolve_lora_id = StreamingSession._resolve_lora_id @@ -187,9 +184,9 @@ def test_resolve_lora_id_falls_through_on_miss(self, tmp_path): assert ss._resolve_lora_id("Metalcore") == "Metalcore" assert ss._resolve_lora_id("nope") == "nope" - def test_resolve_lora_id_without_engine_is_identity(self, tmp_path): + def test_resolve_lora_id_without_lora_available_is_identity(self, tmp_path): ss = _stub_session(tmp_path) - ss.engine_obj = None + ss.lora_available = False assert ss._resolve_lora_id("Ambient") == "Ambient" @@ -266,7 +263,7 @@ def test_unknown_family_stays_permissive(self): assert ACEStepBackend.lora_compatible(_AceStub("2B"), {}) -class _FamilyBackend: +class _FamilyBackend(_ScaleBackend): """Stub backend running the REAL ACE predicate (family + scale).""" checkpoint_scale = "2B" @@ -276,11 +273,11 @@ class _FamilyBackend: class TestMixedFamilyCatalog: def test_catalog_annotates_sa3_entry_incompatible(self, tmp_path): clear_cache() - ss = _StubSession([ + ss = _StubSession([]) + ss.backend = _FamilyBackend([ _Desc(_write_lora(tmp_path, "ambient-v1", "Ambient", "2B")), _Desc(_write_sa3_file(tmp_path, "sa3style")), ]) - ss.backend = _FamilyBackend() verdicts = { e["id"]: e["compatible"] for e in ss.lora_catalog_payload() } @@ -291,8 +288,8 @@ def test_alias_of_sa3_only_name_misses(self, tmp_path): display-name reference to an SA3-only entry misses; the exact id still passes through (and fails loudly at the enable boundary).""" clear_cache() - ss = _StubSession([_Desc(_write_sa3_file(tmp_path, "sa3style"))]) - ss.backend = _FamilyBackend() + ss = _StubSession([]) + ss.backend = _FamilyBackend([_Desc(_write_sa3_file(tmp_path, "sa3style"))]) assert ss._resolve_lora_id("sa3style") == "sa3style" # exact id assert ss._resolve_lora_id("SA3STYLE") == "SA3STYLE" # alias: miss diff --git a/tests/unit/test_lora_sa3_manager.py b/tests/unit/test_lora_sa3_manager.py new file mode 100644 index 00000000..ca864e5d --- /dev/null +++ b/tests/unit/test_lora_sa3_manager.py @@ -0,0 +1,458 @@ +"""Unit tests for :mod:`acestep.engine.sa3_lora` (SA3LoRAManager). + +CPU-only, but against the REAL vendored ``stable_audio_3`` LoRA code +(parametrize registration, strength buffers, upstream removal) on a +tiny module tree — the manager's whole value is wrapping that engine's +exact math and sharp edges, so stubbing it out would test nothing. +Skips when the managed vendor checkout isn't on disk (CI without +models), same convention as the checkpoint-gated tests. + +Covers the plan's Phase 1 test contract: enable-at-strength atomicity, +rollback on failure, middle-slot removal/re-enable, strength targeting +after churn, the conditioner-dirty flag, ``-xs`` rejection, wrong-family +rejection, lineage validation, and teardown completeness. +""" + +from __future__ import annotations + +import sys +from pathlib import Path + +import pytest +import torch +import torch.nn as nn + +sys.path.insert(0, str(Path(__file__).parent.parent.parent)) + + +def _vendored_or_skip(): + from acestep.engine.sa3_helpers import ensure_sa3_paths + + ensure_sa3_paths() + lora_utils = pytest.importorskip( + "stable_audio_3.models.lora.utils", + reason="SA3 vendored source not on disk", + ) + from stable_audio_3.models.lora import model as lora_model + + return lora_model, lora_utils + + +def _make_manager(): + """Tiny two-root tree mirroring the SA3 layout: a 'DiT' root with a + Linear + Conv1d, and a conditioner root with one Linear under the + conditioners.* prefix (the seconds_total shape).""" + from acestep.engine.sa3_lora import SA3LoRAManager + + torch.manual_seed(7) + model = nn.Module() + blk = nn.Module() + blk.lin = nn.Linear(16, 8, bias=False) + model.blk = blk + model.conv = nn.Conv1d(4, 6, 3, bias=False) + + cond = nn.Module() + conds = nn.Module() + seconds = nn.Module() + seconds.lin = nn.Linear(5, 8, bias=False) + conds.seconds_total = seconds + cond.conditioners = conds + + mgr = SA3LoRAManager( + model_root=model, conditioner_root=cond, checkpoint_id="medium", + ) + return mgr, model, cond + + +def _write_lora( + tmp_path: Path, + stem: str, + entries: dict, + *, + rank: int = 4, + alpha: float = 8.0, + adapter_type: str = "lora", + base_model: str | None = "medium-base", + seed: int = 0, +) -> Path: + """Trainer-faithful file: fp16 tensors, parametrization keys at + index 0, lora_config in the safetensors header (the vendored + ``save_lora_safetensors`` does the fp16 cast + header embed). + + ``entries``: {fqn: (fan_out, fan_in)} — random A/B pairs are built + per module from ``seed``. + """ + _, lora_utils = _vendored_or_skip() + g = torch.Generator().manual_seed(seed) + sd = {} + for fqn, (fan_out, fan_in) in entries.items(): + prefix = f"{fqn}.parametrizations.weight.0" + sd[f"{prefix}.lora_A"] = torch.randn(rank, fan_in, generator=g) + sd[f"{prefix}.lora_B"] = torch.randn(fan_out, rank, generator=g) + config = {"rank": rank, "alpha": alpha, "adapter_type": adapter_type} + if base_model is not None: + config["base_model"] = base_model + p = tmp_path / f"{stem}.safetensors" + lora_utils.save_lora_safetensors(sd, config, p) + return p + + +def _adapters_on(mod): + """(lora_index, scaling, strength, A, B) per parametrization, in + physical order.""" + out = [] + plist = getattr(getattr(mod, "parametrizations", None), "weight", None) + for p in plist or []: + out.append(( + p.lora_index, p.scaling, float(p.lora_strength), + p.lora_A.detach(), p.lora_B.detach(), + )) + return out + + +def _expected_weight(base: torch.Tensor, mod) -> torch.Tensor: + """Replay the vendored lora_forward chain from the LIVE adapter + tensors (physical order, same op order).""" + w = base.clone() + for _idx, scaling, strength, A, B in _adapters_on(mod): + delta = (B @ A).view(w.shape) + w = w + (scaling * strength * delta).to(w.dtype) + return w + + +# --------------------------------------------------------------------------- +# Enable / strength / disable basics +# --------------------------------------------------------------------------- + + +def test_enable_at_strength_applies_scaled_delta(tmp_path): + mgr, model, _cond = _make_manager() + base = model.blk.lin.weight.detach().clone() + p = _write_lora(tmp_path, "styleA", {"blk.lin": (8, 16)}) + + lid = mgr.register_lora(str(p)) + mgr.enable_lora(lid, strength=0.5) + + assert mgr.get_lora(lid).state == "enabled" + assert mgr.get_lora(lid).strength == 0.5 + adapters = _adapters_on(model.blk.lin) + assert len(adapters) == 1 + _idx, scaling, strength, A, B = adapters[0] + assert scaling == pytest.approx(8.0 / 4.0) # alpha/rank, not ACE's raw B@A + assert strength == 0.5 + # Non-trivial delta actually landed in the effective weight. + assert torch.allclose( + model.blk.lin.weight, _expected_weight(base, model.blk.lin), + ) + assert not torch.allclose(model.blk.lin.weight, base) + + +def test_conv1d_target_applies(tmp_path): + mgr, model, _cond = _make_manager() + base = model.conv.weight.detach().clone() + # Conv1d weight [6, 4, 3] flattens to (6, 12). + p = _write_lora(tmp_path, "convy", {"conv": (6, 12)}) + lid = mgr.register_lora(str(p)) + mgr.enable_lora(lid, strength=1.0) + assert torch.allclose( + model.conv.weight, _expected_weight(base, model.conv), + ) + assert not torch.allclose(model.conv.weight, base) + + +def test_set_strength_is_live_buffer_write(tmp_path): + mgr, model, _cond = _make_manager() + base = model.blk.lin.weight.detach().clone() + p = _write_lora(tmp_path, "s", {"blk.lin": (8, 16)}) + lid = mgr.register_lora(str(p)) + mgr.enable_lora(lid, strength=1.0) + w_at_1 = model.blk.lin.weight.detach().clone() + + mgr.set_lora_strength(lid, 0.25) + assert float(_adapters_on(model.blk.lin)[0][2]) == 0.25 + w_at_quarter = model.blk.lin.weight.detach() + assert torch.allclose( + w_at_quarter - base, (w_at_1 - base) * 0.25, atol=1e-6, + ) + + +def test_set_strength_requires_enabled(tmp_path): + mgr, _model, _cond = _make_manager() + p = _write_lora(tmp_path, "s2", {"blk.lin": (8, 16)}) + lid = mgr.register_lora(str(p)) + with pytest.raises(ValueError, match="not enabled"): + mgr.set_lora_strength(lid, 0.5) + + +def test_disable_restores_base_bitwise_and_frees_state(tmp_path): + mgr, model, _cond = _make_manager() + base = model.blk.lin.weight.detach().clone() + p = _write_lora(tmp_path, "d", {"blk.lin": (8, 16)}) + lid = mgr.register_lora(str(p)) + mgr.enable_lora(lid, strength=1.0) + + mgr.disable_lora(lid) + + # Upstream removal restores the ORIGINAL tensor — bitwise. + assert torch.equal(model.blk.lin.weight, base) + assert not hasattr(model.blk.lin, "parametrizations") + assert mgr.get_lora(lid).state == "registered" + # Strength survives the cycle (slider-position convention). + assert mgr.get_lora(lid).strength == 1.0 + mgr.enable_lora(lid) + assert mgr.get_lora(lid).strength == 1.0 + + +def test_prewarm_stages_without_touching_model(tmp_path): + mgr, model, _cond = _make_manager() + p = _write_lora(tmp_path, "pw", {"blk.lin": (8, 16)}) + lid = mgr.register_lora(str(p)) + + f = mgr.prewarm_lora(lid) + f.result(timeout=10) + + assert mgr.get_lora(lid).state == "materialized" + assert mgr.get_lora(lid).materialized_bytes > 0 + assert not hasattr(model.blk.lin, "parametrizations") + assert mgr.touches_conditioner(lid) is False + + +# --------------------------------------------------------------------------- +# Hard validation at materialize/enable (D3) +# --------------------------------------------------------------------------- + + +def test_ace_peft_file_fails_loudly(tmp_path): + from safetensors.torch import save_file + + mgr, _model, _cond = _make_manager() + p = tmp_path / "acefile.safetensors" + save_file( + { + "base_model.model.q.lora_A.weight": torch.zeros(4, 16), + "base_model.model.q.lora_B.weight": torch.zeros(8, 4), + }, + str(p), + ) + lid = mgr.register_lora(str(p)) + with pytest.raises(RuntimeError, match="not an SA3 LoRA") as ei: + mgr.enable_lora(lid, strength=1.0) + assert "ACE" in str(ei.value) + + +def test_xs_variant_rejected_with_clear_error(tmp_path): + _, lora_utils = _vendored_or_skip() + mgr, _model, _cond = _make_manager() + sd = { + "blk.lin.parametrizations.weight.0.M_xs": torch.zeros(4, 4), + } + p = tmp_path / "xsfile.safetensors" + lora_utils.save_lora_safetensors( + sd, {"rank": 4, "alpha": 4, "adapter_type": "lora-xs"}, p, + ) + lid = mgr.register_lora(str(p)) + with pytest.raises(RuntimeError, match="not .*supported|not yet"): + mgr.enable_lora(lid, strength=1.0) + + +def test_lineage_mismatch_fails_with_checkpoint_names(tmp_path): + mgr, _model, _cond = _make_manager() # checkpoint_id="medium" + p = _write_lora( + tmp_path, "wronglineage", {"blk.lin": (8, 16)}, + base_model="small-music-base", + ) + lid = mgr.register_lora(str(p)) + with pytest.raises(RuntimeError, match="medium") as ei: + mgr.enable_lora(lid, strength=1.0) + assert "small-music" in str(ei.value) + + +def test_unknown_lineage_is_permissive(tmp_path): + mgr, _model, _cond = _make_manager() + p = _write_lora( + tmp_path, "custom", {"blk.lin": (8, 16)}, + base_model="my-custom-finetune", + ) + lid = mgr.register_lora(str(p)) + mgr.enable_lora(lid, strength=1.0) + assert mgr.get_lora(lid).state == "enabled" + + +def test_shape_mismatch_everywhere_fails(tmp_path): + mgr, _model, _cond = _make_manager() + # fan_in 32 doesn't fit blk.lin's (8, 16). + p = _write_lora(tmp_path, "wrongsize", {"blk.lin": (8, 32)}) + lid = mgr.register_lora(str(p)) + with pytest.raises(RuntimeError, match="does not fit"): + mgr.enable_lora(lid, strength=1.0) + + +def test_ckpt_extension_rejected(tmp_path): + mgr, _model, _cond = _make_manager() + p = tmp_path / "old.ckpt" + p.write_bytes(b"\x80\x04") # never actually loaded + lid = mgr.register_lora(str(p)) + with pytest.raises(RuntimeError, match="safetensors"): + mgr.enable_lora(lid, strength=1.0) + + +# --------------------------------------------------------------------------- +# Transactional enable + rollback (D4) +# --------------------------------------------------------------------------- + + +def test_failed_enable_rolls_back_bitwise(tmp_path, monkeypatch): + lora_model, _ = _vendored_or_skip() + mgr, model, cond = _make_manager() + pA = _write_lora(tmp_path, "okA", {"blk.lin": (8, 16)}, seed=1) + lidA = mgr.register_lora(str(pA)) + mgr.enable_lora(lidA, strength=1.0) + + pre_params = { + n: t.detach().clone() + for root in (model, cond) + for n, t in root.named_parameters() + } + pre_effective = model.blk.lin.weight.detach().clone() + + pB = _write_lora( + tmp_path, "failB", {"blk.lin": (8, 16), "conv": (6, 12)}, seed=2, + ) + lidB = mgr.register_lora(str(pB)) + + # Fail AFTER registration + install, at the strength-buffer step — + # the deepest point before commit. + real_set = lora_model.set_lora_strength + + def _boom(*a, **k): + raise RuntimeError("simulated failure at strength write") + + monkeypatch.setattr(lora_model, "set_lora_strength", _boom) + with pytest.raises(RuntimeError, match="simulated failure"): + mgr.enable_lora(lidB, strength=0.7) + monkeypatch.setattr(lora_model, "set_lora_strength", real_set) + + # Bit-identical parameter set (names + values) vs pre-enable. + now_params = { + n: t.detach() + for root in (model, cond) + for n, t in root.named_parameters() + } + assert set(now_params) == set(pre_params) + for n, t in now_params.items(): + assert torch.equal(t, pre_params[n]), n + assert torch.equal(model.blk.lin.weight, pre_effective) + # B is not enabled; A still is and stays targetable. + assert mgr.get_lora(lidB).state == "materialized" + assert mgr.get_lora(lidA).state == "enabled" + mgr.set_lora_strength(lidA, 0.3) + assert float(_adapters_on(model.blk.lin)[0][2]) == pytest.approx(0.3) + # The failed transaction's index is never reused, and a retry works. + mgr.enable_lora(lidB, strength=0.7) + assert mgr.get_lora(lidB).state == "enabled" + + +# --------------------------------------------------------------------------- +# Slot churn (D4) — the remove_lora_by_index physical-shift hazard +# --------------------------------------------------------------------------- + + +def test_middle_slot_removal_then_enable_stays_id_accurate(tmp_path): + mgr, model, _cond = _make_manager() + base = model.blk.lin.weight.detach().clone() + lids = [] + for i, stem in enumerate(("a", "b", "c")): + p = _write_lora(tmp_path, stem, {"blk.lin": (8, 16)}, seed=10 + i) + lid = mgr.register_lora(str(p)) + mgr.enable_lora(lid, strength=1.0) + lids.append(lid) + + # Remove the MIDDLE adapter — physical positions shift under + # upstream removal. + mgr.disable_lora("b") + assert len(_adapters_on(model.blk.lin)) == 2 + + # Enable a fourth AFTER the shift (direct-copy install path). + pD = _write_lora(tmp_path, "d", {"blk.lin": (8, 16)}, seed=99) + mgr.enable_lora(mgr.register_lora(str(pD)), strength=1.0) + assert len(_adapters_on(model.blk.lin)) == 3 + assert torch.allclose( + model.blk.lin.weight, _expected_weight(base, model.blk.lin), + ) + + # Strength targeting by id lands on the right adapter objects. + mgr.set_lora_strength("c", 0.5) + mgr.set_lora_strength("d", 0.25) + by_index = { + idx: s for idx, _sc, s, _A, _B in _adapters_on(model.blk.lin) + } + assert by_index[mgr._index_by_id["a"]] == 1.0 + assert by_index[mgr._index_by_id["c"]] == 0.5 + assert by_index[mgr._index_by_id["d"]] == 0.25 + assert torch.allclose( + model.blk.lin.weight, _expected_weight(base, model.blk.lin), + ) + + # Full teardown of the churned state restores base bitwise. + for lid in ("a", "c", "d"): + mgr.disable_lora(lid) + assert torch.equal(model.blk.lin.weight, base) + assert not hasattr(model.blk.lin, "parametrizations") + + +# --------------------------------------------------------------------------- +# Conditioner flag (D5) + teardown (close) +# --------------------------------------------------------------------------- + + +def test_touches_conditioner_flag(tmp_path): + mgr, _model, cond = _make_manager() + base = cond.conditioners.seconds_total.lin.weight.detach().clone() + p = _write_lora( + tmp_path, "condy", + {"blk.lin": (8, 16), "conditioners.seconds_total.lin": (8, 5)}, + ) + lid = mgr.register_lora(str(p)) + mgr.prewarm_lora(lid).result(timeout=10) + assert mgr.touches_conditioner(lid) is True + + mgr.enable_lora(lid, strength=1.0) + assert not torch.allclose( + cond.conditioners.seconds_total.lin.weight, base, + ) + mgr.disable_lora(lid) + assert torch.equal(cond.conditioners.seconds_total.lin.weight, base) + # Staged payload dropped with the disable — flag is gone too. + assert mgr.touches_conditioner(lid) is False + + +def test_close_returns_model_to_pristine(tmp_path): + mgr, model, cond = _make_manager() + pre = { + n: t.detach().clone() + for root in (model, cond) + for n, t in root.named_parameters() + } + for i, stem in enumerate(("x", "y")): + p = _write_lora( + tmp_path, stem, + {"blk.lin": (8, 16), "conditioners.seconds_total.lin": (8, 5)}, + seed=20 + i, + ) + mgr.enable_lora(mgr.register_lora(str(p)), strength=1.0) + + mgr.close() + + now = { + n: t.detach() + for root in (model, cond) + for n, t in root.named_parameters() + } + assert set(now) == set(pre) + for n, t in now.items(): + assert torch.equal(t, pre[n]), n + for root in (model, cond): + for _n, mod in root.named_modules(): + assert not getattr(mod, "parametrizations", None) + assert mgr.list_loras() == [] + mgr.close() # idempotent diff --git a/tests/unit/test_sa3_backend.py b/tests/unit/test_sa3_backend.py index 9025cb5d..26b90d73 100644 --- a/tests/unit/test_sa3_backend.py +++ b/tests/unit/test_sa3_backend.py @@ -447,3 +447,216 @@ def test_decode_seed_follows_the_emerged_request_seed(): b.produce(_knobs(b, seed=1234), CTX, "reuse") b.render_window(0.0) assert b.codec.decode_seeds[-1] == 99 + + +# --------------------------------------------------------------------------- +# LoRA surface (notes/SA3_LORA_PLAN.md phase 1) — fake manager, real wiring +# --------------------------------------------------------------------------- + +from acestep.streaming.sa3_backend import sa3_lora_compatible + + +class _FakeLoraDesc: + def __init__(self, lora_id, state="registered", strength=0.0): + self.id = lora_id + self.name = lora_id + self.path = f"/nonexistent/{lora_id}.safetensors" + self.state = state + self.strength = strength + self.materialized_bytes = 0 + + +class _FakeLoraManager: + def __init__(self, ids=(), touches=()): + self._descs = {i: _FakeLoraDesc(i) for i in ids} + self._touches = set(touches) + self.calls: list = [] + self.closed = False + + def list_loras(self): + return list(self._descs.values()) + + @property + def has_active_loras(self): + return any(d.state == "enabled" for d in self._descs.values()) + + def enable_lora(self, lora_id, strength=None): + self.calls.append(("enable", lora_id, strength)) + d = self._descs.setdefault(lora_id, _FakeLoraDesc(lora_id)) + d.state = "enabled" + if strength is not None: + d.strength = float(strength) + + def disable_lora(self, lora_id): + self.calls.append(("disable", lora_id)) + if lora_id in self._descs: + self._descs[lora_id].state = "registered" + + def set_lora_strength(self, lora_id, strength): + self.calls.append(("strength", lora_id, strength)) + self._descs[lora_id].strength = float(strength) + + def touches_conditioner(self, lora_id): + return lora_id in self._touches + + def close(self): + self.closed = True + + +def test_lora_capability_bit_requires_manager_and_toggle(): + assert _backend().capabilities().lora is False + assert _backend( + lora_manager=_FakeLoraManager(), use_lora=False, + ).capabilities().lora is False + assert _backend(use_lora=True).capabilities().lora is False # no manager + assert _backend( + lora_manager=_FakeLoraManager(), use_lora=True, + ).capabilities().lora is True + + +def test_lora_knob_specs_expand_per_id(): + b = _backend(lora_manager=_FakeLoraManager(), use_lora=True) + names = [s.name for s in b.knob_specs(["ambient", "phonk"])] + assert "lora_str_ambient" in names + assert "lora_str_phonk" in names + # Base manifest unchanged and first (display order). + assert names[0] == "sa3_denoise" + + +def test_sa3_lora_compatible_family_and_lineage(): + assert sa3_lora_compatible({"lora_family": "ace"}, "medium") is False + assert sa3_lora_compatible( + {"lora_family": "sa3", "base_model": "medium-base"}, "medium", + ) is True + assert sa3_lora_compatible( + {"lora_family": "sa3", "base_model": "small-music-base"}, "medium", + ) is False + # Unknown on either side stays permissive. + assert sa3_lora_compatible({}, "medium") is True + assert sa3_lora_compatible( + {"lora_family": "sa3", "base_model": "mystery"}, "medium", + ) is True + assert sa3_lora_compatible( + {"lora_family": "sa3", "base_model": "medium-base"}, "", + ) is True + + +def test_prepare_tick_drives_strength_from_knob(): + mgr = _FakeLoraManager(ids=["x"]) + mgr.enable_lora("x", strength=1.0) + b = _backend(lora_manager=mgr, use_lora=True) + + b.produce(_knobs(b, **{"lora_str_x": 0.5}), CTX, "skip") + assert ("strength", "x", 0.5) in mgr.calls + + # Within the 0.02 slider-delta gate: no engine call. + n_calls = len(mgr.calls) + b.produce(_knobs(b, **{"lora_str_x": 0.51}), CTX, "skip") + assert len(mgr.calls) == n_calls + + # Non-enabled entries are ignored even if a stray knob rides in. + mgr.disable_lora("x") + n_calls = len(mgr.calls) + b.produce(_knobs(b, **{"lora_str_x": 1.7}), CTX, "skip") + assert not any(c[0] == "strength" for c in mgr.calls[n_calls:]) + + +def test_dit_swaps_to_eager_while_lora_active(): + mgr = _FakeLoraManager(ids=["x"]) + b = _backend(lora_manager=mgr, use_lora=True, eager_dit=_ZeroDit()) + accel = b._dit_accel + eager = b._dit_eager + assert b.adapter.dit is accel + + b.enable_lora("x", strength=1.0) + assert b.adapter.dit is eager + + b.disable_lora("x") + assert b.adapter.dit is accel + + +def test_eager_session_dit_swap_is_noop(): + mgr = _FakeLoraManager(ids=["x"]) + b = _backend(lora_manager=mgr, use_lora=True) # no separate eager dit + dit = b.adapter.dit + b.enable_lora("x", strength=1.0) + assert b.adapter.dit is dit + b.disable_lora("x") + assert b.adapter.dit is dit + + +def test_conditioner_lora_triggers_cond_rebuild(monkeypatch): + mgr = _FakeLoraManager(ids=["c"], touches=["c"]) + b = _backend(lora_manager=mgr, use_lora=True) + rebuilds: list = [] + monkeypatch.setattr( + b, "_rebuild_conditioning_after_lora", + lambda op, lid: rebuilds.append((op, lid)), + ) + + b.enable_lora("c", strength=1.0) + b.set_lora_strength("c", 0.5) + b.disable_lora("c") + + assert rebuilds == [ + ("enable_lora", "c"), + ("set_lora_strength", "c"), + ("disable_lora", "c"), + ] + + +def test_dit_only_lora_skips_cond_rebuild(monkeypatch): + mgr = _FakeLoraManager(ids=["d"]) + b = _backend(lora_manager=mgr, use_lora=True) + rebuilds: list = [] + monkeypatch.setattr( + b, "_rebuild_conditioning_after_lora", + lambda op, lid: rebuilds.append((op, lid)), + ) + b.enable_lora("d", strength=1.0) + b.disable_lora("d") + assert rebuilds == [] + + +def test_cond_rebuild_goes_through_prompt_swap_path(monkeypatch): + import types + + mgr = _FakeLoraManager(ids=["c"], touches=["c"]) + b = _backend(lora_manager=mgr, use_lora=True) + b._prompt_rebuilder = object() # present: rebuild not skipped + b.state = types.SimpleNamespace( + prompt_text="warm tags", prompt_text_b="other tags", + ) + swaps: list = [] + monkeypatch.setattr( + b, "handle_set_prompt", + lambda tags, tags_b=None: swaps.append((tags, tags_b)), + ) + + b._rebuild_conditioning_after_lora("enable_lora", "c") + assert swaps == [("warm tags", "other tags")] + + +def test_lora_pending_gates_has_pending_refit(): + import threading + import types + + mgr = _FakeLoraManager(ids=["x"]) + b = _backend(lora_manager=mgr, use_lora=True) + b.state = types.SimpleNamespace( + _lock=threading.Lock(), pending_enable=[], pending_disable=[], + interp_feedback="slerp", + ) + assert b.has_pending_refit() is False + b.state.pending_enable.append(("x", 1.0)) + assert b.has_pending_refit() is True + b.state.pending_enable.clear() + b.state.pending_disable.append("x") + assert b.has_pending_refit() is True + + +def test_close_tears_down_lora_manager(): + mgr = _FakeLoraManager(ids=["x"]) + b = _backend(lora_manager=mgr, use_lora=True) + b.close() + assert mgr.closed is True From 8f3761e277a507adddf87d8a8ef8edbd5c14172a Mon Sep 17 00:00:00 2001 From: RyanOnTheInside <7623207+ryanontheinside@users.noreply.github.com> Date: Wed, 22 Jul 2026 18:22:42 -0400 Subject: [PATCH 4/6] =?UTF-8?q?feat(sa3-lora):=20TRT=20refit=20path=20?= =?UTF-8?q?=E2=80=94=20refittable=20engines,=20manifest=20tool,=20refit=20?= =?UTF-8?q?mirror=20(phase=202=20code)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - trt/refit_core (new): shared typed-Weights push + commit/missing- weights diagnostics; ACE's TRTLoRAManager now routes its pushes through it (byte-identical, its unit tests untouched and green) - sa3_build: --refit builds sa3_m_dit_refit_l* engines with BuilderFlag.REFIT via a separate config dataclass so existing engine metadata identity is untouched (dry-run verified) - sa3_trt: refit engines are excluded from the shared deserialization cache (exclusive per-session ownership; dropping the engine with the session is the base-weight rollback); find_dit_engine gains want_refittable (LoRA sessions prefer refit engines, else fp16mixed over fp8, logged); SA3TRTDit exposes engine/engine_path/refittable - scripts/sa3/gen_sa3_refit_manifest (new): torch-FQN -> ONNX initializer mapping by shape + value fingerprint with transpose orientation; refuses partial manifests; --validate-engine runs the bit-identity gate (base-value refit must leave outputs bitwise unchanged) - sa3_trt_lora (new): SA3TRTRefitMirror — the eager SA3LoRAManager stays the source of truth; the mirror reads each parametrized module's merged weight (exact for every adapter variant by construction) and refits the exclusively-owned engine; dirty-set pushes base weights back after a disable - sa3_backend: mirror wiring (no eager swap in mirror mode; sync on enable/disable); knob strength changes under the mirror are stashed in rebuild_imminent and surfaced via has_pending_refit, applied post-pre-coverage in _prepare_tick — never a refit unannounced; eager sessions keep the inline buffer write - tests: mirror manifest/dirty-set/orientation vs a fake refitter with real vendored parametrizations; backend routing; upstream-loader parity for lora/dora-rows/dora-cols/bora and stack-order determinism GPU-gated (idle GPU required): the refit engine build, the manifest bit-identity validation run, and eager-vs-refit output parity. --- acestep/engine/sa3_context.py | 11 +- acestep/engine/sa3_trt.py | 98 ++++++++++- acestep/engine/sa3_trt_lora.py | 241 ++++++++++++++++++++++++++ acestep/engine/trt/lora_refit.py | 64 ++----- acestep/engine/trt/refit_core.py | 83 +++++++++ acestep/engine/trt/sa3_build.py | 72 ++++++++ acestep/streaming/sa3_backend.py | 170 ++++++++++++++++-- scripts/sa3/gen_sa3_refit_manifest.py | 231 ++++++++++++++++++++++++ tests/unit/test_lora_sa3_manager.py | 104 +++++++++++ tests/unit/test_sa3_backend.py | 81 +++++++++ tests/unit/test_sa3_trt_lora.py | 196 +++++++++++++++++++++ 11 files changed, 1281 insertions(+), 70 deletions(-) create mode 100644 acestep/engine/sa3_trt_lora.py create mode 100644 acestep/engine/trt/refit_core.py create mode 100644 scripts/sa3/gen_sa3_refit_manifest.py create mode 100644 tests/unit/test_sa3_trt_lora.py diff --git a/acestep/engine/sa3_context.py b/acestep/engine/sa3_context.py index 26fe8b40..681ec49e 100644 --- a/acestep/engine/sa3_context.py +++ b/acestep/engine/sa3_context.py @@ -85,6 +85,7 @@ def latent_rate_hz(self) -> float: def make_dit( self, *, latent_frames: int, seconds_total: float, backend: str = "eager", + prefer_refittable: bool = False, ): """The per-step velocity callable for one session: with ``backend="tensorrt"``, the built TRT engine when one covers @@ -98,7 +99,10 @@ def make_dit( ``backend`` is the resolved acceleration value the session creator threads through from the serving layer's accel param (compile is already normalized to eager there: SA3 has no - torch.compile path).""" + torch.compile path). ``prefer_refittable`` is the LoRA-session + preference (notes/SA3_LORA_PLAN.md D6b): pick a refit-built + engine when one covers the window, and avoid fp8 (whose refit + story is unproven) otherwise.""" from acestep.engine.sa3_trt import SA3TRTDit, find_dit_engine if backend != "tensorrt": @@ -108,7 +112,10 @@ def make_dit( ) return self.dit - engine_path = find_dit_engine(self.model_id, int(latent_frames)) + engine_path = find_dit_engine( + self.model_id, int(latent_frames), + want_refittable=bool(prefer_refittable), + ) if engine_path is None: logger.info( "sa3_dit_eager model_id={} latent_frames={} reason=no_trt_engine", diff --git a/acestep/engine/sa3_trt.py b/acestep/engine/sa3_trt.py index c71139f7..468f4803 100644 --- a/acestep/engine/sa3_trt.py +++ b/acestep/engine/sa3_trt.py @@ -85,11 +85,24 @@ _DIT_FP8_DIR_RE = re.compile( r"^(?P.+_dit)_fp8_l(?P\d+)_(?P\d+)_(?P\d+)$" ) +# Refittable fp16mixed DiT engines (BuilderFlag.REFIT; built by +# ``sa3_build --refit``): ``sa3_m_dit_refit_l{min}_{opt}_{max}``. The +# LoRA refit path (notes/SA3_LORA_PLAN.md D6b) mutates these in place, +# so they are EXCLUDED from the shared deserialization cache below — +# every consumer gets an exclusively-owned instance that dies with its +# session (which is also the base-weight rollback guarantee: a mutated +# engine can never be handed to a later session). +_DIT_REFIT_DIR_RE = re.compile( + r"^(?P.+_dit)_refit_l(?P\d+)_(?P\d+)_(?P\d+)$" +) _SAME_L_DIR_RE = re.compile(r"^same_l_decode_window_t(?P\d+)_(?P\d+)_(?P\d+)$") # Deserialized-engine process cache. Engines are immutable post-load and # support multiple execution contexts, so sharing one deserialization # across sessions is safe; the per-wrapper state is the context+buffers. +# The IMMUTABILITY ASSUMPTION is the cache's contract: refittable +# engines (see _DIT_REFIT_DIR_RE) must never enter it — they are +# deserialized per-session with exclusive ownership instead. _ENGINE_CACHE: dict = {} _ENGINE_CACHE_LOCK = threading.Lock() _SAME_PLUGIN_REGISTERED = False @@ -99,9 +112,31 @@ def trt_engines_dir() -> Path: return paths.models_dir() / "sa3" / "trt_engines" -def _deserialize_engine(path: Path): +def is_refittable_engine_path(path: Path) -> bool: + """Whether ``path`` names a refit-built DiT engine (by the naming + marker ``sa3_build --refit`` stamps).""" + return _DIT_REFIT_DIR_RE.match(Path(path).parent.name) is not None + + +def _deserialize_engine(path: Path, *, exclusive: bool = False): import tensorrt as trt + exclusive = exclusive or is_refittable_engine_path(path) + if exclusive: + # Refittable (or explicitly exclusive) engines bypass the cache: + # in-place refit would otherwise mutate every consumer and + # outlive the session that made it. The load cost (~seconds for + # the 2.8 GB DiT) is the accepted price of exclusive ownership. + logger.info( + "sa3_trt_engine_load path={} size_gb={:.1f} ownership=exclusive", + path, path.stat().st_size / 1e9, + ) + runtime = trt.Runtime(trt.Logger(trt.Logger.WARNING)) + engine = runtime.deserialize_cuda_engine(path.read_bytes()) + if engine is None: + raise RuntimeError(f"failed to deserialize TRT engine {path}") + return engine + with _ENGINE_CACHE_LOCK: engine = _ENGINE_CACHE.get(str(path)) if engine is None: @@ -144,19 +179,44 @@ def _register_same_plugin() -> None: # --------------------------------------------------------------------------- -def find_dit_engine(model_id: str, latent_frames: int) -> Optional[Path]: +def find_dit_engine( + model_id: str, latent_frames: int, *, want_refittable: bool = False, +) -> Optional[Path]: """Smallest-profile built DiT engine covering ``latent_frames`` for - ``model_id``'s weights, or None (caller falls back to eager).""" + ``model_id``'s weights, or None (caller falls back to eager). + + ``want_refittable`` is the LoRA-session preference (plan D6b): + + * A covering refit-built engine wins when one exists — LoRA strength + changes then run as in-place refits instead of the eager-DiT swap. + * Otherwise selection prefers **fp16mixed over fp8** (D6c: FP8 refit + on Stability's graph is its own unproven investigation; until then + a LoRA session must not land on an engine the refit path can't + serve at higher fidelity than the fp16mixed one), logged, and the + interim eager swap covers actual enables. + + Without it, selection is unchanged: fp8 preferred when covering, + else fp16mixed; refit-built engines are ignored (their refit + support costs a little optimization freedom, and non-LoRA sessions + shouldn't pay it). + """ prefix = DIT_ENGINE_PREFIX.get(model_id) base = trt_engines_dir() if prefix is None or not base.is_dir(): return None - best = None # smallest-covering fp16mixed engine - best_fp8 = None # smallest-covering fp8 engine (preferred when present) + best = None # smallest-covering fp16mixed engine + best_fp8 = None # smallest-covering fp8 engine + best_refit = None # smallest-covering refit-built engine for sub in base.iterdir(): f = sub / f"{sub.name}.trt" if not f.is_file(): continue + mr = _DIT_REFIT_DIR_RE.match(sub.name) + if mr and mr.group("prefix") == prefix: + lo, hi = int(mr.group("lo")), int(mr.group("hi")) + if lo <= latent_frames <= hi and (best_refit is None or hi < best_refit[0]): + best_refit = (hi, f) + continue mf = _DIT_FP8_DIR_RE.match(sub.name) if mf and mf.group("prefix") == prefix: lo, hi = int(mf.group("lo")), int(mf.group("hi")) @@ -168,6 +228,20 @@ def find_dit_engine(model_id: str, latent_frames: int) -> Optional[Path]: lo, hi = int(m.group("lo")), int(m.group("hi")) if lo <= latent_frames <= hi and (best is None or hi < best[0]): best = (hi, f) + if want_refittable: + if best_refit is not None: + logger.info( + "sa3_dit_refit_selected engine={} latent_frames={}", + best_refit[1].parent.name, latent_frames, + ) + return best_refit[1] + if best_fp8 is not None and best is not None: + logger.info( + "sa3_dit_fp8_skipped_for_lora engine={} using={} " + "reason=fp8_refit_unproven", + best_fp8[1].parent.name, best[1].parent.name, + ) + return best[1] if best else (best_fp8[1] if best_fp8 else None) # fp8 is ~1.8x faster; prefer it when one covers the window. if best_fp8 is not None: logger.info( @@ -192,7 +266,11 @@ def max_dit_engine_latents(model_id: str) -> Optional[int]: # install would otherwise report no cap and skip the clamp. his = [] for sub in base.iterdir(): - m = _DIT_DIR_RE.match(sub.name) or _DIT_FP8_DIR_RE.match(sub.name) + m = ( + _DIT_REFIT_DIR_RE.match(sub.name) + or _DIT_DIR_RE.match(sub.name) + or _DIT_FP8_DIR_RE.match(sub.name) + ) if ( m and m.group("prefix") == prefix @@ -259,6 +337,14 @@ class SA3TRTDit: def __init__(self, engine_path: Path, *, latent_frames: int, seconds_total: float): engine = _deserialize_engine(engine_path) + # Exposed for the LoRA refit mirror: on a refit-built engine + # (exclusively owned, never cached) the mirror attaches an + # IRefitter to this same instance. None of the wrapper's own + # state cares about refits — buffers and context stay valid + # across refit_cuda_engine(). + self.engine = engine + self.engine_path = Path(engine_path) + self.refittable = is_refittable_engine_path(engine_path) self._ctx = engine.create_execution_context() self._stream = torch.cuda.Stream() L = int(latent_frames) diff --git a/acestep/engine/sa3_trt_lora.py b/acestep/engine/sa3_trt_lora.py new file mode 100644 index 00000000..7eff23fa --- /dev/null +++ b/acestep/engine/sa3_trt_lora.py @@ -0,0 +1,241 @@ +"""SA3TRTRefitMirror: LoRA weights onto a refittable SA3 DiT engine. + +The Phase-2 endgame of notes/SA3_LORA_PLAN.md (D6b): instead of the +interim eager-DiT swap, a LoRA-enabled TRT session runs a refit-built +engine (``sa3_m_dit_refit_l*``, exclusively owned — see +``acestep.engine.sa3_trt._deserialize_engine``) and mirrors the eager +parametrization state onto it after every LoRA mutation. + +Composition is **merged-weight by construction**: the source of truth +stays the eager modules that :class:`~acestep.engine.sa3_lora.SA3LoRAManager` +parametrizes — reading ``module.weight`` evaluates the vendored +parametrization chain (alpha/rank scaling, DoRA/BoRA magnitude +renormalization, live strength buffers), so the mirrored value is exact +for every adapter variant with zero re-implemented math. The mirror +D2Hs each parametrized module's merged weight into a staging buffer, +pushes it under the ONNX initializer name the manifest maps it to +(transposing when the graph stores the MatMul ``[in, out]`` +orientation), and commits one ``refit_cuda_engine()``. + +The manifest (``gen_sa3_refit_manifest.py``) is generated offline by +shape + value fingerprinting over Stability's ONNX and validated by a +bit-identity refit (every mapped weight refit with its own base value +must leave the engine output bit-identical) — a wrong manifest cannot +ship silently. + +Sync scope: modules that currently carry parametrizations (their merged +value), plus previously-synced modules whose adapters were since +removed (their restored base goes back — the dirty set). Untouched +weights already hold their build-time values, and on session close the +engine is simply dropped (exclusive ownership — the rollback +guarantee), so no full base-restore pass ever runs. + +Sync stalls are full refits (hundreds of ms class). They are always +announced before they run: enables/disables ride ``has_pending_refit`` +via the session pending queues, and knob-driven strength changes are +detected in ``rebuild_imminent`` and stashed (never applied +mid-tick unannounced) — see ``SA3Backend`` (plan D6b.5). +""" + +from __future__ import annotations + +import json +import time +from pathlib import Path +from typing import Dict, Optional + +from loguru import logger +import torch + +from acestep.engine.trt.refit_core import ( + commit_refit, + np_view_for_push, + set_typed_weights, +) + +MANIFEST_VERSION = 1 +# Shared (per-ONNX) manifest filename; the per-engine sidecar +# ``.refit_manifest.json`` wins when present. +SHARED_MANIFEST_NAME = "sa3_m_dit_refit_manifest.json" + + +def find_refit_manifest(engine_path: Path) -> Optional[Path]: + """Manifest lookup: per-engine sidecar first, then the shared + per-ONNX manifest in the engines root.""" + engine_path = Path(engine_path) + sidecar = Path(str(engine_path) + ".refit_manifest.json") + if sidecar.is_file(): + return sidecar + shared = engine_path.parent.parent / SHARED_MANIFEST_NAME + if shared.is_file(): + return shared + return None + + +class SA3TRTRefitMirror: + """See module docstring. ``engine`` is the exclusively-owned + refittable engine; ``model_root`` the live DiT tree + (``sam.model.model``) the SA3LoRAManager parametrizes.""" + + def __init__( + self, + engine, + model_root: torch.nn.Module, + manifest_path: Path, + *, + _refitter=None, + _trt=None, + ) -> None: + # _refitter/_trt are test injection points: the dirty-set and + # staging logic is CPU-testable against a recording fake, while + # production always constructs the real IRefitter here. + if _trt is None: + import tensorrt as _trt # noqa: PLC0415 + + trt = _trt + self._trt = trt + self._model_root = model_root + self._refitter = _refitter if _refitter is not None else ( + trt.Refitter(engine, trt.Logger(trt.Logger.WARNING)) + ) + if not hasattr(self._refitter, "get_all_weights"): + raise RuntimeError("TRT engine refitting requires TensorRT 10.0+") + refittable_names = set(self._refitter.get_all_weights()) + if not refittable_names: + raise RuntimeError( + "engine has no refittable weights; was it built with " + "sa3_build --refit?" + ) + + manifest = json.loads(Path(manifest_path).read_text(encoding="utf-8")) + if manifest.get("version") != MANIFEST_VERSION: + raise RuntimeError( + f"refit manifest {manifest_path} has version " + f"{manifest.get('version')}, expected {MANIFEST_VERSION}" + ) + # fqn -> (initializer_name, transposed) + self._map: Dict[str, tuple] = {} + missing = 0 + for fqn, rec in manifest.get("weights", {}).items(): + init = rec["initializer"] + if init not in refittable_names: + missing += 1 + continue + self._map[fqn] = (init, bool(rec.get("transposed", False))) + if not self._map: + raise RuntimeError( + f"refit manifest {manifest_path} maps no weight that this " + f"engine exposes as refittable ({missing} mapped to " + f"non-refittable names)" + ) + if missing: + logger.warning( + "sa3_refit_manifest_partial mapped={} not_refittable={}", + len(self._map), missing, + ) + # Staging buffers, allocated lazily per synced fqn (pinned when + # possible) and reused across syncs; each must stay alive until + # commit (TRT dereferences the host pointers there). + self._staging: Dict[str, torch.Tensor] = {} + self._dtype_cache: Dict[str, tuple] = {} + # Fqns whose engine slot currently holds a LoRA-merged value. A + # later sync that finds such a module UNparametrized (adapter + # disabled) must push its base weight back — skipping it would + # leave the stale merged value in the engine. + self._dirty: set = set() + logger.info( + "sa3_trt_refit_mirror_ready mapped_weights={} manifest={}", + len(self._map), manifest_path, + ) + + def _resolve_module(self, fqn: str): + mod = self._model_root + for part in fqn.split("."): + if not hasattr(mod, part): + return None + mod = getattr(mod, part) + return mod + + def _prototype(self, init_name: str): + """(trt dtype, torch staging dtype) for one initializer.""" + cached = self._dtype_cache.get(init_name) + if cached is not None: + return cached + trt = self._trt + torch_dt = torch.float32 + trt_dt = trt.float32 + try: + proto = self._refitter.get_weights_prototype(init_name) + trt_dt = proto.dtype + torch_dt = { + trt.float32: torch.float32, + trt.float16: torch.float16, + }.get(proto.dtype, torch.float32) + if hasattr(trt, "bfloat16") and proto.dtype == trt.bfloat16: + torch_dt = torch.bfloat16 + except Exception: + pass + self._dtype_cache[init_name] = (trt_dt, torch_dt) + return trt_dt, torch_dt + + @torch.no_grad() + def sync(self, *, reason: str = "") -> int: + """Push the merged weight of every currently-parametrized mapped + module and commit one refit. Returns the number of weights + pushed (0 = nothing parametrized; no commit issued).""" + import torch.nn.utils.parametrize as parametrize + + t0 = time.perf_counter() + to_push: list = [] # (fqn, init_name) + now_clean: list = [] # dirty fqns whose base value goes back + with parametrize.cached(): + for fqn, (init_name, transposed) in self._map.items(): + mod = self._resolve_module(fqn) + if mod is None: + continue + parametrized = bool(getattr(mod, "parametrizations", None)) + if not parametrized and fqn not in self._dirty: + continue # engine already holds this base value + # Parametrized: reading .weight evaluates the vendored + # chain (the merged value). Unparametrized-but-dirty: + # .weight IS the restored base — push it back. + merged = mod.weight.detach() + if transposed and merged.dim() == 2: + merged = merged.transpose(0, 1) + trt_dt, torch_dt = self._prototype(init_name) + buf = self._staging.get(fqn) + if buf is None or buf.shape != merged.shape or buf.dtype != torch_dt: + buf = torch.empty( + merged.shape, dtype=torch_dt, device="cpu", + ) + try: + buf = buf.pin_memory() + except RuntimeError: + pass + self._staging[fqn] = buf + buf.copy_(merged.to(torch_dt), non_blocking=True) + to_push.append((fqn, init_name)) + if parametrized: + self._dirty.add(fqn) + else: + now_clean.append(fqn) + if not to_push: + return 0 + # All D2H copies must land before TRT reads the host pointers. + if torch.cuda.is_available(): + torch.cuda.synchronize() + for fqn, init_name in to_push: + trt_dt, _torch_dt = self._prototype(init_name) + set_typed_weights( + self._refitter, self._trt, init_name, + np_view_for_push(self._staging[fqn]), trt_dt, + context=f"sa3 refit mirror fqn={fqn}", + ) + commit_refit(self._refitter) + self._dirty.difference_update(now_clean) + logger.info( + "sa3_trt_refit_sync pushed={} restored={} reason={} refit_ms={:.1f}", + len(to_push), len(now_clean), reason or "-", + (time.perf_counter() - t0) * 1000, + ) + return len(to_push) diff --git a/acestep/engine/trt/lora_refit.py b/acestep/engine/trt/lora_refit.py index 1962e1bc..3890992b 100644 --- a/acestep/engine/trt/lora_refit.py +++ b/acestep/engine/trt/lora_refit.py @@ -22,6 +22,11 @@ LoRADescriptor, # re-exported for back-compat _LoRAEntry, # re-exported for tests ) +from acestep.engine.trt.refit_core import ( + commit_refit, + np_view_for_push, + set_typed_weights, +) # numpy dtype -> torch dtype _NP_TO_TORCH = { @@ -629,62 +634,31 @@ def _apply_to_engine( if trt_name is None: continue buf = self._refit_bufs[param_name] - if buf.dtype == torch.bfloat16: - arr = buf.view(torch.uint16).numpy() - else: - arr = buf.numpy() - weights = self._trt.Weights( + arr = np_view_for_push(buf) + set_typed_weights( + refitter, self._trt, trt_name, arr, self._trt_dtype[param_name], - int(arr.ctypes.data), - int(arr.size), - ) - ok = refitter.set_named_weights(trt_name, weights) - if not ok: - proto_desc = "unknown" - if hasattr(refitter, "get_weights_prototype"): - try: - proto = refitter.get_weights_prototype(trt_name) - proto_desc = ( - f"dtype={proto.dtype}, size={proto.size}" - ) - except Exception: - pass - raise RuntimeError( - f"TRT rejected refit weights for {trt_name}: " - f"buf dtype={buf.dtype}, arr dtype={arr.dtype} " - f"size={arr.size}; engine prototype {proto_desc}, " + context=( + f"buf dtype={buf.dtype}, " f"fp8={self._is_fp8.get(param_name, False)}" - ) + ), + ) for name, arr, dtype in getattr(self, "_co_refit_static", ()): - weights = self._trt.Weights( - dtype, int(arr.ctypes.data), int(arr.size), + set_typed_weights( + refitter, self._trt, name, arr, dtype, + context="co-refit activation scale", ) - if not refitter.set_named_weights(name, weights): - raise RuntimeError( - f"TRT rejected co-refit weight {name}: " - f"dtype={dtype}, size={arr.size}" - ) for param_key, (name, _scale_buf, arr, dtype) in getattr( self, "_weight_scale_co_refit", {} ).items(): - weights = self._trt.Weights( - dtype, int(arr.ctypes.data), int(arr.size), + set_typed_weights( + refitter, self._trt, name, arr, dtype, + context=f"co-refit weight scale (param={param_key})", ) - if not refitter.set_named_weights(name, weights): - raise RuntimeError( - f"TRT rejected co-refit weight scale {name} " - f"(param={param_key}): dtype={dtype}, size={arr.size}" - ) if count > 0: - ok = refitter.refit_cuda_engine() - if not ok: - missing = refitter.get_missing_weights() - if _refit_ok_required: - raise RuntimeError( - f"TRT refit failed. Missing weights: {missing}" - ) + commit_refit(refitter, required=_refit_ok_required) # Re-export the lifecycle public symbols so existing imports diff --git a/acestep/engine/trt/refit_core.py b/acestep/engine/trt/refit_core.py new file mode 100644 index 00000000..6309c319 --- /dev/null +++ b/acestep/engine/trt/refit_core.py @@ -0,0 +1,83 @@ +"""Shared TRT refit primitives (notes/SA3_LORA_PLAN.md D6b). + +The pieces of in-place ``IRefitter`` writeback that are family-agnostic: +typed ``trt.Weights`` construction for dtypes numpy cannot represent +(BF16 / FP8 storage travels as uint16/uint8 views, and the ndarray +overload of ``set_named_weights`` would mistype them — TRT then rejects +the refit with a dtype-mismatch error), the loud set/commit error +shapes, and the missing-weights diagnostics. + +Consumers: + +* :class:`acestep.engine.trt.lora_refit.TRTLoRAManager` (ACE) — its + delta-composition and FP8 co-refit machinery stay family-private; the + typed push + commit go through here (byte-identical behavior, guarded + by its existing unit tests). +* :class:`acestep.engine.sa3_trt_lora.SA3TRTRefitMirror` — composes + MERGED weights from the parametrized torch modules and pushes them + through the same primitives. +""" + +from __future__ import annotations + +from typing import Optional + +import torch + + +def np_view_for_push(buf: torch.Tensor): + """A CPU numpy view of ``buf`` suitable for a typed weights push. + + BF16 has no numpy dtype; its storage travels as a uint16 view (the + ``trt.Weights(dtype, ptr, count)`` wrapper re-types it). FP8 + storage is expected pre-viewed as uint8 by the caller. The returned + ndarray aliases ``buf``'s memory — the caller must keep ``buf`` + alive until the refit commits. + """ + if buf.dtype == torch.bfloat16: + return buf.view(torch.uint16).numpy() + return buf.numpy() + + +def set_typed_weights( + refitter, + trt_mod, + name: str, + arr, + trt_dtype, + *, + context: str = "", +) -> None: + """``set_named_weights`` with an explicit TRT dtype, raising with + full diagnostics on rejection. ``arr`` must stay alive until + :func:`commit_refit` returns (TRT dereferences the pointer then).""" + weights = trt_mod.Weights(trt_dtype, int(arr.ctypes.data), int(arr.size)) + if refitter.set_named_weights(name, weights): + return + proto_desc = "unknown" + if hasattr(refitter, "get_weights_prototype"): + try: + proto = refitter.get_weights_prototype(name) + proto_desc = f"dtype={proto.dtype}, size={proto.size}" + except Exception: + pass + raise RuntimeError( + f"TRT rejected refit weights for {name}: " + f"arr dtype={getattr(arr, 'dtype', '?')} size={arr.size}; " + f"engine prototype {proto_desc}" + + (f"; {context}" if context else "") + ) + + +def commit_refit(refitter, *, required: bool = True) -> bool: + """``refit_cuda_engine()`` with the missing-weights diagnostic. + Returns the TRT result; raises when ``required`` and it failed.""" + ok = refitter.refit_cuda_engine() + if not ok and required: + missing: Optional[list] = None + try: + missing = refitter.get_missing_weights() + except Exception: + pass + raise RuntimeError(f"TRT refit failed. Missing weights: {missing}") + return bool(ok) diff --git a/acestep/engine/trt/sa3_build.py b/acestep/engine/trt/sa3_build.py index be59e456..b10c5b05 100644 --- a/acestep/engine/trt/sa3_build.py +++ b/acestep/engine/trt/sa3_build.py @@ -168,6 +168,33 @@ def engine_name(self) -> str: ) +@dataclass +class SA3DiTRefitBuildConfig: + """Build parameters for one REFITTABLE sa3-m DiT engine + (notes/SA3_LORA_PLAN.md D6b: the LoRA refit endgame). + + A separate dataclass (same rationale as the fp8 one: the metadata + skip gate hashes the whole config, so folding a ``refit`` field into + :class:`SA3DiTBuildConfig` would change the existing engines' + identity and force a needless rebuild). Same fp16mixed ONNX, the + ``BuilderFlag.REFIT`` builder flag, and the ``_refit`` name marker + that :func:`acestep.engine.sa3_trt.is_refittable_engine_path` (and + the exclusive-ownership deserialization policy) keys on. + """ + + min_latents: int + opt_latents: int + max_latents: int + workspace_gb: float = 16.0 + onnx_files: list[str] = field(default_factory=lambda: list(DIT_ONNX_FILES)) + + def engine_name(self) -> str: + return ( + f"sa3_m_dit_refit_l{self.min_latents}" + f"_{self.opt_latents}_{self.max_latents}" + ) + + @dataclass class SameLWindowBuildConfig: """Build parameters for the SAME-L window decoder engine.""" @@ -208,6 +235,7 @@ def _build_strongly_typed_engine( engine_path: str, workspace_gb: float, profile_shapes: dict[str, tuple[tuple, tuple, tuple]], + refit: bool = False, ) -> None: """Parse + build one STRONGLY_TYPED engine and serialize it to disk. @@ -233,6 +261,8 @@ def _build_strongly_typed_engine( config.set_memory_pool_limit( trt.MemoryPoolType.WORKSPACE, int(workspace_gb * (1 << 30)), ) + if refit: + config.set_flag(trt.BuilderFlag.REFIT) profile = builder.create_optimization_profile() for input_name, (lo, opt, hi) in profile_shapes.items(): profile.set_shape(input_name, lo, opt, hi) @@ -257,6 +287,7 @@ def _build_dit_engine( component: str = "sa3_m_dit", precision_label: str = "fp16mixed", local_onnx: str | None = None, + refit: bool = False, ) -> tuple[str, str, float, str]: """Build one sa3-m DiT engine. Returns (label, path, elapsed, status). @@ -323,6 +354,7 @@ def _build_dit_engine( "seconds_total": ((1,), (1,), (1,)), "local_add_cond": ((1, 257, lo), (1, 257, opt), (1, 257, hi)), }, + refit=refit, ) _write_metadata(engine_path=engine_path, expected=expected, env=env) elapsed = time.time() - t0 @@ -426,6 +458,14 @@ def _matrix_jobs(args) -> list[tuple[str, str]]: f" (~{hi * SAMPLES_PER_LATENT / SA3_SAMPLE_RATE:.0f}s window)", cfg.engine_name(), )) + if args.refit: + for lo, opt, hi in dit_profiles: + cfg = SA3DiTRefitBuildConfig(lo, opt, hi) + jobs.append(( + f"SA3-M DiT refit l{lo}_{opt}_{hi}" + f" (~{hi * SAMPLES_PER_LATENT / SA3_SAMPLE_RATE:.0f}s window)", + cfg.engine_name(), + )) if not args.dit_only: lo, opt, hi = CANONICAL_SAME_L_WINDOW cfg = SameLWindowBuildConfig(lo, opt, hi) @@ -534,6 +574,12 @@ def main() -> int: "(~1.8x/step; preferred at runtime when present, " "fp16mixed fallback). Needs the published " "dit_fp8.onnx or --fp8-onnx.") + single.add_argument("--refit", action="store_true", + help="Also build the REFITTABLE fp16mixed DiT " + "variant(s) (BuilderFlag.REFIT; the LoRA " + "in-place-refit engines, preferred by " + "LoRA-enabled sessions and exclusively owned " + "at runtime — never process-cached).") single.add_argument("--fp8-onnx", default=None, help="Path to a producer-built dit_fp8.onnx (with its " ".onnx.data sidecar alongside) to compile instead " @@ -585,6 +631,18 @@ def main() -> int: precision_label="fp8", local_onnx=args.fp8_onnx, )) + if args.refit: + for lo, opt, hi in dit_profiles: + results.append(_build_dit_engine( + output_dir=args.output_dir, + config=SA3DiTRefitBuildConfig( + lo, opt, hi, workspace_gb=args.workspace_gb), + env=env, + force_rebuild=args.force_rebuild, + component="sa3_m_dit_refit", + precision_label="fp16mixed refit", + refit=True, + )) if not args.dit_only: lo, opt, hi = CANONICAL_SAME_L_WINDOW results.append(_build_same_l_window_engine( @@ -631,6 +689,20 @@ def main() -> int: component="sa3_m_dit_fp8", precision_label="fp8", local_onnx=args.fp8_onnx, )) + if args.refit: + refit_cfg = SA3DiTRefitBuildConfig( + min_latents=config.min_latents, + opt_latents=config.opt_latents, + max_latents=config.max_latents, + workspace_gb=args.workspace_gb, + ) + built.append(_build_dit_engine( + output_dir=args.output_dir, config=refit_cfg, env=env, + force_rebuild=args.force_rebuild, + component="sa3_m_dit_refit", + precision_label="fp16mixed refit", + refit=True, + )) if args.same_l_window: d_lo, d_opt, d_hi = CANONICAL_SAME_L_WINDOW config = SameLWindowBuildConfig( diff --git a/acestep/streaming/sa3_backend.py b/acestep/streaming/sa3_backend.py index 84e59737..4c6e0a09 100644 --- a/acestep/streaming/sa3_backend.py +++ b/acestep/streaming/sa3_backend.py @@ -231,6 +231,12 @@ def __init__( # Runtime checkpoint id ("medium"/"small-music"/…) for the # lineage axis of lora_compatible. model_id: str = "", + # Phase-2 TRT endgame: an SA3TRTRefitMirror over a refit-built, + # exclusively-owned engine. When present, LoRA mutations refit + # the engine in place instead of swapping to the eager DiT, and + # knob-driven strength changes route through the pending stash + # (D6b.5) so the refit stall is always announced. + refit_mirror=None, ): super().__init__(adapter=adapter, codec=codec) self._cond = cond @@ -315,6 +321,14 @@ def __init__( # flips adapter.dit between them by LoRA activity. self._dit_accel = adapter.dit if adapter is not None else None self._dit_eager = eager_dit if eager_dit is not None else self._dit_accel + self._refit_mirror = refit_mirror + # Knob-driven strength changes stashed by rebuild_imminent under + # the refit mirror (a strength change is a full engine refit + # there, and must be announced via rebuild_imminent / + # has_pending_refit before it runs — never applied mid-tick + # unannounced). Empty and unused in eager mode, where a + # strength change is a ~5 ms buffer write applied inline. + self._pending_lora_strengths: dict = {} # Rendered-audio cache: one full decode+resample per fresh # latent (SAME-S decodes the whole window in ~11 ms); window @@ -396,17 +410,52 @@ def from_context( else context.encode_source(source_audio, cond.audio_sample_size) if source_audio is not None else None ) + # LoRA sessions prefer a refit-built engine (and avoid fp8, + # whose refit story is unproven) — see find_dit_engine. + want_lora = bool(kwargs.get("use_lora")) and ( + kwargs.get("lora_manager") is not None + ) adapter = SA3Adapter( context.make_dit( latent_frames=cond.latent_frames, seconds_total=duration_s, backend=dit_backend, + prefer_refittable=want_lora, ), schedule_builder=context.make_schedule_builder(cond, steps), device=context.device, dtype=context.dtype, ) + # Phase-2 refit mirror: engaged only when the selected DiT is a + # refit-built engine AND its validated manifest exists; + # otherwise the interim eager swap (D6a) covers LoRA. + refit_mirror = None + if want_lora and getattr(adapter.dit, "refittable", False): + from acestep.engine.sa3_trt_lora import ( + SA3TRTRefitMirror, + find_refit_manifest, + ) + + manifest = find_refit_manifest(adapter.dit.engine_path) + if manifest is None: + logger.warning( + "sa3_refit_manifest_missing engine={} — LoRA will use " + "the eager-DiT swap; generate the manifest with " + "scripts/sa3/gen_sa3_refit_manifest.py", + adapter.dit.engine_path.parent.name, + ) + else: + try: + refit_mirror = SA3TRTRefitMirror( + adapter.dit.engine, context.sam.model.model, manifest, + ) + except Exception as exc: + logger.warning( + "sa3_refit_mirror_init_failed error={} — LoRA will " + "use the eager-DiT swap", exc, + ) + def _prompt_rebuilder(tags: str, steps_now: int): # Per-prompt re-conditioning (handle_set_prompt): same fixed # duration, fresh T5Gemma capture + a schedule-builder @@ -447,6 +496,7 @@ def _source_encoder(waveform, sample_rate, sample_size): # LoRA fallback swaps to it. eager_dit=context.dit, model_id=str(context.model_id), + refit_mirror=refit_mirror, **kwargs, ) @@ -508,7 +558,13 @@ def has_pending_refit(self) -> bool: the same pending-queue read as ACE's implementation, so the runner pre-covers the enable stall (which can synchronously materialize + register 200+ parametrizations).""" - if not self._use_lora or self.state is None: + if not self._use_lora: + return False + # Stashed TRT strength refits count too (D6b.5): they apply at + # the next prepare and stall exactly like a queued enable. + if self._pending_lora_strengths: + return True + if self.state is None: return False try: with self.state._lock: @@ -552,6 +608,8 @@ def enable_lora(self, lora_id: str, strength=None) -> None: with self._conditioner_lock: mgr.enable_lora(lora_id, strength=strength) self._sync_dit_for_lora() + if self._refit_mirror is not None: + self._refit_mirror.sync(reason="enable_lora") if mgr.touches_conditioner(lora_id): self._rebuild_conditioning_after_lora("enable_lora", lora_id) @@ -563,6 +621,11 @@ def disable_lora(self, lora_id: str) -> None: with self._conditioner_lock: mgr.disable_lora(lora_id) self._sync_dit_for_lora() + if self._refit_mirror is not None: + # The mirror's dirty-set pushes this adapter's base weights + # back — the engine must not keep the merged values. + self._refit_mirror.sync(reason="disable_lora") + self._pending_lora_strengths.pop(lora_id, None) if touched: self._rebuild_conditioning_after_lora("disable_lora", lora_id) @@ -570,6 +633,11 @@ def set_lora_strength(self, lora_id: str, strength: float) -> None: mgr = self._require_lora_manager() with self._conditioner_lock: mgr.set_lora_strength(lora_id, strength) + if self._refit_mirror is not None: + # Direct facade calls sync immediately; the knob-driven path + # batches instead (_apply_pending_lora_strengths calls the + # manager directly and syncs once for the whole stash). + self._refit_mirror.sync(reason="set_lora_strength") # A strength change on a conditioner-targeting LoRA changes the # conditioner's output too — the cached cond bundle must follow. # (Per-tick slider sweeps on such files pay a T5Gemma capture @@ -578,11 +646,41 @@ def set_lora_strength(self, lora_id: str, strength: float) -> None: if mgr.touches_conditioner(lora_id): self._rebuild_conditioning_after_lora("set_lora_strength", lora_id) + def _apply_pending_lora_strengths(self) -> None: + """Drain the D6b.5 stash: apply each strength through the facade + (manager buffer write + conditioner rebuild where the file + targets it), then ONE mirror sync covers the whole batch.""" + if not self._pending_lora_strengths: + return + pending, self._pending_lora_strengths = ( + self._pending_lora_strengths, {}, + ) + mgr = self._lora_mgr + for lora_id, strength in pending.items(): + try: + with self._conditioner_lock: + mgr.set_lora_strength(lora_id, strength) + if mgr.touches_conditioner(lora_id): + self._rebuild_conditioning_after_lora( + "set_lora_strength", lora_id, + ) + except Exception as exc: + logger.exception( + "sa3_lora_pending_strength_failed id={} error={}", + lora_id, exc, + ) + if self._refit_mirror is not None: + self._refit_mirror.sync(reason="strength") + def _sync_dit_for_lora(self) -> None: """D6a interim: LoRA weights live on the eager torch modules, so a TRT-DiT session swaps to the eager DiT while any LoRA is active and back when the last disables. Loud log both ways; - no-op for sessions already on the eager DiT.""" + no-op for sessions already on the eager DiT — and for refit- + mirror sessions, where the engine itself carries the LoRA + weights (the Phase-2 endgame) and must NOT be swapped away.""" + if self._refit_mirror is not None: + return if self.adapter is None or self._dit_eager is None: return if self._dit_accel is self._dit_eager: @@ -642,7 +740,35 @@ def read_knobs(self) -> dict: return self.knob_state.get_all_values() def rebuild_imminent(self, knobs: dict) -> bool: - return int(knobs.get("steps_override", self._steps)) != self._steps + steps_changed = ( + int(knobs.get("steps_override", self._steps)) != self._steps + ) + return steps_changed or self._detect_pending_lora_strengths(knobs) + + def _detect_pending_lora_strengths(self, knobs: dict) -> bool: + """D6b.5: under the refit mirror a strength change is a full + engine refit, and ``rebuild_imminent`` is the sanctioned + announcement point for knob-driven stalls (called once per tick, + after the knob read, before produce). Deltas are stashed here + and applied by the SAME tick's ``_prepare_tick`` — after the + runner has pre-covered. Eager mode returns False: its strength + writes are ~5 ms buffer updates applied inline.""" + if ( + self._refit_mirror is None + or not self._use_lora + or self._lora_mgr is None + ): + return False + for desc in self._lora_mgr.list_loras(): + if desc.state != "enabled": + continue + try: + v = float(knobs.get(f"lora_str_{desc.id}", desc.strength)) + except (TypeError, ValueError): + continue + if abs(v - desc.strength) > 0.02: + self._pending_lora_strengths[desc.id] = v + return bool(self._pending_lora_strengths) # ---- control (universal): per-prompt re-conditioning ------------------------ @@ -831,21 +957,25 @@ def _prepare_tick(self, knobs: dict, ctx: TickContext) -> dict: # Per-LoRA live strength (the ACE per-tick convention): iterate # the catalog so the active set can change at runtime; strength # only flows for ENABLED entries, gated by the shared 0.02 - # slider-delta threshold. Under the parametrization engine this - # is a buffer write (no recompute); the gate also pre-guards the - # future TRT refit cost. + # slider-delta threshold. Eager mode applies inline (a buffer + # write, no recompute). Refit-mirror mode instead applies the + # deltas rebuild_imminent stashed THIS tick, after the runner + # pre-covered — the D6b.5 route: never a refit unannounced. if self._use_lora and self._lora_mgr is not None: - for desc in self._lora_mgr.list_loras(): - if desc.state != "enabled": - continue - try: - lora_str = float( - knobs.get(f"lora_str_{desc.id}", desc.strength), - ) - except (TypeError, ValueError): - continue - if abs(lora_str - desc.strength) > 0.02: - self.set_lora_strength(desc.id, lora_str) + if self._refit_mirror is not None: + self._apply_pending_lora_strengths() + else: + for desc in self._lora_mgr.list_loras(): + if desc.state != "enabled": + continue + try: + lora_str = float( + knobs.get(f"lora_str_{desc.id}", desc.strength), + ) + except (TypeError, ValueError): + continue + if abs(lora_str - desc.strength) > 0.02: + self.set_lora_strength(desc.id, lora_str) # Source-lock strength rides the shared override so a strength # bump engages the blend on in-flight slots submitted while it @@ -1113,6 +1243,12 @@ def close(self) -> None: self._lora_mgr.close() except Exception as exc: logger.warning("sa3_lora_teardown_raised error={}", exc) + # The refit mirror holds an IRefitter over the exclusively-owned + # engine; dropping both with the session IS the base-weight + # rollback (the mutated engine can never reach another session + # because refittable engines are never process-cached). + self._refit_mirror = None + self._pending_lora_strengths.clear() # ---- bookkeeping ------------------------------------------------------------- diff --git a/scripts/sa3/gen_sa3_refit_manifest.py b/scripts/sa3/gen_sa3_refit_manifest.py new file mode 100644 index 00000000..edd37543 --- /dev/null +++ b/scripts/sa3/gen_sa3_refit_manifest.py @@ -0,0 +1,231 @@ +"""Generate (and validate) the SA3 DiT refit manifest. + +Maps every LoRA-targetable torch weight of the loaded SA3 DiT +(``sam.model.model`` Linear/Conv1d weights — the modules +``SA3LoRAManager`` can parametrize) to its ONNX initializer name in +Stability's ``dit_fp16mixed.onnx`` graph, by shape + value fingerprint, +recording the transpose orientation. The manifest is what +:class:`acestep.engine.sa3_trt_lora.SA3TRTRefitMirror` consumes. + +Matching: for each torch weight, candidate initializers are those whose +shape equals the weight's 2D view directly or transposed; ambiguity is +resolved by full value comparison (both sides cast to fp32; the ONNX +trunk stores fp16, the torch model loads fp16, so an exact-bits check +decides). A weight matching in BOTH orientations (only possible for +symmetric-shape symmetric-value degenerates) or in neither is reported, +never silently guessed. + +Validation (``--validate-engine``): deserialize the refit-built engine +EXCLUSIVELY, refit every mapped weight with its own base value, and +compare a fixed forward against the unrefit output — bit-identical or +the manifest (or orientation) is wrong. This is the gate that makes a +wrong manifest impossible to ship silently (plan D6b.3). + +Run: + .venv/Scripts/python.exe scripts/sa3/gen_sa3_refit_manifest.py \ + --onnx /dit_fp16mixed.onnx + .venv/Scripts/python.exe scripts/sa3/gen_sa3_refit_manifest.py \ + --onnx <...> --validate-engine <...>/sa3_m_dit_refit_l1_646_646.trt + +The ONNX ships on HF (stabilityai/stable-audio-3-optimized, +onnx/sa3-m/dit_fp16mixed.onnx + .data). Pass --fetch to download it via +huggingface_hub (cached; ~3 GB on first use), or point --onnx at an +existing copy (the sa3_build cache has one after any DiT engine build). +""" + +from __future__ import annotations + +import argparse +import json +import sys +import time +from pathlib import Path + +_HERE = Path(__file__).resolve().parent +_REPO_ROOT = next(p for p in (_HERE, *_HERE.parents) if (p / "pyproject.toml").exists()) +for _p in (str(_REPO_ROOT),): + while _p in sys.path: + sys.path.remove(_p) + sys.path.insert(0, _p) + +import torch # noqa: E402 + + +def _iter_target_weights(model_root): + import torch.nn as nn + + for name, mod in model_root.named_modules(): + if isinstance(mod, (nn.Linear, nn.Conv1d)): + yield name, mod.weight.detach() + + +def build_manifest(model_root, onnx_path: Path) -> dict: + import onnx + from onnx import numpy_helper + + print(f"loading ONNX graph {onnx_path} (external data alongside)...") + model = onnx.load(str(onnx_path), load_external_data=True) + inits = list(model.graph.initializer) + by_shape: dict[tuple, list] = {} + for init in inits: + by_shape.setdefault(tuple(init.dims), []).append(init) + print(f"{len(inits)} initializers indexed") + + weights: dict[str, dict] = {} + unmatched: list[str] = [] + ambiguous: list[str] = [] + used: set[str] = set() + for fqn, w in _iter_target_weights(model_root): + w2 = w.view(w.shape[0], -1).float().cpu() + direct_shape = tuple(w.shape) + w2_shape = tuple(w2.shape) + t_shape = (w2_shape[1], w2_shape[0]) + candidates = [] + for shape, transposed in ( + (direct_shape, False), + (w2_shape, False), + (t_shape, True), + ): + for init in by_shape.get(shape, []): + if init.name in used: + continue + arr = torch.from_numpy( + numpy_helper.to_array(init).copy() + ).float() + arr2 = arr.view(arr.shape[0], -1) if not transposed else ( + arr.view(arr.shape[0], -1).transpose(0, 1) + ) + ref = w2 + if arr2.shape != ref.shape: + continue + # fp16-trunk graphs store the same fp16 the torch model + # loads; exact equality decides. A slack allclose backs + # up fp32-island initializers that round differently. + if torch.equal(arr2, ref) or torch.allclose( + arr2, ref, rtol=1e-3, atol=1e-4, + ): + candidates.append((init.name, transposed)) + names = {c[0] for c in candidates} + if not names: + unmatched.append(fqn) + continue + if len(names) > 1: + ambiguous.append(fqn) + continue + init_name, transposed = candidates[0] + used.add(init_name) + weights[fqn] = {"initializer": init_name, "transposed": transposed} + + print( + f"mapped {len(weights)} weights; unmatched={len(unmatched)} " + f"ambiguous={len(ambiguous)}" + ) + for fqn in unmatched[:10]: + print(f" UNMATCHED: {fqn}") + for fqn in ambiguous[:10]: + print(f" AMBIGUOUS: {fqn}") + return { + "version": 1, + "onnx": onnx_path.name, + "weights": weights, + "unmatched": unmatched, + "ambiguous": ambiguous, + } + + +def validate_engine(manifest_path: Path, engine_path: Path, ctx) -> bool: + """Refit every mapped weight with its own base value; the engine + output must stay bit-identical.""" + from acestep.engine.sa3_trt import SA3TRTDit + from acestep.engine.sa3_trt_lora import SA3TRTRefitMirror + from acestep.engine.sa3_stream_helpers import stack_sa3_cond_bundles + + cond = ctx.prepare_cond( + prompt="warm analog house groove, 124 bpm", duration=20.0, steps=8, + ) + dit = SA3TRTDit( + engine_path, latent_frames=cond.latent_frames, seconds_total=20.0, + ) + if not dit.refittable: + print(f"FAIL: {engine_path} is not a refit-built engine") + return False + stacked = stack_sa3_cond_bundles([cond.cond_bundle]) + x = torch.randn( + 1, ctx.latent_channels, cond.latent_frames, device="cuda", + generator=torch.Generator(device="cuda").manual_seed(7), + ) + ref = dit.step_bundle(x, 0.5, stacked).clone() + + mirror = SA3TRTRefitMirror(dit.engine, ctx.sam.model.model, manifest_path) + # Force-push every mapped weight at its BASE value: mark all dirty + # so sync pushes them even though nothing is parametrized. + mirror._dirty = set(mirror._map.keys()) + t0 = time.perf_counter() + pushed = mirror.sync(reason="validate") + print(f"validation refit: {pushed} weights in " + f"{(time.perf_counter() - t0) * 1000:.0f}ms") + + out = dit.step_bundle(x, 0.5, stacked).clone() + if torch.equal(out, ref): + print("PASS: base-value refit is bit-identical") + return True + err = (out - ref).abs().max().item() + print(f"FAIL: output changed after base-value refit (max_err={err:.3e})") + return False + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--model", default="medium") + ap.add_argument("--onnx", default=None, + help="path to dit_fp16mixed.onnx (with its .data sidecar)") + ap.add_argument("--fetch", action="store_true", + help="fetch the ONNX from HF (cached; ~3 GB first use)") + ap.add_argument("--out", default=None, + help="manifest output path (default: " + "/sa3_m_dit_refit_manifest.json)") + ap.add_argument("--validate-engine", default=None, + help="refit-built engine to bit-identity-validate against") + args = ap.parse_args() + + from acestep.engine.sa3_context import SA3Context + from acestep.engine.sa3_trt import trt_engines_dir + from acestep.engine.sa3_trt_lora import SHARED_MANIFEST_NAME + + if args.onnx: + onnx_path = Path(args.onnx) + elif args.fetch: + from acestep.engine.trt.sa3_build import DIT_ONNX_FILES, _fetch_onnx + + onnx_path = Path(_fetch_onnx(DIT_ONNX_FILES)) + else: + ap.error("pass --onnx or --fetch") + if not onnx_path.is_file(): + ap.error(f"ONNX not found: {onnx_path}") + + out = Path(args.out) if args.out else ( + trt_engines_dir() / SHARED_MANIFEST_NAME + ) + + print(f"loading SA3Context({args.model!r})...") + ctx = SA3Context(model_id=args.model) + + manifest = build_manifest(ctx.sam.model.model, onnx_path) + if manifest["unmatched"] or manifest["ambiguous"]: + print("REFUSING to write a partial manifest " + "(unmatched/ambiguous weights above)") + return 1 + out.parent.mkdir(parents=True, exist_ok=True) + out.write_text(json.dumps(manifest, indent=2), encoding="utf-8") + print(f"manifest written: {out} ({len(manifest['weights'])} weights)") + + if args.validate_engine: + ok = validate_engine(out, Path(args.validate_engine), ctx) + return 0 if ok else 1 + print("NOTE: run again with --validate-engine to " + "complete the bit-identity gate before first production use.") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/unit/test_lora_sa3_manager.py b/tests/unit/test_lora_sa3_manager.py index ca864e5d..117d943a 100644 --- a/tests/unit/test_lora_sa3_manager.py +++ b/tests/unit/test_lora_sa3_manager.py @@ -456,3 +456,107 @@ def test_close_returns_model_to_pristine(tmp_path): assert not getattr(mod, "parametrizations", None) assert mgr.list_loras() == [] mgr.close() # idempotent + + +# --------------------------------------------------------------------------- +# Phase 3: upstream-reference parity per adapter variant +# --------------------------------------------------------------------------- +# +# The strongest correctness property the manager can have: applying a +# file through SA3LoRAManager must yield exactly the weights upstream's +# own load_and_apply_loras produces from the same file — for every +# adapter variant, at every strength, in every enable order. + + +def _write_variant_lora(tmp_path: Path, stem: str, adapter_type: str, + fan=(8, 16), rank: int = 4, seed: int = 0) -> Path: + _, lora_utils = _vendored_or_skip() + g = torch.Generator().manual_seed(seed) + fan_out, fan_in = fan + prefix = "blk.lin.parametrizations.weight.0" + sd = { + f"{prefix}.lora_A": torch.randn(rank, fan_in, generator=g) * 0.5, + f"{prefix}.lora_B": torch.randn(fan_out, rank, generator=g) * 0.5, + } + if adapter_type in ("dora-rows",): + sd[f"{prefix}.magnitude"] = torch.rand(fan_out, generator=g) + 0.5 + elif adapter_type in ("dora-cols",): + sd[f"{prefix}.magnitude"] = torch.rand(fan_in, generator=g) + 0.5 + elif adapter_type == "bora": + sd[f"{prefix}.magnitude_r"] = torch.rand(fan_out, generator=g) + 0.5 + sd[f"{prefix}.magnitude_c"] = torch.rand(fan_in, generator=g) + 0.5 + p = tmp_path / f"{stem}.safetensors" + lora_utils.save_lora_safetensors( + sd, {"rank": rank, "alpha": 6.0, "adapter_type": adapter_type}, p, + ) + return p + + +def _upstream_model_with(paths): + """A twin tiny model with the same seed, loaded through the vendored + load_and_apply_loras (the trainer/inference reference path).""" + from stable_audio_3.models.lora.loader import load_and_apply_loras + + torch.manual_seed(7) + model = nn.Module() + blk = nn.Module() + blk.lin = nn.Linear(16, 8, bias=False) + model.blk = blk + # Twin of _make_manager's model root (same torch.manual_seed(7) + # draw order for the first Linear). + load_and_apply_loras(model, [str(p) for p in paths], "diffusion_uncond") + return model + + +@pytest.mark.parametrize( + "adapter_type", ["lora", "dora-rows", "dora-cols", "bora"], +) +def test_variant_parity_with_upstream_loader(tmp_path, adapter_type): + lora_model, _ = _vendored_or_skip() + p = _write_variant_lora(tmp_path, f"v_{adapter_type}", adapter_type) + + upstream = _upstream_model_with([p]) + mgr, ours, _cond = _make_manager() + mgr.enable_lora(mgr.register_lora(str(p)), strength=1.0) + + assert torch.equal( + ours.blk.lin.weight.detach(), upstream.blk.lin.weight.detach(), + ), f"{adapter_type}: manager != upstream loader at strength 1" + + # Strength parity away from 1 (nonlinear for DoRA/BoRA). + mgr.set_lora_strength(f"v_{adapter_type}", 0.7) + lora_model.set_lora_strength(upstream, 0.7, lora_index=0) + assert torch.equal( + ours.blk.lin.weight.detach(), upstream.blk.lin.weight.detach(), + ), f"{adapter_type}: manager != upstream loader at strength 0.7" + + +def test_multi_lora_composition_order_matches_upstream(tmp_path): + """Stacked adapters compose sequentially (order-dependent for + DoRA); the manager's enable order must reproduce upstream's load + order exactly — and a different order is genuinely different.""" + pA = _write_variant_lora(tmp_path, "stackA", "dora-rows", seed=21) + pB = _write_variant_lora(tmp_path, "stackB", "lora", seed=22) + + upstream_ab = _upstream_model_with([pA, pB]) + mgr, ours, _cond = _make_manager() + mgr.enable_lora(mgr.register_lora(str(pA)), strength=1.0) + mgr.enable_lora(mgr.register_lora(str(pB)), strength=1.0) + assert torch.equal( + ours.blk.lin.weight.detach(), upstream_ab.blk.lin.weight.detach(), + ) + + # The reverse enable order equals upstream's reverse load order... + upstream_ba = _upstream_model_with([pB, pA]) + mgr2, ours2, _cond2 = _make_manager() + mgr2.enable_lora(mgr2.register_lora(str(pB)), strength=1.0) + mgr2.enable_lora(mgr2.register_lora(str(pA)), strength=1.0) + assert torch.equal( + ours2.blk.lin.weight.detach(), upstream_ba.blk.lin.weight.detach(), + ) + # ...and DoRA-in-the-stack makes the two orders genuinely differ, + # which is why enable order is part of the parity contract. + assert not torch.equal( + upstream_ab.blk.lin.weight.detach(), + upstream_ba.blk.lin.weight.detach(), + ) diff --git a/tests/unit/test_sa3_backend.py b/tests/unit/test_sa3_backend.py index 26b90d73..92744825 100644 --- a/tests/unit/test_sa3_backend.py +++ b/tests/unit/test_sa3_backend.py @@ -660,3 +660,84 @@ def test_close_tears_down_lora_manager(): b = _backend(lora_manager=mgr, use_lora=True) b.close() assert mgr.closed is True + + +# --------------------------------------------------------------------------- +# Phase-2 refit-mirror routing (fake mirror; the real one is GPU-gated) +# --------------------------------------------------------------------------- + + +class _FakeMirror: + def __init__(self): + self.syncs: list = [] + + def sync(self, *, reason=""): + self.syncs.append(reason) + return 1 + + +def test_mirror_mode_never_swaps_dit_and_syncs_on_enable_disable(): + mgr = _FakeLoraManager(ids=["x"]) + mirror = _FakeMirror() + b = _backend( + lora_manager=mgr, use_lora=True, eager_dit=_ZeroDit(), + refit_mirror=mirror, + ) + accel = b._dit_accel + assert accel is not b._dit_eager # distinct objects: swap WOULD fire + + b.enable_lora("x", strength=1.0) + assert b.adapter.dit is accel # no eager swap in mirror mode + assert mirror.syncs == ["enable_lora"] + + b.disable_lora("x") + assert b.adapter.dit is accel + assert mirror.syncs == ["enable_lora", "disable_lora"] + + +def test_mirror_strength_routes_through_pending_stash(): + mgr = _FakeLoraManager(ids=["x"]) + mgr.enable_lora("x", strength=1.0) + mirror = _FakeMirror() + b = _backend(lora_manager=mgr, use_lora=True, refit_mirror=mirror) + mirror.syncs.clear() + + knobs = _knobs(b, **{"lora_str_x": 0.4}) + # The announcement point: rebuild_imminent stashes and reports. + assert b.rebuild_imminent(knobs) is True + assert b._pending_lora_strengths == {"x": 0.4} + assert b.has_pending_refit() is True + + # The same tick's produce applies the stash: one manager write, one + # batched mirror sync, stash drained. + b.produce(knobs, CTX, "skip") + assert ("strength", "x", 0.4) in mgr.calls + assert mirror.syncs == ["strength"] + assert b._pending_lora_strengths == {} + assert b.has_pending_refit() is False + + # Within the 0.02 gate: no announcement, no stash. + assert b.rebuild_imminent(_knobs(b, **{"lora_str_x": 0.41})) is False + assert b._pending_lora_strengths == {} + + +def test_eager_mode_strength_is_not_stashed(): + mgr = _FakeLoraManager(ids=["x"]) + mgr.enable_lora("x", strength=1.0) + b = _backend(lora_manager=mgr, use_lora=True) # no mirror + + knobs = _knobs(b, **{"lora_str_x": 0.4}) + assert b.rebuild_imminent(knobs) is False # buffer write, no stall + b.produce(knobs, CTX, "skip") + assert ("strength", "x", 0.4) in mgr.calls # applied inline + + +def test_close_drops_mirror_and_pending(): + mgr = _FakeLoraManager(ids=["x"]) + mirror = _FakeMirror() + b = _backend(lora_manager=mgr, use_lora=True, refit_mirror=mirror) + b._pending_lora_strengths["x"] = 0.5 + b.close() + assert b._refit_mirror is None + assert b._pending_lora_strengths == {} + assert mgr.closed is True diff --git a/tests/unit/test_sa3_trt_lora.py b/tests/unit/test_sa3_trt_lora.py new file mode 100644 index 00000000..275a8ffd --- /dev/null +++ b/tests/unit/test_sa3_trt_lora.py @@ -0,0 +1,196 @@ +"""SA3 TRT LoRA refit mirror (phase 2) — CPU tests. + +The mirror's contract pieces that don't need an engine: manifest +consumption, merged-weight composition via the live parametrizations +(real vendored code on a tiny tree; skips without the vendor checkout), +transpose orientation, the dirty-set (a disabled adapter's BASE weights +must go back to the engine on the next sync), and single-commit +batching — all against a recording fake refitter. The real-engine +bit-identity gate lives in scripts/sa3/gen_sa3_refit_manifest.py +--validate-engine and runs on the GPU with a refit-built engine. +""" + +from __future__ import annotations + +import json +import sys +import types +from pathlib import Path + +import pytest +import torch +import torch.nn as nn + +sys.path.insert(0, str(Path(__file__).parent.parent.parent)) + +from acestep.engine.sa3_trt_lora import SA3TRTRefitMirror + + +class _FakeTrt(types.SimpleNamespace): + def __init__(self): + super().__init__( + float32="F32", float16="F16", + Weights=lambda dtype, ptr, size: (dtype, ptr, size), + ) + + +class _FakeRefitter: + def __init__(self, names): + self._names = list(names) + self.sets: list = [] + self.commits = 0 + + def get_all_weights(self): + return list(self._names) + + def set_named_weights(self, name, weights): + self.sets.append(name) + return name in self._names + + def refit_cuda_engine(self): + self.commits += 1 + return True + + +def _vendored_manager_or_skip(model_root): + from acestep.engine.sa3_helpers import ensure_sa3_paths + + ensure_sa3_paths() + pytest.importorskip( + "stable_audio_3.models.lora.model", + reason="SA3 vendored source not on disk", + ) + from acestep.engine.sa3_lora import SA3LoRAManager + + return SA3LoRAManager( + model_root=model_root, conditioner_root=None, checkpoint_id="medium", + ) + + +def _write_lora_file(tmp_path, fqn="blk.lin", fan=(8, 16), rank=4): + from stable_audio_3.models.lora.utils import save_lora_safetensors + + g = torch.Generator().manual_seed(3) + sd = { + f"{fqn}.parametrizations.weight.0.lora_A": + torch.randn(rank, fan[1], generator=g), + f"{fqn}.parametrizations.weight.0.lora_B": + torch.randn(fan[0], rank, generator=g), + } + p = tmp_path / "adapter.safetensors" + save_lora_safetensors( + sd, {"rank": rank, "alpha": rank, "adapter_type": "lora"}, p, + ) + return p + + +def _mirror(tmp_path, model_root, *, transposed=True): + manifest = tmp_path / "manifest.json" + manifest.write_text(json.dumps({ + "version": 1, + "weights": { + "blk.lin": { + "initializer": "onnx::MatMul_7", "transposed": transposed, + }, + }, + }), encoding="utf-8") + refitter = _FakeRefitter(["onnx::MatMul_7"]) + m = SA3TRTRefitMirror( + object(), model_root, manifest, + _refitter=refitter, _trt=_FakeTrt(), + ) + return m, refitter + + +def _model(): + torch.manual_seed(11) + model = nn.Module() + blk = nn.Module() + blk.lin = nn.Linear(16, 8, bias=False) + model.blk = blk + return model + + +def test_manifest_version_and_coverage_validated(tmp_path): + model = _model() + bad = tmp_path / "bad.json" + bad.write_text(json.dumps({"version": 99, "weights": {}}), encoding="utf-8") + with pytest.raises(RuntimeError, match="version"): + SA3TRTRefitMirror( + object(), model, bad, + _refitter=_FakeRefitter(["x"]), _trt=_FakeTrt(), + ) + unmapped = tmp_path / "unmapped.json" + unmapped.write_text(json.dumps({ + "version": 1, + "weights": {"blk.lin": {"initializer": "not_in_engine"}}, + }), encoding="utf-8") + with pytest.raises(RuntimeError, match="refittable"): + SA3TRTRefitMirror( + object(), model, unmapped, + _refitter=_FakeRefitter(["other"]), _trt=_FakeTrt(), + ) + + +def test_sync_noop_without_parametrizations(tmp_path): + model = _model() + mirror, refitter = _mirror(tmp_path, model) + assert mirror.sync(reason="test") == 0 + assert refitter.commits == 0 + assert refitter.sets == [] + + +def test_sync_pushes_merged_weight_and_restores_base_on_disable(tmp_path): + model = _model() + base = model.blk.lin.weight.detach().clone() + mgr = _vendored_manager_or_skip(model) + lora_file = _write_lora_file(tmp_path) + mirror, refitter = _mirror(tmp_path, model, transposed=True) + + lid = mgr.register_lora(str(lora_file)) + mgr.enable_lora(lid, strength=1.0) + merged = model.blk.lin.weight.detach().clone() + assert not torch.equal(merged, base) + + # Enable-sync: pushes the merged value (transposed orientation). + assert mirror.sync(reason="enable") == 1 + assert refitter.commits == 1 + assert refitter.sets == ["onnx::MatMul_7"] + staged = mirror._staging["blk.lin"] + assert torch.equal(staged, merged.transpose(0, 1).to(staged.dtype)) + assert "blk.lin" in mirror._dirty + + # Strength change re-pushes. + mgr.set_lora_strength(lid, 0.5) + assert mirror.sync(reason="strength") == 1 + assert refitter.commits == 2 + half = model.blk.lin.weight.detach() + assert torch.equal( + mirror._staging["blk.lin"], half.transpose(0, 1), + ) + + # Disable: the module is unparametrized again, but the engine slot + # still holds the merged value — the dirty-set forces a base push. + mgr.disable_lora(lid) + assert mirror.sync(reason="disable") == 1 + assert refitter.commits == 3 + assert torch.equal( + mirror._staging["blk.lin"], base.transpose(0, 1), + ) + assert "blk.lin" not in mirror._dirty + + # Nothing left to push. + assert mirror.sync(reason="idle") == 0 + assert refitter.commits == 3 + + +def test_untransposed_mapping_pushes_torch_orientation(tmp_path): + model = _model() + mgr = _vendored_manager_or_skip(model) + lora_file = _write_lora_file(tmp_path) + mirror, _refitter = _mirror(tmp_path, model, transposed=False) + mgr.enable_lora(mgr.register_lora(str(lora_file)), strength=1.0) + assert mirror.sync() == 1 + assert torch.equal( + mirror._staging["blk.lin"], model.blk.lin.weight.detach(), + ) From 3567c7c5f4ead4dc871136e5c7805ce20c909c8c Mon Sep 17 00:00:00 2001 From: RyanOnTheInside <7623207+ryanontheinside@users.noreply.github.com> Date: Wed, 22 Jul 2026 22:09:28 -0400 Subject: [PATCH 5/6] feat(lora): accept underfit's short sa3-* lineage spelling A real underfit-trained LoRA found in the wild writes base_model="sa3-medium" in its lora_config header; canonical_sa3_lineage previously fell back to permissive-unknown on that spelling instead of a positive lineage match. --- acestep/lora_metadata.py | 3 +++ tests/unit/test_lora_metadata.py | 3 +++ 2 files changed, 6 insertions(+) diff --git a/acestep/lora_metadata.py b/acestep/lora_metadata.py index a9771030..928d3c18 100644 --- a/acestep/lora_metadata.py +++ b/acestep/lora_metadata.py @@ -261,6 +261,9 @@ def canonical_sa3_lineage(value: Optional[str]) -> Optional[str]: s = s.rsplit("/", 1)[-1] if s.startswith("stable-audio-3-"): s = s[len("stable-audio-3-"):] + elif s.startswith("sa3-"): + # underfit writes e.g. base_model="sa3-medium" (seen in the wild) + s = s[len("sa3-"):] if s.endswith("-base"): s = s[: -len("-base")] return s if s in _SA3_RUNTIME_LINEAGES else None diff --git a/tests/unit/test_lora_metadata.py b/tests/unit/test_lora_metadata.py index 37068950..566450ee 100644 --- a/tests/unit/test_lora_metadata.py +++ b/tests/unit/test_lora_metadata.py @@ -535,6 +535,9 @@ def test_to_wire_carries_family_fields(tmp_path): ("stabilityai/stable-audio-3-medium", "medium"), ("stabilityai/stable-audio-3-medium-base", "medium"), ("stable-audio-3-small-music-base", "small-music"), + # underfit's short trainer spelling (seen in the wild) + ("sa3-medium", "medium"), + ("sa3-small-music-base", "small-music"), # case-insensitive, whitespace-tolerant ("Medium-Base", "medium"), (" small-music ", "small-music"), From bb279b00d41edf38f9e078d1e333ebcabdd32af2 Mon Sep 17 00:00:00 2001 From: RyanOnTheInside <7623207+ryanontheinside@users.noreply.github.com> Date: Wed, 22 Jul 2026 22:09:48 -0400 Subject: [PATCH 6/6] feat(sa3-lora): eager-vs-refit parity harness + listening renderer GPU validation of the TRT refit path: refit engine built idle-GPU; manifest bit-identity gate PASS (228/228 weights, 0 unmatched or ambiguous); lora_smoke PASS twice with a real underfit-trained LoRA (goa_trance_sa3m); eager-vs-refit per-step cos >= 0.99989 at strengths 0/0.5/1.0 and with a 2-LoRA stack; strength-0 and both restore phases value-identical; full 8-step denoise final-latent cos 0.99947. Status ledgers in notes/SA3_LORA_PLAN.md (local) updated to match. --- scripts/sa3/lora_listening_render.py | 156 ++++++++++++++++++ scripts/sa3/lora_refit_parity.py | 238 +++++++++++++++++++++++++++ 2 files changed, 394 insertions(+) create mode 100644 scripts/sa3/lora_listening_render.py create mode 100644 scripts/sa3/lora_refit_parity.py diff --git a/scripts/sa3/lora_listening_render.py b/scripts/sa3/lora_listening_render.py new file mode 100644 index 00000000..b8c5433a --- /dev/null +++ b/scripts/sa3/lora_listening_render.py @@ -0,0 +1,156 @@ +"""Render SA3 LoRA listening WAVs (notes/SA3_LORA_PLAN.md Phase 3 signoff). + +Fixed prompt + seed through upstream's own ``StableAudioModel.generate`` +(the real SA3 pipeline: T5Gemma conditioning, pingpong sampler, VAE +decode), producing: + +* baseline (no LoRA), eager DiT; +* the LoRA at chosen strengths, eager DiT (parametrized weights); +* the LoRA at 1.0 with the DiT routed through the refit-mirrored TRT + engine — the same upstream sampler loop, ``sam.model.model`` swapped + for a shim that calls ``SA3TRTDit.step_bundle``. + +Duration defaults to 54 s (the refit engine's all-valid L=646 window) +with the saved WAV trimmed to ``--save-seconds``, so the eager and TRT +renders share identical latent geometry and noise draws. + +Run: + .venv/Scripts/python.exe scripts/sa3/lora_listening_render.py \ + --lora --out-dir notes/sa3_lora_listening +""" + +from __future__ import annotations + +import argparse +import sys +from pathlib import Path + +_HERE = Path(__file__).resolve().parent +_REPO_ROOT = next(p for p in (_HERE, *_HERE.parents) if (p / "pyproject.toml").exists()) +for _p in (str(_REPO_ROOT),): + while _p in sys.path: + sys.path.remove(_p) + sys.path.insert(0, _p) + +import torch # noqa: E402 + + +class _TRTShim(torch.nn.Module): + """Routes DiT forwards inside upstream's sampler to the TRT engine. + Holds the eager module so ``parameters()``-based dtype probes in + ``generate`` keep working.""" + + def __init__(self, trt_dit, eager_dit): + super().__init__() + self._trt = trt_dit + self._eager = eager_dit + + def forward(self, x, t, **kwargs): + t_val = float(t.flatten()[0].item()) if torch.is_tensor(t) else float(t) + # Persistent output buffer -> materialize; match eager dtype so + # sampler arithmetic and randn_like draws stay path-identical. + return self._trt.step_bundle(x, t_val, kwargs).to(dtype=x.dtype, copy=True) + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--model", default="medium") + ap.add_argument("--lora", required=True) + ap.add_argument("--prompt", + default="oldschool goa trance, 145 bpm, hypnotic acid leads, " + "rolling bassline") + ap.add_argument("--seed", type=int, default=42) + ap.add_argument("--steps", type=int, default=8) + ap.add_argument("--duration", type=float, default=54.0) + ap.add_argument("--save-seconds", type=float, default=30.0) + ap.add_argument("--strengths", type=float, nargs="*", default=[0.5, 1.0, 1.5]) + ap.add_argument("--trt-strength", type=float, default=1.0) + ap.add_argument("--out-dir", default=str(_REPO_ROOT / "notes" / "sa3_lora_listening")) + args = ap.parse_args() + + import soundfile as sf + import torch.nn.utils.parametrize as parametrize + + from acestep.engine.sa3_context import SA3Context + from acestep.engine.sa3_lora import SA3LoRAManager + + out_dir = Path(args.out_dir) + out_dir.mkdir(parents=True, exist_ok=True) + + print(f"loading SA3Context({args.model!r})...") + ctx = SA3Context(model_id=args.model) + sam = ctx.sam + sr = sam.model.sample_rate + save_samples = int(args.save_seconds * sr) + + def render(tag: str) -> Path: + with torch.no_grad(), parametrize.cached(): + audio = sam.generate( + prompt=args.prompt, + duration=args.duration, + steps=args.steps, + seed=args.seed, + cfg_scale=1.0, + ) + wav = audio[0, :, :save_samples].T.float().cpu().numpy() + path = out_dir / f"{tag}.wav" + sf.write(str(path), wav, samplerate=sr) + peak = float(abs(audio).max()) + print(f"[render] {path.name} peak={peak:.3f}") + return path + + stem = Path(args.lora).stem + + # --- baseline (no LoRA), eager ------------------------------------- + render("01_baseline_eager") + + # --- eager strengths ---------------------------------------------- + mgr = SA3LoRAManager( + model_root=sam.model.model, + conditioner_root=sam.model.conditioner, + checkpoint_id=ctx.model_id, + ) + lid = mgr.register_lora(args.lora) + mgr.prewarm_lora(lid).result(timeout=180) + first = True + for i, s in enumerate(args.strengths): + if first: + mgr.enable_lora(lid, strength=s) + first = False + else: + mgr.set_lora_strength(lid, s) + render(f"{i + 2:02d}_{stem}_s{s:g}_eager") + + # --- TRT refit render ---------------------------------------------- + from acestep.engine.sa3_trt import SA3TRTDit, find_dit_engine + from acestep.engine.sa3_trt_lora import SA3TRTRefitMirror, find_refit_manifest + + # Latent geometry of the padded window generate() will use. + cond = ctx.prepare_cond(prompt=args.prompt, duration=args.duration, + steps=args.steps) + L = cond.latent_frames + engine_path = find_dit_engine(args.model, L, want_refittable=True) + if engine_path is None: + raise RuntimeError(f"no refittable TRT DiT engine for L={L}") + trt_dit = SA3TRTDit(engine_path, latent_frames=L, seconds_total=args.duration) + manifest = find_refit_manifest(engine_path) + if manifest is None: + raise RuntimeError("no refit manifest; run gen_sa3_refit_manifest.py") + mirror = SA3TRTRefitMirror(trt_dit.engine, sam.model.model, manifest) + + mgr.set_lora_strength(lid, args.trt_strength) + mirror.sync(reason="listening-render") + real_dit = sam.model.model + sam.model.model = _TRTShim(trt_dit, real_dit) + try: + render(f"{len(args.strengths) + 2:02d}_{stem}_s{args.trt_strength:g}_trt_refit") + finally: + sam.model.model = real_dit + + mgr.close() + print("done") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/sa3/lora_refit_parity.py b/scripts/sa3/lora_refit_parity.py new file mode 100644 index 00000000..96bd04b6 --- /dev/null +++ b/scripts/sa3/lora_refit_parity.py @@ -0,0 +1,238 @@ +"""Eager-vs-refit LoRA output parity (notes/SA3_LORA_PLAN.md Phase 2 item 5). + +Runs the eager parametrized DiT against the refit-mirrored TRT engine on +identical inputs through the LoRA lifecycle: + +* baseline (no LoRA) per-step cos — the engine-numerics floor; +* one LoRA at strengths 0.0 / 0.5 / 1.0 (strength 0 must leave the TRT + outputs value-identical to baseline: a zero delta merges to the exact + base weights); +* a second stacked LoRA; +* disable of the second (the mirror's dirty-set must push the first + LoRA's merged weights back — TRT outputs value-identical to the + single-LoRA phase); +* disable of the first (full restore — TRT outputs value-identical to + baseline); +* full 8-step pingpong denoise (upstream's own sampler) with the LoRA + at strength 1.0, eager vs TRT-shimmed, same seed: final-latent cos. + +Per-step comparisons use the same cond bundle on both paths, so +conditioner-side LoRA keys (which can never reach the TRT graph — the +seconds_total tail is baked in as constants, plan D5) do not skew the +DiT-level verdict. + +Run: + .venv/Scripts/python.exe scripts/sa3/lora_refit_parity.py --lora + (add --lora2 for the stacking phase; defaults to the smoke + synthetic adapter, synthesizing it if absent) +""" + +from __future__ import annotations + +import argparse +import sys +from pathlib import Path + +_HERE = Path(__file__).resolve().parent +_REPO_ROOT = next(p for p in (_HERE, *_HERE.parents) if (p / "pyproject.toml").exists()) +for _p in (str(_REPO_ROOT),): + while _p in sys.path: + sys.path.remove(_p) + sys.path.insert(0, _p) + +import torch # noqa: E402 + +PROMPT = "warm analog house groove, 124 bpm" +DURATION = 54.0 +STEPS = 8 +T_VALUES = (1.0, 0.7, 0.4, 0.1) +SEED = 1528 +# Engine-numerics floor from the pre-LoRA TRT parity signoff +# (scripts/sa3/sa3_trt_dit_cond_parity.py): fp16mixed cos >= 0.9998/step +# on all-valid windows. +COS_FLOOR = 0.9998 + + +class _TRTShim(torch.nn.Module): + """Routes DiT forwards to the TRT engine inside upstream's sampler. + Holds the eager module so ``parameters()`` keeps reporting the model + dtype to ``generate``-style callers.""" + + def __init__(self, trt_dit, eager_dit): + super().__init__() + self._trt = trt_dit + self._eager = eager_dit + + def forward(self, x, t, **kwargs): + t_val = float(t.flatten()[0].item()) if torch.is_tensor(t) else float(t) + # step_bundle returns its persistent fp32 output buffer — + # materialize and match the eager dtype so sampler arithmetic + # (and randn_like draws) stay identical across paths. + return self._trt.step_bundle(x, t_val, kwargs).to(dtype=x.dtype, copy=True) + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--model", default="medium") + ap.add_argument("--lora", required=True, help="trained SA3 LoRA (primary)") + ap.add_argument("--lora2", default=None, + help="second LoRA for the stacking phase " + "(default: the smoke synthetic adapter)") + args = ap.parse_args() + + from acestep.engine.sa3_context import SA3Context + from acestep.engine.sa3_lora import SA3LoRAManager + from acestep.engine.sa3_trt import SA3TRTDit, find_dit_engine + from acestep.engine.sa3_trt_lora import SA3TRTRefitMirror, find_refit_manifest + + print(f"loading SA3Context({args.model!r})...") + ctx = SA3Context(model_id=args.model) + sam = ctx.sam + + cond = ctx.prepare_cond(prompt=PROMPT, duration=DURATION, steps=STEPS) + L = cond.latent_frames + from acestep.engine.sa3_stream_helpers import stack_sa3_cond_bundles + stacked = stack_sa3_cond_bundles([cond.cond_bundle]) + print(f"[cond] latent_frames={L}") + + engine_path = find_dit_engine(args.model, L, want_refittable=True) + if engine_path is None: + raise RuntimeError(f"no refittable TRT DiT engine for L={L}; " + f"build with sa3_build --dit --refit") + trt_dit = SA3TRTDit(engine_path, latent_frames=L, seconds_total=DURATION) + if not trt_dit.refittable: + raise RuntimeError(f"selected engine is not refit-built: {engine_path}") + manifest = find_refit_manifest(engine_path) + if manifest is None: + raise RuntimeError("no refit manifest; run gen_sa3_refit_manifest.py") + mirror = SA3TRTRefitMirror(trt_dit.engine, sam.model.model, manifest) + + mgr = SA3LoRAManager( + model_root=sam.model.model, + conditioner_root=sam.model.conditioner, + checkpoint_id=ctx.model_id, + ) + + lora2 = args.lora2 + if lora2 is None: + from acestep.paths import models_dir + import lora_smoke as smoke # scripts/sa3 on sys.path via ensure_sa3_paths + + default2 = models_dir() / "sa3" / "smoke" / "smoke_synthetic_lora.safetensors" + if not default2.is_file(): + default2.parent.mkdir(parents=True, exist_ok=True) + smoke._synthesize_lora(ctx, default2.parent) + lora2 = str(default2) + + import torch.nn.utils.parametrize as parametrize + + def compare(tag: str): + """Per-step eager-vs-TRT at fixed t values on identical x. + Returns (worst_cos, {t: trt_out}).""" + worst = 1.0 + trt_outs = {} + g = torch.Generator(device=ctx.device.type).manual_seed(SEED) + for t_val in T_VALUES: + x = torch.randn(1, ctx.latent_channels, L, device=ctx.device, + dtype=ctx.dtype, generator=g) + t_b = torch.full((1,), t_val, device=ctx.device, dtype=ctx.dtype) + with torch.no_grad(), parametrize.cached(): + v_eager = ctx.dit(x, t_b, **stacked).float() + v_trt = trt_dit.step_bundle(x, t_val, stacked).float().clone() + cos = torch.nn.functional.cosine_similarity( + v_eager.flatten(), v_trt.flatten(), dim=0).item() + rel = ((v_trt - v_eager).norm() / v_eager.norm()).item() + worst = min(worst, cos) + trt_outs[t_val] = v_trt + print(f" [{tag}] t={t_val:.2f} cos={cos:.6f} rel_rms={rel:.4e}") + return worst, trt_outs + + def outs_equal(a: dict, b: dict) -> bool: + return all(torch.equal(a[t], b[t]) for t in T_VALUES) + + failures = [] + + def gate(name: str, ok: bool): + print(f" -> {name}: {'PASS' if ok else 'FAIL'}") + if not ok: + failures.append(name) + + print("\n=== baseline (no LoRA) ===") + worst, base_outs = compare("base") + gate(f"baseline worst cos {worst:.6f} >= {COS_FLOOR}", worst >= COS_FLOOR) + + lid1 = mgr.register_lora(args.lora) + mgr.prewarm_lora(lid1).result(timeout=180) + print(f"\nlora1={Path(args.lora).name} " + f"touches_conditioner={mgr.touches_conditioner(lid1)}") + + print(f"\n=== lora1 @ 0.0 ===") + mgr.enable_lora(lid1, strength=0.0) + mirror.sync(reason="parity-enable-s0") + worst, s0_outs = compare("s0.0") + gate("strength-0 TRT outputs identical to baseline", outs_equal(s0_outs, base_outs)) + gate(f"s0 worst cos {worst:.6f} >= {COS_FLOOR}", worst >= COS_FLOOR) + + s1_outs = None + for s in (0.5, 1.0): + print(f"\n=== lora1 @ {s} ===") + mgr.set_lora_strength(lid1, s) + mirror.sync(reason=f"parity-s{s}") + worst, outs = compare(f"s{s}") + gate(f"s{s} worst cos {worst:.6f} >= {COS_FLOOR}", worst >= COS_FLOOR) + if s == 1.0: + s1_outs = outs + + print(f"\n=== + lora2 (stacked) ===") + lid2 = mgr.register_lora(lora2) + mgr.prewarm_lora(lid2).result(timeout=180) + mgr.enable_lora(lid2, strength=0.7) + mirror.sync(reason="parity-stack") + worst, _ = compare("stack") + gate(f"stacked worst cos {worst:.6f} >= {COS_FLOOR}", worst >= COS_FLOOR) + + print(f"\n=== disable lora2 (dirty-set restore) ===") + mgr.disable_lora(lid2) + mirror.sync(reason="parity-drop2") + worst, outs = compare("drop2") + gate("post-drop2 TRT outputs identical to lora1@1.0 phase", + outs_equal(outs, s1_outs)) + + print(f"\n=== disable lora1 (full restore) ===") + mgr.disable_lora(lid1) + mirror.sync(reason="parity-drop1") + worst, outs = compare("drop1") + gate("post-drop-all TRT outputs identical to baseline", + outs_equal(outs, base_outs)) + + print(f"\n=== full {STEPS}-step pingpong denoise, lora1 @ 1.0, seed {SEED} ===") + from stable_audio_3.inference.sampling import sample_flow_pingpong + + mgr.enable_lora(lid1, strength=1.0) + mirror.sync(reason="parity-fullloop") + sigmas = ctx.make_schedule_builder(cond, STEPS)(1.0).detach().float() + + def full_loop(model) -> torch.Tensor: + torch.manual_seed(SEED) + x = torch.randn(1, ctx.latent_channels, L, device=ctx.device, dtype=ctx.dtype) + with torch.no_grad(), parametrize.cached(): + return sample_flow_pingpong( + model, x, sigmas.to(ctx.device), disable_tqdm=True, **stacked, + ).float().clone() + + lat_eager = full_loop(ctx.dit) + lat_trt = full_loop(_TRTShim(trt_dit, ctx.dit)) + cos = torch.nn.functional.cosine_similarity( + lat_eager.flatten(), lat_trt.flatten(), dim=0).item() + rel = ((lat_trt - lat_eager).norm() / lat_eager.norm()).item() + print(f" final-latent cos={cos:.6f} rel_rms={rel:.4e}") + gate(f"final-latent cos {cos:.6f} >= 0.99", cos >= 0.99) + + mgr.close() + verdict = "PASS" if not failures else f"FAIL ({', '.join(failures)})" + print(f"\n=== lora_refit_parity: {verdict} ===") + return 0 if not failures else 1 + + +if __name__ == "__main__": + raise SystemExit(main())