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

Merge remote-tracking branch 'origin/diffusion-controlnet' into diffusion-sdxl

Browse source
This commit is contained in:
Daniel Han 2026-07-02 02:28:33 +00:00
commit e6bf4c4cd6
16 changed files with 2286 additions and 1408 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -520,6 +520,9 @@ class DiffusionBackend:
model_kind: Optional[str] = None,
) -> dict[str, Any]:
"""Validate, then run the (slow) load on a daemon thread. Returns at once."""
# A blank token (the Studio default when none is configured) must mean
# "anonymous", not an explicit empty credential the Hub rejects with 401.
hf_token = (hf_token.strip() if isinstance(hf_token, str) else hf_token) or None
fam = self.validate_load_request(
repo_id,
gguf_filename = gguf_filename,
@ -649,6 +652,18 @@ class DiffusionBackend:
fraction = min(downloaded / expected, 1.0) if expected > 0 else 0.0
return _progress("downloading", downloaded, expected, fraction)
def loading_repo_ids(self) -> tuple[str, ...]:
"""Repo ids an in-flight background load is downloading (empty when idle).
The delete-cached guard needs this: during a load ``status()["loaded"]`` is
still False, but deleting the target repo (or its companion base) would yank
blobs and snapshot files from under the download/assembly."""
with self._lock:
loading = self._loading
if loading is None or loading.error is not None:
return ()
return tuple(r for r in (loading.repo_id, loading.base_repo) if r)
@staticmethod
def _estimate_download_bytes(
repo_id: str,
@ -760,7 +775,10 @@ class DiffusionBackend:
hf_token = hf_token or None
# Validate first (cheap, no torch/diffusers) so a direct call with a bad
# family fails with ValueError even in a no-diffusers runtime.
# family fails with ValueError even in a no-diffusers runtime. Sanitize the
# token here too (direct callers bypass begin_load): a blank string must
# load anonymously, not 401 as an explicit empty credential.
hf_token = (hf_token.strip() if isinstance(hf_token, str) else hf_token) or None
fam = self.validate_load_request(
repo_id,
gguf_filename = gguf_filename,
@ -784,12 +802,18 @@ class DiffusionBackend:
# The cancel makes that wait ~one step (or the rest of the denoise for a
# pipeline that ignores the step callback).
with self._lock:
# Bail BEFORE signalling any cancel if this load was already superseded (an
# unload/eviction or a newer load bumped the token while we were resolving /
# downloading). Otherwise a stale worker would abort an unrelated, still-live
# generation from the CURRENT model and only then discover it has nothing to do.
if _load_token is not None and _load_token != self._load_token:
raise RuntimeError("Diffusion load was cancelled.")
if self._active_generate_cancel is not None:
self._active_generate_cancel.set()
with self._generate_lock:
with self._lock:
# Bail before the (slow, VRAM-heavy) build if an unload/eviction or a
# newer load superseded this one while we were resolving/downloading.
# Re-check under the generate lock: a newer load/unload may have superseded
# this one while we waited for the in-flight denoise to exit.
if _load_token is not None and _load_token != self._load_token:
raise RuntimeError("Diffusion load was cancelled.")
@ -858,6 +882,12 @@ class DiffusionBackend:
)
pipe = None
transformer_quant_engaged = None
# Drop the exception (and its traceback) BEFORE clearing the cache:
# exc.__traceback__ keeps _load_dense_quant_pipeline's frame -- and
# thus its partially-built dense bf16 transformer/pipe -- alive, so
# clear_gpu_cache() could not otherwise reclaim that VRAM before the
# GGUF build (the OOM-fallback path this cleanup exists for).
del exc
clear_gpu_cache()
if pipe is None:
@ -1435,6 +1465,26 @@ class DiffusionBackend:
raise ValueError(f"Failed to apply LoRA: {exc}") from exc
pipe._unsloth_loras = desired
@staticmethod
def _reset_step_cache(pipe: Any) -> None:
"""Clear the transformer's stateful step cache (FBCache) before a generation.
diffusers keys FBCache residuals by cache context ("cond"/"uncond") on the
long-lived transformer, and neither the pipeline nor the context exit resets
them (``StateManager`` only clears via ``reset_stateful_hooks``, which no
pipeline calls). This backend reuses one resident pipe across generations, so
without a reset the next generation's first step compares its first-block
residual against the PREVIOUS request's -- a tensor-shape mismatch when the
resolution/batch changed, or a stale-cache reuse otherwise. Best-effort: a
transformer without the hook (uncached load) is a silent no-op."""
transformer = getattr(pipe, "transformer", None)
reset = getattr(transformer, "reset_stateful_hooks", None)
if callable(reset):
try:
reset()
except Exception: # noqa: BLE001 — reset is best-effort, never fail a generation
pass
def generate(
self,
*,
@ -1744,6 +1794,13 @@ class DiffusionBackend:
if "callback_on_step_end" in call_params:
kwargs["callback_on_step_end"] = _on_step
# Start each generation from a clean step cache: FBCache residuals from
# a prior request on this resident pipe would otherwise be compared
# against this generation's first step (shape mismatch on a resolution/
# batch change, or stale reuse). No-op when no cache is engaged.
if state.transformer_cache:
self._reset_step_cache(state.pipe)
self._gen = gen
try:
# inference_mode is strictly faster than the no_grad diffusers

6
studio/backend/core/inference/diffusion_memory.py
View file

@ -434,20 +434,20 @@ def apply_memory_plan(
# is in the low-VRAM situation where the decode-time spike can OOM, so turn VAE
# tiling on now (if not already engaged) to cap it.
nonlocal tiling_engaged
pipe.enable_model_cpu_offload()
pipe.enable_model_cpu_offload(device = device)
if not tiling_engaged:
tiling_engaged = _enable_vae_saver(pipe, "enable_vae_tiling", "enable_tiling", logger)
policy = plan.offload_policy
if policy == OFFLOAD_MODEL:
pipe.enable_model_cpu_offload()
pipe.enable_model_cpu_offload(device = device)
elif policy == OFFLOAD_GROUP:
if not _apply_group_offload(pipe, device, logger):
_fallback_to_model_offload()
policy = OFFLOAD_MODEL
elif policy == OFFLOAD_SEQUENTIAL:
try:
pipe.enable_sequential_cpu_offload()
pipe.enable_sequential_cpu_offload(device = device)
except Exception as exc: # noqa: BLE001 — keep the model loadable
if logger is not None:
logger.warning(

9
studio/backend/core/inference/diffusion_prequant.py
View file

@ -193,6 +193,15 @@ def load_prequantized_transformer(
transformer.load_state_dict(state_dict, strict = True, assign = True)
transformer = transformer.to(device)
# Built via from_config (not from_pretrained), so it starts in TRAIN mode; the
# dense and GGUF paths load through from_pretrained, which diffusers documents as
# returning an eval()'d module. Match that here so any train/eval-sensitive layer
# (e.g. dropout) can't make prequant inference nondeterministic or diverge from
# the other load paths.
try:
transformer.eval()
except Exception: # noqa: BLE001 — eval() is best-effort
pass
try: # diagnostic marker, mirrors the runtime-quant path
transformer._unsloth_runtime_quant = scheme
except Exception: # noqa: BLE001 — marker is best-effort

9
studio/backend/core/inference/sd_cpp_args.py
View file

@ -207,8 +207,15 @@ def build_sd_cpp_command(
"""
if not files.diffusion_model:
raise ValueError("diffusion_model path is required")
if not str(params.prompt).strip():
# ``(prompt or "")`` so a None prompt is rejected here rather than slipping past
# ``str(None)`` == "None" (truthy) and landing in argv as a literal "None".
if not (params.prompt or "").strip():
raise ValueError("prompt is required")
# sd-cli inpaint needs the source image too: a --mask with no --init-img is an
# invalid invocation (sd-cli has nothing to inpaint into), so reject it here with a
# clear error instead of emitting a command that fails deep in sd-cli.
if params.mask and not params.init_img:
raise ValueError("init_img is required when mask is set (inpaint needs a source image)")
cmd: list[str] = [binary, "--mode", DEFAULT_MODE, "--diffusion-model", files.diffusion_model]
for flag, value in (

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

@ -457,6 +457,16 @@ class SdCppDiffusionBackend:
fraction = min(downloaded / expected, 1.0) if expected > 0 else 0.0
return _progress("downloading", downloaded, expected, fraction)
def loading_repo_ids(self) -> tuple[str, ...]:
"""Repo ids an in-flight background load is downloading (empty when idle).
Mirrors the diffusers backend so the delete-cached guard can query whichever
engine is active without caring which one it got."""
with self._lock:
loading = self._loading
if loading is None or loading.error is not None:
return ()
return tuple(r for r in (loading.repo_id, loading.base_repo) if r)
# ── Generate ───────────────────────────────────────────────────────────
def generate(

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

@ -1878,8 +1878,11 @@ class DiffusionGenerateRequest(BaseModel):
)
steps: int = Field(9, ge = 1, le = 100, description = "Number of denoising steps")
guidance: float = Field(0.0, ge = 0.0, le = 20.0, description = "Classifier-free guidance scale")
# le = 2**53-1: seeds round-trip through JSON gallery recipes, where JavaScript
# rounds integers above Number.MAX_SAFE_INTEGER -- a restored recipe would then
# generate a different image. Random seeds are already masked to this range.
seed: Optional[int] = Field(
None, ge = 0, le = 2**64 - 1, description = "Seed for reproducibility (random if omitted)"
None, ge = 0, le = 2**53 - 1, description = "Seed for reproducibility (random if omitted)"
)
batch_size: int = Field(
1, ge = 1, le = 32, description = "Images generated in one forward pass (VRAM-heavy)"

24
studio/backend/routes/models.py
View file

@ -27,6 +27,10 @@ class CachedModelRepo(BaseModel):
repo_id: str
size_bytes: int
last_modified: Optional[float] = None
# "text-to-image" for cached diffusers image repos; response_model would silently
# drop the value the handler sets, letting image-only repos pass the chat picker's
# task gate.
task: Optional[str] = None
class CachedModelsResponse(BaseModel):
@ -3399,8 +3403,13 @@ async def delete_cached_model(
# delete guard is otherwise chat-only, so its GGUF could be removed from
# under a live pipeline. Repo-level match, like the chat guards above.
try:
from core.inference.diffusion import get_diffusion_backend
diffusion_status = get_diffusion_backend().status()
# The ACTIVE engine (diffusers or native sd_cpp): on a native selection the
# diffusers singleton reports unloaded while sd-cli still generates from the
# cached GGUF, so checking it alone would let the files be deleted mid-use.
from core.inference.diffusion_engine_router import get_active_diffusion_engine
engine = get_active_diffusion_engine()
diffusion_status = engine.status()
if diffusion_status.get("loaded") and diffusion_status.get("repo_id"):
loaded_id = str(diffusion_status["repo_id"]).lower()
if loaded_id == repo_id.lower() or loaded_id.startswith(repo_id.lower()):
@ -3408,6 +3417,17 @@ async def delete_cached_model(
status_code = 400,
detail = "Unload the model before deleting",
)
# Also refuse while a background image load is DOWNLOADING this repo (or its
# companion base): status().loaded is still False in that window, but deleting
# would remove blobs from under the in-flight download/assembly.
loading_ids = getattr(engine, "loading_repo_ids", tuple)()
for lid in loading_ids:
lid = str(lid).lower()
if lid == repo_id.lower() or lid.startswith(repo_id.lower()):
raise HTTPException(
status_code = 400,
detail = "An Images model load is using this repo; wait for it to finish",
)
except HTTPException:
raise
except Exception:

9
studio/backend/routes/training.py
View file

@ -372,9 +372,14 @@ async def start_training(
# release the arbiter so it doesn't think the gone pipeline owns
# the GPU. Must precede the chat block, which early-returns.
from core.inference import gpu_arbiter
from core.inference.diffusion import get_diffusion_backend
from core.inference.diffusion_engine_router import (
get_active_diffusion_engine,
)
diffusion = get_diffusion_backend()
# The ACTIVE engine, not the diffusers singleton: on a native
# (sd_cpp) selection the diffusers backend reports unloaded while
# the native engine still holds model state / a live generation.
diffusion = get_active_diffusion_engine()
if diffusion.is_loaded:
logger.info(
"Unloading diffusion (Images) model to free GPU memory for training"

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

@ -170,11 +170,13 @@ class _FakePipe:
self.moved_to = device
return self
def enable_model_cpu_offload(self) -> None:
def enable_model_cpu_offload(self, device = None) -> None:
self.offloaded = True
self.offload_device = device
def enable_sequential_cpu_offload(self) -> None:
def enable_sequential_cpu_offload(self, device = None) -> None:
self.sequential_offloaded = True
self.offload_device = device
def enable_vae_tiling(self) -> None:
self.vae_tiled = True
@ -1214,6 +1216,29 @@ def test_unload_cancels_in_flight_load(fake_runtime):
)
def test_superseded_load_does_not_cancel_live_generation(fake_runtime):
# A superseded background load (its token was bumped by a newer load/unload) that
# finally reaches load_pipeline must bail WITHOUT signalling the current model's
# in-flight generation: the token check has to run before the cancel is set, or a
# stale worker aborts an unrelated, still-live denoise.
import threading as _threading
backend = DiffusionBackend()
fam = detect_family("unsloth/Z-Image-Turbo-GGUF")
live_cancel = _threading.Event()
backend._active_generate_cancel = live_cancel # a generation from the CURRENT model
token = 11
backend._load_token = token + 1 # this load has already been superseded
with pytest.raises(RuntimeError, match = "cancelled"):
backend.load_pipeline(
"unsloth/Z-Image-Turbo-GGUF",
gguf_filename = "z-image-turbo-Q4_K_S.gguf",
base_repo = fam.base_repo,
_load_token = token,
)
assert not live_cancel.is_set() # the live generation was left untouched
def test_pick_dtype_bf16_only_on_ampere(fake_runtime, monkeypatch):
# BF16 only on Ampere+ (cc >= 8); pre-Ampere cards must fall back to FP16.
torch = sys.modules["torch"]
@ -1863,3 +1888,41 @@ def test_transformer_quant_skipped_when_plan_offloads(fake_runtime, tmp_path, mo
assert status["transformer_quant"] is None
assert status["offload_policy"] == "model"
assert _FakeTransformer.last["path"] # GGUF path used
def test_reset_step_cache_helper_is_best_effort():
# Calls the transformer's reset hook when present.
calls = []
pipe = types.SimpleNamespace(
transformer = types.SimpleNamespace(reset_stateful_hooks = lambda: calls.append(True))
)
DiffusionBackend._reset_step_cache(pipe)
assert calls == [True]
# No transformer, or a transformer without the hook -> silent no-op (never raises).
DiffusionBackend._reset_step_cache(types.SimpleNamespace())
DiffusionBackend._reset_step_cache(types.SimpleNamespace(transformer = object()))
def test_generate_resets_step_cache_only_when_engaged(fake_runtime, tmp_path):
# FBCache residuals live on the resident transformer across generations, so each
# generate() must reset the stateful cache first -- but only when a cache is engaged.
(tmp_path / "model.gguf").write_bytes(b"weights")
backend = DiffusionBackend()
backend.load_pipeline(
str(tmp_path),
gguf_filename = "model.gguf",
base_repo = "base/repo",
family_override = "z-image",
)
resets = []
backend._state.pipe.transformer = types.SimpleNamespace(
reset_stateful_hooks = lambda: resets.append(True)
)
# No cache engaged (transformer_cache is None) -> reset must NOT run.
backend.generate(prompt = "a sloth")
assert resets == []
# Engage a cache; every subsequent generation resets the stateful cache first.
object.__setattr__(backend._state, "transformer_cache", "fbcache")
backend.generate(prompt = "a sloth")
backend.generate(prompt = "another sloth")
assert resets == [True, True]

22
studio/backend/tests/test_diffusion_memory.py
View file

@ -327,16 +327,19 @@ def test_snapshot_never_raises_on_probe_failure(monkeypatch):
class _RecordingPipe:
def __init__(self) -> None:
self.calls: list[str] = []
self.offload_device = None
def to(self, device):
self.calls.append(f"to:{device}")
return self
def enable_model_cpu_offload(self):
def enable_model_cpu_offload(self, device = None):
self.calls.append("model_offload")
self.offload_device = device
def enable_sequential_cpu_offload(self):
def enable_sequential_cpu_offload(self, device = None):
self.calls.append("sequential_offload")
self.offload_device = device
def enable_vae_tiling(self):
self.calls.append("vae_tiling")
@ -385,6 +388,16 @@ def test_apply_model_offload_engages_offload_and_tiling():
assert "to:cuda" not in pipe.calls # offload owns placement; never both
assert "vae_tiling" in pipe.calls and "vae_slicing" in pipe.calls
assert effective == OFFLOAD_MODEL and tiled is True
assert pipe.offload_device == "cuda" # device threaded to enable_model_cpu_offload
def test_apply_model_offload_passes_target_device():
# enable_model_cpu_offload defaults to CUDA in diffusers; on a non-CUDA accelerator
# (e.g. Intel XPU, which this backend supports) the target device must be forwarded
# or diffusers offloads to the wrong backend and the load fails.
pipe = _RecordingPipe()
apply_memory_plan(pipe, _plan(OFFLOAD_MODEL, tiling = False), device = "xpu")
assert pipe.offload_device == "xpu"
def test_apply_vae_tiling_falls_back_to_vae_submodule():
@ -404,7 +417,7 @@ def test_apply_vae_tiling_falls_back_to_vae_submodule():
def _slice(self):
self.vae.sliced = True
def enable_model_cpu_offload(self):
def enable_model_cpu_offload(self, device = None):
self.offloaded = True
pipe = _VaeOnly()
@ -439,13 +452,14 @@ def test_apply_sequential_offload():
)
assert "sequential_offload" in pipe.calls and "to:cuda" not in pipe.calls
assert effective == OFFLOAD_SEQUENTIAL
assert pipe.offload_device == "cuda" # device threaded to sequential offload too
def test_apply_sequential_falls_back_to_model_offload_when_unsupported():
# Sequential offload is unreliable for GGUF on some diffusers versions; the
# applier must fall back to whole-module offload and report what actually ran.
class _NoSeqPipe(_RecordingPipe):
def enable_sequential_cpu_offload(self):
def enable_sequential_cpu_offload(self, device = None):
raise RuntimeError("sequential offload not supported for this transformer")
pipe = _NoSeqPipe()

13
studio/backend/tests/test_diffusion_prequant.py
View file

@ -66,6 +66,7 @@ class _FakeTransformer:
def __init__(self):
self.assigned = None
self.moved = None
self.eval_called = False
@classmethod
def load_config(cls, base, **kw):
@ -101,6 +102,10 @@ class _FakeTransformer:
self.moved = device
return self
def eval(self):
self.eval_called = True
return self
def _stub_torch_accelerate(
monkeypatch,
@ -182,6 +187,14 @@ def test_load_meta_init_and_assign(monkeypatch, tmp_path):
assert t._unsloth_runtime_quant == "fp8"
def test_load_puts_transformer_in_eval_mode(monkeypatch, tmp_path):
# Built via from_config (not from_pretrained), so the loader must eval() it to match
# the dense/GGUF paths; otherwise train-mode dropout makes inference nondeterministic.
t = _load(monkeypatch, tmp_path, _good_ckpt())
assert t is not None
assert t.eval_called is True
def test_load_missing_file_is_none(monkeypatch, tmp_path):
assert _load(monkeypatch, tmp_path, _good_ckpt(), exists = False) is None

19
studio/backend/tests/test_sd_cpp_args.py
View file

@ -222,6 +222,25 @@ def test_build_inpaint_adds_mask():
assert _pair(cmd, "--init-img") == "/in/src.png"
def test_build_inpaint_mask_without_init_img_rejected():
# sd-cli inpaint needs a source image; a --mask with no --init-img is invalid argv,
# so the builder must reject it up front instead of emitting a doomed command.
files = SdCppModelFiles(diffusion_model = "/m/z.gguf")
params = SdCppGenParams(prompt = "x", mask = "/in/mask.png")
with pytest.raises(ValueError, match = "init_img is required"):
build_sd_cpp_command("/bin/sd-cli", files, params, output_path = "/o.png")
def test_build_rejects_none_prompt():
# A None prompt must be rejected, not coerced to the literal string "None" and
# forwarded into argv.
files = SdCppModelFiles(diffusion_model = "/m/z.gguf")
with pytest.raises(ValueError, match = "prompt is required"):
build_sd_cpp_command(
"/bin/sd-cli", files, SdCppGenParams(prompt = None), output_path = "/o.png"
)
def test_build_edit_repeats_ref_image():
files = SdCppModelFiles(diffusion_model = "/m/flux.gguf")
params = SdCppGenParams(prompt = "add a hat", ref_images = ("/r/a.png", "/r/b.png"))

3
studio/frontend/src/components/app-sidebar.tsx
View file

@ -1203,10 +1203,7 @@ export function AppSidebar() {
icon={PaintBrush02Icon}
label={t("shell.navigation.images")}
active={pathname === "/images" || pathname.startsWith("/images/")}
disabled={chatOnly}
tooltip={trainExportDisabledHint}
onClick={() => {
if (chatOnly) return;
navigate({ to: "/images" });
closeMobileIfOpen();
}}

24
studio/frontend/src/features/images/images-page.tsx
View file

@ -1272,13 +1272,15 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
}, [refreshStatus, dismissLoadToast, pollLoadProgress]);
const handleLoad = useCallback(
// Resolves true when the background load STARTED (callers may revert
// optimistic picker state on false); poll outcomes are handled internally.
async (
repoId: string,
opts: {
kind: "gguf" | "single_file" | "pipeline";
filename?: string;
},
) => {
): Promise<boolean> => {
// Cancel any prior poll loop so two can't run at once.
if (pollTimer.current) clearTimeout(pollTimer.current);
setBusy("loading");
@ -1312,9 +1314,10 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
toast.error(err instanceof Error ? err.message : "Failed to start load");
setBusy(null);
void refreshStatus();
return;
return false;
}
void pollLoadProgress();
return true;
},
[
pollLoadProgress,
@ -1362,13 +1365,19 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
void handleLoad(id, { kind: spec.kind, filename: spec.filename });
return;
}
// GGUF quant pick from the variant expander.
// GGUF quant pick from the variant expander. Optimistic for instant picker
// feedback, but revert if the load fails to START (400/409/network): the
// selector must not advertise a quant that is not the loaded one. Poll-phase
// failures re-sync via refreshStatus.
if (meta.ggufVariant && meta.ggufFilename) {
const prevQuant = quant;
setQuant(meta.ggufVariant);
const dq = defaultsFor(id);
setSteps(dq.steps);
setGuidance(dq.guidance);
void handleLoad(id, { kind: "gguf", filename: meta.ggufFilename });
void handleLoad(id, { kind: "gguf", filename: meta.ggufFilename }).then((started) => {
if (!started) setQuant(prevQuant);
});
return;
}
// A direct local .gguf file picked without a variant isn't wired for Images.
@ -1386,7 +1395,7 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
setGuidance(d.guidance);
void handleLoad(id, { kind: "pipeline" });
},
[busy, handleLoad],
[busy, handleLoad, quant],
);
const handleUnload = useCallback(async () => {
@ -1546,7 +1555,10 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
height: h,
steps,
guidance,
seed: baseSeed + i,
// Offset runs by the batch size: the native engine seeds image j of a
// run at seed+j, so a +1 run offset would regenerate the previous run's
// batch-mates. Unique per image on both engines, reproducible via recipes.
seed: baseSeed + i * batchSize,
batch_size: batchSize,
// Transform/Inpaint/Extend send the source image (+ mask for inpaint/extend) and
// a denoise strength, resolved above. The backend derives output size from the