From a9e5a806541ccf1474000a7be66383b6db519d2a Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 2 Jul 2026 03:26:46 +0000 Subject: [PATCH 1/2] Address the round of Codex review findings on the merged diffusion phases Memory planning and dense-quant path: size a local diffusers base's resident companions from its on-disk VAE and text-encoder weights instead of folding them to zero, feed the distilled variant hint into the runtime headroom estimate so turbo and schnell models are not over-reserved, place group-offload companions resident before attaching the transformer hooks so a failed placement falls back to whole-module offload instead of crashing, and bail out of the dense transformer download before it starts when the requested quant scheme is unsupported so the load falls back to GGUF cleanly. sd.cpp stack: scrub the native path lease secret from sd-cli child env, redact native load-progress errors, forward the resolved accelerator when auto-installing a forced-native binary, release stale diffusion GPU ownership on CPU-native loads, and remove the sd.cpp install tree on uninstall. Prequant and scripts: reject prequant artifacts missing base_model_id when a base is requested, expanduser before checkpoint existence checks, record and validate the int8 exclusion filter and fp8 fast-accum in checkpoint metadata, make verify_prequant_backend allowlist its local checkpoint and fail on missing or bad LPIPS and on load-peak regressions, average only finite PSNR values in diffusion_quality, and reset the process-wide attention backend between perf probe variants. API and UI: normalize attention_backend casing before Literal validation, close hidden popovers when leaving the Images page, and clear the stale quant label when loading a direct local GGUF file. --- scripts/build_prequant_checkpoint.py | 15 ++++- scripts/diffusion_quality.py | 24 +++++-- scripts/perf_levers_probe.py | 9 +++ scripts/uninstall.ps1 | 5 ++ scripts/uninstall.sh | 4 ++ scripts/verify_prequant_backend.py | 57 ++++++++++++++-- studio/backend/core/inference/diffusion.py | 57 ++++++++++++++-- .../core/inference/diffusion_engine_router.py | 17 ++++- .../core/inference/diffusion_memory.py | 11 +++- .../core/inference/diffusion_prequant.py | 65 +++++++++++++++++-- .../backend/core/inference/sd_cpp_backend.py | 6 +- .../backend/core/inference/sd_cpp_engine.py | 8 ++- studio/backend/models/inference.py | 8 +++ studio/backend/routes/inference.py | 10 ++- .../backend/tests/test_diffusion_backend.py | 47 ++++++++++++++ .../tests/test_diffusion_engine_router.py | 25 +++++++ .../backend/tests/test_diffusion_prequant.py | 59 +++++++++++++++++ .../tests/test_inference_model_validation.py | 28 ++++++++ studio/backend/tests/test_sd_cpp_backend.py | 36 ++++++++++ studio/backend/tests/test_sd_cpp_engine.py | 13 ++++ .../src/features/images/images-page.tsx | 23 ++++++- 21 files changed, 495 insertions(+), 32 deletions(-) diff --git a/scripts/build_prequant_checkpoint.py b/scripts/build_prequant_checkpoint.py index 366822de94..aebc12e983 100644 --- a/scripts/build_prequant_checkpoint.py +++ b/scripts/build_prequant_checkpoint.py @@ -55,8 +55,10 @@ def main(argv = None) -> int: # Reuse the runtime quant factory + filter so offline == runtime (the LPIPS-0 invariant). from core.inference.diffusion_transformer_quant import ( + TQ_FP8, TQ_SCHEMES, _make_quant_config, + _resolve_fast_accum, exclude_tokens_for_scheme, make_filter_fn, ) @@ -83,12 +85,14 @@ def main(argv = None) -> int: # skip the M=1 AdaLN-modulation / conditioning-embedder projections, else the saved checkpoint # bakes them as int8 and crashes (torch._int_mm needs M>16) at the first denoise step on # Flux / Qwen. fp8 / fp4 / mx use scaled_mm (no M limit) -> exclude_tokens_for_scheme returns (). + exclude_name_tokens = exclude_tokens_for_scheme(scheme) + # fp8 bakes the accumulate mode into the saved kernels; record the resolved choice so the + # loader can refuse a checkpoint whose baked value contradicts an explicit runtime request. + fast_accum = _resolve_fast_accum(None) if scheme == TQ_FP8 else None quantize_( transformer, _make_quant_config(scheme), - filter_fn = make_filter_fn( - args.min_features, exclude_name_tokens = exclude_tokens_for_scheme(scheme) - ), + filter_fn = make_filter_fn(args.min_features, exclude_name_tokens = exclude_name_tokens), ) # Move the state dict to CPU for a portable, GPU-free artifact. @@ -103,6 +107,11 @@ def main(argv = None) -> int: "family": fam.name, "scheme": scheme, "min_features": args.min_features, + # The layers skipped for this scheme (int8's M=1 modulation projections; () for + # the scaled_mm schemes) and, for fp8, the baked accumulate mode. Both let the + # loader reject a checkpoint that would not match the runtime path. + "exclude_name_tokens": list(exclude_name_tokens), + "fast_accum": fast_accum, "torch_dtype": args.dtype, "quant_backend": "torchao", "transformer_class": fam.transformer_class, diff --git a/scripts/diffusion_quality.py b/scripts/diffusion_quality.py index 5b241085df..2d0df87e0e 100644 --- a/scripts/diffusion_quality.py +++ b/scripts/diffusion_quality.py @@ -67,6 +67,12 @@ def _to_rgb(path_or_img: Any) -> Any: return np.asarray(img.convert("RGB"), dtype = np.float64) +# Finite PSNR (dB) a perfect (inf) sample is capped to when averaged with imperfect ones, +# so a lossless render counts as excellent without hiding diverged samples. Well above the +# ~37 dB compile and ~21 dB quant noise floors this harness reports. +_PERFECT_MATCH_PSNR = 100.0 + + def psnr(a_img: Any, b_img: Any) -> float: """PSNR (dB) between two images; inf when identical, 0 when shapes differ.""" a, b = _to_rgb(a_img), _to_rgb(b_img) @@ -280,13 +286,19 @@ def _compare( clip_sim.append(clip.image_similarity(img, ref)) def _mean(xs: list[float]) -> Optional[float]: - # Preserve +inf: an identical render (reference vs itself, or a lossless - # quant/offload) scores PSNR=inf, which is exactly the case this harness - # verifies; dropping it as non-finite would print "-" instead of "inf". - if xs and any(x == math.inf for x in xs): + # +inf marks an identical render (reference vs itself, or a lossless quant/offload) + # scoring PSNR=inf -- the case this harness verifies. Report inf ONLY when every + # sample is inf; a mix of inf and finite means some renders diverged, so a bare inf + # would mask those bad samples. Cap the perfect ones to a high finite PSNR and + # average so the drift still shows. (Only PSNR is ever inf; SSIM/CLIP stay finite.) + if not xs: + return None + if all(x == math.inf for x in xs): return math.inf - finite = [x for x in xs if math.isfinite(x)] - return round(sum(finite) / len(finite), 4) if finite else None + vals = [ + _PERFECT_MATCH_PSNR if x == math.inf else x for x in xs if math.isfinite(x) or x == math.inf + ] + return round(sum(vals) / len(vals), 4) if vals else None return { "mean_psnr": _mean(psnrs), diff --git a/scripts/perf_levers_probe.py b/scripts/perf_levers_probe.py index 3964ab46ad..1f33608060 100644 --- a/scripts/perf_levers_probe.py +++ b/scripts/perf_levers_probe.py @@ -148,6 +148,15 @@ def run( del pipe # free the resident pipe so a skipped variant doesn't leak VRAM torch.cuda.empty_cache() return None + else: + # set_attention_backend pins diffusers' PROCESS-WIDE active backend, and a fresh + # transformer's processors (backend None) inherit it. Force native for the no-attn + # variants so they aren't silently measured under a prior variant's kernel (e.g. + # fbcache running with a leftover sage backend). + try: + pipe.transformer.set_attention_backend("native") + except Exception as exc: # noqa: BLE001 — best-effort isolation + print(f" [{tag}] attn(native-reset)={type(exc).__name__}:{str(exc)[:60]}", flush = True) if fbcache is not None: try: from diffusers.hooks import FirstBlockCacheConfig, apply_first_block_cache diff --git a/scripts/uninstall.ps1 b/scripts/uninstall.ps1 index 88defb9ea0..5d60bacf90 100644 --- a/scripts/uninstall.ps1 +++ b/scripts/uninstall.ps1 @@ -335,6 +335,10 @@ function Uninstall-UnslothStudio { # with it). A user-set UNSLOTH_LLAMA_CPP_PATH is left alone. $defaultUnslothHome = if ($env:USERPROFILE) { Join-Path $env:USERPROFILE ".unsloth" } else { $null } $defaultLlamaCpp = if ($defaultUnslothHome) { Join-Path $defaultUnslothHome "llama.cpp" } else { $null } + # Default-mode native diffusion build (install_sd_cpp_prebuilt.default_install_dir()), + # a sibling of studio like llama.cpp. No-op in env/custom mode and when absent. A + # user-set UNSLOTH_SD_CPP_PATH is left alone. + $defaultSdCpp = if ($defaultUnslothHome) { Join-Path $defaultUnslothHome "stable-diffusion.cpp" } else { $null } $defaultCache = if ($defaultUnslothHome) { Join-Path $defaultUnslothHome ".cache" } else { $null } # Isolated Node.js runtime (install_node_prebuilt.py), a sibling of studio in # default mode. No-op in env/custom mode (nested under the custom root) and absent. @@ -384,6 +388,7 @@ function Uninstall-UnslothStudio { # Default-mode shared llama.cpp build + cache (siblings of studio under # ~/.unsloth). No-op in env/custom mode and when absent. if ($defaultLlamaCpp) { _RemovePath $defaultLlamaCpp } + if ($defaultSdCpp) { _RemovePath $defaultSdCpp } if ($defaultCache) { _RemovePath $defaultCache } # Isolated Node.js runtime (sibling of studio under ~/.unsloth). No-op in env/ # custom mode (nested under the custom root, removed with it) and when absent. diff --git a/scripts/uninstall.sh b/scripts/uninstall.sh index 31e851fcbb..f68fd37131 100755 --- a/scripts/uninstall.sh +++ b/scripts/uninstall.sh @@ -216,6 +216,10 @@ _remove_path "$HOME/.unsloth/studio" # by deleting it). No-op in env/custom mode (they nest under the custom root) and # when absent. A user-set UNSLOTH_LLAMA_CPP_PATH is intentionally kept. _remove_path "$HOME/.unsloth/llama.cpp" +# Default-mode native diffusion (stable-diffusion.cpp / sd-cli) build, a sibling of +# studio like llama.cpp (install_sd_cpp_prebuilt.default_install_dir()). No-op in +# env/custom mode and when absent. A user-set UNSLOTH_SD_CPP_PATH is kept. +_remove_path "$HOME/.unsloth/stable-diffusion.cpp" _remove_path "$HOME/.unsloth/.cache" # Isolated Node.js runtime (install_node_prebuilt.py), a sibling of studio in # default mode. No-op in env/custom mode (nested under the custom root) and absent. diff --git a/scripts/verify_prequant_backend.py b/scripts/verify_prequant_backend.py index e160597599..226fb1aaf1 100644 --- a/scripts/verify_prequant_backend.py +++ b/scripts/verify_prequant_backend.py @@ -37,6 +37,14 @@ OUT = Path(os.environ.get("PREQUANT_OUT_DIR", str(_RESEARCH / "prequant_verify_i logging.basicConfig(level = logging.INFO, format = "%(message)s") LOGGER = logging.getLogger("verify_prequant") +# Prequant and runtime produce the SAME quantized weights, so their images must be +# near-identical; anything above this LPIPS means the prequant path diverged and the run +# fails. The prequant load peak must also sit clearly below the dense runtime peak (the +# whole point of the path); require at least this fractional headroom. +LPIPS_MAX = 0.02 +PREQUANT_PEAK_MAX_FRACTION = 0.75 +_RUNTIME_PEAK_FILE = OUT / "runtime_peak.txt" + def _target(dtype): import types @@ -81,7 +89,11 @@ def run(mode, steps, seed, res): sys.path.insert(0, str(BACKEND)) import torch import diffusers - from core.inference.diffusion_prequant import PrequantSource, load_prequantized_transformer + from core.inference.diffusion_prequant import ( + ALLOW_LOCAL_PREQUANT_PATH_ENV, + PrequantSource, + load_prequantized_transformer, + ) from core.inference.diffusion_transformer_quant import quantize_transformer OUT.mkdir(parents = True, exist_ok = True) @@ -90,6 +102,14 @@ def run(mode, steps, seed, res): torch.cuda.empty_cache() if mode == "prequant": + # A local checkpoint is refused unless its directory is allowlisted (unpickling an + # arbitrary file is unsafe). This verifier's CKPT is operator-supplied and trusted, + # so allowlist its directory here or the load returns None and measures nothing. + ckpt_dir = os.path.dirname(os.path.realpath(CKPT)) + existing = os.environ.get(ALLOW_LOCAL_PREQUANT_PATH_ENV, "") + os.environ[ALLOW_LOCAL_PREQUANT_PATH_ENV] = ( + ckpt_dir if not existing else existing + os.pathsep + ckpt_dir + ) source = PrequantSource(kind = "path", location = CKPT, filename = None) transformer = load_prequantized_transformer( transformer_cls, @@ -122,17 +142,44 @@ def run(mode, steps, seed, res): scheme = quantize_transformer(pipe, _target(torch.bfloat16), mode = "fp8", logger = LOGGER) load_peak = torch.cuda.max_memory_allocated() / 1e9 print(f"[runtime] engaged={scheme} load_gpu_peak={load_peak:.1f} GB", flush = True) + # Persist the dense reference peak so a later prequant run can enforce its VRAM win. + _RUNTIME_PEAK_FILE.write_text(f"{load_peak:.6f}") img, dt = _gen(pipe, steps, seed, res) # warmup img, dt = _gen(pipe, steps, seed, res) img.save(OUT / f"{mode}.png") print(f"[{mode}] gen={dt:.3f}s saved {mode}.png", flush = True) + if mode != "prequant": + return 0 + + # Enforce the two invariants this verifier exists to check, so a broken prequant + # checkpoint fails loudly instead of passing just because generation completed. ref_path = OUT / "runtime.png" - if mode == "prequant" and ref_path.exists(): - from PIL import Image - lp = _lpips(np.array(Image.open(ref_path).convert("RGB")), np.array(img)) - print(f"[prequant] LPIPS_vs_runtime={lp}", flush = True) + if not ref_path.exists(): + print("FAIL: runtime reference image missing; run --mode runtime first", flush = True) + return 1 + from PIL import Image + lp = _lpips(np.array(Image.open(ref_path).convert("RGB")), np.array(img)) + print(f"[prequant] LPIPS_vs_runtime={lp}", flush = True) + if lp is None: + print("FAIL: LPIPS could not be computed (install lpips)", flush = True) + return 1 + if lp > LPIPS_MAX: + print(f"FAIL: LPIPS {lp:.4f} > {LPIPS_MAX} (prequant diverged from runtime)", flush = True) + return 1 + if _RUNTIME_PEAK_FILE.exists(): + try: + runtime_peak = float(_RUNTIME_PEAK_FILE.read_text().strip()) + except ValueError: + runtime_peak = 0.0 + if runtime_peak > 0.0 and load_peak > runtime_peak * PREQUANT_PEAK_MAX_FRACTION: + print( + f"FAIL: prequant load peak {load_peak:.1f} GB not below " + f"{PREQUANT_PEAK_MAX_FRACTION:.0%} of dense {runtime_peak:.1f} GB", + flush = True, + ) + return 1 return 0 diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 0defec3fd1..6323f3dd52 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -475,6 +475,34 @@ class DiffusionBackend: return 0 # repo not in cache yet return total + @staticmethod + def _companion_cache_bytes(base: str) -> int: + """Resident companion (VAE + text-encoder) size for the memory plan. + + For a hub base repo this is the cached blob total (``_cache_bytes``). For a + LOCAL diffusers base directory the blob cache is empty, so sum the on-disk + component weights instead, excluding ``transformer/`` (the GGUF supplies the + transformer). Without this a local base folds its multi-GB VAE / text-encoder + weights to zero and auto planning can pick a resident placement that OOMs.""" + local = Path(base).expanduser() + if local.is_dir(): + total = 0 + for f in local.rglob("*"): + if f.suffix.lower() not in (".safetensors", ".bin", ".pt", ".ckpt"): + continue + try: + rel = f.relative_to(local) + except ValueError: + continue + if rel.parts and rel.parts[0] == "transformer": + continue # supplied by the GGUF single-file; not resident here + try: + total += f.stat().st_size + except OSError: + continue + return total + return DiffusionBackend._cache_bytes(base) + # ── Synchronous load / generate / unload ─────────────────────────────── def load_pipeline( @@ -773,7 +801,14 @@ class DiffusionBackend: compile -> placement.""" # 1. Pre-quantized checkpoint, when one is configured for the resolved scheme. scheme = select_transformer_quant_scheme(target, mode) - if scheme is not None and fam is not None: + if scheme is None: + # Bail BEFORE the (multi-GB) dense download: an explicit unsupported scheme + # (e.g. fp8 on Ampere, nvfp4 off Blackwell) would otherwise materialise the + # dense transformer and move the pipe to CUDA only to fail at quantize below -- + # a long finalization under the load lock after the old model was already + # evicted. load_pipeline catches this and builds the GGUF pipeline instead. + raise RuntimeError("transformer quant unsupported for this device/scheme") + if fam is not None: source = resolve_prequant_source(fam, scheme, path_override = prequant_path) if source is not None: transformer = load_prequantized_transformer( @@ -787,6 +822,10 @@ class DiffusionBackend: # Reject a checkpoint built with a different Linear filter than the # dense path uses, so the prequant and runtime-quant models match. min_features = DEFAULT_MIN_LINEAR_FEATURES, + # Only enforced when the caller forces fp8 fast-accum: a checkpoint that + # baked the other choice would ignore the request, so fall to the dense + # path (which applies it) instead of silently using the baked kernels. + fast_accum = fast_accum, logger = logger, ) if transformer is not None: @@ -839,13 +878,23 @@ class DiffusionBackend: device_memory = snapshot_device_memory(target) transformer_resident = estimate_gguf_resident_mib(file_size_mib(gguf_path)) # The companion components (VAE + text encoders) load near their on-disk - # size; sum whatever the prefetch already placed in the base-repo cache. - companion = self._cache_bytes(base) + # size; sum whatever the prefetch placed in the base-repo cache, or -- for a + # LOCAL diffusers base -- the on-disk component weights (the blob cache is + # empty for a local path, which would otherwise fold multi-GB companions to 0 + # and let auto planning pick a resident placement that OOMs). + companion = self._companion_cache_bytes(base) companion_mib = int(companion // (1024 * 1024)) if companion else None model_dense_mib = None if transformer_resident is not None: model_dense_mib = transformer_resident + (companion_mib or 0) - runtime_headroom = estimate_image_runtime_mib(width = None, height = None, family = fam.name) + # Feed the variant hint (gguf filename + base repo) next to the family name so + # estimate_image_runtime_mib sees distilled markers ("turbo"/"schnell") that + # detect_family normalizes out of fam.name -- distilled models need ~15% less + # activation headroom, and over-reserving can force needless offload / tiling. + variant_hint = " ".join( + p for p in (fam.name, Path(gguf_path).name if gguf_path else "", base or "") if p + ) + runtime_headroom = estimate_image_runtime_mib(width = None, height = None, family = variant_hint) return plan_diffusion_memory( target = target, device_memory = device_memory, diff --git a/studio/backend/core/inference/diffusion_engine_router.py b/studio/backend/core/inference/diffusion_engine_router.py index c8db7f81d2..7c1fb22fbf 100644 --- a/studio/backend/core/inference/diffusion_engine_router.py +++ b/studio/backend/core/inference/diffusion_engine_router.py @@ -43,6 +43,18 @@ logger = get_logger(__name__) _DISABLE_TOKENS = frozenset({"0", "off", "false", "no"}) _ENABLE_TOKENS = frozenset({"1", "on", "true", "yes"}) +# Resolved device backend -> the prebuilt sd-cli accelerator to install. Only used +# when a *force-native* load on a GPU host has to install the binary: without this the +# installer defaults to "cpu" and downloads the plain build, so an sd_cpp generation +# forced on a ROCm/Intel box would silently run on CPU. Unknown/GPU-less backends fall +# back to "auto" (the CPU/Metal plain build), matching the installer's own default. +# (install_sd_cpp_prebuilt has no CUDA-Linux asset, so "cuda" only differs on Windows.) +_INSTALL_ACCELERATOR = {"rocm": "rocm", "cuda": "cuda", "xpu": "vulkan"} + + +def _install_accelerator_for(backend: str) -> str: + return _INSTALL_ACCELERATOR.get(backend, "auto") + # The engine the current (or most recent) load committed to, and why a non-native # choice was made. Mutated only under _lock during selection. _lock = threading.Lock() @@ -125,7 +137,10 @@ def select_and_activate_engine(fam: DiffusionFamily, *, hf_token: Optional[str] binary = None if policy_eligible and fam_ok: - binary = ensure_sd_cpp_binary(allow_install = _install_allowed()) + binary = ensure_sd_cpp_binary( + allow_install = _install_allowed(), + accelerator = _install_accelerator_for(backend), + ) # Probe runnability here, before committing the route to native: a present but # non-runnable binary (wrong arch, missing shared libs, no execute bit) would # otherwise pass as available and only fail inside the background load, instead diff --git a/studio/backend/core/inference/diffusion_memory.py b/studio/backend/core/inference/diffusion_memory.py index d3bfb4c175..a39565f1fc 100644 --- a/studio/backend/core/inference/diffusion_memory.py +++ b/studio/backend/core/inference/diffusion_memory.py @@ -502,14 +502,19 @@ def _apply_group_offload(pipe: Any, device: str, logger: Any) -> bool: gkwargs["non_blocking"] = True if "record_stream" in _params: gkwargs["record_stream"] = True - apply_group_offloading(transformer, **gkwargs) - # Place the remaining (smaller) components resident; the streamed - # transformer manages its own placement via the offloading hooks. + # Place the remaining (smaller) components resident BEFORE attaching the + # transformer's group-offload hooks. If a companion .to() OOMs we return False + # with NO hooks installed, so the caller's whole-module offload fallback works: + # diffusers REJECTS enable_model_cpu_offload on a pipeline that already carries + # group-offload hooks, which would otherwise turn the intended fallback into a + # load-time crash. The streamed transformer manages its own placement via the + # offloading hooks applied next. for name, comp in getattr(pipe, "components", {}).items(): if name == "transformer": continue if isinstance(comp, torch.nn.Module): comp.to(onload) + apply_group_offloading(transformer, **gkwargs) return True except Exception as exc: # noqa: BLE001 — fall back to whole-module offload if logger is not None: diff --git a/studio/backend/core/inference/diffusion_prequant.py b/studio/backend/core/inference/diffusion_prequant.py index 2232e726c8..cd7a726dc1 100644 --- a/studio/backend/core/inference/diffusion_prequant.py +++ b/studio/backend/core/inference/diffusion_prequant.py @@ -136,6 +136,7 @@ def load_prequantized_transformer( hf_token: Optional[str] = None, scheme: str, min_features: Optional[int] = None, + fast_accum: Optional[bool] = None, logger: Any = None, ) -> Optional[Any]: """Load the pre-quantized transformer described by ``source`` onto ``device``. @@ -171,7 +172,9 @@ def load_prequantized_transformer( # a torch.save pickle. weights_only=False is required to rebuild those subclasses. # The local-path branch is gated above; the repo branch is a first-party artifact. ckpt = torch.load(path, weights_only = False, map_location = "cpu") - if not _validate_checkpoint(ckpt, scheme, base, logger, min_features = min_features): + if not _validate_checkpoint( + ckpt, scheme, base, logger, min_features = min_features, fast_accum = fast_accum + ): return None state_dict = ckpt["state_dict"] @@ -223,7 +226,11 @@ def _resolve_checkpoint_path(source: PrequantSource, hf_token: Optional[str]) -> """The local file path for ``source``, downloading from the Hub if needed; None if absent.""" if source.kind == "path": import os - return source.location if os.path.isfile(source.location) else None + # Expand ~ once: the allowlist gate (_local_prequant_path_allowed) already + # expands it, so a "~/..." path that passed the gate must be expanded here too + # or os.path.isfile() sees the literal "~" and silently skips a real checkpoint. + expanded = os.path.expanduser(source.location) + return expanded if os.path.isfile(expanded) else None if source.kind == "repo": from huggingface_hub import hf_hub_download return hf_hub_download(repo_id = source.location, filename = source.filename, token = hf_token) @@ -236,13 +243,20 @@ def _validate_checkpoint( base: str, logger: Any, min_features: Optional[int] = None, + fast_accum: Optional[bool] = None, ) -> bool: """Reject a checkpoint that is the wrong format / scheme / base model / filter. ``min_features`` (when given) is the runtime Linear-feature threshold: a checkpoint built with a different ``--min-features`` quantises a different set of Linear layers, so ``load_state_dict(assign=True)`` would silently install a model that does not match - what the dense path produces while status still reports the requested scheme. Reject it.""" + what the dense path produces while status still reports the requested scheme. Reject it. + + ``fast_accum`` (fp8 only) is the runtime accumulate choice: when the caller forces it + explicitly (not None) and the checkpoint recorded a different baked value, the loaded + fp8 kernels would ignore the request while status still reports fp8, so reject and let + the dense path honor it. A checkpoint that predates the metadata (field absent) is + accepted unchanged for backward compatibility.""" if not isinstance(ckpt, dict) or ckpt.get("format") != PREQUANT_FORMAT: _warn(logger, scheme, ValueError("unrecognised pre-quant checkpoint format")) return False @@ -254,9 +268,21 @@ def _validate_checkpoint( _warn(logger, scheme, ValueError(f"checkpoint scheme {meta.get('scheme')!r} != {scheme!r}")) return False ckpt_base = meta.get("base_model_id") - if ckpt_base and base and not _same_base_model(ckpt_base, base): - _warn(logger, scheme, ValueError(f"checkpoint base {ckpt_base!r} != {base!r}")) - return False + if base: + # A checkpoint whose keys happen to match a different base can load strict=True and + # then generate from the wrong weights while status reports the requested scheme. + # Our builder always records base_model_id, so a checkpoint that omits it against a + # requested base is untrustworthy -- refuse rather than silently accept it. + if not ckpt_base: + _warn( + logger, + scheme, + ValueError(f"checkpoint metadata missing base_model_id; refusing for base {base!r}"), + ) + return False + if not _same_base_model(ckpt_base, base): + _warn(logger, scheme, ValueError(f"checkpoint base {ckpt_base!r} != {base!r}")) + return False if min_features is not None: ckpt_min = meta.get("min_features") if ckpt_min is not None and int(ckpt_min) != int(min_features): @@ -266,6 +292,33 @@ def _validate_checkpoint( ValueError(f"checkpoint min_features {ckpt_min!r} != runtime {min_features!r}"), ) return False + # The int8 exclusion set (M=1 modulation / conditioning-embedder projections) is derived + # from the scheme, but a future change to that token list would leave older checkpoints + # with a stale baked set that still passes scheme+min_features and then crashes at the + # first denoise step. When the checkpoint records the set, reject a mismatch; absent + # (older artifact) is accepted since scheme+min_features already pin today's filter. + ckpt_excludes = meta.get("exclude_name_tokens") + if ckpt_excludes is not None: + from .diffusion_transformer_quant import exclude_tokens_for_scheme + + expected = tuple(exclude_tokens_for_scheme(scheme)) + if tuple(ckpt_excludes) != expected: + _warn( + logger, + scheme, + ValueError(f"checkpoint exclude_name_tokens {tuple(ckpt_excludes)!r} != {expected!r}"), + ) + return False + # fp8 fast-accum is baked into the saved kernels; only enforce when the caller forces it. + if fast_accum is not None: + ckpt_fa = meta.get("fast_accum") + if ckpt_fa is not None and bool(ckpt_fa) != bool(fast_accum): + _warn( + logger, + scheme, + ValueError(f"checkpoint fast_accum {ckpt_fa!r} != requested {bool(fast_accum)!r}"), + ) + return False return True diff --git a/studio/backend/core/inference/sd_cpp_backend.py b/studio/backend/core/inference/sd_cpp_backend.py index cdd2b1c543..20e63459d1 100644 --- a/studio/backend/core/inference/sd_cpp_backend.py +++ b/studio/backend/core/inference/sd_cpp_backend.py @@ -366,9 +366,13 @@ class SdCppDiffusionBackend: if self._load_token != _load_token: return logger.error("sd_cpp.load_failed: %s", exc) + # Redact filesystem paths before this reaches /images/load-progress: an + # asset-fetch / local-path / cache-IO failure can embed absolute paths + # (e.g. /home//...), and the diffusers load path scrubs the same way. + from utils.native_path_leases import redact_native_paths with self._lock: if self._load_token == _load_token and self._loading is not None: - self._loading.error = str(exc) + self._loading.error = redact_native_paths(str(exc)) def _asset_specs( self, repo_id: str, gguf_filename: str, fam: DiffusionFamily diff --git a/studio/backend/core/inference/sd_cpp_engine.py b/studio/backend/core/inference/sd_cpp_engine.py index 557f6fcb3e..77c3b457b0 100644 --- a/studio/backend/core/inference/sd_cpp_engine.py +++ b/studio/backend/core/inference/sd_cpp_engine.py @@ -35,6 +35,7 @@ from pathlib import Path from typing import Callable, Optional from utils.process_lifetime import child_popen_kwargs +from utils.native_path_leases import child_env_without_native_path_secret from core.inference.sd_cpp_args import ( SdCppGenParams, SdCppModelFiles, @@ -96,8 +97,13 @@ def runtime_env(binary: str, base_env: Optional[dict[str, str]] = None) -> dict[ next to ``sd-cli``, so prepend the binary's own directory to the platform library path. A locally-built binary that is already linked finds its libs regardless, so this is harmless there. + + Every ``sd-cli`` launch (version probe + generate/upscale) funnels through here, + so this is also the chokepoint that strips the native-path lease secret from the + child env -- the sd-cli binary is an external process that must not be able to + mint/verify native-path grants, matching the other subprocess launchers. """ - env = dict(os.environ if base_env is None else base_env) + env = child_env_without_native_path_secret(os.environ if base_env is None else base_env) var = _lib_path_var() bindir = str(Path(binary).resolve().parent) existing = env.get(var, "") diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index e2e57a4648..11d3e5025c 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -1792,6 +1792,14 @@ class DiffusionLoadRequest(BaseModel): "shifts the residual distribution).", ) + @field_validator("attention_backend", mode = "before") + @classmethod + def _normalize_attention_backend(cls, value): + # The dispatcher accepts case/whitespace variants ("CuDNN", " sage "), but the + # Literal above is validated before any normaliser runs, so fold a string to its + # canonical lower/stripped form here -- otherwise valid casing gets a 422. + return value.strip().lower() if isinstance(value, str) else value + class DiffusionGenerateRequest(BaseModel): """Request to generate one image from the loaded diffusion model.""" diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 4358344be8..8c946484a2 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -11055,7 +11055,7 @@ async def load_diffusion_model( annotate_status, select_and_activate_engine, ) - from core.inference.gpu_arbiter import acquire_for, DIFFUSION + from core.inference.gpu_arbiter import acquire_for, release, DIFFUSION from core.inference.sd_cpp_engine import ENGINE_SD_CPP from utils.native_path_leases import redact_native_paths @@ -11088,6 +11088,14 @@ async def load_diffusion_model( # Then kick the (slow) load onto a background thread and return at once -- # the client polls images/load-progress. await asyncio.to_thread(acquire_for, DIFFUSION) + else: + # A CPU-only native load never touches the GPU, so it neither acquires nor is + # tracked by the arbiter. But switching here FROM a previous diffusers/GPU load + # (select_and_activate_engine unloaded it above) leaves DIFFUSION still marked + # as the arbiter owner; a later chat acquire would then "evict" this CPU model + # for no reason. Release that stale ownership -- release() is owner-guarded, so + # it is a no-op when diffusion never owned the GPU. + await asyncio.to_thread(release, DIFFUSION) status_dict = await asyncio.to_thread( engine.begin_load, request.model_path, diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index bc9b3d8e69..7ea7bb1cc5 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/studio/backend/tests/test_diffusion_backend.py @@ -1161,6 +1161,53 @@ def test_transformer_quant_skipped_when_plan_offloads(fake_runtime, tmp_path, mo assert _FakeTransformer.last["path"] # GGUF path used +def test_transformer_quant_unsupported_scheme_skips_dense_download( + fake_runtime, tmp_path, monkeypatch +): + # An explicit unsupported scheme (select_transformer_quant_scheme -> None) must fail + # the dense path BEFORE materialising the multi-GB dense transformer, then fall back + # to GGUF -- otherwise the download runs under the load lock during finalization + # after the old model was already evicted, only to fail at quantize. + from core.inference import diffusion as dmod + + backend = DiffusionBackend() + _force_cuda_target(backend, monkeypatch) + monkeypatch.setattr(dmod, "dense_transformer_supported", lambda target: True) + monkeypatch.setattr(dmod, "select_transformer_quant_scheme", lambda target, mode: None) + monkeypatch.setattr(dmod, "resolve_prequant_source", lambda fam, scheme, **kw: None) + + @classmethod + def _fp_fail(cls, *a, **k): + pytest.fail("dense transformer must not download when the scheme is unsupported") + + monkeypatch.setattr(_FakeTransformer, "from_pretrained", _fp_fail, raising = False) + (tmp_path / "m.gguf").write_bytes(b"x") + status = backend.load_pipeline( + str(tmp_path), + gguf_filename = "m.gguf", + family_override = "z-image", + transformer_quant = "fp8", + ) + assert status["loaded"] is True + assert status["transformer_quant"] is None # fell back to GGUF + assert _FakeTransformer.last["path"] # GGUF from_single_file used + + +def test_companion_cache_bytes_local_dir_excludes_transformer(tmp_path): + # A LOCAL diffusers base: sum the on-disk VAE / text-encoder weights so auto memory + # planning sees the resident companions, but exclude transformer/ (the GGUF supplies + # it) and non-weight files. A folded-to-zero companion could OOM a resident plan. + (tmp_path / "vae").mkdir() + (tmp_path / "vae" / "diffusion_pytorch_model.safetensors").write_bytes(b"x" * 100) + (tmp_path / "text_encoder").mkdir() + (tmp_path / "text_encoder" / "model.safetensors").write_bytes(b"y" * 50) + (tmp_path / "transformer").mkdir() + (tmp_path / "transformer" / "diffusion_pytorch_model.safetensors").write_bytes(b"z" * 9999) + (tmp_path / "model_index.json").write_bytes(b"{}") # non-weight file, ignored + total = DiffusionBackend._companion_cache_bytes(str(tmp_path)) + assert total == 150 # vae + text_encoder only; transformer/ and json excluded + + def test_reset_step_cache_helper_is_best_effort(): # Calls the transformer's reset hook when present. calls = [] diff --git a/studio/backend/tests/test_diffusion_engine_router.py b/studio/backend/tests/test_diffusion_engine_router.py index 07d495757b..898711e8bd 100644 --- a/studio/backend/tests/test_diffusion_engine_router.py +++ b/studio/backend/tests/test_diffusion_engine_router.py @@ -145,6 +145,31 @@ def test_force_sd_cpp_without_binary_falls_back(monkeypatch): assert _select() == ENGINE_DIFFUSERS +@pytest.mark.parametrize( + "backend, expected", + [("rocm", "rocm"), ("cuda", "cuda"), ("xpu", "vulkan"), ("cpu", "auto"), ("mps", "auto")], +) +def test_install_accelerator_maps_backend(backend, expected): + assert r._install_accelerator_for(backend) == expected + + +def test_force_native_install_uses_gpu_accelerator(monkeypatch): + # Forcing sd_cpp on a ROCm host with no binary must install the ROCm build, not the + # default CPU one -- otherwise the forced-native generation silently runs on CPU. + _set_device(monkeypatch, "rocm") + _set_runnable(monkeypatch) + seen = {} + + def _fake_ensure(**kwargs): + seen.update(kwargs) + return "/usr/bin/sd-cli" + + monkeypatch.setattr(r, "ensure_sd_cpp_binary", _fake_ensure) + monkeypatch.setenv("UNSLOTH_DIFFUSION_ENGINE", "sd_cpp") + assert _select() == ENGINE_SD_CPP + assert seen.get("accelerator") == "rocm" + + # ── active_status annotation ────────────────────────────────────────────────── diff --git a/studio/backend/tests/test_diffusion_prequant.py b/studio/backend/tests/test_diffusion_prequant.py index 2f1a37f8d6..aebbb40af6 100644 --- a/studio/backend/tests/test_diffusion_prequant.py +++ b/studio/backend/tests/test_diffusion_prequant.py @@ -149,6 +149,7 @@ def _load( load_raises = False, exists = True, allow_local = True, + fast_accum = None, ): _FakeTransformer.calls = {} _stub_torch_accelerate(monkeypatch, ckpt, load_raises = load_raises) @@ -171,6 +172,7 @@ def _load( dtype = "bfloat16", hf_token = None, scheme = scheme, + fast_accum = fast_accum, logger = None, ) @@ -218,6 +220,63 @@ def test_load_base_mismatch_is_none(monkeypatch, tmp_path): assert _load(monkeypatch, tmp_path, _good_ckpt(base = "other/model")) is None +def test_load_missing_base_metadata_is_none(monkeypatch, tmp_path): + # A checkpoint whose keys happen to match a different base can load strict=True and then + # render from the wrong weights, so a base was requested but none recorded must be refused. + ckpt = _good_ckpt() + del ckpt["metadata"]["base_model_id"] + assert _load(monkeypatch, tmp_path, ckpt) is None + + +def test_load_fast_accum_mismatch_is_none(monkeypatch, tmp_path): + # fp8 fast-accum is baked into the saved kernels; an explicit request that contradicts + # the recorded value must fall to the dense path (which honors it), not silently use it. + ckpt = _good_ckpt() + ckpt["metadata"]["fast_accum"] = True + assert _load(monkeypatch, tmp_path, ckpt, fast_accum = False) is None + + +def test_load_fast_accum_match_ok(monkeypatch, tmp_path): + ckpt = _good_ckpt() + ckpt["metadata"]["fast_accum"] = True + assert _load(monkeypatch, tmp_path, ckpt, fast_accum = True) is not None + + +def test_load_fast_accum_auto_ignores_baked(monkeypatch, tmp_path): + # An auto (None) request must accept whatever the checkpoint baked, on any GPU class. + ckpt = _good_ckpt() + ckpt["metadata"]["fast_accum"] = True + assert _load(monkeypatch, tmp_path, ckpt, fast_accum = None) is not None + + +def test_load_exclude_tokens_mismatch_is_none(monkeypatch, tmp_path): + # An int8 checkpoint recording a stale exclusion set (would bake M=1 modulation linears + # as int8 and crash) must be rejected rather than loaded. + ckpt = _good_ckpt(scheme = "int8") + ckpt["metadata"]["exclude_name_tokens"] = ["stale_token"] + assert _load(monkeypatch, tmp_path, ckpt, scheme = "int8") is None + + +def test_load_exclude_tokens_match_ok(monkeypatch, tmp_path): + from core.inference.diffusion_transformer_quant import exclude_tokens_for_scheme + + ckpt = _good_ckpt(scheme = "int8") + ckpt["metadata"]["exclude_name_tokens"] = list(exclude_tokens_for_scheme("int8")) + assert _load(monkeypatch, tmp_path, ckpt, scheme = "int8") is not None + + +def test_resolve_checkpoint_path_expands_user(monkeypatch, tmp_path): + # The allowlist gate expands ~, so the existence check must too, or a "~/..." checkpoint + # that passed the gate is silently skipped. + import os + + real = tmp_path / "transformer_fp8.pt" + real.write_bytes(b"x") + monkeypatch.setattr(os.path, "expanduser", lambda p: str(real) if p == "~/ckpt.pt" else p) + source = PrequantSource(kind = "path", location = "~/ckpt.pt", filename = None) + assert pq._resolve_checkpoint_path(source, None) == str(real) + + # ── local-path opt-in gate (RCE guard) ─────────────────────────────────────────── def test_load_local_path_refused_by_default(monkeypatch, tmp_path): # A valid checkpoint at a real file is still refused: torch.load must never run on a diff --git a/studio/backend/tests/test_inference_model_validation.py b/studio/backend/tests/test_inference_model_validation.py index 2427ad35fa..aa6891d4f8 100644 --- a/studio/backend/tests/test_inference_model_validation.py +++ b/studio/backend/tests/test_inference_model_validation.py @@ -220,3 +220,31 @@ def test_walkback_handles_malformed_function_string(): ] ) assert req.messages[-1].tool_call_id == "call_a" + + +# ── DiffusionLoadRequest.attention_backend casing (Literal validated before normalizer) ── +import pytest +from pydantic import ValidationError + +from models.inference import DiffusionLoadRequest + + +def _diff_load(**kw): + return DiffusionLoadRequest(model_path = "repo", gguf_filename = "m.gguf", **kw) + + +def test_attention_backend_casing_and_whitespace_normalized(): + # The dispatcher accepts case/whitespace variants; the before-validator must fold them so + # the lowercase Literal does not 422 an otherwise-valid request. + assert _diff_load(attention_backend = "CuDNN").attention_backend == "cudnn" + assert _diff_load(attention_backend = " sage ").attention_backend == "sage" + + +def test_attention_backend_none_preserved(): + assert _diff_load(attention_backend = None).attention_backend is None + assert _diff_load().attention_backend is None + + +def test_attention_backend_unknown_still_rejected(): + with pytest.raises(ValidationError): + _diff_load(attention_backend = "bogus") diff --git a/studio/backend/tests/test_sd_cpp_backend.py b/studio/backend/tests/test_sd_cpp_backend.py index 84e1b37dd2..46a0514332 100644 --- a/studio/backend/tests/test_sd_cpp_backend.py +++ b/studio/backend/tests/test_sd_cpp_backend.py @@ -319,3 +319,39 @@ def test_run_load_cancels_and_waits_for_inflight_generation(monkeypatch): b._generate_lock.release() assert committed.wait(5) # only now does the commit run assert b._state is not None and b._state.repo_id == "unsloth/Z-Image-Turbo-GGUF" + + +def test_run_load_redacts_paths_in_progress_error(monkeypatch): + # A load failure surfaced via load_progress() must run through redact_native_paths, the + # same scrub the diffusers load path applies, so a registered native path can't leak. + from utils import native_path_leases as npl + + secret_root = "/managed/native/root" + npl._remember_native_path_for_redaction(secret_root, "model dir") + try: + b = SdCppDiffusionBackend(engine = _FakeEngine()) + fam = detect_family("z-image") + monkeypatch.setattr(b, "_asset_specs", lambda *a, **k: []) + monkeypatch.setattr(b, "_set_expected_bytes", lambda *a, **k: None) + + def _boom(*a, **k): + raise RuntimeError(f"failed to read {secret_root}/z.gguf") + + monkeypatch.setattr(b, "_fetch_assets", _boom) + + b._load_token = 1 + b._loading = bk._SdLoading(repo_id = "unsloth/Z-Image-Turbo-GGUF", base_repo = fam.base_repo) + b._run_load( + repo_id = "unsloth/Z-Image-Turbo-GGUF", + gguf_filename = "z.gguf", + base = fam.base_repo, + fam = fam, + hf_token = None, + _load_token = 1, + ) + err = b.load_progress()["error"] + assert err and secret_root not in err and "" in err + finally: + with npl._REDACTION_LOCK: + if secret_root in npl._NATIVE_PATH_REDACTIONS: + npl._NATIVE_PATH_REDACTIONS.remove(secret_root) diff --git a/studio/backend/tests/test_sd_cpp_engine.py b/studio/backend/tests/test_sd_cpp_engine.py index 054574d905..2c10699be1 100644 --- a/studio/backend/tests/test_sd_cpp_engine.py +++ b/studio/backend/tests/test_sd_cpp_engine.py @@ -116,6 +116,19 @@ def test_runtime_env_prepends_binary_dir_to_lib_path(): assert "/existing" in env[var] +def test_runtime_env_scrubs_native_path_lease_secret(monkeypatch): + # The sd-cli child is an external process and must never receive the native-path + # lease secret; every launch (version + generate/upscale) funnels through runtime_env. + monkeypatch.setenv("UNSLOTH_STUDIO_NATIVE_PATH_LEASE_SECRET", "top-secret") + from_os = runtime_env("/opt/sdcpp/bin/sd-cli") + assert "UNSLOTH_STUDIO_NATIVE_PATH_LEASE_SECRET" not in from_os + from_base = runtime_env( + "/opt/sdcpp/bin/sd-cli", + {"UNSLOTH_STUDIO_NATIVE_PATH_LEASE_SECRET": "top-secret"}, + ) + assert "UNSLOTH_STUDIO_NATIVE_PATH_LEASE_SECRET" not in from_base + + def test_runtime_env_handles_missing_lib_path(): var = eng._lib_path_var() env = runtime_env("/opt/sdcpp/bin/sd-cli", {}) diff --git a/studio/frontend/src/features/images/images-page.tsx b/studio/frontend/src/features/images/images-page.tsx index 14174985de..326564c907 100644 --- a/studio/frontend/src/features/images/images-page.tsx +++ b/studio/frontend/src/features/images/images-page.tsx @@ -357,6 +357,10 @@ function RecipePopover({ // Controlled + force-closed off-tab: PopoverContent portals to body, so the // hidden/inert page wrapper can't contain it when the page is kept mounted. const [open, setOpen] = useState(false); + // Also clear the flag when leaving the tab so it does not reopen on return. + useEffect(() => { + if (!active) setOpen(false); + }, [active]); return ( setOpen(active && o)}> @@ -620,6 +624,16 @@ export function ImagesPage({ active = true }: { active?: boolean }) { })(); }, [active, refreshStatus]); + // Collapse the body-ported popovers when leaving the tab. Their open state is + // controlled and force-closed via `active && open` while off-tab, but the + // underlying flag stays set, so returning to /images would otherwise pop them + // back open unprompted. Reset it so the page comes back in a neutral state. + useEffect(() => { + if (active) return; + setSelectorOpen(false); + setAspectOpen(false); + }, [active]); + // Poll load-progress until the background load reaches "ready" or "error", // updating the persistent toast in place each tick. const pollLoadProgress = useCallback(async () => { @@ -758,10 +772,17 @@ export function ImagesPage({ active = true }: { active?: boolean }) { const filename = slash >= 0 ? norm.slice(slash + 1) : norm; const dir = slash >= 0 ? norm.slice(0, slash) : "."; if (!filename.toLowerCase().endsWith(".gguf")) return; + // A direct pick carries no curated variant label; surface the filename so + // the selector stops advertising the previously loaded quant. Optimistic, + // reverted if the load fails to start (mirrors the curated branch above). + const prevQuant = quant; + setQuant(filename); const d = defaultsFor(id); setSteps(d.steps); setGuidance(d.guidance); - void handleLoad(dir, filename); + void handleLoad(dir, filename).then((started) => { + if (!started) setQuant(prevQuant); + }); } }, [busy, handleLoad, quant], From dd792c6312040d5284214dd8f97286e028a159b3 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 2 Jul 2026 03:30:12 +0000 Subject: [PATCH 2/2] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- scripts/diffusion_quality.py | 4 +++- scripts/perf_levers_probe.py | 4 +++- scripts/verify_prequant_backend.py | 1 + .../backend/core/inference/diffusion_engine_router.py | 1 + studio/backend/core/inference/diffusion_prequant.py | 10 +++++++--- studio/backend/core/inference/sd_cpp_backend.py | 1 + 6 files changed, 16 insertions(+), 5 deletions(-) diff --git a/scripts/diffusion_quality.py b/scripts/diffusion_quality.py index 2d0df87e0e..a8d504713d 100644 --- a/scripts/diffusion_quality.py +++ b/scripts/diffusion_quality.py @@ -296,7 +296,9 @@ def _compare( if all(x == math.inf for x in xs): return math.inf vals = [ - _PERFECT_MATCH_PSNR if x == math.inf else x for x in xs if math.isfinite(x) or x == math.inf + _PERFECT_MATCH_PSNR if x == math.inf else x + for x in xs + if math.isfinite(x) or x == math.inf ] return round(sum(vals) / len(vals), 4) if vals else None diff --git a/scripts/perf_levers_probe.py b/scripts/perf_levers_probe.py index 1f33608060..565e28c6b9 100644 --- a/scripts/perf_levers_probe.py +++ b/scripts/perf_levers_probe.py @@ -156,7 +156,9 @@ def run( try: pipe.transformer.set_attention_backend("native") except Exception as exc: # noqa: BLE001 — best-effort isolation - print(f" [{tag}] attn(native-reset)={type(exc).__name__}:{str(exc)[:60]}", flush = True) + print( + f" [{tag}] attn(native-reset)={type(exc).__name__}:{str(exc)[:60]}", flush = True + ) if fbcache is not None: try: from diffusers.hooks import FirstBlockCacheConfig, apply_first_block_cache diff --git a/scripts/verify_prequant_backend.py b/scripts/verify_prequant_backend.py index 226fb1aaf1..66dc89b0c6 100644 --- a/scripts/verify_prequant_backend.py +++ b/scripts/verify_prequant_backend.py @@ -160,6 +160,7 @@ def run(mode, steps, seed, res): print("FAIL: runtime reference image missing; run --mode runtime first", flush = True) return 1 from PIL import Image + lp = _lpips(np.array(Image.open(ref_path).convert("RGB")), np.array(img)) print(f"[prequant] LPIPS_vs_runtime={lp}", flush = True) if lp is None: diff --git a/studio/backend/core/inference/diffusion_engine_router.py b/studio/backend/core/inference/diffusion_engine_router.py index 7c1fb22fbf..0c758501f8 100644 --- a/studio/backend/core/inference/diffusion_engine_router.py +++ b/studio/backend/core/inference/diffusion_engine_router.py @@ -55,6 +55,7 @@ _INSTALL_ACCELERATOR = {"rocm": "rocm", "cuda": "cuda", "xpu": "vulkan"} def _install_accelerator_for(backend: str) -> str: return _INSTALL_ACCELERATOR.get(backend, "auto") + # The engine the current (or most recent) load committed to, and why a non-native # choice was made. Mutated only under _lock during selection. _lock = threading.Lock() diff --git a/studio/backend/core/inference/diffusion_prequant.py b/studio/backend/core/inference/diffusion_prequant.py index cd7a726dc1..088ebc3575 100644 --- a/studio/backend/core/inference/diffusion_prequant.py +++ b/studio/backend/core/inference/diffusion_prequant.py @@ -226,6 +226,7 @@ def _resolve_checkpoint_path(source: PrequantSource, hf_token: Optional[str]) -> """The local file path for ``source``, downloading from the Hub if needed; None if absent.""" if source.kind == "path": import os + # Expand ~ once: the allowlist gate (_local_prequant_path_allowed) already # expands it, so a "~/..." path that passed the gate must be expanded here too # or os.path.isfile() sees the literal "~" and silently skips a real checkpoint. @@ -277,7 +278,9 @@ def _validate_checkpoint( _warn( logger, scheme, - ValueError(f"checkpoint metadata missing base_model_id; refusing for base {base!r}"), + ValueError( + f"checkpoint metadata missing base_model_id; refusing for base {base!r}" + ), ) return False if not _same_base_model(ckpt_base, base): @@ -300,13 +303,14 @@ def _validate_checkpoint( ckpt_excludes = meta.get("exclude_name_tokens") if ckpt_excludes is not None: from .diffusion_transformer_quant import exclude_tokens_for_scheme - expected = tuple(exclude_tokens_for_scheme(scheme)) if tuple(ckpt_excludes) != expected: _warn( logger, scheme, - ValueError(f"checkpoint exclude_name_tokens {tuple(ckpt_excludes)!r} != {expected!r}"), + ValueError( + f"checkpoint exclude_name_tokens {tuple(ckpt_excludes)!r} != {expected!r}" + ), ) return False # fp8 fast-accum is baked into the saved kernels; only enforce when the caller forces it. diff --git a/studio/backend/core/inference/sd_cpp_backend.py b/studio/backend/core/inference/sd_cpp_backend.py index 20e63459d1..3c79b14f2d 100644 --- a/studio/backend/core/inference/sd_cpp_backend.py +++ b/studio/backend/core/inference/sd_cpp_backend.py @@ -370,6 +370,7 @@ class SdCppDiffusionBackend: # asset-fetch / local-path / cache-IO failure can embed absolute paths # (e.g. /home//...), and the diffusers load path scrubs the same way. from utils.native_path_leases import redact_native_paths + with self._lock: if self._load_token == _load_token and self._loading is not None: self._loading.error = redact_native_paths(str(exc))