diff --git a/installer_files b/installer_files new file mode 120000 index 0000000000..7c77e509a9 --- /dev/null +++ b/installer_files @@ -0,0 +1 @@ +/home/rafael/unsloth/installer_files \ No newline at end of file diff --git a/start.sh b/start.sh new file mode 100755 index 0000000000..803fe1b749 --- /dev/null +++ b/start.sh @@ -0,0 +1,51 @@ +#!/usr/bin/env bash +# Build the studio frontend from local source, then serve it via the REPO's +# studio backend (studio/backend/run.py) run under the installed studio venv, +# so local BACKEND edits take effect (plain `unsloth studio` runs the installed +# package's backend instead). Uses the self-contained Node in +# installer_files/node so the system Node is never touched. +set -euo pipefail + +PORT="${PORT:-8888}" +HOST="${HOST:-0.0.0.0}" + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +FRONTEND_DIR="$SCRIPT_DIR/studio/frontend" +NODE_BIN="$SCRIPT_DIR/installer_files/node/bin" + +if [ ! -x "$NODE_BIN/node" ]; then + echo "ERROR: local Node not found at $NODE_BIN" >&2 + echo "Install it once, e.g.:" >&2 + echo " curl -fL https://nodejs.org/dist/v22.12.0/node-v22.12.0-linux-x64.tar.xz | tar -xJ -C /tmp" >&2 + echo " mv /tmp/node-v22.12.0-linux-x64 $SCRIPT_DIR/installer_files/node" >&2 + exit 1 +fi + +export PATH="$NODE_BIN:$PATH" +echo "==> Using $(node -v) / npm $(npm -v)" + +cd "$FRONTEND_DIR" + +if [ ! -d node_modules ] || ! cmp -s package-lock.json node_modules/.lockfile-stamp; then + echo "==> Installing frontend dependencies" + npm ci + cp package-lock.json node_modules/.lockfile-stamp +fi + +echo "==> Building frontend ($FRONTEND_DIR)" +npm run build + +STUDIO_VENV_PY="${UNSLOTH_STUDIO_HOME:-$HOME/.unsloth/studio}/unsloth_studio/bin/python" +RUN_PY="$SCRIPT_DIR/studio/backend/run.py" + +if [ ! -x "$STUDIO_VENV_PY" ]; then + echo "ERROR: studio venv python not found at $STUDIO_VENV_PY" >&2 + echo "Run 'unsloth studio' once to create the venv, or set UNSLOTH_STUDIO_HOME." >&2 + exit 1 +fi + +echo "==> Starting repo backend on $HOST:$PORT (frontend: $FRONTEND_DIR/dist)" +echo " backend: $RUN_PY" +echo " python: $STUDIO_VENV_PY" +cd "$SCRIPT_DIR/studio/backend" +exec "$STUDIO_VENV_PY" "$RUN_PY" --host "$HOST" --port "$PORT" --frontend "$FRONTEND_DIR/dist" diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index eb4bb33034..c06e3cc359 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -246,6 +246,8 @@ class _LoadState: # Pre-warmed torch.compile cache context (diffusion_compile_cache.CacheContext) when a # compiled tier ran, else None. Carries the per-key inductor dir + bundle for save/restore. compile_cache_ctx: Any = None + # Token kept so LoRA adapters selected at generate time can be fetched from the Hub. + hf_token: Optional[str] = None @dataclass @@ -1145,6 +1147,7 @@ class DiffusionBackend: transformer_cache = cache_engaged, eager_patched = eager_patched, compile_cache_ctx = compile_ctx, + hf_token = hf_token, ) state_committed = True finally: @@ -1398,6 +1401,87 @@ class DiffusionBackend: except (StopIteration, AttributeError, RuntimeError): pass + def _apply_loras( + self, state: Any, loras: Optional[list[tuple[str, float]]], cancel: threading.Event + ) -> None: + """Load + activate requested LoRA adapters on ``state.pipe`` (non-fused), or clear + them when none are requested. + + The applied set is recorded on the pipe object, so an unchanged selection is a no-op + and a model swap (a fresh pipe with no marker) resets naturally. Never fuses: fusing + breaks on quantized (bnb-4bit / torchao) transformers and blocks live weight tweaks. + """ + from core.inference import diffusion_lora + + pipe = state.pipe + current = getattr(pipe, "_unsloth_loras", ()) + specs = [(i, w) for (i, w) in (loras or []) if w != 0] + + if not specs: + if current: + try: + pipe.unload_lora_weights() + except Exception: # noqa: BLE001 -- best-effort clear + pass + pipe._unsloth_loras = () + return + + if not diffusion_lora.supports_lora( + engine = "diffusers", + family = getattr(state.family, "name", None), + model_kind = state.kind, + transformer_quant = state.transformer_quant, + compiled = "compiled" in (getattr(state, "speed_optims", ()) or ()), + ): + raise ValueError( + "LoRA is not supported for this model/quantisation on the diffusers engine " + "(GGUF-via-diffusers, torchao fp8/int8, or a torch.compile'd Speed=default/max " + "load). Use a bf16 or bnb-4bit load at Speed=off/eager, or the native engine " + "for GGUF models." + ) + + resolved = diffusion_lora.resolve_specs(specs, hf_token = state.hf_token, cancel_event = cancel) + # The shared catalog scans both .safetensors and .gguf, but diffusers' + # load_lora_weights only takes safetensors; a .gguf adapter would otherwise fail + # deep in generation. Reject it here as a clean 400 before touching the pipe. + bad = [r.id for r in resolved if r.fmt != "safetensors"] + if bad: + raise ValueError( + "GGUF LoRA adapters are not supported on the diffusers engine " + f"({', '.join(bad)}); use a .safetensors adapter, or the native engine." + ) + # Unique adapter names (diffusers requires distinct names; sanitized stems can collide). + uniq: list[tuple[str, str, float]] = [] + seen: set[str] = set() + for r in resolved: + name = r.alias + n = 1 + while name in seen: + n += 1 + name = f"{r.alias}_{n}" + seen.add(name) + uniq.append((name, r.path, r.weight)) + + desired = tuple(uniq) + if desired == current: + return + try: + if current: + pipe.unload_lora_weights() + for name, path, _weight in uniq: + pipe.load_lora_weights(path, adapter_name = name) + pipe.set_adapters( + [name for name, _p, _w in uniq], adapter_weights = [w for _n, _p, w in uniq] + ) + except Exception as exc: # noqa: BLE001 -- surface as a clean 400 + try: + pipe.unload_lora_weights() + except Exception: # noqa: BLE001 + pass + pipe._unsloth_loras = () + raise ValueError(f"Failed to apply LoRA: {exc}") from exc + pipe._unsloth_loras = desired + @staticmethod def _reset_step_cache(pipe: Any) -> None: """Clear the transformer's stateful step cache (FBCache) before a generation. @@ -1445,6 +1529,9 @@ class DiffusionBackend: # pipeline accepts a list, so multiple references can be combined (subject + style, # character + scene). Ignored by non-reference workflows. reference_images: Optional[list[str]] = None, + # LoRA adapters as (id, weight) pairs; loaded onto the pipe (non-fused) and activated + # with set_adapters for this generation. None/empty = no LoRA (adapters cleared). + loras: Optional[list[tuple[str, float]]] = None, ) -> dict[str, Any]: import torch from PIL import Image @@ -1477,6 +1564,10 @@ class DiffusionBackend: seed = int(seed) generator.manual_seed(seed) + # Apply/adjust LoRA adapters on the resident pipe (non-fused) before picking + # the workflow pipe; from_pipe pipes share the transformer, so it propagates. + self._apply_loras(state, loras, cancel) + # Select the pipeline for this workflow. txt2img uses the loaded pipe; # img2img/inpaint reuse its resident modules via from_pipe (no reload); # an edit model's OWN loaded pipe is already the edit pipeline. @@ -1745,6 +1836,13 @@ class DiffusionBackend: uninstall_patches() uninstall_arch_patches() + # NOTE: we deliberately do NOT call state.pipe.unload_lora_weights() here. unload() + # sets the cancel event but does not take _generate_lock, so a LoRA-backed denoise + # can still be running on this same pipe for up to one more callback; mutating its + # adapter layers now would race that in-flight generation. The whole pipe is dropped + # just below (self._state = None; del state; clear_gpu_cache()), so the adapter + # tensors are freed with it -- no explicit unload is needed for memory or for a + # later load (which builds a fresh pipe). # Drop the workflow pipes built around this load's modules so they don't pin the # freed pipeline (they only re-wire its components, but holding the wrappers # would keep the modules alive past unload). @@ -1775,7 +1873,10 @@ class DiffusionBackend: "attention_backend": None, "transformer_cache": None, "workflows": [], + "supports_lora": False, } + from core.inference import diffusion_lora + return { "loaded": True, "repo_id": state.repo_id, @@ -1797,6 +1898,13 @@ class DiffusionBackend: # Image-conditioned workflows the loaded family supports, so the UI can gate # its tabs. txt2img is always available on the diffusers engine. "workflows": _family_workflows(state.family), + "supports_lora": diffusion_lora.supports_lora( + engine = "diffusers", + family = state.family.name, + model_kind = state.kind, + transformer_quant = state.transformer_quant, + compiled = "compiled" in (getattr(state, "speed_optims", ()) or ()), + ), } diff --git a/studio/backend/core/inference/diffusion_lora.py b/studio/backend/core/inference/diffusion_lora.py new file mode 100644 index 0000000000..7db07dcd60 --- /dev/null +++ b/studio/backend/core/inference/diffusion_lora.py @@ -0,0 +1,378 @@ +"""Shared LoRA support for the Studio diffusion backends. + +Both engines apply LoRA differently -- the native stable-diffusion.cpp CLI selects +adapters by `` prompt tags resolved against a `--lora-model-dir`, +while diffusers loads them with `load_lora_weights()` + `set_adapters()`. This module +holds the parts they share: a curated + local catalog, id->file resolution (with HF +download), a managed directory materialiser for the native tier, deterministic native +alias naming, and the single `supports_lora()` gate the UI/status and both backends use. + +The request layer only ever passes a LoRA *id* (a discovery id, local stem, or HF repo +id) plus a weight -- never a raw filesystem path -- so a client cannot make the backend +read an arbitrary file. Resolution validates the id against the catalog / the local LoRA +directory / the HF hub before anything is loaded. +""" + +from __future__ import annotations + +import os +import re +import threading +from dataclasses import dataclass, field +from pathlib import Path +from typing import Optional + +from utils.hf_xet_fallback import hf_hub_download_with_xet_fallback +from utils.paths.storage_roots import studio_root + +from .diffusion_families import DIFFUSION_CANCELLED_MSG + +# LoRA file formats we accept. sd-cli probes .safetensors/.gguf/.pt; diffusers loads +# .safetensors. We expose safetensors + gguf (pt is legacy/pickled -> excluded for safety). +_NATIVE_EXTS = (".safetensors", ".gguf") +_DIFFUSERS_EXTS = (".safetensors",) +_ALL_EXTS = (".safetensors", ".gguf") + + +@dataclass(frozen = True) +class LoraCatalogEntry: + """One discoverable LoRA adapter.""" + + id: str + display_name: str + source: str # "local" | "hub" + fmt: str # "safetensors" | "gguf" + # Family names this LoRA is compatible with (see diffusion_families). Empty = unknown + # (shown but not family-gated). Used by the UI to grey out incompatible adapters. + families: tuple[str, ...] = () + repo_id: Optional[str] = None # for source == "hub" + weight_name: Optional[str] = None # file within the repo (source == "hub") + local_path: Optional[str] = None # for source == "local" + size_bytes: int = 0 + weight_default: float = 1.0 + + +@dataclass(frozen = True) +class ResolvedLora: + """A LoRA resolved to a concrete local file, ready to apply.""" + + id: str + alias: str # sanitized stem used for the native tag / diffusers adapter name + path: str + fmt: str + weight: float + + +# Curated, family-tagged catalog of known-good diffusion LoRAs. Kept intentionally small +# and data-driven; extend as unsloth hosts/curates more. Entries are HF repos with a +# single-file weight. (Left minimal on purpose -- local discovery is the primary source, +# and users can also reference any public HF LoRA repo id directly.) +_CURATED: tuple[LoraCatalogEntry, ...] = () + + +def loras_dir() -> Path: + """Local directory Studio scans for user-provided diffusion LoRA files.""" + d = studio_root() / "loras" / "diffusion" + d.mkdir(parents = True, exist_ok = True) + return d + + +def sanitize_alias(raw: str) -> str: + """Deterministic, filesystem- and prompt-tag-safe alias from an id/stem. + + The native `` tag resolves NAME as a filename stem, so the alias must + contain no path separators, spaces, colons, or angle brackets. It is also used as the + diffusers PEFT adapter name, which additionally forbids "." (PEFT treats it as a module + path separator), so dots are replaced too -- many real LoRA filenames carry a version + like "V1.0". Collisions across sources are broken by the caller (materialize_native_dir + / the diffusers manager) with a numeric suffix. + """ + stem = raw.rsplit("/", 1)[-1] + for ext in _ALL_EXTS: + if stem.lower().endswith(ext): + stem = stem[: -len(ext)] + break + stem = re.sub(r"[^A-Za-z0-9_-]+", "_", stem).strip("_-") + return stem or "lora" + + +def _scan_local() -> list[LoraCatalogEntry]: + root = loras_dir() + try: + children = sorted(root.iterdir()) + except OSError: + return [] + files = [p for p in children if p.is_file() and p.suffix.lower() in _ALL_EXTS] + # Two files that share a stem but differ in extension (foo.safetensors + foo.gguf) + # would collide on id (== stem), so the frontend select value and resolve_one's + # id->entry lookup could only ever address one of them. Disambiguate a colliding + # stem by keeping the full filename as the id; a unique stem stays the clean stem. + stem_counts: dict[str, int] = {} + for p in files: + stem_counts[p.stem] = stem_counts.get(p.stem, 0) + 1 + entries: list[LoraCatalogEntry] = [] + for p in files: + ext = p.suffix.lower() + try: + size = p.stat().st_size + except OSError: + size = 0 + entry_id = p.name if stem_counts.get(p.stem, 0) > 1 else p.stem + entries.append( + LoraCatalogEntry( + id = entry_id, + display_name = entry_id, + source = "local", + fmt = "gguf" if ext == ".gguf" else "safetensors", + local_path = str(p), + size_bytes = size, + ) + ) + return entries + + +def list_loras(*, family: Optional[str] = None) -> list[LoraCatalogEntry]: + """Return the merged catalog (curated + local), optionally family-filtered. + + Cheap: a single directory scan plus the in-memory curated list. Network is only + touched later, on resolve(), when a hub adapter is actually selected. + """ + merged = list(_CURATED) + _scan_local() + if family: + fam = family.strip().lower() + merged = [e for e in merged if not e.families or fam in {f.lower() for f in e.families}] + # Stable order: local first (user intent), then by display name. + merged.sort(key = lambda e: (e.source != "local", e.display_name.lower())) + return merged + + +def _catalog_by_id() -> dict[str, LoraCatalogEntry]: + return {e.id: e for e in (list(_CURATED) + _scan_local())} + + +def resolve_one( + spec_id: str, + weight: float, + *, + hf_token: Optional[str] = None, + cancel_event: Optional[threading.Event] = None, +) -> ResolvedLora: + """Resolve a request LoRA id + weight to a concrete local file. + + Accepts, in order: a catalog/local id, or a bare HF repo id (``owner/name`` with an + optional ``owner/name:weight_file.safetensors`` suffix). Downloads hub weights via the + shared xet-fallback helper. Raises FileNotFoundError/ValueError on an unresolvable or + unsupported id -- the caller maps that to a clear 400. + """ + # An empty / whitespace token sent verbatim to HfApi triggers an auth error instead + # of falling back to anonymous access; normalise it to None. + hf_token = hf_token.strip() if hf_token and hf_token.strip() else None + entry = _catalog_by_id().get(spec_id) + if entry is not None: + if entry.source == "local": + path = entry.local_path or "" + if not path or not os.path.exists(path): + raise FileNotFoundError(f"LoRA '{spec_id}' is no longer present on disk") + return ResolvedLora(spec_id, sanitize_alias(spec_id), path, entry.fmt, weight) + # hub catalog entry + if not entry.repo_id or not entry.weight_name: + raise ValueError(f"LoRA '{spec_id}' has no downloadable weight") + path = hf_hub_download_with_xet_fallback( + entry.repo_id, entry.weight_name, hf_token, cancel_event = cancel_event + ) + return ResolvedLora(spec_id, sanitize_alias(spec_id), path, entry.fmt, weight) + + # Not in the catalog: allow a bare public HF repo id (owner/name[:weight_file]). + if "/" in spec_id: + repo_id, _, weight_name = spec_id.partition(":") + weight_name = weight_name or None + if weight_name is not None: + # A client-supplied weight file must stay a plain filename inside the repo: + # reject traversal / absolute paths so it can never resolve outside the HF + # cache dir once handed to the downloader. + if ( + ".." in weight_name + or weight_name.startswith(("/", "\\", "~")) + or "\\" in weight_name + or os.path.isabs(weight_name) + ): + raise ValueError(f"invalid LoRA weight file path '{weight_name}'") + if weight_name is None: + weight_name = _pick_repo_weight_file(repo_id, hf_token) + ext = os.path.splitext(weight_name)[1].lower() + if ext not in _ALL_EXTS: + raise ValueError(f"unsupported LoRA file '{weight_name}' (need .safetensors/.gguf)") + path = hf_hub_download_with_xet_fallback( + repo_id, weight_name, hf_token, cancel_event = cancel_event + ) + fmt = "gguf" if ext == ".gguf" else "safetensors" + return ResolvedLora(spec_id, sanitize_alias(repo_id), path, fmt, weight) + + raise FileNotFoundError( + f"unknown LoRA '{spec_id}': not a local adapter, catalog entry, or HF repo id" + ) + + +def _pick_repo_weight_file(repo_id: str, hf_token: Optional[str]) -> str: + """Pick the single LoRA weight file in an HF repo (prefer safetensors).""" + from huggingface_hub import HfApi + + files = HfApi(token = hf_token).list_repo_files(repo_id) + safes = [f for f in files if f.lower().endswith(".safetensors") and "/" not in f] + if len(safes) == 1: + return safes[0] + # Prefer a filename hinting at a lora, else the first safetensors, else a gguf. + for f in safes: + if "lora" in f.lower(): + return f + if safes: + return safes[0] + ggufs = [f for f in files if f.lower().endswith(".gguf") and "/" not in f] + if ggufs: + return ggufs[0] + raise FileNotFoundError(f"no .safetensors/.gguf LoRA file found in '{repo_id}'") + + +def resolve_specs( + specs: list[tuple[str, float]], + *, + hf_token: Optional[str] = None, + cancel_event: Optional[threading.Event] = None, +) -> list[ResolvedLora]: + """Resolve request (id, weight) pairs, dropping zero-weight entries. + + A stale / unknown id raises FileNotFoundError inside resolve_one; convert it to + ValueError so the route (which maps only ValueError to a 400) reports bad client + input instead of a generic 500. A Hub download can also raise + ``RuntimeError("Cancelled")`` when the user unloads / starts a superseding load + mid-download; convert that to the diffusion cancellation sentinel so the route + maps it to a 409 instead of a generic server error toast.""" + out: list[ResolvedLora] = [] + try: + for spec_id, weight in specs: + if weight == 0: + continue + out.append(resolve_one(spec_id, weight, hf_token = hf_token, cancel_event = cancel_event)) + except FileNotFoundError as exc: + raise ValueError(str(exc)) from exc + except RuntimeError as exc: + if str(exc) == "Cancelled": + raise RuntimeError(DIFFUSION_CANCELLED_MSG) from exc + raise + return out + + +def materialize_native_dir(resolved: list[ResolvedLora], dest: Path) -> list[ResolvedLora]: + """Populate ``dest`` with symlinks (copy fallback) to the resolved LoRA files. + + sd-cli scans ``--lora-model-dir`` and resolves ```` against filenames in + it, so each selected adapter needs a uniquely-named file there. We use a dedicated + managed directory (never a broad cache dir), which keeps the directory scan small and + safe. Returns the resolved list with aliases updated to the (collision-broken) stems + actually written, so the caller injects matching prompt tags. + """ + dest.mkdir(parents = True, exist_ok = True) + used: set[str] = set() + out: list[ResolvedLora] = [] + for r in resolved: + alias = r.alias + n = 1 + while alias in used: + n += 1 + alias = f"{r.alias}_{n}" + used.add(alias) + ext = os.path.splitext(r.path)[1].lower() or ( + ".gguf" if r.fmt == "gguf" else ".safetensors" + ) + link = dest / f"{alias}{ext}" + try: + if link.exists() or link.is_symlink(): + link.unlink() + os.symlink(os.path.realpath(r.path), link) + except OSError: + import shutil + shutil.copy2(r.path, link) + out.append(ResolvedLora(r.id, alias, str(link), r.fmt, r.weight)) + return out + + +_TAG_RE = re.compile(r"]+):([^>]+)>") + + +def inject_prompt_tags(prompt: str, resolved: list[ResolvedLora]) -> str: + """Append `` tags for the selected adapters, using the backend- + validated weights. + + sd-cli strips these tags before they reach the model, so appending them is safe and + deterministic. A selected adapter's weight is validated (0-2) and recorded in the + request/gallery, so the injected tag must WIN over any `` the user + typed. Strip ALL user-typed tags first: only the selected adapters are materialized in + the managed `--lora-model-dir`, so a tag for an unselected alias can never resolve + anyway (sd-cli's extract_and_remove_lora silently removes unresolved tags), and a tag + for a selected alias must not override the validated weight. Then append the validated + tags for the selected adapters. + """ + # Drop every user-typed tag: unselected ones are dead (not in the managed dir) and + # selected ones must not override the validated weight / 0-2 bounds. + cleaned = _TAG_RE.sub("", prompt) + # Collapse whitespace left by stripped tags without disturbing the user's text. + cleaned = re.sub(r"[ \t]{2,}", " ", cleaned).strip() + tags = [f"" for r in resolved] + if not tags: + return cleaned + sep = "" if not cleaned or cleaned.endswith(" ") else " " + return f"{cleaned}{sep}{' '.join(tags)}" + + +def _fmt_weight(w: float) -> str: + # Stable, compact float formatting (no trailing zeros): 1.0 -> "1", 0.75 -> "0.75". + s = f"{w:.4f}".rstrip("0").rstrip(".") + return s or "0" + + +# Families the native sd-cli LoRA name-conversion supports (SD1.5/SD2/SDXL/SD3/FLUX/ +# z-image). Qwen-Image has no LoRA branch in stable-diffusion.cpp -> excluded until +# validated. Matched by substring against the resolved family name. +_NATIVE_LORA_FAMILY_TOKENS = ( + "flux.1", + "flux.2", + "z-image", + "sd1", + "sd2", + "sdxl", + "sd3", + "stable-diffusion", +) +# Diffusers quant schemes that cannot take LoRA cleanly (torchao tensor-subclass weights). +_DIFFUSERS_LORA_BLOCKED_QUANT = ("int8", "fp8", "nvfp4", "mxfp8") + + +def supports_lora( + *, + engine: Optional[str], + family: Optional[str], + model_kind: Optional[str], + transformer_quant: Optional[str], + compiled: bool = False, +) -> bool: + """Single gate for whether the current load can apply LoRA (used by status + backends). + + Native (sd_cpp): GGUF via sd-cli, for the LoRA-capable families only (Qwen excluded). + Diffusers: bf16 or bnb-4bit transformers, but NOT the dense torchao fp8/int8 fast path + (tensor-subclass weights) and NOT GGUF-via-diffusers. A diffusers transformer that was + torch.compile'd at load (Speed=default/max) also can't take a non-hotswap adapter: + diffusers requires the adapter to be loaded BEFORE compilation, so applying one to the + already-compiled module fails with adapter-key mismatches. ``compiled`` is diffusers-only + (the native sd-cli path has no torch compile). + """ + fam = (family or "").lower() + if engine == "sd_cpp": + return any(tok in fam for tok in _NATIVE_LORA_FAMILY_TOKENS) + # diffusers + if model_kind == "gguf": + return False # GGUF diffusers transformer: use the native engine for LoRA + if transformer_quant and transformer_quant.lower() in _DIFFUSERS_LORA_BLOCKED_QUANT: + return False + if compiled: + return False # can't load an adapter onto an already-compiled transformer + return True diff --git a/studio/backend/core/inference/sd_cpp_args.py b/studio/backend/core/inference/sd_cpp_args.py index ae729a881a..2edd230f26 100644 --- a/studio/backend/core/inference/sd_cpp_args.py +++ b/studio/backend/core/inference/sd_cpp_args.py @@ -412,6 +412,7 @@ def build_img_gen_request( cfg_scale: Optional[float] = None, distilled_guidance: Optional[float] = None, output_format: str = "png", + lora: Optional[list[dict]] = None, ) -> dict: """Build the ``POST /sdcpp/v1/img_gen`` JSON body for one text-to-image request. @@ -453,6 +454,11 @@ def build_img_gen_request( req["seed"] = int(seed) if sample_params: req["sample_params"] = sample_params + # Structured LoRA list: the sdcpp API resolves each ``path`` against the server's + # scanned ``--lora-model-dir`` (prompt-embedded ```` tags are intentionally + # unsupported server-side), so LoRAs must be staged there and named here. + if lora: + req["lora"] = lora return req diff --git a/studio/backend/core/inference/sd_cpp_backend.py b/studio/backend/core/inference/sd_cpp_backend.py index 757a28a3a0..059b85974c 100644 --- a/studio/backend/core/inference/sd_cpp_backend.py +++ b/studio/backend/core/inference/sd_cpp_backend.py @@ -213,6 +213,8 @@ class _SdState: flow_shift: Optional[float] = None server: Optional[SdCppServer] = None mode: str = "server" + # Token kept so LoRA adapters selected at generate time can be fetched from the Hub. + hf_token: Optional[str] = None def _memory_policy(memory_mode: Optional[str], cpu_offload: bool) -> str: @@ -571,6 +573,7 @@ class SdCppDiffusionBackend: flow_shift = fam.sd_cpp_flow_shift, server = server, mode = mode, + hf_token = hf_token, ) with self._lock: if self._load_token != _load_token: @@ -706,11 +709,17 @@ class SdCppDiffusionBackend: upscale: Optional[float] = None, # Reference workflow is GPU/diffusers-only (FLUX.2); accepted for interface parity. reference_images: Optional[list[str]] = None, + # LoRA adapters as (id, weight) pairs; resolved up front, then applied per engine + # path: prompt tags for one-shot sd-cli, structured `lora` entries + # for the resident sd-server. None/empty = no LoRA. + loras: Optional[list[tuple[str, float]]] = None, ) -> dict[str, Any]: import tempfile from PIL import Image + from core.inference import diffusion_lora + if ( init_image is not None or mask_image is not None @@ -748,6 +757,27 @@ class SdCppDiffusionBackend: else: seed = int(seed) cfg_scale, flux_guidance = _map_guidance(state.family, guidance) + # Resolve any selected LoRA adapters up front (downloads land in the HF + # cache; a bad id fails here as a clear 400 before we generate). Drop + # weight-0 rows BEFORE the support gate: weight 0 disables an adapter, so a + # request carrying only disabled rows stays a no-op even on a family where + # native LoRA is unsupported, rather than 400 on a dead selection. + lora_resolved: list = [] + active_loras = [(i, w) for (i, w) in (loras or []) if w != 0] + if active_loras: + if not diffusion_lora.supports_lora( + engine = "sd_cpp", + family = state.family.name, + model_kind = "gguf", + transformer_quant = None, + ): + raise ValueError( + f"LoRA is not supported for {state.family.name} on the native " + "sd.cpp engine." + ) + lora_resolved = diffusion_lora.resolve_specs( + active_loras, hf_token = state.hf_token, cancel_event = cancel + ) self._gen = _SdGen(total_steps = int(steps)) if state.mode == "server" and state.server is not None: images, seeds = self._generate_server( @@ -761,6 +791,7 @@ class SdCppDiffusionBackend: batch_size = batch_size, cfg_scale = cfg_scale, flux_guidance = flux_guidance, + lora_resolved = lora_resolved, cancel = cancel, ) else: @@ -775,6 +806,7 @@ class SdCppDiffusionBackend: batch_size = batch_size, cfg_scale = cfg_scale, flux_guidance = flux_guidance, + lora_resolved = lora_resolved, cancel = cancel, ) if cancel.is_set(): @@ -808,6 +840,7 @@ class SdCppDiffusionBackend: batch_size: int, cfg_scale: Optional[float], flux_guidance: Optional[float], + lora_resolved: list, cancel: threading.Event, ) -> tuple[list, list[int]]: """Generate via the resident sd-server (no model reload). @@ -818,11 +851,20 @@ class SdCppDiffusionBackend: signed-int64 range (the request model / diffusers accept larger seeds), and each chunk is submitted at base+offset so the per-image seeds stay reproducible. Each chunk gets a timeout proportional to its image count so a slow CPU batch is not - cancelled partway through on one fixed deadline.""" + cancelled partway through on one fixed deadline. + + LoRA on the server goes through the structured ``lora`` request field, NOT prompt + tags (the sdcpp API intentionally ignores ```` in the prompt). Selected + adapters are staged into the server's ``--lora-model-dir`` scratch dir, which the + server rescans per request, and referenced by their staged filename.""" import io + import os + import shutil from PIL import Image + from core.inference import diffusion_lora + assert state.server is not None total = max(1, int(batch_size)) # sd.cpp's image seed is signed int64; mask the base (and every derived seed) so a @@ -830,40 +872,60 @@ class SdCppDiffusionBackend: base_seed = int(seed) & ((1 << 63) - 1) images: list = [] seeds: list[int] = [] - for offset in range(0, total, _MAX_SERVER_BATCH): - if cancel.is_set(): - raise SdCppCancelled("sd-server generation was cancelled.") - count = min(_MAX_SERVER_BATCH, total - offset) - chunk_seed = (base_seed + offset) & ((1 << 63) - 1) - payload = build_img_gen_request( - prompt = prompt, - negative_prompt = negative_prompt or None, - width = int(width), - height = int(height), - steps = int(steps), - seed = chunk_seed, - batch_count = count, - sample_method = state.sampling_method, - flow_shift = state.flow_shift, - cfg_scale = cfg_scale, - distilled_guidance = flux_guidance, - ) - blobs = state.server.img_gen( - payload, - on_step = self._on_log, - cancel_event = cancel, - total_timeout = _SERVER_PER_IMAGE_TIMEOUT_S * count, - ) - # All-or-nothing per chunk, like the one-shot path: if the server returns fewer - # blobs than requested (e.g. one image in the batch failed to encode), fail - # rather than silently dropping images from the user's requested batch. - if not cancel.is_set() and len(blobs) != count: - raise RuntimeError( - f"sd-server returned {len(blobs)} of {count} requested images in the batch." + # Stage selected LoRAs into a per-request subdir of the server's lora-model-dir so a + # previous request's adapters can't leak into this one; reference them by the path + # relative to that dir (what the server's recursive scan resolves against). The + # subdir is removed after the batch. supports_lora already gated the family upstream. + lora_payload: Optional[list[dict]] = None + lora_stage: Optional[Path] = None + if lora_resolved: + server_lora_dir = state.server.lora_dir + if server_lora_dir: + lora_stage = Path(server_lora_dir) / f"gen_{os.urandom(6).hex()}" + materialized = diffusion_lora.materialize_native_dir(lora_resolved, lora_stage) + lora_payload = [ + {"path": f"{lora_stage.name}/{Path(m.path).name}", "multiplier": float(m.weight)} + for m in materialized + ] + try: + for offset in range(0, total, _MAX_SERVER_BATCH): + if cancel.is_set(): + raise SdCppCancelled("sd-server generation was cancelled.") + count = min(_MAX_SERVER_BATCH, total - offset) + chunk_seed = (base_seed + offset) & ((1 << 63) - 1) + payload = build_img_gen_request( + prompt = prompt, + negative_prompt = negative_prompt or None, + width = int(width), + height = int(height), + steps = int(steps), + seed = chunk_seed, + batch_count = count, + sample_method = state.sampling_method, + flow_shift = state.flow_shift, + cfg_scale = cfg_scale, + distilled_guidance = flux_guidance, + lora = lora_payload, ) - images.extend(Image.open(io.BytesIO(b)).convert("RGB") for b in blobs) - # sd.cpp advances the seed per image within a job, so report chunk_seed+i. - seeds.extend((chunk_seed + i) & ((1 << 63) - 1) for i in range(len(blobs))) + blobs = state.server.img_gen( + payload, + on_step = self._on_log, + cancel_event = cancel, + total_timeout = _SERVER_PER_IMAGE_TIMEOUT_S * count, + ) + # All-or-nothing per chunk, like the one-shot path: if the server returns fewer + # blobs than requested (e.g. one image in the batch failed to encode), fail + # rather than silently dropping images from the user's requested batch. + if not cancel.is_set() and len(blobs) != count: + raise RuntimeError( + f"sd-server returned {len(blobs)} of {count} requested images in the batch." + ) + images.extend(Image.open(io.BytesIO(b)).convert("RGB") for b in blobs) + # sd.cpp advances the seed per image within a job, so report chunk_seed+i. + seeds.extend((chunk_seed + i) & ((1 << 63) - 1) for i in range(len(blobs))) + finally: + if lora_stage is not None: + shutil.rmtree(lora_stage, ignore_errors = True) return images, seeds def _generate_oneshot( @@ -879,13 +941,21 @@ class SdCppDiffusionBackend: batch_size: int, cfg_scale: Optional[float], flux_guidance: Optional[float], + lora_resolved: list, cancel: threading.Event, ) -> tuple[list, list[int]]: - """Fallback path: re-run one-shot sd-cli per image (reloads the model each time).""" + """Fallback path: re-run one-shot sd-cli per image (reloads the model each time). + + LoRA on the one-shot path uses sd-cli's own mechanism: materialize the selected + adapters into a ``--lora-model-dir`` and inject matching ```` tags + into the prompt (sd-cli parses and strips them). supports_lora already gated the + family upstream, so a non-empty ``lora_resolved`` is safe to apply here.""" import tempfile from PIL import Image + from core.inference import diffusion_lora + engine = self._resolve_engine() extra_args: list[str] = [] if state.vae_format: @@ -896,6 +966,17 @@ class SdCppDiffusionBackend: images = [] seeds: list[int] = [] with tempfile.TemporaryDirectory(prefix = "sdcpp_gen_") as tmpdir: + # Materialize selected LoRAs into a managed dir sd-cli can scan, and inject + # matching tags into the prompt (deduped against any the user + # typed). Empty -> prompt/dir unchanged. + eff_prompt = prompt + lora_dir: Optional[str] = None + if lora_resolved: + materialized = diffusion_lora.materialize_native_dir( + lora_resolved, Path(tmpdir) / "loras" + ) + eff_prompt = diffusion_lora.inject_prompt_tags(prompt, materialized) + lora_dir = str(Path(tmpdir) / "loras") for index in range(max(1, int(batch_size))): if cancel.is_set(): raise RuntimeError(DIFFUSION_CANCELLED_MSG) @@ -906,7 +987,7 @@ class SdCppDiffusionBackend: seed_i = (seed + index) & ((1 << 63) - 1) out_path = str(Path(tmpdir) / f"img_{index}.png") params = SdCppGenParams( - prompt = prompt, + prompt = eff_prompt, negative_prompt = negative_prompt or None, width = int(width), height = int(height), @@ -916,6 +997,8 @@ class SdCppDiffusionBackend: seed = seed_i, sampling_method = state.sampling_method, batch_count = 1, + lora_dir = lora_dir, + lora_apply_mode = "auto" if lora_dir else None, ) engine.generate( state.files, @@ -1024,8 +1107,11 @@ class SdCppDiffusionBackend: "transformer_cache": None, "engine": "sd_cpp", "native_mode": None, + "supports_lora": False, "workflows": [], } + from core.inference import diffusion_lora + return { "loaded": True, "repo_id": state.repo_id, @@ -1048,6 +1134,12 @@ class SdCppDiffusionBackend: "attention_backend": None, "transformer_cache": None, "engine": "sd_cpp", + "supports_lora": diffusion_lora.supports_lora( + engine = "sd_cpp", + family = state.family.name, + model_kind = "gguf", + transformer_quant = None, + ), # "server" = resident sd-server (load once); "oneshot" = legacy per-image sd-cli. "native_mode": state.mode, # The native engine supports plain text-to-image only (generate() rejects diff --git a/studio/backend/core/inference/sd_cpp_server.py b/studio/backend/core/inference/sd_cpp_server.py index 17b269e71d..d66f3cdded 100644 --- a/studio/backend/core/inference/sd_cpp_server.py +++ b/studio/backend/core/inference/sd_cpp_server.py @@ -121,6 +121,14 @@ class SdCppServer: def base_url(self) -> str: return f"http://{self.host}:{self.port}" + @property + def lora_dir(self) -> Optional[str]: + """The server's ``--lora-model-dir`` scratch dir. Adapters staged here are picked + up per request (the sdcpp ``img_gen`` route refreshes its LoRA scan each call), so + server-mode LoRA works by materializing here and referencing the files via the + structured ``lora`` request field (the server ignores ```` prompt tags).""" + return self._scratch_dir + def is_alive(self) -> bool: return self._process is not None and self._process.poll() is None diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 41faf29303..d3feff32f7 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -1813,6 +1813,22 @@ class DiffusionLoadRequest(BaseModel): return value.strip().lower() if isinstance(value, str) else value +class LoraSpec(BaseModel): + """One LoRA adapter to apply for a generation, referenced by its discovery id. + + The id is resolved against the backend's own LoRA catalog + local scan (see + core/inference/diffusion_lora.py); the client never supplies a raw filesystem + path, so an arbitrary file can't be loaded. Weight 0 disables the adapter. + """ + + id: str = Field( + ..., min_length = 1, max_length = 512, description = "LoRA discovery id (repo id or local stem)" + ) + weight: float = Field( + 1.0, ge = 0.0, le = 2.0, description = "Adapter strength; 0 disables, 1.0 is full strength" + ) + + class DiffusionGenerateRequest(BaseModel): """Request to generate one image from the loaded diffusion model.""" @@ -1873,6 +1889,30 @@ class DiffusionGenerateRequest(BaseModel): description = "Additional reference images (base64/data-URL) for the FLUX.2 reference " "workflow, combined with init_image. Up to 3; ignored by other workflows.", ) + loras: Optional[list[LoraSpec]] = Field( + None, + max_length = 8, + description = "LoRA adapters to apply for this generation (by discovery id + weight). " + "Omitted/empty applies none and behaves exactly as before. Rejected with a clear " + "message when the loaded model or its quantisation can't apply LoRA.", + ) + + @field_validator("loras") + @classmethod + def _unique_lora_ids(cls, value: Optional[list[LoraSpec]]) -> Optional[list[LoraSpec]]: + # Both apply paths break alias collisions by suffixing the adapter name/file, so a + # repeated id would load the SAME adapter as several distinct adapters and stack + # its effect past the per-adapter weight bound. The UI already prevents duplicates; + # reject them for API clients too so each adapter takes effect at most once. + if value: + seen: set[str] = set() + for spec in value: + if spec.id in seen: + raise ValueError( + f"duplicate LoRA id '{spec.id}'; list each adapter at most once" + ) + seen.add(spec.id) + return value @field_validator("reference_images") @classmethod @@ -1913,6 +1953,9 @@ class GalleryImage(BaseModel): 1, description = "Batch size used; with batch_index it lets restore replay this image" ) model: Optional[str] = Field(None, description = "Model repo id that produced it") + loras: list[str] = Field( + default_factory = list, description = "LoRA adapters applied, formatted as 'id:weight'" + ) created_at: float = Field(..., description = "Creation time (epoch seconds)") @@ -2002,6 +2045,12 @@ class DiffusionStatusResponse(BaseModel): None, description = "Why diffusers was chosen over the native sd.cpp engine (null when none)", ) + supports_lora: bool = Field( + False, + description = "Whether the loaded model + quantisation can apply LoRA adapters (drives the " + "LoRA picker's enabled state). False on unsupported families/quant (e.g. torchao fp8/int8 " + "dense, GGUF-via-diffusers, or Qwen-Image on the native engine).", + ) # ── OpenAI-compatible images API (POST /v1/images/generations) ── diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 2c0c2bf6b0..929aa19afa 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -11161,6 +11161,7 @@ async def generate_diffusion_image( strength = request.strength, upscale = request.upscale, reference_images = request.reference_images, + loras = [(l.id, l.weight) for l in request.loras] if request.loras else None, ) except ValueError as exc: # Bad client input (undecodable image/mask, or a workflow the loaded family @@ -11220,6 +11221,9 @@ async def generate_diffusion_image( # needs the original batch_size: persist it so restore can replay. "batch_size": request.batch_size, "model": result.get("repo_id"), + "loras": ( + [f"{l.id}:{l.weight:g}" for l in request.loras] if request.loras else [] + ), "created_at": created_at, }, ) diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index cfba8a1181..b22d47c545 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -2031,6 +2031,40 @@ async def scan_loras( ) +@router.get("/diffusion-loras") +async def scan_diffusion_loras( + family: Optional[str] = Query( + default = None, description = "Filter to LoRAs compatible with this diffusion family" + ), + current_subject: str = Depends(get_current_subject), +): + """List diffusion image LoRA adapters for the Images workflow. + + Merges the curated catalog with local files in ``/loras/diffusion``, + optionally filtered to the loaded model's family. Cheap: one directory scan, no network + (a hub adapter is only downloaded when actually selected for a generation). Distinct from + ``/loras`` above, which lists trained/exported TEXT adapters. + """ + from core.inference import diffusion_lora + + entries = diffusion_lora.list_loras(family = family) + return { + "loras": [ + { + "id": e.id, + "display_name": e.display_name, + "source": e.source, + "format": e.fmt, + "families": list(e.families), + "size_bytes": e.size_bytes, + "weight_default": e.weight_default, + } + for e in entries + ], + "loras_dir": str(diffusion_lora.loras_dir()), + } + + def _is_path_under(path: Path, root: Path) -> bool: try: path.resolve().relative_to(root.resolve()) diff --git a/studio/backend/tests/test_diffusion_lora.py b/studio/backend/tests/test_diffusion_lora.py new file mode 100644 index 0000000000..f99739eff9 --- /dev/null +++ b/studio/backend/tests/test_diffusion_lora.py @@ -0,0 +1,375 @@ +"""Tests for diffusion LoRA support: the shared helpers, request-model validation, the +native prompt-tag/dir wiring, and the diffusers set_adapters manager.""" + +from __future__ import annotations + +import os +import types +from pathlib import Path + +import pytest + +from core.inference import diffusion_lora as dl + + +# ── Pure helpers ──────────────────────────────────────────────────────────── + + +def test_sanitize_alias_strips_path_ext_and_unsafe_chars(): + assert dl.sanitize_alias("My Cool/LoRA v2.safetensors") == "LoRA_v2" + assert dl.sanitize_alias("owner/repo-name") == "repo-name" + assert dl.sanitize_alias("weird:<>chars.gguf") == "weird_chars" + assert dl.sanitize_alias("") == "lora" + # Internal dots (version tags like "V1.0") must be replaced: the alias becomes a + # diffusers PEFT adapter name and PEFT rejects "." in module/adapter names. + assert ( + dl.sanitize_alias("Qwen-Image-2512-Lightning-8steps-V1.0-bf16") + == "Qwen-Image-2512-Lightning-8steps-V1_0-bf16" + ) + assert "." not in dl.sanitize_alias("model.v1.0.safetensors") + + +def test_inject_prompt_tags_appends_with_spacing(): + r = dl.ResolvedLora("id", "style", "/p.safetensors", "safetensors", 0.8) + assert dl.inject_prompt_tags("a cat", [r]) == "a cat " + # weight formatting: 1.0 -> "1", trailing zeros trimmed + r1 = dl.ResolvedLora("id", "s", "/p", "safetensors", 1.0) + assert dl.inject_prompt_tags("x", [r1]) == "x " + + +def test_inject_prompt_tags_validated_weight_overrides_user_typed(): + r = dl.ResolvedLora("id", "style", "/p", "safetensors", 0.8) + # A user-typed tag for a SELECTED adapter is replaced by the backend-validated weight + # (so the recorded/validated 0-2 weight wins over whatever was typed), not duplicated. + assert dl.inject_prompt_tags("a cat ", [r]) == "a cat " + + +def test_inject_prompt_tags_strips_unselected_user_tags(): + r = dl.ResolvedLora("id", "style", "/p", "safetensors", 0.8) + # A user tag for an alias that is NOT selected is stripped: only selected adapters are + # materialized in the managed --lora-model-dir, so sd-cli would drop the dead tag anyway; + # removing it keeps the prompt clean and unambiguous. + out = dl.inject_prompt_tags("a cat ", [r]) + assert "" not in out + assert out == "a cat " + + +def test_inject_prompt_tags_empty_returns_prompt(): + assert dl.inject_prompt_tags("hello", []) == "hello" + + +def test_supports_lora_matrix(): + # native: flux/z-image yes, qwen no + assert dl.supports_lora( + engine = "sd_cpp", family = "flux.1", model_kind = "gguf", transformer_quant = None + ) + assert dl.supports_lora( + engine = "sd_cpp", family = "z-image", model_kind = "gguf", transformer_quant = None + ) + assert not dl.supports_lora( + engine = "sd_cpp", family = "qwen-image", model_kind = "gguf", transformer_quant = None + ) + # diffusers: bf16 yes, fp8/int8 dense no, gguf-diffusers no + assert dl.supports_lora( + engine = "diffusers", family = "flux.1", model_kind = "pipeline", transformer_quant = None + ) + assert dl.supports_lora( + engine = "diffusers", family = "flux.1", model_kind = "single_file", transformer_quant = None + ) + assert not dl.supports_lora( + engine = "diffusers", family = "flux.1", model_kind = "single_file", transformer_quant = "fp8" + ) + assert not dl.supports_lora( + engine = "diffusers", family = "flux.1", model_kind = "single_file", transformer_quant = "int8" + ) + assert not dl.supports_lora( + engine = "diffusers", family = "flux.1", model_kind = "gguf", transformer_quant = None + ) + # A torch.compile'd diffusers transformer (Speed=default/max) can't take a non-hotswap + # adapter: diffusers needs the adapter loaded before compilation. + assert not dl.supports_lora( + engine = "diffusers", + family = "flux.1", + model_kind = "pipeline", + transformer_quant = None, + compiled = True, + ) + # compiled is diffusers-only; the native path ignores it. + assert dl.supports_lora( + engine = "sd_cpp", + family = "flux.1", + model_kind = "gguf", + transformer_quant = None, + compiled = True, + ) + + +def test_resolve_specs_maps_cancelled_to_diffusion_sentinel(tmp_path, monkeypatch): + # A Hub download cancelled mid-flight raises RuntimeError("Cancelled"); resolve_specs + # must convert it to the diffusion cancellation sentinel so the route maps it to 409, + # not a generic 500 server-error toast. + def _boom(spec_id, weight, **kw): + raise RuntimeError("Cancelled") + + monkeypatch.setattr(dl, "resolve_one", _boom) + with pytest.raises(RuntimeError) as ei: + dl.resolve_specs([("a", 1.0)]) + assert str(ei.value) == dl.DIFFUSION_CANCELLED_MSG + + # A non-cancellation RuntimeError is left untouched. + def _other(spec_id, weight, **kw): + raise RuntimeError("disk full") + + monkeypatch.setattr(dl, "resolve_one", _other) + with pytest.raises(RuntimeError) as ei2: + dl.resolve_specs([("a", 1.0)]) + assert str(ei2.value) == "disk full" + + +def test_materialize_native_dir_symlinks_and_breaks_collisions(tmp_path): + a = tmp_path / "a.safetensors" + a.write_bytes(b"x") + b = tmp_path / "sub" + b.mkdir() + b2 = b / "a.safetensors" # same stem as `a` -> alias collision + b2.write_bytes(b"y") + resolved = [ + dl.ResolvedLora("a", "a", str(a), "safetensors", 1.0), + dl.ResolvedLora("a2", "a", str(b2), "safetensors", 0.5), + ] + dest = tmp_path / "managed" + out = dl.materialize_native_dir(resolved, dest) + aliases = [r.alias for r in out] + assert aliases == ["a", "a_2"] # collision broken + for r in out: + assert os.path.exists(r.path) + assert Path(r.path).parent == dest + + +def test_list_loras_scans_local(tmp_path, monkeypatch): + d = tmp_path / "loras" + d.mkdir() + (d / "mystyle.safetensors").write_bytes(b"x") + (d / "other.gguf").write_bytes(b"y") + (d / "ignore.txt").write_bytes(b"z") + monkeypatch.setattr(dl, "loras_dir", lambda: d) + ids = {e.id for e in dl.list_loras()} + assert ids == {"mystyle", "other"} + fmts = {e.id: e.fmt for e in dl.list_loras()} + assert fmts["other"] == "gguf" and fmts["mystyle"] == "safetensors" + + +def test_resolve_one_local_and_unknown(tmp_path, monkeypatch): + d = tmp_path / "loras" + d.mkdir() + (d / "mystyle.safetensors").write_bytes(b"x") + monkeypatch.setattr(dl, "loras_dir", lambda: d) + r = dl.resolve_one("mystyle", 0.7) + assert r.path.endswith("mystyle.safetensors") and r.weight == 0.7 + with pytest.raises(FileNotFoundError): + dl.resolve_one("does-not-exist", 1.0) + + +def test_resolve_specs_drops_zero_weight(tmp_path, monkeypatch): + d = tmp_path / "loras" + d.mkdir() + (d / "a.safetensors").write_bytes(b"x") + monkeypatch.setattr(dl, "loras_dir", lambda: d) + out = dl.resolve_specs([("a", 0.0), ("a", 1.0)]) + assert len(out) == 1 and out[0].weight == 1.0 + + +def test_resolve_specs_maps_unknown_id_to_valueerror(tmp_path, monkeypatch): + # An unknown / stale id raises FileNotFoundError in resolve_one; resolve_specs must + # surface it as ValueError so the route returns 400, not a generic 500. + d = tmp_path / "loras" + d.mkdir() + monkeypatch.setattr(dl, "loras_dir", lambda: d) + with pytest.raises(ValueError): + dl.resolve_specs([("nope", 1.0)]) + + +def test_scan_local_disambiguates_identical_stems(tmp_path, monkeypatch): + # foo.safetensors and foo.gguf must get distinct ids so each is addressable; a + # unique stem keeps its clean stem id. + d = tmp_path / "loras" + d.mkdir() + (d / "foo.safetensors").write_bytes(b"x") + (d / "foo.gguf").write_bytes(b"y") + (d / "solo.safetensors").write_bytes(b"z") + monkeypatch.setattr(dl, "loras_dir", lambda: d) + by_id = {e.id: e for e in dl.list_loras()} + assert "foo.safetensors" in by_id and "foo.gguf" in by_id + assert by_id["foo.safetensors"].fmt == "safetensors" + assert by_id["foo.gguf"].fmt == "gguf" + assert "solo" in by_id # unique stem is untouched + + +def test_resolve_one_rejects_traversal_weight_name(tmp_path, monkeypatch): + # A client-supplied weight file with traversal / absolute path is rejected before it + # can reach the downloader (it must stay a plain filename inside the repo). + monkeypatch.setattr(dl, "loras_dir", lambda: tmp_path) + for bad in ("owner/name:../secret.safetensors", "owner/name:/etc/x.safetensors"): + with pytest.raises(ValueError): + dl.resolve_one(bad, 1.0) + + +# ── Request-model validation ──────────────────────────────────────────────── + + +def test_lora_spec_and_request_validation(): + from models.inference import DiffusionGenerateRequest, LoraSpec + + # empty / missing loras -> unchanged behaviour + assert DiffusionGenerateRequest(prompt = "x").loras is None + req = DiffusionGenerateRequest( + prompt = "x", loras = [{"id": "a", "weight": 0.5}, {"id": "b", "weight": 1.0}] + ) + assert [l.id for l in req.loras] == ["a", "b"] + # weight bounds enforced + with pytest.raises(Exception): + LoraSpec(id = "a", weight = 3.0) + with pytest.raises(Exception): + LoraSpec(id = "a", weight = -0.1) + # default weight + assert LoraSpec(id = "a").weight == 1.0 + # duplicate ids are rejected: repeating an id would load the same adapter as several + # distinct suffixed adapters and stack its effect past the per-adapter weight bound. + with pytest.raises(Exception): + DiffusionGenerateRequest( + prompt = "x", loras = [{"id": "a", "weight": 0.5}, {"id": "a", "weight": 1.0}] + ) + + +# ── Diffusers apply manager ───────────────────────────────────────────────── + + +class _FakePipe: + def __init__(self): + self.loaded: list[tuple[str, str]] = [] + self.active = None + self.unloaded = 0 + + def load_lora_weights( + self, + path, + adapter_name = None, + ): + self.loaded.append((path, adapter_name)) + + def set_adapters( + self, + names, + adapter_weights = None, + ): + self.active = (list(names), list(adapter_weights) if adapter_weights else None) + + def unload_lora_weights(self): + self.unloaded += 1 + self.loaded = [] + self.active = None + + +def _fake_state( + pipe, + *, + kind = "pipeline", + quant = None, +): + fam = types.SimpleNamespace(name = "flux.1") + return types.SimpleNamespace( + pipe = pipe, family = fam, kind = kind, transformer_quant = quant, hf_token = None + ) + + +def _backend(): + from core.inference.diffusion import DiffusionBackend + return DiffusionBackend() + + +def test_diffusers_apply_loads_and_sets_adapters(monkeypatch): + import threading + + monkeypatch.setattr( + dl, + "resolve_specs", + lambda specs, **_: [ + dl.ResolvedLora(i, dl.sanitize_alias(i), f"/{i}.safetensors", "safetensors", w) + for i, w in specs + ], + ) + pipe = _FakePipe() + _backend()._apply_loras( + _fake_state(pipe), [("styleA", 0.8), ("styleB", 1.0)], threading.Event() + ) + assert [n for _p, n in pipe.loaded] == ["styleA", "styleB"] + assert pipe.active == (["styleA", "styleB"], [0.8, 1.0]) + assert getattr(pipe, "_unsloth_loras") # marker recorded + + +def test_diffusers_apply_noop_when_unchanged(monkeypatch): + import threading + + monkeypatch.setattr( + dl, + "resolve_specs", + lambda specs, **_: [ + dl.ResolvedLora(i, dl.sanitize_alias(i), f"/{i}.safetensors", "safetensors", w) + for i, w in specs + ], + ) + pipe = _FakePipe() + b = _backend() + b._apply_loras(_fake_state(pipe), [("styleA", 0.8)], threading.Event()) + first_loaded = list(pipe.loaded) + b._apply_loras(_fake_state(pipe), [("styleA", 0.8)], threading.Event()) + assert pipe.loaded == first_loaded # not reloaded + assert pipe.unloaded == 0 + + +def test_diffusers_apply_clears_when_empty(monkeypatch): + import threading + + monkeypatch.setattr( + dl, + "resolve_specs", + lambda specs, **_: [ + dl.ResolvedLora(i, dl.sanitize_alias(i), f"/{i}.safetensors", "safetensors", w) + for i, w in specs + ], + ) + pipe = _FakePipe() + b = _backend() + b._apply_loras(_fake_state(pipe), [("styleA", 0.8)], threading.Event()) + b._apply_loras(_fake_state(pipe), [], threading.Event()) + assert pipe.unloaded == 1 + assert pipe._unsloth_loras == () + + +def test_diffusers_apply_rejects_unsupported_quant(): + import threading + pipe = _FakePipe() + with pytest.raises(ValueError, match = "not supported"): + _backend()._apply_loras( + _fake_state(pipe, kind = "single_file", quant = "fp8"), + [("styleA", 1.0)], + threading.Event(), + ) + + +def test_diffusers_apply_rejects_gguf_adapter(monkeypatch): + # A .gguf adapter (discoverable in the shared catalog) cannot load on the diffusers + # engine; it must be rejected as a clean 400 before touching the pipe. + import threading + + monkeypatch.setattr( + dl, + "resolve_specs", + lambda specs, **_: [ + dl.ResolvedLora(i, dl.sanitize_alias(i), f"/{i}.gguf", "gguf", w) for i, w in specs + ], + ) + pipe = _FakePipe() + with pytest.raises(ValueError, match = "GGUF LoRA"): + _backend()._apply_loras(_fake_state(pipe), [("styleA", 1.0)], threading.Event()) + assert pipe.loaded == [] # never touched the pipe diff --git a/studio/backend/tests/test_sd_cpp_backend.py b/studio/backend/tests/test_sd_cpp_backend.py index fb64f1324b..109724ca0b 100644 --- a/studio/backend/tests/test_sd_cpp_backend.py +++ b/studio/backend/tests/test_sd_cpp_backend.py @@ -92,6 +92,7 @@ class _FakeServer: self.payloads = [] self.timeouts = [] self.alive = True + self.lora_dir = None # set by a test to the server's --lora-model-dir scratch dir def is_alive(self): return self.alive and not self.stopped @@ -668,3 +669,86 @@ def test_run_load_redacts_paths_in_progress_error(monkeypatch): with npl._REDACTION_LOCK: if secret_root in npl._NATIVE_PATH_REDACTIONS: npl._NATIVE_PATH_REDACTIONS.remove(secret_root) + + +# ── LoRA (native engine) ──────────────────────────────────────────────────────── + + +def _fake_materialize(resolved, dest): + """Stand-in for diffusion_lora.materialize_native_dir: write a stub file per adapter + into ``dest`` and return the resolved list pointing at the written paths (mirroring the + real helper's contract without touching the Hub / real weights).""" + from pathlib import Path as _P + + from core.inference import diffusion_lora as dl + + dest.mkdir(parents = True, exist_ok = True) + out = [] + for r in resolved: + p = _P(dest) / f"{r.alias}.safetensors" + p.write_bytes(b"stub") + out.append(dl.ResolvedLora(r.id, r.alias, str(p), r.fmt, r.weight)) + return out + + +def _patch_lora(monkeypatch, resolved, supported = True): + from core.inference import diffusion_lora as dl + + monkeypatch.setattr(dl, "supports_lora", lambda **k: supported) + monkeypatch.setattr(dl, "resolve_specs", lambda specs, **k: list(resolved)) + monkeypatch.setattr(dl, "materialize_native_dir", _fake_materialize) + + +def test_generate_oneshot_applies_loras_via_prompt_tags(monkeypatch): + # One-shot sd-cli LoRA: adapters materialized into a --lora-model-dir and selected with + # tags injected into the prompt (the real inject_prompt_tags runs here). + from core.inference import diffusion_lora as dl + + eng = _FakeEngine() + b = _loaded_backend(engine = eng) # mode = "oneshot" + _patch_lora(monkeypatch, [dl.ResolvedLora("id1", "myalias", "/x/a.safetensors", "safetensors", 0.8)]) + b.generate(prompt = "a fox", steps = 4, seed = 1, loras = [("id1", 0.8)]) + _, params, _, _ = eng.calls[0] + assert params.lora_dir is not None and params.lora_apply_mode == "auto" + assert "" in params.prompt + + +def test_generate_server_stages_loras_and_sends_structured_field(monkeypatch, tmp_path): + # Server-mode LoRA rides the structured `lora` request field (the sdcpp API ignores + # prompt tags): adapters staged into the server's --lora-model-dir, referenced + # by their path relative to it + the validated multiplier. + from pathlib import Path as _P + + from core.inference import diffusion_lora as dl + + b = SdCppDiffusionBackend() + servers: list = [] + _run_server_load(monkeypatch, b, servers) + servers[0].lora_dir = str(tmp_path) + _patch_lora(monkeypatch, [dl.ResolvedLora("id1", "myalias", "/x/a.safetensors", "safetensors", 0.7)]) + b.generate(prompt = "x", steps = 4, seed = 1, batch_size = 1, loras = [("id1", 0.7)]) + payload = servers[0].payloads[0] + assert "lora" in payload and len(payload["lora"]) == 1 + assert payload["lora"][0]["multiplier"] == 0.7 + assert payload["lora"][0]["path"].endswith("myalias.safetensors") + assert " { return parseJson(await authFetch("/api/inference/images/unload", { method: "POST" })); } +/** List diffusion LoRA adapters, optionally filtered to a model family. */ +export async function listDiffusionLoras(family?: string): Promise { + const qs = family ? `?family=${encodeURIComponent(family)}` : ""; + const data = await parseJson<{ loras: DiffusionLoraInfo[] }>( + await authFetch(`/api/models/diffusion-loras${qs}`), + ); + return data.loras ?? []; +} + export interface GalleryPage { images: GalleryImage[]; has_more: boolean; diff --git a/studio/frontend/src/features/images/images-page.tsx b/studio/frontend/src/features/images/images-page.tsx index a563bbd909..58340003c9 100644 --- a/studio/frontend/src/features/images/images-page.tsx +++ b/studio/frontend/src/features/images/images-page.tsx @@ -48,8 +48,10 @@ import { toast } from "@/lib/toast"; import { type DiffusionGenerateProgress, type DiffusionLoadProgress, + type DiffusionLoraInfo, type DiffusionStatus, type GalleryImage, + type LoraSpecInput, deleteGalleryImage, fetchGalleryObjectUrl, generateDiffusionImage, @@ -57,6 +59,7 @@ import { getDiffusionStatus, getGallery, getGenerateProgress, + listDiffusionLoras, loadDiffusionModel, unloadDiffusionModel, } from "./api"; @@ -910,6 +913,10 @@ export function ImagesPage({ active = true }: { active?: boolean }) { // Reference (FLUX.2): up to 3 ADDITIONAL reference images beyond the primary one, combined // by the model (subject + style, character + scene). const [referenceImages, setReferenceImages] = useState([]); + // LoRA adapters selected for the next generation (id + weight), plus the list the picker + // offers. Applied at generate time; available adapters are refreshed per loaded family. + const [loras, setLoras] = useState([]); + const [availableLoras, setAvailableLoras] = useState([]); // Advanced options live in a right-docked panel (like Chat's settings panel). Closed by // default; a single fixed toggle in the top bar opens/closes it (the icon never moves). const [advancedOpen, setAdvancedOpen] = useState(false); @@ -985,6 +992,36 @@ export function ImagesPage({ active = true }: { active?: boolean }) { galleryCache.quant = quant; }, [images, hasMore, selectedId, quant]); + // Refresh the LoRA picker's options when the loaded model (family) changes, and drop any + // selected adapters the new model can't use so a stale/incompatible LoRA is never sent. + const loraCapable = Boolean(status?.loaded && status?.supports_lora); + useEffect(() => { + if (!loraCapable) { + setAvailableLoras([]); + setLoras([]); + return; + } + let cancelled = false; + listDiffusionLoras(status?.family ?? undefined) + .then((list) => { + if (cancelled) return; + setAvailableLoras(list); + const ids = new Set(list.map((l) => l.id)); + setLoras((prev) => prev.filter((s) => ids.has(s.id))); + }) + .catch(() => { + if (cancelled) return; + // Clear the SELECTED adapters too, not just the options: leaving a stale `loras` + // selection in state (with the picker now hidden/empty) would still be posted by + // handleGenerate and could apply adapters from the previous model, or fail. + setAvailableLoras([]); + setLoras([]); + }); + return () => { + cancelled = true; + }; + }, [loraCapable, status?.family]); + const selected = useMemo( () => images.find((i) => i.id === selectedId) ?? images[0] ?? null, [images, selectedId], @@ -1081,6 +1118,19 @@ export function ImagesPage({ active = true }: { active?: boolean }) { const m = matchAspect(image.width, image.height); setAspect(m.key); setPortrait(m.portrait); + // Restore selected LoRA adapters from the recipe ("id:weight" strings); split on the + // LAST colon so an id that itself contains ':' is preserved. Unparseable entries are + // skipped, and a recipe with no LoRAs clears the current selection so the restore + // reproduces the image faithfully rather than leaking a stale form selection. + const restoredLoras: LoraSpecInput[] = []; + for (const entry of image.loras ?? []) { + const idx = entry.lastIndexOf(":"); + if (idx <= 0) continue; + const id = entry.slice(0, idx); + const weight = Number(entry.slice(idx + 1)); + if (id && Number.isFinite(weight)) restoredLoras.push({ id, weight }); + } + setLoras(restoredLoras); toast.success("Settings restored to inputs"); }, []); @@ -1570,6 +1620,8 @@ export function ImagesPage({ active = true }: { active?: boolean }) { strength: condStrength, upscale: condUpscale, reference_images: condRefImages, + // Drop zero-weight rows so the recipe records only adapters that actually applied. + loras: loras.length ? loras.filter((l) => l.weight > 0) : undefined, }); if (!isMounted.current) break; // Prepend this run's records (newest first) and load their blobs. @@ -1587,7 +1639,7 @@ export function ImagesPage({ active = true }: { active?: boolean }) { setGenDone(null); setGenStep(null); } - }, [prompt, negativePrompt, width, height, steps, guidance, seed, batchSize, count, workflow, initImage, maskImage, strength, extendPct, extendSides, upscaleFactor, upscaleStrength, referenceImages, ensureSrc]); + }, [prompt, negativePrompt, width, height, steps, guidance, seed, batchSize, count, workflow, initImage, maskImage, strength, extendPct, extendSides, upscaleFactor, upscaleStrength, referenceImages, loras, ensureSrc]); // Keep the active workflow valid for the loaded model: an edit-only model (Qwen-Image- // Edit) has no Create/Transform tabs, a base model has no Edit tab. Snap to the first @@ -2025,6 +2077,89 @@ export function ImagesPage({ active = true }: { active?: boolean }) { onChange={(e) => setPrompt(e.target.value)} /> + {/* LoRA adapters: shown only when the loaded model + quant can apply them and at + least one adapter is discoverable. Stack multiple, each with a 0-2 weight. The + backend owns how they apply (native prompt tags / diffusers set_adapters); the + UI only sends {id, weight}. */} + {loraCapable && availableLoras.length > 0 && ( + +
+ {loras.map((sel, i) => ( +
+
+ + +
+ + setLoras((prev) => prev.map((p, j) => (j === i ? { ...p, weight: v } : p))) + } + /> +
+ ))} + {loras.length < Math.min(availableLoras.length, 8) && ( + + )} +
+ + )} {/* A negative prompt only does anything with guidance on, so hide it at guidance 0 (Z-Image-Turbo's default) instead of showing a dead field. */} {guidance > 0 && (