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

Merge branch 'diffusion-auto-badges' into diffusion-more-families

Browse source 
# Conflicts:
#	studio/backend/core/inference/diffusion_families.py
This commit is contained in:
Daniel Han 2026-07-05 02:17:52 +00:00
commit 5cccb44ca2
19 changed files with 256 additions and 25 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -70,6 +70,7 @@ from .diffusion_speed import (
)
from .diffusion_attention import (
apply_attention_backend,
normalize_attention_backend,
select_attention_backend,
_ensure_attention_backend_installed,
)
@ -83,7 +84,7 @@ from .diffusion_cache import (
maybe_toggle_step_cache,
normalize_transformer_cache,
)
from .diffusion_precision import quantize_text_encoders
from .diffusion_precision import normalize_te_quant, quantize_text_encoders
from .diffusion_prequant import (
load_prequantized_transformer,
resolve_prequant_source,
@ -746,6 +747,15 @@ class DiffusionBackend:
if self._load_token != token:
return
logger.error("diffusion.load_failed: %s", exc)
# Free the debris of a failed construction (e.g. a load-time OOM): _state was
# never committed, and the next load's _unload_locked early-returns on a None
# state, so nothing else releases the reserved VRAM. Guarded: a sticky CUDA
# error makes synchronize() raise, which would skip stamping the REAL error
# below and leave the client polling forever.
try:
clear_gpu_cache()
except Exception: # noqa: BLE001
pass
# Redact native paths: this error is surfaced verbatim via the
# load-progress poll, and Studio can run as a shared server.
from utils.native_path_leases import redact_native_paths
@ -951,6 +961,15 @@ class DiffusionBackend:
model_kind = model_kind,
)
kind = resolve_model_kind(gguf_filename, model_kind)
# Validate every mode string that can raise NOW, before this load evicts the
# previous pipeline below: their first in-line uses all sit past _unload_locked,
# where a bad request would cost the user their working model. Validate-only for
# transformer_quant: the raw value keeps the unset/auto vs explicit-off tri-state.
normalize_transformer_quant(transformer_quant)
normalize_speed_mode(speed_mode)
normalize_attention_backend(attention_backend)
normalize_transformer_cache(transformer_cache)
normalize_te_quant(text_encoder_quant)
# For a full pipeline the repo itself supplies every component, so it is its
# own base; the single-file kinds resolve the companion base diffusers repo.
base = (
@ -1117,7 +1136,12 @@ class DiffusionBackend:
# 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()
# Guarded: after an OOM/sticky CUDA error synchronize() can
# raise, and this fallback path must still reach the GGUF build.
try:
clear_gpu_cache()
except Exception: # noqa: BLE001
pass
if transformer_quant_engaged is not None and quant_plan is not None:
# The re-planned resident placement is the one the engaged dense build
# actually uses; the GGUF-size plan stays in force for the fallback.
@ -2238,7 +2262,8 @@ class DiffusionBackend:
gen = _GenState(total_steps = steps)
def _on_step(pipe, step_index, timestep, callback_kwargs):
now = time.time()
# Monotonic: a wall-clock adjustment (NTP) mid-denoise would skew the ETA.
now = time.monotonic()
gen.step = step_index + 1
if gen.first_step_at == 0.0:
gen.first_step_at = now
@ -2340,9 +2365,8 @@ class DiffusionBackend:
self._cancel_event.set()
with self._lock:
# Abort an in-flight denoise too by setting ITS cancel event, so the step
# callback stops it. unload does NOT take _generate_lock — it must return
# promptly; the running generate keeps its own pipe reference, so freeing
# _state here can't crash it, and its VRAM is reclaimed when it returns
# callback stops it. The running generate keeps its own pipe reference, so
# freeing _state here can't crash it; its VRAM is reclaimed when it exits
# (within ~one step thanks to the cancel).
if self._active_generate_cancel is not None:
self._active_generate_cancel.set()
@ -2351,6 +2375,14 @@ class DiffusionBackend:
# committing) and drop the marker so the next load starts clean.
self._load_token += 1
self._loading = None
# Wait for the signalled denoise to actually exit before reporting unloaded:
# callers treat this return as "VRAM is free" (the GPU arbiter hands the GPU
# to chat next; the training routes size their run against it), and the
# denoise holds its pipe until the next step callback. generate() holds
# _generate_lock for its full body, so a bare acquire is the exit barrier
# (never while holding _lock -- generate takes _lock inside _generate_lock).
with self._generate_lock:
pass
return self.status()
def _unload_locked(self) -> None:
@ -2375,7 +2407,7 @@ class DiffusionBackend:
uninstall_patches()
uninstall_arch_patches()
# NOTE: we deliberately do NOT call state.pipe.unload_lora_weights() here. unload()
# sets the cancel event but does not take _generate_lock, so a LoRA-backed denoise
# only acquires _generate_lock AFTER this teardown, so a LoRA-backed denoise
# can still be running on this same pipe for up to one more callback; mutating its
# adapter layers now would race that in-flight generation. The whole pipe is dropped
# just below (self._state = None; del state; clear_gpu_cache()), so the adapter