mirror of
https://github.com/unslothai/unsloth.git
synced 2026-10-21 22:12:39 +00:00
Studio Images: per-load engine choice for GGUF picks (native sd.cpp for LoRA) (#14232)
Closes #9295. ## Summary On a GPU host a GGUF image model always loads on diffusers, and a GGUF run as-is there cannot take a LoRA, so Studio hides the LoRA controls. The native sd.cpp engine applies LoRAs to the same GGUF, but a GPU host could only reach it through `UNSLOTH_DIFFUSION_ENGINE`. This adds an engine pick per load. ### Backend - `/images/load` and `/images/download-plan` take `engine`: `auto` (default, unchanged), `diffusers` or `sd_cpp`. It applies to that load only, and `UNSLOTH_DIFFUSION_ENGINE` still wins over it. - A pick the engine cannot serve is refused with a 400 before anything is unloaded, never swapped for the other engine: - native for a non-GGUF kind - native for a family without sd.cpp assets - native with no runnable binary - diffusers without the pipeline class - Switching engines goes through the existing safe unload. - Status adds `engine_forced`, `native_engine_eligible`, `lora_via_native`, and a `resolved.engine` provenance entry (Auto vs explicit, with the reason). ### Frontend - Images Advanced has an Engine select (Auto / Diffusers / Native (sd.cpp)), shown for a loaded GGUF the native engine can run, with the usual provenance badge. - The choice is remembered per checkpoint once a load lands with it, and recalled for that checkpoint's next load. A refused pick is never replayed. - When a GGUF runs as-is on diffusers (no LoRA) and native would apply LoRAs, a note under the prompt offers "Reload with native engine", which reloads the same pick with `engine=sd_cpp`. - The hint covers the tradeoffs: native applies LoRAs and usually needs less VRAM, but has no torch.compile or quantised fast paths and is often slower on CUDA. With precision on Auto a GPU host already loads the full base as dense int8, which bakes LoRAs in at load, so the note only appears when the GGUF runs as-is (precision Off). ## Screenshots FLUX.2 klein 4B GGUF loaded as-is on diffusers. The note under the prompt that offers the native reload:  The Engine select in Advanced, with its provenance badge:  ## Testing - Backend: `test_diffusion_engine_router.py` gains cases for the field, env precedence, refusals before unload, status eligibility and provenance. Router, MLX router, load-contract and routes suites: 579 passed. - Frontend: `tests/image-engine-choice.test.ts` (7 tests); image and diffusion suites 291/291; `tsc --noEmit -p .` and `scripts/typecheck.mjs --tests` clean. - End to end through the real `/images` routes: one CUDA GPU, FLUX.2 klein 4B `Q4_K_M` GGUF, 512x512, 4 steps, seed 7, sd.cpp `master-968` CUDA prebuilt, LoRA `adirik/pixel-art-lora-flux.2-klein-4B`. | Load | engine | supports_lora | lora_via_native | resolved.engine | load | gen (2nd) | LoRA gen | |---|---|---|---|---|---|---|---| | Auto, precision Off (GGUF as-is) | diffusers | false | true | auto | 50.6 s | 0.3 s | 400 (unsupported) | | engine=sd_cpp | sd_cpp | true | false | explicit | 12.0 s | 0.9 s | 200, 1.6 s | | Auto again | diffusers | false | true | auto | 37.2 s | | | | Auto, precision Auto (dense int8) | diffusers | true | false | auto | 63.4 s | 0.7 s | 400 (bake at load) | | engine=sd_cpp | sd_cpp | true | false | explicit | 41.5 s | 2.1 s | 200, 3.3 s | - The native LoRA applies: mean absolute pixel difference vs the same seed without it is 28.86 (0-255), and the next generation without the LoRA matches the no-LoRA image exactly. - Refusals: native for a pipeline pick returns 400 "The native sd.cpp engine loads GGUF models only."; an unknown engine returns 422. - Timings are single runs on a shared GPU, not a speed comparison. Not covered: the CPU sd.cpp bundle on a GPU host (one managed sd.cpp install per host, so a CPU pick would replace the CUDA build), and Mac / Windows / ROCm runs. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Daniel Han
GitHub
parent
f5e91f1ca5
commit
6ccf0e4e74
12 changed files with 1404 additions and 40 deletions
|
|
@ -11177,8 +11177,8 @@ class DiffusionBackend:
|
|||
"(GGUF-via-diffusers, or a torch.compile'd Speed=default/max load). Reload with "
|
||||
"transformer_quant int8 or fp8, which rebuilds the GGUF into a LoRA-capable dense "
|
||||
"transformer, or use a bf16 / bnb-4bit load at Speed=off/eager. To keep the GGUF "
|
||||
"weights themselves, run the native engine, which a GPU host selects only when "
|
||||
"UNSLOTH_DIFFUSION_ENGINE=sd_cpp is set."
|
||||
"weights themselves, reload it on the native engine (Advanced > Engine > Native, "
|
||||
"or engine=sd_cpp on the load), which a GPU host does not select on its own."
|
||||
)
|
||||
|
||||
desired = self._resolve_lora_set(
|
||||
|
|
|
|||
|
|
@ -11,8 +11,11 @@ and remembered, so ``generate`` / ``unload`` / ``status`` / progress all act on
|
|||
Built on the pure ``select_diffusion_engine`` decision; this module adds the policy (env opt-out,
|
||||
MPS gating, per-family native-asset support, lazy binary availability) and records fallbacks.
|
||||
|
||||
Per load, ``/images/load``'s ``engine`` field (auto | diffusers | sd_cpp) asks for an engine; the env override below
|
||||
wins over it. An explicit pick the engine cannot serve is refused, never silently swapped for the other engine.
|
||||
|
||||
Env knobs:
|
||||
UNSLOTH_DIFFUSION_ENGINE=auto|diffusers|sd_cpp force an engine (auto = decide)
|
||||
UNSLOTH_DIFFUSION_ENGINE=auto|diffusers|sd_cpp force an engine (auto = decide); wins over a per-load pick
|
||||
UNSLOTH_DIFFUSION_SD_CPP=auto|0|1 enable/disable the native route
|
||||
UNSLOTH_DIFFUSION_SD_CPP_MPS=0|1 allow native on Apple MPS (default off)
|
||||
UNSLOTH_DIFFUSION_SD_CPP_INSTALL=auto|0|1 allow lazy binary install (in sd_cpp_backend)
|
||||
|
|
@ -362,6 +365,8 @@ before_engine_switch: ContextVar[Optional[Callable[[], None]]] = ContextVar(
|
|||
"before_engine_switch", default = None
|
||||
)
|
||||
_mlx_fallback_reason: Optional[str] = None
|
||||
# The explicit engine the active engine's last load asked for (diffusers | sd_cpp), None when it was left to auto.
|
||||
_engine_pick: Optional[str] = None
|
||||
|
||||
_MLX_MIN_MACOS = (14, 0)
|
||||
|
||||
|
|
@ -428,14 +433,17 @@ def mlx_selection(
|
|||
model_kind: Optional[str],
|
||||
gguf_filename: Optional[str],
|
||||
backend: Optional[str] = None,
|
||||
engine: Optional[str] = None,
|
||||
) -> tuple[bool, Optional[str]]:
|
||||
"""``(use_mlx, refusal)``. Shared by selection and ``predict_engine`` so the download plan and the load
|
||||
cannot disagree about MLX."""
|
||||
cannot disagree about MLX. ``engine`` is the load's own pick (``/images/load``'s ``engine`` field)."""
|
||||
if not _mlx_requested():
|
||||
return False, None
|
||||
forced = os.environ.get("UNSLOTH_DIFFUSION_ENGINE", "auto").strip().lower()
|
||||
if forced in (ENGINE_DIFFUSERS, ENGINE_SD_CPP):
|
||||
return False, f"UNSLOTH_DIFFUSION_ENGINE={forced} overrides UNSLOTH_DIFFUSION_MLX"
|
||||
if engine in (ENGINE_DIFFUSERS, ENGINE_SD_CPP):
|
||||
return False, f"the load asked for the {engine} engine"
|
||||
if model_kind and model_kind != "gguf":
|
||||
return False, f"the MLX engine loads GGUF picks only (got {model_kind})"
|
||||
if fam is None or not family_mlx_supported(fam):
|
||||
|
|
@ -467,8 +475,31 @@ def _note_mlx_refusal(reason: Optional[str]) -> None:
|
|||
logger.warning("UNSLOTH_DIFFUSION_MLX=1 but the MLX engine is not used: %s", reason)
|
||||
|
||||
|
||||
def _engine_config() -> tuple[str, str, bool]:
|
||||
# Values of ``/images/load``'s ``engine`` field. "auto" (or absent) keeps the host's own choice.
|
||||
ENGINE_CHOICES: tuple[str, ...] = ("auto", ENGINE_DIFFUSERS, ENGINE_SD_CPP)
|
||||
|
||||
|
||||
def env_forced_engine() -> Optional[str]:
|
||||
"""The engine UNSLOTH_DIFFUSION_ENGINE forces, else None."""
|
||||
forced = os.environ.get("UNSLOTH_DIFFUSION_ENGINE", "auto").strip().lower()
|
||||
return forced if forced in (ENGINE_DIFFUSERS, ENGINE_SD_CPP) else None
|
||||
|
||||
|
||||
def _requested_engine(engine: Optional[str]) -> Optional[str]:
|
||||
"""The load's own pick when it names an engine and the env does not force one."""
|
||||
pick = (engine or "auto").strip().lower()
|
||||
if pick not in ENGINE_CHOICES:
|
||||
raise ValueError(f"engine must be one of {', '.join(ENGINE_CHOICES)} (got {engine!r}).")
|
||||
if pick == "auto" or env_forced_engine() is not None:
|
||||
return None
|
||||
return pick
|
||||
|
||||
|
||||
def _engine_config(engine: Optional[str] = None) -> tuple[str, str, bool]:
|
||||
forced = os.environ.get("UNSLOTH_DIFFUSION_ENGINE", "auto").strip().lower()
|
||||
requested = _requested_engine(engine)
|
||||
if requested is not None:
|
||||
forced = requested
|
||||
sd_cpp = os.environ.get("UNSLOTH_DIFFUSION_SD_CPP", "auto").strip().lower()
|
||||
mps = os.environ.get("UNSLOTH_DIFFUSION_SD_CPP_MPS", "0").strip().lower() in _ENABLE_TOKENS
|
||||
return forced, sd_cpp, mps
|
||||
|
|
@ -618,6 +649,53 @@ def _selected_card(gpu_ordinal) -> Optional[str]:
|
|||
return None
|
||||
|
||||
|
||||
# A started load's pick, committed to ``_engine_pick`` only once that load's state is published: (pick, engine, the
|
||||
# engine's state when the load started).
|
||||
_pending_engine_pick: Optional[tuple[Optional[str], Any, Any]] = None
|
||||
|
||||
|
||||
def _settle_engine_pick_locked() -> None:
|
||||
"""Commit the pending pick once its engine holds a NEW loaded state. A background load that fails before it
|
||||
publishes keeps the old state resident, and with it the old pick; one that tore the old state down leaves nothing
|
||||
loaded to label, and a later load replaces the pending entry."""
|
||||
global _engine_pick, _pending_engine_pick
|
||||
if _pending_engine_pick is None:
|
||||
return
|
||||
pick, engine, before = _pending_engine_pick
|
||||
current = loaded_state_of(engine)
|
||||
if current is not None and current is not before:
|
||||
_engine_pick = pick
|
||||
_pending_engine_pick = None
|
||||
|
||||
|
||||
_NO_SNAPSHOT = object()
|
||||
|
||||
|
||||
def loaded_state_of(engine: Any) -> Any:
|
||||
"""The engine's published load state (identity only), for ``note_engine_pick``'s before-load snapshot."""
|
||||
return getattr(engine, "_state", None)
|
||||
|
||||
|
||||
def note_engine_pick(engine: Optional[str], loaded_by: Any = None, before: Any = _NO_SNAPSHOT) -> None:
|
||||
"""Record the load's engine pick for the status provenance. The load route calls it once ``loaded_by`` (the
|
||||
engine) has accepted the load, inside ``begin_load_on``'s transition lock, so a pick refused by admission,
|
||||
preflight or the arbiter never relabels the resident model and concurrent loads register in order. With an engine
|
||||
the pick is committed only when that load lands (``_settle_engine_pick_locked``); without one, at once.
|
||||
``before`` is ``loaded_state_of(loaded_by)`` taken before the load's worker started: a cached load can publish
|
||||
before ``begin_load`` returns, and a snapshot taken after it would wait for a state that never comes."""
|
||||
global _engine_pick, _pending_engine_pick
|
||||
pick = _requested_engine(engine)
|
||||
with _lock:
|
||||
_settle_engine_pick_locked()
|
||||
if loaded_by is None:
|
||||
_engine_pick = pick
|
||||
_pending_engine_pick = None
|
||||
else:
|
||||
prior = loaded_state_of(loaded_by) if before is _NO_SNAPSHOT else before
|
||||
_pending_engine_pick = (pick, loaded_by, prior)
|
||||
_settle_engine_pick_locked()
|
||||
|
||||
|
||||
def select_and_activate_engine(
|
||||
fam: DiffusionFamily,
|
||||
*,
|
||||
|
|
@ -626,8 +704,12 @@ def select_and_activate_engine(
|
|||
gpu_ordinal: Optional[int] = None,
|
||||
before_fallback: Optional[Callable[[], None]] = None,
|
||||
gguf_filename: Optional[str] = None,
|
||||
engine: Optional[str] = None,
|
||||
) -> Any:
|
||||
"""Pick + activate the engine for loading ``fam`` on this host; return the engine.
|
||||
|
||||
``engine`` is the load's own pick (auto | diffusers | sd_cpp). An explicit one the host or the pick cannot honour
|
||||
raises ValueError BEFORE anything is unloaded (``refuse_engine_pick``); UNSLOTH_DIFFUSION_ENGINE wins over it.
|
||||
``before_fallback`` may raise to refuse a diffusers fallback before it unloads the resident model.
|
||||
|
||||
Falls back to diffusers (recording a reason) when the native route is disabled, the device has
|
||||
|
|
@ -637,17 +719,26 @@ def select_and_activate_engine(
|
|||
With UNSLOTH_DIFFUSION_MLX=1 the MLX engine is tried first (``mlx_selection``, which needs
|
||||
``gguf_filename`` to judge the quant); a refusal is recorded and the choice below runs unchanged.
|
||||
"""
|
||||
use_mlx, mlx_refusal = mlx_selection(fam, model_kind = model_kind, gguf_filename = gguf_filename)
|
||||
requested = _requested_engine(engine)
|
||||
refuse_engine_pick(fam, engine = engine, model_kind = model_kind)
|
||||
use_mlx, mlx_refusal = mlx_selection(
|
||||
fam, model_kind = model_kind, gguf_filename = gguf_filename, engine = requested
|
||||
)
|
||||
_note_mlx_refusal(mlx_refusal)
|
||||
if use_mlx:
|
||||
return _activate(ENGINE_MLX, None)
|
||||
if model_kind and model_kind != "gguf":
|
||||
return _activate(ENGINE_DIFFUSERS, f"non-GGUF load ({model_kind}) requires diffusers")
|
||||
|
||||
forced, sd_cpp_pref, mps_enabled = _engine_config()
|
||||
forced, sd_cpp_pref, mps_enabled = _engine_config(engine)
|
||||
|
||||
if forced == ENGINE_DIFFUSERS:
|
||||
return _activate(ENGINE_DIFFUSERS, "forced (UNSLOTH_DIFFUSION_ENGINE=diffusers)")
|
||||
return _activate(
|
||||
ENGINE_DIFFUSERS,
|
||||
"selected for this load"
|
||||
if requested == ENGINE_DIFFUSERS
|
||||
else "forced (UNSLOTH_DIFFUSION_ENGINE=diffusers)",
|
||||
)
|
||||
|
||||
prefer_native = forced == ENGINE_SD_CPP
|
||||
if sd_cpp_pref in _DISABLE_TOKENS and not prefer_native:
|
||||
|
|
@ -775,11 +866,41 @@ def select_and_activate_engine(
|
|||
if off_torch is not None:
|
||||
reason = f"{reason}; UNSLOTH_DIFFUSION_SD_CPP_DEVICE={off_torch.label} not honoured"
|
||||
logger.warning("image load stays on torch's device: %s", reason)
|
||||
if requested == ENGINE_SD_CPP:
|
||||
# The load asked for native: say why it cannot run, rather than load it on diffusers behind its back.
|
||||
raise ValueError(f"The native sd.cpp engine cannot load this model here: {reason}.")
|
||||
if before_fallback is not None:
|
||||
before_fallback()
|
||||
return _activate(ENGINE_DIFFUSERS, reason)
|
||||
|
||||
|
||||
def refuse_engine_pick(
|
||||
fam: Optional[DiffusionFamily], *, engine: Optional[str], model_kind: Optional[str]
|
||||
) -> None:
|
||||
"""ValueError for an explicit engine pick that cannot serve this load, checked without any side effect. The
|
||||
binary-dependent native refusal comes later, from selection itself, still before anything is unloaded."""
|
||||
requested = _requested_engine(engine)
|
||||
if requested == ENGINE_SD_CPP:
|
||||
if model_kind and model_kind != "gguf":
|
||||
raise ValueError("The native sd.cpp engine loads GGUF models only.")
|
||||
# Only the administrator's UNSLOTH_DIFFUSION_ENGINE=sd_cpp overrides the disable, and that pin already made
|
||||
# ``requested`` None above.
|
||||
if _engine_config()[1] in _DISABLE_TOKENS:
|
||||
raise ValueError(
|
||||
"The native sd.cpp engine is disabled on this server (UNSLOTH_DIFFUSION_SD_CPP=0)."
|
||||
)
|
||||
if fam is not None and not family_sd_cpp_supported(fam):
|
||||
raise ValueError(
|
||||
f"The native sd.cpp engine has no support for the {fam.name} model family."
|
||||
)
|
||||
elif requested == ENGINE_DIFFUSERS:
|
||||
if fam is not None and not family_pipeline_available(fam):
|
||||
raise ValueError(
|
||||
f"Diffusers cannot build the {fam.name} model family with the installed version; "
|
||||
"use the Auto or native engine."
|
||||
)
|
||||
|
||||
|
||||
def native_binary_installed(
|
||||
*, gpu_ordinal: Optional[int] = None, fam: Optional[DiffusionFamily] = None
|
||||
) -> bool:
|
||||
|
|
@ -841,8 +962,10 @@ def predict_engine(
|
|||
model_kind: Optional[str] = None,
|
||||
gpu_ordinal: Optional[int] = None,
|
||||
gguf_filename: Optional[str] = None,
|
||||
engine: Optional[str] = None,
|
||||
) -> str:
|
||||
"""The engine a load of ``fam`` would select on this host, WITHOUT any side effect.
|
||||
"""The engine a load of ``fam`` would select on this host, WITHOUT any side effect. ``engine`` is the load's own
|
||||
pick, as in ``select_and_activate_engine``; a pick selection would refuse raises the same ValueError here.
|
||||
|
||||
Same policy as ``select_and_activate_engine`` -- and it has to be, because the download plan
|
||||
is built from it: the two engines need different files (sharded diffusers components vs
|
||||
|
|
@ -856,12 +979,15 @@ def predict_engine(
|
|||
diffusers for the very first native load, when no binary is on disk yet) would mispredict the
|
||||
common case: a fresh CPU host.
|
||||
"""
|
||||
if mlx_selection(fam, model_kind = model_kind, gguf_filename = gguf_filename)[0]:
|
||||
requested = _requested_engine(engine)
|
||||
# The plan must refuse what the load will refuse, rather than plan (and stage) for the other engine.
|
||||
refuse_engine_pick(fam, engine = engine, model_kind = model_kind)
|
||||
if mlx_selection(fam, model_kind = model_kind, gguf_filename = gguf_filename, engine = requested)[0]:
|
||||
return ENGINE_MLX
|
||||
if model_kind and model_kind != "gguf":
|
||||
return ENGINE_DIFFUSERS
|
||||
|
||||
forced, sd_cpp_pref, mps_enabled = _engine_config()
|
||||
forced, sd_cpp_pref, mps_enabled = _engine_config(engine)
|
||||
if forced == ENGINE_DIFFUSERS:
|
||||
return ENGINE_DIFFUSERS
|
||||
prefer_native = forced == ENGINE_SD_CPP
|
||||
|
|
@ -888,9 +1014,15 @@ def predict_engine(
|
|||
or _resident_build_upgrades_on_load(gpu_ordinal)
|
||||
)
|
||||
)
|
||||
return select_diffusion_engine(
|
||||
choice = select_diffusion_engine(
|
||||
backend, native_available = native_available, prefer_native = prefer_native
|
||||
)
|
||||
if requested == ENGINE_SD_CPP and choice != ENGINE_SD_CPP:
|
||||
# Selection raises the same refusal; planning diffusers here would stage GB for a load that cannot run.
|
||||
raise ValueError(
|
||||
"The native sd.cpp engine cannot load this model here: native sd.cpp binary unavailable."
|
||||
)
|
||||
return choice
|
||||
|
||||
|
||||
def family_buildable_here(
|
||||
|
|
@ -952,9 +1084,81 @@ def annotate_status(status: dict[str, Any]) -> dict[str, Any]:
|
|||
from .flux3_prompt_upsampler import family_capabilities
|
||||
|
||||
out.update(family_capabilities(out.get("family") if out.get("loaded") else None))
|
||||
out.update(engine_choice_status(out))
|
||||
return out
|
||||
|
||||
|
||||
def engine_choice_status(status: dict[str, Any]) -> dict[str, Any]:
|
||||
"""What the Images page needs to offer the engine choice for the loaded pick: the env pin (which wins over a pick),
|
||||
whether the native engine can serve it (a GGUF of a family with sd.cpp assets; the binary is checked on reload, not
|
||||
on a status poll), and whether reloading it there would enable LoRA where this engine cannot apply one."""
|
||||
from .diffusion_lora import supports_lora
|
||||
|
||||
fam = None
|
||||
if status.get("loaded") and status.get("model_kind") == "gguf" and status.get("family"):
|
||||
try:
|
||||
from .diffusion_families import detect_family
|
||||
fam = detect_family("", override = status.get("family"))
|
||||
except Exception: # noqa: BLE001 -- an unknown family reads as not eligible
|
||||
fam = None
|
||||
# Separate text-encoder / VAE files load on diffusers only; a native reload of this build is refused.
|
||||
# UNSLOTH_DIFFUSION_SD_CPP=0 refuses a native pick (``refuse_engine_pick``) unless the server pins sd_cpp. The
|
||||
# MPS knob does not apply: an explicit pick is native-eligible on any backend.
|
||||
native_disabled = (
|
||||
_engine_config()[1] in _DISABLE_TOKENS and env_forced_engine() != ENGINE_SD_CPP
|
||||
)
|
||||
native_eligible = (
|
||||
fam is not None
|
||||
and not native_disabled
|
||||
and family_sd_cpp_supported(fam)
|
||||
and not status.get("component_files")
|
||||
)
|
||||
lora_via_native = bool(
|
||||
native_eligible
|
||||
and status.get("engine") != ENGINE_SD_CPP
|
||||
and not status.get("supports_lora")
|
||||
and supports_lora(
|
||||
engine = ENGINE_SD_CPP, family = fam.name, model_kind = "gguf", transformer_quant = None
|
||||
)
|
||||
)
|
||||
out: dict[str, Any] = {
|
||||
"engine_forced": env_forced_engine(),
|
||||
"native_engine_eligible": native_eligible,
|
||||
"lora_via_native": lora_via_native,
|
||||
}
|
||||
if status.get("loaded") and status.get("engine"):
|
||||
# The Advanced Engine select's provenance, in the shape of the other resolved controls.
|
||||
resolved = dict(status.get("resolved") or {})
|
||||
resolved["engine"] = _engine_resolved(status.get("engine"))
|
||||
out["resolved"] = resolved
|
||||
return out
|
||||
|
||||
|
||||
def _engine_resolved(active: str) -> dict[str, Any]:
|
||||
forced = env_forced_engine()
|
||||
with _lock:
|
||||
_settle_engine_pick_locked()
|
||||
committed = _engine_pick
|
||||
pick = None if forced else committed
|
||||
if pick:
|
||||
reason = "selected for this load"
|
||||
elif forced:
|
||||
reason = f"pinned by UNSLOTH_DIFFUSION_ENGINE={forced}"
|
||||
elif active == ENGINE_SD_CPP:
|
||||
reason = "native sd.cpp engine chosen for this host"
|
||||
elif active == ENGINE_MLX:
|
||||
reason = "MLX engine (UNSLOTH_DIFFUSION_MLX=1)"
|
||||
else:
|
||||
reason = _fallback_reason or "diffusers selected"
|
||||
return {
|
||||
"value": active,
|
||||
"requested": pick,
|
||||
"source": "explicit" if pick else "auto",
|
||||
"status": "applied",
|
||||
"reason": reason,
|
||||
}
|
||||
|
||||
|
||||
def active_status() -> dict[str, Any]:
|
||||
"""The active engine's status, annotated with which engine + any fallback reason."""
|
||||
return annotate_status(get_active_diffusion_engine().status())
|
||||
|
|
|
|||
|
|
@ -4574,6 +4574,13 @@ class DiffusionLoadRequest(BaseModel):
|
|||
None, description = "Force a family when it can't be inferred from the repo id"
|
||||
)
|
||||
hf_token: Optional[str] = Field(None, description = "HuggingFace token for gated repos")
|
||||
engine: Optional[Literal["auto", "diffusers", "sd_cpp"]] = Field(
|
||||
None,
|
||||
description = "Inference engine for this load: auto (default: diffusers on a GPU, native "
|
||||
"sd.cpp without one), diffusers, or sd_cpp (native; GGUF only, applies LoRA to GGUF "
|
||||
"models). UNSLOTH_DIFFUSION_ENGINE wins over it. A pick the engine cannot serve is "
|
||||
"refused, not swapped.",
|
||||
)
|
||||
cpu_offload: bool = Field(False, description = "Enable model CPU offload to fit low-VRAM cards")
|
||||
memory_mode: Optional[Literal["auto", "fast", "balanced", "low_vram"]] = Field(
|
||||
None,
|
||||
|
|
@ -5592,6 +5599,21 @@ class DiffusionStatusResponse(BaseModel):
|
|||
description = "Why UNSLOTH_DIFFUSION_MLX=1 did not get the MLX engine on the last load (null when "
|
||||
"MLX was not requested or is in use)",
|
||||
)
|
||||
engine_forced: Optional[str] = Field(
|
||||
None,
|
||||
description = "Engine UNSLOTH_DIFFUSION_ENGINE pins (diffusers | sd_cpp); a load's own engine pick "
|
||||
"is ignored while set. Null when unset.",
|
||||
)
|
||||
native_engine_eligible: bool = Field(
|
||||
False,
|
||||
description = "Whether the loaded pick is a GGUF of a family the native sd.cpp engine can run "
|
||||
"(the binary itself is checked when a load asks for it).",
|
||||
)
|
||||
lora_via_native: bool = Field(
|
||||
False,
|
||||
description = "LoRA is unavailable on the current engine for this GGUF, but reloading it on the "
|
||||
"native sd.cpp engine would apply LoRAs.",
|
||||
)
|
||||
supports_lora: bool = Field(
|
||||
False,
|
||||
description = "Whether the loaded model + quantisation can apply LoRA adapters (drives the "
|
||||
|
|
|
|||
|
|
@ -50689,6 +50689,10 @@ async def diffusion_download_plan(
|
|||
base_repo = request.base_repo,
|
||||
**_component_file_kwargs(request),
|
||||
)
|
||||
# The same refusal the load gives an engine pick that cannot serve this model, before anything is staged.
|
||||
from core.inference.diffusion_engine_router import refuse_engine_pick
|
||||
|
||||
refuse_engine_pick(fam, engine = request.engine, model_kind = kind)
|
||||
planner = backend
|
||||
# BEFORE the plan is handed back and staged. The load route refuses a precision this
|
||||
# host cannot honour, but the UI plans and downloads first, so an explicit FP8 on an
|
||||
|
|
@ -50718,6 +50722,7 @@ async def diffusion_download_plan(
|
|||
model_kind = kind,
|
||||
gpu_ordinal = gpu_ordinal,
|
||||
gguf_filename = request.gguf_filename,
|
||||
engine = request.engine,
|
||||
)
|
||||
if fam is not None
|
||||
else None
|
||||
|
|
@ -50733,6 +50738,12 @@ async def diffusion_download_plan(
|
|||
from core.inference.mlx_diffusion_backend import get_mlx_diffusion_backend
|
||||
|
||||
planner = get_mlx_diffusion_backend()
|
||||
# The native refusal is value-only (no GPU), so it runs during training too: only the diffusers probe waits.
|
||||
if fam is not None and planner is not backend:
|
||||
_assert_native_precision_unset(
|
||||
transformer_quant = request.transformer_quant,
|
||||
text_encoder_quant = request.text_encoder_quant,
|
||||
)
|
||||
if fam is not None and not training:
|
||||
if planner is backend:
|
||||
await asyncio.to_thread(
|
||||
|
|
@ -50755,11 +50766,6 @@ async def diffusion_download_plan(
|
|||
text_encoder_files = request.supplied_text_encoder_files(),
|
||||
gguf_filename = request.gguf_filename,
|
||||
)
|
||||
else:
|
||||
_assert_native_precision_unset(
|
||||
transformer_quant = request.transformer_quant,
|
||||
text_encoder_quant = request.text_encoder_quant,
|
||||
)
|
||||
plan = await asyncio.to_thread(
|
||||
planner.download_plan,
|
||||
request.model_path,
|
||||
|
|
@ -50925,6 +50931,8 @@ async def load_diffusion_model_gated(
|
|||
begin_load_on,
|
||||
engine_for,
|
||||
image_install_accelerator,
|
||||
loaded_state_of,
|
||||
note_engine_pick,
|
||||
off_torch_sd_cpp_device,
|
||||
predict_engine,
|
||||
select_and_activate_engine,
|
||||
|
|
@ -50981,6 +50989,10 @@ async def load_diffusion_model_gated(
|
|||
base_repo = request.base_repo,
|
||||
**_component_file_kwargs(request),
|
||||
)
|
||||
# An explicit engine pick that cannot serve this load is refused here, before anything is evicted.
|
||||
from core.inference.diffusion_engine_router import refuse_engine_pick
|
||||
|
||||
refuse_engine_pick(fam, engine = request.engine, model_kind = kind)
|
||||
# Off-torch native shares no VRAM with training or chat; re-settled after selection (diffusers lands on torch's card).
|
||||
off_torch = await asyncio.to_thread(off_torch_sd_cpp_device)
|
||||
# Refuse while training is running: a multi-GB pipeline would compete with the training subprocess for VRAM.
|
||||
|
|
@ -51021,6 +51033,7 @@ async def load_diffusion_model_gated(
|
|||
model_kind = kind,
|
||||
gpu_ordinal = gpu_ordinal,
|
||||
gguf_filename = request.gguf_filename,
|
||||
engine = request.engine,
|
||||
)
|
||||
if fam is not None
|
||||
else None
|
||||
|
|
@ -51163,6 +51176,8 @@ async def load_diffusion_model_gated(
|
|||
),
|
||||
# Only the opt-in MLX engine reads it (to judge the GGUF quant).
|
||||
gguf_filename = request.gguf_filename,
|
||||
# The load's own engine pick (auto when absent); refused there rather than swapped.
|
||||
engine = request.engine,
|
||||
),
|
||||
)
|
||||
# predict_engine is selection's read-only twin: it never installs, so a host whose sd-cli
|
||||
|
|
@ -51239,6 +51254,8 @@ async def load_diffusion_model_gated(
|
|||
from core.inference.flux3_prompt_upsampler import stop_for_handover
|
||||
|
||||
stop_for_handover()
|
||||
# Snapshotted BEFORE the worker starts: a cached load can publish before begin_load returns.
|
||||
prior_state = loaded_state_of(engine)
|
||||
# Kicks the slow load onto a background thread and returns at once (the client polls images/load-progress).
|
||||
started = engine.begin_load(
|
||||
request.model_path,
|
||||
|
|
@ -51271,6 +51288,9 @@ async def load_diffusion_model_gated(
|
|||
# A warm start that took the generation slot between the stop above and begin_load publishing the load would
|
||||
# bring the old model's upsampler up beside this one: cancel it now that the load is visible.
|
||||
_cancel_upsampler_warm_start()
|
||||
# Only once the engine accepted the load, committed when it lands, so a refused or failed pick never
|
||||
# relabels the resident model.
|
||||
note_engine_pick(request.engine, engine, prior_state)
|
||||
return started
|
||||
|
||||
def _begin_load():
|
||||
|
|
|
|||
|
|
@ -102,7 +102,9 @@ def test_the_two_begin_load_signatures_declare_local_files_only_alike():
|
|||
# ── The route, driven onto the native engine ───────────────────────────────
|
||||
|
||||
|
||||
def _drive_the_images_load(monkeypatch, *, user_initiated: bool):
|
||||
def _drive_the_images_load(
|
||||
monkeypatch, *, user_initiated: bool, engine_choice = None, begin_load_error = None, notes = None
|
||||
):
|
||||
"""Run ``POST /images/load``'s body with the NATIVE engine selected; return the mock engine.
|
||||
|
||||
Autospec'd off the real class, so the call the route makes is bound against the real
|
||||
|
|
@ -116,6 +118,22 @@ def _drive_the_images_load(monkeypatch, *, user_initiated: bool):
|
|||
|
||||
engine = create_autospec(SdCppDiffusionBackend, instance = True)
|
||||
engine.begin_load.return_value = {"loaded": False, "repo_id": None}
|
||||
if begin_load_error is not None:
|
||||
engine.begin_load.side_effect = begin_load_error
|
||||
if notes is not None:
|
||||
# The before-load snapshot is taken before begin_load starts the worker.
|
||||
monkeypatch.setattr(
|
||||
router_module,
|
||||
"loaded_state_of",
|
||||
lambda e: notes.append(("snapshot", engine.begin_load.call_count, e is engine)),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
router_module,
|
||||
"note_engine_pick",
|
||||
lambda pick, loaded_by = None, before = None: notes.append(
|
||||
(pick, engine.begin_load.call_count, loaded_by is engine)
|
||||
),
|
||||
)
|
||||
engine.preflight_base_access.return_value = None
|
||||
|
||||
monkeypatch.setattr(router_module, "predict_engine", lambda *a, **k: ENGINE_SD_CPP)
|
||||
|
|
@ -142,6 +160,7 @@ def _drive_the_images_load(monkeypatch, *, user_initiated: bool):
|
|||
DiffusionLoadRequest(
|
||||
model_path = "unsloth/FLUX.1-dev-GGUF",
|
||||
gguf_filename = "flux1-dev-Q4_K_M.gguf",
|
||||
engine = engine_choice,
|
||||
),
|
||||
"test-user",
|
||||
user_initiated = user_initiated,
|
||||
|
|
@ -161,6 +180,23 @@ def test_the_images_page_can_load_on_the_native_engine(monkeypatch, user_initiat
|
|||
assert engine.begin_load.call_args.kwargs["local_files_only"] is (not user_initiated)
|
||||
|
||||
|
||||
def test_the_engine_pick_is_recorded_only_once_the_engine_accepted_the_load(monkeypatch):
|
||||
# Recorded after begin_load, not at selection: a refusal past selection keeps the resident model's provenance.
|
||||
notes = []
|
||||
_drive_the_images_load(monkeypatch, user_initiated = True, engine_choice = "sd_cpp", notes = notes)
|
||||
assert notes == [("snapshot", 0, True), ("sd_cpp", 1, True)]
|
||||
notes.clear()
|
||||
with pytest.raises(Exception):
|
||||
_drive_the_images_load(
|
||||
monkeypatch,
|
||||
user_initiated = True,
|
||||
engine_choice = "sd_cpp",
|
||||
begin_load_error = RuntimeError("busy"),
|
||||
notes = notes,
|
||||
)
|
||||
assert notes == [("snapshot", 0, True)]
|
||||
|
||||
|
||||
# ── The native loader honours it ───────────────────────────────────────────
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -42,11 +42,15 @@ def _clean_env_and_state(monkeypatch):
|
|||
# get_active_diffusion_engine() returning the sd.cpp backend for the rest of the process, 503ing eight tests in a full-suite run.
|
||||
saved_engine = r._active_engine_name
|
||||
saved_reason = r._fallback_reason
|
||||
saved_pick = r._engine_pick
|
||||
saved_pending = r._pending_engine_pick
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
r._active_engine_name = saved_engine
|
||||
r._fallback_reason = saved_reason
|
||||
r._engine_pick = saved_pick
|
||||
r._pending_engine_pick = saved_pending
|
||||
r._supported_family_capabilities.cache_clear()
|
||||
|
||||
|
||||
|
|
@ -489,3 +493,349 @@ def test_the_prediction_counts_the_upgrade_a_load_performs_for_a_moved_pin(monke
|
|||
assert moved == [old]
|
||||
monkeypatch.setenv("UNSLOTH_DIFFUSION_SD_CPP_INSTALL", "0")
|
||||
assert r.predict_engine(fam, model_kind = "gguf") == ENGINE_DIFFUSERS
|
||||
|
||||
|
||||
# ── per-load engine pick (/images/load's ``engine``) ──────────────────────────
|
||||
|
||||
|
||||
def _select_with(fam_name = "z-image", **kwargs):
|
||||
r.select_and_activate_engine(detect_family(fam_name), **kwargs)
|
||||
return r.active_engine_name()
|
||||
|
||||
|
||||
def test_a_load_can_ask_for_native_on_a_gpu_host(monkeypatch):
|
||||
_set_device(monkeypatch, "cuda")
|
||||
_set_binary(monkeypatch, "/usr/bin/sd-cli")
|
||||
_set_runnable(monkeypatch)
|
||||
# Auto keeps the GPU host's own choice; the pick moves only this load.
|
||||
assert _select_with(model_kind = "gguf") == ENGINE_DIFFUSERS
|
||||
assert _select_with(model_kind = "gguf", engine = "auto") == ENGINE_DIFFUSERS
|
||||
assert _select_with(model_kind = "gguf", engine = "sd_cpp") == ENGINE_SD_CPP
|
||||
assert (
|
||||
r.predict_engine(detect_family("z-image"), model_kind = "gguf", engine = "sd_cpp")
|
||||
== ENGINE_SD_CPP
|
||||
)
|
||||
assert _select_with(model_kind = "gguf") == ENGINE_DIFFUSERS
|
||||
|
||||
|
||||
def test_a_load_can_ask_for_diffusers_on_a_cpu_host(monkeypatch):
|
||||
_set_device(monkeypatch, "cpu")
|
||||
_set_binary(monkeypatch, "/usr/bin/sd-cli")
|
||||
_set_runnable(monkeypatch)
|
||||
assert _select_with(model_kind = "gguf") == ENGINE_SD_CPP
|
||||
assert _select_with(model_kind = "gguf", engine = "diffusers") == ENGINE_DIFFUSERS
|
||||
assert r._fallback_reason == "selected for this load"
|
||||
assert (
|
||||
r.predict_engine(detect_family("z-image"), model_kind = "gguf", engine = "diffusers")
|
||||
== ENGINE_DIFFUSERS
|
||||
)
|
||||
|
||||
|
||||
def test_a_native_pick_that_cannot_run_is_refused_before_any_unload(monkeypatch):
|
||||
_set_device(monkeypatch, "cuda")
|
||||
_set_binary(monkeypatch, None)
|
||||
unloads = []
|
||||
monkeypatch.setattr(
|
||||
r,
|
||||
"get_active_diffusion_engine",
|
||||
lambda: SimpleNamespace(status = lambda: {"loaded": True}, unload = lambda: unloads.append(1)),
|
||||
)
|
||||
r._active_engine_name = ENGINE_DIFFUSERS
|
||||
fallbacks = []
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match = "native sd.cpp engine cannot load this model here: native sd.cpp binary unavailable",
|
||||
):
|
||||
_select_with(
|
||||
model_kind = "gguf", engine = "sd_cpp", before_fallback = lambda: fallbacks.append(1)
|
||||
)
|
||||
assert unloads == [] and fallbacks == [] and r.active_engine_name() == ENGINE_DIFFUSERS
|
||||
# Without a pick the old silent fallback stands.
|
||||
assert _select_with(model_kind = "gguf") == ENGINE_DIFFUSERS
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"fam_name, kind, engine, match",
|
||||
[
|
||||
("z-image", "pipeline", "sd_cpp", "GGUF models only"),
|
||||
("lumina-2", "gguf", "sd_cpp", "no support for the lumina-2 model family"),
|
||||
],
|
||||
)
|
||||
def test_native_picks_the_engine_cannot_serve_are_refused(
|
||||
monkeypatch, fam_name, kind, engine, match
|
||||
):
|
||||
_set_device(monkeypatch, "cuda")
|
||||
_set_binary(monkeypatch, "/usr/bin/sd-cli")
|
||||
_set_runnable(monkeypatch)
|
||||
with pytest.raises(ValueError, match = match):
|
||||
r.refuse_engine_pick(detect_family(fam_name), engine = engine, model_kind = kind)
|
||||
with pytest.raises(ValueError, match = match):
|
||||
_select_with(fam_name, model_kind = kind, engine = engine)
|
||||
|
||||
|
||||
def test_a_diffusers_pick_needs_the_pipeline_class(monkeypatch):
|
||||
monkeypatch.setattr(r, "family_pipeline_available", lambda fam: False)
|
||||
with pytest.raises(ValueError, match = "Diffusers cannot build the z-image model family"):
|
||||
r.refuse_engine_pick(detect_family("z-image"), engine = "diffusers", model_kind = "gguf")
|
||||
r.refuse_engine_pick(detect_family("z-image"), engine = "auto", model_kind = "gguf")
|
||||
r.refuse_engine_pick(detect_family("z-image"), engine = None, model_kind = "gguf")
|
||||
|
||||
|
||||
def test_an_unknown_engine_pick_is_refused():
|
||||
with pytest.raises(ValueError, match = "engine must be one of auto, diffusers, sd_cpp"):
|
||||
r.refuse_engine_pick(detect_family("z-image"), engine = "vulkan", model_kind = "gguf")
|
||||
|
||||
|
||||
def test_the_env_override_wins_over_a_load_pick(monkeypatch):
|
||||
_set_device(monkeypatch, "cuda")
|
||||
_set_binary(monkeypatch, "/usr/bin/sd-cli")
|
||||
_set_runnable(monkeypatch)
|
||||
monkeypatch.setenv("UNSLOTH_DIFFUSION_ENGINE", "diffusers")
|
||||
assert _select_with(model_kind = "gguf", engine = "sd_cpp") == ENGINE_DIFFUSERS
|
||||
# Not refused either: the pin decides, and the status says so.
|
||||
r.refuse_engine_pick(detect_family("lumina-2"), engine = "sd_cpp", model_kind = "gguf")
|
||||
assert r.engine_choice_status({"loaded": False})["engine_forced"] == ENGINE_DIFFUSERS
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"status, eligible, lora",
|
||||
[
|
||||
# A FLUX.2 klein GGUF on diffusers: LoRA hidden there, applies natively.
|
||||
(
|
||||
{
|
||||
"loaded": True,
|
||||
"model_kind": "gguf",
|
||||
"family": "flux.2-klein",
|
||||
"engine": "diffusers",
|
||||
"supports_lora": False,
|
||||
},
|
||||
True,
|
||||
True,
|
||||
),
|
||||
# Already native.
|
||||
(
|
||||
{
|
||||
"loaded": True,
|
||||
"model_kind": "gguf",
|
||||
"family": "flux.2-klein",
|
||||
"engine": "sd_cpp",
|
||||
"supports_lora": True,
|
||||
},
|
||||
True,
|
||||
False,
|
||||
),
|
||||
# Qwen-Image runs natively but sd.cpp applies no LoRA to it.
|
||||
(
|
||||
{
|
||||
"loaded": True,
|
||||
"model_kind": "gguf",
|
||||
"family": "qwen-image",
|
||||
"engine": "diffusers",
|
||||
"supports_lora": False,
|
||||
},
|
||||
True,
|
||||
False,
|
||||
),
|
||||
# A pipeline pick is diffusers-only.
|
||||
(
|
||||
{
|
||||
"loaded": True,
|
||||
"model_kind": "pipeline",
|
||||
"family": "flux.2-klein",
|
||||
"engine": "diffusers",
|
||||
"supports_lora": True,
|
||||
},
|
||||
False,
|
||||
False,
|
||||
),
|
||||
# No sd.cpp assets for the family.
|
||||
(
|
||||
{
|
||||
"loaded": True,
|
||||
"model_kind": "gguf",
|
||||
"family": "lumina-2",
|
||||
"engine": "diffusers",
|
||||
"supports_lora": False,
|
||||
},
|
||||
False,
|
||||
False,
|
||||
),
|
||||
({"loaded": False}, False, False),
|
||||
],
|
||||
)
|
||||
def test_engine_choice_status(status, eligible, lora):
|
||||
out = r.engine_choice_status(status)
|
||||
assert {k: out[k] for k in ("engine_forced", "native_engine_eligible", "lora_via_native")} == {
|
||||
"engine_forced": None,
|
||||
"native_engine_eligible": eligible,
|
||||
"lora_via_native": lora,
|
||||
}
|
||||
|
||||
|
||||
def test_the_engine_pick_is_reported_as_resolved_provenance(monkeypatch):
|
||||
_set_device(monkeypatch, "cuda")
|
||||
_set_binary(monkeypatch, "/usr/bin/sd-cli")
|
||||
_set_runnable(monkeypatch)
|
||||
loaded = {
|
||||
"loaded": True,
|
||||
"model_kind": "gguf",
|
||||
"family": "z-image",
|
||||
"resolved": {"speed_mode": {"value": "off"}},
|
||||
}
|
||||
|
||||
def resolved_engine():
|
||||
out = r.engine_choice_status({**loaded, "engine": r.active_engine_name()})
|
||||
assert out["resolved"]["speed_mode"] == {"value": "off"}
|
||||
return out["resolved"]["engine"]
|
||||
|
||||
def load_with(**kwargs):
|
||||
# The route records the pick once the engine accepted the load.
|
||||
_select_with(**kwargs)
|
||||
r.note_engine_pick(kwargs.get("engine"))
|
||||
|
||||
load_with(model_kind = "gguf", engine = "sd_cpp")
|
||||
entry = resolved_engine()
|
||||
assert (
|
||||
entry["value"] == ENGINE_SD_CPP
|
||||
and entry["requested"] == ENGINE_SD_CPP
|
||||
and entry["source"] == "explicit"
|
||||
)
|
||||
load_with(model_kind = "gguf")
|
||||
entry = resolved_engine()
|
||||
assert (
|
||||
entry["value"] == ENGINE_DIFFUSERS
|
||||
and entry["source"] == "auto"
|
||||
and "GPU backend" in entry["reason"]
|
||||
)
|
||||
# A refused pick leaves the resident model's provenance alone.
|
||||
_set_binary(monkeypatch, None)
|
||||
with pytest.raises(ValueError):
|
||||
_select_with(model_kind = "gguf", engine = "sd_cpp")
|
||||
assert resolved_engine()["source"] == "auto"
|
||||
monkeypatch.setenv("UNSLOTH_DIFFUSION_ENGINE", "diffusers")
|
||||
entry = resolved_engine()
|
||||
assert entry["source"] == "auto" and "UNSLOTH_DIFFUSION_ENGINE=diffusers" in entry["reason"]
|
||||
assert "resolved" not in r.engine_choice_status({"loaded": False})
|
||||
|
||||
|
||||
def test_selection_alone_does_not_relabel_the_resident_model(monkeypatch):
|
||||
# A same-engine reload whose admission / preflight / arbiter then fails keeps the old model resident: the
|
||||
# pick is recorded only by note_engine_pick, after the engine accepted the load.
|
||||
_set_device(monkeypatch, "cpu")
|
||||
_set_binary(monkeypatch, "/usr/bin/sd-cli")
|
||||
_set_runnable(monkeypatch)
|
||||
r._engine_pick = None
|
||||
assert _select_with(model_kind = "gguf", engine = "sd_cpp") == ENGINE_SD_CPP
|
||||
assert r._engine_pick is None
|
||||
r.note_engine_pick("sd_cpp")
|
||||
assert r._engine_pick == ENGINE_SD_CPP
|
||||
r.note_engine_pick(None)
|
||||
assert r._engine_pick is None
|
||||
|
||||
|
||||
def test_the_plan_refuses_a_native_pick_the_load_would_refuse(monkeypatch):
|
||||
# No runnable binary and no install allowed: the load raises, so the plan must not stage a diffusers plan.
|
||||
_set_device(monkeypatch, "cuda")
|
||||
_set_binary(monkeypatch, None)
|
||||
monkeypatch.setenv("UNSLOTH_DIFFUSION_SD_CPP_INSTALL", "0")
|
||||
fam = detect_family("z-image")
|
||||
with pytest.raises(ValueError, match = "native sd.cpp engine cannot load this model here"):
|
||||
r.predict_engine(fam, model_kind = "gguf", engine = "sd_cpp")
|
||||
with pytest.raises(ValueError, match = "native sd.cpp engine cannot load this model here"):
|
||||
_select_with(model_kind = "gguf", engine = "sd_cpp")
|
||||
# Auto still plans the fallback it will take.
|
||||
assert r.predict_engine(fam, model_kind = "gguf") == ENGINE_DIFFUSERS
|
||||
assert r.predict_engine(fam, model_kind = "gguf", engine = "auto") == ENGINE_DIFFUSERS
|
||||
|
||||
|
||||
def test_a_native_pick_honours_the_native_disable(monkeypatch):
|
||||
_set_device(monkeypatch, "cpu")
|
||||
_set_binary(monkeypatch, "/usr/bin/sd-cli")
|
||||
_set_runnable(monkeypatch)
|
||||
monkeypatch.setenv("UNSLOTH_DIFFUSION_SD_CPP", "0")
|
||||
fam = detect_family("z-image")
|
||||
with pytest.raises(ValueError, match = "disabled on this server"):
|
||||
r.refuse_engine_pick(fam, engine = "sd_cpp", model_kind = "gguf")
|
||||
with pytest.raises(ValueError, match = "disabled on this server"):
|
||||
r.predict_engine(fam, model_kind = "gguf", engine = "sd_cpp")
|
||||
with pytest.raises(ValueError, match = "disabled on this server"):
|
||||
_select_with(model_kind = "gguf", engine = "sd_cpp")
|
||||
assert _select_with(model_kind = "gguf") == ENGINE_DIFFUSERS
|
||||
# The administrator's pin still overrides the disable.
|
||||
monkeypatch.setenv("UNSLOTH_DIFFUSION_ENGINE", "sd_cpp")
|
||||
r.refuse_engine_pick(fam, engine = "sd_cpp", model_kind = "gguf")
|
||||
assert _select_with(model_kind = "gguf", engine = "sd_cpp") == ENGINE_SD_CPP
|
||||
|
||||
|
||||
def test_component_file_loads_are_not_native_eligible():
|
||||
# Separate text-encoder / VAE files are diffusers-only: the plan and the load refuse them on native.
|
||||
out = r.engine_choice_status(
|
||||
{
|
||||
"loaded": True,
|
||||
"model_kind": "gguf",
|
||||
"family": "flux.2-klein",
|
||||
"engine": "diffusers",
|
||||
"supports_lora": False,
|
||||
"component_files": {"vae": "ae.safetensors"},
|
||||
}
|
||||
)
|
||||
assert out["native_engine_eligible"] is False and out["lora_via_native"] is False
|
||||
|
||||
|
||||
def test_the_engine_pick_commits_only_when_the_background_load_lands():
|
||||
# begin_load only starts the worker; a same-engine reload keeps the old state resident while it fetches, and a
|
||||
# worker failure before publishing keeps it for good. The pick follows the published state, not the start.
|
||||
old_state = object()
|
||||
engine = SimpleNamespace(_state = old_state)
|
||||
r._engine_pick = None
|
||||
r._pending_engine_pick = None
|
||||
r.note_engine_pick("sd_cpp", engine)
|
||||
assert r._engine_resolved(ENGINE_SD_CPP)["source"] == "auto"
|
||||
# Torn down mid-load: nothing loaded to label yet.
|
||||
engine._state = None
|
||||
assert r._engine_resolved(ENGINE_SD_CPP)["source"] == "auto"
|
||||
engine._state = object()
|
||||
entry = r._engine_resolved(ENGINE_SD_CPP)
|
||||
assert entry["source"] == "explicit" and entry["requested"] == ENGINE_SD_CPP
|
||||
# A later reload that fails before publishing leaves the landed pick in place.
|
||||
r.note_engine_pick(None, engine)
|
||||
assert r._engine_resolved(ENGINE_SD_CPP)["requested"] == ENGINE_SD_CPP
|
||||
# A load that landed before the next one started is committed, not lost.
|
||||
r.note_engine_pick("diffusers", engine)
|
||||
engine._state = object()
|
||||
r.note_engine_pick(None, engine)
|
||||
assert r._engine_pick == ENGINE_DIFFUSERS
|
||||
|
||||
|
||||
def test_a_load_that_publishes_before_begin_load_returns_still_commits_its_pick():
|
||||
# begin_load starts the worker before it returns; a cached load can publish in that gap. The route snapshots the
|
||||
# state before starting it, so the published state still reads as new.
|
||||
old_state = object()
|
||||
engine = SimpleNamespace(_state = old_state)
|
||||
r._engine_pick = None
|
||||
r._pending_engine_pick = None
|
||||
before = r.loaded_state_of(engine)
|
||||
engine._state = object() # the worker published before the route noted the pick
|
||||
r.note_engine_pick("sd_cpp", engine, before)
|
||||
assert r._engine_pick == ENGINE_SD_CPP
|
||||
assert r._engine_resolved(ENGINE_SD_CPP)["source"] == "explicit"
|
||||
|
||||
|
||||
def test_the_native_disable_hides_the_native_actions(monkeypatch):
|
||||
# A pick refuse_engine_pick would 400 is never advertised: no Native select, no "Reload with native engine".
|
||||
status = {
|
||||
"loaded": True,
|
||||
"model_kind": "gguf",
|
||||
"family": "flux.2-klein",
|
||||
"engine": "diffusers",
|
||||
"supports_lora": False,
|
||||
}
|
||||
assert r.engine_choice_status(status)["lora_via_native"] is True
|
||||
monkeypatch.setenv("UNSLOTH_DIFFUSION_SD_CPP", "0")
|
||||
out = r.engine_choice_status(status)
|
||||
assert out["native_engine_eligible"] is False and out["lora_via_native"] is False
|
||||
with pytest.raises(ValueError, match = "disabled on this server"):
|
||||
r.refuse_engine_pick(detect_family("flux.2-klein"), engine = "sd_cpp", model_kind = "gguf")
|
||||
# The administrator's sd_cpp pin overrides the disable, as it does for the load.
|
||||
monkeypatch.setenv("UNSLOTH_DIFFUSION_ENGINE", "sd_cpp")
|
||||
assert r.engine_choice_status(status)["native_engine_eligible"] is True
|
||||
|
|
|
|||
|
|
@ -1852,6 +1852,25 @@ def test_the_native_plan_refuses_the_same_request_the_native_load_would(client,
|
|||
assert "native engine" in resp.json()["detail"]
|
||||
|
||||
|
||||
def test_the_native_plan_refuses_an_explicit_precision_while_training_runs(client, monkeypatch):
|
||||
# Value-only, so it is not gated on training like the diffusers probe: the plan must not stage a native load
|
||||
# /images/load refuses once training ends.
|
||||
import core.inference.diffusion_engine_router as engine_router
|
||||
import routes.inference as inference_routes
|
||||
from core.inference.sd_cpp_engine import ENGINE_SD_CPP
|
||||
|
||||
monkeypatch.setattr(inference_routes, "_training_is_active", lambda: True)
|
||||
monkeypatch.setattr(engine_router, "predict_engine", lambda fam, **kw: ENGINE_SD_CPP)
|
||||
resp = _post_download_plan(
|
||||
client,
|
||||
model_path = "unsloth/Z-Image-Turbo-GGUF",
|
||||
gguf_filename = "z-image-turbo-Q4_K_M.gguf",
|
||||
transformer_quant = "fp8",
|
||||
)
|
||||
assert resp.status_code == 409, resp.text
|
||||
assert "native engine" in resp.json()["detail"]
|
||||
|
||||
|
||||
def test_the_plan_still_answers_when_nothing_was_promised(client):
|
||||
resp = _post_download_plan(
|
||||
client,
|
||||
|
|
|
|||
|
|
@ -73,6 +73,15 @@ export interface DiffusionStatus {
|
|||
workflows?: string[];
|
||||
// Absent on an older backend: callers keep the historical limits (4 images, RGB, 16 px, 2048).
|
||||
conditioning?: DiffusionConditioning | null;
|
||||
// Engine that ran the load ("diffusers" | "sd_cpp" | "mlx") and, for diffusers, why it was chosen.
|
||||
engine?: string | null;
|
||||
fallback_reason?: string | null;
|
||||
// Engine UNSLOTH_DIFFUSION_ENGINE pins; a load's own engine pick is ignored while set.
|
||||
engine_forced?: string | null;
|
||||
// A loaded GGUF the native sd.cpp engine can run (drives the Advanced Engine select).
|
||||
native_engine_eligible?: boolean;
|
||||
// LoRA is unavailable on this engine, but a reload on the native engine would apply LoRAs.
|
||||
lora_via_native?: boolean;
|
||||
// Whether the loaded model + quantisation can apply LoRA adapters (drives the LoRA picker enabled state).
|
||||
supports_lora?: boolean;
|
||||
// Whether the loaded model can apply a ControlNet. Diffusers only, for families with a ControlNet pipeline.
|
||||
|
|
@ -161,6 +170,8 @@ export interface DiffusionLoadRequest {
|
|||
base_repo?: string;
|
||||
family_override?: string;
|
||||
hf_token?: string;
|
||||
// Engine for this load (GGUF only); omit for auto. Refused with a 400 when the engine cannot serve the pick.
|
||||
engine?: "auto" | "diffusers" | "sd_cpp";
|
||||
cpu_offload?: boolean;
|
||||
// Separate .safetensors files (e.g. ComfyUI models/text_encoders, models/vae); gguf / single_file only.
|
||||
text_encoder_file?: string | string[];
|
||||
|
|
|
|||
|
|
@ -341,6 +341,26 @@ import {
|
|||
withEnhanceChoice,
|
||||
} from "./lib/prompt-enhance";
|
||||
import { EnhancedPromptRecipeRow, EnhancePromptControls } from "./prompt-enhance-panel";
|
||||
import {
|
||||
ENGINE_HINT,
|
||||
ENGINE_OPTIONS,
|
||||
type EngineChoice,
|
||||
engineModelKey,
|
||||
engineRequestField,
|
||||
engineSelectVisible,
|
||||
loadedEngineChoice,
|
||||
nativeLoadAdvanced,
|
||||
nativeReloadTarget,
|
||||
offersNativeLoraReload,
|
||||
reapplyChangesEngine,
|
||||
reapplyTarget,
|
||||
readEngineChoices,
|
||||
reloadSource,
|
||||
residentReloadTarget,
|
||||
rememberLoadedEngine,
|
||||
sendsEngine,
|
||||
withEngineChoice,
|
||||
} from "./lib/engine-choice";
|
||||
import { useNavigate, useSearch } from "@tanstack/react-router";
|
||||
import { useStagedDownload, type StagedDownloadEntry } from "@/features/hub/download-manager";
|
||||
import { DiffusionTrainPanel } from "./train/diffusion-train-panel";
|
||||
|
|
@ -1415,7 +1435,8 @@ function reportLoadFailure(failure: unknown, fallback: string): void {
|
|||
}
|
||||
|
||||
type Busy = "loading" | "unloading" | "generating" | null;
|
||||
type ImageLoadOptions = { kind: "gguf" | "single_file" | "pipeline"; filename?: string; displayRepoId?: string };
|
||||
// baseRepo: only a reload of a resident model the page did not load sends one, to rebuild the same companions.
|
||||
type ImageLoadOptions = { kind: "gguf" | "single_file" | "pipeline"; filename?: string; displayRepoId?: string; baseRepo?: string };
|
||||
type LastLoad = { repoId: string } & ImageLoadOptions & Pick<RememberedImageModel, "textEncoderFiles" | "vaeFile">;
|
||||
|
||||
// What a pick optimistically replaced, so a load that never takes can put it all back. The
|
||||
|
|
@ -1447,6 +1468,7 @@ type LoadAdvanced = Pick<
|
|||
| "gpu_ids"
|
||||
| "text_encoder_file"
|
||||
| "vae_file"
|
||||
| "engine"
|
||||
>;
|
||||
|
||||
function openImageLabel(t: ReturnType<typeof useT>, prompt: string): string {
|
||||
|
|
@ -1545,6 +1567,8 @@ export function ImagesPage({
|
|||
const [steps, setSteps] = useState(DEFAULT_GEN.steps);
|
||||
const [guidance, setGuidance] = useState(DEFAULT_GEN.guidance);
|
||||
const [samplerChoices, setSamplerChoices] = useState(readSamplerChoices);
|
||||
// Engine pick per GGUF checkpoint, recorded only by a load that landed with it.
|
||||
const [engineChoices, setEngineChoices] = useState(readEngineChoices);
|
||||
const [enhanceChoices, setEnhanceChoices] = useState(readEnhanceChoices);
|
||||
// The family it was set under ("" = before a model loaded, which applies to whichever family loads next).
|
||||
const [enhanceOverrideState, setEnhanceOverrideState] = useState<{ family: string; on: boolean } | null>(null);
|
||||
|
|
@ -1668,6 +1692,10 @@ export function ImagesPage({
|
|||
// while loaded shows "Reapply".
|
||||
const [modelSelectionAction, setModelSelectionAction] = useState<"load" | "download">("load");
|
||||
const [speedMode, setSpeedMode] = useState<"auto" | "off" | "eager" | "default" | "max">("auto");
|
||||
// The Engine select edits the LOADED checkpoint's next load (Reapply); any other pick uses its stored choice.
|
||||
const [engineChoice, setEngineChoice] = useState<EngineChoice>("auto");
|
||||
// The loaded build the Engine select was last seeded from; cleared by a failed load to force a reseed.
|
||||
const [engineSeed, setEngineSeed] = useState("");
|
||||
const [transformerQuant, setTransformerQuant] = useState<
|
||||
"none" | "auto" | "int8" | "fp8" | "nvfp4" | "mxfp8"
|
||||
>("auto");
|
||||
|
|
@ -2832,7 +2860,10 @@ export function ImagesPage({
|
|||
const statusTicket = useRef(0);
|
||||
const setStatusIfNewest = useCallback(
|
||||
(ticket: number, next: DiffusionStatus) => {
|
||||
if (ticket === statusTicket.current) setStatus(next);
|
||||
if (ticket !== statusTicket.current) return;
|
||||
setStatus(next);
|
||||
// Any accepted status, not only the poll's ready branch: a load that finished while the page was closed.
|
||||
setEngineChoices((all) => rememberLoadedEngine(all, next));
|
||||
},
|
||||
[],
|
||||
);
|
||||
|
|
@ -2933,6 +2964,10 @@ export function ImagesPage({
|
|||
return;
|
||||
}
|
||||
setStatusIfNewest(ticket, loaded);
|
||||
if (loaded.loaded && loaded.model_kind === "gguf") {
|
||||
const key = engineModelKey(loaded.repo_id, loaded.gguf_filename);
|
||||
setEngineChoices((all) => withEngineChoice(all, key, loadedEngineChoice(loaded)));
|
||||
}
|
||||
const evicted = pendingEvicted.current;
|
||||
pendingEvicted.current = [];
|
||||
toast.success("Model loaded", evictedToastOptions(evicted));
|
||||
|
|
@ -2980,6 +3015,9 @@ export function ImagesPage({
|
|||
setCanReapply(lastLoadRevert.current.prev != null);
|
||||
lastLoadRevert.current = null;
|
||||
}
|
||||
// And for the Engine select: a worker that failed before publishing leaves the resident provenance, so the
|
||||
// seed would not move and the next Reapply would replay the failed pick.
|
||||
setEngineSeed("");
|
||||
// A failed load may have freed a previously-loaded model, so resync to the real backend state.
|
||||
void refreshStatus();
|
||||
return;
|
||||
|
|
@ -3125,12 +3163,17 @@ export function ImagesPage({
|
|||
const seedKey = `${repoId}\0${residentDefaults}\0${residentSteps}\0${residentGuidance}`;
|
||||
if (seededResident.current === seedKey) return;
|
||||
seededResident.current = seedKey;
|
||||
// Wire Reapply to the resident model too. Only a full pipeline is reloadable by repo id
|
||||
// alone; a resident GGUF carries no checkpoint filename, so the target stays null and the
|
||||
// button hidden. Set before the recipe decision below, which is a separate question.
|
||||
if (status?.model_kind === "pipeline") {
|
||||
lastLoad.current = { repoId, kind: "pipeline", displayRepoId: status.display_repo_id ?? undefined };
|
||||
}
|
||||
// Wire Reapply to the resident model too: a pipeline by repo id, a GGUF by repo id + file (so an
|
||||
// Engine change on it can take effect). Set before the recipe decision below, a separate question.
|
||||
const resident = residentReloadTarget({
|
||||
loaded: true,
|
||||
repo_id: repoId,
|
||||
model_kind: status?.model_kind,
|
||||
gguf_filename: status?.gguf_filename,
|
||||
display_repo_id: status?.display_repo_id,
|
||||
component_files: status?.component_files,
|
||||
});
|
||||
if (resident) lastLoad.current = resident;
|
||||
// A stored recipe is the user's own choice, so it outranks the resident model's defaults on the first seed.
|
||||
if (!residentSeeded.current) {
|
||||
residentSeeded.current = true;
|
||||
|
|
@ -3163,10 +3206,38 @@ export function ImagesPage({
|
|||
status?.repo_id,
|
||||
status?.base_repo,
|
||||
status?.model_kind,
|
||||
status?.gguf_filename,
|
||||
status?.component_files,
|
||||
status?.transformer_quant,
|
||||
status?.resolved?.transformer_quant?.source,
|
||||
]);
|
||||
|
||||
const loadedEngineKey =
|
||||
status?.loaded && status.model_kind === "gguf" ? engineModelKey(status.repo_id, status.gguf_filename) : "";
|
||||
const loadedEngine = loadedEngineChoice(status);
|
||||
// Reseed the select when the loaded build changes, during render rather than in an effect.
|
||||
if (engineSeed !== `${loadedEngineKey}|${loadedEngine}`) {
|
||||
setEngineSeed(`${loadedEngineKey}|${loadedEngine}`);
|
||||
setEngineChoice(loadedEngine);
|
||||
}
|
||||
const engineForced = status?.engine_forced ?? null;
|
||||
const engineSelection = useRef({ key: loadedEngineKey, choice: engineChoice, stored: engineChoices, forced: engineForced });
|
||||
useEffect(() => {
|
||||
engineSelection.current = { key: loadedEngineKey, choice: engineChoice, stored: engineChoices, forced: engineForced };
|
||||
}, [loadedEngineKey, engineChoice, engineChoices, engineForced]);
|
||||
// A load's engine: an explicit one it was handed, else the Engine select for the loaded checkpoint, else
|
||||
// the choice that checkpoint last loaded with. GGUF only; Auto omits the field.
|
||||
const engineFor = useCallback(
|
||||
(repoId: string, opts: ImageLoadOptions, advanced: LoadAdvanced) => {
|
||||
if (!sendsEngine(opts.kind, opts.filename)) return undefined;
|
||||
if (advanced.engine) return engineRequestField(advanced.engine);
|
||||
const { key, choice, stored } = engineSelection.current;
|
||||
const target = engineModelKey(repoId, opts.filename);
|
||||
return engineRequestField(target && target === key ? choice : stored[target]);
|
||||
},
|
||||
[],
|
||||
);
|
||||
|
||||
// Reseed the Advanced selects from the LOADED build, so a declined request snaps to what
|
||||
// engaged and Precision never advertises a scheme the model is not running. Keyed on the
|
||||
// LOAD-TIME half of the record: the backend rewrites the speed/attention/cache entries at
|
||||
|
|
@ -3305,7 +3376,14 @@ export function ImagesPage({
|
|||
const prevLastLoad = lastLoad.current;
|
||||
// A torchao int8/fp8 transformer takes adapters only at LOAD time and /images/generate then
|
||||
// rejects a new set, so a reload must keep the selection. Ignored by bf16 / bnb-4bit.
|
||||
const advanced = pinned ?? currentLoadAdvanced(repoId);
|
||||
const picked = pinned ?? currentLoadAdvanced(repoId);
|
||||
// A native pick drops what the native engine refuses (explicit precision, separate component files).
|
||||
const advanced = nativeLoadAdvanced(
|
||||
picked,
|
||||
engineFor(repoId, opts, picked),
|
||||
engineSelection.current.forced,
|
||||
sendsEngine(opts.kind, opts.filename),
|
||||
);
|
||||
const bakeLoras = advanced.loras ?? [];
|
||||
// Whether THIS load carries the selection into the build, so a quantized load that did not can drop it.
|
||||
bakedLorasOnLoad.current = bakeLoras.length > 0;
|
||||
|
|
@ -3318,13 +3396,16 @@ export function ImagesPage({
|
|||
textEncoderFiles: componentFiles.text_encoder_file,
|
||||
vaeFile: componentFiles.vae_file,
|
||||
});
|
||||
// A redirect loads another checkpoint, whose base the backend infers.
|
||||
const baseRepo = redirected.repoId === repoId ? opts.baseRepo : undefined;
|
||||
repoId = redirected.repoId;
|
||||
opts = { kind: redirected.kind, filename: redirected.filename, displayRepoId: redirected.displayRepoId };
|
||||
opts = { kind: redirected.kind, filename: redirected.filename, displayRepoId: redirected.displayRepoId, baseRepo };
|
||||
lastLoad.current = {
|
||||
repoId,
|
||||
kind: opts.kind,
|
||||
filename: opts.filename,
|
||||
displayRepoId: opts.displayRepoId,
|
||||
baseRepo: opts.baseRepo,
|
||||
textEncoderFiles: componentFiles.text_encoder_file,
|
||||
vaeFile: componentFiles.vae_file,
|
||||
};
|
||||
|
|
@ -3345,7 +3426,9 @@ export function ImagesPage({
|
|||
display_repo_id: opts.displayRepoId,
|
||||
model_kind: opts.kind,
|
||||
gguf_filename: opts.filename,
|
||||
base_repo: opts.baseRepo,
|
||||
hf_token: hfApiToken(getHfToken()),
|
||||
engine: engineFor(repoId, opts, advanced),
|
||||
cpu_offload: advanced.cpu_offload,
|
||||
speed_mode: advanced.speed_mode,
|
||||
// Do not carry a saved precision into a known incompatible artifact.
|
||||
|
|
@ -3369,6 +3452,9 @@ export function ImagesPage({
|
|||
lastLoad.current = prevLastLoad;
|
||||
setCanReapply(prevLastLoad != null);
|
||||
lastLoadRevert.current = null;
|
||||
// A refused load keeps the resident model and its provenance, so the seed would not move: reseed the
|
||||
// Engine select from the resident build rather than replay the refused pick on the next Reapply.
|
||||
setEngineSeed("");
|
||||
dismissLoadToast();
|
||||
reportLoadFailure(err, "Failed to start load");
|
||||
setBusy(null);
|
||||
|
|
@ -3398,7 +3484,7 @@ export function ImagesPage({
|
|||
void pollLoadProgress();
|
||||
return settle(true);
|
||||
},
|
||||
[pollLoadProgress, refreshStatus, dismissLoadToast, currentLoadAdvanced, cancelLoadFromToast, pickToast],
|
||||
[pollLoadProgress, refreshStatus, dismissLoadToast, currentLoadAdvanced, cancelLoadFromToast, pickToast, engineFor],
|
||||
);
|
||||
|
||||
// Set or clear the Transform/Inpaint source image; always drop the painted mask, which is
|
||||
|
|
@ -3630,16 +3716,25 @@ export function ImagesPage({
|
|||
(
|
||||
repoId: string,
|
||||
opts: ImageLoadOptions,
|
||||
advanced: LoadAdvanced,
|
||||
) =>
|
||||
getDiffusionDownloadPlan({
|
||||
picked: LoadAdvanced,
|
||||
) => {
|
||||
// Normalised exactly as handleLoad does, so the plan describes the load that will run.
|
||||
const advanced = nativeLoadAdvanced(
|
||||
picked,
|
||||
engineFor(repoId, opts, picked),
|
||||
engineSelection.current.forced,
|
||||
sendsEngine(opts.kind, opts.filename),
|
||||
);
|
||||
return getDiffusionDownloadPlan({
|
||||
model_path: repoId,
|
||||
gguf_filename: opts.filename,
|
||||
base_repo: opts.baseRepo,
|
||||
model_kind: opts.kind,
|
||||
// The same token and Advanced values handleLoad sends, so the plan describes the load that
|
||||
// will run: without the token a gated base plans no companion, and without the controls it
|
||||
// stages shards the load never opens.
|
||||
hf_token: hfApiToken(getHfToken()),
|
||||
engine: engineFor(repoId, opts, advanced),
|
||||
cpu_offload: advanced.cpu_offload,
|
||||
speed_mode: advanced.speed_mode,
|
||||
// Keep the planned precision identical to the load request.
|
||||
|
|
@ -3656,8 +3751,9 @@ export function ImagesPage({
|
|||
// use, so a selection the load carries has to reach the plan.
|
||||
gpu_ids: advanced.gpu_ids,
|
||||
...componentFileFields(opts.kind, advanced.text_encoder_file, advanced.vae_file),
|
||||
}),
|
||||
[],
|
||||
});
|
||||
},
|
||||
[engineFor],
|
||||
);
|
||||
|
||||
const loadOrStage = useCallback(
|
||||
|
|
@ -3668,8 +3764,11 @@ export function ImagesPage({
|
|||
token?: number,
|
||||
familyOverrideRequired = false,
|
||||
downloadSnapshot?: LoadAdvanced,
|
||||
// A reload of the resident model loads whatever the picker's Load / Download only action says.
|
||||
forceLoad = false,
|
||||
): Promise<boolean> => {
|
||||
const downloadOnly = downloadSnapshot !== undefined || modelSelectionAction === "download";
|
||||
const downloadOnly =
|
||||
downloadSnapshot !== undefined || (!forceLoad && modelSelectionAction === "download");
|
||||
// Staging never sets `busy`, so a second pick passes the guard while this plan is in flight, and
|
||||
// plans resolve in response order rather than pick order. Bumped before the non-hub return too, so
|
||||
// a local pick invalidates an in-flight hub plan. Never for a download-only pick: it supersedes
|
||||
|
|
@ -3709,6 +3808,10 @@ export function ImagesPage({
|
|||
// A local snapshot of the copy is now a Hub load of the base: plan and stage it like one.
|
||||
source = "hub";
|
||||
}
|
||||
// Resolve the engine into the snapshot, so the plan and the load it fires cannot read different selects.
|
||||
if (!advanced.engine && sendsEngine(opts.kind, opts.filename)) {
|
||||
advanced = { ...advanced, engine: engineFor(repoId, opts, advanced) ?? "auto" };
|
||||
}
|
||||
if (source !== "hub" && !downloadOnly) return handleLoadRef.current(repoId, opts, advanced);
|
||||
// Show feedback before the potentially slow Hub metadata request.
|
||||
const pickToastId = downloadOnly ? undefined : pickToast.show();
|
||||
|
|
@ -3797,7 +3900,7 @@ export function ImagesPage({
|
|||
}
|
||||
return handleLoadRef.current(repoId, opts, advanced, pickToastId);
|
||||
},
|
||||
[stage, currentLoadAdvanced, requestDownloadPlan, modelSelectionAction, pickGuard, revertPick, pickToast],
|
||||
[stage, currentLoadAdvanced, requestDownloadPlan, modelSelectionAction, pickGuard, revertPick, pickToast, engineFor],
|
||||
);
|
||||
|
||||
const resolveDownloadFootprint = useCallback(
|
||||
|
|
@ -4010,9 +4113,37 @@ export function ImagesPage({
|
|||
|
||||
// Reload the current model with the current advanced options.
|
||||
const handleReapply = useCallback(() => {
|
||||
const l = lastLoad.current;
|
||||
if (l) void handleLoad(l.repoId, { kind: l.kind, filename: l.filename, displayRepoId: l.displayRepoId });
|
||||
}, [handleLoad]);
|
||||
// Against the current status: another client may have replaced the model this page last loaded.
|
||||
const l = reapplyTarget(lastLoad.current, status);
|
||||
if (!l) return;
|
||||
const opts = { kind: l.kind, filename: l.filename, displayRepoId: l.displayRepoId, baseRepo: l.baseRepo };
|
||||
// An engine change reads the other engine's companion files: plan and stage them first.
|
||||
if (reapplyChangesEngine(engineChoice, status)) {
|
||||
void loadOrStage(l.repoId, opts, reloadSource(l.repoId), undefined, true, undefined, true);
|
||||
return;
|
||||
}
|
||||
void handleLoad(l.repoId, opts);
|
||||
}, [engineChoice, handleLoad, loadOrStage, status]);
|
||||
|
||||
// The LoRA note's action: the same GGUF pick, on the native engine. A resident model the page did not load
|
||||
// reloads from its status, which names the checkpoint.
|
||||
const handleNativeReload = useCallback(() => {
|
||||
const l = nativeReloadTarget(lastLoad.current, status);
|
||||
if (!l) return;
|
||||
setEngineChoice("sd_cpp");
|
||||
// Now, not on the next render: the plan and the staged snapshot read the select through this ref.
|
||||
engineSelection.current = { ...engineSelection.current, choice: "sd_cpp" };
|
||||
// Through the plan, which stages the native companions (VAE / text encoders) the diffusers load never fetched.
|
||||
void loadOrStage(
|
||||
l.repoId,
|
||||
{ kind: l.kind, filename: l.filename, displayRepoId: l.displayRepoId, baseRepo: l.baseRepo },
|
||||
reloadSource(l.repoId),
|
||||
undefined,
|
||||
true,
|
||||
undefined,
|
||||
true,
|
||||
);
|
||||
}, [loadOrStage, status]);
|
||||
|
||||
// Every pick supersedes the one before it, whichever route it takes: a staged download
|
||||
// outlives its pick, and the direct-local branches call handleLoad rather than loadOrStage,
|
||||
|
|
@ -5099,6 +5230,20 @@ export function ImagesPage({
|
|||
["download", "Download only"],
|
||||
]}
|
||||
/>
|
||||
{!external && engineSelectVisible(status) && (
|
||||
<AdvancedSelect
|
||||
label="Engine"
|
||||
hint={
|
||||
status?.engine_forced
|
||||
? `${ENGINE_HINT} The server pins UNSLOTH_DIFFUSION_ENGINE=${status.engine_forced}, which wins over this choice.`
|
||||
: ENGINE_HINT
|
||||
}
|
||||
badge={<ResolvedBadge status={status} controlKey="engine" />}
|
||||
value={engineChoice}
|
||||
onValueChange={(v) => setEngineChoice(v as EngineChoice)}
|
||||
options={ENGINE_OPTIONS}
|
||||
/>
|
||||
)}
|
||||
<AdvancedSelect
|
||||
label="Speed"
|
||||
hint="Auto picks per model: GGUF compiles at load; a dense model keeps the first two images exact and eager, then compiles from the 3rd (~2x from there). eager = fused kernels, no compile. default/max add torch.compile (max also TF32 + fused QKV, plus the step cache on 20+ step models)."
|
||||
|
|
@ -5914,6 +6059,14 @@ export function ImagesPage({
|
|||
hint="What to steer the image away from. Only used when guidance is above 0."
|
||||
/>
|
||||
)}
|
||||
{!external && !loraCapable && offersNativeLoraReload(status) && (
|
||||
<div className="flex items-center justify-between gap-2 rounded-md border border-dashed px-3 py-2 text-xs text-muted-foreground">
|
||||
<span>LoRAs need the native engine for this GGUF; diffusers runs it without them.</span>
|
||||
<Button size="sm" variant="secondary" className="h-7 shrink-0 text-xs" disabled={busy !== null} onClick={handleNativeReload}>
|
||||
Reload with native engine
|
||||
</Button>
|
||||
</div>
|
||||
)}
|
||||
{/* LoRA adapters: shown whenever the loaded model + quant can apply them. Each carries a 0-2 weight. */}
|
||||
{loraCapable && (
|
||||
<Field
|
||||
|
|
@ -6200,7 +6353,7 @@ export function ImagesPage({
|
|||
>
|
||||
Generate
|
||||
</Button>
|
||||
{!external && status?.loaded && (canReapply || status?.model_kind === "pipeline") && (
|
||||
{!external && status?.loaded && (canReapply || residentReloadTarget(status) !== null) && (
|
||||
<Tooltip>
|
||||
<TooltipTrigger asChild={true}>
|
||||
<Button
|
||||
|
|
|
|||
231
studio/frontend/src/features/images/lib/engine-choice.ts
Normal file
231
studio/frontend/src/features/images/lib/engine-choice.ts
Normal file
|
|
@ -0,0 +1,231 @@
|
|||
// SPDX-License-Identifier: AGPL-3.0-only
|
||||
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
// The Images Advanced "Engine" choice for GGUF picks: Auto (the host's own choice), diffusers, or the native
|
||||
// sd.cpp engine, remembered per checkpoint. Only a load that landed records its choice, so a refused pick
|
||||
// (no native binary, an unsupported family) is never replayed on the next load of that model.
|
||||
|
||||
export type EngineChoice = "auto" | "diffusers" | "sd_cpp";
|
||||
|
||||
export const ENGINE_OPTIONS: Array<[EngineChoice, string]> = [
|
||||
["auto", "Auto"],
|
||||
["diffusers", "Diffusers"],
|
||||
["sd_cpp", "Native (sd.cpp)"],
|
||||
];
|
||||
|
||||
export const ENGINE_HINT =
|
||||
"Which engine runs this GGUF. Auto keeps the host's choice: diffusers on a GPU, native sd.cpp on CPU. Native applies LoRAs to GGUF models and usually needs less VRAM, but has no torch.compile or quantised fast paths and is often slower on CUDA. Diffusers runs a GGUF without LoRA. Applies on the next load (Reapply).";
|
||||
|
||||
const KEY = "unsloth:images:engine-by-model";
|
||||
|
||||
export function engineModelKey(repoId: string | null | undefined, filename: string | null | undefined): string {
|
||||
return repoId && filename ? `${repoId}|${filename}` : "";
|
||||
}
|
||||
|
||||
function isPick(value: unknown): value is Exclude<EngineChoice, "auto"> {
|
||||
return value === "diffusers" || value === "sd_cpp";
|
||||
}
|
||||
|
||||
export function readEngineChoices(): Record<string, Exclude<EngineChoice, "auto">> {
|
||||
try {
|
||||
const value = JSON.parse(localStorage.getItem(KEY) ?? "{}");
|
||||
if (!value || typeof value !== "object" || Array.isArray(value)) return {};
|
||||
const out: Record<string, Exclude<EngineChoice, "auto">> = {};
|
||||
for (const [k, v] of Object.entries(value)) if (isPick(v)) out[k] = v;
|
||||
return out;
|
||||
} catch {
|
||||
return {};
|
||||
}
|
||||
}
|
||||
|
||||
/** Store the choice for one checkpoint; Auto drops the entry. */
|
||||
export function withEngineChoice(
|
||||
choices: Record<string, Exclude<EngineChoice, "auto">>,
|
||||
modelKey: string,
|
||||
choice: EngineChoice,
|
||||
): Record<string, Exclude<EngineChoice, "auto">> {
|
||||
if (!modelKey) return choices;
|
||||
const next = { ...choices };
|
||||
if (isPick(choice)) next[modelKey] = choice;
|
||||
else delete next[modelKey];
|
||||
try {
|
||||
localStorage.setItem(KEY, JSON.stringify(next));
|
||||
} catch {
|
||||
// storage can be unavailable in private browsing
|
||||
}
|
||||
return next;
|
||||
}
|
||||
|
||||
/** Whether a load of this kind may carry an engine pick: the native engine loads GGUF only. */
|
||||
export function sendsEngine(kind: string | null | undefined, filename: string | null | undefined): boolean {
|
||||
if (kind) return kind === "gguf";
|
||||
return (filename ?? "").toLowerCase().endsWith(".gguf");
|
||||
}
|
||||
|
||||
/** The request field for a load, omitted for Auto so an untouched load stays byte-identical. */
|
||||
export function engineRequestField(choice: EngineChoice | undefined): Exclude<EngineChoice, "auto"> | undefined {
|
||||
return isPick(choice) ? choice : undefined;
|
||||
}
|
||||
|
||||
interface EngineStatus {
|
||||
loaded: boolean;
|
||||
engine?: string | null;
|
||||
engine_forced?: string | null;
|
||||
native_engine_eligible?: boolean;
|
||||
lora_via_native?: boolean;
|
||||
supports_lora?: boolean;
|
||||
resolved?: Record<string, { source: string; requested?: string | boolean | null }> | null;
|
||||
}
|
||||
|
||||
/** The select shows only for a loaded GGUF the native engine can run. */
|
||||
export function engineSelectVisible(status: EngineStatus | null | undefined): boolean {
|
||||
return Boolean(status?.loaded && status.native_engine_eligible);
|
||||
}
|
||||
|
||||
/** The select value for the loaded build: what its load asked for, Auto when it asked for nothing. */
|
||||
export function loadedEngineChoice(status: EngineStatus | null | undefined): EngineChoice {
|
||||
const record = status?.loaded ? status.resolved?.engine : undefined;
|
||||
if (record?.source === "explicit" && isPick(record.requested)) return record.requested;
|
||||
return "auto";
|
||||
}
|
||||
|
||||
/** The note under the LoRA slot: LoRA is off on this engine, and a native reload would turn it on. A server pin
|
||||
* (UNSLOTH_DIFFUSION_ENGINE) would ignore the reload's pick, so it offers nothing then. */
|
||||
export function offersNativeLoraReload(status: EngineStatus | null | undefined): boolean {
|
||||
return Boolean(
|
||||
status?.loaded && status.lora_via_native && !status.supports_lora && !status.engine_forced,
|
||||
);
|
||||
}
|
||||
|
||||
interface ReloadStatus {
|
||||
loaded: boolean;
|
||||
repo_id?: string | null;
|
||||
model_kind?: string | null;
|
||||
gguf_filename?: string | null;
|
||||
display_repo_id?: string | null;
|
||||
base_repo?: string | null;
|
||||
}
|
||||
|
||||
/** The resident GGUF's effective base repo, so a reload from status rebuilds the same companions (an API load may have
|
||||
* named a non-default base). Omitted when status names none. */
|
||||
function baseRepoOf(status: ReloadStatus): { baseRepo?: string } {
|
||||
return status.base_repo ? { baseRepo: status.base_repo } : {};
|
||||
}
|
||||
|
||||
/** The GGUF the native-reload note names: the page's own last load only while it is still the resident model (another
|
||||
* client may have replaced it), else the checkpoint status reports. */
|
||||
export function nativeReloadTarget<T extends { repoId: string; kind?: string | null; filename?: string | null }>(
|
||||
last: T | null | undefined,
|
||||
status: ReloadStatus | null | undefined,
|
||||
): T | { repoId: string; kind: "gguf"; filename: string; displayRepoId?: string; baseRepo?: string } | null {
|
||||
if (!status?.loaded || status.model_kind !== "gguf" || !status.repo_id || !status.gguf_filename) return null;
|
||||
if (last && last.repoId === status.repo_id && last.kind === "gguf" && last.filename === status.gguf_filename) {
|
||||
return last;
|
||||
}
|
||||
return {
|
||||
repoId: status.repo_id,
|
||||
kind: "gguf",
|
||||
filename: status.gguf_filename,
|
||||
displayRepoId: status.display_repo_id ?? undefined,
|
||||
...baseRepoOf(status),
|
||||
};
|
||||
}
|
||||
|
||||
interface EngineAdvanced {
|
||||
transformer_quant?: string;
|
||||
text_encoder_quant?: string;
|
||||
text_encoder_file?: string | string[];
|
||||
vae_file?: string;
|
||||
}
|
||||
|
||||
/** The Advanced values an explicit native load can carry. The load and plan routes refuse an explicit precision
|
||||
* (``_assert_native_precision_unset``, 409) and separate text-encoder / VAE files (``_refuse_native_component_files``,
|
||||
* 400) on sd.cpp, so a native pick drops them rather than fail. Only an explicit pick, which the router serves natively
|
||||
* or refuses: under a server pin (UNSLOTH_DIFFUSION_ENGINE) the pick is ignored, and a pin to sd_cpp still falls back
|
||||
* to diffusers when native cannot run the pick, which must then build with the requested options. */
|
||||
export function nativeLoadAdvanced<T extends EngineAdvanced>(
|
||||
advanced: T,
|
||||
engine: EngineChoice | undefined,
|
||||
forced?: string | null,
|
||||
// Whether the pick is a GGUF, the only kind the native engine loads.
|
||||
gguf = true,
|
||||
): T {
|
||||
if (!gguf || forced || engine !== "sd_cpp") return advanced;
|
||||
return { ...advanced, transformer_quant: undefined, text_encoder_quant: undefined, text_encoder_file: undefined, vae_file: undefined };
|
||||
}
|
||||
|
||||
interface LoadedEngineStatus extends EngineStatus {
|
||||
model_kind?: string | null;
|
||||
repo_id?: string | null;
|
||||
gguf_filename?: string | null;
|
||||
}
|
||||
|
||||
/** Record the choice a resident GGUF loaded with, from any accepted status (a load that finished while the page was
|
||||
* closed is never polled to ready). The latest load wins: Auto drops a stored explicit entry, as the ready branch does.
|
||||
* Returns ``choices`` itself when there is nothing new, so an unchanged status re-renders and writes nothing. */
|
||||
export function rememberLoadedEngine(
|
||||
choices: Record<string, Exclude<EngineChoice, "auto">>,
|
||||
status: LoadedEngineStatus | null | undefined,
|
||||
): Record<string, Exclude<EngineChoice, "auto">> {
|
||||
if (!status?.loaded || status.model_kind !== "gguf") return choices;
|
||||
const key = engineModelKey(status.repo_id, status.gguf_filename);
|
||||
const choice = loadedEngineChoice(status);
|
||||
if (!key || (isPick(choice) ? choices[key] === choice : !(key in choices))) return choices;
|
||||
return withEngineChoice(choices, key, choice);
|
||||
}
|
||||
|
||||
interface ResidentStatus extends ReloadStatus {
|
||||
component_files?: Record<string, string> | null;
|
||||
}
|
||||
|
||||
/** A Reapply target for a resident model this page did not load: a pipeline by repo id, a GGUF by repo id + file (so
|
||||
* an Engine change on it can take effect). None for a GGUF with separate component files, whose paths status names
|
||||
* only by basename, so a reload from status would silently drop them. */
|
||||
export function residentReloadTarget(
|
||||
status: ResidentStatus | null | undefined,
|
||||
):
|
||||
| { repoId: string; kind: "pipeline"; filename?: undefined; displayRepoId?: string; baseRepo?: undefined }
|
||||
| { repoId: string; kind: "gguf"; filename: string; displayRepoId?: string; baseRepo?: string }
|
||||
| null {
|
||||
if (!status?.loaded || !status.repo_id) return null;
|
||||
const displayRepoId = status.display_repo_id ?? undefined;
|
||||
if (status.model_kind === "pipeline") return { repoId: status.repo_id, kind: "pipeline", displayRepoId };
|
||||
if (
|
||||
status.model_kind === "gguf" &&
|
||||
status.gguf_filename &&
|
||||
Object.keys(status.component_files ?? {}).length === 0
|
||||
) {
|
||||
return { repoId: status.repo_id, kind: "gguf", filename: status.gguf_filename, displayRepoId, ...baseRepoOf(status) };
|
||||
}
|
||||
return null;
|
||||
}
|
||||
|
||||
/** What Reapply reloads: the page's own last load while it is still the resident model, else the checkpoint status
|
||||
* reports (another client may have replaced it), never a stale pick the page is no longer showing. */
|
||||
export function reapplyTarget<T extends { repoId: string; kind?: string | null; filename?: string | null }>(
|
||||
last: T | null | undefined,
|
||||
status: ResidentStatus | null | undefined,
|
||||
): T | ReturnType<typeof residentReloadTarget> {
|
||||
if (!status?.loaded) return last ?? null;
|
||||
if (
|
||||
last &&
|
||||
last.repoId === status.repo_id &&
|
||||
last.kind === status.model_kind &&
|
||||
(last.kind === "pipeline" || last.filename === status.gguf_filename)
|
||||
) {
|
||||
return last;
|
||||
}
|
||||
return residentReloadTarget(status);
|
||||
}
|
||||
|
||||
/** Whether a Reapply of the resident GGUF would move it to another engine. The two engines read different companion
|
||||
* files, so such a reload goes through the download plan first. A server pin ignores the pick, so it moves nothing. */
|
||||
export function reapplyChangesEngine(choice: EngineChoice, status: EngineStatus & { model_kind?: string | null } | null | undefined): boolean {
|
||||
if (!status?.loaded || status.model_kind !== "gguf" || status.engine_forced) return false;
|
||||
return choice !== loadedEngineChoice(status);
|
||||
}
|
||||
|
||||
/** The picker source a reload of ``repoId`` takes: a local path loads directly, a Hub repo id is planned and staged. */
|
||||
export function reloadSource(repoId: string): "hub" | "local" {
|
||||
return /^(\/|~|[A-Za-z]:[\\/]|\\\\)/.test(repoId) ? "local" : "hub";
|
||||
}
|
||||
|
|
@ -64,9 +64,14 @@ function isOff(value: string | boolean | null | undefined): boolean {
|
|||
return text === "" || text === "none" || text === "off" || text === "0";
|
||||
}
|
||||
|
||||
const ENGINE_LABELS: Record<string, string> = { sd_cpp: "Native", diffusers: "Diffusers", mlx: "MLX" };
|
||||
|
||||
/** The engaged value of a resolved Advanced control, formatted for its badge. */
|
||||
export function formatResolvedValue(key: string, value: string | boolean | null | undefined): string {
|
||||
if (key === "cpu_offload") return value ? "On" : "Off";
|
||||
if (key === "engine" && typeof value === "string") {
|
||||
return ENGINE_LABELS[value] ?? value.toUpperCase();
|
||||
}
|
||||
if (value === null || value === undefined || value === "") return "Off";
|
||||
if (typeof value === "boolean") return value ? "On" : "Off";
|
||||
if (value === "_native_cudnn" || value.toLowerCase() === "cudnn") return "cuDNN";
|
||||
|
|
|
|||
313
studio/frontend/tests/image-engine-choice.test.ts
Normal file
313
studio/frontend/tests/image-engine-choice.test.ts
Normal file
|
|
@ -0,0 +1,313 @@
|
|||
// SPDX-License-Identifier: AGPL-3.0-only
|
||||
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
import assert from "node:assert/strict";
|
||||
import test from "node:test";
|
||||
import {
|
||||
ENGINE_OPTIONS,
|
||||
engineModelKey,
|
||||
engineRequestField,
|
||||
engineSelectVisible,
|
||||
loadedEngineChoice,
|
||||
nativeLoadAdvanced,
|
||||
nativeReloadTarget,
|
||||
offersNativeLoraReload,
|
||||
reapplyChangesEngine,
|
||||
reapplyTarget,
|
||||
readEngineChoices,
|
||||
reloadSource,
|
||||
residentReloadTarget,
|
||||
rememberLoadedEngine,
|
||||
sendsEngine,
|
||||
withEngineChoice,
|
||||
} from "../src/features/images/lib/engine-choice.ts";
|
||||
import { formatResolvedValue, resolvedBadge } from "../src/lib/resolved-precision.ts";
|
||||
import { installLocalStorageFake, readSrc } from "./helpers/kit.ts";
|
||||
|
||||
const { storage, store } = installLocalStorageFake();
|
||||
|
||||
const ggufOnDiffusers = {
|
||||
loaded: true,
|
||||
engine: "diffusers",
|
||||
engine_forced: null,
|
||||
native_engine_eligible: true,
|
||||
lora_via_native: true,
|
||||
supports_lora: false,
|
||||
};
|
||||
|
||||
test("the select offers Auto, Diffusers and Native, and shows only for a GGUF native can run", () => {
|
||||
assert.deepEqual(
|
||||
ENGINE_OPTIONS.map(([v]) => v),
|
||||
["auto", "diffusers", "sd_cpp"],
|
||||
);
|
||||
assert.equal(engineSelectVisible(ggufOnDiffusers), true);
|
||||
assert.equal(engineSelectVisible({ ...ggufOnDiffusers, native_engine_eligible: false }), false);
|
||||
assert.equal(engineSelectVisible({ loaded: false }), false);
|
||||
// An older backend sends no eligibility.
|
||||
assert.equal(engineSelectVisible({ loaded: true, engine: "diffusers" }), false);
|
||||
assert.equal(engineSelectVisible(null), false);
|
||||
});
|
||||
|
||||
test("only a GGUF load carries the pick, and Auto omits it", () => {
|
||||
assert.equal(sendsEngine("gguf", "x.gguf"), true);
|
||||
assert.equal(sendsEngine("pipeline", undefined), false);
|
||||
assert.equal(sendsEngine("single_file", "x.safetensors"), false);
|
||||
assert.equal(sendsEngine(undefined, "Model-Q4_K_M.GGUF"), true);
|
||||
assert.equal(engineRequestField("auto"), undefined);
|
||||
assert.equal(engineRequestField(undefined), undefined);
|
||||
assert.equal(engineRequestField("sd_cpp"), "sd_cpp");
|
||||
assert.equal(engineRequestField("diffusers"), "diffusers");
|
||||
});
|
||||
|
||||
test("the choice is remembered per checkpoint and Auto forgets it", () => {
|
||||
storage.removeItem("unsloth:images:engine-by-model");
|
||||
const key = engineModelKey("unsloth/FLUX.2-klein-4B-GGUF", "flux-2-klein-4b-Q4_K_M.gguf");
|
||||
assert.equal(key, "unsloth/FLUX.2-klein-4B-GGUF|flux-2-klein-4b-Q4_K_M.gguf");
|
||||
assert.equal(engineModelKey("repo", null), "");
|
||||
let choices = withEngineChoice({}, key, "sd_cpp");
|
||||
assert.deepEqual(readEngineChoices(), { [key]: "sd_cpp" });
|
||||
choices = withEngineChoice(choices, "other|b.gguf", "diffusers");
|
||||
choices = withEngineChoice(choices, key, "auto");
|
||||
assert.deepEqual(choices, { "other|b.gguf": "diffusers" });
|
||||
assert.deepEqual(readEngineChoices(), { "other|b.gguf": "diffusers" });
|
||||
// No key, nothing stored.
|
||||
assert.deepEqual(withEngineChoice(choices, "", "sd_cpp"), choices);
|
||||
// Junk in storage reads as no choice.
|
||||
storage.setItem("unsloth:images:engine-by-model", JSON.stringify({ a: "vulkan", b: 3, c: "sd_cpp" }));
|
||||
assert.deepEqual(readEngineChoices(), { c: "sd_cpp" });
|
||||
storage.setItem("unsloth:images:engine-by-model", "not json");
|
||||
assert.deepEqual(readEngineChoices(), {});
|
||||
});
|
||||
|
||||
test("the loaded build's select value follows the backend's engine provenance", () => {
|
||||
const resolved = (entry: { source: string; requested?: string | null; value?: string }) => ({
|
||||
...ggufOnDiffusers,
|
||||
resolved: { engine: entry },
|
||||
});
|
||||
assert.equal(loadedEngineChoice(resolved({ source: "explicit", requested: "sd_cpp", value: "sd_cpp" })), "sd_cpp");
|
||||
assert.equal(loadedEngineChoice(resolved({ source: "explicit", requested: "diffusers" })), "diffusers");
|
||||
assert.equal(loadedEngineChoice(resolved({ source: "auto", requested: null, value: "diffusers" })), "auto");
|
||||
assert.equal(loadedEngineChoice(ggufOnDiffusers), "auto");
|
||||
assert.equal(loadedEngineChoice({ loaded: false, resolved: { engine: { source: "explicit", requested: "sd_cpp" } } }), "auto");
|
||||
});
|
||||
|
||||
test("the LoRA note offers a native reload only where it would turn LoRA on", () => {
|
||||
assert.equal(offersNativeLoraReload(ggufOnDiffusers), true);
|
||||
assert.equal(offersNativeLoraReload({ ...ggufOnDiffusers, lora_via_native: false }), false);
|
||||
assert.equal(offersNativeLoraReload({ ...ggufOnDiffusers, supports_lora: true }), false);
|
||||
// The server pin would ignore the reload's pick.
|
||||
assert.equal(offersNativeLoraReload({ ...ggufOnDiffusers, engine_forced: "diffusers" }), false);
|
||||
assert.equal(offersNativeLoraReload({ loaded: false }), false);
|
||||
});
|
||||
|
||||
test("the engine badge names the engine", () => {
|
||||
assert.equal(formatResolvedValue("engine", "sd_cpp"), "Native");
|
||||
assert.equal(formatResolvedValue("engine", "diffusers"), "Diffusers");
|
||||
assert.equal(formatResolvedValue("engine", "mlx"), "MLX");
|
||||
assert.deepEqual(resolvedBadge("engine", { value: "diffusers", requested: null, source: "auto", status: "applied", reason: "GPU backend 'cuda' uses diffusers" }), {
|
||||
label: "Auto: Diffusers",
|
||||
tone: "auto",
|
||||
tooltip: "GPU backend 'cuda' uses diffusers",
|
||||
});
|
||||
assert.equal(resolvedBadge("engine", { value: "sd_cpp", requested: "sd_cpp", source: "explicit", status: "applied", reason: "selected for this load" }), null);
|
||||
});
|
||||
|
||||
test("the page sends the pick on the load and the plan, and wires the select and the note", () => {
|
||||
const source = readSrc("features/images/images-page.tsx");
|
||||
assert.equal(source.match(/engine: engineFor\(repoId, opts, advanced\),/g)?.length, 2);
|
||||
assert.match(source, /label="Engine"[\s\S]{0,400}controlKey="engine"/);
|
||||
assert.match(source, /offersNativeLoraReload\(status\)[\s\S]{0,600}onClick=\{handleNativeReload\}[\s\S]{0,80}Reload with native engine/);
|
||||
assert.match(source, /setEngineChoice\("sd_cpp"\);[\s\S]{0,120}engineSelection\.current = \{ \.\.\.engineSelection\.current, choice: "sd_cpp" \};/);
|
||||
// Recorded from the load that landed, never from the select alone.
|
||||
assert.match(source, /setEngineChoices\(\(all\) => withEngineChoice\(all, key, loadedEngineChoice\(loaded\)\)\)/);
|
||||
});
|
||||
|
||||
test("the native reload targets the resident GGUF, not a stale page load another client replaced", () => {
|
||||
const resident = { loaded: true, model_kind: "gguf", repo_id: "org/new-GGUF", gguf_filename: "new-Q4.gguf", display_repo_id: null };
|
||||
const stale = { repoId: "org/old-GGUF", kind: "gguf" as const, filename: "old-Q4.gguf" };
|
||||
assert.deepEqual(nativeReloadTarget(stale, resident), {
|
||||
repoId: "org/new-GGUF",
|
||||
kind: "gguf",
|
||||
filename: "new-Q4.gguf",
|
||||
displayRepoId: undefined,
|
||||
});
|
||||
const own = { repoId: "org/new-GGUF", kind: "gguf" as const, filename: "new-Q4.gguf", displayRepoId: "Shown" };
|
||||
assert.equal(nativeReloadTarget(own, resident), own);
|
||||
assert.equal(nativeReloadTarget(own, { loaded: false }), null);
|
||||
assert.equal(nativeReloadTarget(own, { ...resident, model_kind: "pipeline" }), null);
|
||||
const source = readSrc("features/images/images-page.tsx");
|
||||
assert.match(source, /const l = nativeReloadTarget\(lastLoad\.current, status\);/);
|
||||
});
|
||||
|
||||
test("a refused load reseeds the Engine select from the resident build", () => {
|
||||
const source = readSrc("features/images/images-page.tsx");
|
||||
assert.match(source, /setEngineSeed\(""\);[\s\S]{0,80}dismissLoadToast\(\);\s*reportLoadFailure\(err, "Failed to start load"\)/);
|
||||
});
|
||||
|
||||
test("a staged pick snapshots its resolved engine for both the plan and the load", () => {
|
||||
const source = readSrc("features/images/images-page.tsx");
|
||||
assert.match(
|
||||
source,
|
||||
/if \(!advanced\.engine && sendsEngine\(opts\.kind, opts\.filename\)\) \{\s*advanced = \{ \.\.\.advanced, engine: engineFor\(repoId, opts, advanced\) \?\? "auto" \};\s*\}\s*if \(source !== "hub" && !downloadOnly\)/,
|
||||
);
|
||||
// "auto" in the snapshot omits the field, so an untouched load stays byte-identical, and a later select edit
|
||||
// cannot leak into it.
|
||||
assert.equal(engineRequestField("auto"), undefined);
|
||||
});
|
||||
|
||||
test("a native load drops the precision and component files the native routes refuse", () => {
|
||||
const advanced = {
|
||||
cpu_offload: false,
|
||||
speed_mode: "eager",
|
||||
transformer_quant: "fp8",
|
||||
text_encoder_quant: "fp8",
|
||||
text_encoder_file: ["/m/te.safetensors"],
|
||||
vae_file: "/m/vae.safetensors",
|
||||
loras: [{ id: "a", weight: 1 }],
|
||||
};
|
||||
const native = nativeLoadAdvanced(advanced, "sd_cpp");
|
||||
assert.equal(native.transformer_quant, undefined);
|
||||
assert.equal(native.text_encoder_quant, undefined);
|
||||
assert.equal(native.text_encoder_file, undefined);
|
||||
assert.equal(native.vae_file, undefined);
|
||||
// Everything else, LoRAs included (the point of the reload), is carried.
|
||||
assert.deepEqual(native.loras, advanced.loras);
|
||||
assert.equal(native.speed_mode, "eager");
|
||||
// Auto, diffusers, or a server pin to diffusers (which ignores the pick) keep the values.
|
||||
assert.equal(nativeLoadAdvanced(advanced, undefined), advanced);
|
||||
assert.equal(nativeLoadAdvanced(advanced, "diffusers"), advanced);
|
||||
assert.equal(nativeLoadAdvanced(advanced, "sd_cpp", "diffusers"), advanced);
|
||||
assert.equal(nativeLoadAdvanced(advanced, "sd_cpp", "sd_cpp"), advanced);
|
||||
// Applied to both the load and its plan.
|
||||
const source = readSrc("features/images/images-page.tsx");
|
||||
assert.equal(
|
||||
source.match(/const advanced = nativeLoadAdvanced\(\s*picked,\s*engineFor\(repoId, opts, picked\),\s*engineSelection\.current\.forced,/g)?.length,
|
||||
2,
|
||||
);
|
||||
});
|
||||
|
||||
test("an explicit choice is persisted from any accepted loaded status, not only the poll's ready branch", () => {
|
||||
store.clear();
|
||||
const explicit = {
|
||||
loaded: true,
|
||||
model_kind: "gguf",
|
||||
repo_id: "org/klein-GGUF",
|
||||
gguf_filename: "klein-Q4.gguf",
|
||||
resolved: { engine: { source: "explicit", requested: "sd_cpp" } },
|
||||
};
|
||||
const stored = rememberLoadedEngine({}, explicit);
|
||||
assert.deepEqual(stored, { "org/klein-GGUF|klein-Q4.gguf": "sd_cpp" });
|
||||
assert.deepEqual(readEngineChoices(), stored);
|
||||
// Unchanged: the same object, so nothing re-renders.
|
||||
assert.equal(rememberLoadedEngine(stored, explicit), stored);
|
||||
// An unloaded status and a non-GGUF never write or clear.
|
||||
assert.equal(rememberLoadedEngine(stored, { loaded: false }), stored);
|
||||
assert.equal(rememberLoadedEngine(stored, { ...explicit, model_kind: "pipeline" }), stored);
|
||||
const source = readSrc("features/images/images-page.tsx");
|
||||
assert.match(source, /setStatus\(next\);[\s\S]{0,160}setEngineChoices\(\(all\) => rememberLoadedEngine\(all, next\)\)/);
|
||||
});
|
||||
|
||||
test("a resident load that reports Auto clears the stored explicit choice: the latest load wins", () => {
|
||||
store.clear();
|
||||
const key = "org/klein-GGUF|klein-Q4.gguf";
|
||||
const auto = {
|
||||
loaded: true,
|
||||
model_kind: "gguf",
|
||||
repo_id: "org/klein-GGUF",
|
||||
gguf_filename: "klein-Q4.gguf",
|
||||
resolved: { engine: { source: "auto", requested: null } },
|
||||
};
|
||||
const stored = { [key]: "sd_cpp" as const, "other|x.gguf": "diffusers" as const };
|
||||
const cleared = rememberLoadedEngine(stored, auto);
|
||||
assert.deepEqual(cleared, { "other|x.gguf": "diffusers" });
|
||||
assert.deepEqual(readEngineChoices(), cleared);
|
||||
// Nothing stored for it: unchanged, no write.
|
||||
assert.equal(rememberLoadedEngine(cleared, auto), cleared);
|
||||
});
|
||||
|
||||
test("a resident GGUF the page did not load gets a Reapply target from status", () => {
|
||||
const gguf = { loaded: true, model_kind: "gguf", repo_id: "org/klein-GGUF", gguf_filename: "klein-Q4.gguf", display_repo_id: null };
|
||||
assert.deepEqual(residentReloadTarget(gguf), { repoId: "org/klein-GGUF", kind: "gguf", filename: "klein-Q4.gguf", displayRepoId: undefined });
|
||||
assert.deepEqual(residentReloadTarget({ loaded: true, model_kind: "pipeline", repo_id: "org/p", display_repo_id: "Shown" }), {
|
||||
repoId: "org/p",
|
||||
kind: "pipeline",
|
||||
displayRepoId: "Shown",
|
||||
});
|
||||
// Status names component files by basename only: reloading from it would drop them.
|
||||
assert.equal(residentReloadTarget({ ...gguf, component_files: { vae: "ae.safetensors" } }), null);
|
||||
assert.equal(residentReloadTarget({ ...gguf, gguf_filename: null }), null);
|
||||
assert.equal(residentReloadTarget({ loaded: false }), null);
|
||||
const source = readSrc("features/images/images-page.tsx");
|
||||
assert.match(source, /const resident = residentReloadTarget\(\{[\s\S]{0,300}\}\);\s*if \(resident\) lastLoad\.current = resident;/);
|
||||
// The Reapply button shows for that target too, not only for a pipeline.
|
||||
assert.match(source, /status\?\.loaded && \(canReapply \|\| residentReloadTarget\(status\) !== null\) &&/);
|
||||
});
|
||||
|
||||
test("an engine-changing reload is planned and staged, a same-engine Reapply loads directly", () => {
|
||||
const onDiffusers = { loaded: true, model_kind: "gguf", engine: "diffusers", resolved: { engine: { source: "auto", requested: null } } };
|
||||
assert.equal(reapplyChangesEngine("sd_cpp", onDiffusers), true);
|
||||
assert.equal(reapplyChangesEngine("diffusers", onDiffusers), true);
|
||||
assert.equal(reapplyChangesEngine("auto", onDiffusers), false);
|
||||
const onNative = { ...onDiffusers, engine: "sd_cpp", resolved: { engine: { source: "explicit", requested: "sd_cpp" } } };
|
||||
assert.equal(reapplyChangesEngine("sd_cpp", onNative), false);
|
||||
assert.equal(reapplyChangesEngine("auto", onNative), true);
|
||||
// A server pin ignores the pick; a pipeline has no engine choice.
|
||||
assert.equal(reapplyChangesEngine("sd_cpp", { ...onDiffusers, engine_forced: "diffusers" }), false);
|
||||
assert.equal(reapplyChangesEngine("sd_cpp", { ...onDiffusers, model_kind: "pipeline" }), false);
|
||||
assert.equal(reloadSource("unsloth/FLUX.2-klein-GGUF"), "hub");
|
||||
assert.equal(reloadSource("/models/klein"), "local");
|
||||
assert.equal(reloadSource("C:\\models\\klein"), "local");
|
||||
const source = readSrc("features/images/images-page.tsx");
|
||||
// Both the note's native reload and an engine-changing Reapply go through loadOrStage, loading even under Download only.
|
||||
assert.equal(source.match(/reloadSource\(l\.repoId\),?\s*undefined,\s*true,\s*undefined,\s*true,?\s*\)/g)?.length, 2);
|
||||
assert.match(source, /if \(reapplyChangesEngine\(engineChoice, status\)\) \{\s*void loadOrStage\(/);
|
||||
assert.match(source, /downloadSnapshot !== undefined \|\| \(!forceLoad && modelSelectionAction === "download"\)/);
|
||||
});
|
||||
|
||||
test("Reapply resolves its target against the checkpoint status reports", () => {
|
||||
const own = { repoId: "org/A-GGUF", kind: "gguf" as const, filename: "a.gguf" };
|
||||
const onA = { loaded: true, model_kind: "gguf", repo_id: "org/A-GGUF", gguf_filename: "a.gguf" };
|
||||
assert.equal(reapplyTarget(own, onA), own);
|
||||
// Another client replaced A with B: Reapply reloads B, never swaps back to A.
|
||||
const onB = { loaded: true, model_kind: "gguf", repo_id: "org/B-GGUF", gguf_filename: "b.gguf", base_repo: "org/B-base" };
|
||||
assert.deepEqual(reapplyTarget(own, onB), { repoId: "org/B-GGUF", kind: "gguf", filename: "b.gguf", displayRepoId: undefined, baseRepo: "org/B-base" });
|
||||
// B is not reconstructable from status: nothing, rather than A.
|
||||
assert.equal(reapplyTarget(own, { ...onB, component_files: { vae: "ae.safetensors" } }), null);
|
||||
// A pipeline matches by repo id alone; nothing resident keeps the page's own target.
|
||||
const pipe = { repoId: "org/P", kind: "pipeline" as const };
|
||||
assert.equal(reapplyTarget(pipe, { loaded: true, model_kind: "pipeline", repo_id: "org/P" }), pipe);
|
||||
assert.equal(reapplyTarget(own, { loaded: false }), own);
|
||||
const source = readSrc("features/images/images-page.tsx");
|
||||
assert.match(source, /const l = reapplyTarget\(lastLoad\.current, status\);/);
|
||||
});
|
||||
|
||||
test("a server pin never strips options: its native route falls back to diffusers when native cannot run", () => {
|
||||
const advanced = { transformer_quant: "fp8", text_encoder_quant: "fp8", text_encoder_file: ["/m/te.safetensors"], vae_file: "/m/vae.safetensors" };
|
||||
// The pin ignores the pick, and its fallback must build with the requested options.
|
||||
for (const engine of [undefined, "auto", "diffusers", "sd_cpp"] as const) {
|
||||
assert.equal(nativeLoadAdvanced(advanced, engine, "sd_cpp", true), advanced, String(engine));
|
||||
}
|
||||
// An explicit pick (refused, never swapped, when native cannot run) still drops them; a non-GGUF never does.
|
||||
assert.equal(nativeLoadAdvanced(advanced, "sd_cpp", null, true).text_encoder_quant, undefined);
|
||||
assert.equal(nativeLoadAdvanced(advanced, "sd_cpp", null, false), advanced);
|
||||
const source = readSrc("features/images/images-page.tsx");
|
||||
assert.equal(source.match(/engineSelection\.current\.forced,\s*sendsEngine\(opts\.kind, opts\.filename\),/g)?.length, 2);
|
||||
});
|
||||
|
||||
test("a background load that fails reseeds the Engine select, like a refused start", () => {
|
||||
const source = readSrc("features/images/images-page.tsx");
|
||||
assert.match(source, /if \(p\.phase === "error"\) \{[\s\S]{0,1200}setEngineSeed\(""\);[\s\S]{0,200}void refreshStatus\(\);\s*return;/);
|
||||
});
|
||||
|
||||
test("a reload from status keeps the resident GGUF's base repo, on the load and the plan", () => {
|
||||
const gguf = { loaded: true, model_kind: "gguf", repo_id: "org/klein-GGUF", gguf_filename: "klein-Q4.gguf", base_repo: "org/custom-base" };
|
||||
assert.equal(residentReloadTarget(gguf)?.baseRepo, "org/custom-base");
|
||||
assert.equal(nativeReloadTarget(null, gguf)?.baseRepo, "org/custom-base");
|
||||
assert.equal("baseRepo" in (residentReloadTarget({ ...gguf, base_repo: null }) ?? {}), false);
|
||||
const source = readSrc("features/images/images-page.tsx");
|
||||
assert.equal(source.match(/base_repo: opts\.baseRepo,/g)?.length, 2);
|
||||
assert.equal(source.match(/displayRepoId: l\.displayRepoId, baseRepo: l\.baseRepo \}/g)?.length, 2);
|
||||
// Carried into the next Reapply target, dropped by a retired-copy redirect.
|
||||
assert.match(source, /const baseRepo = redirected\.repoId === repoId \? opts\.baseRepo : undefined;/);
|
||||
assert.match(source, /displayRepoId: opts\.displayRepoId,\s*baseRepo: opts\.baseRepo,/);
|
||||
});
|
||||
Loading…
Add table
Add a link
Reference in a new issue