ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Support LoRA adapters on torchao int8/fp8 quantized image pipelines

Browse source 
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.
This commit is contained in:
Daniel Han 2026-07-17 09:37:17 +00:00
commit 362cacc448
9 changed files with 394 additions and 52 deletions
Download patch file Download diff file Expand all files Collapse all files

214
studio/backend/core/inference/diffusion.py
View file

@ -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.

23
studio/backend/core/inference/diffusion_auto_policy.py
View file

@ -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
)

30
studio/backend/core/inference/diffusion_lora.py
View file

@ -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

7
studio/backend/core/inference/diffusion_transformer_quant.py
View file

@ -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),

3
studio/backend/core/inference/sd_cpp_backend.py
View file

@ -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.

18
studio/backend/models/inference.py
View file

@ -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,

1
studio/backend/routes/inference.py
View file

@ -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:

144
studio/backend/tests/test_diffusion_backend.py
View file

@ -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)

30
studio/backend/tests/test_diffusion_lora.py
View file

@ -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):