From 362cacc448189ec5c7472cfe81aeb97b5c95e547 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 17 Jul 2026 09:37:17 +0000 Subject: [PATCH] Support LoRA adapters on torchao int8/fp8 quantized image pipelines Adapters are baked at load time: they attach to the dense transformer, then quantize_ converts only the frozen base linears (the lora_ side path is excluded by name), then the loader compiles. Post-quant PEFT injection is not possible on a manually quantized module, so the prequant shortcut is skipped for a baked load and the memory plan is sized for the dense build (force_dense on the quant candidate). At generation time the baked topology is frozen: weight tweaks and disabling (scale 0 reproduces the quantized base exactly) go through set_adapters, while adding or removing adapters returns a clean 400 telling the client to reload with the new selection. supports_lora now returns True for int8/fp8 diffusers loads (checked before the gguf-kind early return, since the quant fast path keeps the picker kind); nvfp4/mxfp8 and GGUF-via-diffusers stay blocked. The load request model takes an optional loras list, threaded through begin_load on both engines (native ignores it and keeps applying LoRA at generation). Verified end to end on GPU: Z-Image GGUF picker + int8 + trained adapter loads through the API, bake marker logged, weight 1.0 vs 0 renders differ visibly, weight 0.5 accepted live, unknown adapter rejected as 400. Affected suites: 296 passed. --- studio/backend/core/inference/diffusion.py | 214 ++++++++++++++---- .../core/inference/diffusion_auto_policy.py | 23 +- .../backend/core/inference/diffusion_lora.py | 30 ++- .../inference/diffusion_transformer_quant.py | 7 +- .../backend/core/inference/sd_cpp_backend.py | 3 + studio/backend/models/inference.py | 18 +- studio/backend/routes/inference.py | 1 + .../backend/tests/test_diffusion_backend.py | 144 ++++++++++++ studio/backend/tests/test_diffusion_lora.py | 30 ++- 9 files changed, 406 insertions(+), 64 deletions(-) diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index c58692f1bb..e6fc778196 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -558,6 +558,7 @@ class DiffusionBackend: requested = mode, base_repo = kwargs.get("base_repo"), prequant_path = kwargs.get("transformer_prequant_path"), + force_dense = bool(kwargs.get("loras")), logger = None, ) # A prequant candidate loads a small checkpoint, not the dense transformer/ shards, @@ -734,6 +735,7 @@ class DiffusionBackend: transformer_cache: Optional[str] = None, transformer_cache_threshold: Optional[float] = None, model_kind: Optional[str] = None, + loras: Optional[list[tuple[str, float]]] = None, ) -> dict[str, Any]: """Validate, then run the (slow) load on a daemon thread. Returns at once.""" # A blank token must mean "anonymous", not an empty credential the Hub 401s. @@ -777,6 +779,7 @@ class DiffusionBackend: transformer_cache = transformer_cache, transformer_cache_threshold = transformer_cache_threshold, model_kind = model_kind, + loras = loras, _load_token = token, ), daemon = True, @@ -1073,6 +1076,10 @@ class DiffusionBackend: transformer_cache: Optional[str] = None, transformer_cache_threshold: Optional[float] = None, model_kind: Optional[str] = None, + # LoRA adapters to BAKE into a torchao int8/fp8 build (attached on the dense + # transformer before quantize_ + compile). Ignored by every other load kind: bf16 / + # bnb loads take adapters at generation time, GGUF-as-is has no dense transformer. + loras: Optional[list[tuple[str, float]]] = None, _load_token: Optional[int] = None, _base_local_dir: Optional[str] = None, ) -> dict[str, Any]: @@ -1205,6 +1212,9 @@ class DiffusionBackend: requested = transformer_quant, base_repo = base, prequant_path = transformer_prequant_path, + # A LoRA bake skips the prequant shortcut, so size the candidate + # for the dense build it will actually run. + force_dense = bool(loras), logger = logger, ) if candidate is not None: @@ -1273,7 +1283,12 @@ class DiffusionBackend: # path must NOT count as prequant here, or it skips the dense-fit re-check # and OOMs materialising the dense transformer after eviction. prequant = ( - usable_prequant_source( + # A LoRA bake skips the prequant shortcut (adapters attach on the + # dense transformer), so the dense-fit re-check must gate the fast + # path exactly as if no prequant source existed. + None + if loras + else usable_prequant_source( fam, scheme, path_override = transformer_prequant_path ) if scheme is not None @@ -1324,6 +1339,7 @@ class DiffusionBackend: base_local_dir = _base_local_dir, prequant_path = transformer_prequant_path, allow_dense_fallback = dense_fallback_allowed, + lora_specs = loras, ) except Exception as exc: # noqa: BLE001 — fall back to the GGUF build logger.warning( @@ -1702,6 +1718,7 @@ class DiffusionBackend: prequant_path: Optional[str] = None, base_local_dir: Optional[str] = None, allow_dense_fallback: bool = True, + lora_specs: Optional[list[tuple[str, float]]] = None, ) -> tuple[Any, str]: """Build the opt-in fast pipeline and return ``(pipe, engaged_scheme)``. @@ -1714,6 +1731,13 @@ class DiffusionBackend: 2. Dense + quantise (fallback): load the DENSE bf16 transformer from the base repo, place it on the device, and torchao-quantise it in place. + ``lora_specs`` bakes LoRA adapters into the build: they attach on the DENSE + transformer (peft's post-quant torchao dispatch needs quantizer metadata a manual + quantize_ never has), then quantize_ converts only the frozen base linears (the + ``lora_`` side path is excluded by name), then the loader compiles. That forces the + dense path -- the prequant shortcut is skipped -- so a baked-LoRA load pays the dense + peak. Verified on the Studio stack: scale 0 reproduces the quantized base exactly. + Raises if the scheme is unsupported or quantisation fails, so ``load_pipeline`` catches it and falls back to the GGUF build. Quantisation runs ON the device and BEFORE the loader compiles the repeated block, so the order stays quantize -> @@ -1725,7 +1749,9 @@ class DiffusionBackend: # nvfp4 off Blackwell) would otherwise materialise the transformer only to fail at # quantize, after eviction. load_pipeline catches this and builds the GGUF pipeline. raise RuntimeError("transformer quant unsupported for this device/scheme") - if fam is not None: + if fam is not None and not lora_specs: + # A LoRA bake needs the DENSE transformer (adapters attach before quantize_), so + # the prequant shortcut is skipped when adapters were requested. source = resolve_prequant_source(fam, scheme, path_override = prequant_path) if source is not None: transformer = load_prequantized_transformer( @@ -1763,6 +1789,29 @@ class DiffusionBackend: pipe = self._assemble_pipe( pipeline_cls, base, transformer, dtype, hf_token, device, base_local_dir, fam = fam ) + if lora_specs: + # Bake the adapters BEFORE quantize_: peft injects its wrappers on the dense + # Linears (the post-quant torchao dispatch would TypeError on a manually + # quantized module), then quantize_ converts only each wrapper's frozen + # base_layer while the "lora_" side path stays high precision. + baked = self._resolve_lora_set( + [(i, w) for (i, w) in lora_specs if w != 0], + family = getattr(fam, "name", None), + hf_token = hf_token, + ) + for name, path, _weight in baked: + pipe.load_lora_weights(path, adapter_name = name) + pipe.set_adapters( + [n for (n, _p, _w) in baked], + adapter_weights = [w for (_n, _p, w) in baked], + ) + pipe._unsloth_loras = baked + pipe._unsloth_loras_baked = True + logger.info( + "diffusion.lora_bake: %d adapter(s) attached before %s quantize", + len(baked), + scheme, + ) scheme = quantize_transformer( pipe, target, @@ -2064,49 +2113,25 @@ class DiffusionBackend: except (StopIteration, AttributeError, RuntimeError, TypeError): 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. + @staticmethod + def _resolve_lora_set( + specs: list[tuple[str, float]], + *, + family: Optional[str], + hf_token: Optional[str], + cancel: Optional[threading.Event] = None, + ) -> tuple[tuple[str, str, float], ...]: + """Resolve (id, weight) specs to a ``(name, path, weight)`` tuple set for diffusers. - 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. + Shared by the generation-time apply path and the quant load-time bake so both produce + IDENTICAL tuples for the same request (the no-op / weight-only comparisons depend on it). """ 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, - family = getattr(state.family, "name", None), - hf_token = state.hf_token, + family = family, + hf_token = hf_token, cancel_event = cancel, ) # diffusers load_lora_weights takes safetensors only; reject a .gguf adapter as a clean 400. @@ -2127,8 +2152,64 @@ class DiffusionBackend: name = f"{r.alias}_{n}" seen.add(name) uniq.append((name, r.path, r.weight)) + return tuple(uniq) - desired = tuple(uniq) + 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. + + A torchao int8/fp8 pipe carries its adapters from the load-time BAKE (attached before + quantize_ + compile). Its module topology is frozen: weight-only changes go through + set_adapters (value-level, compile-guard safe); adding/removing adapters needs a reload + with the new selection, surfaced as a clean 400 here. + """ + 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] + + quant_baked = bool(getattr(pipe, "_unsloth_loras_baked", False)) + quant = (state.transformer_quant or "").lower() + if quant in ("int8", "fp8", "nvfp4", "mxfp8"): + self._adjust_baked_loras(state, pipe, specs, current, quant_baked, cancel) + return + + 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, 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." + ) + + desired = self._resolve_lora_set( + specs, + family = getattr(state.family, "name", None), + hf_token = state.hf_token, + cancel = cancel, + ) + uniq = list(desired) if desired == current: return try: @@ -2148,6 +2229,59 @@ class DiffusionBackend: raise ValueError(f"Failed to apply LoRA: {exc}") from exc pipe._unsloth_loras = desired + def _adjust_baked_loras( + self, + state: Any, + pipe: Any, + specs: list[tuple[str, float]], + current: tuple, + quant_baked: bool, + cancel: threading.Event, + ) -> None: + """Generation-time LoRA handling for a torchao-quantized pipe. + + The adapters (if any) were baked at load time, before quantize_ + compile, so the + module topology is immutable here. Allowed without a reload: weight tweaks on the + baked set and disabling everything (scale 0 reproduces the quantized base exactly; + set_adapters is value-level, so torch.compile guards absorb it). Anything that would + change topology (adding adapters to a bake-less load, or a different adapter set) + raises a clean 400 telling the client to reload with the new selection. + """ + if not quant_baked: + if not specs: + return # no adapters baked, none requested + raise ValueError( + "This quantized (int8/fp8) load was built without LoRA adapters. Reload the " + "model with the adapter selection to bake it into the quantized transformer." + ) + if not specs: + # Disable every baked adapter: scale 0 reproduces the quantized base exactly. + names = [n for (n, _p, _w) in current] + if any(w != 0 for (_n, _p, w) in current): + pipe.set_adapters(names, adapter_weights = [0.0] * len(names)) + pipe._unsloth_loras = tuple((n, p, 0.0) for (n, p, _w) in current) + return + desired = self._resolve_lora_set( + specs, + family = getattr(state.family, "name", None), + hf_token = state.hf_token, + cancel = cancel, + ) + if desired == current: + return + if [(n, p) for (n, p, _w) in desired] == [(n, p) for (n, p, _w) in current]: + # Same adapters, new weights: value-level change on the baked topology. + pipe.set_adapters( + [n for (n, _p, _w) in desired], + adapter_weights = [w for (_n, _p, w) in desired], + ) + pipe._unsloth_loras = desired + return + raise ValueError( + "The LoRA selection changed, but a quantized (int8/fp8) transformer bakes its " + "adapters at load time. Reload the model with the new adapter selection." + ) + @staticmethod def _reset_step_cache(pipe: Any) -> None: """Clear the transformer's stateful step cache (FBCache) before a generation. diff --git a/studio/backend/core/inference/diffusion_auto_policy.py b/studio/backend/core/inference/diffusion_auto_policy.py index c3c6c2ed3f..b7d805c8c6 100644 --- a/studio/backend/core/inference/diffusion_auto_policy.py +++ b/studio/backend/core/inference/diffusion_auto_policy.py @@ -152,6 +152,7 @@ def resolve_dense_quant_candidate( requested: Optional[str], base_repo: Optional[str] = None, prequant_path: Optional[str] = None, + force_dense: bool = False, logger: Optional[logging.Logger] = None, ) -> Optional[DenseQuantEstimate]: """The dense-quant candidate the loader should re-plan memory against, or None. @@ -172,16 +173,20 @@ def resolve_dense_quant_candidate( if scheme is None: return None prequant_available = False - try: - from .diffusion_prequant import usable_prequant_source + # force_dense: the loader will SKIP the prequant shortcut (e.g. a LoRA bake attaches + # adapters on the dense transformer), so the candidate must be sized for the dense build. + if not force_dense: + try: + from .diffusion_prequant import usable_prequant_source - # usable_ (not resolve_): a local path override counts only when the loader will accept it - # (allowlisted AND present), else load_prequantized_transformer refuses it and rebuilds - # dense after the resident pipe is unloaded (the evict-then-OOM this prefetch avoids). - src = usable_prequant_source(fam, scheme, path_override = prequant_path) - prequant_available = src is not None - except Exception: # noqa: BLE001 -- prequant probing must never sink the candidate - prequant_available = False + # usable_ (not resolve_): a local path override counts only when the loader will + # accept it (allowlisted AND present), else load_prequantized_transformer refuses it + # and rebuilds dense after the resident pipe is unloaded (the evict-then-OOM this + # prefetch avoids). + src = usable_prequant_source(fam, scheme, path_override = prequant_path) + prequant_available = src is not None + except Exception: # noqa: BLE001 -- prequant probing must never sink the candidate + prequant_available = False estimate = estimate_dense_quant( fam, scheme, base_repo = base_repo, prequant_available = prequant_available ) diff --git a/studio/backend/core/inference/diffusion_lora.py b/studio/backend/core/inference/diffusion_lora.py index ecaeb688a7..546f0d254e 100644 --- a/studio/backend/core/inference/diffusion_lora.py +++ b/studio/backend/core/inference/diffusion_lora.py @@ -413,8 +413,15 @@ _NATIVE_LORA_FAMILY_TOKENS = ( "sd3", "stable-diffusion", ) -# Diffusers quant schemes that cannot take LoRA cleanly (torchao tensor-subclass weights). -_DIFFUSERS_LORA_BLOCKED_QUANT = ("int8", "fp8", "nvfp4", "mxfp8") +# Diffusers quant schemes whose LoRA path is the load-time BAKE (adapters attach on the +# dense transformer BEFORE torchao quantize_ + compile; peft's post-quant TorchaoLoraLinear +# dispatch needs quantizer metadata a manual quantize_ never has). Verified on the Studio +# stack (peft 0.18.1 / torchao 0.17 / torch 2.10): adapter-first, quantize-base-second is +# clean for both schemes -- scale 0 reproduces the quantized base bit-exactly and the wrapped +# transformer compiles. +_DIFFUSERS_LORA_BAKED_QUANT = ("int8", "fp8") +# Prototype schemes with no validated LoRA path (and no shipped families needing one). +_DIFFUSERS_LORA_BLOCKED_QUANT = ("nvfp4", "mxfp8") def supports_lora( @@ -427,19 +434,26 @@ def supports_lora( ) -> bool: """Single gate for whether the current load can apply LoRA (status + backends). - Native (sd_cpp): GGUF via sd-cli, LoRA-capable families only (Qwen excluded). Diffusers: bf16 - or bnb-4bit, but NOT the dense torchao fp8/int8 path (tensor-subclass weights), NOT - GGUF-via-diffusers, and NOT a torch.compile'd transformer (diffusers requires the adapter - loaded BEFORE compilation). ``compiled`` is diffusers-only. + Native (sd_cpp): GGUF via sd-cli, LoRA-capable families only (Qwen excluded). Diffusers: + bf16 / bnb-4bit apply at generation time (but NOT once the transformer is torch.compile'd: + diffusers needs the adapter loaded before compilation); torchao int8/fp8 apply via the + load-time bake (select adapters when loading; a different selection needs a reload), so + ``compiled`` does not gate them -- the bake precedes compilation by construction. The quant + check runs BEFORE the gguf-kind check because the quant fast path keeps the PICKER kind + ("gguf") while the effective transformer is a dense torchao build. nvfp4/mxfp8 stay + unsupported; GGUF-via-diffusers stays on the native engine for LoRA. """ fam = (family or "").lower() if engine == "sd_cpp": return any(tok in fam for tok in _NATIVE_LORA_FAMILY_TOKENS) # diffusers + quant = (transformer_quant or "").lower() + if quant in _DIFFUSERS_LORA_BAKED_QUANT: + return True # load-time bake; adapters ride inside the compiled quantized build + if quant in _DIFFUSERS_LORA_BLOCKED_QUANT: + return False 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/diffusion_transformer_quant.py b/studio/backend/core/inference/diffusion_transformer_quant.py index e16a4dbff8..964d34b465 100644 --- a/studio/backend/core/inference/diffusion_transformer_quant.py +++ b/studio/backend/core/inference/diffusion_transformer_quant.py @@ -469,7 +469,12 @@ def quantize_transformer( # int8 skips the M=1 projections; scaled_mm schemes have no M limit but fp8/mxfp8 assert # a bf16 weight, so on a mixed-precision DiT (Wan/Hunyuan) they must skip non-bf16 ones or # the pass raises. nvfp4 quantises fp32 fine, so it is not gated (see _REQUIRE_BF16_SCHEMES). - exclude = exclude_tokens_for_scheme(scheme, family) + # "lora_" keeps a baked adapter's side path (lora_A/lora_B/lora_embedding) high + # precision when adapters were attached before this pass; the tiny ranks usually fall + # under min_features anyway, but an explicit token does not depend on the rank. Runtime + # only: NOT part of exclude_tokens_for_scheme, whose list is baked into prequant + # checkpoint metadata (adding it there would reject every existing checkpoint). + exclude = exclude_tokens_for_scheme(scheme, family) + ("lora_",) quantize_( transformer, _make_quant_config(scheme, fast_accum = fast_accum), diff --git a/studio/backend/core/inference/sd_cpp_backend.py b/studio/backend/core/inference/sd_cpp_backend.py index 5d4a5aa4c0..020e63cb5a 100644 --- a/studio/backend/core/inference/sd_cpp_backend.py +++ b/studio/backend/core/inference/sd_cpp_backend.py @@ -378,6 +378,9 @@ class SdCppDiffusionBackend: transformer_cache_threshold: Optional[float] = None, # Accepted for interface parity; native is GGUF-only (router forces diffusers otherwise). model_kind: Optional[str] = None, + # Parity with the diffusers engine's load-time LoRA bake; native applies LoRA at + # generation time through sd-cli, so a load-time selection is ignored here. + loras: Optional[list[tuple[str, float]]] = None, ) -> dict[str, Any]: """Validate, then fetch assets on a daemon thread. Returns at once.""" # Empty/whitespace token = "no token"; "" verbatim breaks the anonymous fallback. diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 86fdcac7a2..a2f4cf7c79 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -1849,6 +1849,18 @@ class DiffusionLoadRequest(BaseModel): "the OS path separator). A bare on/off value such as '1' is deliberately not " "accepted -- it must name an allowed directory.", ) + loras: Optional[list[LoraSpec]] = Field( + None, + max_length = 8, + description = "LoRA adapters to BAKE into a torchao int8/fp8 build: they attach on " + "the dense transformer before quantisation and compilation (the only ordering the " + "quantized fast path supports), so they ride inside the compiled build. Weight " + "changes and disabling apply live at generation time; CHANGING the adapter set " + "needs a reload with the new selection. Ignored by every other load kind (bf16 / " + "bnb-4bit loads take adapters at generation time; GGUF-as-is has no dense " + "transformer). Also forces the dense build path: a baked-LoRA load skips the " + "hosted pre-quantized checkpoint and pays the dense load peak.", + ) attention_backend: Optional[ Literal[ "auto", @@ -2222,8 +2234,10 @@ class DiffusionStatusResponse(BaseModel): 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).", + "LoRA picker's enabled state). torchao int8/fp8 builds support LoRA via the load-time " + "bake: select the adapters when loading; weight changes apply live, a different adapter " + "set needs a reload. False on unsupported families/quant (e.g. nvfp4/mxfp8, " + "GGUF-via-diffusers, or Qwen-Image on the native engine).", ) supports_controlnet: bool = Field( False, diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 7a01cb3c76..b929a07519 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -14359,6 +14359,7 @@ async def load_diffusion_model( transformer_cache = request.transformer_cache, transformer_cache_threshold = request.transformer_cache_threshold, model_kind = kind, + loras = [(s.id, s.weight) for s in request.loras] if request.loras else None, ) if needs_gpu: diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index f483396fc1..4d6891f5b4 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/studio/backend/tests/test_diffusion_backend.py @@ -12,6 +12,7 @@ from __future__ import annotations import contextlib import sys +import threading import types import pytest @@ -2800,6 +2801,148 @@ def test_dense_quant_replan_no_retry_when_capacity_truly_short( assert replan_calls == [True] # genuine capacity shortfall: declined without a retry +class _BakePipe: + def __init__(self): + self.calls: list = [] + + def load_lora_weights(self, path, adapter_name = None): + self.calls.append(("load", path, adapter_name)) + + def set_adapters(self, names, adapter_weights = None): + self.calls.append(("set", tuple(names), tuple(adapter_weights))) + + +def test_dense_quant_lora_bake_attaches_before_quantize(fake_runtime, monkeypatch): + # A LoRA bake must (a) skip the prequant shortcut (adapters need the DENSE transformer), + # (b) attach the adapters BEFORE quantize_transformer (peft's post-quant torchao dispatch + # TypeErrors on a manually quantized module), and (c) mark the pipe as baked. + from core.inference import diffusion as dmod + + backend = DiffusionBackend() + monkeypatch.setattr( + dmod, "select_transformer_quant_scheme", lambda target, mode, family = None: "int8" + ) + prequant_consulted = [] + monkeypatch.setattr( + dmod, + "resolve_prequant_source", + lambda *a, **k: prequant_consulted.append(True) or None, + ) + order: list = [] + + class FakeTransformerCls: + @staticmethod + def from_pretrained(*a, **k): + order.append("dense_load") + return object() + + pipe = _BakePipe() + monkeypatch.setattr( + DiffusionBackend, "_assemble_pipe", staticmethod(lambda *a, **k: pipe) + ) + monkeypatch.setattr( + DiffusionBackend, + "_resolve_lora_set", + staticmethod(lambda specs, **k: (("sloth", "/adapters/sloth.safetensors", 0.8),)), + ) + + def fake_quantize(p, target, **k): + order.append("quantize") + assert any(c[0] == "load" for c in p.calls), "adapters must attach before quantize" + return "int8" + + monkeypatch.setattr(dmod, "quantize_transformer", fake_quantize) + got_pipe, scheme = backend._load_dense_quant_pipeline( + FakeTransformerCls, + object, + "base/repo", + "cuda", + "bf16", + None, + types.SimpleNamespace(device = "cuda", dtype = "bf16"), + "int8", + fam = types.SimpleNamespace(name = "z-image"), + lora_specs = [("sloth", 0.8)], + ) + assert scheme == "int8" + assert prequant_consulted == [] # prequant shortcut skipped for the bake + assert order == ["dense_load", "quantize"] + assert pipe.calls[0] == ("load", "/adapters/sloth.safetensors", "sloth") + assert pipe.calls[1] == ("set", ("sloth",), (0.8,)) + assert pipe._unsloth_loras == (("sloth", "/adapters/sloth.safetensors", 0.8),) + assert pipe._unsloth_loras_baked is True + + +def _quant_lora_state(pipe, quant = "int8"): + return types.SimpleNamespace( + pipe = pipe, + transformer_quant = quant, + kind = "gguf", + family = types.SimpleNamespace(name = "z-image"), + hf_token = None, + speed_optims = ("compiled",), + ) + + +def test_apply_loras_quant_unbaked_requires_reload(monkeypatch): + # A quantized pipe built WITHOUT adapters cannot take one at generation time (topology is + # frozen after quantize_ + compile): clean 400 telling the client to reload. + backend = DiffusionBackend() + pipe = _BakePipe() + with pytest.raises(ValueError, match = "Reload the model with the adapter selection"): + backend._apply_loras( + _quant_lora_state(pipe), [("sloth", 1.0)], threading.Event() + ) + # ...but a no-adapter generation on the same pipe stays a plain no-op. + backend._apply_loras(_quant_lora_state(pipe), [], threading.Event()) + assert pipe.calls == [] + + +def test_apply_loras_quant_baked_matrix(monkeypatch): + # Baked pipe: same set -> no-op; weight-only change -> set_adapters (value-level, no + # topology change); empty -> all scales 0 (reproduces the quantized base); different + # adapter set -> reload error. + backend = DiffusionBackend() + monkeypatch.setattr( + DiffusionBackend, + "_resolve_lora_set", + staticmethod( + lambda specs, **k: tuple( + (i, f"/adapters/{i}.safetensors", w) for (i, w) in specs + ) + ), + ) + + def baked_pipe(): + pipe = _BakePipe() + pipe._unsloth_loras = (("sloth", "/adapters/sloth.safetensors", 0.8),) + pipe._unsloth_loras_baked = True + return pipe + + ev = threading.Event() + # same set: no-op + pipe = baked_pipe() + backend._apply_loras(_quant_lora_state(pipe), [("sloth", 0.8)], ev) + assert pipe.calls == [] + # weight-only change: live set_adapters + marker update + pipe = baked_pipe() + backend._apply_loras(_quant_lora_state(pipe), [("sloth", 1.4)], ev) + assert pipe.calls == [("set", ("sloth",), (1.4,))] + assert pipe._unsloth_loras == (("sloth", "/adapters/sloth.safetensors", 1.4),) + # empty: scale everything to 0 (quantized base output), marker keeps paths + pipe = baked_pipe() + backend._apply_loras(_quant_lora_state(pipe), [], ev) + assert pipe.calls == [("set", ("sloth",), (0.0,))] + assert pipe._unsloth_loras == (("sloth", "/adapters/sloth.safetensors", 0.0),) + # empty again after zeroing: no further calls + backend._apply_loras(_quant_lora_state(pipe), [], ev) + assert len(pipe.calls) == 1 + # different adapter set: topology change -> reload error + pipe = baked_pipe() + with pytest.raises(ValueError, match = "Reload the model with the new adapter selection"): + backend._apply_loras(_quant_lora_state(pipe), [("other", 1.0)], ev) + + def test_assemble_pipe_routes_krea2_per_component(monkeypatch): # krea's repo ships transformers-5.x configs and no top-level tokenizer files, so # Pipeline.from_pretrained dies in the tokenizer (vocab_file = None). The quant fast @@ -2961,6 +3104,7 @@ def test_dense_quant_prefetch_needed_gates(fake_runtime, monkeypatch): requested, base_repo = None, prequant_path = None, + force_dense = False, logger = None, ): seen.append(requested) diff --git a/studio/backend/tests/test_diffusion_lora.py b/studio/backend/tests/test_diffusion_lora.py index 8b400f9b3e..036c54ac64 100644 --- a/studio/backend/tests/test_diffusion_lora.py +++ b/studio/backend/tests/test_diffusion_lora.py @@ -72,19 +72,36 @@ def test_supports_lora_matrix(): 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 + # diffusers: bf16 yes, torchao int8/fp8 yes (load-time bake), nvfp4/mxfp8 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( + assert dl.supports_lora( engine = "diffusers", family = "flux.1", model_kind = "single_file", transformer_quant = "fp8" ) - assert not dl.supports_lora( + assert dl.supports_lora( engine = "diffusers", family = "flux.1", model_kind = "single_file", transformer_quant = "int8" ) + # The quant fast path keeps the PICKER kind ("gguf") while the effective transformer is a + # dense torchao build, so the quant check must decide BEFORE the gguf-kind check; and the + # bake precedes compilation by construction, so compiled does not gate quant builds. + assert dl.supports_lora( + engine = "diffusers", + family = "z-image", + model_kind = "gguf", + transformer_quant = "int8", + compiled = True, + ) + assert not dl.supports_lora( + engine = "diffusers", family = "flux.1", model_kind = "single_file", transformer_quant = "nvfp4" + ) + assert not dl.supports_lora( + engine = "diffusers", family = "flux.1", model_kind = "single_file", transformer_quant = "mxfp8" + ) assert not dl.supports_lora( engine = "diffusers", family = "flux.1", model_kind = "gguf", transformer_quant = None ) @@ -392,14 +409,19 @@ def test_diffusers_apply_clears_when_empty(monkeypatch): def test_diffusers_apply_rejects_unsupported_quant(): + # int8/fp8 pipes bake adapters at load time; a bake-less quant pipe cannot take one at + # generation time (frozen topology) and must direct the client to reload with the + # selection. nvfp4/mxfp8 are never baked, so the same reload error is unreachable there + # via the API (supports_lora blocks the load), but the backend path is shared. import threading pipe = _FakePipe() - with pytest.raises(ValueError, match = "not supported"): + with pytest.raises(ValueError, match = "Reload the model with the adapter selection"): _backend()._apply_loras( _fake_state(pipe, kind = "single_file", quant = "fp8"), [("styleA", 1.0)], threading.Event(), ) + assert pipe.loaded == [] # rejected before touching the pipe def test_diffusers_apply_rejects_gguf_adapter(monkeypatch):