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

Trim the comments across the diffusion backend

Browse source 
Comment-only pass over the Python this PR touches: drop what the code already
says, collapse multi-line explanations that still read on one line, and keep
the reasoning that is not recoverable from the code. No code, docstring
semantics or behaviour changes; verified with an AST comparison against the
previous revision, and the backend suite is unchanged (same 37 environment
failures as before: the API integration tests that need a live keyed server,
the flash-attn install hooks, and the GPU memory fields).
This commit is contained in:
Daniel Han 2026-07-26 20:30:46 +00:00
commit 36df317293
113 changed files with 3389 additions and 4087 deletions
Download patch file Download diff file Expand all files Collapse all files

34
scripts/build_prequant_checkpoint.py
View file

@ -83,24 +83,15 @@ def main(argv = None) -> int:
args.base, subfolder = "transformer", torch_dtype = torch.bfloat16, token = args.hf_token
).to("cuda")
print(f" quantising in place ({scheme}) ...", flush = True)
# Mirror the runtime path EXACTLY (offline == runtime, LPIPS-0 invariant): for int8 also skip
# the M=1 AdaLN-modulation / conditioning-embedder projections, else the 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 (). Pass the
# family: int8 also carries PER-FAMILY exclusions (Qwen-Image's unpadded text stream runs at
# M = prompt tokens, so a short prompt breaks _int_mm), and the loader validates the baked
# list against exclude_tokens_for_scheme(scheme, metadata["family"]) -- so building with
# family=None both bakes the crashing text-stream linears and yields an artifact the runtime
# then rejects (silently falling back to the dense quantise this script exists to avoid).
# Mirror the runtime exclusions exactly: int8 also skips the M=1 modulation projections
# (torch._int_mm needs M>16) plus per-family ones; scaled_mm schemes skip none. family=None
# bakes the crashing linears and yields an artifact the runtime then rejects.
exclude_name_tokens = exclude_tokens_for_scheme(scheme, fam.name)
# fp8 and mxfp8 assert a bf16 weight, so their filter must skip any non-bf16 Linear the
# transformer keeps: a mixed-precision DiT (Wan / Hunyuan) keeps its _keep_in_fp32_modules in
# fp32 even under torch_dtype=bf16, so quantising one raises inside quantize_ and aborts the
# pass. nvfp4 quantises fp32 fine, so it isn't gated. Runtime quantize_transformer gates on
# scheme membership; mirror it here so the offline checkpoint quantises the same layer set.
# fp8 / mxfp8 assert bf16 weights, so skip any non-bf16 Linear (mixed-precision DiTs keep some
# in fp32); nvfp4 handles fp32. Mirrors the runtime quantize_transformer gate.
require_bf16 = scheme in _REQUIRE_BF16_SCHEMES
# 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.
# fp8 bakes the accumulate mode into the saved kernels; record it so the loader can reject a
# checkpoint contradicting an explicit runtime request.
fast_accum = _resolve_fast_accum(None) if scheme == TQ_FP8 else None
quantize_(
transformer,
@ -112,7 +103,7 @@ def main(argv = None) -> int:
),
)
# Move the state dict to CPU for a portable, GPU-free artifact.
# CPU state dict for a portable, GPU-free artifact.
state_dict = {
k: (v.detach().to("cpu") if hasattr(v, "detach") else v)
for k, v in transformer.state_dict().items()
@ -122,10 +113,8 @@ 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), whether non-bf16 Linears were skipped (the scaled_mm bf16 gate), and,
# for fp8, the baked accumulate mode. All let the loader reject a checkpoint that wouldn't
# match the runtime path.
# Skipped layers, whether non-bf16 Linears were skipped, and the fp8 accumulate mode: all let
# the loader reject a checkpoint that would not match the runtime path.
"exclude_name_tokens": list(exclude_name_tokens),
"require_bf16": require_bf16,
"fast_accum": fast_accum,
@ -136,8 +125,7 @@ def main(argv = None) -> int:
"torchao_version": getattr(torchao, "__version__", "?"),
"diffusers_version": diffusers.__version__,
}
# Record the fp8 granularity so the loader can reject a stale per-tensor checkpoint
# (the runtime now requires per-row; see FP8_GRANULARITY).
# fp8 granularity: lets the loader reject a stale per-tensor checkpoint (runtime needs per-row).
if scheme == TQ_FP8:
metadata["fp8_granularity"] = FP8_GRANULARITY
ckpt = {

11
scripts/build_te_prequant_checkpoint.py
View file

@ -57,8 +57,7 @@ def main(argv = None) -> int:
from core.inference.diffusion_precision import _cast_fp8
from core.inference.diffusion_te_prequant import TE_PREQUANT_FORMAT
# The family is metadata for forensics; detection lives in different modules per
# branch (diffusion_families vs video_families), so resolve best-effort by name.
# Family is forensic metadata; detection differs per branch, so resolve best-effort by name.
family = args.family.strip().lower()
subfolder = args.component if args.config_subfolder is None else args.config_subfolder
@ -70,9 +69,8 @@ def main(argv = None) -> int:
print(f" loading dense encoder from {args.base} (subfolder={subfolder!r}) ...", flush = True)
t0 = time.time()
config = transformers.AutoConfig.from_pretrained(args.base, **from_pretrained_kwargs)
# Prefer the checkpoint's own architecture (what the diffusers pipeline instantiates,
# e.g. Gemma3ForConditionalGeneration); AutoModel.from_config would give the bare base
# class and record a te_class whose state dict the pipeline cannot use.
# Prefer the checkpoint's own architecture; AutoModel.from_config gives the bare base class,
# whose state dict the pipeline cannot use.
arch = (getattr(config, "architectures", None) or [None])[0]
if arch and hasattr(transformers, arch):
encoder_cls_name = arch
@ -104,8 +102,7 @@ def main(argv = None) -> int:
"te_class": encoder_cls_name,
"torch_dtype": args.dtype,
"cast_backend": "diffusers_layerwise",
# str(): torch.__version__ is a TorchVersion object; pickling it into the
# checkpoint makes torch.load(weights_only=True) reject the whole artifact.
# str(): a pickled TorchVersion makes torch.load(weights_only=True) reject the artifact.
"torch_version": str(torch.__version__),
"transformers_version": str(transformers.__version__),
}

12
scripts/diffusion_bench.py
View file

@ -41,10 +41,8 @@ from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Optional
# The backend package lives at unsloth/studio/backend; this file is at
# unsloth/scripts/diffusion_bench.py. Put the backend root on sys.path so
# ``core.inference.diffusion`` imports as the server does. (The backend import is deferred
# into main() so --help never triggers torch.)
# Put the backend root on sys.path so ``core.inference.diffusion`` imports as the server does.
# (The backend import is deferred into main() so --help never triggers torch.)
_BACKEND_ROOT = Path(__file__).resolve().parent.parent / "studio" / "backend"
if str(_BACKEND_ROOT) not in sys.path:
sys.path.insert(0, str(_BACKEND_ROOT))
@ -370,10 +368,8 @@ def _compare(args: argparse.Namespace) -> int:
print(" refusing noisy comparison (pass --force-compare to override).", flush = True)
return 2
# PSNR vs the stored reference image. The baseline stores an absolute reference_png, which
# breaks if the baseline dir was copied/moved, so fall back to reference.png next to the
# baseline JSON. A still-missing reference is a failure below, not a silent pass -- else the
# benchmark reports PASS having done no image comparison.
# PSNR vs the stored reference. reference_png is absolute, so fall back to reference.png beside
# the baseline JSON; a still-missing reference fails below rather than passing silently.
ref_png = Path(baseline.get("accuracy", {}).get("reference_png", ""))
if not ref_png.is_file():
ref_png = baseline_path.parent / "reference.png"

14
scripts/diffusion_quality.py
View file

@ -67,9 +67,8 @@ 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.
# Finite PSNR (dB) cap for perfect (inf) samples, so a lossless render averages as excellent
# without hiding diverged ones. Well above the ~37 dB compile and ~21 dB quant noise floors.
_PERFECT_MATCH_PSNR = 100.0
@ -192,8 +191,7 @@ def _wait_for_load(backend: Any, timeout_s: int = 3600) -> None:
def _hf_file_size_mib(repo: str, filename: str) -> Optional[int]:
# A local model dir / file: stat it directly. The Hub lookup below returns None for
# a local path, which would drop every candidate from _recommend (file_size_mib None).
# Local paths: stat directly, since the Hub lookup returns None and _recommend would drop them.
try:
local = Path(repo).expanduser()
if local.is_dir():
@ -286,10 +284,8 @@ def _compare(
clip_sim.append(clip.image_similarity(img, ref))
def _mean(xs: list[float]) -> Optional[float]:
# +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 means some renders diverged, so a bare inf would mask them. Cap the perfect
# ones to a high finite PSNR and average so the drift still shows. (Only PSNR is ever inf.)
# +inf marks an identical render. Report inf only when every sample is inf; otherwise cap the
# perfect ones to a high finite PSNR so partial drift still shows. (Only PSNR is ever inf.)
if not xs:
return None
if all(x == math.inf for x in xs):

5
scripts/fbcache_flux_probe.py
View file

@ -101,9 +101,8 @@ def run(
from diffusers.hooks import apply_first_block_cache
apply_first_block_cache(pipe.transformer, FirstBlockCacheConfig(threshold = threshold))
if compile_:
# FBCache's per-step decision is a graph break, so a cached run must compile with
# fullgraph=False (mirroring production); fullgraph=True would fail the warmup compile and
# the row would fall back to an eager cached run, producing misleading speedups.
# FBCache's per-step decision is a graph break, so cached runs compile with fullgraph=False as
# production does; fullgraph=True would fail warmup and silently fall back to eager.
fullgraph = threshold is None
try:
pipe.transformer.compile_repeated_blocks(fullgraph = fullgraph, dynamic = True)

5
scripts/fp8_overflow_check.py
View file

@ -63,9 +63,8 @@ def _run(fast_accum, steps, res, seed, mf):
if m > stats["max_abs"]:
stats["max_abs"] = m
# Hook the quantised linears (where an fp8-accumulation overflow would surface). Run EAGER:
# forward hooks don't trace through torch.compile, and the fp8 fast-accum accumulation is
# identical compiled or eager -- compile only changes scheduling.
# Hook the quantised linears where an fp8-accumulation overflow surfaces. Run eager: forward
# hooks don't trace through torch.compile, and accumulation is identical either way.
for m in pipe.transformer.modules():
if isinstance(m, nn.Linear):
m.register_forward_hook(hook)

6
scripts/image_speedmem_bench.py
View file

@ -50,8 +50,7 @@ for _p in (str(_BACKEND_ROOT), str(_REPO_ROOT / "scripts")):
if _p not in sys.path:
sys.path.insert(0, _p)
# Fixed prompt set (the diffusion_quality.py defaults + one photographic subject) so the
# LPIPS mean is not hostage to a single composition.
# Fixed prompt set so the LPIPS mean is not hostage to a single composition.
PROMPTS = [
"A cozy reading nook by a rain-streaked window, warm lamplight, a cat asleep on a stack of books",
"A lone lighthouse on a rocky cliff at sunset, dramatic clouds, crashing waves, highly detailed",
@ -313,8 +312,7 @@ def _generate(
if "callback_on_step_end" in call_params:
kwargs["callback_on_step_end"] = _cb
last.clear()
# Production resets the step cache before every generation; without this the
# first step compares against the PREVIOUS prompt's final residual.
# Reset the step cache like production, else step 1 compares against the previous prompt.
_reset_step_cache(pipe)
_sync()
t0 = time.perf_counter()

4
scripts/int8_linear_probe.py
View file

@ -45,8 +45,8 @@ def main() -> int:
lins = [(n, m) for n, m in model.named_modules() if isinstance(m, torch.nn.Linear)]
selected = [(n, m) for n, m in lins if m.in_features >= MIN and m.out_features >= MIN]
print(f"\n### {label}: {len(lins)} Linear, {len(selected)} pass min_features={MIN}")
# Heuristic: a modulation/embedder Linear is one OUTSIDE the repeated transformer blocks,
# i.e. its fqn does not contain a numeric block index, OR out==k*in (k>=3) AdaLN shape.
# A modulation/embedder Linear sits outside the repeated blocks (no numeric index in its fqn),
# or has an AdaLN out==k*in (k>=3) shape.
sus = []
for n, m in selected:
depth_idx = any(p.isdigit() for p in n.split("."))

4
scripts/nvfp4_t211_probe.py
View file

@ -47,8 +47,8 @@ def diagnostics() -> None:
f"e8m0={hasattr(torch, 'float8_e8m0fnu')} _scaled_mm={hasattr(torch, '_scaled_mm')}",
flush = True,
)
# torchao prints "Skipping import of cpp extensions ..." to stderr at import on torch<2.11.
# On 2.11 that line is absent -> the CUTLASS FP4 GEMM extension is live.
# torchao prints "Skipping import of cpp extensions ..." on torch<2.11; absence of that line
# on 2.11 means the CUTLASS FP4 GEMM extension is live.
print(
" (no 'Skipping import of cpp extensions' line above => cpp/CUTLASS ext loaded)",
flush = True,

15
scripts/perf_levers_probe.py
View file

@ -36,9 +36,8 @@ def _lpips(ref, arr):
import lpips
import torch
# Keep the metric model on CPU: caching it on CUDA leaves it resident across variants, and
# each run resets peak-memory stats, so its VRAM would be charged to (and reduce headroom
# for) every later variant's measurement.
# Keep the metric model on CPU: cached on CUDA it stays resident across variants and its VRAM
# is charged to every later measurement (each run resets peak-memory stats).
if _LP["fn"] is None:
_LP["fn"] = lpips.LPIPS(net = "alex", verbose = False).eval()
@ -72,8 +71,8 @@ def _reset_inductor_flags():
ic.coordinate_descent_tuning = False
ic.coordinate_descent_check_all_directions = False
ic.epilogue_fusion = True
# Reset the int-mm fusion flag too, or it leaks from the inductor_flags variant into
# every later compiled row and the attention/fbcache measurements stop being isolated.
# Reset the int-mm fusion flag too, else it leaks from the inductor_flags variant into every
# later compiled row.
try:
ic.force_fuse_int_mm_with_mul = False
except Exception: # noqa: BLE001
@ -149,10 +148,8 @@ def run(
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 measured under a prior variant's kernel (e.g. fbcache with a leftover sage
# backend).
# set_attention_backend pins diffusers' process-wide backend and fresh processors inherit it,
# so force native for the no-attn variants (else they run under a prior variant's kernel).
try:
pipe.transformer.set_attention_backend("native")
except Exception as exc: # noqa: BLE001 — best-effort isolation

15
scripts/verify_prequant_backend.py
View file

@ -37,10 +37,8 @@ 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 point of the
# path); require at least this fractional headroom.
# Prequant and runtime produce the same quantized weights, so above this LPIPS the prequant
# path diverged. The prequant load peak must also sit this fraction below the dense peak.
LPIPS_MAX = 0.02
PREQUANT_PEAK_MAX_FRACTION = 0.75
_RUNTIME_PEAK_FILE = OUT / "runtime_peak.txt"
@ -102,9 +100,8 @@ 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.
# A local checkpoint is refused unless its directory is allowlisted (unpickling is unsafe).
# CKPT is operator-supplied and trusted, so allowlist it or the load returns None.
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] = (
@ -153,8 +150,8 @@ def run(mode, steps, seed, res):
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.
# Enforce both invariants, so a broken prequant checkpoint fails loudly rather than passing
# just because generation completed.
ref_path = OUT / "runtime.png"
if not ref_path.exists():
print("FAIL: runtime reference image missing; run --mode runtime first", flush = True)

17
scripts/video_quality.py
View file

@ -69,8 +69,7 @@ DEFAULT_PROMPT = (
"splashing water, cinematic, camera tracking sideways"
)
# Finite PSNR (dB) an identical clip is capped to when averaged, matching
# scripts/diffusion_quality.py.
# Finite PSNR (dB) cap for an identical clip, matching scripts/diffusion_quality.py.
_PERFECT_MATCH_PSNR = 100.0
@ -159,11 +158,9 @@ def clip_metrics(
"""All frame metrics for one candidate clip vs the reference clip."""
import numpy as np
# The gate holds the requested shape (num_frames) fixed for reference and candidate, so both
# clips must decode to the same frame count. A shorter candidate is a truncated/corrupt render;
# comparing only the shared prefix would let good early frames mask the missing tail, so the
# mismatch is recorded and gated as FAIL (see verdict()) rather than dropped. No off-by-one is
# tolerated.
# num_frames is fixed for reference and candidate, so both clips must decode to the same frame
# count. A shorter candidate is truncated: comparing the shared prefix would let good early
# frames mask the missing tail, so the mismatch is gated as FAIL (see verdict()).
ref_count, cand_count = len(ref_frames), len(cand_frames)
frame_count_mismatch = ref_count != cand_count
n = min(ref_count, cand_count)
@ -212,8 +209,7 @@ def audio_metrics(ref_audio: Optional[Any], cand_audio: Optional[Any]) -> dict[s
return float(np.sqrt((arr**2).mean())) if arr.size else 0.0
ref_rms, cand_rms = _rms(ref_audio), _rms(cand_audio)
# NaN candidate audio compares False against any threshold, so call it out
# explicitly: a NaN track is a collapse, not a pass.
# NaN compares False against any threshold, so call it out: a NaN track is a collapse.
silent_collapse = (
ref_rms is not None
and ref_rms >= 1e-3
@ -459,8 +455,7 @@ def selftest() -> int:
shifted = clip_metrics(ref, make_clip(offset = 0.5))
check(shifted["ssim_mean"] < same["ssim_mean"], "content shift lowers ssim")
# A truncated render whose surviving prefix is pixel-identical must still FAIL
# on the frame-count mismatch alone, not PASS on the good early frames.
# A truncated render with a pixel-identical prefix must still FAIL on the frame-count mismatch.
truncated = clip_metrics(ref, make_clip()[: n // 2])
check(
truncated["frame_count_mismatch"] is True

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

File diff suppressed because it is too large Load diff

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

@ -185,8 +185,8 @@ def _spec_zimage_forward():
# =====================================================================================
# flux.1: FluxTransformerBlock / FluxSingleTransformerBlock
# (block modulation goes through AdaLayerNormZero, handled by the shared patch; here we fuse the
# inline norm2 modulation + the gated residual adds.)
# (block modulation goes via AdaLayerNormZero; here we fuse the inline norm2 modulation
# and the gated residual adds.)
# =====================================================================================
def _flux_double_forward(
self,
@ -315,8 +315,7 @@ def _spec_flux_single():
# =====================================================================================
# flux.2-klein: Flux2TransformerBlock / Flux2SingleTransformerBlock
# (modulation is INLINE, so we fuse both modulation and gated residuals; scale/shift/gate are
# [B,1,dim] so no [:, None].)
# (modulation is INLINE, so fuse both; scale/shift/gate are [B,1,dim] so no [:, None].)
# =====================================================================================
def _flux2_double_forward(
self,

52
studio/backend/core/inference/diffusion_attention.py
View file

@ -59,10 +59,9 @@ def normalize_attention_backend(value: Optional[str]) -> Optional[str]:
return normalized
# Backends diffusers validates only by package at set time but whose kernels need a specific
# CUDA arch at run time (so an explicit request on the wrong card sets fine then crashes
# mid-generation). Gate by a (min, max-exclusive) capability range: FA3 is Hopper-SM90 only
# (upper bound, so flash3 on a B200 drops to native), FA4 is Blackwell+ (no upper bound).
# Backends diffusers validates only by package at set time but whose kernels need a specific CUDA
# arch at run time. Gate by a (min, max-exclusive) capability range: FA3 is Hopper-SM90 only (so
# flash3 on a B200 drops to native), FA4 is Blackwell+.
_ARCH_CAPABILITY: dict[str, tuple[tuple[int, int], Optional[tuple[int, int]]]] = {
"_flash_3_hub": ((9, 0), (10, 0)), # FlashAttention 3 -> Hopper (SM90) only
"flash_4_hub": ((10, 0), None), # FlashAttention 4 -> Blackwell (SM100)+
@ -117,8 +116,8 @@ def select_attention_backend(
backend = _ALIASES[alias]
if backend == "native":
return None
# AITER is the AMD ROCm kernel: honor it on a ROCm CUDA target, drop it elsewhere (else
# the NVIDIA-only guard below would drop the one backend that only works on ROCm).
# AITER is the AMD ROCm kernel: honor it on a ROCm target, drop it elsewhere (else the
# NVIDIA-only guard below would drop the one backend that only works on ROCm).
if backend == "aiter":
if getattr(target, "device", None) == "cuda" and not _is_cuda_nvidia(target):
return backend
@ -146,9 +145,9 @@ def _cudnn_attention_supported() -> bool:
return have is None or have >= (8, 0)
# Optional-kernel backends installable on demand: dispatcher name -> (probe module, pip
# package). Wheels only (--only-binary=:all:): a source build needs a CUDA toolchain a Studio
# host may lack; no wheel means a native fallback. cuDNN/native ship with torch.
# Optional-kernel backends installable on demand: dispatcher name -> (probe module, pip package).
# Wheels only (--only-binary=:all:): a source build needs a CUDA toolchain a Studio host may
# lack. cuDNN/native ship with torch.
_INSTALLABLE_BACKENDS: dict[str, tuple[str, str]] = {
"sage": ("sageattention", "sageattention"),
"flash": ("flash_attn", "flash-attn"),
@ -162,10 +161,9 @@ _INSTALLABLE_BACKENDS: dict[str, tuple[str, str]] = {
# 0 - never install; a missing kernel falls back to native
_ATTENTION_INSTALL_ENV = "UNSLOTH_DIFFUSION_ATTENTION_INSTALL"
# Packages a pip install was already attempted for in THIS process (success or failure). The
# loader pre-installs outside its locks, then re-resolves under _generate_lock where apply would
# otherwise call pip a SECOND time -- a no-wheel/offline host would re-run the full 600s install
# holding the load lock, blocking unload/cancel. A recorded attempt makes the retry a no-op.
# Packages a pip install was already attempted for in THIS process. The loader pre-installs
# outside its locks, then re-resolves under _generate_lock where apply would otherwise re-run
# the full 600s install holding the load lock; a recorded attempt makes the retry a no-op.
_INSTALL_ATTEMPTED: set[str] = set()
@ -189,8 +187,8 @@ def _ensure_attention_backend_installed(backend: str, logger: Any = None) -> Non
return
except Exception: # noqa: BLE001 — a broken install probes as missing; try the install
pass
# Attempt each install once per process (see _INSTALL_ATTEMPTED): else the in-lock apply path
# re-runs the whole install under _generate_lock and blocks unload/cancel.
# Attempt each install once per process, else the in-lock apply path re-runs the whole install
# under _generate_lock and blocks unload/cancel.
if package in _INSTALL_ATTEMPTED:
return
_INSTALL_ATTEMPTED.add(package)
@ -203,9 +201,8 @@ def _ensure_attention_backend_installed(backend: str, logger: Any = None) -> Non
)
try:
subprocess.run(
# --no-deps: install ONLY this kernel wheel. xformers/flash-attn pin an exact torch,
# so normal resolution would replace the running torch/triton. Without deps an
# ABI-incompatible kernel just fails to import -> native fallback.
# --no-deps: install ONLY this kernel wheel, since xformers/flash-attn pin an exact torch and
# normal resolution would replace the running one. An ABI mismatch just fails to import.
[
sys.executable,
"-m",
@ -220,13 +217,13 @@ def _ensure_attention_backend_installed(backend: str, logger: Any = None) -> Non
timeout = 600,
check = True,
)
# The import system caches directory listings, so the next find_spec can miss the wheel
# just installed (mtime resolution). Invalidate the finder caches so it's picked up now.
# The import system caches directory listings, so invalidate the finder caches or the next
# find_spec can miss the wheel just installed.
importlib.invalidate_caches()
except Exception as exc: # noqa: BLE001 — no wheel / no network -> native fallback
if logger is not None:
# CalledProcessError.str() shows only the exit code; the real reason is in stderr.
# Surface it so the native fallback is diagnosable.
# CalledProcessError.str() shows only the exit code; surface stderr so the fallback is
# diagnosable.
stderr = getattr(exc, "stderr", None)
if stderr:
if isinstance(stderr, bytes):
@ -288,15 +285,14 @@ def apply_attention_backend(
except Exception as exc: # noqa: BLE001 — unavailable kernel -> restore native below
_warn(logger, backend, exc)
if engaged:
# set_attention_backend also pins the backend process-wide. Each DiT's processors now
# keep it locally, so reset the global to native ONCE, else a later unconfigured
# component inherits this kernel.
# set_attention_backend also pins the backend process-wide. Each DiT's processors keep it
# locally, so reset the global to native ONCE, else a later component inherits this kernel.
_reset_global_backend_to_native(logger)
if logger is not None:
logger.info("diffusion.attention: backend=%s", backend)
return backend
# No backend requested, or every set failed: pin native so a stale process-wide backend can't
# leak in. Fresh DiTs follow the global, so one reset via any setter covers them all.
# leak in. Fresh DiTs follow the global, so one reset covers them all.
_restore_native_backend(setters[0], logger)
return None
@ -306,8 +302,8 @@ def _active_attention_backend() -> Optional[str]:
try:
from diffusers.models.attention_dispatch import _AttentionBackendRegistry
# get_active_backend() returns (AttentionBackendName, fn) or None; take element 0 and
# read its .value ("native"), not off the tuple (which never compares equal to a name).
# get_active_backend() returns (AttentionBackendName, fn) or None; read element 0's .value,
# not the tuple (which never compares equal to a name).
active = _AttentionBackendRegistry.get_active_backend()
if active is None:
return None

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

@ -28,8 +28,8 @@ from typing import Any, Optional
_MIB_PER_GB = 1000.0**3 / (1024.0 * 1024.0) # component sizes below are decimal GB
# Steady size of a torchao-quantised transformer relative to bf16: int8/fp8 store one byte per
# param plus per-row scales (~0.52x, plus slack for bf16 norms/embeddings/proj_out); nvfp4 packs
# two params per byte plus block scales. Measured on live int8/fp8 loads.
# param plus per-row scales (~0.52x with slack for bf16 norms/embeddings); nvfp4 packs two per
# byte plus block scales. Measured on live int8/fp8 loads.
_QUANT_STEADY_FACTOR: dict[str, float] = {
"int8": 0.55,
"fp8": 0.55,
@ -37,10 +37,9 @@ _QUANT_STEADY_FACTOR: dict[str, float] = {
"nvfp4": 0.33,
}
# bf16-RESIDENT component sizes in decimal GB: (transformer, text encoders, VAE). What the
# components occupy on device after the dtype cast, NOT the download size (Z-Image ships fp32:
# 24.6 GB of shards -> 12.3 GB bf16). From each base repo's HF sibling metadata, cross-checked
# against the training-side dense_bf16_gb table and measured loads.
# bf16-RESIDENT component sizes in decimal GB: (transformer, text encoders, VAE). What they
# occupy on device after the dtype cast, NOT the download size (Z-Image ships fp32: 24.6 GB of
# shards -> 12.3 GB bf16). From HF sibling metadata, cross-checked against measured loads.
_FAMILY_BF16_GB: dict[str, tuple[float, float, float]] = {
"flux.1": (23.8, 9.8, 0.2),
"flux.1-kontext": (23.8, 9.8, 0.2),
@ -54,18 +53,16 @@ _FAMILY_BF16_GB: dict[str, tuple[float, float, float]] = {
"lumina-2": (5.2, 5.2, 0.2),
# 17B dual-stream DiT (32.5 GB bf16 on disk) + Qwen2.5-VL 15.5 GB + ByT5 0.8 GB.
"hunyuanimage-2.1": (32.5, 16.3, 0.8),
# 17B MoE DiT (34.2 GB bf16) + FOUR text encoders: CLIP-L 0.5 + CLIP-G 2.8 + T5-XXL 9.5
# from the repo, plus the Llama-3.1-8B text_encoder_4 (~16 GB bf16) assembled from the
# open mirror at load time (diffusion_hidream.py).
# 17B MoE DiT (34.2 GB bf16) + FOUR text encoders: CLIP-L 0.5 + CLIP-G 2.8 + T5-XXL 9.5 from
# the repo, plus Llama-3.1-8B text_encoder_4 (~16 GB bf16) from the open mirror at load time.
"hidream-i1": (34.2, 28.8, 0.2),
# Two ~9.3B DiTs (conditional + unconditional_transformer for Ideogram's dual-branch CFG),
# both resident, plus a Qwen3-VL encoder. The vendor stores them as raw float8; these are the
# bf16-resident sizes after the dtype cast, so each doubles (37.2 = 2 x 18.6, encoder 16.3).
# Two ~9.3B DiTs (Ideogram's dual-branch CFG), both resident, plus a Qwen3-VL encoder. The
# vendor stores them as raw float8; these are the bf16-resident sizes, so each doubles.
"ideogram-4": (37.2, 16.3, 0.2),
}
# Base-repo overrides for families offering multiple sizes under one entry (the table carries the
# family default).
# Base-repo overrides for families offering multiple sizes under one entry (the table carries
# the family default).
_BASE_REPO_BF16_GB: dict[str, tuple[float, float, float]] = {
"black-forest-labs/FLUX.2-klein-9B": (18.2, 16.4, 0.2),
}
@ -181,16 +178,14 @@ def resolve_dense_quant_candidate(
if scheme is None:
return None
prequant_available = False
# force_dense: the loader will SKIP the prequant shortcut (e.g. a LoRA bake attaches
# adapters on the dense transformer), so the candidate must be sized for the dense build.
# force_dense: the loader will SKIP the prequant shortcut (e.g. a LoRA bake), so size the
# candidate for the dense build.
if not force_dense:
try:
from .diffusion_prequant import usable_prequant_source
# usable_ (not resolve_): a local path override counts only when the loader will
# accept it (allowlisted AND present), else load_prequantized_transformer refuses it
# and rebuilds dense after the resident pipe is unloaded (the evict-then-OOM this
# prefetch avoids).
# usable_ (not resolve_): a local path override counts only when the loader will accept it
# (allowlisted AND present), else it rebuilds dense after eviction -- the OOM this avoids.
src = usable_prequant_source(
fam, scheme, path_override = prequant_path, base_repo = base_repo
)
@ -211,9 +206,9 @@ def resolve_dense_quant_candidate(
prequant_available,
)
if estimate is not None:
# The dense path may DOWNLOAD the artifact into the HF cache; this must never wedge a
# nearly-full disk. An already-cached re-download is a no-op, so the gate only
# false-positives on an already-critically-full disk, where the GGUF fallback is right anyway.
# The dense path may DOWNLOAD the artifact into the HF cache, which must never wedge a nearly
# full disk. A cached re-download is a no-op, so this only trips on an already-critical disk,
# where the GGUF fallback is right anyway.
needed_mib = (
estimate.steady_transformer_mib
if estimate.prequant

8
studio/backend/core/inference/diffusion_batched.py
View file

@ -30,12 +30,12 @@ from __future__ import annotations
from typing import Any, Callable, Optional
# Upper bound on images per generation call (mirrors the route's batch_size cap):
# a prompt/seed list beyond this is a client error, not an OOM to back off from.
# Upper bound on images per generation call (mirrors the route's cap): a longer prompt/seed
# list is a client error, not an OOM to back off from.
MAX_BATCH_IMAGES = 32
# Seeds stay in JS's safe-integer range so they round-trip through the JSON
# gallery recipes and reproduce the image (a raw 64-bit seed loses precision).
# Seeds stay in JS's safe-integer range so they round-trip through the JSON gallery recipes
# (a raw 64-bit seed loses precision).
SEED_MASK = (1 << 53) - 1

34
studio/backend/core/inference/diffusion_cache.py
View file

@ -28,13 +28,12 @@ TC_FBCACHE = "fbcache"
TC_MODES = (TC_FBCACHE,)
# FBCache residual thresholds: higher skips more steps (faster, lower quality). Quantised
# transformers shift the residual distribution, so they need a higher threshold to trigger at
# all (per ParaAttention's fp8 guidance).
# transformers shift the residual distribution, so they need a higher threshold to trigger.
DEFAULT_FBCACHE_THRESHOLD = 0.08
QUANT_FBCACHE_THRESHOLD = 0.12
# Auto step-count bar: FBCache's win scales with step count, so auto engages only at 20+ steps
# ("dev" schedules 28+ qualify, distilled turbo 4-9 never do).
# ("dev" schedules qualify, distilled turbo never does).
FBCACHE_MIN_STEPS = 20
@ -119,9 +118,8 @@ def _compile_hooked_block_inners(transformer: Any, logger: Any = None) -> int:
continue
if getattr(orig, "__self__", None) is None:
continue # not a plain bound method; arming would miss the block
# fullgraph=False / dynamic=True: a cache is active (its decision graph-breaks) and
# this matches the default tier. Dynamo caches per code object, so re-arming after
# a toggle is ~free (~0.03 s).
# fullgraph=False / dynamic=True: a cache is active (its decision graph-breaks) and this matches
# the default tier. Dynamo caches per code object, so re-arming after a toggle is ~free.
fn_ref.original_forward = torch.compile(orig, fullgraph = False, dynamic = True)
hook._unsloth_orig_inner = orig
armed += 1
@ -204,16 +202,15 @@ def apply_step_cache(
else (QUANT_FBCACHE_THRESHOLD if quant_active else DEFAULT_FBCACHE_THRESHOLD)
)
# Engage only via the native enable_cache (CacheMixin path): the lower-level
# apply_first_block_cache hook would also install on a non-CacheMixin transformer (e.g.
# Z-Image) whose pipeline opens no cache_context and crashes generation. So a model without
# enable_cache runs uncached instead of being reported cached then failing.
# apply_first_block_cache hook would also install on a non-CacheMixin transformer whose pipeline
# opens no cache_context and crashes generation. Such a model runs uncached instead.
enable_cache = getattr(transformer, "enable_cache", None)
if not callable(enable_cache):
_warn(logger, mode, RuntimeError("transformer has no cache_context (not a CacheMixin)"))
return None
# A CacheMixin transformer is necessary but not sufficient: the hook raises "No context is set"
# unless the PIPELINE wraps its denoise loop in cache_context(...). Flux Kontext / img2img /
# inpaint / controlnet reuse FluxTransformer2DModel yet open none, so run uncached instead.
# inpaint / controlnet reuse FluxTransformer2DModel yet open none, so run uncached.
if not _pipeline_opens_cache_context(pipe):
_warn(
logger, mode, RuntimeError("pipeline __call__ opens no cache_context; running uncached")
@ -227,12 +224,11 @@ def apply_step_cache(
config = FirstBlockCacheConfig(threshold = thr)
enable_cache(config)
# enable_cache after the pipe ran leaves a stale cached child-registry list; the new block
# enable_cache after the pipe ran leaves a stale cached child-registry list, so the new block
# hooks would never receive the cache context. Must follow every enable_cache.
_invalidate_child_registry_cache(transformer)
# If blocks are already regionally compiled (toggle path: compile ran at load), re-point
# the fresh hooks' compute branch at compiled inners; the load path is armed by
# _compile_repeated_blocks. No-op when nothing is compiled.
# If blocks are already regionally compiled (toggle path), re-point the fresh hooks' compute
# branch at compiled inners; the load path is armed by _compile_repeated_blocks.
_compile_hooked_block_inners(transformer, logger)
try:
transformer._unsloth_step_cache = f"{mode}@{thr}"
@ -242,9 +238,9 @@ def apply_step_cache(
logger.info("diffusion.cache: %s engaged (threshold=%s)", mode, thr)
return mode
except Exception as exc: # noqa: BLE001 — incompatible model -> run uncached
# enable_cache can fail after hooking some blocks; drop partial hooks so the
# reported-uncached model isn't half-cached. Restore armed compiled inners FIRST
# (remove_hook splices original_forward back into module.forward).
# enable_cache can fail after hooking some blocks; drop partial hooks so the reported-uncached
# model isn't half-cached. Restore armed compiled inners FIRST (remove_hook splices
# original_forward back into module.forward).
_restore_hooked_block_inners(transformer)
try:
transformer.disable_cache()
@ -317,8 +313,8 @@ def maybe_toggle_step_cache(
disable_cache = getattr(transformer, "disable_cache", None)
if callable(disable_cache):
try:
# Restore before remove_hook splices original_forward back, so compiled
# wrappers don't leak onto the uncached path.
# Restore before remove_hook splices original_forward back, so compiled wrappers don't leak
# onto the uncached path.
_restore_hooked_block_inners(transformer)
disable_cache()
transformer._unsloth_step_cache = None

16
studio/backend/core/inference/diffusion_compile_cache.py
View file

@ -43,10 +43,9 @@ from typing import Any, Optional
# --------------------------------------------------------------------------------- env knobs
# UNSLOTH_DIFFUSION_COMPILE_CACHE: auto (default) | 0 | 1
# auto -> load a matching bundle AND save one after the first compiled generation.
# Measured on Qwen-Image (B200, deferred 3rd-gen engage, FBCache armed): compile
# hitch drops 29.1 -> 22.2 s, bit-identical output, 7.9 MB bundle, ~0.5 s save.
# Residual warmup is dynamo tracing + guards, which Mega-cache does not capture.
# auto -> load a matching bundle AND save one after the first compiled generation. Measured on
# Qwen-Image (B200): compile hitch 29.1 -> 22.2 s, bit-identical, 7.9 MB bundle. The
# residual warmup is dynamo tracing + guards, which Mega-cache does not capture.
# 1 -> same as auto, also re-saves on a hit (distributor refresh).
# 0 -> disabled (plain local compile, no cache dir override).
# UNSLOTH_DIFFUSION_COMPILE_CACHE_DIR: root dir for bundles (default under the workspace).
@ -77,8 +76,8 @@ def _save_enabled(mode: str) -> bool:
return False
if mode == "on":
return True
# auto: save by default (without a saved bundle no user gets a warm restart). SAVE
# env overrides: "0" -> load-only, "1" -> keep on.
# auto: save by default (without a saved bundle no user gets a warm restart). SAVE env
# overrides: "0" -> load-only, "1" -> keep on.
return (os.environ.get(_ENV_SAVE) or "").strip().lower() not in ("0", "off", "false", "no")
@ -256,9 +255,8 @@ def begin(
if ctx.bundle.exists() and ctx.manifest_path.exists():
ctx.hit = _try_load(ctx, logger)
if ctx.hit and mode != "on":
# Loaded artifacts == on-disk artifacts, so nothing to save (~0.5 s for no
# change). A new static-compile shape re-dirties via register_shape; mode
# "on" (distributor refresh) keeps saving.
# Loaded artifacts == on-disk artifacts, so nothing to save. A new static-compile shape
# re-dirties via register_shape; mode "on" (distributor refresh) keeps saving.
ctx.saved = True
else:
_info(logger, f"compile-cache: no bundle for key {key} (will compile locally)")

25
studio/backend/core/inference/diffusion_cond_cache.py
View file

@ -40,8 +40,8 @@ from typing import Any, Optional
_ENV_DIR = "UNSLOTH_DIFFUSION_COND_CACHE_DIR"
# Bound arguments that never change the returned embeddings: the target device
# is a placement detail (the hit is moved there) and no encode path draws RNG.
# Bound arguments that never change the returned embeddings: the target device is a placement
# detail (the hit is moved there) and no encode path draws RNG.
_KEY_EXCLUDED_ARGS = frozenset({"device", "generator"})
@ -60,9 +60,8 @@ def _json_safe(value: Any) -> bool:
return False
# Per-slot layout codes for _flatten/_unflatten (stored as the leading int64 tensor):
# -1 = None slot, -2 = a bare tensor, n >= 0 = a LIST of n tensors (Z-Image returns
# its per-prompt embeddings as a list, so one nesting level must round-trip).
# Per-slot layout codes for _flatten/_unflatten (the leading int64 tensor): -1 = None slot,
# -2 = a bare tensor, n >= 0 = a LIST of n tensors (Z-Image returns per-prompt lists).
_SLOT_NONE = -1
_SLOT_TENSOR = -2
@ -134,8 +133,7 @@ def install(
if not callable(encode):
return False
try:
# Lazy: core.training imports parts of core.inference, so the module-level
# import would be circular; the extras module itself is stdlib-only.
# Lazy: core.training imports parts of core.inference, so a module-level import would be circular.
from core.training.diffusion_train_extras import PersistentConditioningCache
signature = inspect.signature(encode)
cache = PersistentConditioningCache(root, family, 0)
@ -144,14 +142,11 @@ def install(
logger.warning("diffusion.cond_cache: install failed: %s", exc)
return False
# Everything beyond the call arguments that changes the embedding numerics. A GGUF /
# single-file checkpoint takes its TEXT ENCODERS from the companion base repo, so the
# base identity must key the cache too: the same checkpoint reloaded against a different
# base would otherwise hit entries encoded by the previous base's encoder. Defaults to
# the checkpoint itself (a full pipeline is its own base).
# The identifiers alone are not versions: both are paired with a revision marker so a
# Hub repo advancing to a new commit, or a local directory updated in place, misses
# instead of returning embeddings from the previous text encoder.
# Everything beyond the call arguments that changes the embedding numerics. A GGUF/single-file
# checkpoint takes its TEXT ENCODERS from the companion base repo, so the base identity keys the
# cache too, else the same checkpoint against a different base would hit the previous base's
# entries. Both identifiers are paired with a revision marker, so a Hub repo advancing a commit
# or a local directory updated in place misses instead of reusing the old encoder.
base_ref = base_repo if base_repo else repo_id
load_fp = {
"repo": str(repo_id),

24
studio/backend/core/inference/diffusion_controlnet.py
View file

@ -27,8 +27,8 @@ from utils.paths.storage_roots import studio_root
# edge map here (no heavy detector dependency).
CONTROL_TYPES = ("passthrough", "canny")
# Diffusers quant schemes that cannot host ControlNet cleanly (torchao tensor-subclass weights);
# gated off like LoRA, along with GGUF-via-diffusers.
# Diffusers quant schemes that cannot host ControlNet cleanly (torchao tensor subclasses); gated
# off like LoRA, along with GGUF-via-diffusers.
_DIFFUSERS_BLOCKED_QUANT = ("int8", "fp8", "nvfp4", "mxfp8")
@ -166,11 +166,10 @@ def resolve_controlnet(spec_id: str, *, family: Optional[str] = None) -> Resolve
entry = _catalog_by_id().get(spec_id)
if entry is None:
# A curated entry named by its full repo id must still hit the family gate below, not slip
# through the bare-repo fallback and load through the wrong family's class.
# through the bare-repo fallback into the wrong family's class.
entry = next((e for e in _CURATED if e.repo_id and e.repo_id == spec_id), None)
if entry is not None:
# A direct API call could bypass the UI filter and send an entry for another family; reject
# it before any download so it never reaches the wrong pipeline.
# A direct API call could send an entry for another family; reject it before any download.
fam = (family or "").strip().lower()
if entry.families and fam and fam not in {f.lower() for f in entry.families}:
raise ValueError(
@ -186,8 +185,8 @@ def resolve_controlnet(spec_id: str, *, family: Optional[str] = None) -> Resolve
raise ValueError(f"ControlNet '{spec_id}' has no repo")
return ResolvedControlNet(spec_id, entry.repo_id, is_local = False)
# A bare HF repo id (owner/name). STRICT shape (one slash, alphanumeric-leading segments) so a
# filesystem-looking id can never reach from_pretrained and bypass the no-raw-path contract.
# A bare HF repo id (owner/name). STRICT shape so a filesystem-looking id can never reach
# from_pretrained and bypass the no-raw-path contract.
if re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9_.-]*/[A-Za-z0-9][A-Za-z0-9_.-]*", spec_id):
return ResolvedControlNet(spec_id, spec_id, is_local = False)
@ -196,8 +195,7 @@ def resolve_controlnet(spec_id: str, *, family: Optional[str] = None) -> Resolve
)
# Union ControlNet mode indices: a union model selects the active mode via an integer
# ``control_mode``. Standard indices for the FLUX.1 / Qwen-Image union ControlNets.
# Union ControlNet mode indices: a union model selects its head via an integer ``control_mode``.
_UNION_CONTROL_MODES: dict[str, int] = {
"canny": 0,
"tile": 1,
@ -217,8 +215,8 @@ def union_control_mode(spec_id: str, control_type: str) -> Optional[int]:
returns a 400 instead of running the wrong head. A non-union entry returns None."""
entry = _catalog_by_id().get(spec_id)
if entry is None:
# A union model may be named by its bare repo id; the catalog is keyed by short id, so
# match on repo_id too, else the mode is dropped and the union runs the wrong head.
# A union model may be named by its bare repo id, so match on repo_id too (the catalog is keyed
# by short id), else the mode is dropped and the union runs the wrong head.
entry = next((e for e in _CURATED if e.repo_id and e.repo_id == spec_id), None)
if entry is None or not entry.is_union:
return None
@ -250,8 +248,8 @@ def preprocess_control(image: Any, control_type: str) -> Any:
mag = np.hypot(gx, gy)
peak = float(mag.max())
if peak <= 1e-6:
# Flat image -> no edges, which is an all-black map. Returning the source would
# instead condition the ControlNet on its raw luminance.
# A flat image has no edges, so the map is all black. Returning the source would instead
# condition the ControlNet on its raw luminance.
return Image.new("RGB", image.size, (0, 0, 0))
mag = mag / peak * 255.0
edges = (mag > 40.0).astype(np.uint8) * 255 # white edges on black (ControlNet convention)

18
studio/backend/core/inference/diffusion_device.py
View file

@ -151,16 +151,15 @@ def diffusion_device_target_from_torch_device(
def _cuda_or_rocm_target(torch: Any, *, is_rocm: bool) -> DiffusionDeviceTarget:
if is_rocm:
# ROCm lacks NVIDIA's pre-Ampere bf16-emulation quirk, so is_bf16_supported() is
# trustworthy; bf16 only when it proves it.
# ROCm lacks NVIDIA's pre-Ampere bf16-emulation quirk, so is_bf16_supported() is trustworthy.
try:
bf16_ok = bool(torch.cuda.is_bf16_supported())
except Exception:
bf16_ok = False
dtype = torch.bfloat16 if bf16_ok else torch.float16
else:
# NVIDIA: bf16 needs Ampere+ (major >= 8), by capability NOT is_bf16_supported()
# (pre-Ampere cards emulate bf16 slowly but report it supported; the #6658 fix).
# NVIDIA: bf16 needs Ampere+ (major >= 8), by capability NOT is_bf16_supported() (pre-Ampere
# cards emulate bf16 slowly but report it supported; the #6658 fix).
try:
major = torch.cuda.get_device_capability()[0]
except Exception:
@ -219,12 +218,12 @@ def _mps_or_cpu_target(torch: Any) -> DiffusionDeviceTarget:
if mps_available:
# torch reads PYTORCH_MPS_HIGH_WATERMARK_RATIO once, at the first MPS allocation (the probe
# below), so relax it before that or the allocator caps at ~1.7x recommendedMaxWorkingSet
# and can OOM a model that would fit in unified RAM. setdefault respects an override.
# below), so relax it first or the allocator caps at ~1.7x recommendedMaxWorkingSet and can OOM
# a model that would fit. setdefault respects an override.
os.environ.setdefault("PYTORCH_MPS_HIGH_WATERMARK_RATIO", "0.0")
# Prefer bfloat16, else float32, NEVER silent float16: modern DiTs produce activations far
# outside fp16's range (Z-Image MLP peaks near 9e5 -> inf -> NaN -> black image). bf16
# (macOS 14+) shares fp32's exponent range; on older macOS float32 keeps output correct.
# outside fp16's range (Z-Image MLP peaks near 9e5 -> inf -> NaN -> black image). bf16 (macOS
# 14+) shares fp32's exponent range; on older macOS float32 keeps output correct.
dtype = torch.bfloat16 if _mps_supports_bfloat16(torch) else torch.float32
return DiffusionDeviceTarget(
device = "mps",
@ -239,8 +238,7 @@ def _mps_or_cpu_target(torch: Any) -> DiffusionDeviceTarget:
def _cpu_target(torch: Any, dtype: Any = None) -> DiffusionDeviceTarget:
# torch is None on the no-torch CPU fallback; leave dtype=None then rather than crash on
# torch.float32.
# torch is None on the no-torch CPU fallback; leave dtype=None rather than crash.
if dtype is None and torch is not None:
dtype = torch.float32
return DiffusionDeviceTarget(

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

@ -110,8 +110,7 @@ def _rmsnorm_forward(self, hidden_states):
# Fall back to the exact original where F.rms_norm is NOT equivalent to diffusers:
# * NPU / bias / fp32-weight -> special handling / an fp32 quirk;
# * tuple `dim` -> diffusers reduces only the LAST dim, F.rms_norm reduces every dim;
# * dtype mismatch -> diffusers computes variance in fp32 from the ORIGINAL tensor, so
# casting first would change the variance.
# * dtype mismatch -> diffusers computes variance in fp32 from the ORIGINAL tensor.
if _NPU or self.bias is not None or _orig_rmsnorm_forward is None or len(tuple(self.dim)) != 1:
return _orig_rmsnorm_forward(self, hidden_states) # type: ignore[misc]
weight = self.weight
@ -124,8 +123,8 @@ def _rmsnorm_forward(self, hidden_states):
return _orig_rmsnorm_forward(self, hidden_states) # type: ignore[misc]
# Install / uninstall via the shared patch backend: the live original is fingerprinted
# (can_safely_patch, relaxed) so a changed forward is left UNPATCHED, and stashed for exact restore.
# Install / uninstall via the shared patch backend: the live original is fingerprinted so a
# changed forward is left UNPATCHED, and stashed for exact restore.
def _specs():
# (class, patched_fn)
return [
@ -154,8 +153,7 @@ def install_compile_safe_patches() -> int:
for cls, new_fn in _specs():
if cls is None:
continue
# torch < 2.4 has no F.rms_norm: leave the original in place, don't install a patch whose
# fast path would AttributeError.
# torch < 2.4 has no F.rms_norm: leave the original rather than install an AttributeError.
if cls is _RMSNorm and not hasattr(F, "rms_norm"):
logger.info("eager-patch: skipping RMSNorm (this torch has no F.rms_norm)")
continue

32
studio/backend/core/inference/diffusion_engine_router.py
View file

@ -47,8 +47,7 @@ _ENABLE_TOKENS = frozenset({"1", "on", "true", "yes"})
# Resolved device backend -> the prebuilt sd-cli accelerator to install. Used only for a
# force-native load on a GPU host: without it the installer defaults to "cpu" and a forced
# ROCm/Intel sd_cpp generation would silently run on CPU. Unknown backends -> "auto" (CPU/Metal),
# matching the installer default. (No CUDA-Linux asset, so "cuda" only differs on Windows.)
# ROCm/Intel generation would silently run on CPU. Unknown backends -> "auto" (CPU/Metal).
_INSTALL_ACCELERATOR = {"rocm": "rocm", "cuda": "cuda", "xpu": "vulkan"}
@ -59,8 +58,8 @@ def _install_accelerator_for(backend: str) -> str:
# The engine the current load committed to, and why a non-native choice was made. Mutated only
# under _lock during selection.
_lock = threading.Lock()
# Serializes a whole engine switch (check -> unload -> publish); _lock alone is released during the
# slow unload(), letting two overlapping selections load onto the engine the other is unloading.
# Serializes a whole engine switch (check -> unload -> publish); _lock alone is released during
# the slow unload(), letting two selections load onto the engine the other is unloading.
_transition_lock = threading.Lock()
_active_engine_name: str = ENGINE_DIFFUSERS
_fallback_reason: Optional[str] = None
@ -90,12 +89,12 @@ def active_engine_name() -> str:
def _activate(name: str, reason: Optional[str]) -> Any:
global _active_engine_name, _fallback_reason
# Serialize the whole check -> unload -> publish transition without holding _lock across the
# slow unload(), closing the window where a second _activate reads the still-old active engine,
# takes the "no change" branch, and loads onto the engine this call is unloading.
# slow unload(), closing the window where a second _activate reads the still-old active engine
# and loads onto the engine this call is unloading.
with _transition_lock:
# Switching engines: unload the deactivated one first, else its model stays resident but
# unreachable (the evictor only targets the active engine), leaking 10+ GB. The unload is
# slow, so resolve the engine under _lock but run unload() OUTSIDE it.
# unreachable (the evictor only targets the active engine), leaking 10+ GB. The unload is slow,
# so resolve under _lock but run unload() OUTSIDE it.
engine_to_unload = None
old_name = None
with _lock:
@ -107,9 +106,8 @@ def _activate(name: str, reason: Optional[str]) -> Any:
_fallback_reason = reason if name == ENGINE_DIFFUSERS else None
if engine_to_unload is not None:
# Publish the new engine only AFTER the old one unloads. The evictor unloads
# get_active_diffusion_engine(), so flipping the name first would let a concurrent
# acquire_for evict the new (empty) engine while the old model is still freeing VRAM.
# Keeping the OLD engine as the evict target grants the GPU only once it is freed.
# get_active_diffusion_engine(), so flipping the name first would let a concurrent acquire_for
# evict the new (empty) engine while the old model is still freeing VRAM.
try:
engine_to_unload.unload()
except Exception as exc: # noqa: BLE001 -- best-effort; never block the switch
@ -176,9 +174,9 @@ def select_and_activate_engine(
binary = None
server_binary = None
if policy_eligible and fam_ok:
# Probe the resident sd-server FIRST (the backend prefers it): a server-only install must
# still route to native, and a server-only host shouldn't pay an sd-cli download. Install
# the accelerator-matched build so a forced-native GPU load gets the GPU server.
# Probe the resident sd-server FIRST (the backend prefers it): a server-only install must still
# route to native, and a server-only host shouldn't pay an sd-cli download. Install the
# accelerator-matched build so a forced-native GPU load gets the GPU server.
server_binary = ensure_sd_server_binary(
allow_install = _install_allowed(),
accelerator = _install_accelerator_for(backend),
@ -188,9 +186,9 @@ def select_and_activate_engine(
"sd-server at %s is present but not runnable; not using it", server_binary
)
server_binary = None
# sd-cli is the one-shot fallback. Always LOCATE an existing binary, but auto-INSTALL only
# when there is no usable server. Probe runnability before committing native: a present but
# non-runnable binary would otherwise pass as available and fail inside the background load.
# sd-cli is the one-shot fallback. Always LOCATE an existing binary, but auto-INSTALL only when
# there is no usable server. Probe runnability first: a present but non-runnable binary would
# otherwise pass as available and fail inside the background load.
binary = ensure_sd_cpp_binary(
allow_install = _install_allowed() and server_binary is None,
accelerator = _install_accelerator_for(backend),

240
studio/backend/core/inference/diffusion_families.py
View file

@ -36,35 +36,33 @@ class DiffusionFamily:
# Pipeline kwarg carrying guidance. Most use "guidance_scale"; Qwen-Image's real CFG is
# "true_cfg_scale" (its distilled guidance is off).
cfg_kwarg: str = "guidance_scale"
# The pipe attribute holding the denoiser: DiT families ``pipe.transformer`` (default), U-Net
# families (SDXL) ``pipe.unet``. Read wherever the backend touches the denoiser generically.
# The pipe attribute holding the denoiser: ``pipe.transformer`` for DiT families (default),
# ``pipe.unet`` for U-Net families (SDXL).
denoiser_attr: str = "transformer"
# True when a single-file ``.safetensors`` is the WHOLE pipeline (SDXL), so the loader calls
# ``pipeline_class.from_single_file`` directly. DiT families leave this False (transformer-only).
# ``pipeline_class.from_single_file``. DiT families leave this False (transformer-only).
single_file_is_pipeline: bool = False
# True for families whose pipeline needs MULTIPLE denoisers no single file carries (Ideogram 4:
# conditional + unconditional_transformer), so only a full ``pipeline`` load is valid;
# validate_load_request rejects single-file / GGUF kinds up front.
# True for families needing MULTIPLE denoisers no single file carries (Ideogram 4), so only a
# full ``pipeline`` load is valid; validate_load_request rejects single-file / GGUF up front.
pipeline_only: bool = False
# Optional diffusers pipeline classes for image-conditioned workflows, built around the resident
# modules via ``Pipeline.from_pipe`` (no reload). None = unsupported (UI gates it off).
img2img_pipeline_class: Optional[str] = None
inpaint_pipeline_class: Optional[str] = None
# ControlNet pipeline + model classes: the backend loads the model via from_pretrained and
# builds the pipeline via ``from_pipe(base, controlnet=model)`` (no reload). None on both =
# no support (UI gates it off).
# ControlNet pipeline + model classes: the model loads via from_pretrained and the pipeline via
# ``from_pipe(base, controlnet=model)`` (no reload). None on both = no support.
controlnet_pipeline_class: Optional[str] = None
controlnet_model_class: Optional[str] = None
# True when the inpaint pipeline keeps the canvas size, so it can also drive outpaint. False for
# FLUX.2 (it scales >1MP inputs to ~1MP, shrinking an outpaint canvas) -> Inpaint but not Extend.
# FLUX.2 (it scales >1MP inputs to ~1MP, shrinking the canvas) -> Inpaint but not Extend.
inpaint_preserves_size: bool = True
# True for instruction-editing families (Qwen-Image-Edit / FLUX Kontext): the OWN pipeline IS the
# edit pipeline (image + instruction, no plain text-to-image), used directly (no from_pipe).
# True for instruction-editing families (Qwen-Image-Edit / FLUX Kontext): the OWN pipeline IS
# the edit pipeline (image + instruction, no plain text-to-image), used directly (no from_pipe).
# ``base_repo`` supplies the VAE / text-encoder / processor / scheduler for the GGUF transformer.
edit: bool = False
# True for families whose text-to-image pipeline ALSO accepts reference image(s) (FLUX.2's
# ``image`` arg). Unlike ``edit`` they still do plain text-to-image; unlike img2img the
# conditioning is reference-based (no ``strength``, output size from width/height). Used directly.
# conditioning is reference-based (no ``strength``, output size from width/height).
reference: bool = False
# Extra lowercased substrings (besides ``name``) that map a repo id here.
aliases: tuple[str, ...] = field(default_factory = tuple)
@ -72,29 +70,24 @@ class DiffusionFamily:
# promotes a resolved float16 to float32 for these.
fp16_incompatible: bool = False
# False only for a family whose denoiser block doesn't compile cleanly with regional
# torch.compile. Consulted on the GGUF path too; all current families compile, so this stays True.
# torch.compile. Consulted on the GGUF path too; all current families compile.
supports_torch_compile: bool = True
# Optional pre-quantized transformer checkpoints as (scheme, repo_id) pairs. When the fast quant
# path resolves a scheme with a hosted checkpoint, the loader fetches the already-quantized
# weights instead of the dense bf16 (lower load VRAM + smaller download). Empty -> unchanged.
# Optional pre-quantized transformer checkpoints as (scheme, repo_id) pairs: the loader fetches
# already-quantized weights instead of the dense bf16 (lower load VRAM + smaller download).
prequant_repos: tuple[tuple[str, str], ...] = field(default_factory = tuple)
# Hosted checkpoints for NON-DEFAULT bases of the family, as (base_repo, scheme, repo_id)
# triples with base_repo lowercased. One family entry covers several published variants
# (flux.1: schnell/dev/Krea-dev) whose weights differ, so each variant needs its own baked
# checkpoint; the loader's base_model_id validation correctly refuses the default entry for
# them. Resolution prefers an exact variant match, then falls back to ``prequant_repos``.
# (flux.1: schnell/dev/Krea-dev) whose weights differ, so each needs its own baked checkpoint.
# Resolution prefers an exact variant match, then falls back to ``prequant_repos``.
prequant_variant_repos: tuple[tuple[str, str, str], ...] = field(default_factory = tuple)
# Hosted PRE-CAST text-encoder checkpoints as (scheme, component, repo_id) triples
# (component is the pipeline attribute, e.g. "text_encoder"). Serves the layerwise-fp8
# storage scheme only: the cast is a deterministic transform, so the stored artifact is
# bit-identical to dense-load-then-cast while skipping the multi-GB dense TE download
# (see diffusion_te_prequant.py). Empty -> the TE loads dense and casts as before.
# Hosted PRE-CAST text-encoder checkpoints as (scheme, component, repo_id) triples. Serves the
# layerwise-fp8 storage scheme only: the cast is deterministic, so the artifact is bit-identical
# to dense-load-then-cast while skipping the multi-GB dense TE download. Empty -> load dense.
te_prequant_repos: tuple[tuple[str, str, str], ...] = field(default_factory = tuple)
# Native (sd.cpp) single-file assets, used only on the no-GPU sd.cpp engine. The transformer GGUF
# is shared with diffusers; sd-cli also needs a single-file VAE + text encoder(s) (the base repo
# ships those sharded). Each is a (repo_id, filename); ``sd_cpp_text_encoders`` carries a trailing
# SdCppModelFiles field name (clip_l / t5xxl / llm / qwen2vl / clip_g) for the sd-cli flag. Empty
# -> no native mapping (sd.cpp route falls back to diffusers).
# Native (sd.cpp) single-file assets, used only on the no-GPU sd.cpp engine. The transformer
# GGUF is shared with diffusers; sd-cli also needs a single-file VAE + text encoder(s). Each is
# a (repo_id, filename); ``sd_cpp_text_encoders`` carries a trailing SdCppModelFiles field name
# (clip_l / t5xxl / llm / qwen2vl / clip_g) for the sd-cli flag. Empty -> no native mapping.
sd_cpp_vae: Optional[tuple[str, str]] = None
# VAE latent-format override for sd-cli (--vae-format): "flux2" for FLUX.2, None otherwise.
sd_cpp_vae_format: Optional[str] = None
@ -103,47 +96,44 @@ class DiffusionFamily:
# invocation (e.g. Qwen-Image needs euler + flow-shift 3). None leaves sd-cli defaults.
sd_cpp_sampling_method: Optional[str] = None
sd_cpp_flow_shift: Optional[float] = None
# True when Studio can TRAIN a LoRA on this family (a trainer is registered). Opt-in per family
# (each arch needs its own loop); the training-start path refuses a non-trainable family up front.
# True when Studio can TRAIN a LoRA on this family (a trainer is registered). Opt-in per family;
# the training-start path refuses a non-trainable family up front.
trainable: bool = False
# Recommended base repos to train FROM, most-preferred first (e.g. a QLoRA prequant repo, then
# bf16). Surfaced by the Train UI.
train_base_repos: tuple[str, ...] = field(default_factory = tuple)
# When set, deploying a LoRA trained on this family loads THIS repo instead of the trained-on
# checkpoint (Krea: train on Raw, preview on Turbo). Both sides must be the same precision so the
# swap never enlarges the load. Unset elsewhere.
# checkpoint (Krea: train on Raw, preview on Turbo). Both sides must be the same precision.
deploy_base_repo: Optional[str] = None
# Keyed by architecture, not per variant: a checkpoint's specific base repo is read from its HF
# base_model tag at load time, so one entry covers Turbo/full, schnell/dev, etc. (base_repo here is
# a fallback). Only archs whose diffusers transformer supports from_single_file load here.
# Keyed by architecture, not per variant: a checkpoint's base repo is read from its HF
# base_model tag at load time, so one entry covers Turbo/full, schnell/dev, etc. Only archs
# whose diffusers transformer supports from_single_file load here.
_FAMILIES: tuple[DiffusionFamily, ...] = (
DiffusionFamily(
name = "flux.1",
pipeline_class = "FluxPipeline",
transformer_class = "FluxTransformer2DModel",
base_repo = "black-forest-labs/FLUX.1-schnell",
# Hosted pre-quantized DiT checkpoints (gate-validated vs same-seed bf16). The loader
# verifies the checkpoint's baked base_model_id against the repo actually being loaded,
# so a non-default base (e.g. FLUX.1-dev under this family) safely falls back to the
# dense-quantize path instead of loading schnell weights.
# Hosted pre-quantized DiT checkpoints (gate-validated vs same-seed bf16). The loader verifies
# the baked base_model_id against the repo being loaded, so a non-default base safely falls back
# to the dense-quantize path instead of loading schnell weights.
prequant_repos = (
("int8", "unsloth/FLUX.1-schnell-FP8"),
("fp8", "unsloth/FLUX.1-schnell-FP8"),
),
# Gate-validated checkpoints baked from the dev / Krea-dev weights (same arch, different
# weights): without these entries the default schnell checkpoint is refused for those
# bases and every int8/fp8 load pays the dense download + on-the-fly quantise.
# weights): without these the default schnell checkpoint is refused and every int8/fp8 load pays
# the dense download + on-the-fly quantise.
prequant_variant_repos = (
("black-forest-labs/flux.1-dev", "int8", "unsloth/FLUX.1-dev-FP8"),
("black-forest-labs/flux.1-dev", "fp8", "unsloth/FLUX.1-dev-FP8"),
("black-forest-labs/flux.1-krea-dev", "int8", "unsloth/FLUX.1-Krea-dev-FP8"),
("black-forest-labs/flux.1-krea-dev", "fp8", "unsloth/FLUX.1-Krea-dev-FP8"),
),
# Pre-cast T5-XXL (9.52 -> 5.90 GB; CLIP-L stays dense, 0.25 GB). One artifact
# serves schnell/dev/Krea-dev: the T5 shards are byte-identical across all three
# (verified sha256, see diffusion_te_prequant._TE_EQUIVALENT_BASES).
# Pre-cast T5-XXL (9.52 -> 5.90 GB; CLIP-L stays dense). One artifact serves schnell/dev/
# Krea-dev: the T5 shards are byte-identical across all three (verified sha256).
te_prequant_repos = (("fp8", "text_encoder_2", "unsloth/FLUX.1-schnell-FP8"),),
aliases = ("flux1", "flux-1"),
# LoRA training targets FLUX.1-dev via the DiT trainer (QLoRA nf4); the dev repo is gated.
@ -159,7 +149,7 @@ _FAMILIES: tuple[DiffusionFamily, ...] = (
("comfyanonymous/flux_text_encoders", "t5xxl_fp16.safetensors", "t5xxl"),
),
),
# FLUX.2-klein is Flux2KleinPipeline (Qwen3 encoder), not the Mistral Flux2Pipeline; must
# FLUX.2-klein is Flux2KleinPipeline (Qwen3 encoder), not the Mistral Flux2Pipeline, so it must
# precede a generic flux match. The Mistral Flux2Pipeline is the flux.2-dev family below.
DiffusionFamily(
name = "flux.2-klein",
@ -174,8 +164,8 @@ _FAMILIES: tuple[DiffusionFamily, ...] = (
# LoRA training via the DiT trainer (QLoRA nf4 by default); klein-4B is not gated.
trainable = True,
train_base_repos = ("black-forest-labs/FLUX.2-klein-4B",),
# Flux2KleinPipeline takes reference image(s) via `image`, so it exposes a "reference"
# workflow atop text-to-image. It has an inpaint pipeline (no img2img) -> inpaint + extend.
# Flux2KleinPipeline takes reference image(s) via `image`, so it exposes a "reference" workflow
# atop text-to-image. It has an inpaint pipeline (no img2img) -> inpaint + extend.
reference = True,
inpaint_pipeline_class = "Flux2KleinInpaintPipeline",
# FLUX.2 scales >1MP inputs to ~1MP, so outpaint can't grow.
@ -188,9 +178,9 @@ _FAMILIES: tuple[DiffusionFamily, ...] = (
("Comfy-Org/z_image_turbo", "split_files/text_encoders/qwen_3_4b.safetensors", "llm"),
),
),
# FLUX.2-dev: full (non-distilled) FLUX.2 on the Mistral Flux2Pipeline (distinct from klein), so
# its own entry. Base repo is gated. Text-to-image only (no Flux2 img2img/inpaint in diffusers
# 0.38). VAE + Mistral encoder come from the open Comfy-Org/flux2-dev mirror for sd-cli.
# FLUX.2-dev: full (non-distilled) FLUX.2 on the Mistral Flux2Pipeline, so its own entry. Base
# repo is gated. Text-to-image only (no Flux2 img2img/inpaint in diffusers 0.38). VAE + Mistral
# encoder come from the open Comfy-Org/flux2-dev mirror for sd-cli.
DiffusionFamily(
name = "flux.2-dev",
pipeline_class = "Flux2Pipeline",
@ -203,8 +193,8 @@ _FAMILIES: tuple[DiffusionFamily, ...] = (
# Pre-cast Mistral-Small-24B conditioner (bf16 ~48 GB dense, ~24.7 GB pre-cast).
te_prequant_repos = (("fp8", "text_encoder", "unsloth/FLUX.2-dev-FP8"),),
aliases = ("flux2-dev", "flux2dev"),
# LoRA training via the DiT trainer (QLoRA nf4 by default); the base repo is gated, so
# training requires an HF token with the FLUX.2-dev license accepted.
# LoRA training via the DiT trainer (QLoRA nf4 by default); the base repo is gated, so training
# requires an HF token with the FLUX.2-dev license accepted.
trainable = True,
train_base_repos = ("black-forest-labs/FLUX.2-dev",),
sd_cpp_vae = ("Comfy-Org/flux2-dev", "split_files/vae/flux2-vae.safetensors"),
@ -219,8 +209,8 @@ _FAMILIES: tuple[DiffusionFamily, ...] = (
),
DiffusionFamily(
# FLUX instruction editing: FluxKontextPipeline takes an image + edit instruction; the GGUF
# transformer is standard FluxTransformer2DModel. Specific aliases first so detect_family
# prefers this over "flux.1" and un-rejects the "kontext" keyword.
# transformer is standard FluxTransformer2DModel. Specific aliases first so detect_family prefers
# this over "flux.1" and un-rejects the "kontext" keyword.
name = "flux.1-kontext",
pipeline_class = "FluxKontextPipeline",
transformer_class = "FluxTransformer2DModel",
@ -230,8 +220,8 @@ _FAMILIES: tuple[DiffusionFamily, ...] = (
),
DiffusionFamily(
# Qwen instruction editing: the 2511 checkpoint ships as QwenImageEditPlusPipeline
# (multi-image); the GGUF transformer is standard QwenImageTransformer2DModel. Specific
# aliases first so detect_family prefers this over "qwen-image".
# (multi-image); the GGUF transformer is standard QwenImageTransformer2DModel. Specific aliases
# first so detect_family prefers this over "qwen-image".
name = "qwen-image-edit",
pipeline_class = "QwenImageEditPlusPipeline",
transformer_class = "QwenImageTransformer2DModel",
@ -253,8 +243,8 @@ _FAMILIES: tuple[DiffusionFamily, ...] = (
base_repo = "Qwen/Qwen-Image",
# int8 only: fp8 is family-denied (_FAMILY_SCHEME_DENY) so a repo entry would be dead.
prequant_repos = (("int8", "unsloth/Qwen-Image-FP8"),),
# Pre-cast Qwen2.5-VL-7B (bf16 ~16.6 GB dense, ~8.8 GB pre-cast). The DiT fp8 denial
# is a transformer-scheme rule; the layerwise TE cast is unaffected.
# Pre-cast Qwen2.5-VL-7B (bf16 ~16.6 GB dense, ~8.8 GB pre-cast). The DiT fp8 denial is a
# transformer-scheme rule; the layerwise TE cast is unaffected.
te_prequant_repos = (("fp8", "text_encoder", "unsloth/Qwen-Image-FP8"),),
cfg_kwarg = "true_cfg_scale",
aliases = ("qwen_image", "qwenimage"),
@ -288,8 +278,8 @@ _FAMILIES: tuple[DiffusionFamily, ...] = (
("int8", "unsloth/Z-Image-Turbo-FP8"),
("fp8", "unsloth/Z-Image-Turbo-FP8"),
),
# Pre-cast Qwen3-4B TE (8.04 -> 4.41 GB). NOT shared with flux.2-klein-4B: klein's
# TE retrained layer 35's MLP (up/down_proj maxdiff 0.86 vs this checkpoint).
# Pre-cast Qwen3-4B TE (8.04 -> 4.41 GB). NOT shared with flux.2-klein-4B: klein's TE retrained
# layer 35's MLP (up/down_proj maxdiff 0.86 vs this checkpoint).
te_prequant_repos = (("fp8", "text_encoder", "unsloth/Z-Image-Turbo-FP8"),),
aliases = ("zimage", "z_image"),
# LoRA training via the DiT trainer (bf16); defaults to the prequant nf4 repo for QLoRA.
@ -316,13 +306,13 @@ _FAMILIES: tuple[DiffusionFamily, ...] = (
("int8", "unsloth/Krea-2-Turbo-FP8"),
("fp8", "unsloth/Krea-2-Turbo-FP8"),
),
# Pre-cast Qwen3-VL-4B TE (8.88 -> 4.83 GB); handed into load_krea2_pipeline
# directly (constructor assembly never sees pipe_kwargs).
# Pre-cast Qwen3-VL-4B TE (8.88 -> 4.83 GB); handed into load_krea2_pipeline directly
# (constructor assembly never sees pipe_kwargs).
te_prequant_repos = (("fp8", "text_encoder", "unsloth/Krea-2-Turbo-FP8"),),
aliases = ("krea2",),
# LoRA training via the DiT trainer (no prequant repo yet, so nf4 quantizes on the fly).
# Krea's guidance: train on the undistilled Raw, run adapters on Turbo, so Raw is the
# default training base and Turbo the inference/base repo.
# LoRA training via the DiT trainer (no prequant repo yet, so nf4 quantizes on the fly). Krea's
# guidance: train on the undistilled Raw, run adapters on Turbo, so Raw is the default training
# base and Turbo the inference/base repo.
trainable = True,
train_base_repos = ("krea/Krea-2-Raw", "krea/Krea-2-Turbo"),
# Adapters trained on Raw run on Turbo; deploy previews them there (same bf16 precision).
@ -330,50 +320,43 @@ _FAMILIES: tuple[DiffusionFamily, ...] = (
# Exported bf16-only; fp16 unvalidated upstream, so keep the fp16 fallback off like z-image.
fp16_incompatible = True,
),
# Lumina Image 2.0: a 2.6B single-stream DiT with a Gemma2-2B encoder and a standard
# 16-channel AutoencoderKL, all transformers-4.x-compatible, so the generic
# from_pretrained pipeline path loads it. No GGUF/sd.cpp mapping exists upstream.
# NOT aliased to bare "lumina": Lumina-Next checkpoints are a different arch
# (LuminaText2ImgPipeline) and must stay unknown rather than crash mid-load.
# Lumina Image 2.0: a 2.6B single-stream DiT with a Gemma2-2B encoder and a standard 16-channel
# AutoencoderKL, all transformers-4.x-compatible, so the generic from_pretrained path loads it.
# No GGUF/sd.cpp mapping upstream. NOT aliased to bare "lumina": Lumina-Next checkpoints are a
# different arch and must stay unknown rather than crash mid-load.
DiffusionFamily(
name = "lumina-2",
pipeline_class = "Lumina2Pipeline",
transformer_class = "Lumina2Transformer2DModel",
base_repo = "Alpha-VLLM/Lumina-Image-2.0",
# Gate-validated hosted checkpoints (28/28 pairs each; LPIPS mean 0.146 int8 /
# 0.116 fp8 vs same-seed bf16).
# Gate-validated hosted checkpoints (28/28 pairs each; LPIPS mean 0.146 int8 / 0.116 fp8).
prequant_repos = (
("int8", "unsloth/Lumina-Image-2.0-FP8"),
("fp8", "unsloth/Lumina-Image-2.0-FP8"),
),
# Pre-cast Gemma2-2B TE. The Hub stores it fp32 (10.46 GB), so the 3.20 GB
# artifact is a 3.3x download cut even though the model is small.
# Pre-cast Gemma2-2B TE. The Hub stores it fp32 (10.46 GB), so the 3.20 GB artifact is a 3.3x
# download cut even though the model is small.
te_prequant_repos = (("fp8", "text_encoder", "unsloth/Lumina-Image-2.0-FP8"),),
aliases = ("lumina-image-2.0", "lumina-image-2", "lumina2"),
# Published and validated bf16-only upstream; keep the fp16 fallback off like z-image.
fp16_incompatible = True,
),
# HunyuanImage 2.1 (diffusers >= 0.39): a 17B dual-stream DiT with a Qwen2.5-VL text
# encoder, a ByT5 glyph encoder, and the 32x-compression HunyuanImage VAE. The community
# mirror also ships guider/ocr_guider components (AdaptiveProjectedMixGuidance), which
# 0.39 loads natively, so the generic from_pretrained pipeline path covers the whole
# stack. 2K-native (the card recipe renders 2048x2048); classifier-free guidance runs
# inside the repo's guider at its baked scale, and the call's own guidance knob is
# distilled_guidance_scale (there is no guidance_scale kwarg). Distinct from
# HunyuanImage-3.0, which stays excluded above: 2.1 has a real diffusers pipeline.
# HunyuanImage 2.1 (diffusers >= 0.39): a 17B dual-stream DiT with a Qwen2.5-VL text encoder, a
# ByT5 glyph encoder, and the 32x-compression HunyuanImage VAE. The community mirror also ships
# guider/ocr_guider components 0.39 loads natively, so the generic from_pretrained path covers
# the stack. 2K-native; CFG runs inside the repo's guider at its baked scale and the call's own
# knob is distilled_guidance_scale. Distinct from the excluded HunyuanImage-3.0.
DiffusionFamily(
name = "hunyuanimage-2.1",
# Hosted checkpoints, verified bit-identical to on-the-fly quantize (the family's
# guider pipeline is not run-to-run deterministic, so same-seed LPIPS vs bf16 blends
# trajectory divergence with harness noise; per-case hard checks pass and the drift is
# compositional, reviewed visually).
# Hosted checkpoints, verified bit-identical to on-the-fly quantize (the guider pipeline is not
# run-to-run deterministic, so same-seed LPIPS vs bf16 blends trajectory divergence with harness
# noise; per-case hard checks pass and the drift is compositional, reviewed visually).
prequant_repos = (
("int8", "unsloth/HunyuanImage-2.1-FP8"),
("fp8", "unsloth/HunyuanImage-2.1-FP8"),
),
# The Qwen2.5-VL TE is byte-identical to Qwen-Image's (verified sha256, see
# _TE_EQUIVALENT_BASES), so the family reuses the Qwen-Image artifact: zero new
# hosting, 16.58 -> 8.84 GB download. ByT5 (text_encoder_2) stays dense.
# The Qwen2.5-VL TE is byte-identical to Qwen-Image's (verified sha256), so the family reuses
# that artifact: zero new hosting, 16.58 -> 8.84 GB download. ByT5 stays dense.
te_prequant_repos = (("fp8", "text_encoder", "unsloth/Qwen-Image-FP8"),),
pipeline_class = "HunyuanImagePipeline",
transformer_class = "HunyuanImageTransformer2DModel",
@ -385,21 +368,19 @@ _FAMILIES: tuple[DiffusionFamily, ...] = (
),
# HiDream-I1: a 17B MoE DiT (16 double + 32 single layers, 4 routed experts) with FOUR text
# encoders. The repos ship CLIP-L/CLIP-G/T5-XXL but NOT the Llama-3.1-8B text_encoder_4 their
# model_index names: the loader assembles it from the open unsloth mirror
# (diffusion_hidream.py). Full / Dev / Fast share the arch, so one family covers all three
# (per-variant step/guidance defaults below). city96 publishes a GGUF but the GGUF path would
# need the same TE4 assembly for tiny demand, so no GGUF artifact is wired yet.
# model_index names: the loader assembles it from the open unsloth mirror. Full / Dev / Fast
# share the arch, so one family covers all three. A GGUF path would need the same TE4 assembly
# for tiny demand, so none is wired yet.
DiffusionFamily(
name = "hidream-i1",
# Hosted checkpoints: 28/28 per-case gate pairs per scheme (LPIPS suite means 0.291
# int8 / 0.278 fp8, the 50-step trajectory band); int8 verified bit-identical to
# on-the-fly quantize across all 1615 state dict tensors.
# Hosted checkpoints: 28/28 per-case gate pairs per scheme (LPIPS suite means 0.291 int8 /
# 0.278 fp8, the 50-step trajectory band); int8 verified bit-identical to on-the-fly quantize.
prequant_repos = (
("int8", "unsloth/HiDream-I1-Full-FP8"),
("fp8", "unsloth/HiDream-I1-Full-FP8"),
),
# Pre-cast Llama-3.1-8B TE4 (16.1 GB bf16 -> 8.1 GB). The generic TE pass only covers
# text_encoder.._3, so TE4 engages via hidream_te4_kwargs, not te_prequant_pipe_kwargs.
# text_encoder.._3, so TE4 engages via hidream_te4_kwargs.
te_prequant_repos = (("fp8", "text_encoder_4", "unsloth/HiDream-I1-Full-FP8"),),
pipeline_class = "HiDreamImagePipeline",
transformer_class = "HiDreamImageTransformer2DModel",
@ -410,10 +391,9 @@ _FAMILIES: tuple[DiffusionFamily, ...] = (
),
# Ideogram 4 (diffusers >= 0.39): a 34-layer DiT PAIR (conditional + unconditional_transformer
# for dual-branch CFG, both ~9B, so memory planning counts two DiTs) with a Qwen3-VL encoder.
# No bf16 checkpoint: ideogram-4-fp8 (raw float8, upcast on load) is the highest-precision
# artifact and the family base; the -nf4 repos carry bnb-4bit quantization_configs. All gated.
# No GGUF/sd.cpp mapping. CFG quirk: the pipeline takes guidance_scale OR a per-step
# guidance_schedule (see the loader's IDEOGRAM4 branch).
# No bf16 checkpoint: ideogram-4-fp8 (raw float8, upcast on load) is the family base; the -nf4
# repos carry bnb-4bit configs. All gated, no GGUF/sd.cpp mapping. CFG quirk: the pipeline takes
# guidance_scale OR a per-step guidance_schedule (see the loader's IDEOGRAM4 branch).
DiffusionFamily(
name = "ideogram-4",
pipeline_class = "Ideogram4Pipeline",
@ -424,9 +404,8 @@ _FAMILIES: tuple[DiffusionFamily, ...] = (
pipeline_only = True,
),
# SDXL is the one U-Net family: the denoiser is ``pipe.unet`` and a single-file ``.safetensors``
# is the WHOLE pipeline, so it sets ``denoiser_attr="unet"`` + ``single_file_is_pipeline=True``
# and loads via the pipeline class. img2img / inpaint / ControlNet are the standard SDXL
# pipelines via from_pipe. No GGUF/single-file transformer path and no sd.cpp mapping.
# is the WHOLE pipeline, so it sets ``denoiser_attr="unet"`` + ``single_file_is_pipeline=True``.
# img2img / inpaint / ControlNet are the standard SDXL pipelines via from_pipe. No GGUF path.
DiffusionFamily(
name = "sdxl",
pipeline_class = "StableDiffusionXLPipeline",
@ -454,8 +433,8 @@ def trainable_family_names() -> tuple[str, ...]:
return tuple(fam.name for fam in _FAMILIES if fam.trainable)
# The family whose CFG uses a guidance_scale/guidance_schedule pair (the loader special-cases the
# call). Named here so the two modules can't drift.
# The family whose CFG uses a guidance_scale/guidance_schedule pair (the loader special-cases
# the call). Named here so the two modules can't drift.
IDEOGRAM4_FAMILY_NAME = "ideogram-4"
# The family whose generate call carries the card's CFG-truncation ratio (the loader
@ -463,10 +442,9 @@ IDEOGRAM4_FAMILY_NAME = "ideogram-4"
LUMINA2_FAMILY_NAME = "lumina-2"
# Models Studio deliberately does NOT support, reason surfaced verbatim in the load error (vs the
# generic unknown-family message). Keyed by a lowercase repo-id substring. The bar is a diffusers
# pipeline: HunyuanImage-3.0 is an 80B MoE needing AutoModelForCausalLM + trust_remote_code (RCE
# out of the question).
# Models Studio deliberately does NOT support, reason surfaced verbatim in the load error.
# Keyed by a lowercase repo-id substring. The bar is a diffusers pipeline: HunyuanImage-3.0 is
# an 80B MoE needing AutoModelForCausalLM + trust_remote_code (RCE out of the question).
_EXCLUDED_MODELS: tuple[tuple[str, str], ...] = (
(
# "-3" scoped so a future HunyuanImage 2.x with a diffusers pipeline falls through normally.
@ -486,9 +464,9 @@ def excluded_model_reason(repo_id: str) -> Optional[str]:
return None
# Editing / inpaint checkpoints share an arch keyword but need a different pipeline + input image.
# "layered" rejects Qwen-Image-Layered (its transformer expects an extra addition_t_cond input
# the standard pipeline never supplies, so it crashes at the first denoise). Fails the load fast.
# Editing / inpaint checkpoints share an arch keyword but need a different pipeline + input
# image. "layered" rejects Qwen-Image-Layered, whose transformer expects an extra
# addition_t_cond input the standard pipeline never supplies. Fails the load fast.
_EDIT_KEYWORDS = ("edit", "kontext", "inpaint", "layered")
@ -526,10 +504,9 @@ def detect_family(repo_id: str, override: Optional[str] = None) -> Optional[Diff
needle = repo_id.lower()
match = _best_family_match(needle)
if match is not None:
# Don't let a generic family (qwen-image) swallow a variant it can't run
# (qwen-image-LAYERED): if the id carries a reject keyword the matched family doesn't
# declare, reject. Scope the check to the LAST path component so a parent folder named
# `edit` doesn't reject a valid file (the repo_id/filename fallback passes the filename last).
# Don't let a generic family (qwen-image) swallow a variant it can't run (qwen-image-LAYERED):
# if the id carries a reject keyword the matched family doesn't declare, reject. Scoped to the
# LAST path component so a parent folder named `edit` doesn't reject a valid file.
basename = re.split(r"[/\\]+", needle)[-1]
matched_tokens = (match.name, *match.aliases)
if any(
@ -568,13 +545,12 @@ def resolve_base_repo(fam: DiffusionFamily, base_repo: Optional[str]) -> str:
# Default (steps, guidance) per model for callers that can't pass them (the OpenAI
# /v1/images/generations endpoint has no step/guidance knobs). Matched by substring, most specific
# first -- same values as the UI's MODEL_DEFAULTS table (images-page.tsx); keep in sync.
# /v1/images/generations endpoint has no such knobs). Matched by substring, most specific first;
# same values as the UI's MODEL_DEFAULTS table (images-page.tsx), keep in sync.
_GENERATION_DEFAULTS: tuple[tuple[str, int, float], ...] = (
("z-image-turbo", 9, 0.0),
# FLUX.1 Krea dev is a FLUX.1-dev finetune (flux.1 family), NOT a Krea-2: its card runs
# 28 steps at guidance 4.5. Must precede the generic "krea" key below, which would
# otherwise hand it Krea-2-Turbo's 8-step no-CFG recipe.
# FLUX.1 Krea dev is a FLUX.1-dev finetune (flux.1 family), NOT a Krea-2: its card runs 28 steps
# at guidance 4.5. Must precede the generic "krea" key, which would hand it Turbo's recipe.
("flux.1-krea", 28, 4.5),
# Krea 2 Raw (undistilled): 52 steps / guidance 3.5. Must precede the generic "krea" key.
("krea-2-raw", 52, 3.5),
@ -587,19 +563,19 @@ _GENERATION_DEFAULTS: tuple[tuple[str, int, float], ...] = (
("flux.2-dev", 28, 4.0), # full (non-distilled)
("qwen-image", 20, 4.0),
("z-image", 20, 4.0),
# Lumina Image 2.0 model-card: 50 steps, guidance 4 (plus cfg_trunc_ratio 0.25, which the
# loader passes itself; see LUMINA2_FAMILY_NAME).
# Lumina Image 2.0 model-card: 50 steps, guidance 4 (plus cfg_trunc_ratio 0.25, which the loader
# passes itself; see LUMINA2_FAMILY_NAME).
("lumina", 50, 4.0),
# HunyuanImage 2.1 model-card: 50 steps; the guidance value feeds the call's
# distilled_guidance_scale (default 3.25), while real CFG runs inside the repo guiders.
# HunyuanImage 2.1 model-card: 50 steps; the guidance value feeds distilled_guidance_scale
# (default 3.25), while real CFG runs inside the repo guiders.
("hunyuanimage", 50, 3.25),
# HiDream-I1 upstream inference.py: Full 50 steps / guidance 5; the distilled Dev (28) and
# Fast (16) run guidance-free. Specific keys precede the generic "hidream" (Full + fallback).
# HiDream-I1 upstream inference.py: Full 50 steps / guidance 5; the distilled Dev (28) and Fast
# (16) run guidance-free. Specific keys precede the generic "hidream".
("hidream-i1-dev", 28, 0.0),
("hidream-i1-fast", 16, 0.0),
("hidream", 50, 5.0),
# Ideogram 4 model-card: 48 steps, guidance 7 (its schedule tapers the last 3 steps to 3.0;
# the loader keeps that taper when the request matches these defaults exactly).
# Ideogram 4 model-card: 48 steps, guidance 7 (its schedule tapers the last 3 steps to 3.0; the
# loader keeps that taper when the request matches these defaults exactly).
("ideogram", 48, 7.0),
# SDXL: Turbo distilled; base wants ~30 steps + CFG ~7. "sdxl-turbo" precedes "sdxl".
("sdxl-turbo", 3, 0.0),

8
studio/backend/core/inference/diffusion_hidream.py
View file

@ -85,8 +85,8 @@ def hidream_te4_kwargs(
hf_token = hf_token,
scheme = "fp8",
logger = logger,
# The Llama TE4 lives in its own standalone repo (config at the root), and
# the pipeline needs hidden states/attentions from its forward.
# The Llama TE4 lives in its own standalone repo (config at the root), and the pipeline needs
# hidden states/attentions from its forward.
config_subfolder = "",
config_overrides = {
"output_hidden_states": True,
@ -116,8 +116,8 @@ def hidream_te4_kwargs(
_cast_fp8(text_encoder_4, cast_target)
logger.info("diffusion.hidream: TE4 layerwise fp8 cast engaged")
except Exception as exc: # noqa: BLE001 -- best-effort like the generic TE pass
# A mid-pass failure can leave fp8 storage / upcast hooks behind; a half-cast
# encoder cannot run as dense, so rebuild it fresh instead of shipping partial state.
# A mid-pass failure can leave fp8 storage / upcast hooks behind; a half-cast encoder cannot
# run as dense, so rebuild it fresh instead of shipping partial state.
logger.warning("diffusion.hidream: TE4 fp8 cast failed, reloading dense: %s", exc)
text_encoder_4 = LlamaForCausalLM.from_pretrained(
HIDREAM_LLAMA_REPO,

26
studio/backend/core/inference/diffusion_ideogram4.py
View file

@ -60,8 +60,8 @@ def _patch_create_causal_mask() -> None:
# The pipeline calls this by keyword. Rename inputs_embeds -> input_embeds for the 5.x spelling.
if "inputs_embeds" in kwargs and "inputs_embeds" not in params and "input_embeds" in params:
kwargs["input_embeds"] = kwargs.pop("inputs_embeds")
# Supply cache_position when required and omitted: past_key_values is None, so positions
# run 0..seq_len-1.
# Supply cache_position when required and omitted: past_key_values is None, so positions run
# 0..seq_len-1.
if "cache_position" in params and "cache_position" not in kwargs:
embeds = kwargs.get("input_embeds", kwargs.get("inputs_embeds"))
if embeds is not None:
@ -73,8 +73,7 @@ def _patch_create_causal_mask() -> None:
# The fp8 attention is a fused ``qkv`` matrix (Q/K/V stacked, each ``hidden_size`` rows).
# hidden_size = attention_head_dim * num_attention_heads, read from config so a future change
# can't mis-split it.
# hidden_size is read from config so a future change can't mis-split it.
_QKV_SPLIT = ("to_q", "to_k", "to_v")
@ -265,8 +264,8 @@ def load_ideogram4_text_encoder(
# Construct normally (so __init__ computes the non-persistent rotary inv_freq the checkpoint
# omits) then copy the dequantized weights in. Build at the target dtype: this ~8B Qwen3-VL
# scaffold is ~2x at the fp32 default (~33 vs ~16 GB) and loads FIRST, so fp32 can OOM a 64 GB
# host. inv_freq is computed in explicit fp32, so a bf16 default leaves it correct.
# scaffold is ~2x at the fp32 default and loads FIRST, so fp32 can OOM a 64 GB host. inv_freq
# is computed in explicit fp32, so a bf16 default leaves it correct.
default_dtype = torch.get_default_dtype()
torch.set_default_dtype(dtype)
try:
@ -333,8 +332,7 @@ def load_ideogram4_transformer(
is_fp8 = True
break
if not is_fp8:
# Already the diffusers split layout (-nf4): let from_pretrained re-apply its
# quantization_config.
# Already the diffusers split layout (-nf4): let from_pretrained re-apply its quantization_config.
model_kwargs: dict[str, Any] = {"subfolder": subfolder, "torch_dtype": dtype}
if token:
model_kwargs["token"] = token
@ -348,10 +346,10 @@ def load_ideogram4_transformer(
config.pop("quantization_config", None)
hidden_size = int(config["attention_head_dim"]) * int(config["num_attention_heads"])
# from_config materializes the full ~9B module before the weights copy in. At fp32 that's ~2x
# the bf16 model (~37 vs ~18 GB), and the second DiT builds while the first + encoder are
# resident, so fp32 can OOM smaller hosts. Build at the target dtype; only rotary_emb.inv_freq
# is absent from the checkpoint (computed in explicit fp32), so a bf16 default leaves it correct.
# from_config materializes the full ~9B module before the weights copy in. At fp32 that is ~2x
# the bf16 model, and the second DiT builds while the first + encoder are resident, so fp32 can
# OOM smaller hosts. Build at the target dtype; only rotary_emb.inv_freq is absent from the
# checkpoint (computed in explicit fp32), so a bf16 default leaves it correct.
default_dtype = torch.get_default_dtype()
torch.set_default_dtype(dtype)
try:
@ -391,8 +389,8 @@ def load_ideogram4_pipeline(
text_encoder = load_ideogram4_text_encoder(repo_id, dtype, hf_token = token)
tokenizer = load_krea2_tokenizer(repo_id, hf_token = token)
transformer = load_ideogram4_transformer(repo_id, "transformer", dtype, hf_token = token)
# The second DiT drives the unconditional branch of Ideogram's dual-branch CFG (same class/size,
# always required).
# The second DiT drives the unconditional branch of Ideogram's dual-branch CFG (same class and
# size, always required).
unconditional_transformer = load_ideogram4_transformer(
repo_id, "unconditional_transformer", dtype, hf_token = token
)

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

@ -88,9 +88,8 @@ def _load_model_index(repo_id: str, hf_token: Optional[str] = None) -> dict[str,
except OSError:
pass
if is_local_dir:
# A local checkpoint dir without the file must fail clearly here: falling through
# to hf_hub_download with a filesystem path as the repo id would die with an
# opaque HFValidationError instead.
# A local checkpoint dir without the file must fail clearly here, else hf_hub_download dies
# with an opaque HFValidationError on the filesystem path.
raise FileNotFoundError(f"model_index.json not found in local model dir {repo_id}")
from huggingface_hub import hf_hub_download
@ -117,8 +116,8 @@ def load_krea2_pipeline(
"""
import diffusers
# diffusers gained Krea2Pipeline in 0.39; on an older install the getattr chain below
# would die with a bare AttributeError mid-load, so fail first with the actionable fix.
# diffusers gained Krea2Pipeline in 0.39; on an older install the getattr chain below would die
# with a bare AttributeError mid-load, so fail first with the actionable fix.
if not hasattr(diffusers, "Krea2Pipeline"):
raise RuntimeError(
f"Krea 2 needs diffusers >= 0.39.0 (Krea2Pipeline); this environment has "

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

@ -42,8 +42,7 @@ class LoraCatalogEntry:
display_name: str
source: str # "local" | "hub"
fmt: str # "safetensors" | "gguf"
# Compatible family names (empty = unknown, shown but not family-gated). UI greys out
# incompatible adapters.
# Compatible family names (empty = unknown, shown but not family-gated).
families: tuple[str, ...] = ()
repo_id: Optional[str] = None # source == "hub"
weight_name: Optional[str] = None # file within the repo (hub)
@ -125,7 +124,7 @@ def _scan_local() -> list[LoraCatalogEntry]:
return []
files = [p for p in children if p.is_file() and p.suffix.lower() in _ALL_EXTS]
# Two files sharing a stem but differing in extension collide on id (== stem), so a colliding
# stem keeps the full filename as its id; a unique stem stays the clean stem.
# stem keeps the full filename as its id.
stem_counts: dict[str, int] = {}
for p in files:
stem_counts[p.stem] = stem_counts.get(p.stem, 0) + 1
@ -137,9 +136,8 @@ def _scan_local() -> list[LoraCatalogEntry]:
except OSError:
size = 0
entry_id = p.name if stem_counts.get(p.stem, 0) > 1 else p.stem
# A ``<stem>.json`` sidecar (written by the trainer on publish) records the adapter's
# family + default weight so it is family-gated instead of "unknown". Best-effort: a
# missing/bad sidecar leaves the defaults.
# A ``<stem>.json`` sidecar (written by the trainer on publish) records the adapter's family +
# default weight so it is family-gated instead of "unknown". Best-effort.
families, weight_default = _read_lora_sidecar(p)
entries.append(
LoraCatalogEntry(
@ -217,7 +215,7 @@ def resolve_one(
so a direct API client cannot load a LoRA tagged for another family through the wrong pipeline.
An untagged catalog entry (empty ``families``) stays unrestricted.
"""
# An empty/whitespace token triggers an auth error instead of anonymous access; normalise to None.
# An empty token triggers an auth error instead of anonymous access; normalise to None.
hf_token = hf_token.strip() if hf_token and hf_token.strip() else None
entry = _catalog_by_id().get(spec_id)
if entry is not None:
@ -245,8 +243,8 @@ def resolve_one(
repo_id, _, weight_name = spec_id.partition(":")
weight_name = weight_name or None
if weight_name is not None:
# A client-supplied weight file must stay a plain filename inside the repo: reject
# traversal / absolute paths so it can't resolve outside the HF cache dir.
# A client-supplied weight file must stay a plain filename inside the repo: reject traversal /
# absolute paths so it can't resolve outside the HF cache dir.
if (
".." in weight_name
or weight_name.startswith(("/", "\\", "~"))
@ -401,8 +399,8 @@ def _fmt_weight(w: float) -> str:
return s or "0"
# Families sd-cli's LoRA name-conversion supports (Qwen-Image has no branch -> excluded).
# Matched by substring against the resolved family name.
# Families sd-cli's LoRA name-conversion supports (Qwen-Image has no branch). Matched by
# substring against the resolved family name.
_NATIVE_LORA_FAMILY_TOKENS = (
"flux.1",
"flux.2",
@ -413,12 +411,10 @@ _NATIVE_LORA_FAMILY_TOKENS = (
"sd3",
"stable-diffusion",
)
# Diffusers quant schemes whose LoRA path is the load-time BAKE (adapters attach on the
# dense transformer BEFORE torchao quantize_ + compile; peft's post-quant TorchaoLoraLinear
# dispatch needs quantizer metadata a manual quantize_ never has). Verified on the Studio
# stack (peft 0.18.1 / torchao 0.17 / torch 2.10): adapter-first, quantize-base-second is
# clean for both schemes -- scale 0 reproduces the quantized base bit-exactly and the wrapped
# transformer compiles.
# Diffusers quant schemes whose LoRA path is the load-time BAKE (adapters attach on the dense
# transformer BEFORE torchao quantize_ + compile; peft's post-quant TorchaoLoraLinear dispatch
# needs quantizer metadata a manual quantize_ never has). Verified on the Studio stack (peft
# 0.18.1 / torchao 0.17 / torch 2.10): scale 0 reproduces the quantized base bit-exactly.
_DIFFUSERS_LORA_BAKED_QUANT = ("int8", "fp8")
# Prototype schemes with no validated LoRA path (and no shipped families needing one).
_DIFFUSERS_LORA_BLOCKED_QUANT = ("nvfp4", "mxfp8")

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

@ -34,11 +34,11 @@ MEMORY_MODES = (
)
# ── offload policies (what the loader does) ──────────────────────────
# none -> all weights resident (fastest; fits only with room).
# model -> enable_model_cpu_offload(): one top-level module on the GPU at a time.
# group -> apply_group_offloading() on the transformer: stream a few blocks at a time with a
# none -- all weights resident (fastest; fits only with room).
# model -- enable_model_cpu_offload(): one top-level module on the GPU at a time.
# group -- apply_group_offloading() on the transformer: stream a few blocks at a time with a
# prefetch stream (lowest practical VRAM for the dominant module).
# sequential -> enable_sequential_cpu_offload(): submodule-level (broken for GGUF on diffusers
# sequential -- enable_sequential_cpu_offload(): submodule-level (broken for GGUF on diffusers
# 0.38, kept as an explicit escape hatch).
OFFLOAD_NONE = "none"
OFFLOAD_MODEL = "model"
@ -46,7 +46,7 @@ OFFLOAD_GROUP = "group"
OFFLOAD_SEQUENTIAL = "sequential"
# Transformer blocks resident per group under group offloading: fewer = lower VRAM, more
# host<->device traffic. One is the lowest-VRAM setting.
# host-to-device traffic.
DEFAULT_GROUP_BLOCKS = 1
DEFAULT_IMAGE_WIDTH = 1024
@ -195,8 +195,8 @@ def _cuda_memory(backend: str) -> tuple[Optional[int], Optional[int], str]:
free, total = torch.cuda.mem_get_info()
kind = "discrete_vram"
try:
# Query the CURRENT device (mem_get_info reports it); hardcoding 0 would inspect the
# wrong GPU and misclassify discrete vs unified when the active device isn't 0.
# Query the CURRENT device (mem_get_info reports it); hardcoding 0 would inspect the wrong GPU
# and misclassify discrete vs unified.
props = torch.cuda.get_device_properties(torch.cuda.current_device())
if bool(getattr(props, "integrated", False) or getattr(props, "is_integrated", False)):
kind = "unified_memory" # e.g. Jetson / integrated SoC
@ -408,8 +408,8 @@ def plan_diffusion_memory(
return group_floor is not None and budget is not None and group_floor <= budget
if not can_offload or device_memory.is_unified:
# MPS / CPU can't stream to a separate device; on unified memory offload just shuffles
# bytes within the same pool.
# MPS / CPU can't stream to a separate device; on unified memory offload just shuffles bytes
# within the same pool.
policy = OFFLOAD_NONE
if device_memory.is_unified:
reasons.append("unified/system memory: CPU offload frees no device memory")
@ -442,8 +442,8 @@ def plan_diffusion_memory(
policy = OFFLOAD_MODEL
reasons.append("companions exceed budget; whole-module offload of every component")
# The legacy cpu_offload flag applies only when no memory_mode was supplied (memory_mode
# overrides it), so an explicit `fast` request stays resident even with the old flag on.
# The legacy cpu_offload flag applies only when no memory_mode was supplied, so an explicit
# `fast` request stays resident even with the old flag on.
if (
explicit_offload
and normalize_memory_mode(requested_mode) is None
@ -454,11 +454,10 @@ def plan_diffusion_memory(
policy = OFFLOAD_MODEL
reasons.append("explicit cpu_offload overrides resident placement")
# VAE savers cap the high-res decode spike. Slicing (one image at a time) is EXACT, so enable
# it on any offload tier / non-discrete backend. Tiling (spatial chunks) is only bit-identical
# for a single tile (<=1MP), so restrict it to the lowest tiers (model / sequential) or no
# spare device pool (MPS / CPU). Group offload keeps the VAE resident -> exact full-image
# decode. On a roomy discrete GPU both stay off.
# VAE savers cap the high-res decode spike. Slicing (one image at a time) is EXACT, so enable it
# on any offload tier / non-discrete backend. Tiling (spatial chunks) is only bit-identical for
# a single tile (<=1MP), so restrict it to the lowest tiers or no spare device pool. Group
# offload keeps the VAE resident for an exact full-image decode; on a roomy GPU both stay off.
any_offload = policy != OFFLOAD_NONE or device_memory.backend in ("mps", "cpu")
tile = policy in (OFFLOAD_MODEL, OFFLOAD_SEQUENTIAL) or device_memory.backend in ("mps", "cpu")
return MemoryPlan(
@ -495,8 +494,8 @@ def apply_memory_plan(
_enable_vae_saver(pipe, "enable_vae_slicing", "enable_slicing", logger)
def _fallback_to_model_offload() -> None:
# The GROUP plan set vae_tiling=False (VAE stays resident). Dropping to whole-module
# offload is the low-VRAM case where the decode spike can OOM, so turn tiling on now.
# The GROUP plan set vae_tiling=False (VAE stays resident). Dropping to whole-module offload is
# the low-VRAM case where the decode spike can OOM, so turn tiling on now.
nonlocal tiling_engaged
pipe.enable_model_cpu_offload(device = device)
if not tiling_engaged:
@ -555,8 +554,8 @@ def _apply_group_offload(pipe: Any, device: str, logger: Any) -> bool:
import torch
from diffusers.hooks import apply_group_offloading
# A dual-DiT pipeline (e.g. Ideogram 4) carries a second denoiser as large as the first;
# leaving it resident defeats this tier. Stream every DiT, keep only smaller companions.
# A dual-DiT pipeline (Ideogram 4) carries a second denoiser as large as the first; leaving it
# resident defeats this tier. Stream every DiT, keep only smaller companions.
streamed: dict[str, Any] = {"transformer": transformer}
for extra in ("transformer_2", "unconditional_transformer"):
module = getattr(pipe, extra, None)
@ -572,19 +571,18 @@ def _apply_group_offload(pipe: Any, device: str, logger: Any) -> bool:
"num_blocks_per_group": DEFAULT_GROUP_BLOCKS,
"use_stream": use_stream,
}
# On the CUDA stream path, overlap each block's H2D copy with compute: non_blocking
# issues it async, record_stream defers the free until the copy's stream is done. Lossless
# (only transfer scheduling changes). Gated on the signature so older diffusers still works.
# On the CUDA stream path, overlap each block's H2D copy with compute: non_blocking issues it
# async, record_stream defers the free until the copy's stream is done. Lossless, and gated on
# the signature so older diffusers still works.
if use_stream:
_params = inspect.signature(apply_group_offloading).parameters
if "non_blocking" in _params:
gkwargs["non_blocking"] = True
if "record_stream" in _params:
gkwargs["record_stream"] = True
# Place the 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 fallback works (diffusers REJECTS enable_model_cpu_offload on a pipe that
# already has group-offload hooks). The streamed transformer places itself via the hooks next.
# Place the 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
# fallback works (diffusers REJECTS enable_model_cpu_offload once group hooks exist).
for name, comp in getattr(pipe, "components", {}).items():
if name in streamed:
continue
@ -596,10 +594,9 @@ def _apply_group_offload(pipe: Any, device: str, logger: Any) -> bool:
return True
except Exception as exc: # noqa: BLE001 — fall back to whole-module offload
if installed:
# An earlier streamed module already has hooks but a later one failed: the pipe is in a
# PARTIAL group-offload state that enable_model_cpu_offload rejects, so the caller's
# fallback would crash. Propagate the real failure (e.g. the OOM) instead of a
# misleading hook error; the "no hooks installed" cases below fall back cleanly.
# An earlier streamed module already has hooks but a later one failed: the pipe is in a PARTIAL
# group-offload state that enable_model_cpu_offload rejects, so propagate the real failure (e.g.
# the OOM) instead of a misleading hook error. The "no hooks installed" cases fall back cleanly.
if logger is not None:
logger.warning(
"diffusion.memory: group offload failed after installing hooks on %d "

70
studio/backend/core/inference/diffusion_precision.py
View file

@ -37,12 +37,11 @@ TE_QUANT_MODES = (TE_QUANT_FP8, TE_QUANT_NVFP4, TE_QUANT_INT8, TE_QUANT_FP8_DYNA
_TEXT_ENCODER_ATTRS = ("text_encoder", "text_encoder_2", "text_encoder_3")
# int8 degrades on large text encoders unless the quant-sensitive decoder blocks stay bf16.
# Per-family (skip_first, skip_last) blocks to keep dense, from measured hidden-state fidelity
# (per-token cosine vs bf16 at the consumed layer): keeping first blocks stops early-layer error
# seeding, last blocks protect the read layer. Families absent have no schedule clearing the bar,
# so int8 falls back to fp8.
# qwen-image (Qwen2.5-VL-7B): first+last 6 -> ~0.997 cosine (both ends; outlier-bound).
# flux.2-dev (Mistral-Small-24B): first 3 -> ~0.98 cosine (early-layer seeding).
# Per-family (skip_first, skip_last) blocks to keep dense, from measured hidden-state fidelity:
# keeping first blocks stops early-layer error seeding, last blocks protect the read layer.
# Families absent have no schedule clearing the bar, so int8 falls back to fp8.
# qwen-image (Qwen2.5-VL-7B): first+last 6 gives ~0.997 cosine (both ends; outlier-bound).
# flux.2-dev (Mistral-Small-24B): first 3 gives ~0.98 cosine (early-layer seeding).
_TE_INT8_SKIP: dict[str, tuple[int, int]] = {
"qwen-image": (6, 6),
"qwen-image-edit": (6, 6),
@ -114,8 +113,8 @@ def quantize_text_encoders(
if skip is None:
_note(logger, f"int8 has no keep-bf16 schedule for family '{family}'; using fp8")
mode = TE_QUANT_FP8
# torchao modes produce subclasses that reject Module.to(), which an offload placement uses
# (the DiT path skips torchao under offload for the same reason). Layerwise fp8 streams fine.
# torchao modes produce subclasses that reject Module.to(), which an offload placement uses.
# Layerwise fp8 streams fine.
if offload_active and mode in (TE_QUANT_INT8, TE_QUANT_FP8_DYNAMIC, TE_QUANT_NVFP4):
_note(
logger,
@ -176,7 +175,7 @@ def _keep_bf16_block_fqns(encoder: Any, skip_first: int, skip_last: int) -> set[
def _cast_int8_selective(encoder: Any, target: Any, skip_first: int, skip_last: int) -> None:
# torchao dynamic int8 on the FLOP-heavy Linears, keeping the first/last decoder blocks (and
# vision tower / lm_head / T5 wo) bf16. Reuses the transformer-quant factory so config never drifts.
# vision tower / lm_head / T5 wo) bf16. Reuses the transformer-quant factory so config can't drift.
from torchao.quantization import quantize_
from .diffusion_transformer_quant import (
TQ_INT8,
@ -217,8 +216,8 @@ def _weight_has_zero_output_row(module: Any) -> bool:
def _cast_fp8_dynamic(encoder: Any, target: Any) -> None:
# torchao dynamic fp8 COMPUTE, per-row (torch._scaled_mm on the fp8 cores). Unlike layerwise
# `fp8` this keeps the matmul in fp8 instead of upcasting. Robust across encoder sizes, so no
# per-layer keep-bf16; only the vision tower / lm_head / T5 wo are excluded.
# `fp8` this keeps the matmul in fp8. Robust across encoder sizes, so no per-layer keep-bf16;
# only the vision tower / lm_head / T5 wo are excluded.
from torchao.quantization import quantize_
from .diffusion_transformer_quant import (
TQ_FP8,
@ -227,8 +226,8 @@ def _cast_fp8_dynamic(encoder: Any, target: Any) -> None:
make_filter_fn,
)
# require_bf16: scaled_mm asserts a bf16 weight, so skip any stray non-bf16 Linear rather than
# aborting the pass (belt-and-suspenders over the named T5 wo exclusion).
# require_bf16: scaled_mm asserts a bf16 weight, so skip a stray non-bf16 Linear rather than
# aborting the pass.
base = make_filter_fn(
DEFAULT_MIN_LINEAR_FEATURES, _te_exclude_tokens(encoder), require_bf16 = True
)
@ -246,11 +245,10 @@ def _cast_fp8(encoder: Any, target: Any) -> None:
from diffusers.hooks import apply_layerwise_casting
from diffusers.hooks.layerwise_casting import DEFAULT_SKIP_MODULES_PATTERN
# Idempotent: a pre-cast encoder (diffusion_te_prequant) arrives with the layerwise hooks
# already installed, and re-registering the same hook name raises -- which would make
# quantize_text_encoders report the (actually engaged) cast as failed. Keyed on the explicit
# completion marker this function sets, NOT on hook presence alone: leftover hooks from a
# cast that failed mid-pass must still fail closed, not read as "already cast".
# Idempotent: a pre-cast encoder arrives with the layerwise hooks already installed, and
# re-registering the same hook name raises, which would report the engaged cast as failed. Keyed
# on the explicit completion marker, NOT hook presence: leftover hooks from a cast that failed
# mid-pass must still fail closed.
if getattr(encoder, "_unsloth_te_cast_complete", False) and _has_layerwise_hooks(encoder):
return
@ -264,10 +262,9 @@ def _cast_fp8(encoder: Any, target: Any) -> None:
# racing the upcast hook so F.linear sees fp8 input vs bf16 weight. Literal substrings.
skip += tuple(re.escape(m) for m in (getattr(encoder, "_keep_in_fp32_modules", None) or ()))
# (2) an output projection tied to the input embedding. A CausalLM encoder (FLUX.2's Qwen3)
# ties lm_head.weight to embed_tokens.weight; casting lm_head (an nn.Linear) to fp8 drags the
# shared embedding down, which then emits fp8 activations that crash the first RMSNorm. Skip
# the tied projection so the shared tensor stays dense (lm_head is unused for prompt encoding).
# (2) an output projection tied to the input embedding. A CausalLM encoder (FLUX.2's Qwen3) ties
# lm_head.weight to embed_tokens.weight; casting lm_head to fp8 drags the shared embedding down,
# which then emits fp8 activations that crash the first RMSNorm. lm_head is unused here anyway.
get_out, get_in = (
getattr(encoder, "get_output_embeddings", None),
getattr(encoder, "get_input_embeddings", None),
@ -284,27 +281,23 @@ def _cast_fp8(encoder: Any, target: Any) -> None:
storage_dtype = torch.float8_e4m3fn,
compute_dtype = target.dtype,
skip_modules_pattern = skip,
# Keep token-embedding tables full precision: the diffusers default only skips vision
# pos/patch embeds, and fp8'ing nn.Embedding quantizes every prompt token to the coarse
# fp8 grid, hurting fidelity.
# Keep token-embedding tables full precision: the diffusers default only skips vision pos/patch
# embeds, and fp8'ing nn.Embedding quantizes every prompt token to the coarse fp8 grid.
skip_modules_classes = (torch.nn.Embedding,),
)
# Module.dtype reports the first floating parameter, which is now fp8 STORAGE; pipelines
# derive tensor dtypes from encoder.dtype (Flux2 casts prompt embeds to it and feeds the
# result to randn_tensor, which has no fp8 kernel; VLM pipelines cast pixel_values to it,
# racing the upcast hooks). The encoder computes in target.dtype, so report that -- via a
# property shadowed on the ORIGINAL class reading a per-instance override. Swapping
# __class__ to a dynamic subclass instead breaks transformers' kwargs-based output
# recording (Qwen3VLModel returned hidden_states=None and krea-2 crashed at encode).
# Module.dtype reports the first floating parameter, which is now fp8 STORAGE; pipelines derive
# tensor dtypes from encoder.dtype (Flux2 feeds it to randn_tensor, which has no fp8 kernel;
# VLM pipelines cast pixel_values to it, racing the upcast hooks). The encoder computes in
# target.dtype, so report that via a property shadowed on the ORIGINAL class reading a
# per-instance override. A dynamic __class__ swap instead breaks transformers' output recording.
compute_dtype = getattr(target, "dtype", None)
try:
if compute_dtype is not None:
_install_dtype_override(type(encoder))
encoder._unsloth_te_compute_dtype = compute_dtype
# Marks the cast COMPLETE (hooks fully installed), enabling the idempotent early return
# above. Best-effort like the dtype override: a non-Module double without settable
# attributes still counts as cast, it just re-casts on a repeat call.
# Marks the cast COMPLETE (hooks fully installed), enabling the idempotent early return above.
# Best-effort: a non-Module double without settable attributes just re-casts on a repeat call.
encoder._unsloth_te_cast_complete = True
except Exception: # noqa: BLE001 — real HF encoders are heap-type nn.Modules; only doubles fail
pass
@ -346,10 +339,9 @@ def _has_layerwise_hooks(encoder: Any) -> bool:
def _cast_nvfp4(encoder: Any, target: Any) -> None:
# Weight-only NVFP4: linear weights become 4-bit NVFP4 on Blackwell FP4 cores; norms /
# embeddings untouched. Exclude the VLM vision tower / lm_head / T5 wo and sub-512 projections
# like the int8/fp8 TE modes (4-bit-ing a VLM image tower degrades the edit conditioning);
# require_bf16 skips non-bf16 Linears so the cast engages instead of aborting.
# Weight-only NVFP4: linear weights become 4-bit NVFP4 on Blackwell FP4 cores; norms / embeddings
# untouched. Exclude the VLM vision tower / lm_head / T5 wo and sub-512 projections like the
# int8/fp8 TE modes; require_bf16 skips non-bf16 Linears so the cast engages instead of aborting.
from torchao.quantization import quantize_
from torchao.prototype.mx_formats import NVFP4WeightOnlyConfig
from .diffusion_transformer_quant import DEFAULT_MIN_LINEAR_FEATURES, make_filter_fn

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

@ -27,11 +27,10 @@ from typing import Any, Optional
# torch.save dict layout tag; bump on an on-disk change so old/foreign artifacts are rejected.
PREQUANT_FORMAT = "unsloth_prequant_transformer_state_dict_v1"
# Loading ends in ``torch.load(weights_only=False)``, which executes pickle code. A hosted
# repo checkpoint is first-party; a ``kind == "path"`` can come from a request's
# ``transformer_prequant_path``, so unpickling it is RCE. A request-supplied path is
# unpickled ONLY when it resolves inside an operator-configured ALLOWLIST of directories; a
# bare on/off toggle is never a wildcard. The hosted-repo path is unaffected.
# Loading ends in ``torch.load(weights_only=False)``, which executes pickle code. A hosted repo
# checkpoint is first-party; a ``kind == "path"`` can come from a request's
# ``transformer_prequant_path``, so it is unpickled ONLY when it resolves inside an
# operator-configured ALLOWLIST of directories. A bare on/off toggle is never a wildcard.
ALLOW_LOCAL_PREQUANT_PATH_ENV = "UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH"
_PREQUANT_TOGGLE_TOKENS = {"1", "true", "yes", "on", "0", "false", "no", "off"}
@ -185,8 +184,8 @@ def load_prequantized_transformer(
dense-quantise. Best-effort: never raises for an unavailable artifact.
"""
try:
# weights_only=False executes pickle code, so a local path is unpickled ONLY when
# allowlisted. The hosted family repo is first-party and always allowed.
# weights_only=False executes pickle code, so a local path is unpickled ONLY when allowlisted.
# The hosted family repo is first-party and always allowed.
if source.kind == "path" and not _local_prequant_path_allowed(source.location):
_warn(
logger,
@ -205,9 +204,8 @@ def load_prequantized_transformer(
import torch
# torchao weight subclasses aren't safetensors-serializable, so the checkpoint is a
# torch.save pickle; weights_only=False rebuilds those subclasses. Local path gated
# above; repo branch is first-party.
# torchao weight subclasses aren't safetensors-serializable, so the checkpoint is a torch.save
# pickle and weights_only=False rebuilds them. Local path gated above.
ckpt = torch.load(path, weights_only = False, map_location = "cpu")
if not _validate_checkpoint(
ckpt, scheme, base, logger, min_features = min_features, fast_accum = fast_accum
@ -220,21 +218,19 @@ def load_prequantized_transformer(
with init_empty_weights():
transformer = transformer_cls.from_config(config)
# assign=True swaps in the loaded tensors rather than copying into meta (a copy into
# meta is a no-op); strict=True since the saved dict is the full state dict of the
# same class.
# assign=True swaps in the loaded tensors rather than copying into meta (a no-op); strict=True
# since the saved dict is the full state dict of the same class.
transformer.load_state_dict(state_dict, strict = True, assign = True)
if _has_meta_tensors(transformer):
# Non-persistent buffers (built in __init__, absent from the state dict) stay on
# meta. Rebuild on CPU so they hold real values, then re-assign the quantized
# weights; dense bf16 lives in CPU RAM only, the GPU gets just the quant footprint.
# Non-persistent buffers (built in __init__, absent from the state dict) stay on meta. Rebuild on
# CPU so they hold real values, then re-assign the quantized weights; dense bf16 lives in CPU RAM
# only, so the GPU gets just the quant footprint.
transformer = transformer_cls.from_config(config)
transformer.load_state_dict(state_dict, strict = True, assign = True)
transformer = transformer.to(device)
# from_config starts in TRAIN mode; the dense/GGUF paths use from_pretrained, which
# returns an eval()'d module. Match that so train/eval-sensitive layers (e.g.
# dropout) can't make prequant inference diverge from the other paths.
# from_config starts in TRAIN mode while the dense/GGUF paths use from_pretrained (eval()'d).
# Match that so train/eval-sensitive layers can't make prequant inference diverge.
try:
transformer.eval()
except Exception: # noqa: BLE001 — eval() is best-effort
@ -313,9 +309,8 @@ def _validate_checkpoint(
if meta.get("scheme") != scheme:
_warn(logger, scheme, ValueError(f"checkpoint scheme {meta.get('scheme')!r} != {scheme!r}"))
return False
# fp8 REQUIRES per-row granularity (per-tensor collapses outlier-heavy DiTs to noise). An
# old checkpoint omits ``fp8_granularity`` or records non-per-row; reject so the loader
# re-quantises instead of installing a broken fp8 transformer.
# fp8 REQUIRES per-row granularity (per-tensor collapses outlier-heavy DiTs to noise). An old
# checkpoint omits ``fp8_granularity`` or records non-per-row; reject so the loader re-quantises.
from .diffusion_transformer_quant import FP8_GRANULARITY, TQ_FP8
if scheme == TQ_FP8 and meta.get("fp8_granularity") != FP8_GRANULARITY:
@ -330,9 +325,9 @@ def _validate_checkpoint(
return False
ckpt_base = meta.get("base_model_id")
if base:
# Keys matching a different base can load strict=True and generate from the wrong
# weights. Our builder always records base_model_id, so one that omits it against a
# requested base is untrustworthy -- refuse it.
# Keys matching a different base can load strict=True and generate from the wrong weights. Our
# builder always records base_model_id, so one that omits it against a requested base is
# untrustworthy.
if not ckpt_base:
_warn(
logger,
@ -354,16 +349,15 @@ def _validate_checkpoint(
ValueError(f"checkpoint min_features {ckpt_min!r} != runtime {min_features!r}"),
)
return False
# The int8 exclusion set is scheme-derived, but a future change to the token list would
# leave old checkpoints with a stale baked set that passes scheme+min_features then
# crashes at the first denoise. Reject a recorded mismatch; absent is accepted.
# The int8 exclusion set is scheme-derived, but a future change to the token list would leave old
# checkpoints with a stale baked set that passes scheme+min_features then crashes at the first
# denoise. Reject a recorded mismatch; absent is accepted.
ckpt_excludes = meta.get("exclude_name_tokens")
if ckpt_excludes is not None:
from .diffusion_transformer_quant import exclude_tokens_for_scheme
# The exclude set derives from scheme AND family; use the recorded family so an artifact
# baked under an older token list (e.g. a Qwen int8 checkpoint from before the
# text-stream exclude) is rejected and re-quantised, not loaded crashing.
# The exclude set derives from scheme AND family; use the recorded family so an artifact baked
# under an older token list is rejected and re-quantised, not loaded crashing.
expected = tuple(exclude_tokens_for_scheme(scheme, meta.get("family")))
if tuple(ckpt_excludes) != expected:
_warn(
@ -374,9 +368,8 @@ def _validate_checkpoint(
),
)
return False
# require_bf16 (skip non-bf16 Linears) is scheme-pinned (fp8/mxfp8 need bf16; nvfp4/int8
# take fp32). Recording and verifying it guards against a future _REQUIRE_BF16_SCHEMES
# change loading an old-filter checkpoint (different quantised layer set). Absent accepted.
# require_bf16 (skip non-bf16 Linears) is scheme-pinned. Recording and verifying it guards
# against a future _REQUIRE_BF16_SCHEMES change loading an old-filter checkpoint. Absent accepted.
ckpt_require_bf16 = meta.get("require_bf16")
if ckpt_require_bf16 is not None:
from .diffusion_transformer_quant import _REQUIRE_BF16_SCHEMES

82
studio/backend/core/inference/diffusion_speed.py
View file

@ -187,7 +187,7 @@ def apply_speed_optims(
}
mode = normalize_speed_mode(speed_mode)
# TF32 (max) and cudnn.benchmark (any non-off CUDA load) are process-global; the caller
# snapshots/restores them via snapshot_backend_flags so a later `off` load never inherits them.
# snapshots/restores them so a later `off` load never inherits them.
if mode == SPEED_OFF:
return applied
@ -197,16 +197,16 @@ def apply_speed_optims(
# Lossless: a channels-last VAE speeds up its convs with no numeric change.
applied["channels_last"] = _vae_channels_last(pipe, logger)
# Near-lossless: cuDNN autotunes the fixed-shape VAE convs (CUDA only). May pick a different
# Near-lossless: cuDNN autotunes the fixed-shape VAE convs (CUDA only). It may pick a different
# conv algorithm, so it is a "default"-tier (not bit-identical) win.
if on_cuda:
applied["cudnn_benchmark"] = _enable_cudnn_benchmark(logger)
# Consumer-only: fp16 GEMMs accumulate in fp16 (~2x on GeForce-class parts; datacenter HBM
# parts gain nothing and keep fp32 accumulate). bf16 loads measured bit-identical with the flag
# on (36/36 same-seed cases), so on the neutral tiers it engages only when compute dtype is NOT
# fp16. fp16 pipelines showed same-seed drift (mean 2-5% on SDXL/FLUX), so fp16 compute gets it
# only under ``max``. Guarded by _FP16_ACCUM_DENY and the UNSLOTH_DISABLE_FP16_ACCUM kill switch.
# Consumer-only: fp16 GEMMs accumulate in fp16 (~2x on GeForce-class parts; datacenter parts
# gain nothing). bf16 loads measured bit-identical with the flag on (36/36 same-seed cases), so
# on the neutral tiers it engages only when compute dtype is NOT fp16. fp16 pipelines showed
# same-seed drift (mean 2-5%), so fp16 compute gets it only under ``max``. Guarded by
# _FP16_ACCUM_DENY and the UNSLOTH_DISABLE_FP16_ACCUM kill switch.
if on_cuda:
applied["fp16_accum"] = _enable_fp16_accumulation(
family, logger, dtype = getattr(target, "dtype", None), speed_mode = mode
@ -214,18 +214,15 @@ def apply_speed_optims(
# --- the compile lever, per tier ---
# default = LIGHT: GGUF compiles ONLY the dequant op chain (cheap, VRAM-free,
# resolution-invariant; block stays eager); dense has no dequant, so falls back to the
# regional block compile.
# max = FULL: regional max-autotune compile of the repeated block (subsumes the dequant fusion,
# so no standalone compiled dequant here).
# resolution-invariant); dense has no dequant, so falls back to the regional block compile.
# max = FULL: regional max-autotune compile of the repeated block (subsumes the dequant fusion).
# eager = no compile.
if mode == SPEED_DEFAULT:
if is_gguf and on_cuda and family_allows_compile:
applied["compiled_dequant"] = gguf_compile.install_compiled_dequant(logger)
elif compile_eligible(target, is_gguf = is_gguf, family = family):
# A U-Net (SDXL) fuses QKV BEFORE its whole-module compile: 36.3 vs 39.3 ms/step
# (LPIPS 0.033). DiTs were neutral under the regional compile (Qwen-Image 6.53 vs
# 6.52 s), so they keep the fuse on the max tier only.
# A U-Net (SDXL) fuses QKV BEFORE its whole-module compile: 36.3 vs 39.3 ms/step (LPIPS 0.033).
# DiTs were neutral under the regional compile, so they keep the fuse on the max tier only.
if _denoiser_unet(pipe) is not None:
applied["fused_qkv"] = _fuse_qkv(pipe, logger)
applied["compiled"] = _compile_repeated_blocks(
@ -245,8 +242,8 @@ def apply_speed_optims(
)
# A compiled U-Net family also compiles the VAE decode (a real share at SDXL's step rate:
# 4.98 -> 4.25 s over 4 images, LPIPS unchanged). DiT families skip it (a few % of their
# generation). dynamic=True keeps it resolution-robust; fullgraph=False tolerates offload hooks.
# 4.98 to 4.25 s over 4 images, LPIPS unchanged). DiT families skip it. dynamic=True keeps it
# resolution-robust; fullgraph=False tolerates offload hooks.
if applied["compiled"] and _denoiser_unet(pipe) is not None:
applied["compiled_vae_decode"] = _compile_vae_decode(pipe, logger)
@ -274,10 +271,10 @@ def _vae_channels_last(pipe: Any, logger: Any) -> bool:
# U-Net denoisers ship no ``_repeated_blocks`` (heterogeneous block mix), so the regional compile
# can't reach them; these classes get a WHOLE-module STATIC ``torch.compile`` instead. On SDXL
# (B200, 30 steps / 1024px): static whole-UNet runs 26.9 ms/step vs the 45.9 ms/step bit-exact
# reference -- 1.61x end-to-end (6.16 -> 3.83 s) at LPIPS 0.034 -- while dynamic=True compiles 5x
# slower for less win, and a regional block compile only reaches 45.0 ms/step. Static shapes mean
# a recompile per new (height, width, batch); the Mega-cache bundle carries each across restarts.
# (B200, 30 steps / 1024px): 26.9 ms/step vs the 45.9 ms/step bit-exact reference, 1.61x
# end-to-end at LPIPS 0.034, while dynamic=True compiles 5x slower for less win. Static shapes
# mean a recompile per new (height, width, batch); the Mega-cache bundle carries each across
# restarts.
_UNET_WHOLE_COMPILE: frozenset[str] = frozenset({"UNet2DConditionModel"})
@ -334,8 +331,8 @@ def _compile_repeated_blocks(
# the regional block (its static output buffer is overwritten across steps).
#
# fullgraph drops to False under a step cache OR offloading: both insert an
# ``@torch.compiler.disable``d function (FBCache's per-step decision, offload's onload hook),
# i.e. a graph break fullgraph=True rejects. The break is cheap; the rest still compiles.
# ``@torch.compiler.disable``d function, i.e. a graph break fullgraph=True rejects. The break is
# cheap; the rest still compiles.
kwargs: dict[str, Any] = {
"fullgraph": not (cache_active or offload_active),
"dynamic": not max_autotune,
@ -345,21 +342,20 @@ def _compile_repeated_blocks(
try:
import torch
# Heterogeneous-block DiTs (e.g. Z-Image) compile ~one graph per distinct block shape;
# Z-Image needs ~11, above dynamo's default recompile_limit of 8. Past the limit a resident
# load hard-errors under fullgraph (offload/cache silently drops blocks to eager), so raise
# it to 64 (diffusers' documented regional-compile fix). NOT
# force_parameter_static_shapes=False: no variant-count win here and ~6x slower (24 -> 143s).
# Heterogeneous-block DiTs (e.g. Z-Image) compile ~one graph per distinct block shape, and
# Z-Image needs ~11, above dynamo's default recompile_limit of 8. Past the limit a resident load
# hard-errors under fullgraph, so raise it to 64 (diffusers' documented regional-compile fix).
# NOT force_parameter_static_shapes=False: no variant-count win and ~6x slower.
dynamo_cfg = getattr(getattr(torch, "_dynamo", None), "config", None)
if dynamo_cfg is not None:
for _limit_attr in ("recompile_limit", "cache_size_limit"): # name varies by torch ver
if hasattr(dynamo_cfg, _limit_attr):
setattr(dynamo_cfg, _limit_attr, max(getattr(dynamo_cfg, _limit_attr) or 0, 64))
# Match eager's intermediate rounding in inductor's fused pointwise kernels: they keep
# chains in fp32 where eager materialises bf16 between ops, a per-forward delta a multi-step
# denoise amplifies. Measured LPIPS vs eager: Qwen-Image 0.019 -> 0.006, FLUX.1-dev
# 0.046 -> 0.029 (+2% step), FLUX.2-klein 0.018 -> 0.017, HunyuanVideo-1.5-720p 0.221 ->
# 0.052, all at ~zero cost. Process-global, so snapshot_backend_flags restores it on unload.
# Match eager's intermediate rounding in inductor's fused pointwise kernels: they keep chains in
# fp32 where eager materialises bf16 between ops, a per-forward delta a multi-step denoise
# amplifies. Measured LPIPS vs eager: Qwen-Image 0.019 to 0.006, FLUX.1-dev 0.046 to 0.029 (+2%
# step), FLUX.2-klein 0.018 to 0.017, HunyuanVideo-1.5-720p 0.221 to 0.052, all at ~zero cost.
# Process-global, so snapshot_backend_flags restores it on unload.
inductor_cfg = _inductor_config()
if inductor_cfg is not None and hasattr(inductor_cfg, "emulate_precision_casts"):
inductor_cfg.emulate_precision_casts = True
@ -368,10 +364,9 @@ def _compile_repeated_blocks(
return False
if unet is not None:
# Whole-module static compile for the U-Net classes above. fullgraph mirrors the regional
# decision (in practice only offload lowers it, U-Nets have no CacheMixin); dynamic is
# ALWAYS False, so each new (height, width, batch) pays its own compile (Mega-cache carries
# it across restarts). ``Module.compile`` keeps the module identity, so unload/status/LoRA
# see the same object.
# decision (in practice only offload lowers it); dynamic is ALWAYS False, so each new
# (height, width, batch) pays its own compile. ``Module.compile`` keeps the module identity, so
# unload/status/LoRA see the same object.
unet_kwargs: dict[str, Any] = {"fullgraph": kwargs["fullgraph"], "dynamic": False}
if max_autotune:
unet_kwargs["mode"] = "max-autotune-no-cudagraphs"
@ -392,10 +387,9 @@ def _compile_repeated_blocks(
_warn(logger, "compile_repeated_blocks", exc)
continue
# A step cache engaged BEFORE this compile has already wrapped each block's forward in a
# @torch.compiler.disable'd hook, so the compute branch would run eager and forfeit the
# regional compile. Re-point the hooks' inner forward at compiled wrappers (no-op with no
# cache hooks); the toggle path is armed by apply_step_cache. Lazy import keeps the
# dependency one-directional.
# @torch.compiler.disable'd hook, so the compute branch would run eager and forfeit the regional
# compile. Re-point the hooks' inner forward at compiled wrappers (no-op without cache hooks);
# the toggle path is armed by apply_step_cache.
try:
from .diffusion_cache import _compile_hooked_block_inners
_compile_hooked_block_inners(transformer, logger)
@ -443,8 +437,8 @@ def _enable_tf32(logger: Any) -> bool:
# Families the overflow harness found to produce non-finite activations / NEW black frames under
# fp16 accumulation. Empty by measurement: no overflow across all six families (bf16 bit-identical,
# fp16 finite; the fp16 same-seed drift is why fp16 compute is gated to ``max`` below).
# fp16 accumulation. Empty by measurement: no overflow across all six families (bf16
# bit-identical, fp16 finite; the fp16 same-seed drift is why fp16 compute is gated to ``max``).
_FP16_ACCUM_DENY: frozenset[str] = frozenset()
@ -492,8 +486,8 @@ def _enable_fp16_accumulation(
def _fuse_qkv(pipe: Any, logger: Any) -> bool:
# Prefer the pipe-level fuse (covers every component); else fuse each denoiser DiT so a
# dual-DiT family fuses BOTH experts.
# Prefer the pipe-level fuse (covers every component); else fuse each denoiser DiT so a dual-DiT
# family fuses BOTH experts.
fn = getattr(pipe, "fuse_qkv_projections", None)
if callable(fn):
try:

65
studio/backend/core/inference/diffusion_te_prequant.py
View file

@ -41,15 +41,14 @@ TE_PREQUANT_FORMAT = "unsloth_prequant_text_encoder_state_dict_v1"
# The one scheme hosted in v1 (see module docstring).
TE_PREQUANT_SCHEMES = ("fp8",)
# Components the pipeline-assembly injection covers (the attrs quantize_text_encoders
# casts; text_encoder_4 is family-assembled separately, see diffusion_hidream.py).
# Components the pipeline-assembly injection covers (text_encoder_4 is family-assembled
# separately, see diffusion_hidream.py).
TE_PREQUANT_COMPONENTS = ("text_encoder", "text_encoder_2", "text_encoder_3")
# Bases whose text-encoder weights are VERIFIED byte-identical, so one hosted artifact
# serves all of them. Verification: every safetensors shard's LFS sha256 compared across
# repos on 2026-07-18 (huggingface_hub list_repo_tree; no local download needed). The
# checkpoint validator accepts a base_model_id from the same group as the loading base;
# everything else keeps the strict refusal. Ids are lowercased.
# Bases whose text-encoder weights are VERIFIED byte-identical, so one hosted artifact serves
# all of them (every shard's LFS sha256 compared across repos on 2026-07-18). The checkpoint
# validator accepts a base_model_id from the same group; everything else keeps the strict
# refusal. Ids are lowercased.
_TE_EQUIVALENT_BASES: tuple[frozenset[str], ...] = (
# Qwen2.5-VL-7B text encoder: 4 shards, 16,584,414,544 bytes, identical sha256 set.
frozenset(
@ -58,9 +57,9 @@ _TE_EQUIVALENT_BASES: tuple[frozenset[str], ...] = (
"hunyuanvideo-community/hunyuanimage-2.1-diffusers",
}
),
# T5-XXL (text_encoder_2): 2 shards, 9,524,648,584 bytes, identical sha256 set across
# every FLUX.1 release; HiDream-I1 ships the same bytes as text_encoder_3 (cross-
# component filename/metadata mapping is not wired yet, entry documents the identity).
# T5-XXL (text_encoder_2): 2 shards, 9,524,648,584 bytes, identical sha256 set across every
# FLUX.1 release; HiDream-I1 ships the same bytes as text_encoder_3 (cross-component mapping is
# not wired yet, this entry documents the identity).
frozenset(
{
"black-forest-labs/flux.1-schnell",
@ -189,10 +188,9 @@ def load_prequant_text_encoder(
import torch
# The layerwise-fp8 state dict is plain tensors (fp8 storage for cast leaves, the
# original dtype for skipped modules), so weights_only=True suffices: no pickle code
# runs even for a local-path artifact. A future torchao-subclass scheme needs a
# format bump AND weights_only=False behind the same allowlist as the DiT module.
# The layerwise-fp8 state dict is plain tensors, so weights_only=True suffices: no pickle code
# runs even for a local-path artifact. A future torchao-subclass scheme needs a format bump AND
# weights_only=False behind the same allowlist as the DiT module.
ckpt = torch.load(path, weights_only = True, map_location = "cpu")
if not _validate_checkpoint(ckpt, scheme, component, base, logger):
return None
@ -214,10 +212,9 @@ def load_prequant_text_encoder(
if subfolder:
config_kwargs["subfolder"] = subfolder
config = transformers.AutoConfig.from_pretrained(base, **config_kwargs)
# Krea-2 ships transformers-5.x configs whose rope lives under rope_parameters;
# the runtime component loader remaps it for a 4.x runtime, and the meta-init
# here must match or the rebuilt encoder forwards with a broken rope. No-op for
# every other family (and on a 5.x runtime).
# Krea-2 ships transformers-5.x configs whose rope lives under rope_parameters; the runtime
# component loader remaps it for a 4.x runtime, and the meta-init here must match or the rebuilt
# encoder forwards with a broken rope. No-op for every other family.
from .diffusion_krea2 import remap_rope_parameters
remap_rope_parameters(getattr(config, "text_config", config))
@ -227,28 +224,26 @@ def load_prequant_text_encoder(
with init_empty_weights():
encoder = encoder_cls(config)
# assign=True swaps in the loaded tensors rather than copying into meta; strict=True
# since the saved dict is the full state dict of the same class.
# assign=True swaps in the loaded tensors rather than copying into meta; strict=True since the
# saved dict is the full state dict of the same class.
encoder.load_state_dict(state_dict, strict = True, assign = True)
if _has_meta_tensors(encoder):
# Non-persistent buffers (built in __init__, absent from the state dict) stay on
# meta. Rebuild on CPU so they hold real values, then re-assign the cast weights.
# Non-persistent buffers (built in __init__, absent from the state dict) stay on meta. Rebuild on
# CPU so they hold real values, then re-assign the cast weights.
encoder = encoder_cls(config)
encoder.load_state_dict(state_dict, strict = True, assign = True)
# assign=True swaps in SEPARATE tensors for tied weights (the saved dict carries a
# copy per key), untying e.g. Qwen3's lm_head from embed_tokens. An untied head
# defeats _cast_fp8's tied-projection skip below (the head would get cast while the
# builder's did not, breaking bit-identity and duplicating the embedding). Re-tie to
# the builder-identical structure; a no-op for untied configs.
# assign=True swaps in SEPARATE tensors for tied weights (the saved dict carries a copy per key),
# untying e.g. Qwen3's lm_head from embed_tokens. An untied head defeats _cast_fp8's
# tied-projection skip below, breaking bit-identity and duplicating the embedding. Re-tie to the
# builder-identical structure; a no-op for untied configs.
tie = getattr(encoder, "tie_weights", None)
if callable(tie):
tie()
encoder.eval()
# Install the SAME upcast hooks the runtime cast applies. The weight cast inside is
# idempotent (fp8 -> fp8), so this only arms the per-layer upcast; without it the
# fp8 storage weights would meet bf16 activations at the first forward. A hook
# failure means the encoder cannot run; fall back to the dense path.
# Install the SAME upcast hooks the runtime cast applies. The weight cast inside is idempotent,
# so this only arms the per-layer upcast; without it the fp8 storage weights would meet bf16
# activations at the first forward. A hook failure means the encoder cannot run.
from .diffusion_precision import _cast_fp8
class _Target:
@ -301,8 +296,8 @@ def te_prequant_pipe_kwargs(
if mode != TE_QUANT_FP8:
return {}
family = getattr(fam, "name", None)
# The per-family TE deny table ships on the video branch's precision module; the
# image branch has no denials. Resolve lazily so one module serves both.
# The per-family TE deny table ships on the video branch's precision module; the image branch has
# no denials. Resolve lazily so one module serves both.
denied = getattr(precision, "_te_family_denied", None)
if callable(denied) and denied(family, mode):
return {}
@ -367,8 +362,8 @@ def _validate_checkpoint(ckpt: Any, scheme: str, component: str, base: str, logg
return False
ckpt_base = meta.get("base_model_id")
if base:
# Keys matching a different base can load strict=True and encode prompts with the
# wrong weights. The builder always records base_model_id; refuse one that omits it.
# Keys matching a different base can load strict=True and encode prompts with the wrong weights.
# The builder always records base_model_id; refuse one that omits it.
if not ckpt_base:
_warn(
logger,

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

@ -33,10 +33,9 @@ TQ_AUTO = "auto"
TQ_SCHEMES = (TQ_INT8, TQ_FP8, TQ_NVFP4, TQ_MXFP8)
TQ_MODES = (TQ_AUTO,) + TQ_SCHEMES
# Schemes whose torchao path asserts a bf16 weight, so their filter must skip non-bf16
# Linears (make_filter_fn's require_bf16) rather than aborting the whole pass on a stray fp32
# Linear (e.g. T5's `wo`). Verified on torchao 0.17 / B200: fp8 per-row and mxfp8 assert bf16;
# nvfp4 and int8 quantise fp32/fp16 fine, so they are not gated (keeps those projections quant).
# Schemes whose torchao path asserts a bf16 weight, so their filter must skip non-bf16 Linears
# rather than aborting the whole pass on a stray fp32 Linear (e.g. T5's `wo`). Verified on
# torchao 0.17 / B200: fp8 per-row and mxfp8 assert bf16; nvfp4 and int8 quantise fp32/fp16 fine.
_REQUIRE_BF16_SCHEMES = (TQ_FP8, TQ_MXFP8)
# fp8 granularity the runtime uses: per-ROW is REQUIRED for correctness on outlier-heavy DiTs.
@ -47,14 +46,13 @@ FP8_GRANULARITY = "per_row"
# Skip linears below this feature size: a small FLOP share, so leaving them bf16 costs ~nothing.
DEFAULT_MIN_LINEAR_FEATURES = 512
# int8-ONLY name exclusions. int8 uses torch._int_mm, which needs activation rows M > 16. A
# DiT's AdaLN modulation projections and timestep / guidance / pooled-text conditioning
# embedders run once from the [batch, dim] vector (M = batch = 1), not per token, so they crash
# _int_mm despite large feature dims (Flux norm1.linear 3072->18432, Qwen img_mod.1, Flux.2
# *_modulation.linear); min_features misses them (they are big), so int8 also skips any Linear
# whose fqn matches a token here. Negligible FLOPs (M=1, once per block), so int8 keeps the full
# speedup on attention/FFN (M = seq). fp8/nvfp4/mxfp8 use scaled_mm (no M limit) and quantise
# these fine, so the exclusion is int8-only. Sequence embedders (M = seq) are NOT excluded.
# int8-ONLY name exclusions. int8 uses torch._int_mm, which needs activation rows M above 16. A
# DiT's AdaLN modulation projections and timestep / guidance / pooled-text conditioning embedders
# run once from the [batch, dim] vector (M = batch = 1), not per token, so they crash _int_mm
# despite large feature dims; min_features misses them (they are big), so int8 also skips any
# Linear whose fqn matches a token here. Negligible FLOPs, so int8 keeps the full speedup on
# attention/FFN (M = seq). fp8/nvfp4/mxfp8 use scaled_mm (no M limit), so this is int8-only.
# Sequence embedders (M = seq) are NOT excluded.
_INT8_EXCLUDE_NAME_TOKENS = (
"norm", # AdaLN modulation .linear
"_mod", # Qwen img_mod / txt_mod
@ -63,19 +61,17 @@ _INT8_EXCLUDE_NAME_TOKENS = (
"guidance_embed",
"time_text_embed", # Flux/Qwen (pooled-text + timestep); NOT context_embedder
"pooled",
# Krea 2's time_embed.linear_2 (6144->6144, M = batch); its time_mod_proj is caught by
# "_mod", and img_in / final_layer.linear / text_fusion.projector fall under min_features.
# Krea 2's time_embed.linear_2 (M = batch); its time_mod_proj is caught by "_mod", and img_in /
# final_layer.linear / text_fusion.projector fall under min_features.
"time_embed",
)
# int8 PER-FAMILY name exclusions, on top of _INT8_EXCLUDE_NAME_TOKENS. Qwen-Image's MMDiT
# runs every TEXT-stream Linear at M = actual prompt tokens (the Qwen2.5-VL embeds are not
# padded to a fixed length like FLUX's 512-token T5), so a short prompt ("Cute sloth writing
# on a paper" = 13 tokens, or the near-empty negative prompt) drives torch._int_mm below its
# M > 16 floor and the denoise crashes (measured on B200: "self.size(0) needs to be greater
# than 16, but got 13"). Keep the text stream bf16: it runs at M = tens vs the image stream's
# M ~ 4k+, so the exclusion costs ~nothing and the image stream keeps full int8 coverage.
# int8 PER-FAMILY name exclusions, on top of _INT8_EXCLUDE_NAME_TOKENS. Qwen-Image's MMDiT runs
# every TEXT-stream Linear at M = actual prompt tokens (the Qwen2.5-VL embeds are not padded to a
# fixed length like FLUX's 512-token T5), so a short prompt or the near-empty negative prompt
# drives torch._int_mm below its M floor of 16 and the denoise crashes. Keep the text stream
# bf16: it runs at M = tens vs the image stream's M ~ 4k+, so the exclusion costs ~nothing.
# txt_mod is already covered by "_mod" in the base list; txt_in is the context embedder.
_QWENIMAGE_INT8_EXCLUDES = (
"txt_in",
@ -108,11 +104,10 @@ def exclude_tokens_for_scheme(scheme: str, family: Optional[str] = None) -> tupl
# Per-arch preference for ``auto`` -- best first, lower-precision schemes as fallbacks. On
# Blackwell fp8 leads: measured on B200, plain fp8 dynamic is faster AND more accurate than the
# alternatives at the DiT's shapes. mxfp8's block scaling adds overhead with no speed win.
# nvfp4 is also below fp8: the FP4 GEMM is real with torch>=2.11 + torchao's CUTLASS kernel
# (16384^3 GEMM ~3826 TFLOPS, 1.37x fp8) but only beats fp8 on very large GEMMs; at the DiT's
# shapes (hidden ~3072, MLP ~12288, M~4096) it is slower (0.81x on Z-Image 1024px) and less
# accurate (LPIPS 0.166 vs fp8 0.044), so it stays an explicit opt-in, never the auto pick.
# alternatives at the DiT's shapes. mxfp8's block scaling adds overhead with no speed win. nvfp4
# is also below fp8: its FP4 GEMM is real with torch>=2.11 + torchao's CUTLASS kernel but only
# wins on very large GEMMs; at DiT shapes it is slower (0.81x on Z-Image 1024px) and less
# accurate (LPIPS 0.166 vs fp8 0.044), so it stays an explicit opt-in.
#
# DATA-CENTER order. On a consumer / workstation GPU int8 moves first (_prefer_consumer_scheme):
# consumer cards halve fp8/fp16 FP32-accumulate, while int8 (int32 accumulate) runs full-rate.
@ -124,12 +119,12 @@ _AUTO_LADDER: tuple[tuple[tuple[int, int], tuple[str, ...]], ...] = (
# Families whose activation ranges break specific schemes at the MODEL level (the smoke probe
# only proves the GEMM runs). Measured with the 28-pair prequant accuracy gate on B200:
# qwen-image + fp8 -> every frame black (luma 0.0000, SSIM 0.016): Qwen's outliers exceed
# qwen-image + fp8 -- every frame black (luma 0.0000, SSIM 0.016): Qwen's outliers exceed
# even per-row fp8's range (the same fp8 matches bf16 on Z-Image/FLUX).
# qwen-image + mxfp8 -> semantic damage at 1024px (CLIP delta mean 0.0146, worst 0.064/0.102).
# qwen-image + nvfp4 -> LPIPS mean 0.51 vs bf16: unusable.
# int8 dynamic is excellent on Qwen (LPIPS 0.069 / SSIM 0.958), so auto falls through to it. The
# deny also applies to an EXPLICIT request (returning None gives the same GGUF fallback).
# qwen-image + mxfp8 -- semantic damage at 1024px (CLIP delta mean 0.0146, worst 0.064/0.102).
# qwen-image + nvfp4 -- LPIPS mean 0.51 vs bf16: unusable.
# int8 dynamic is excellent on Qwen, so auto falls through to it. The deny also applies to an
# EXPLICIT request (returning None gives the same GGUF fallback).
_FAMILY_SCHEME_DENY: dict[str, frozenset[str]] = {
"qwen-image": frozenset({TQ_FP8, TQ_MXFP8, TQ_NVFP4}),
"qwen-image-edit": frozenset({TQ_FP8, TQ_MXFP8, TQ_NVFP4}), # same DiT
@ -144,8 +139,8 @@ def _family_denied(family, scheme: str) -> bool:
_SMOKE_CACHE: dict[tuple[str, str], bool] = {}
# Data-center GPU tokens (un-nerfed FP32 accumulate). Matched as whole tokens of
# get_device_name() so workstation "A4000" isn't mistaken for data-center "A40". Anything else
# is treated as consumer-class (FP32-accumulate halved). See developer.nvidia.com/cuda/gpus.
# get_device_name() so workstation "A4000" isn't mistaken for data-center "A40". Anything else is
# treated as consumer-class (FP32-accumulate halved).
_DATACENTER_GPU_TOKENS = frozenset(
{
"B200",
@ -342,25 +337,21 @@ def _make_quant_config(scheme: str, fast_accum: Optional[bool] = None) -> Any:
return Int8DynamicActivationInt8WeightConfig()
if scheme == TQ_FP8:
# Per-ROW granularity (per-token activation + per-channel weight scale) is REQUIRED for
# correctness. torchao defaults to per-TENSOR (one scale): a DiT with extreme outliers
# breaks -- Z-Image MLP activations peak near 6.6e4, so one outlier forces a tensor-wide
# scale that pushes normal values (~1-30) below fp8 resolution to ~0 and the denoise
# collapses to noise (B200: per-tensor = noise, per-row = matches bf16). Per-row confines
# each outlier to its token/channel (also why int8, per-token by default, was always
# fine). _smoke_probe checks per-row scaled_mm, so a build without it falls to int8.
# correctness. torchao defaults to per-TENSOR (one scale), which breaks a DiT with extreme
# outliers: Z-Image MLP activations peak near 6.6e4, so one outlier forces a tensor-wide scale
# that pushes normal values below fp8 resolution and the denoise collapses to noise. Per-row
# confines each outlier to its token/channel. _smoke_probe checks per-row scaled_mm, so a build
# without it falls to int8.
#
# fast accumulate (fp8 only) is chosen by GPU class unless forced: consumer cards run fp8
# ~2x faster with FP16 accumulate than FP32 (~838 vs ~419 TFLOPS on RTX 50xx); data-center
# keeps precise accumulate.
# fast accumulate (fp8 only) is chosen by GPU class unless forced: consumer cards run fp8 ~2x
# faster with FP16 accumulate than FP32; data-center keeps precise accumulate.
#
# activation_value_lb floors the dynamic per-row activation scale: an ALL-ZERO token row
# otherwise yields scale 0 -> NaN qdata -> black frames on torchao's plain-torch kernel
# path (fused fbgemm/mslk quantize kernels clamp internally, which masks the bug on boxes
# that have them). Zero rows are real: Wan 2.2 zero-pads its text conditioning and
# Hunyuan-1.5 / Qwen-Image regenerate zero rows inside their blocks. Weight scales are
# untouched (weights are never all-zero rows in practice and the floor is 1e-12), so
# pre-quantized fp8 checkpoints stay valid. The knob exists since the Float8Tensor rework
# (torchao >= 0.13); older versions keep today's behaviour via the signature check.
# otherwise yields scale 0, NaN qdata and black frames on torchao's plain-torch kernel path
# (fused kernels clamp internally, masking the bug where they exist). Zero rows are real: Wan 2.2
# zero-pads its text conditioning and Hunyuan-1.5 / Qwen-Image regenerate zero rows inside their
# blocks. Weight scales are untouched, so pre-quantized fp8 checkpoints stay valid. The knob
# exists since the Float8Tensor rework (torchao >= 0.13); older versions keep today's behaviour.
import inspect
from torchao.quantization import PerRow
@ -368,12 +359,11 @@ def _make_quant_config(scheme: str, fast_accum: Optional[bool] = None) -> Any:
config_params = inspect.signature(Float8DynamicActivationFloat8WeightConfig).parameters
if "activation_value_lb" in config_params:
fp8_kwargs["activation_value_lb"] = 1e-12
# Pin the plain-torch quantize kernel. The default AUTO silently switches to the MSLK
# kernel whenever an mslk package is importable (sm90+), which changes fp8 scale
# rounding BITWISE (measured: 8/8 FLUX matrices differ, scales ~55% of bytes) -- so a
# box that merely gains mslk would break the hosted-prequant bit-identity invariant.
# Measured on B200 the mslk path is also SLOWER compiled (opaque extern call blocks
# inductor's quantize fusion: FLUX.1 fp8 e2e 1.149 -> 1.624 s), so the pin costs nothing.
# Pin the plain-torch quantize kernel. The default AUTO silently switches to the MSLK kernel
# whenever an mslk package is importable (sm90+), which changes fp8 scale rounding BITWISE, so a
# box that merely gains mslk would break the hosted-prequant bit-identity invariant. Measured on
# B200 the mslk path is also slower compiled (an opaque extern call blocks inductor's quantize
# fusion: FLUX.1 fp8 e2e 1.149 to 1.624 s), so the pin costs nothing.
if "kernel_preference" in config_params:
try:
from torchao.quantization.quantize_.common.kernel_preference import (
@ -393,9 +383,9 @@ def _make_quant_config(scheme: str, fast_accum: Optional[bool] = None) -> Any:
if scheme == TQ_NVFP4:
from torchao.prototype.mx_formats import NVFP4DynamicActivationNVFP4WeightConfig
# Select the CUTLASS FP4 path, not the default Triton kernel (use_triton_kernel=True
# needs MSLK): on a Blackwell box with CUTLASS FP4 but no MSLK, the default fails the
# smoke probe and falls back to GGUF instead of using the FP4 tensor cores.
# Select the CUTLASS FP4 path, not the default Triton kernel (use_triton_kernel=True needs MSLK):
# on a Blackwell box with CUTLASS FP4 but no MSLK, the default fails the smoke probe and falls
# back to GGUF instead of using the FP4 tensor cores.
try:
return NVFP4DynamicActivationNVFP4WeightConfig(use_triton_kernel = False)
except TypeError: # older torchao without the knob
@ -491,19 +481,16 @@ def quantize_transformer(
try:
from torchao.quantization import quantize_
# int8 skips the M=1 projections; scaled_mm schemes have no M limit but fp8/mxfp8 assert
# a bf16 weight, so on a mixed-precision DiT (Wan/Hunyuan) they must skip non-bf16 ones or
# the pass raises. nvfp4 quantises fp32 fine, so it is not gated (see _REQUIRE_BF16_SCHEMES).
# "lora_" keeps a baked adapter's side path (lora_A/lora_B/lora_embedding) high
# precision when adapters were attached before this pass; the tiny ranks usually fall
# under min_features anyway, but an explicit token does not depend on the rank. Runtime
# only: NOT part of exclude_tokens_for_scheme, whose list is baked into prequant
# checkpoint metadata (adding it there would reject every existing checkpoint).
# int8 skips the M=1 projections; scaled_mm schemes have no M limit but fp8/mxfp8 assert a bf16
# weight, so on a mixed-precision DiT (Wan/Hunyuan) they must skip non-bf16 ones or the pass
# raises. nvfp4 quantises fp32 fine, so it is not gated. "lora_" keeps a baked adapter's side
# path high precision when adapters were attached before this pass. Runtime only: NOT part of
# exclude_tokens_for_scheme, whose list is baked into prequant metadata (adding it there would
# reject every existing checkpoint).
exclude = exclude_tokens_for_scheme(scheme, family) + ("lora_",)
# GEMM tiling floors per scheme (see make_filter_fn): scaled_mm needs 16-aligned dims
# (fp8/nvfp4), MX block scaling needs 32. int8's _int_mm has no such floor and keeps
# the historical filter -- and DiT dims are 16/32-divisible in practice, so existing
# prequant checkpoints quantise the same layer set as before.
# GEMM tiling floors per scheme (see make_filter_fn): scaled_mm needs 16-aligned dims, MX block
# scaling needs 32. int8's _int_mm has no such floor and keeps the historical filter -- and DiT
# dims are 16/32-divisible in practice, so existing prequant checkpoints are unaffected.
divisible = {TQ_FP8: 16, TQ_NVFP4: 16, TQ_MXFP8: 32}.get(scheme, 0)
quantize_(
transformer,

26
studio/backend/core/inference/gpu_arbiter.py
View file

@ -35,34 +35,32 @@ def _evict_chat() -> None:
from core.inference.llama_cpp import chat_load_active
llama = get_llama_cpp_backend()
# is_active (process exists), not is_loaded (exists AND healthy): a chat model still starting
# up holds VRAM but isn't healthy, so is_loaded would skip it and let the load race diffusion.
# chat_load_active too: an HF load has no process until its GGUF downloaded, so is_active
# alone found nothing to cancel and let that load spawn onto the GPU we just granted away.
# unload_model sets the cancel event the download loop polls, so the pending load aborts.
# is_active (process exists), not is_loaded (exists AND healthy): a chat model still starting up
# holds VRAM but isn't healthy, so is_loaded would skip it and let the load race diffusion.
# chat_load_active too: an HF load has no process until its GGUF downloaded, so is_active alone
# found nothing to cancel. unload_model sets the cancel event the download loop polls, so the
# pending load aborts.
if llama.is_active or chat_load_active():
llama.unload_model()
orchestrator = get_inference_backend()
if orchestrator.active_model_name:
orchestrator.unload_model(orchestrator.active_model_name)
# An in-flight safetensors load has no active_model_name yet (it is published only once the
# worker reports success), so the unload above misses it and the load would finish onto the
# GPU we just granted away. cancel_load discards the loading marker BEFORE tearing the worker
# down, so a load parked between retries (or in _wait_response) observes the removal and
# aborts instead of publishing. It runs off the lifecycle gate, which the load itself holds
# for its whole duration, so this cannot deadlock.
# An in-flight safetensors load has no active_model_name yet (published only once the worker
# reports success), so the unload above misses it and it would finish onto the GPU we just
# granted away. cancel_load discards the loading marker BEFORE tearing the worker down, so a load
# parked between retries observes the removal and aborts. It runs off the lifecycle gate, which
# the load itself holds throughout, so this cannot deadlock.
for pending in list(getattr(orchestrator, "loading_models", ()) or ()):
orchestrator.cancel_load(pending)
# Kill the subprocess too: its base CUDA context holds VRAM diffusion needs.
orchestrator._shutdown_subprocess(timeout = 5.0)
# The driver reclaims the killed VRAM asynchronously; wait for it to settle before diffusion
# allocates, else a warm chat->diffusion handoff can transiently OOM.
# allocates, else a warm chat-to-diffusion handoff can transiently OOM.
llama._wait_for_vram_settle(since_kill = time.monotonic())
def _evict_diffusion() -> None:
# Unload whichever engine the router has active (diffusers or native sd.cpp), so a
# chat acquire frees the right one.
# Unload whichever engine the router has active (diffusers or native sd.cpp).
from core.inference.diffusion_engine_router import get_active_diffusion_engine
get_active_diffusion_engine().unload()

13
studio/backend/core/inference/image_gallery.py
View file

@ -70,7 +70,7 @@ def save(image: Any, meta: dict[str, Any]) -> dict[str, Any]:
directory = gallery_dir()
final_path = directory / f"{image_id}.png"
# Write to a dotted temp (skipped by the *.png glob) then atomically rename, so a crash mid-write
# never leaves a truncated {id}.png that the listing would surface as a corrupt record.
# never leaves a truncated {id}.png the listing would surface as a corrupt record.
tmp_path = directory / f".{image_id}.png.tmp"
try:
tmp_path.write_bytes(_png_bytes(image, meta))
@ -176,10 +176,9 @@ def list_images(
return []
paths.sort(key = _mtime, reverse = True)
# Page over READABLE records, not raw files: filtering a foreign PNG out of an already-sliced
# window would drop valid images and make has_more wrong. Read only as far as needed.
# Known limit: this re-reads headers from newest down to `offset+limit` per page, so a deep
# scroll is O(offset) header-opens. PIL opens are lazy (header only) and off the event loop, so
# no freeze; a later phase can switch to cursor-based paging if it bites.
# window would drop valid images and make has_more wrong. Read only as far as needed. Known
# limit: this re-reads headers from newest down to `offset+limit` per page, so a deep scroll is
# O(offset) header-opens; PIL opens are lazy and off the event loop, so nothing freezes.
want = None if limit is None else offset + limit
records = []
for path in paths:
@ -199,8 +198,8 @@ def delete(image_id: str) -> bool:
path = image_path(image_id)
if path is None:
return False
# Only delete files we own (a readable recipe chunk); a hand-dropped foreign PNG is invisible
# to list_images, so a guessed id must not destroy it.
# Only delete files we own (a readable recipe chunk); a hand-dropped foreign PNG is invisible to
# list_images, so a guessed id must not destroy it.
if _read_meta(path) is None:
return False
try:

8
studio/backend/core/inference/llama_cpp.py
View file

@ -1383,8 +1383,8 @@ def gguf_load_in_flight(hf_repo: Optional[str]):
# Chat loads in flight, repo-agnostic (local paths and safetensors included). The repo-keyed
# counter above answers the download manager; this one answers the GPU arbiter, which has to know
# a chat load exists before llama-server is spawned.
# counter above answers the download manager; this one answers the GPU arbiter, which must know a
# chat load exists before llama-server is spawned.
_CHAT_LOADS_IN_FLIGHT = 0
@ -1430,8 +1430,8 @@ def zero_vram_chat_load(
"""
if gpu_memory_mode != "manual" or gpu_layers != 0:
return False
# Any speculative mode may launch a GPU drafter; only the request's own knobs are known
# here, so treat every non-empty selection as GPU-bearing rather than guess.
# Any speculative mode may launch a GPU drafter, and only the request's own knobs are known here,
# so treat every non-empty selection as GPU-bearing.
if needs_mmproj or speculative_type:
return False
if LlamaCppBackend._is_vulkan_backend():

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

@ -40,7 +40,7 @@ def text_encoder_flags_for_family(family_name: str) -> tuple[str, ...]:
# sd-cli's image-gen mode (``img_gen``; older ``txt2img``). img2img is the same mode with
# --init-img, so this one token covers both text- and image-conditioned generation.
# --init-img, so this one token covers both.
DEFAULT_MODE = "img_gen"
@ -99,10 +99,10 @@ class SdCppUpscaleParams:
tile_size: Optional[int] = None
# Native (sd.cpp) speed profiles (engine-side analogue of diffusion_speed). off: nothing. default:
# --diffusion-fa (flash attention) + --diffusion-conv-direct (numerically exact). On the CPU tier
# this serves, direct conv measured z-image Q8_0 sampling 56.1 -> 51.3s (~9%) with decode/RSS
# unchanged, so it's in the default profile. max keeps it (profiles are a superset chain).
# Native (sd.cpp) speed profiles (engine-side analogue of diffusion_speed). off: nothing.
# default: --diffusion-fa + --diffusion-conv-direct (numerically exact). On the CPU tier this
# serves, direct conv measured z-image Q8_0 sampling 56.1 to 51.3s (~9%) with decode/RSS
# unchanged. max keeps it (profiles are a superset chain).
NATIVE_SPEED_OFF = "off"
NATIVE_SPEED_DEFAULT = "default"
NATIVE_SPEED_MAX = "max"
@ -211,8 +211,8 @@ def build_sd_cpp_command(
cmd += ["--lora-model-dir", params.lora_dir]
if params.lora_apply_mode:
cmd += ["--lora-apply-mode", params.lora_apply_mode]
# Emit explicit dims when given. An image-conditioned run that leaves them unset omits the
# flags so sd.cpp derives the size from the input; a plain txt2img keeps the 1024 default.
# Emit explicit dims when given. An image-conditioned run that leaves them unset omits the flags
# so sd.cpp derives the size from the input; a plain txt2img keeps the 1024 default.
if params.width is not None or params.height is not None:
w = int(params.width) if params.width is not None else 1024
h = int(params.height) if params.height is not None else 1024
@ -262,7 +262,7 @@ def build_sd_cpp_upscale_command(
raise ValueError("input_image is required for upscale")
if not params.upscale_model:
raise ValueError("upscale_model is required for upscale")
# Reject repeats < 1 explicitly: a truthiness guard would swallow repeats=0 into sd-cli's
# Reject repeats below 1 explicitly: a truthiness guard would swallow repeats=0 into sd-cli's
# one-pass default, quietly changing the caller's intent.
if params.repeats < 1:
raise ValueError("repeats must be >= 1 for upscale")
@ -405,7 +405,7 @@ def build_img_gen_request(
if sample_params:
req["sample_params"] = sample_params
# Structured LoRA list: the API resolves each ``path`` against the server's ``--lora-model-dir``
# (prompt-embedded ``<lora:>`` tags are unsupported server-side), so LoRAs are staged and named here.
# (prompt-embedded ``<lora:>`` tags are unsupported server-side), so LoRAs are staged here.
if lora:
req["lora"] = lora
return req

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

@ -70,8 +70,8 @@ from utils.subprocess_compat import windows_hidden_subprocess_kwargs
logger = get_logger(__name__)
# A sampling-progress line ("4/4", "[ 12/ 28]", "sampling: 50%|...| 14/28"). Only a match
# whose denominator equals the requested step count is trusted, so a stray "1/100" can't move the bar.
# A sampling-progress line ("4/4", "[ 12/ 28]", "sampling: 50%|...| 14/28"). Only a match whose
# denominator equals the requested step count is trusted, so a stray "1/100" can't move the bar.
_STEP_RE = re.compile(r"(\d+)\s*/\s*(\d+)")
# Serialises the one-time binary install so concurrent first-loads don't race.
@ -118,8 +118,8 @@ def _server_binary_runnable(binary: str) -> bool:
return False # cannot exec at all (wrong arch / no execute bit / missing loader)
except Exception: # noqa: BLE001 -- don't block on a flaky probe (timeout etc.)
return True
# Negative return code = signal death (e.g. -4 SIGILL from an incompatible prebuilt on
# an older CPU): launches then crashes, so treat as unavailable and fall back to diffusers.
# Negative return code = signal death (e.g. -4 SIGILL from an incompatible prebuilt on an older
# CPU): launches then crashes, so treat as unavailable and fall back to diffusers.
return proc.returncode >= 0 and proc.returncode not in (126, 127)
@ -221,8 +221,8 @@ class _SdState:
# Token kept so LoRA adapters selected at generate time can be fetched from the Hub.
hf_token: Optional[str] = None
# The GGUF basename this load committed, so companion resolution reproduces the load identity:
# some variants pick their encoder by filename (FLUX.2-klein-9B -> Qwen3-8B) and a local
# *klein-9B*.gguf carries that keyword only in the basename, not the repo id.
# some variants pick their encoder by filename and a local *klein-9B*.gguf carries that keyword
# only in the basename, not the repo id.
gguf_filename: Optional[str] = None
@ -297,16 +297,16 @@ class SdCppDiffusionBackend:
self._lock = threading.Lock()
self._generate_lock = threading.Lock()
self._engine = engine # resolved lazily on first load so import stays cheap
# An injected engine (test seam / escape hatch) pins one-shot mode; a fallback-cached
# engine must NOT, so a now-available server can still be used on the next load.
# An injected engine (test seam / escape hatch) pins one-shot mode; a fallback-cached engine
# must NOT, so a now-available server can still be used on the next load.
self._engine_injected = engine is not None
self._state: Optional[_SdState] = None
self._loading: Optional[_SdLoading] = None
self._load_token = 0
self._cancel_event = threading.Event()
self._active_generate_cancel: Optional[threading.Event] = None
# sd-server started for an in-flight load, before it commits to _state; tracked so an
# unload / superseding load can stop it mid-startup instead of waiting out the timeout.
# sd-server started for an in-flight load, before it commits to _state; tracked so an unload /
# superseding load can stop it mid-startup instead of waiting out the timeout.
self._pending_server: Optional[SdCppServer] = None
self._gen: Optional[_SdGen] = None
@ -337,8 +337,8 @@ class SdCppDiffusionBackend:
"""
if self._engine_injected and self._engine is not None:
return "oneshot", None, self._resolve_engine()
# Install the server build matching the resolved backend (ROCm/Vulkan/CUDA), not the
# default CPU build. Lazy import avoids an import cycle with the router.
# Install the server build matching the resolved backend (ROCm/Vulkan/CUDA), not the default CPU
# build. Lazy import avoids an import cycle with the router.
from core.inference.diffusion_engine_router import _install_accelerator_for
accelerator = _install_accelerator_for(
@ -367,8 +367,8 @@ class SdCppDiffusionBackend:
cpu_offload: bool = False,
memory_mode: Optional[str] = None,
speed_mode: Optional[str] = None,
# diffusers-only knobs accepted for a uniform call and ignored (sd.cpp has no
# torchao quant / SDPA dispatcher / fbcache).
# diffusers-only knobs accepted for a uniform call and ignored (sd.cpp has no torchao quant /
# SDPA dispatcher / fbcache).
text_encoder_quant: Optional[str] = None,
transformer_quant: Optional[str] = None,
transformer_quant_fast_accum: Optional[bool] = None,
@ -378,8 +378,8 @@ class SdCppDiffusionBackend:
transformer_cache_threshold: Optional[float] = None,
# Accepted for interface parity; native is GGUF-only (router forces diffusers otherwise).
model_kind: Optional[str] = None,
# Parity with the diffusers engine's load-time LoRA bake; native applies LoRA at
# generation time through sd-cli, so a load-time selection is ignored here.
# Parity with the diffusers engine's load-time LoRA bake; native applies LoRA at generation time
# through sd-cli, so a load-time selection is ignored here.
loras: Optional[list[tuple[str, float]]] = None,
) -> dict[str, Any]:
"""Validate, then fetch assets on a daemon thread. Returns at once."""
@ -389,8 +389,8 @@ class SdCppDiffusionBackend:
raise ValueError(
"gguf_filename is required: the native engine loads single-file GGUF checkpoints only."
)
# Filename-fallback detector (as the route validated) so a local .gguf whose family
# keyword lives only in the basename doesn't dead-end here on a native-routed host.
# Filename-fallback detector (as the route validated) so a local .gguf whose family keyword lives
# only in the basename doesn't dead-end here on a native-routed host.
fam = detect_family_for_pick(repo_id, gguf_filename, family_override)
if fam is None:
raise ValueError(
@ -405,8 +405,8 @@ class SdCppDiffusionBackend:
with self._lock:
if self._loading is not None and self._loading.error is None:
raise RuntimeError("A diffusion load is already in progress.")
# A superseding load must stop any in-flight generation, else the old run can
# still persist an image after the new load starts (matches unload()'s cancel).
# A superseding load must stop any in-flight generation, else the old run can still persist an
# image after the new load starts (matches unload()'s cancel).
if self._active_generate_cancel is not None:
self._active_generate_cancel.set()
self._load_token += 1
@ -455,12 +455,12 @@ class SdCppDiffusionBackend:
_load_token: int,
) -> None:
try:
# Resolve mode (server preferred, one-shot fallback) + binary up front so an
# install / missing-binary failure surfaces before the multi-GB asset pull.
# Resolve mode (server preferred, one-shot fallback) + binary up front so an install /
# missing-binary failure surfaces before the multi-GB asset pull.
mode, server_binary, engine = self._resolve_backend()
if mode == "server":
# Probe the server binary before the pull: a present-but-unrunnable build
# would download everything then fail. Fall back to one-shot if usable.
# Probe the server binary before the pull: a present-but-unrunnable build would download
# everything then fail. Fall back to one-shot if usable.
assert server_binary is not None
if not _server_binary_runnable(server_binary):
logger.warning(
@ -475,8 +475,7 @@ class SdCppDiffusionBackend:
raise RuntimeError("sd-server binary is present but not runnable.")
mode, server_binary, engine = "oneshot", None, self._resolve_engine()
if mode == "oneshot":
# version() is None when a present binary can't run; fail now, not on the
# first generation.
# version() is None when a present binary can't run; fail now, not on the first generation.
assert engine is not None
if engine.version() is None:
raise RuntimeError("sd-cli binary is present but not runnable.")
@ -495,17 +494,17 @@ class SdCppDiffusionBackend:
qwen2vl = paths.get("qwen2vl"),
)
device = resolve_diffusion_device_target().device
# Honor speed everywhere; offload only off-CPU (on CPU weights are resident,
# so the flags are no-ops).
# Honor speed everywhere; offload only off-CPU (on CPU weights are resident, so the flags are
# no-ops).
offload: tuple[str, ...] = ()
if device != "cpu":
offload = tuple(offload_flags(_memory_policy(memory_mode, cpu_offload)))
native_speed = _native_speed_for(speed_mode)
# Tear down the old model then commit the new one under _generate_lock: abort and
# WAIT for any generation that started during the download, so a stale run can't
# persist an image afterward and two resident servers never coexist. The lock is
# taken only now (not during the fetch), so the long download never serialises generation.
# Tear down the old model then commit the new one under _generate_lock: abort and WAIT for any
# generation that started during the download, so a stale run can't persist an image and two
# resident servers never coexist. The lock is taken only now, so the download never serialises
# generation.
with self._lock:
if self._load_token != _load_token:
return # superseded / cancelled
@ -523,9 +522,8 @@ class SdCppDiffusionBackend:
if mode == "server":
assert server_binary is not None
server = SdCppServer(server_binary)
# Publish the uncommitted server so unload() / a superseding load can stop
# it mid-startup (stop() aborts the readiness wait) instead of waiting out
# the full startup timeout while holding the generate lock.
# Publish the uncommitted server so unload() / a superseding load can stop it mid-startup
# (stop() aborts the readiness wait) instead of waiting out the full startup timeout.
with self._lock:
self._pending_server = server
try:
@ -705,8 +703,8 @@ class SdCppDiffusionBackend:
if fam.sd_cpp_vae:
repos.append(fam.sd_cpp_vae[0])
# Same per-variant selection as _asset_specs (keyed on repo id AND GGUF filename) so the
# cache-deletion guard protects the encoder repo this load actually downloaded; dropping
# the filename would fall back to the 4B default and protect the wrong repo.
# cache-deletion guard protects the encoder repo this load downloaded; dropping the filename
# would fall back to the 4B default and protect the wrong repo.
repos.extend(
terepo
for terepo, _f, _k in sd_cpp_text_encoders_for(
@ -728,20 +726,19 @@ class SdCppDiffusionBackend:
guidance: float = 0.0,
seed: Optional[int] = None,
batch_size: int = 1,
# Batched prompt/seed lists are diffusers-engine features (one batched DiT forward);
# accepted for interface parity and rejected clearly below (sd-cli renders serially,
# so a "batched" list here would silently be a slow loop).
# Batched prompt/seed lists are diffusers-engine features (one batched DiT forward); accepted for
# interface parity and rejected clearly below, since sd-cli would silently render them serially.
prompts: Optional[list[str]] = None,
seeds: Optional[list[int]] = None,
# Accepted for interface parity; native is text-to-image only, so image-conditioned
# requests are rejected clearly below rather than silently dropped.
# Accepted for interface parity; native is text-to-image only, so image-conditioned requests are
# rejected clearly below rather than silently dropped.
init_image: Optional[str] = None,
mask_image: Optional[str] = None,
strength: Optional[float] = None,
upscale: Optional[float] = None, # needs an init image; rejected by the guard below
reference_images: Optional[list[str]] = None, # GPU/diffusers-only (FLUX.2)
# LoRA (id, weight) pairs; resolved up front then applied per path: prompt tags for
# one-shot sd-cli, structured `lora` for sd-server. None/empty = no LoRA.
# LoRA (id, weight) pairs; resolved up front then applied per path: prompt tags for one-shot
# sd-cli, structured `lora` for sd-server. None/empty = no LoRA.
loras: Optional[list[tuple[str, float]]] = None,
# ControlNet is diffusers-only; rejected by the guard below (accepted for parity).
controlnet: Optional[tuple[str, str, str, float, float, float]] = None,
@ -783,8 +780,8 @@ class SdCppDiffusionBackend:
state = self._state
if state is None:
raise RuntimeError(DIFFUSION_NOT_LOADED_MSG)
# A resident server can exit while idle; drop stale state and report not-loaded
# so the client gets the recoverable reload path, not a 500 from img_gen.
# A resident server can exit while idle; drop stale state and report not-loaded so the client
# gets the recoverable reload path, not a 500 from img_gen.
if (
state.mode == "server"
and state.server is not None
@ -793,9 +790,8 @@ class SdCppDiffusionBackend:
self._state = None
raise RuntimeError(DIFFUSION_NOT_LOADED_MSG)
self._active_generate_cancel = cancel
# Publish an active (step 0) state before the slow pre-generate setup (LoRA
# listing/download) so a reload's progress probe doesn't read idle while this
# generation holds _generate_lock and let a second generate queue behind it.
# Publish an active (step 0) state before the slow pre-generate setup (LoRA listing/download) so
# a reload's progress probe doesn't read idle while this generation holds _generate_lock.
# Mirrors DiffusionBackend.generate; sd-cli progress lines advance this count.
self._gen = _SdGen(total_steps = int(steps))
try:
@ -804,9 +800,8 @@ class SdCppDiffusionBackend:
else:
seed = int(seed)
cfg_scale, flux_guidance = _map_guidance(state.family, guidance)
# Resolve selected LoRAs up front (a bad id -> clear 400 before generating).
# Drop weight-0 rows BEFORE the support gate so a request of only-disabled
# rows stays a no-op even where native LoRA is unsupported.
# Resolve selected LoRAs up front (a bad id gives a clear 400 before generating). Drop weight-0
# rows BEFORE the support gate so an only-disabled request stays a no-op even without native LoRA.
lora_resolved: list = []
active_loras = [(i, w) for (i, w) in (loras or []) if w != 0]
if active_loras:
@ -917,8 +912,8 @@ class SdCppDiffusionBackend:
base_seed = int(seed) & ((1 << 63) - 1)
images: list = []
seeds: list[int] = []
# Stage LoRAs into a per-request subdir of the server's lora-model-dir (so a prior
# request's adapters can't leak in), referenced by path relative to that dir; removed after.
# Stage LoRAs into a per-request subdir of the server's lora-model-dir (so a prior request's
# adapters can't leak in), referenced by path relative to that dir; removed after.
lora_payload: Optional[list[dict]] = None
lora_stage: Optional[Path] = None
if lora_resolved:
@ -1022,8 +1017,8 @@ class SdCppDiffusionBackend:
for index in range(max(1, int(batch_size))):
if cancel.is_set():
raise RuntimeError(DIFFUSION_CANCELLED_MSG)
# Distinct reproducible seed per image; mask to int64 (not 53 bits, which
# would truncate large explicit seeds and collide distinct ones).
# Distinct reproducible seed per image; mask to int64 (not 53 bits, which would truncate large
# explicit seeds and collide distinct ones).
seed_i = (seed + index) & ((1 << 63) - 1)
out_path = str(Path(tmpdir) / f"img_{index}.png")
params = SdCppGenParams(
@ -1097,25 +1092,25 @@ class SdCppDiffusionBackend:
self._state = None
self._load_token += 1
self._loading = None
# Grab a mid-start() uncommitted server too so we can stop it (startup is abortable).
# Grab a mid-start() uncommitted server too so we can stop it (startup is abortable).
pending = self._pending_server
self._pending_server = None
# Stop the resident server outside the lock (terminate can take seconds); a mid-flight
# generation had its cancel set above and unwinds as the process goes away.
# Stop the resident server outside the lock (terminate can take seconds); a mid-flight generation
# had its cancel set above and unwinds as the process goes away.
if state is not None and state.server is not None:
state.server.stop()
if pending is not None and pending is not (state.server if state else None):
pending.stop()
# Barrier: wait for a signalled one-shot generation to exit before reporting unloaded,
# since callers treat this return as "device is free" (same pattern as DiffusionBackend.unload).
# Barrier: wait for a signalled one-shot generation to exit before reporting unloaded, since
# callers treat this return as "device is free" (as DiffusionBackend.unload does).
with self._generate_lock:
pass
return self.status()
def status(self) -> dict[str, Any]:
state = self._state
# A resident sd-server can exit after load (OOM/crash while idle); drop stale state so
# status reports not-loaded and clients reload, not a 500 per generation on a dead process.
# A resident sd-server can exit after load (OOM/crash while idle); drop stale state so status
# reports not-loaded and clients reload, instead of a 500 per generation on a dead process.
if (
state is not None
and state.mode == "server"

21
studio/backend/core/inference/sd_cpp_engine.py
View file

@ -65,8 +65,8 @@ def _terminate(proc: "subprocess.Popen") -> None:
proc.kill()
except Exception: # noqa: BLE001 -- best-effort teardown
pass
# Reap the killed child so it doesn't linger as a zombie: callers raise right after
# _terminate, so without this a burst of cancellations leaks process-table entries.
# Reap the killed child so it doesn't linger as a zombie: callers raise right after _terminate,
# so without this a burst of cancellations leaks process-table entries.
try:
proc.wait(timeout = 5)
except Exception: # noqa: BLE001 -- best-effort reap; never block teardown
@ -155,8 +155,8 @@ def _find_binary(
if hit:
return hit
# 3. Default install root. Honors UNSLOTH_STUDIO_HOME / STUDIO_HOME like the installer (base =
# the Studio home's parent), so side-by-side Studios stay isolated; else ~/.unsloth/....
# 3. Default install root. Honors UNSLOTH_STUDIO_HOME / STUDIO_HOME like the installer, so
# side-by-side Studios stay isolated; else ~/.unsloth/....
studio_home = os.environ.get("UNSLOTH_STUDIO_HOME") or os.environ.get("STUDIO_HOME")
default_root = (
Path(studio_home).parent / "stable-diffusion.cpp"
@ -364,13 +364,13 @@ class SdCppEngine:
env = run_env,
# Own session/process group so cancellation/timeout kills the whole tree (POSIX).
start_new_session = (os.name == "posix"),
# Bind the child to the parent's lifetime (PR_SET_PDEATHSIG) so a parent crash can't
# orphan sd-cli holding VRAM/RAM. Composes with start_new_session.
# Bind the child to the parent's lifetime (PR_SET_PDEATHSIG) so a parent crash can't orphan
# sd-cli holding VRAM/RAM. Composes with start_new_session.
**child_popen_kwargs(),
)
# Drain stdout on a reader thread so the timeout holds even when the child hangs WITHOUT
# printing (a plain `for line in proc.stdout` blocks until EOF). The reader pushes lines
# (then a None sentinel) to a queue the main loop polls against a wall-clock deadline.
# Drain stdout on a reader thread so the timeout holds even when the child hangs WITHOUT printing
# (a plain `for line in proc.stdout` blocks until EOF). The reader pushes lines (then a None
# sentinel) to a queue the main loop polls against a wall-clock deadline.
tail: list[str] = []
line_q: "queue.Queue[Optional[str]]" = queue.Queue()
@ -433,7 +433,8 @@ class SdCppEngine:
ENGINE_DIFFUSERS = "diffusers"
ENGINE_SD_CPP = "sd_cpp"
# Backends diffusers serves well with GPU acceleration; everything else is native-engine territory.
# Backends diffusers serves well with GPU acceleration; everything else is native-engine
# territory.
_GPU_BACKENDS = frozenset({"cuda", "rocm", "xpu"})

43
studio/backend/core/inference/sd_cpp_server.py
View file

@ -61,8 +61,8 @@ _TRANSPORT_ERRORS = (
httpx.WriteError,
)
# Readiness probe: port binds only after the model loads, so any 200 means ready. Use
# trivial /v1/models, not /sdcpp/v1/capabilities (can block enumerating metadata).
# Readiness probe: the port binds only after the model loads, so any 200 means ready. Use trivial
# /v1/models, not /sdcpp/v1/capabilities (which can block enumerating metadata).
_READY_PATH = "/v1/models"
# Native async sdcpp API.
_IMG_GEN_PATH = "/sdcpp/v1/img_gen"
@ -72,8 +72,8 @@ _TERMINAL_OK = "completed"
_TERMINAL_FAIL = "failed"
_TERMINAL_CANCELLED = "cancelled"
# Grace for the best-effort native cancel to show in job status before abandoning the
# poll; without the cap a lost cancel would hold the generate lock until the job ends.
# Grace for the best-effort native cancel to show in job status before abandoning the poll;
# without the cap a lost cancel would hold the generate lock until the job ends.
_CANCEL_GRACE_S = 5.0
@ -145,8 +145,8 @@ class SdCppServer:
a concurrent start/stop can't interleave.
"""
with self._lifecycle_lock:
# A stop()/unload that raced in before start() took the lock already set _abort
# and closed the client; honor it rather than leak a spawned model process.
# A stop()/unload that raced in before start() took the lock already set _abort and closed the
# client; honor it rather than leak a spawned model process.
if self._stopped or self._abort.is_set():
raise SdCppCancelled("sd-server start was cancelled before launch.")
self._abort.clear()
@ -174,9 +174,8 @@ class SdCppServer:
self._spawn_error: Optional[Exception] = None
spawned = threading.Event()
# Spawn INSIDE the long-lived drain thread: child_popen_kwargs() sets
# PR_SET_PDEATHSIG, bound to the creating thread on Linux, so the creator must
# outlive the child (a transient spawner ending would kill the server).
# Spawn INSIDE the long-lived drain thread: child_popen_kwargs() sets PR_SET_PDEATHSIG, bound to
# the creating thread on Linux, so the creator must outlive the child.
def _own_process() -> None:
try:
proc = subprocess.Popen(
@ -233,8 +232,8 @@ class SdCppServer:
deadline = time.monotonic() + timeout
url = f"{self.base_url}{_READY_PATH}"
while time.monotonic() < deadline:
# A concurrent stop() sets _abort so this wait bails without holding the
# model-load hostage for the full startup_timeout.
# A concurrent stop() sets _abort so this wait bails without holding the model load hostage for
# the full startup_timeout.
if self._abort.is_set():
logger.info("sd-server startup aborted before ready")
return False
@ -274,8 +273,8 @@ class SdCppServer:
def stop(self) -> None:
"""Terminate the server (SIGTERM -> SIGKILL), join the drain, and release the HTTP
client + atexit handler. Idempotent."""
# Signal abort BEFORE contending for the lock so a start() readiness wait (which
# holds the lock up to startup_timeout) bails immediately instead of blocking stop().
# Signal abort BEFORE contending for the lock so a start() readiness wait (which holds the lock
# up to startup_timeout) bails immediately instead of blocking stop().
self._abort.set()
self._stopped = True
with self._lifecycle_lock:
@ -342,8 +341,8 @@ class SdCppServer:
Raises ``RuntimeError`` on submit/poll failures (including the server dying), with
the log tail attached.
"""
# Already stopped with the cancel event set -> report cancellation (route -> 409),
# not a generic "server died" 500.
# Already stopped with the cancel event set: report cancellation (route 409), not a generic
# "server died" 500.
if self._stopped or not self.is_alive():
if cancel_event is not None and cancel_event.is_set():
raise SdCppCancelled("sd-server generation was cancelled.")
@ -388,17 +387,17 @@ class SdCppServer:
self.cancel(job_id)
cancel_sent_at = time.monotonic()
elif time.monotonic() - cancel_sent_at > _CANCEL_GRACE_S:
# Cancel not reflected within the grace window; abandon the poll so
# the caller can stop the server instead of holding the generate lock.
# Cancel not reflected within the grace window; abandon the poll so the caller can stop the
# server instead of holding the generate lock.
raise SdCppCancelled("sd-server generation was cancelled.")
if not self.is_alive():
# Unwinding a cancel (e.g. unload killed the server) -> clean cancellation.
# Unwinding a cancel (e.g. unload killed the server): clean cancellation.
if cancel_event is not None and cancel_event.is_set():
raise SdCppCancelled("sd-server generation was cancelled.")
raise RuntimeError(self._died_message("img_gen poll", None))
if time.monotonic() > deadline:
# sd-server won't interrupt an in-flight job (cancel_generating=false), so
# cancel + stop to free the slot; the backend reloads on the next generate.
# sd-server won't interrupt an in-flight job (cancel_generating=false), so cancel + stop to free
# the slot; the backend reloads on the next generate.
self.cancel(job_id)
self.stop()
raise RuntimeError(f"sd-server generation timed out after {total_timeout}s")
@ -408,8 +407,8 @@ class SdCppServer:
time.sleep(poll_interval)
continue
except RuntimeError as exc:
# A concurrent stop() closes the shared client -> plain RuntimeError
# ("client has been closed"), not a transport error; map cancel -> 409.
# A concurrent stop() closes the shared client, giving a plain RuntimeError ("client has been
# closed") rather than a transport error; map a cancel to 409.
if cancel_event is not None and cancel_event.is_set():
raise SdCppCancelled("sd-server generation was cancelled.") from exc
raise

346
studio/backend/core/inference/video.py
View file

@ -98,12 +98,12 @@ from core.inference.diffusion import hub_cache_dir
logger = get_logger(__name__)
# Load kinds (mirror the image backend): gguf (single-file GGUF DiT + base repo),
# single_file (safetensors DiT, e.g. fp8 LTX-2.3), pipeline (full diffusers repo).
# Load kinds (mirror the image backend): gguf (single-file GGUF DiT + base repo), single_file
# (safetensors DiT, e.g. fp8 LTX-2.3), pipeline (full diffusers repo).
_MODEL_KINDS = frozenset({"gguf", "single_file", "pipeline"})
# Vendor base repos allowed to load as full (non-GGUF) artifacts despite not being
# under unsloth/. Exact-match, lowercased, safetensors-only, no remote code.
# Vendor base repos allowed to load as full (non-GGUF) artifacts despite not being under
# unsloth/. Exact-match, lowercased, safetensors-only, no remote code.
_TRUSTED_NON_GGUF_VIDEO_REPOS = frozenset(
{
"lightricks/ltx-2",
@ -112,8 +112,7 @@ _TRUSTED_NON_GGUF_VIDEO_REPOS = frozenset(
# Wan2.2 official diffusers base repos: safetensors-only, no remote code.
"wan-ai/wan2.2-ti2v-5b-diffusers",
"wan-ai/wan2.2-t2v-a14b-diffusers",
# HunyuanVideo-1.5 community Diffusers repacks (tencent's own repo is the
# non-diffusers layout with no model_index.json, unloadable here).
# HunyuanVideo-1.5 community Diffusers repacks (tencent's own repo has no model_index.json).
"hunyuanvideo-community/hunyuanvideo-1.5-diffusers-480p_t2v",
"hunyuanvideo-community/hunyuanvideo-1.5-diffusers-720p_t2v",
}
@ -157,8 +156,8 @@ def _picked_gguf_arch(repo_id: str, gguf_filename: str) -> Optional[str]:
path = Path(repo_id).expanduser() / gguf_filename
if not path.is_file():
# Not a local dir: resolve a cached HUB blob (no network). Probe active, legacy,
# AND default cache roots (as the picker's listing does) or a non-active-root GGUF 400s.
# Not a local dir: resolve a cached HUB blob (no network). Probe active, legacy AND default
# cache roots (as the picker's listing does) or a non-active-root GGUF 400s.
from huggingface_hub import try_to_load_from_cache
cached = try_to_load_from_cache(repo_id, gguf_filename)
@ -228,8 +227,8 @@ def _detect_load_family(
else None
)
if fam is None and gguf_filename and not family_override:
# A renamed GGUF carries no family token in its name; resolve via general.architecture
# (its string, e.g. "ltxv", is a family alias). No-backend archs still yield None -> 400.
# A renamed GGUF carries no family token in its name; resolve via general.architecture (its
# string, e.g. "ltxv", is a family alias). No-backend archs still yield None, so a 400.
arch = _picked_gguf_arch(repo_id, gguf_filename)
if arch:
fam = detect_video_family(repo_id, override = arch)
@ -268,17 +267,17 @@ class _VideoLoadState:
backend_flags: Optional[dict] = None
attention_backend: Optional[str] = None
transformer_cache: Optional[str] = None
# AUTO on a cache-capable DiT: generate() re-checks the step count and toggles FBCache
# across FBCACHE_MIN_STEPS. An explicit request (off / fbcache) is never toggled.
# AUTO on a cache-capable DiT: generate() toggles FBCache across FBCACHE_MIN_STEPS; an explicit
# request is never toggled.
cache_auto: bool = False
# Inputs the generation-time toggle re-applies (quantised threshold + override).
cache_quant_active: bool = False
cache_threshold: Optional[float] = None
# Dense transformer quant engaged ("int8"|"fp8"|"nvfp4"|"mxfp8") or None (loaded bf16).
# Pipeline-kind only; torchao-quantised in place onto the low-precision tensor cores.
# Dense transformer quant engaged ("int8"|"fp8"|"nvfp4"|"mxfp8") or None. Pipeline-kind only;
# torchao-quantised in place onto the low-precision tensor cores.
transformer_quant: Optional[str] = None
# Text-encoder quant engaged ("fp8"|"fp8_dynamic"|"int8"|"nvfp4") or None. The companion
# encoder (UMT5/Gemma3/Qwen2.5-VL) is often the largest resident; shrunk in place.
# Text-encoder quant engaged ("fp8"|"fp8_dynamic"|"int8"|"nvfp4") or None. The companion encoder
# is often the largest resident; shrunk in place.
text_encoder_quant: Optional[str] = None
resolved: Optional[dict] = None
@ -296,11 +295,10 @@ def _progress(phase: Optional[str], **extra: Any) -> dict[str, Any]:
# ── dual-DiT (Wan2.2-A14B MoE) helpers ────────────────────────────────────────
# The optimisation helpers and the quantiser all act on ``pipe.transformer`` -- fine for
# single-DiT families. Wan2.2-A14B is a dual-expert MoE (transformer = high-noise steps,
# transformer_2 = low-noise), so an optimisation on ``transformer`` alone leaves the second
# expert unoptimised for half the schedule. Rather than fork each helper, present the second
# DiT AS ``pipe.transformer`` via a thin proxy (built only for is_moe) and call the helper again.
# The optimisation helpers and the quantiser all act on ``pipe.transformer``. Wan2.2-A14B is a
# dual-expert MoE (transformer = high-noise steps, transformer_2 = low-noise), so rather than
# fork each helper, present the second DiT AS ``pipe.transformer`` via a thin proxy (built only
# for is_moe) and call the helper again.
def _transformer_names(pipe: Any, fam: VideoFamily) -> tuple[str, ...]:
@ -334,8 +332,8 @@ class _SecondDiTView:
return getattr(object.__getattribute__(self, "_pipe"), name)
def __setattr__(self, name: str, value: Any) -> None:
# Writes land on the real pipe (else a helper's reassignment vanishes with the
# view); ``transformer`` mirrors onto the second expert.
# Writes land on the real pipe (else a helper's reassignment vanishes with the view);
# ``transformer`` mirrors onto the second expert.
pipe = object.__getattribute__(self, "_pipe")
setattr(pipe, "transformer_2" if name == "transformer" else name, value)
@ -382,8 +380,8 @@ class VideoBackend:
) -> VideoFamily:
"""Cheap, network-free validation shared by the route and the load path."""
kind = resolve_video_model_kind(gguf_filename, model_kind)
# A -GGUF repo picked without a quant filename resolves to pipeline kind and would
# only fail in from_pretrained (no model_index.json) after the route evicts the owner.
# A -GGUF repo picked without a quant filename resolves to pipeline kind and would only fail in
# from_pretrained (no model_index.json) after the route evicts the owner.
if kind == "pipeline" and repo_id.strip().lower().rstrip("/").endswith("-gguf"):
raise ValueError(
f"'{repo_id}' is a GGUF repo: pick one of its .gguf files "
@ -401,15 +399,15 @@ class VideoBackend:
f"Non-GGUF video loads are limited to unsloth/* repos, the official "
f"family base repos, and local paths; '{repo_id}' is neither."
)
# Companions load with from_pretrained, so a base repo is held to the non-GGUF bar:
# a GGUF pick must not smuggle in an arbitrary remote base.
# Companions load with from_pretrained, so a base repo is held to the non-GGUF bar: a GGUF pick
# must not smuggle in an arbitrary remote base.
if base_repo and (base_repo or "").strip() and not _is_trusted_video_repo(base_repo):
raise ValueError(
f"base_repo is limited to unsloth/* repos, the official family base "
f"repos, and local paths; '{base_repo}' is neither."
)
# A local base_repo loads as a full pipeline (needs model_index.json); reject a
# non-pipeline local base here, before the load. Shared helper keeps image/video/training in sync.
# A local base_repo loads as a full pipeline (needs model_index.json); reject a non-pipeline one
# here, before the load. The shared helper keeps image/video/training in sync.
from core.inference.diffusion import _assert_local_base_is_pipeline
_assert_local_base_is_pipeline(base_repo)
@ -424,8 +422,8 @@ class VideoBackend:
)
# A missing local checkpoint must fail HERE, before the route evicts a resident model.
if kind in ("gguf", "single_file"):
# Fail a kind/extension mismatch before the GPU handoff: gguf needs .gguf,
# single_file needs .safetensors (mirrors the image loader's gate).
# Fail a kind/extension mismatch before the GPU handoff: gguf needs .gguf, single_file needs
# .safetensors (mirrors the image loader's gate).
is_gguf_name = (gguf_filename or "").lower().endswith(".gguf")
if kind == "gguf" and not is_gguf_name:
raise ValueError("a 'gguf' load requires a .gguf checkpoint name.")
@ -437,8 +435,8 @@ class VideoBackend:
f"(expected a .safetensors name; use a .gguf name for a GGUF load)."
)
root = Path(repo_id).expanduser()
# Path-shaped: "."/".." prefix, a backslash (never in "org/name"), or an absolute
# path -- so a missing Windows-shaped local pick fails before the handoff, not as a Hub repo.
# Path-shaped: "."/".." prefix, a backslash (never in "org/name"), or an absolute path, so a
# missing Windows-shaped local pick fails before the handoff, not as a Hub repo.
path_shaped = (
repo_id.startswith(("/", "\\", "~", ".")) or "\\" in repo_id or root.is_absolute()
)
@ -449,8 +447,8 @@ class VideoBackend:
except Exception as exc: # noqa: BLE001 -- surface as client input error
raise ValueError(str(exc)) from exc
elif root.is_file():
# The loader hands a local FILE straight through (ignoring gguf_filename), so
# the file's OWN suffix must match the kind; reject a mismatch before the handoff.
# The loader hands a local FILE straight through (ignoring gguf_filename), so the file's OWN
# suffix must match the kind; reject a mismatch before the handoff.
suffix = root.suffix.lower()
if kind == "gguf" and suffix != ".gguf":
raise ValueError(
@ -464,8 +462,8 @@ class VideoBackend:
)
elif path_shaped:
raise ValueError(f"Local model path '{repo_id}' does not exist.")
# A local pipeline pick must be a diffusers directory (model_index.json), else it would
# only fail in from_pretrained after eviction (mirrors the image loader).
# A local pipeline pick must be a diffusers directory (model_index.json), else it would only
# fail in from_pretrained after eviction (mirrors the image loader).
if kind == "pipeline":
root = Path(repo_id).expanduser()
# Gate on .exists() (not .is_dir()) so a local FILE picked as a pipeline is rejected too.
@ -474,8 +472,8 @@ class VideoBackend:
f"Local pipeline path is not a diffusers directory "
f"(no model_index.json): {repo_id}"
)
# Reject a malformed transformer_quant cheaply, before the handoff (applies on
# pipeline-kind loads; ignored on gguf/single_file, matching the image backend).
# Reject a malformed transformer_quant cheaply, before the handoff (pipeline-kind loads only;
# ignored on gguf/single_file, matching the image backend).
normalize_transformer_quant(transformer_quant)
# Reject a malformed text_encoder_quant the same way (any kind: the encoder is always dense).
normalize_te_quant(text_encoder_quant)
@ -562,8 +560,8 @@ class VideoBackend:
if self._load_token == token and self._loading is not None:
self._loading.base_repo = base
self._loading.expected_bytes = expected
# Checkpoint downloads outside the lock so an unload/eviction can preempt the
# multi-GB pull; companions pre-download the same way (scoped, cancellable, resumable).
# Checkpoint downloads outside the lock so an unload/eviction can preempt the multi-GB pull;
# companions pre-download the same way (scoped, cancellable, resumable).
checkpoint_local: Optional[Path] = None
if kwargs.get("gguf_filename") and not Path(kwargs["repo_id"]).expanduser().exists():
from utils.hf_xet_fallback import hf_hub_download_with_xet_fallback
@ -575,17 +573,17 @@ class VideoBackend:
cancel_event = self._cancel_event,
)
)
# An LTX-2.3 checkpoint supplies the VAEs/vocoder/connectors, so the base pull
# shrinks to scheduler + text encoder + tokenizer; recompute the estimate to match
# (detectable only once the checkpoint header is on disk).
# An LTX-2.3 checkpoint supplies the VAEs/vocoder/connectors, so the base pull shrinks to
# scheduler + text encoder + tokenizer; recompute the estimate to match (detectable only once
# the checkpoint header is on disk).
ltx23 = False
if fam is not None and fam.name == "ltx-2" and kind != "pipeline":
from .video_ltx2 import is_ltx23_checkpoint
probe = checkpoint_local
if probe is None:
# Local repos: a bare file, or a dir child via the same resolver load_pipeline
# uses. Unresolvable -> load_pipeline surfaces the real error; keep the wide pull.
# Local repos: a bare file, or a dir child via the same resolver load_pipeline uses.
# Unresolvable means load_pipeline surfaces the real error; keep the wide pull.
root = Path(kwargs["repo_id"]).expanduser()
if root.is_file():
probe = root
@ -612,23 +610,22 @@ class VideoBackend:
if self._load_token == token and self._loading is not None:
self._loading.expected_bytes = expected
base_local = self._predownload_base(base, kwargs.get("hf_token"), kind, ltx23 = ltx23)
# The 2.3 assembly pulls per component from the hub id (its snapshot lacks the base
# VAEs), so it only gets the warmed cache; generic paths get the full local snapshot.
# The 2.3 assembly pulls per component from the hub id (its snapshot lacks the base VAEs), so it
# only gets the warmed cache; generic paths get the full local snapshot.
kwargs["_base_local_dir"] = None if ltx23 else base_local
self.load_pipeline(**kwargs)
with self._lock:
if self._load_token == token:
self._loading = None
except Exception as exc: # noqa: BLE001 -- surfaced via load_progress
# A failed/cancelled load never commits _VideoLoadState, so roll back the
# process-wide speed globals here (token-scoped, so a superseded load can't clobber a newer one's).
# A failed/cancelled load never commits _VideoLoadState, so roll back the process-wide speed
# globals here (token-scoped, so a superseded load can't clobber a newer one's).
self._rollback_precommit_globals(token)
if self._load_token != token:
return
logger.error("video.load_failed: %s", exc)
# Free the debris of a failed construction (mirrors diffusion.py): no state was
# committed, so nothing else releases the VRAM a partial pipeline reserved. Guarded
# so a sticky CUDA error can't skip stamping the real error below.
# Free the debris of a failed construction (mirrors diffusion.py): no state was committed, so
# nothing else releases the VRAM. Guarded so a sticky CUDA error still stamps the real error.
try:
clear_gpu_cache()
except Exception: # noqa: BLE001 -- cleanup is best-effort
@ -658,8 +655,8 @@ class VideoBackend:
diffusion_gguf_compile.uninstall_all()
# LTX-2.3 gets DiT/connectors/VAEs/vocoder from the checkpoint + extras, so only the
# 2.0 base's scheduler / text encoder / tokenizer are pulled.
# LTX-2.3 gets DiT/connectors/VAEs/vocoder from the checkpoint + extras, so only the 2.0 base's
# scheduler / text encoder / tokenizer are pulled.
_LTX23_BASE_PREFIXES = ("scheduler/", "text_encoder/", "tokenizer/")
@staticmethod
@ -685,8 +682,8 @@ class VideoBackend:
files: list[tuple[str, int]] = []
for sibling in info.siblings or []:
name, size = sibling.rfilename, sibling.size or 0
# .jinja: tokenizer/chat_template.jinja is a standalone file apply_chat_template
# needs at generation time; a snapshot without it crashes the first generation.
# .jinja: tokenizer/chat_template.jinja is a standalone file apply_chat_template needs at
# generation time; a snapshot without it crashes the first generation.
if not name.endswith((".safetensors", ".json", ".model", ".txt", ".jinja")):
continue
if "/" not in name and name.endswith(".safetensors"):
@ -745,14 +742,12 @@ class VideoBackend:
fam = _detect_load_family(repo_id, gguf_filename, family_override)
kind = resolve_video_model_kind(gguf_filename, model_kind)
base = repo_id if kind == "pipeline" else resolve_video_base_repo(fam, base_repo)
# The load narrows the base pull for an LTX-2.3 checkpoint (its VAEs, vocoder and
# connectors come from the checkpoint and the 2.3 extras, not the 2.0 base), but it can
# only tell 2.3 from 2.0 by the checkpoint header, which is not on disk yet. Name-based
# here: a wrong guess costs at most an inline pull at load time, while staging the wide
# base list costs gigabytes of weights the pipeline never opens.
# The load narrows the base pull for an LTX-2.3 checkpoint, but it can only tell 2.3 from 2.0 by
# the checkpoint header, which is not on disk yet. Name-based here: a wrong guess costs at most
# an inline pull at load time, while staging the wide base list costs gigabytes.
ltx23 = self._pick_looks_like_ltx23(fam, repo_id, gguf_filename, kind)
# Keyed by repo so a 2.3 pick's checkpoint and extras (both in the 2.3 repo) stay ONE
# scoped job; two entries for one repo would collide on the job key.
# Keyed by repo so a 2.3 pick's checkpoint and extras stay ONE scoped job; two entries for one
# repo would collide on the job key.
entries: dict[str, dict[str, Any]] = {}
total = 0
@ -790,8 +785,8 @@ class VideoBackend:
]
total += add(repo_id, sizes, gguf = gguf_filename)
if ltx23:
# The 2.3 assembly reads these companion files at load; without them here
# they would be pulled inline, outside the panel and its disk preflight.
# The 2.3 assembly reads these companion files at load; without them here they would be pulled
# inline, outside the panel and its disk preflight.
from .video_ltx2 import ltx23_extras_files, LTX23_EXTRAS_REPO
wanted = set(ltx23_extras_files(gguf_filename))
@ -864,8 +859,8 @@ class VideoBackend:
snapshot_root: Optional[Path] = None
for name, _ in files:
# Explicit check: a cached file returns without consulting the event, so a
# warm-cache sweep would otherwise run to completion after an unload cancelled.
# Explicit check: a cached file returns without consulting the event, so a warm-cache sweep would
# otherwise run to completion after an unload cancelled.
if self._cancel_event.is_set():
raise RuntimeError(VIDEO_CANCELLED_MSG)
local = Path(
@ -902,8 +897,7 @@ class VideoBackend:
for name in files:
try:
path = os.path.join(root, name)
# Snapshot entries are symlinks into blobs/; skip them so a
# blob is not counted twice.
# Snapshot entries are symlinks into blobs/; skip them so a blob is not counted twice.
if not os.path.islink(path):
total += os.path.getsize(path)
except OSError:
@ -926,8 +920,8 @@ class VideoBackend:
phase = "downloading"
if expected and downloaded >= expected:
phase = "finalizing"
# The cache scan counts every blob (incl. files this load never reads), so the raw
# counter can exceed the scoped estimate; clamp to what the bar reports.
# The cache scan counts every blob (incl. files this load never reads), so the raw counter can
# exceed the scoped estimate; clamp to what the bar reports.
downloaded = expected
return _progress(
phase,
@ -990,37 +984,36 @@ class VideoBackend:
# Signal a generation from the PREVIOUS model (the token check above bailed a superseded worker).
if self._active_generate_cancel is not None:
self._active_generate_cancel.set()
# Barrier: wait for the signalled generation to exit before teardown, or two models
# coexist in VRAM (the denoise loop holds its pipe ref until the next callback).
# Barrier: wait for the signalled generation to exit before teardown, or two models coexist in
# VRAM (the denoise loop holds its pipe ref until the next callback).
with self._generate_lock:
pass
# The barrier wait can outlive this load (a newer load / unload superseded it); recheck
# before touching shared state so we don't destroy the current model or build a dead pipe.
# The barrier wait can outlive this load (a newer load / unload superseded it); recheck before
# touching shared state so we don't destroy the current model or build a dead pipe.
if _load_token is not None and _load_token != self._load_token:
raise RuntimeError("Video load was cancelled or superseded.")
self._teardown_state()
target = resolve_diffusion_device_target()
device = target.device
# Video DiTs are bf16-native; fp16 overflows, so a resolved fp16 promotes to float32
# (same rule as fp16-incompatible image families). CPU stays float32.
# Video DiTs are bf16-native; fp16 overflows, so a resolved fp16 promotes to float32 (same rule
# as fp16-incompatible image families). CPU stays float32.
dtype = target.dtype
if fam.fp16_incompatible and dtype is torch.float16:
dtype = torch.float32
# Size tables below are bf16 (2-byte); when the promotion lands fp32 on an accelerator,
# dense estimates double, so scale them (GGUF stays quantised, so only dense scales).
# Size tables below are bf16 (2-byte); when the promotion lands fp32 on an accelerator, dense
# estimates double, so scale them (GGUF stays quantised).
dtype_scale = 2.0 if device != "cpu" and dtype is torch.float32 else 1.0
# Precision tri-state (mirror image backend): unset/"auto" -> hardware ladder picks a
# quantised DiT (int8 min, fp8 on datacenter silicon); "none"/"off" pins dense bf16; an
# explicit scheme pins it. Pipeline-kind only; the offload guard below still skips it.
# Precision tri-state (mirror image backend): unset/"auto" -> hardware ladder picks a quantised
# DiT; "none"/"off" pins dense bf16; an explicit scheme pins it. Pipeline-kind only; the offload
# guard below still skips it.
if transformer_quant is None or str(transformer_quant).strip().lower() in (
"",
"auto",
):
# An explicit Speed="off" (bit-exact) load must stay dense bf16: auto-quant would
# engage int8/fp8 + regional compile and break the bit-exact request. Suppress the
# auto default when speed was pinned off (mirrors diffusion.py); else auto applies.
# An explicit Speed="off" (bit-exact) load must stay dense bf16: auto-quant would engage int8/fp8
# + regional compile and break the request. Suppress the auto default when speed was pinned off.
speed_off = speed_mode is not None and str(speed_mode).strip().lower() == SPEED_OFF
transformer_quant = "off" if speed_off else TQ_AUTO
@ -1048,8 +1041,8 @@ class VideoBackend:
if components is not None
else None
)
# Budget ALL weights (image-backend contract): companions stay resident, so
# budgeting the transformer alone lets auto pick OFFLOAD_NONE and OOM.
# Budget ALL weights (image-backend contract): companions stay resident, so budgeting the
# transformer alone lets auto pick OFFLOAD_NONE and OOM.
model_dense_mib = (
transformer_mib + (companion_mib or 0) if transformer_mib is not None else None
)
@ -1066,9 +1059,9 @@ class VideoBackend:
companion_dense_mib = companion_mib,
requested_mode = normalize_memory_mode(memory_mode),
)
# Parity with the image dense-quant path: the bf16-table plan can force offload a
# quantised DiT would not need. Re-plan with the scheme's steady factor and keep the
# resident placement if it fits; fall back to this bf16 plan if quant later fails.
# Parity with the image dense-quant path: the bf16-table plan can force offload a quantised DiT
# would not need. Re-plan with the scheme's steady factor and keep the resident placement if it
# fits; fall back to this bf16 plan if quant later fails.
bf16_plan = plan
quant_replanned = False
if (
@ -1107,20 +1100,19 @@ class VideoBackend:
# ── build the pipeline.
pipeline_cls = getattr(diffusers, fam.pipeline_class)
# cache_dir pins every loader call to the live cache root, so a mid-session
# change can't split one model across the old and new roots.
# cache_dir pins every loader call to the live cache root, so a mid-session change can't split
# one model across the old and new roots.
pipe_kwargs: dict[str, Any] = {"torch_dtype": dtype, "cache_dir": hub_cache_dir()}
if getattr(fam, "vae_force_fp32", False):
# Wan's VAE must decode in float32. A scalar torch_dtype truncates its fp32 weights
# to bf16 (no _keep_in_fp32_modules); a later .to(float32) only widens lossy values
# (banding / black frames). Use the per-component dtype dict; "default" MUST be set
# or unlisted components fall back to fp32 (over-widening the DiT).
# Wan's VAE must decode in float32. A scalar torch_dtype truncates its fp32 weights to bf16 (no
# _keep_in_fp32_modules) and a later .to(float32) only widens lossy values (banding / black
# frames). Use the per-component dtype dict; "default" MUST be set or unlisted components fall
# back to fp32, over-widening the DiT.
pipe_kwargs["torch_dtype"] = {"vae": torch.float32, "default": dtype}
if hf_token:
pipe_kwargs["token"] = hf_token
# A hosted pre-cast fp8 text encoder (when the family ships one and the runtime cast
# would engage) skips the dense TE download -- for LTX's Gemma3-27B that is the ~50 GB
# heavyweight of the load. quantize_text_encoders below re-applies the cast idempotently.
# A hosted pre-cast fp8 text encoder skips the dense TE download -- for LTX's Gemma3-27B that is
# the ~50 GB heavyweight of the load. quantize_text_encoders re-applies the cast idempotently.
from .diffusion_te_prequant import te_prequant_pipe_kwargs
pipe_kwargs.update(
@ -1135,8 +1127,8 @@ class VideoBackend:
)
)
if kind == "pipeline":
# The pre-downloaded snapshot dir keeps from_pretrained off the hub (its sweep would
# also pull root checkpoints + duplicate shards); hub id when pre-download was skipped.
# The pre-downloaded snapshot dir keeps from_pretrained off the hub (its sweep would also pull
# root checkpoints + duplicate shards); hub id when pre-download was skipped.
pipe = pipeline_cls.from_pretrained(_base_local_dir or repo_id, **pipe_kwargs)
else:
transformer_cls = getattr(diffusers, fam.transformer_class)
@ -1155,8 +1147,8 @@ class VideoBackend:
from .video_ltx2 import is_ltx23_checkpoint, load_ltx23_pipeline
if fam.name == "ltx-2" and is_ltx23_checkpoint(checkpoint_path):
# 2.3 checkpoints need the full assembly: new config flags, key renames the
# stock converter lacks, and the 2.3 connectors/VAEs/vocoder the base lacks.
# 2.3 checkpoints need the full assembly: new config flags, key renames the stock converter
# lacks, and the 2.3 connectors/VAEs/vocoder the base lacks.
pipe = load_ltx23_pipeline(
checkpoint_path,
base_repo = base,
@ -1170,8 +1162,8 @@ class VideoBackend:
_base_local_dir or base, transformer = transformer, **pipe_kwargs
)
# The dtype dict already loads the Wan VAE at float32. Belt-and-suspenders for any path
# that bypassed it (e.g. a passed-in vae=): re-pin an fp32-force VAE that came back lower.
# The dtype dict already loads the Wan VAE at float32. Belt-and-suspenders for any path that
# bypassed it (e.g. a passed-in vae=): re-pin an fp32-force VAE that came back lower.
if getattr(fam, "vae_force_fp32", False):
vae = getattr(pipe, "vae", None)
if vae is not None and getattr(vae, "dtype", None) is not torch.float32:
@ -1182,13 +1174,13 @@ class VideoBackend:
clear_gpu_cache()
raise RuntimeError("Video load was cancelled or superseded.")
# For a dual-DiT MoE (Wan2.2-A14B), every optimisation site below covers BOTH experts:
# ``views`` is (pipe, _SecondDiTView(pipe)); a single-DiT load resolves to (pipe,).
# For a dual-DiT MoE (Wan2.2-A14B), every optimisation site below covers BOTH experts: ``views``
# is (pipe, _SecondDiTView(pipe)); a single-DiT load resolves to (pipe,).
views = _views_for(pipe, fam)
# ── dense transformer quant (opt-in, pipeline-kind only): torchao-quantise the dense
# bf16 DiT in place onto the low-precision tensor cores (image-backend fast path). CUDA +
# bf16 only; best-effort. Quant must precede compile (eager dynamic quant is ~30x slower).
# ── dense transformer quant (opt-in, pipeline-kind only): torchao-quantise the dense bf16 DiT in
# place onto the low-precision tensor cores (image-backend fast path). CUDA + bf16 only,
# best-effort. Quant must precede compile (eager dynamic quant is ~30x slower).
transformer_quant_engaged: Optional[str] = None
quant_skipped_for_offload = False
if (
@ -1197,10 +1189,9 @@ class VideoBackend:
and dense_transformer_supported(target)
and plan.offload_policy != "none"
):
# Offload hooks move modules with Module.to(), which torchao quantized tensors reject
# (aten._has_compatible_shallow_copy_type unimplemented) -- a hard crash on the
# Wan2.2-A14B gate run (114 GB dual DiT plans model offload). Skip quant (dense-under-
# offload beats a crash); surfaced in the resolved record, forceable via a resident mode.
# Offload hooks move modules with Module.to(), which torchao quantized tensors reject -- a hard
# crash on the Wan2.2-A14B gate run. Skip quant (dense-under-offload beats a crash); surfaced in
# the resolved record, forceable via a resident mode.
logger.info(
"video.transformer_quant: skipped (offload policy '%s' moves the "
"DiT via Module.to(), unsupported for torchao quantized tensors); "
@ -1215,8 +1206,8 @@ class VideoBackend:
):
engaged = []
for view in views:
# Pass each expert's view so both DiTs quantise with the same arch-chosen scheme.
# The family name drives the per-family deny table (_FAMILY_SCHEME_DENY).
# Pass each expert's view so both DiTs quantise with the same arch-chosen scheme. The family name
# drives the per-family deny table (_FAMILY_SCHEME_DENY).
scheme = quantize_transformer(
view,
target,
@ -1226,8 +1217,8 @@ class VideoBackend:
)
if scheme is not None:
engaged.append(scheme)
# All experts or none: the first is mutated in place, so a second-expert failure
# can't fall back to dense (mismatched precision). Fail cleanly; a full miss stays dense.
# All experts or none: the first is mutated in place, so a second-expert failure can't fall back
# to dense (mismatched precision). Fail cleanly; a full miss stays dense.
if engaged and len(engaged) < len(views):
del pipe
clear_gpu_cache()
@ -1241,9 +1232,9 @@ class VideoBackend:
if quant_replanned and transformer_quant_engaged is None:
plan = bf16_plan
# ── dense text-encoder quant (opt-in): the companion encoder (Gemma3/UMT5/Qwen2.5-VL)
# loads dense bf16 and is often the largest resident. Quantise in place for every kind,
# before placement (so offload moves the smaller weights). Best-effort; family drives int8's keep-bf16 schedule.
# ── dense text-encoder quant (opt-in): the companion encoder loads dense bf16 and is often the
# largest resident. Quantise in place for every kind, before placement (so offload moves the
# smaller weights). Best-effort; family drives int8's keep-bf16 schedule.
text_encoder_quant_engaged = quantize_text_encoders(
pipe,
target,
@ -1253,15 +1244,15 @@ class VideoBackend:
logger = logger,
)
# ── optimisation layers in the image backend's order: step cache FIRST (compile keys
# its fullgraph decision off an active cache; FBCache hooks graph-break), then attention,
# speed profile, placement last. A clip denoise runs minutes, so even a dense load
# amortises the compile: unset resolves to the near-lossless `default`; "off"/explicit honored.
# ── optimisation layers in the image backend's order: step cache FIRST (compile keys its
# fullgraph decision off an active cache), then attention, speed profile, placement last. A clip
# denoise runs minutes, so even a dense load amortises the compile: unset resolves to the
# near-lossless `default`; "off"/explicit honored.
effective_speed = resolve_speed_mode(
speed_mode, is_gguf = kind == "gguf", dense_default = SPEED_DEFAULT
)
# A torchao-quantised DiT must be compiled (eager is ~30x slower), so force at least
# the regional-compile profile when quant engaged but speed was off (matches diffusion.py).
# A torchao-quantised DiT must be compiled (eager is ~30x slower), so force at least the
# regional-compile profile when quant engaged but speed was off (matches diffusion.py).
if transformer_quant_engaged is not None and effective_speed == SPEED_OFF:
logger.info(
"video.transformer_quant: forcing speed_mode=default "
@ -1269,12 +1260,11 @@ class VideoBackend:
)
effective_speed = SPEED_DEFAULT
backend_flags = snapshot_backend_flags()
# Until the state commit transfers ownership to _teardown_state, a failure must restore
# these globals itself (via _rollback_precommit_globals). Registered BEFORE the first mutation.
# Until the state commit transfers ownership to _teardown_state, a failure must restore these
# globals itself (_rollback_precommit_globals). Registered BEFORE the first mutation.
self._precommit_globals = (_load_token, backend_flags)
# Step cache tri-state: unset/"auto" -> step-count policy decides (engage when the DEFAULT
# schedule reaches FBCACHE_MIN_STEPS, re-checked per generation); "off"/"fbcache" pinned.
# Run per expert so both denoisers cache.
# Step cache tri-state: unset/"auto" -> step-count policy decides (FBCACHE_MIN_STEPS, re-checked
# per generation); "off"/"fbcache" pinned. Run per expert so both denoisers cache.
cache_request = normalize_transformer_cache(transformer_cache)
cache_auto = transformer_cache is None or cache_request == TC_AUTO
# GGUF and torchao-quantised DiTs need the higher threshold to trigger over quant noise.
@ -1289,8 +1279,8 @@ class VideoBackend:
view,
mode = cache_request,
threshold = transformer_cache_threshold,
# A quantized transformer's residuals are larger, needing the higher FBCache
# threshold; both engaged quant and GGUF count as quant-active (cache_quant_active).
# A quantized transformer's residuals are larger, needing the higher FBCache threshold; both
# engaged quant and GGUF count as quant-active (cache_quant_active).
quant_active = cache_quant_active,
logger = logger,
)
@ -1318,9 +1308,9 @@ class VideoBackend:
attention_engaged = None
speed_optims: tuple = ()
for view in views:
# Both helpers act on ``view.transformer``; call once per view to set the kernel and
# compile each expert (engaged values match, so record the first). is_gguf keys off
# kind==gguf AND no quant having engaged (a dense torchao DiT is not a GGUF one).
# Both helpers act on ``view.transformer``; call once per view to set the kernel and compile each
# expert (engaged values match, so record the first). is_gguf keys off kind==gguf AND no quant
# having engaged (a dense torchao DiT is not a GGUF one).
gguf_transformer = kind == "gguf" and transformer_quant_engaged is None
engaged = apply_attention_backend(
view,
@ -1335,8 +1325,8 @@ class VideoBackend:
is_gguf = gguf_transformer,
family = fam,
speed_mode = effective_speed,
# An auto cache that could still engage also drops fullgraph (FBCache under a
# fullgraph-compiled DiT crashes the first cached generation).
# An auto cache that could still engage also drops fullgraph (FBCache under a fullgraph-compiled
# DiT crashes the first cached generation).
cache_active = cache_engaged is not None or cache_may_toggle,
offload_active = plan.offload_policy != "none",
)
@ -1344,15 +1334,15 @@ class VideoBackend:
attention_engaged = engaged
speed_optims = tuple(k for k, v in applied.items() if v)
with self._generate_lock:
# A cancelled/superseded load must not place weights on a GPU the arbiter may have
# reassigned; recheck right before placement (the commit below does the final check).
# A cancelled/superseded load must not place weights on a GPU the arbiter may have reassigned;
# recheck right before placement (the commit below does the final check).
if _load_token is not None and _load_token != self._load_token:
del pipe
clear_gpu_cache()
raise RuntimeError("Video load was cancelled or superseded.")
offload_policy, vae_tiling = apply_memory_plan(pipe, plan, device = device, logger = logger)
# A dual-DiT MoE needs no extra per-expert pass: apply_memory_plan already covers every
# DiT (transformer AND transformer_2) under all tiers; a second pass would duplicate-hook.
# A dual-DiT MoE needs no extra per-expert pass: apply_memory_plan already covers every DiT under
# all tiers; a second pass would duplicate-hook.
if not vae_tiling:
# Whole-clip decode is the video memory peak; tiling is near-free, so always on.
try:
@ -1390,10 +1380,9 @@ class VideoBackend:
"transformer_quant": (
transformer_quant,
transformer_quant_engaged or "off",
# Honest framing: the shipped torchao schemes cut load time (hosted
# prequant) and resident memory ~2x, but measured on B200 the per-step
# GEMMs are at best parity with bf16 (int8 dynamic can be slower); the
# generation-speed lever is a calibrated static-scale fp8 path, not this.
# Honest framing: the shipped torchao schemes cut load time (hosted prequant) and resident memory
# ~2x, but measured on B200 the per-step GEMMs are at best parity with bf16; the generation-speed
# lever is a calibrated static-scale fp8 path, not this.
"DiT(s) quantised (halves resident weights; hosted checkpoints cut "
"load time; per-step speed is roughly bf16 parity)"
if transformer_quant_engaged is not None
@ -1626,8 +1615,8 @@ class VideoBackend:
with self._lock:
self._generate_job_active = False
if cancel_event is not None and self._active_generate_cancel is cancel_event:
# Covers a worker that failed before reaching generate()'s finally; identity-guarded
# so a direct generate() that re-registered keeps its cancel handle.
# Covers a worker that failed before reaching generate()'s finally; identity-guarded so a direct
# generate() that re-registered keeps its cancel handle.
self._active_generate_cancel = None
if error is not None:
self._gen = {
@ -1706,11 +1695,10 @@ class VideoBackend:
"num_frames": frames,
"generator": generator,
}
# The 2.3 distilled DiT was trained against a fixed 8-step sigma curve
# (ltx_core DISTILLED_SIGMA_VALUES); at the distilled default step count
# pass it verbatim, with the scheduler's re-shaping transforms neutralised
# for the call (they distort even explicit sigmas). Any other step count
# keeps the scheduler's own spacing.
# The 2.3 distilled DiT was trained against a fixed 8-step sigma curve (ltx_core
# DISTILLED_SIGMA_VALUES); at the distilled default step count pass it verbatim, with the
# scheduler's re-shaping transforms neutralised (they distort even explicit sigmas). Any other
# step count keeps the scheduler's own spacing.
sigma_ctx: Any = contextlib.nullcontext()
if fam.name == "ltx-2" and "sigmas" in call_params:
from .video_ltx2 import (
@ -1724,8 +1712,8 @@ class VideoBackend:
kwargs["sigmas"] = list(LTX23_DISTILLED_SIGMAS)
sigma_ctx = ltx23_verbatim_sigmas(pipe)
if fam.guidance_via_guider:
# HunyuanVideo-1.5: __call__ has no guidance kwarg; CFG scale is a guider
# attribute set per request (near-1 scales auto-disable CFG in the guider).
# HunyuanVideo-1.5: __call__ has no guidance kwarg; CFG scale is a guider attribute set per
# request (near-1 scales auto-disable CFG in the guider).
pipe.guider.guidance_scale = float(guidance)
else:
kwargs[fam.cfg_kwarg] = guidance
@ -1734,9 +1722,9 @@ class VideoBackend:
# LTX-2 takes frame_rate (shapes audio length); others fix their rate, fps only at export.
if "frame_rate" in call_params:
kwargs["frame_rate"] = float(out_fps)
# Dual-DiT MoE: thread the low-noise expert's guidance kwarg only when the family
# declares one AND the signature accepts it -- WanPipeline raises if guidance_scale_2
# is passed with boundary_ratio=None (pipeline_wan.py:322); TI2V-5B never reaches here.
# Dual-DiT MoE: thread the low-noise expert's guidance kwarg only when the family declares one
# AND the signature accepts it -- WanPipeline raises if guidance_scale_2 is passed with
# boundary_ratio=None (pipeline_wan.py:322); TI2V-5B never reaches here.
if fam.cfg2_kwarg and fam.cfg2_kwarg in call_params and guidance_2 is not None:
kwargs[fam.cfg2_kwarg] = float(guidance_2)
@ -1766,8 +1754,8 @@ class VideoBackend:
return callback_kwargs
def _on_scheduler_step(done: int) -> None:
# No cooperative _interrupt (the pipeline never checks it), so cancellation
# must unwind the denoise loop via an exception.
# No cooperative _interrupt (the pipeline never checks it), so cancellation must unwind the
# denoise loop via an exception.
if cancel.is_set():
raise _VideoGenerationCancelled()
_tick(done)
@ -1776,12 +1764,11 @@ class VideoBackend:
kwargs["callback_on_step_end"] = _on_step
progress_ctx = contextlib.nullcontext()
else:
# HunyuanVideo-1.5 has no step callback; each scheduler.step is one denoise
# step, so wrap it for progress + cancel and restore afterwards.
# HunyuanVideo-1.5 has no step callback; each scheduler.step is one denoise step, so wrap it for
# progress + cancel and restore afterwards.
progress_ctx = _scheduler_step_progress(pipe, _on_scheduler_step)
# Re-check an AUTO cache decision against the ACTUAL step count (a many-step
# request gains FBCache, a few-step drops it); explicit choices never toggle. Per view.
# Re-check an AUTO cache decision against the ACTUAL step count; explicit choices never toggle.
if state.cache_auto:
toggled = state.transformer_cache
for view in _views_for(pipe, fam):
@ -1809,8 +1796,8 @@ class VideoBackend:
with torch.inference_mode(), progress_ctx, sigma_ctx:
output = pipe(**kwargs)
except _VideoGenerationCancelled:
# Unwinding by exception skips the pipeline's end-of-call maybe_free_model_hooks();
# under offload the onloaded modules would stay on the GPU, so free them here.
# Unwinding by exception skips the pipeline's end-of-call maybe_free_model_hooks(); under offload
# the onloaded modules would stay on the GPU, so free them here.
free_hooks = getattr(pipe, "maybe_free_model_hooks", None)
if callable(free_hooks):
try:
@ -1828,8 +1815,8 @@ class VideoBackend:
mp4_bytes = self._encode_mp4(
video_frames, out_fps, audio_track, pipe if fam.has_audio else None
)
# A cancel during the blocking export/mux must still discard the clip; re-check
# before it is returned and persisted.
# A cancel during the blocking export/mux must still discard the clip; re-check before it is
# returned and persisted.
if cancel.is_set():
raise RuntimeError(VIDEO_CANCELLED_MSG)
duration_s = len(video_frames) / float(out_fps) if out_fps else 0.0
@ -1886,15 +1873,14 @@ class VideoBackend:
def generate_progress(self) -> dict[str, Any]:
with self._lock:
gen = dict(self._gen)
# generate() swaps in a bare {"active": False} before the worker records the terminal
# dict; report active across that gap so a poller sees active drop only with a terminal phase.
# generate() swaps in a bare {"active": False} before the worker records the terminal dict;
# report active across that gap so a poller sees active drop only with a terminal phase.
if self._generate_job_active:
gen["active"] = True
gen.setdefault("active", False)
# Mirror the image endpoint's field names (total_steps / fraction) alongside the
# native "total": the two generate-progress APIs used to disagree, so a client
# polling the image shape against video read total_steps=null / fraction=0 while
# the step counter advanced.
# Mirror the image endpoint's field names (total_steps / fraction) alongside the native "total":
# the two generate-progress APIs used to disagree, so a client polling the image shape against
# video read total_steps=null / fraction=0 while the step counter advanced.
total = int(gen.get("total") or 0)
step = int(gen.get("step") or 0)
gen["total_steps"] = total
@ -1918,8 +1904,8 @@ class VideoBackend:
state, self._state = self._state, None
if state is not None:
restore_backend_flags(state.backend_flags)
# A GGUF load may have installed the compiled GGUF dequantizer; restore the stock
# kernels so a later speed=off load gets the bit-identical path (mirrors image unload).
# A GGUF load may have installed the compiled GGUF dequantizer; restore the stock kernels so a
# later speed=off load gets the bit-identical path (mirrors image unload).
from . import diffusion_gguf_compile
diffusion_gguf_compile.uninstall_all()
@ -1933,8 +1919,8 @@ class VideoBackend:
self._loading = None
if self._active_generate_cancel is not None:
self._active_generate_cancel.set()
# Barrier: wait for the signalled generation to exit before freeing the pipeline, or we
# report the VRAM free (and let the arbiter start another load) while the clip still holds it.
# Barrier: wait for the signalled generation to exit before freeing the pipeline, else we report
# the VRAM free (and let the arbiter start another load) while the clip still holds it.
with self._generate_lock:
pass
self._teardown_state()

90
studio/backend/core/inference/video_families.py
View file

@ -39,19 +39,19 @@ class VideoFamily:
denoiser_attr: str = "transformer"
# Extra lowercased substrings (besides ``name``) that map a repo id here.
aliases: tuple[str, ...] = field(default_factory = tuple)
# True when the pipeline returns synchronized audio (LTX-2): export muxes the track
# and size estimates count the audio VAE + vocoder.
# True when the pipeline returns synchronized audio (LTX-2): export muxes the track and size
# estimates count the audio VAE + vocoder.
has_audio: bool = False
# Wan2.2-A14B dual-expert MoE: a second DiT (transformer_2) handles the low-noise
# steps with its own guidance kwarg. None/False for single-DiT.
# Wan2.2-A14B dual-expert MoE: a second DiT (transformer_2) handles the low-noise steps with its
# own guidance kwarg. None/False for single-DiT.
transformer2_class: Optional[str] = None
is_moe: bool = False
cfg2_kwarg: Optional[str] = None
# HunyuanVideo-1.5 guidance: __call__ takes NO guidance kwarg; CFG lives on a ``guider``
# component whose guidance_scale is set per request. When True, generate() writes pipe.guider.
guidance_via_guider: bool = False
# Generation defaults + shape. ``frame_step`` is the temporal compression: a valid frame
# count is k*frame_step + 1, so requests are snapped BEFORE latents are allocated.
# Generation defaults + shape. ``frame_step`` is the temporal compression: a valid frame count is
# k*frame_step + 1, so requests are snapped BEFORE latents are allocated.
default_steps: int = 40
default_guidance: float = 4.0
default_num_frames: int = 121
@ -68,20 +68,20 @@ class VideoFamily:
supports_torch_compile: bool = True
# Video DiTs are bf16-native, so fp16 promotes to float32; defaults True.
fp16_incompatible: bool = True
# Wan's VAE decodes in float32 (loading it bf16 causes banding / black frames), so when True
# the loader pins it back to fp32. Its bf16_components_gb term is already the fp32 size.
# Wan's VAE decodes in float32 (loading it bf16 causes banding / black frames), so when True the
# loader pins it back to fp32. Its bf16_components_gb term is already the fp32 size.
vae_force_fp32: bool = False
# Curated GGUF repo for the picker (the DiT as single-file GGUF quants).
gguf_repo: Optional[str] = None
# Hosted PRE-CAST text-encoder checkpoints as (scheme, component, repo_id) triples;
# same semantics as DiffusionFamily.te_prequant_repos (diffusion_te_prequant.py).
# Hosted PRE-CAST text-encoder checkpoints as (scheme, component, repo_id) triples; same
# semantics as DiffusionFamily.te_prequant_repos.
te_prequant_repos: tuple[tuple[str, str, str], ...] = field(default_factory = tuple)
_FAMILIES: tuple[VideoFamily, ...] = (
# LTX-2 (diffusers >= 0.39): ~19B single-stream video DiT generating synchronized audio +
# video in one pass. The Gemma3-12B text encoder is stored fp32 on the hub (~49 GB
# download; ~24 GB resident once cast to bf16). Base repo carries the dev config (40 steps, CFG 4); distilled runs few-step.
# LTX-2 (diffusers >= 0.39): ~19B single-stream video DiT generating synchronized audio + video
# in one pass. The Gemma3-12B text encoder is stored fp32 on the hub (~49 GB download, ~24 GB
# resident as bf16). Base repo carries the dev config (40 steps, CFG 4); distilled runs few-step.
VideoFamily(
name = "ltx-2",
pipeline_class = "LTX2Pipeline",
@ -97,20 +97,17 @@ _FAMILIES: tuple[VideoFamily, ...] = (
resolution_multiple = 32,
# 768x512 native default; 1216x704 the card's quality target; 704x1216 vertical.
resolution_presets = ((768, 512), (1216, 704), (704, 1216), (512, 768)),
# transformer 37.8 bf16; Gemma3-12B TE ~24.4 bf16 RESIDENT (the hub stores it fp32,
# ~49 GB download, but the pipeline loads torch_dtype=bf16); VAE 2.4 + connectors 2.9
# + audio 0.2. The old 50.4 figure double-counted the fp32 store and pushed the auto
# memory plan toward offload on cards that fit the real footprint.
# transformer 37.8 bf16; Gemma3-12B TE ~24.4 bf16 RESIDENT (the hub stores it fp32, ~49 GB
# download, but the pipeline loads torch_dtype=bf16); VAE 2.4 + connectors 2.9 + audio 0.2. The
# old 50.4 figure double-counted the fp32 store and pushed auto toward offload.
bf16_components_gb = (37.8, 24.4, 5.5),
gguf_repo = "unsloth/LTX-2.3-GGUF",
# Pre-cast Gemma3-12B TE (hub store is fp32 ~49 GB, pre-cast ~13.2 GB): the biggest
# download win of the hosted TE set.
# Pre-cast Gemma3-12B TE (fp32 ~49 GB on the hub, pre-cast ~13.2 GB): the biggest download win.
te_prequant_repos = (("fp8", "text_encoder", "unsloth/LTX-2-FP8"),),
),
# Wan2.2-TI2V-5B (diffusers >= 0.35, verified on 0.39): ~5B single-stream video DiT (UMT5
# text encoder). No audio, no second expert (boundary_ratio null, transformer_2 null), so
# single-DiT. Wan VAE temporal compression 4 -> valid frame counts 4k+1. Pipeline defaults
# 50 steps / CFG 5; UI presets target 720p at 24 fps.
# Wan2.2-TI2V-5B (diffusers >= 0.35, verified on 0.39): ~5B single-stream video DiT (UMT5 text
# encoder). No audio, no second expert, so single-DiT. Wan VAE temporal compression 4 gives valid
# frame counts 4k+1. Pipeline defaults 50 steps / CFG 5; UI presets target 720p at 24 fps.
VideoFamily(
name = "wan2.2-ti2v-5b",
pipeline_class = "WanPipeline",
@ -126,22 +123,21 @@ _FAMILIES: tuple[VideoFamily, ...] = (
default_fps = 24,
# Wan VAE temporal factor 4, so valid counts are 4k+1.
frame_step = 4,
# TI2V-5B VAE is 16x spatial + patch 2, so WanPipeline floors H/W to 32; snap to 32 so
# the recorded size matches the rendered clip (a /16-not-/32 request would render at 704).
# TI2V-5B VAE is 16x spatial + patch 2, so WanPipeline floors H/W to 32; snap to 32 so the
# recorded size matches the rendered clip.
resolution_multiple = 32,
# 720p-class presets (all /32); first is the default the loader plans against.
resolution_presets = ((1280, 704), (704, 1280), (960, 960), (832, 480)),
# bf16-RESIDENT. transformer + VAE ship FP32 on disk (index 20.0 GB = 5B x 4), so
# bf16 transformer ~10.0; UMT5 TE ships bf16 (11.4); VAE runs fp32 (2.8).
# bf16-RESIDENT. transformer + VAE ship FP32 on disk (index 20.0 GB = 5B x 4), so bf16
# transformer ~10.0; UMT5 TE ships bf16 (11.4); VAE runs fp32 (2.8).
bf16_components_gb = (10.0, 11.4, 2.8),
vae_force_fp32 = True,
gguf_repo = "QuantStack/Wan2.2-TI2V-5B-GGUF",
),
# Wan2.2-T2V-A14B (diffusers >= 0.35, verified on 0.39): the dual-expert MoE. Both
# transformer + transformer_2 are WanTransformer3DModel with boundary_ratio 0.875; the pipeline
# routes high-noise steps through transformer (guidance_scale) and low-noise through
# transformer_2 (guidance_scale_2, accepted only when boundary_ratio is set), so cfg2_kwarg is
# threaded ONLY here. boundary_ratio lives in the pipeline config, so no per-generation plumbing.
# Wan2.2-T2V-A14B (diffusers >= 0.35, verified on 0.39): the dual-expert MoE. Both transformers
# are WanTransformer3DModel with boundary_ratio 0.875; the pipeline routes high-noise steps
# through transformer (guidance_scale) and low-noise through transformer_2 (guidance_scale_2,
# accepted only when boundary_ratio is set), so cfg2_kwarg is threaded ONLY here.
VideoFamily(
name = "wan2.2-t2v-a14b",
pipeline_class = "WanPipeline",
@ -161,20 +157,20 @@ _FAMILIES: tuple[VideoFamily, ...] = (
default_fps = 16, # A14B runs at 16 fps (vs TI2V-5B's 24)
frame_step = 4,
resolution_multiple = 16,
# 480p + 720p presets (landscape + vertical). A14B's VAE is 8x so multiple 16 renders
# 720 (=45*16) exactly (unlike TI2V-5B's 16x VAE, which floors 720 to 704).
# 480p + 720p presets (landscape + vertical). A14B's VAE is 8x so multiple 16 renders 720 exactly
# (unlike TI2V-5B's 16x VAE, which floors 720 to 704).
resolution_presets = ((1280, 720), (832, 480), (480, 832), (720, 1280)),
# bf16-RESIDENT. Each expert ships FP32 (index 57.15 GB = 14.3B x 4) -> ~28.6 bf16 each ->
# ~57.2 for BOTH (the headline before offload), NOT the 114.3 fp32 sum. UMT5 TE bf16 (11.4); VAE fp32 (0.5).
# bf16-RESIDENT. Each expert ships FP32 (index 57.15 GB = 14.3B x 4), so ~28.6 bf16 each and
# ~57.2 for BOTH (the headline before offload), NOT the 114.3 fp32 sum. UMT5 TE bf16 (11.4);
# VAE fp32 (0.5).
bf16_components_gb = (57.2, 11.4, 0.5),
vae_force_fp32 = True,
# No gguf_repo: community GGUFs split the experts, and a single-file load covers only one.
),
# HunyuanVideo-1.5 (diffusers >= 0.39): 8.3B DiT, Qwen2.5-VL text encoder + ByT5 glyph
# encoder. Three quirks: (1) __call__ has NO guidance kwarg; CFG on the ``guider``
# (guidance_via_guider); (2) NO callback_on_step_end (generate() uses the scheduler.step
# wrapper); (3) tencent's repo is the original layout (no model_index.json), so only the
# community Diffusers repacks load. The transformer declares _repeated_blocks + CacheMixin.
# HunyuanVideo-1.5 (diffusers >= 0.39): 8.3B DiT, Qwen2.5-VL text encoder + ByT5 glyph encoder.
# Three quirks: (1) __call__ has NO guidance kwarg, CFG lives on the ``guider``; (2) NO
# callback_on_step_end (generate() wraps scheduler.step); (3) tencent's repo has no
# model_index.json, so only the community Diffusers repacks load.
VideoFamily(
name = "hunyuanvideo-1.5",
pipeline_class = "HunyuanVideo15Pipeline",
@ -194,12 +190,12 @@ _FAMILIES: tuple[VideoFamily, ...] = (
resolution_multiple = 16,
# 480p-class presets (the base is the 480p variant): landscape, vertical, square.
resolution_presets = ((832, 480), (480, 832), (624, 624)),
# DiT fp32 on disk (32.0 -> 16.6 bf16); VAE (4.7 -> 2.4); Qwen2.5-VL TE bf16 14.0 + ByT5 0.8.
# DiT fp32 on disk (32.0 to 16.6 bf16); VAE (4.7 to 2.4); Qwen2.5-VL TE bf16 14.0 + ByT5 0.8.
bf16_components_gb = (16.6, 14.8, 2.4),
),
# The 720p t2v repack: same architecture/quirks/footprint as the 480p entry; only the
# trained resolution differs. Own family so a 720p load defaults to 720p sizes. Its full-path
# alias out-lengths (and outranks) the generic "hunyuanvideo-1.5" token for this repo only.
# The 720p t2v repack: same architecture/quirks/footprint as the 480p entry, only the trained
# resolution differs. Own family so a 720p load defaults to 720p sizes. Its full-path alias
# out-lengths (and outranks) the generic "hunyuanvideo-1.5" token for this repo only.
VideoFamily(
name = "hunyuanvideo-1.5-720p",
pipeline_class = "HunyuanVideo15Pipeline",
@ -300,8 +296,8 @@ def default_video_generation_params(
for identifier in identifiers:
needle = (identifier or "").lower()
for key, steps, guidance in _VIDEO_GENERATION_DEFAULTS:
# Match the key as a name segment: reject a preceding ASCII letter so "swan-video"
# or "taiwan-clips" doesn't false-match "wan". Trailing chars stay free.
# Match the key as a name segment: reject a preceding ASCII letter so "swan-video" or
# "taiwan-clips" doesn't false-match "wan". Trailing chars stay free.
if re.search(r"(?<![a-z])" + re.escape(key), needle):
return steps, guidance
return fallback

21
studio/backend/core/inference/video_gallery.py
View file

@ -82,8 +82,8 @@ def transcode(video_id: str, fmt: str) -> Optional[bytes]:
"""Re-encode a stored MP4 for the Download menu: "webm" (VP9) or "gif". Returns the bytes, or
None when the id doesn't resolve. Raises RuntimeError on missing codec/deps (route 501s). MP4
downloads stream the original via /file, not here."""
# Ownership-gate like /file: only transcode a Studio-owned clip (readable sidecar), so a
# guessed stem for a foreign/orphan MP4 the gallery hides can't be re-encoded out either.
# Ownership-gate like /file: only transcode a Studio-owned clip (readable sidecar), so a guessed
# stem for a foreign/orphan MP4 the gallery hides can't be re-encoded out either.
path = owned_video_path(video_id)
if path is None:
return None
@ -113,8 +113,8 @@ def _transcode_webm(path: Path) -> bytes:
out_v.width = in_v.codec_context.width
out_v.height = in_v.codec_context.height
out_v.pix_fmt = "yuv420p"
# Realtime settings: VP9's default "good" profile is slow; cpu-used 8 + row-mt is much
# faster at a small quality cost, right for a download button.
# Realtime settings: VP9's default "good" profile is slow; cpu-used 8 + row-mt is much faster at
# a small quality cost, right for a download button.
out_v.options = {"deadline": "realtime", "cpu-used": "8", "row-mt": "1"}
for frame in src.decode(in_v):
for packet in out_v.encode(frame.reformat(format = "yuv420p")):
@ -174,10 +174,9 @@ def _sidecar_path(video_id: str) -> Path:
return gallery_dir() / f"{video_id}.json"
# Sidecar keys every genuine Studio record carries (save() always writes them). delete()/clear()
# own a pair only when its sidecar has all of these, so a hand-dropped MP4 with an empty ("{}") or
# partial sidecar -- which list_videos already hides -- is neither counted as ours nor destroyed.
# Key-presence only (the route owns full schema validation); mirrors image_gallery._REQUIRED_META.
# Sidecar keys every genuine Studio record carries. delete()/clear() own a pair only when its
# sidecar has all of these, so a hand-dropped MP4 with an empty or partial sidecar is neither
# counted as ours nor destroyed. Key-presence only (the route owns schema validation).
_REQUIRED_META = (
"prompt",
"width",
@ -273,9 +272,9 @@ def delete(video_id: str) -> bool:
# list_videos, so a guessed id must not destroy it.
if _read_meta(_sidecar_path(video_id)) is None:
return False
# Delete the MP4 FIRST: if the sidecar were dropped first and the mp4 unlink then failed (lock /
# permission), the still-present mp4 would vanish from the gallery with no retry. mp4-first means
# the worst case is an orphaned sidecar, which list_videos ignores.
# Delete the MP4 FIRST: dropping the sidecar first and failing to unlink the mp4 would leave a
# clip that vanished from the gallery with no retry. mp4-first leaves at worst an orphan sidecar,
# which list_videos ignores.
try:
path.unlink()
except OSError as exc:

17
studio/backend/core/inference/video_ltx2.py
View file

@ -30,8 +30,8 @@ from loggers import get_logger
logger = get_logger(__name__)
# Companion files (text projections, VAEs incl. vocoder) next to the quants in unsloth's GGUF repo:
# the official Lightricks weights split out of the combined checkpoint. Keyed by variant.
# Companion files (text projections, VAEs incl. vocoder) next to the quants in unsloth's GGUF
# repo: the official Lightricks weights split out of the combined checkpoint. Keyed by variant.
LTX23_EXTRAS_REPO = "unsloth/LTX-2.3-GGUF"
_EXTRAS_TEXT_PROJ = "text_encoders/ltx-2.3-22b-{variant}_embeddings_connectors.safetensors"
_EXTRAS_VIDEO_VAE = "vae/ltx-2.3-22b-{variant}_video_vae.safetensors"
@ -363,9 +363,8 @@ def ltx23_extras_files(checkpoint_path: Path | str) -> tuple[str, ...]:
# Upstream ltx_core's DISTILLED_SIGMA_VALUES: the fixed 8-step sampling curve the 22B distilled
# DiT was trained against (the scheduler appends the terminal 0 itself). The base scheduler's
# resolution-shifted flow-match spacing lands FAR from it at every mu the pipeline can compute
# (measured second sigma 0.945-0.981 vs 0.99375, and a 0.37-0.61 -> 0.1 tail vs 0.725 -> 0.42),
# so the distilled default of 8 steps must pass this list verbatim.
# resolution-shifted spacing lands far from it at every mu the pipeline can compute, so the
# distilled default of 8 steps must pass this list verbatim.
LTX23_DISTILLED_SIGMAS: tuple[float, ...] = (
1.0,
0.99375,
@ -527,8 +526,8 @@ def load_ltx23_audio_vae_and_vocoder(
_AUDIO_VAE_RENAME,
torch_dtype,
)
# The 2.3 vocoder is a composite (base + bandwidth-extension stack + mel STFT buffers); keys
# line up module-for-module after the renames.
# The 2.3 vocoder is a composite (base + bandwidth-extension stack + mel STFT buffers); keys line
# up module-for-module after the renames.
vocoder_state = _apply_rename(_to_plain_dtype(vocoder_state, torch_dtype), _VOCODER_RENAME)
for key in [k for k in vocoder_state if ".ups." in k]:
vocoder_state[key.replace(".ups.", ".upsamplers.")] = vocoder_state.pop(key)
@ -570,8 +569,8 @@ def load_ltx23_pipeline(
del state
# The Lightricks fp8 single files store SCALED float8 weights (.weight_scale/.input_scale
# companions). Casting without the scales corrupts every quantized layer, so refuse loudly;
# use the GGUF quants (Q8_0 for highest fidelity) instead.
# companions). Casting without the scales corrupts every quantized layer, so refuse loudly; use
# the GGUF quants (Q8_0 for highest fidelity) instead.
if any(k.endswith((".weight_scale", ".input_scale")) for k in groups["dit"]):
raise ValueError(
"This LTX checkpoint stores scaled fp8 weights, which this loader does "

51
studio/backend/core/training/diffusion_lora_trainer.py
View file

@ -42,8 +42,8 @@ from dataclasses import replace
from pathlib import Path
from typing import Any, Optional
# Shared, family-agnostic building blocks. Re-exported so callers/tests that import them
# from this module (its historical home) keep working unchanged.
# Shared, family-agnostic building blocks. Re-exported so callers/tests that import them from
# this module (its historical home) keep working unchanged.
from core.training.diffusion_train_common import ( # noqa: F401
DEFAULT_LORA_FILENAME,
DEFAULT_LORA_TARGETS,
@ -223,8 +223,8 @@ def _build_sdxl_latent_cache(
a = _hold(dist.mean * vae_scale)
b = _hold(dist.std * vae_scale)
if not forced and not gated:
# Size-gate the auto cache off the first real variant, before building the rest
# (it can exhaust host/pinned RAM). Over budget: bail with the VAE still resident.
# Size-gate the auto cache off the first real variant, before building the rest (it can exhaust
# host/pinned RAM). Over budget: bail with the VAE still resident.
per_variant = a.numel() * a.element_size() + b.numel() * b.element_size()
if _latent_cache_over_budget(per_variant, total_variants):
_emit(
@ -320,8 +320,8 @@ def run_diffusion_lora_training(
precision = "fp16"
weight_dtype = {"bf16": torch.bfloat16, "fp16": torch.float16, "no": torch.float32}[precision]
# TF32 / cudnn.benchmark for the run, restored on the way out (keeps in-process callers
# clean). Wraps the whole body so every return restores the backend flags.
# TF32 / cudnn.benchmark for the run, restored on the way out. Wraps the whole body so every
# return restores the backend flags.
snap = _apply_perf_flags(cfg, device)
try:
# Preflight the base model against the same trust gate as inference, before any fetch.
@ -330,8 +330,8 @@ def run_diffusion_lora_training(
pairs = discover_image_caption_pairs(
cfg.data_dir, instance_prompt = cfg.instance_prompt, caption_column = cfg.caption_column
)
# Resolve num_epochs -> a concrete train_steps now the dataset size is known, and rebind
# cfg so every downstream read sees the same value.
# Resolve num_epochs into a concrete train_steps now the dataset size is known, and rebind cfg so
# every downstream read sees the same value.
cfg = replace(cfg, train_steps = resolve_train_steps(cfg, len(pairs)), num_epochs = 0)
_emit(on_event, "model_load_started", num_images = len(pairs))
@ -378,8 +378,8 @@ def run_diffusion_lora_training(
if weight_dtype != torch.float32:
cast_training_params(unet, dtype = torch.float32)
# Regionally torch.compile the U-Net's repeated blocks via the DiT trainer's never-fatal
# wrapper (failure falls back to eager with a warning). Dense bf16 base compiles under "auto".
# Regionally torch.compile the U-Net's repeated blocks via the DiT trainer's never-fatal wrapper
# (failure falls back to eager with a warning). Dense bf16 base compiles under "auto".
from core.training.diffusion_dit_trainer import _maybe_compile_transformer
compiled = _maybe_compile_transformer(
@ -388,8 +388,8 @@ def run_diffusion_lora_training(
lora_params = [p for p in unet.parameters() if p.requires_grad]
optimizer = _make_lora_optimizer(lora_params, cfg.learning_rate)
# The scheduler advances once per optimizer update (per opt_step, for cfg.train_steps
# total). Count warmup/decay in optimizer steps; the accumulation factor would stretch warmup.
# The scheduler advances once per optimizer update (cfg.train_steps total), so count warmup/decay
# in optimizer steps; the accumulation factor would stretch warmup.
lr_sched = get_scheduler(
cfg.lr_scheduler,
optimizer = optimizer,
@ -401,8 +401,8 @@ def run_diffusion_lora_training(
prediction_type = noise_scheduler.config.prediction_type
# Precompute text embeddings once per unique caption, then free the ~1.5 GB CLIP encoders.
# Deterministic and RNG-free, so the training math is bit-identical to in-loop encoding.
# The env toggle lets the accuracy guard A/B the two paths.
# Deterministic and RNG-free, so the training math is bit-identical to in-loop encoding. The env
# toggle lets the accuracy guard A/B the two paths.
precompute = os.environ.get("UNSLOTH_DIFFUSION_NO_PRECOMPUTE", "") not in ("1", "true")
caption_embeds: dict[str, tuple] = {}
if precompute:
@ -416,8 +416,8 @@ def run_diffusion_lora_training(
if device == "cuda":
torch.cuda.empty_cache()
# Precompute the VAE latent cache, then free the VAE: it holds the posterior affine pair
# so per-step sampling noise is preserved. The env toggle A/Bs cached vs in-loop encode.
# Precompute the VAE latent cache, then free the VAE: it holds the posterior affine pair so
# per-step sampling noise is preserved. The env toggle A/Bs cached vs in-loop encode.
use_cache = cfg.cache_latents and os.environ.get(
"UNSLOTH_DIFFUSION_NO_LATENT_CACHE", ""
) not in ("1", "true")
@ -463,13 +463,13 @@ def run_diffusion_lora_training(
_emit(on_event, "model_load_completed", compiled = compiled)
# Permutation-cycle index sampler: each image is visited once per cycle before any repeat,
# so a short run doesn't leave a small dataset partly unseen.
# Permutation-cycle index sampler: each image is visited once per cycle before any repeat, so a
# short run doesn't leave a small dataset partly unseen.
index_sampler = PermutationBatchSampler(len(pairs), rng)
def _next_batch() -> tuple[list[int], list[str], list[str]]:
# Draw the full configured batch, not min(batch, n): the sampler refills across cycles
# so a dataset smaller than train_batch_size still yields exactly that many indices.
# Draw the full configured batch, not min(batch, n): the sampler refills across cycles so a
# dataset smaller than train_batch_size still yields exactly that many indices.
idx = index_sampler.next_batch(cfg.train_batch_size)
chosen = [pairs[i] for i in idx]
return idx, [c[0] for c in chosen], [c[1] for c in chosen]
@ -552,8 +552,8 @@ def run_diffusion_lora_training(
step_loss += float(loss.detach()) / cfg.gradient_accumulation_steps
micro += 1
# max_grad_norm <= 0 disables clipping (Studio sends 0.0); passing 0.0 to
# clip_grad_norm_ would zero every gradient (no learning).
# max_grad_norm at or below 0 disables clipping (Studio sends 0.0); passing 0.0 to
# clip_grad_norm_ would zero every gradient.
grad_norm = None
if cfg.max_grad_norm and cfg.max_grad_norm > 0:
# Returned value is the total PRE-clip norm, reported to the UI chart.
@ -565,8 +565,7 @@ def run_diffusion_lora_training(
done = opt_step + 1
now = time.time()
if done == 1:
# Step 1 pays the one-time costs (cudnn autotune, compile warmup), so the rate
# starts after it and reflects steady state.
# Step 1 pays the one-time costs (cudnn autotune, compile warmup), so the rate starts after it.
t_steady = now
if done % cfg.log_every == 0 or done == cfg.train_steps:
# ``learning_rate`` (not ``lr``) is the field the Studio training pump reads.
@ -671,8 +670,8 @@ def run_diffusion_training_process(*, event_queue: Any, stop_queue: Any, config:
return got if saw else False
try:
# normalized() resolves + validates the family; dispatch through the registry so a DiT
# family runs its own trainer while SDXL keeps this loop.
# normalized() resolves + validates the family; dispatch through the registry so a DiT family
# runs its own trainer while SDXL keeps this loop.
cfg = _config_from_dict(config).normalized()
trainer = get_trainer(cfg.resolved_family)
trainer(cfg, on_event = on_event, should_stop = should_stop)

269
studio/backend/core/training/diffusion_train_common.py
View file

@ -34,15 +34,13 @@ from core.inference.diffusion_families import (
# Default LoRA target modules: the attention projections common to the SDXL U-Net and the DiT
# transformers (the diffusers/kohya convention). A family wanting a wider set overrides this in
# its own defaults; kept here so DiffusionLoraConfig has a sane fallback.
# its own defaults.
DEFAULT_LORA_TARGETS: tuple[str, ...] = ("to_k", "to_q", "to_v", "to_out.0")
# diffusers' SchedulerType names (diffusers.optimization.get_scheduler). piecewise_constant is
# excluded: it is the only scheduler needing a `step_rules` string, which the trainers never pass
# (and there is no config field for it). Accepting it would pass normalized(), free the resident
# GPU workloads, then crash in the child (get_piecewise_constant_schedule does
# step_rules.split(",") on None) -- the evict-then-fail this validation prevents. The other six
# run with only warmup/training steps.
# diffusers' SchedulerType names. piecewise_constant is excluded: it is the only scheduler
# needing a `step_rules` string the trainers never pass, so accepting it would pass normalized(),
# free the resident GPU workloads, then crash in the child -- the evict-then-fail this validation
# prevents. The other six run with only warmup/training steps.
_LR_SCHEDULERS: frozenset[str] = frozenset(
{
"linear",
@ -54,8 +52,8 @@ _LR_SCHEDULERS: frozenset[str] = frozenset(
}
)
# DiT families whose fp32 RoPE/embedder overflow fp16, so they train in bf16 only. Must stay
# in sync with the DiT trainer's own specs (kept separate to avoid an import cycle).
# DiT families whose fp32 RoPE/embedder overflow fp16, so they train in bf16 only. Must stay in
# sync with the DiT trainer's own specs (kept separate to avoid an import cycle).
_FORCE_BF16_FAMILIES: frozenset[str] = frozenset(
{"qwen-image", "z-image", "krea-2", "flux.2-klein", "flux.2-dev"}
)
@ -65,16 +63,15 @@ _CAPTION_EXTS = (".txt", ".caption")
# diffusers' canonical single-file LoRA name, so load_lora_weights(dir) finds it.
DEFAULT_LORA_FILENAME = "pytorch_lora_weights.safetensors"
# Architectures Studio can neither train nor load, so they are not in the family registry but are
# recognisable by name. Rejecting them by name turns a confusing mid-run crash into a clear error.
# Registry families (flux / qwen-image / z-image / kontext) are handled by the positive check in
# ``resolve_trainable_family``, so they are intentionally absent here.
# Architectures Studio can neither train nor load: not in the family registry but recognisable by
# name, so rejecting them by name turns a confusing mid-run crash into a clear error. Registry
# families are handled by the positive check in ``resolve_trainable_family``.
_NON_TRAINABLE_RESIDUAL_TOKENS = frozenset({"sd3", "pixart", "sana", "lumina", "cogview"})
_NON_TRAINABLE_RESIDUAL_PHRASES = ("stable-diffusion-3", "hunyuan-dit")
EventCb = Callable[[dict[str, Any]], None]
# Returns a falsy value to keep training, or a truthy stop signal: bare True, or a dict
# that may carry ``save=False`` to cancel without saving a partial adapter.
# Returns a falsy value to keep training, or a truthy stop signal: bare True, or a dict that may
# carry ``save=False`` to cancel without saving a partial adapter.
StopCb = Callable[[], Any]
@ -104,10 +101,10 @@ def resolve_trainable_family(base_model: str, model_family: Optional[str] = None
compatible: a genuinely wrong pick still fails cleanly later in from_pretrained).
"""
name = str(base_model or "").strip().lower()
# GGUF weights (a ``.gguf`` file or ``*-GGUF`` repo) are inference-only: training needs the
# full diffusers pipeline (transformer + VAE + text encoders), which a GGUF repo lacks. Reject
# by name even when the family is trainable. Exempt a local diffusers checkout that merely has
# "gguf" in its path, identified by its ``model_index.json`` marker, not a bare ``is_dir()``.
# GGUF weights (a ``.gguf`` file or ``*-GGUF`` repo) are inference-only: training needs the full
# diffusers pipeline, which a GGUF repo lacks. Reject by name even when the family is trainable.
# Exempt a local diffusers checkout that merely has "gguf" in its path, identified by its
# ``model_index.json`` marker, not a bare ``is_dir()``.
local = Path(base_model).expanduser() if base_model else None
is_local_diffusers = bool(local and (local / "model_index.json").is_file())
if name.endswith(".gguf") or ("gguf" in name and not is_local_diffusers):
@ -247,9 +244,9 @@ def get_trainer(family: str) -> Callable[..., str]:
# Families absent here fall back to the DiffusionLoraConfig defaults.
FAMILY_TRAIN_DEFAULTS: dict[str, dict[str, Any]] = {
"sdxl": {"lora_rank": 16, "learning_rate": 1e-4, "resolution": 1024},
# Warmup defaults: a short LR ramp keeps the first adapter updates from overshooting on
# the big flow-matching DiTs (whose logit-normal timestep draw concentrates loss mass
# mid-schedule); the small warmups below are scaled for Studio's short-run step budgets.
# Warmup defaults: a short LR ramp keeps the first adapter updates from overshooting on the big
# flow-matching DiTs (whose logit-normal timestep draw concentrates loss mass mid-schedule);
# these are scaled for Studio's short-run step budgets.
"flux.1": {"lora_rank": 16, "learning_rate": 1e-4, "resolution": 512, "lr_warmup_steps": 20},
"qwen-image": {
"lora_rank": 16,
@ -261,8 +258,8 @@ FAMILY_TRAIN_DEFAULTS: dict[str, dict[str, Any]] = {
# The Krea 2 authors' recommended starting point (their DreamBooth script defaults):
# rank/alpha 32, lr 3e-4, 512px.
"krea-2": {"lora_rank": 32, "learning_rate": 3e-4, "resolution": 512},
# The upstream FLUX.2 DreamBooth references default to rank 16 / lr 1e-4; FLUX.2's
# uniform timestep draw benefits most from a warmup ramp.
# The upstream FLUX.2 DreamBooth references default to rank 16 / lr 1e-4; FLUX.2's uniform
# timestep draw benefits most from a warmup ramp.
"flux.2-klein": {
"lora_rank": 16,
"learning_rate": 1e-4,
@ -283,8 +280,8 @@ def train_defaults(family: str) -> dict[str, Any]:
return dict(FAMILY_TRAIN_DEFAULTS.get((family or "").strip().lower(), {}))
# Display labels + a short VRAM/access note per trainable family, surfaced by the Train UI
# so users pick a base with realistic expectations. Kept next to the defaults they pair with.
# Display labels + a short VRAM/access note per trainable family, surfaced by the Train UI so
# users pick a base with realistic expectations.
_FAMILY_LABELS = {
"sdxl": "SDXL",
"flux.1": "FLUX.1-dev",
@ -306,8 +303,7 @@ _FAMILY_VRAM_NOTES = {
# The flow-matching DiT families (run by diffusion_dit_trainer). They expose the base_precision /
# compile levers and require bf16 compute on CUDA; SDXL is absent (it uses its own
# mixed_precision path). A set so the UI gate, the bf16 preflight, and any future dispatch stay
# in sync.
# mixed_precision path). A set so the UI gate, the bf16 preflight and any dispatch stay in sync.
_DIT_TRAIN_FAMILIES = frozenset(
{"flux.1", "qwen-image", "z-image", "krea-2", "flux.2-klein", "flux.2-dev"}
)
@ -370,10 +366,9 @@ def dit_accelerator_missing_reason(resolved_family: str) -> Optional[str]:
try:
import torch
def probe(owner: Any) -> bool:
# Each accelerator is probed on its own: one missing or throwing probe must not
# decide the other two. torch.mps.is_available() only exists from torch 2.5 and
# the supported floor is 2.4, so a shared try/except here would swallow the
# AttributeError and wave a CPU-only host through the gate it exists for.
# Each accelerator is probed on its own: one missing or throwing probe must not decide the other
# two. torch.mps.is_available() only exists from torch 2.5 while the supported floor is 2.4, so a
# shared try/except would swallow the AttributeError and wave a CPU-only host through.
try:
fn = getattr(owner, "is_available", None)
return bool(fn()) if callable(fn) else False
@ -414,10 +409,9 @@ def training_precision_preflight_error(resolved_family: str, base_precision: str
fam = (resolved_family or "").strip().lower()
mode = (base_precision or "").strip().lower()
if fam in _DIT_TRAIN_FAMILIES and mode in ("bf16", "int8", "fp8", "mxfp8"):
# The DiT trainer's dense precisions all require CUDA (_resolve_base_precision rejects
# bf16/int8/fp8/mxfp8 on device != "cuda"). bf16_unsupported_reason exempts a CPU-only host
# (the fp32 fallback for tests), so without this a dense request on a GPU-less host would
# pass the preflight, evict residents, then raise only in the child.
# The DiT trainer's dense precisions all require CUDA (_resolve_base_precision rejects them on
# device != "cuda"). bf16_unsupported_reason exempts a CPU-only host, so without this a dense
# request on a GPU-less host would pass the preflight, evict residents, then raise in the child.
try:
import torch
has_cuda = torch.cuda.is_available()
@ -433,10 +427,9 @@ def training_precision_preflight_error(resolved_family: str, base_precision: str
"base_precision='int8' needs a functional torchao install; this host's torchao is "
"missing or the non-functional Windows-ROCm stub. Use 'nf4', 'bf16', or 'auto'."
)
# mxfp8 needs Blackwell (sm100+): its MX GEMM has no kernel below sm100 and raises at the
# first training step, AFTER a full dense load. Re-check here (mirroring
# _resolve_base_precision) so a stale/direct client on an older GPU fails fast before
# eviction instead of crashing mid-run.
# mxfp8 needs Blackwell (sm100+): its MX GEMM has no kernel below sm100 and raises at the first
# training step, AFTER a full dense load. Re-check here (mirroring _resolve_base_precision) so a
# stale/direct client fails fast before eviction.
if mode == "mxfp8":
try:
import torch
@ -465,15 +458,13 @@ def family_train_infos() -> list[dict[str, Any]]:
if fam is None:
continue
repos = list(fam.train_base_repos) or [fam.base_repo]
# base_precision applies to the DiT trainer only; SDXL keeps its mixed_precision lever, so
# the UI hides the precision selector for it. compile applies everywhere: the SDXL trainer
# regionally compiles the U-Net's transformer blocks too.
# base_precision applies to the DiT trainer only; SDXL keeps its mixed_precision lever, so the UI
# hides the precision selector for it. compile applies everywhere.
is_dit = name in _DIT_TRAIN_FAMILIES
# On a non-bf16 CUDA GPU the start preflight rejects EVERY DiT family (even nf4, since the
# DiT trainer requires bf16 on CUDA), so advertise no precision -- else /info offers an nf4
# DiT option that always 400s. Otherwise drop any scheme this family's DiT corrupts (fp8 on
# Qwen-Image: activation outliers exceed fp8's range; the inference path denies the same
# set), so the UI never offers a mode normalized() would reject.
# On a non-bf16 CUDA GPU the start preflight rejects EVERY DiT family (the DiT trainer requires
# bf16 on CUDA), so advertise no precision, else /info offers an nf4 DiT option that always 400s.
# Otherwise drop any scheme this family's DiT corrupts (fp8 on Qwen-Image; the inference path
# denies the same set), so the UI never offers a mode normalized() would reject.
dit_block = (
bf16_unsupported_reason(name) or dit_accelerator_missing_reason(name)
if is_dit
@ -493,8 +484,8 @@ def family_train_infos() -> list[dict[str, Any]]:
"vram_note": dit_block or _FAMILY_VRAM_NOTES.get(name, ""),
"precision_modes": fam_modes,
"recommended_precision": "nf4" if (not is_dit or dit_block) else dit_recommended,
# compile is offered everywhere (SDXL regional U-Net + DiT), except a DiT family
# the GPU can't train in bf16 (dit_block), where training is refused outright.
# compile is offered everywhere (SDXL regional U-Net + DiT), except a DiT family the GPU can't
# train in bf16, where training is refused outright.
"supports_compile": bool(not dit_block),
# Krea trains on Raw but previews adapters on Turbo; None elsewhere.
"deploy_base": fam.deploy_base_repo,
@ -511,13 +502,13 @@ class DiffusionLoraConfig:
base_model: str
data_dir: str
output_dir: str
# Dreambooth-style caption applied to any image without its own caption. Required if
# the dataset has no captions.jsonl / sidecar files.
# Dreambooth-style caption applied to any image without its own caption. Required if the dataset
# has no captions.jsonl / sidecar files.
instance_prompt: Optional[str] = None
resolution: int = 1024
train_steps: int = 500
# 0 = disabled (train for train_steps). > 0 overrides train_steps with a run length of
# num_epochs full passes over the dataset, in optimizer steps (see resolve_train_steps).
# 0 = disabled (train for train_steps). Above 0 it overrides train_steps with num_epochs full
# passes over the dataset, in optimizer steps (see resolve_train_steps).
num_epochs: int = 0
learning_rate: float = 1e-4
train_batch_size: int = 1
@ -539,23 +530,21 @@ class DiffusionLoraConfig:
adapter_name: str = "default"
hf_token: Optional[str] = None
# Precompute the VAE latents once (freeing the VAE for the run) instead of re-encoding every
# step. ``cache_variants`` crop/flip draws are frozen per image; the per-step VAE sampling
# noise itself is preserved (see the DiT trainer docstring).
# step. ``cache_variants`` crop/flip draws are frozen per image; the per-step VAE sampling noise
# itself is preserved.
cache_latents: bool = True
cache_variants: int = 4
# Persistent on-disk conditioning cache directory (DiT trainer only). None (default)
# keeps the in-memory-only behavior; a configured directory turns the cache on: latent
# posterior stats and caption embeddings persist as safetensors keyed by content hash +
# family + resolution, and a fully warm cache skips loading the VAE and the multi-GB
# text encoders entirely on the next run.
# Persistent on-disk conditioning cache directory (DiT trainer only). None (default) keeps the
# in-memory-only behavior; a configured directory persists latent posterior stats and caption
# embeddings as safetensors keyed by content hash + family + resolution, so a fully warm cache
# skips loading the VAE and the multi-GB text encoders entirely.
cond_cache_dir: Optional[str] = None
# LoRA EMA decay (DiT trainer only). 0.0 (default) disables it; > 0 keeps an
# exponential moving average of the trainable LoRA params (warmup-ramped so short runs
# still absorb the trajectory) and exports it as a second adapter under
# ``<output_dir>/ema`` next to the primary one.
# LoRA EMA decay (DiT trainer only). 0.0 (default) disables it; a positive value keeps an
# exponential moving average of the trainable LoRA params (warmup-ramped so short runs still
# absorb the trajectory) and exports it as a second adapter in the ema subdir.
ema_decay: float = 0.0
# Regional torch.compile of the transformer blocks: "off" | "on" | "auto" (auto turns
# it on only for a dense, non-bitsandbytes base where it is a clean win).
# Regional torch.compile of the transformer blocks: "off" | "on" | "auto" (auto turns it on only
# for a dense, non-bitsandbytes base where it is a clean win).
compile_transformer: str = "auto"
# TF32 matmuls + high fp32 matmul precision + cudnn autotuning for the run. Near-lossless;
# disable for strict bit-reproducibility A/Bs.
@ -565,18 +554,16 @@ class DiffusionLoraConfig:
# "fp8" (torchao float8 training on the frozen linears, Ada/Hopper/Blackwell + compile), or
# "auto" (by free VRAM + GPU class). Non-nf4 modes need a dense base repo. SDXL ignores it.
base_precision: str = "nf4"
# Training-time timestep shift applied to the flow-matching sigma draw. None resolves
# per family in normalized(): "auto" for qwen-image (reproduce the family's inference
# sigma distribution: the scheduler's exponential time_shift at mu = max_shift, then the
# shift_terminal stretch), 1.0 (identity, the historical behavior) for every other
# family. A numeric value applies the standard linear shift s*u/(1+(s-1)*u); 1.0 is a
# no-op. "auto" on a family without dynamic shifting falls back to identity.
# Training-time timestep shift applied to the flow-matching sigma draw. None resolves per family
# in normalized(): "auto" for qwen-image (reproducing the family's inference sigma distribution),
# 1.0 (identity, the historical behavior) elsewhere. A numeric value applies the standard linear
# shift s*u/(1+(s-1)*u). "auto" on a family without dynamic shifting falls back to identity.
flow_shift: Optional[Any] = None # float | "auto" | None
# Per-sample probability of replacing the caption conditioning with the empty prompt
# (classifier-free-guidance dropout). 0.0 (default) disables it entirely.
cfg_dropout: float = 0.0
# Per-sample loss weighting over the drawn timestep: "none" (default, unweighted MSE)
# or "bell" (bsmntw-style Gaussian bell centered mid-schedule, normalized to mean 1).
# Per-sample loss weighting over the drawn timestep: "none" (default, unweighted MSE) or "bell"
# (bsmntw-style Gaussian bell centered mid-schedule, normalized to mean 1).
weighting_scheme: str = "none"
# How often to emit a progress event (in optimizer steps).
log_every: int = 1
@ -628,7 +615,7 @@ class DiffusionLoraConfig:
ema_decay = float(self.ema_decay or 0.0)
except (TypeError, ValueError) as exc:
raise ValueError(f"ema_decay must be a number, got {self.ema_decay!r}") from exc
# decay = 1.0 would freeze the shadow at its init forever; the EMA update is
# decay = 1.0 would freeze the shadow at its init forever; the update is
# shadow * decay + param * (1 - decay), so valid decays live in [0, 1).
if not 0.0 <= ema_decay < 1.0:
raise ValueError("ema_decay must be in [0, 1); 0 disables the EMA adapter")
@ -642,10 +629,9 @@ class DiffusionLoraConfig:
base_precision = str(self.base_precision or "nf4").strip().lower()
if base_precision not in ("nf4", "bf16", "int8", "fp8", "mxfp8", "auto"):
raise ValueError("base_precision must be one of nf4 / bf16 / int8 / fp8 / mxfp8 / auto")
# base_precision is a DiT-only lever (transformer load precision); SDXL uses its own
# mixed_precision path and ignores it, so the dense-mode gates (prequant base / non-bf16
# compute) apply only to the DiT families. The mode-name check above still runs for every
# family.
# base_precision is a DiT-only lever; SDXL uses its own mixed_precision path and ignores it, so
# the dense-mode gates (prequant base / non-bf16 compute) apply only to the DiT families. The
# mode-name check above still runs for every family.
if resolved_family != "sdxl" and base_precision in ("bf16", "int8", "fp8", "mxfp8"):
if repo_is_prequantized(self.base_model):
raise ValueError(
@ -658,10 +644,10 @@ class DiffusionLoraConfig:
f"base_precision={base_precision!r} trains in bf16 compute; set "
f"mixed_precision to bf16."
)
# Some DiT families are corrupted by fp8's activation range: outliers exceed even
# per-row fp8's range, so the frozen linears' float8 compute learns against a garbage
# forward pass. The inference path already denies these schemes; mirror it here so the
# run fails fast instead of producing a broken adapter. int8 (per-token) is unaffected.
# Some DiT families are corrupted by fp8's activation range: outliers exceed even per-row fp8's
# range, so the frozen linears' float8 compute learns against a garbage forward pass. The
# inference path already denies these schemes; mirror it here so the run fails fast. int8
# (per-token) is unaffected.
from core.inference.diffusion_transformer_quant import _family_denied
if _family_denied(resolved_family, base_precision):
@ -670,9 +656,9 @@ class DiffusionLoraConfig:
f"{resolved_family}: its activations exceed fp8's range and corrupt the "
f"trained result. Use 'nf4', 'int8', 'bf16', or 'auto'."
)
# flow_shift: None resolves to the family default ("auto" only for qwen-image, whose
# scheduler skips its static shift under use_dynamic_shifting and would otherwise
# train on unshifted uniform sigmas); an explicit value is validated and kept.
# flow_shift: None resolves to the family default ("auto" only for qwen-image, whose scheduler
# skips its static shift under use_dynamic_shifting and would otherwise train on unshifted
# uniform sigmas); an explicit value is validated and kept.
flow_shift = self.flow_shift
if flow_shift is None:
flow_shift = "auto" if resolved_family == "qwen-image" else 1.0
@ -687,11 +673,9 @@ class DiffusionLoraConfig:
) from exc
if not isinstance(flow_shift, str):
flow_shift = float(flow_shift)
# isfinite as well as positive: JSON accepts 1e309, which floats to inf, and
# inf <= 0 is False (NaN fails every comparison), so a positivity-only guard
# let it through to the sigma table, where s * u / (1 + (s - 1) * u) is NaN.
# That poisons every sampled sigma and the run saves a corrupted adapter while
# reporting normal progress.
# isfinite as well as positive: JSON accepts 1e309, which floats to inf, and inf is not caught by
# a positivity-only guard (NaN fails every comparison), so it would reach the sigma table where
# s * u / (1 + (s - 1) * u) is NaN, poisoning every sampled sigma while progress looks normal.
if not math.isfinite(flow_shift) or flow_shift <= 0:
raise ValueError(
"flow_shift must be a finite number > 0 (1.0 disables the shift), or 'auto'"
@ -705,12 +689,12 @@ class DiffusionLoraConfig:
weighting_scheme = str(self.weighting_scheme or "none").strip().lower()
if weighting_scheme not in ("none", "bell"):
raise ValueError("weighting_scheme must be one of none / bell")
# A zero/negative gamma would zero out (or invert) the min-SNR weight and
# silently train on a degenerate loss; None is the documented disable.
# A zero/negative gamma would zero out (or invert) the min-SNR weight and silently train on a
# degenerate loss; None is the documented disable.
if self.snr_gamma is not None and float(self.snr_gamma) <= 0:
raise ValueError("snr_gamma must be > 0, or null to disable min-SNR weighting")
# learning_rate can arrive as a string ("1e-4") from the Studio config path, which
# preserves it as a string after validation; coerce so AdamW receives a float.
# learning_rate can arrive as a string ("1e-4") from the Studio config path, so coerce it before
# AdamW sees it.
try:
learning_rate = float(self.learning_rate)
except (TypeError, ValueError) as exc:
@ -719,8 +703,8 @@ class DiffusionLoraConfig:
raise ValueError("learning_rate must be > 0")
alpha = self.lora_alpha if self.lora_alpha is not None else self.lora_rank
targets = tuple(self.lora_target_modules) or DEFAULT_LORA_TARGETS
# A blank Hub token (the Studio default when none is configured) must load
# anonymously, not as an explicit empty credential.
# A blank Hub token (the Studio default when none is configured) must load anonymously, not as an
# explicit empty credential.
token = self.hf_token.strip() if isinstance(self.hf_token, str) else self.hf_token
return replace(
self,
@ -782,9 +766,8 @@ class PermutationBatchSampler:
self._pos = 0
def next_batch(self, k: int) -> list[int]:
# k may exceed n (batch larger than the dataset): the permutation is refilled across as
# many cycles as needed so the caller always gets exactly k indices and the batch never
# shrinks, matching the old sampler's fixed batch shape.
# k may exceed n (batch larger than the dataset): the permutation is refilled across as many
# cycles as needed so the caller always gets exactly k indices and the batch never shrinks.
out: list[int] = []
while len(out) < k:
if self._pos >= len(self._order):
@ -836,8 +819,8 @@ def discover_image_caption_pairs(
meta_path = root / meta_name
if not meta_path.is_file():
continue
# Tolerate a bad upload (invalid UTF-8, or a line of non-object JSON): skip the record so
# the instance_prompt fallback still applies rather than crashing the trainer.
# Tolerate a bad upload (invalid UTF-8, or a line of non-object JSON): skip the record so the
# instance_prompt fallback still applies rather than crashing the trainer.
try:
meta_lines = meta_path.read_text(encoding = "utf-8").splitlines()
except (OSError, UnicodeError):
@ -862,10 +845,9 @@ def discover_image_caption_pairs(
for img in images:
caption: Optional[str] = None
sidecar_present = False
# 1. per-image sidecar caption file (the user's explicit edit; wins over metadata).
# An EMPTY sidecar is a deliberate tombstone (written when a user clears a caption): it
# suppresses the metadata caption but leaves the image uncaptioned so the instance_prompt
# fallback below still applies, rather than dropping the image.
# 1. per-image sidecar caption file (the user's explicit edit; wins over metadata). An EMPTY
# sidecar is a deliberate tombstone (written when a user clears a caption): it suppresses the
# metadata caption but leaves the image uncaptioned, so the instance_prompt fallback applies.
for ext in _CAPTION_EXTS:
sidecar = img.with_suffix(ext)
if sidecar.is_file():
@ -873,12 +855,12 @@ def discover_image_caption_pairs(
try:
caption = sidecar.read_text(encoding = "utf-8").strip()
except (OSError, UnicodeError):
# Unreadable sidecar reads as the empty tombstone, so the
# instance_prompt fallback applies instead of a 500 preflight.
# An unreadable sidecar reads as the empty tombstone, so the instance_prompt fallback applies
# instead of a 500 preflight.
caption = ""
break
# 2. metadata row keyed by file name (basename or relative path; as_posix so a Windows
# backslash path matches the jsonl's forward-slash keys). A sidecar, even empty, wins.
# 2. metadata row keyed by file name (basename or relative path; as_posix so a Windows backslash
# path matches the jsonl's forward-slash keys). A sidecar, even empty, wins.
if not sidecar_present:
caption = meta_caption.get(img.name) or meta_caption.get(
img.relative_to(root).as_posix()
@ -888,10 +870,9 @@ def discover_image_caption_pairs(
caption = instance_prompt
if caption:
if verify_images:
# Reject a corrupt/zero-byte/truncated image now via a cheap PIL header probe
# (verify() doesn't decode full pixels): otherwise it passes this filename-only
# discovery, the start route frees the resident GPU models, and the trainer only
# then crashes in Image.open -- the eviction this preflight prevents.
# Reject a corrupt/zero-byte/truncated image now via a cheap PIL header probe (verify() doesn't
# decode full pixels): otherwise it passes this filename-only discovery, the start route frees the
# resident GPU models, and the trainer only then crashes in Image.open.
try:
from PIL import Image
with Image.open(img) as _probe:
@ -944,15 +925,13 @@ def _plan_cache_variants(
# Host-memory budget for the AUTOMATIC latent cache. The cache holds two fp32 posterior tensors
# (mean/std, VAE scale folded in) per crop/flip variant per image, pinned on a CUDA host. At
# 1024px an SDXL variant is ~0.5 MiB and a 16-channel DiT variant several times that, so a few
# thousand images x cache_variants can exhaust host or pinned RAM. Over budget the default falls
# back to per-step VAE encoding. A fixed constant (not a psutil RAM fraction) keeps the gate
# dependency-free and identical across hosts; deliberately conservative, well under a typical
# host's RAM.
# thousand images x cache_variants can exhaust host or pinned RAM; over budget the default falls
# back to per-step VAE encoding. A fixed constant keeps the gate dependency-free across hosts.
_LATENT_CACHE_BUDGET_BYTES = 4 * 1024**3 # 4 GiB
# Returned by the cache builders when the estimate exceeds budget: the caller keeps the VAE
# resident and encodes each step's latents in-loop. A distinct sentinel from ``None`` (a stop
# requested mid-build) so the two are not conflated.
# requested mid-build).
LATENT_CACHE_OVER_BUDGET: Any = object()
@ -1005,17 +984,16 @@ def _apply_perf_flags(
torch.backends.cudnn.allow_tf32 = True
torch.set_float32_matmul_precision("high")
else:
# The opt-out is a strict-fp32 A/B mode, so actively clear the flags rather
# than inherit ambient state (cudnn TF32 defaults to ON in torch).
# The opt-out is a strict-fp32 A/B mode, so actively clear the flags rather than inherit ambient
# state (cudnn TF32 defaults to ON in torch).
torch.backends.cuda.matmul.allow_tf32 = False
torch.backends.cudnn.allow_tf32 = False
torch.set_float32_matmul_precision("highest")
if cudnn_benchmark:
torch.backends.cudnn.benchmark = True
# The cuDNN SDPA backend's TRAINING graph is broken for the FLUX attention shapes on
# torch 2.10 + cu130 (B200): mha_graph.execute fails, then poisons the context into
# illegal memory accesses. Flash / mem-efficient SDPA are equivalent, so pin those for the
# run (restored on exit).
# The cuDNN SDPA backend's TRAINING graph is broken for the FLUX attention shapes on torch 2.10 +
# cu130 (B200): mha_graph.execute fails, then poisons the context into illegal memory accesses.
# Flash / mem-efficient SDPA are equivalent, so pin those for the run (restored on exit).
cuda_backends = getattr(torch.backends, "cuda", None)
if cuda_backends is not None and hasattr(cuda_backends, "enable_cudnn_sdp"):
try:
@ -1042,8 +1020,8 @@ def _restore_perf_flags(snap: Optional[dict]) -> None:
if snap.get("matmul_precision"):
torch.set_float32_matmul_precision(snap["matmul_precision"])
# Restore the exact pre-run cudnn SDPA state; None means the flag was unreadable
# (or absent) at apply time and was never touched.
# Restore the exact pre-run cudnn SDPA state; None means the flag was unreadable (or absent) at
# apply time and was never touched.
cuda_backends = getattr(torch.backends, "cuda", None)
if (
snap.get("cudnn_sdp") is not None
@ -1056,11 +1034,9 @@ def _restore_perf_flags(snap: Optional[dict]) -> None:
# Official safetensors-only TRAINING bases trusted in addition to the inference allowlist
# (_TRUSTED_NON_GGUF_REPOS in core/inference/diffusion.py). The FLUX.2 bases are now in that
# list too (deploying a trained FLUX.2 adapter reloads the base as an inference pipeline), and
# are kept here so the training gate stays independent of a future edit to the loader list.
# Exact-match lowercased, same rules as the loader list: extend deliberately; never add pickled
# weights or remote code.
# (_TRUSTED_NON_GGUF_REPOS in core/inference/diffusion.py). The FLUX.2 bases are in that list too
# but kept here so the training gate stays independent of a future edit to it. Exact-match
# lowercased: extend deliberately; never add pickled weights or remote code.
_TRAIN_EXTRA_TRUSTED_REPOS = frozenset(
{
"black-forest-labs/flux.2-dev",
@ -1085,10 +1061,9 @@ def _assert_trusted_base_model(base_model: str) -> None:
f"Refusing to train from untrusted base model '{base_model}'. Use a local path or "
f"a trusted repo (an unsloth/* repo or an official base)."
)
# An existing LOCAL base is loaded as a full pipeline (from_pretrained(base_model)) by the
# trainer, which needs a model_index.json. Any existing path is "trusted" above, so reject a
# non-pipeline local dir here -- before /diffusion/start frees the GPU models -- rather than
# have the child fail after teardown.
# An existing LOCAL base is loaded as a full pipeline by the trainer, which needs a
# model_index.json. Any existing path is "trusted" above, so reject a non-pipeline local dir here,
# before /diffusion/start frees the GPU models, rather than have the child fail after teardown.
_assert_local_base_is_pipeline(base_model)
@ -1111,8 +1086,8 @@ def _publish_to_lora_catalog(lora_path: str, cfg: DiffusionLoraConfig) -> Option
alias = sanitize_alias(base)
src_resolved = Path(lora_path).resolve()
dest = loras_dir() / f"{alias}.safetensors"
# A retrain with the same adapter name must not clobber a prior mirror: pick the next
# free numeric suffix instead.
# A retrain with the same adapter name must not clobber a prior mirror: pick the next free
# numeric suffix instead.
if dest.exists() and dest.resolve() != src_resolved:
n = 2
while True:
@ -1148,14 +1123,12 @@ def _write_lora_sidecar(sidecar_path: Path, cfg: DiffusionLoraConfig) -> None:
# Aliases from the generic Studio training payload onto DiffusionLoraConfig fields, so the
# diffusion trainer can also be driven by the shared training request shape (not only its own
# request model, whose keys already match).
# diffusion trainer can also be driven by the shared training request shape.
_CONFIG_ALIASES = {
"model_name": "base_model",
"max_steps": "train_steps",
# The generic payload's num_epochs already matches the diffusion field name, but list it so
# the epochs override is threaded through the shared-payload path as explicitly as
# max_steps -> train_steps is.
# The generic payload's num_epochs already matches the diffusion field name, but list it so the
# epochs override is threaded through the shared-payload path as explicitly as max_steps is.
"num_epochs": "num_epochs",
"batch_size": "train_batch_size",
"lora_r": "lora_rank",
@ -1195,10 +1168,10 @@ def _config_from_dict(config: dict) -> DiffusionLoraConfig:
for k, v in config.items():
if k in valid:
kwargs[k] = v
# Epoch-mode payloads from the generic UI carry max_steps: 0 as the "use epochs" sentinel,
# which the max_steps -> train_steps alias copies as train_steps: 0. Since normalized()
# rejects train_steps < 1 before resolve_train_steps() applies num_epochs, drop a falsy/0
# train_steps when num_epochs > 0 so the dataclass default stands in until epoch resolution.
# Epoch-mode payloads from the generic UI carry max_steps: 0 as the "use epochs" sentinel, which
# the alias copies as train_steps: 0. normalized() rejects train_steps below 1 before
# resolve_train_steps() applies num_epochs, so drop a falsy train_steps when num_epochs is set and
# let the dataclass default stand in until epoch resolution.
try:
_num_epochs = int(kwargs.get("num_epochs") or 0)
except (TypeError, ValueError):

23
studio/backend/core/training/training.py
View file

@ -770,8 +770,8 @@ class TrainingBackend:
# True while a pump thread should be running; cleared on intended exits.
# Left True after an abnormal death so _ensure_pump_alive spots a crash.
self._pump_running: bool = False
# True from the start_training() guard passing until its spawn finishes; blocks a
# second concurrent start (routes call it from a worker thread, so starts can overlap).
# True from the start_training() guard passing until its spawn finishes; blocks a second
# concurrent start (routes call it from a worker thread, so starts can overlap).
self._start_in_progress: bool = False
self._lock = threading.Lock()
self._run_intent_lock = threading.RLock()
@ -851,10 +851,10 @@ class TrainingBackend:
still letting auto-selection place training against the freed memory.
Hook failures never block the start.
"""
# Compare-and-set start guard: the route runs this method on a worker thread, so two
# overlapping /train/start requests can reach it concurrently. Without the flag both
# would pass the alive-check below (the proc is only assigned at the end) and
# double-spawn. Mirrors the diffusion training service's reserve().
# Compare-and-set start guard: the route runs this method on a worker thread, so two overlapping
# /train/start requests can reach it concurrently. Without the flag both would pass the
# alive-check below (the proc is assigned only at the end) and double-spawn. Mirrors the
# diffusion training service's reserve().
with self._lock:
if self._start_in_progress:
logger.warning("Training start already in progress")
@ -874,8 +874,8 @@ class TrainingBackend:
with self._lock:
self._start_in_progress = False
# Named, not part of **kwargs: the body reads it directly, and it must not reach the
# worker config either (start_training's own signature keeps it out).
# Named, not part of **kwargs: the body reads it directly, and it must not reach the worker
# config either.
def _start_training_impl(
self,
job_id: str,
@ -1574,10 +1574,9 @@ class TrainingBackend:
# training invisibly behind a frozen UI. Cheap enough for per-second polls.
self._ensure_pump_alive()
with self._lock:
# A run reserved in start_training but not yet spawned (before_spawn frees residents,
# then GPU auto-selection, then proc.start()) is already active: the load/start guards
# read this to refuse a concurrent /images/load, /video/load, or /diffusion/start, so an
# idle reading here would let another pipeline race the reserved run for VRAM.
# A run reserved in start_training but not yet spawned (before_spawn frees residents, then GPU
# auto-selection, then proc.start()) is already active: the load/start guards read this to refuse
# a concurrent /images/load, /video/load, or /diffusion/start.
if self._start_in_progress:
return True

11
studio/backend/hub/services/download_lifecycle.py
View file

@ -430,10 +430,9 @@ def _try_http_retry(
cancel_marker_transport = original_metadata.transport,
hub_cache = original_metadata.hub_cache,
xet_cache = original_metadata.xet_cache,
# Carry the scoped file list across the reclaim. The record it overwrites is what a
# later start for this scope slot is compared against, so dropping it makes an
# identical scoped start read as a different file set and 409 instead of adopting
# the running download (the worker args below read the same list).
# Carry the scoped file list across the reclaim. The record it overwrites is what a later start
# for this scope slot is compared against, so dropping it makes an identical scoped start read as
# a different file set and 409 instead of adopting the running download.
scoped_files = original_metadata.scoped_files or None,
)
if claimed:
@ -468,8 +467,8 @@ def _try_http_retry(
args.append("--dataset")
elif variant:
args.extend(["--variant", variant])
# A scoped job must retry as the SAME scoped download; without its file list the
# HTTP worker would fall through to a full snapshot of the repo.
# A scoped job must retry as the SAME scoped download; without its file list the HTTP worker
# would fall through to a full snapshot of the repo.
if original_metadata.scoped_files:
args.extend(["--files-json", write_files_manifest(original_metadata.scoped_files)])

13
studio/backend/hub/services/models/deletion.py
View file

@ -603,13 +603,13 @@ def _diffusion_blocks_delete(repo_id: str) -> Optional[str]:
if status.get("loaded") and status.get("repo_id"):
if _loaded_id_matches_repo(str(status["repo_id"]), repo_id):
return "Unload the model before deleting"
# sd.cpp re-reads companion VAE / text-encoder files every generation, and
# status().repo_id covers only the main GGUF, so refuse the companions too.
# sd.cpp re-reads companion VAE / text-encoder files every generation, and status().repo_id
# covers only the main GGUF, so refuse the companions too.
for lid in getattr(engine, "loaded_repo_ids", tuple)():
if _loaded_id_matches_repo(str(lid), repo_id):
return "Unload the model before deleting"
# A downloading repo still reports loaded=False, but deleting would pull blobs
# from under the in-flight fetch.
# A downloading repo still reports loaded=False, but deleting would pull blobs from under the
# in-flight fetch.
for lid in getattr(engine, "loading_repo_ids", tuple)():
if _loaded_id_matches_repo(str(lid), repo_id):
return "An Images model load is using this repo; wait for it to finish"
@ -630,9 +630,8 @@ def _video_blocks_delete(repo_id: str) -> Optional[str]:
return None
status = backend.status()
if status.get("loaded"):
# repo_id names the checkpoint; for a GGUF / single-file load the companion base
# supplies the VAE and text encoders and is just as much part of the live model,
# so refuse it too (the Images guard above does the same via loaded_repo_ids).
# repo_id names the checkpoint; for a GGUF / single-file load the companion base supplies the VAE
# and text encoders and is just as much part of the live model, so refuse it too.
for key in ("repo_id", "base_repo"):
held = status.get(key)
if held and _loaded_id_matches_repo(str(held), repo_id):

5
studio/backend/hub/services/models/downloads.py
View file

@ -45,9 +45,8 @@ def _download_job_key(repo_id: str, variant: Optional[str]) -> str:
)
# A scope rides the variant slot as "@name". No GGUF quant label starts with "@", so a
# scoped job never collides with a real variant or with the repo's full snapshot, and its
# manifest/cancel marker stay in their own space.
# A scope rides the variant slot as "@name". No GGUF quant label starts with "@", so a scoped job
# never collides with a real variant or with the repo's full snapshot.
_SCOPE_PREFIX = "@"

7
studio/backend/hub/utils/download_registry.py
View file

@ -1155,10 +1155,9 @@ class DownloadRegistry:
return False, conflict_state
current = self._jobs.get(key, DownloadState("idle")).state
if current in _ACTIVE_STATES and not replace_active:
# A scope slot is shared by every file set that rides it (two quants of
# one repo both key as "@diffusion"), so adopting the live job would let
# the caller wait on files it never asked for. Reject instead; checked
# here, under the lock, so a concurrent claim cannot slip past it.
# A scope slot is shared by every file set that rides it (two quants of one repo both key as
# "@diffusion"), so adopting the live job would let the caller wait on files it never asked for.
# Reject instead; checked here, under the lock, so a concurrent claim cannot slip past.
live = self._metadata.get(key)
if (
scoped_files is not None

18
studio/backend/hub/workers/hf_download.py
View file

@ -699,10 +699,9 @@ def _download_scoped_snapshot(
blob_hashes: frozenset[str] = frozenset()
if info is not None:
siblings = [s for s in info.siblings if getattr(s, "rfilename", None) in wanted]
# Every requested file must resolve. Dropping an unmatched name silently would
# shrink the manifest and the completion check to the survivors, and
# snapshot_download also succeeds when an allow pattern matches nothing, so the
# job would report complete and trigger a load with a required file missing.
# Every requested file must resolve. Dropping an unmatched name would shrink the manifest and the
# completion check to the survivors, and snapshot_download also succeeds when an allow pattern
# matches nothing, so the job would report complete with a required file missing.
missing = sorted(set(wanted) - {getattr(s, "rfilename", None) for s in siblings})
if missing:
print(
@ -747,12 +746,11 @@ def _download_scoped_snapshot(
max_workers = 1,
)
if info is None:
# With no metadata there is no manifest, so _verify_completed_download below is a
# no-op -- and snapshot_download RETURNS AN EXISTING SNAPSHOT FOLDER, without
# fetching anything, when its own repo_info call also fails. A repo already on disk
# from a full snapshot job (which ignores *.gguf) would therefore flip this job to
# complete having downloaded no weights, and the page would load against them. The
# requested list needs no network, so check it against the disk directly.
# With no metadata there is no manifest, so _verify_completed_download below is a no-op -- and
# snapshot_download RETURNS AN EXISTING SNAPSHOT FOLDER, fetching nothing, when its own repo_info
# call also fails. A repo already on disk from a full snapshot job (which ignores *.gguf) would
# therefore flip this job to complete having downloaded no weights. The requested list needs no
# network, so check it against the disk directly.
root = Path(snapshot_path)
absent = tuple(f for f in files if not (root / f).exists())
if absent:

29
studio/backend/main.py
View file

@ -769,9 +769,8 @@ from utils.upload_limits import ( # noqa: E402
)
_BODY_PROTECTED_PREFIXES = (
# Blanket-protect the whole /v1 surface, like /api/inference below: every /v1 POST route
# (chat/completions, images/generations, audio, embeddings, responses, ...) buffers a JSON
# body and none is a multipart passthrough, so one prefix caps them all -- an enumerated
# Blanket-protect the whole /v1 surface, like /api/inference below: every /v1 POST route buffers
# a JSON body and none is a multipart passthrough, so one prefix caps them all -- an enumerated
# list would silently leave new routes uncapped.
"/v1",
"/p/",
@ -793,27 +792,24 @@ _DATA_RECIPE_UNSTRUCTURED_UPLOAD_PASSTHROUGH_PREFIX = (
# The diffusion dataset upload route (POST /api/train/diffusion/dataset) is a multipart image
# upload under the protected /api/train prefix. Like /api/datasets/upload it enforces its own
# get_upload_limit_bytes() cap, so it must bypass the default body cap here or the middleware
# would 413 near-limit batches (ignoring a raised max_upload_size_mb) before the handler runs.
# Matched as an EXACT path (not a prefix): its JSON sub-routes (.../caption/{filename},
# .../import-example) share the prefix but must keep the small-JSON cap, or a large
# caption/import body would be buffered up to the far larger upload limit.
# would 413 near-limit batches before the handler runs. Matched as an EXACT path: its JSON
# sub-routes share the prefix but must keep the small-JSON cap.
_DIFFUSION_DATASET_UPLOAD_PATH = "/api/train/diffusion/dataset"
_BODY_UPLOAD_PASSTHROUGH_PREFIXES = (
_DATASET_UPLOAD_PASSTHROUGH_PREFIX,
_DATA_RECIPE_UNSTRUCTURED_UPLOAD_PASSTHROUGH_PREFIX,
)
# Passthrough routes matched by EXACT path (the multipart upload only), so sibling JSON
# sub-routes under the same prefix are not swept into the generous upload cap.
# sub-routes under the same prefix keep the normal body cap.
_BODY_UPLOAD_PASSTHROUGH_EXACT_PATHS = (_DIFFUSION_DATASET_UPLOAD_PATH,)
def _get_upload_passthrough_request_max_bytes(path: str) -> int:
if path.startswith(_DATA_RECIPE_UNSTRUCTURED_UPLOAD_PASSTHROUGH_PREFIX):
return upload_request_limit_bytes(UNSTRUCTURED_RECIPE_UPLOAD_MAX_BYTES)
# The trailing-slash variant reaches this middleware BEFORE the router's redirect_slashes
# 307, so it must resolve to the same upload cap as the canonical path or a large upload
# 413s on the default /api/train cap. Stripping slashes can't promote a JSON sub-route:
# those keep extra path components after normalization and still miss the exact match.
# The trailing-slash variant reaches this middleware BEFORE the router's redirect_slashes 307, so
# it must resolve to the same upload cap or a large upload 413s on the default /api/train cap.
# Stripping slashes can't promote a JSON sub-route: those keep extra path components.
if (
path.startswith(_DATASET_UPLOAD_PASSTHROUGH_PREFIX)
or path.rstrip("/") == _DIFFUSION_DATASET_UPLOAD_PATH
@ -883,14 +879,13 @@ class MaxBodyMiddleware:
self.request_max_bytes_getter = request_max_bytes_getter
self.upload_passthrough_prefixes = upload_passthrough_prefixes
self.upload_passthrough_max_bytes_getter = upload_passthrough_max_bytes_getter
# Passthrough routes matched by exact path, not prefix: an upload route whose prefix
# also covers sibling JSON sub-routes that must keep the normal (small) body cap.
# Passthrough routes matched by exact path, not prefix: an upload route whose prefix also covers
# sibling JSON sub-routes that must keep the normal (small) body cap.
self.upload_passthrough_exact_paths = upload_passthrough_exact_paths
def _is_upload_passthrough(self, path: str) -> bool:
# Exact paths also match their trailing-slash variant: the middleware runs before the
# router's redirect_slashes 307, and a JSON sub-route can never normalize to the exact
# path (it keeps extra components).
# Exact paths also match their trailing-slash variant: the middleware runs before the router's
# redirect_slashes 307, and a JSON sub-route can never normalize to the exact path.
return path.rstrip("/") in self.upload_passthrough_exact_paths or any(
path.startswith(p) for p in self.upload_passthrough_prefixes
)

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

@ -2301,9 +2301,8 @@ class DiffusionLoadRequest(BaseModel):
@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.
# 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 form here.
return value.strip().lower() if isinstance(value, str) else value
@ -2360,8 +2359,8 @@ class ControlNetSpec(BaseModel):
@model_validator(mode = "after")
def _check_guidance_range(self) -> "ControlNetSpec":
# An inverted range (start > end) means "act over no steps"; reject it as a clean 422
# instead of letting the diffusers pipeline 500 deep in the denoise.
# An inverted range means "act over no steps"; reject it as a clean 422 instead of letting the
# diffusers pipeline 500 deep in the denoise.
if self.guidance_start > self.guidance_end:
raise ValueError("guidance_start must be <= guidance_end")
return self
@ -2380,18 +2379,17 @@ 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.
# le = 2**53-1: seeds round-trip through JSON gallery recipes, where JavaScript rounds integers
# above Number.MAX_SAFE_INTEGER and a restored recipe would generate a different image.
seed: Optional[int] = Field(
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)"
)
# Batched multi-image generation (diffusers engine): a prompt list renders one image per
# prompt in a single batched forward (txt2img only); a seed list renders one image per
# seed. Each image carries its OWN generator seed, so any batch member replays alone.
# Batched multi-image generation (diffusers engine): a prompt list renders one image per prompt
# in a single batched forward (txt2img only); a seed list renders one image per seed. Each image
# carries its OWN generator seed, so any batch member replays alone.
prompts: Optional[list[str]] = Field(
None,
min_length = 1,
@ -2419,8 +2417,7 @@ class DiffusionGenerateRequest(BaseModel):
@field_validator("seeds")
@classmethod
def _seeds_json_safe(cls, value: Optional[list[int]]) -> Optional[list[int]]:
# Same JSON safe-integer bound as `seed`, so every per-image seed survives the
# round-trip through the gallery recipe.
# Same JSON safe-integer bound as `seed`, so every per-image seed survives the gallery recipe.
if value is not None and any(s < 0 or s > 2**53 - 1 for s in value):
raise ValueError("every seed must be between 0 and 2**53 - 1")
return value
@ -2438,10 +2435,10 @@ class DiffusionGenerateRequest(BaseModel):
)
return self
# Image-conditioned workflows (base64 or data-URL): init_image alone runs img2img,
# init_image + mask_image runs inpaint. Both require a family with the matching pipeline or
# the load is rejected. Cap each base64 string so one request can't buffer a multi-GB payload
# (decoded dimensions are bounded separately); ~32 MiB fits a full 4096px image yet rejects abuse.
# Image-conditioned workflows (base64 or data-URL): init_image alone runs img2img, init_image +
# mask_image runs inpaint. Both require a family with the matching pipeline. Cap each base64
# string so one request can't buffer a multi-GB payload (decoded dimensions are bounded
# separately); ~32 MiB fits a full 4096px image yet rejects abuse.
init_image: Optional[str] = Field(
None,
max_length = 32 * 1024 * 1024,
@ -2455,12 +2452,9 @@ class DiffusionGenerateRequest(BaseModel):
)
strength: Optional[float] = Field(
None,
# EXCLUSIVE lower bound: strength 0 does not "keep the source". Every diffusers
# img2img/inpaint pipeline derives its step count from it (t_start =
# num_inference_steps - int(num_inference_steps * strength)), so 0 leaves zero
# denoising steps: FLUX/Qwen/Z-Image raise "the number of pipeline steps is 0 which
# is < 1", and SDXL img2img has no such guard and crashes on empty latents (a 500).
# The UI slider already starts at 0.1; reject 0 as a 422 instead of a pipeline error.
# EXCLUSIVE lower bound: strength 0 does not "keep the source". Every diffusers img2img/inpaint
# pipeline derives its step count from it, so 0 leaves zero denoising steps: FLUX/Qwen/Z-Image
# raise "the number of pipeline steps is 0", and SDXL img2img crashes on empty latents (a 500).
gt = 0.0,
le = 1.0,
description = "img2img/inpaint denoise strength: low values stay close to the "
@ -2497,10 +2491,9 @@ class DiffusionGenerateRequest(BaseModel):
@field_validator("loras")
@classmethod
def _unique_lora_ids(cls, value: Optional[list[LoraSpec]]) -> Optional[list[LoraSpec]]:
# Both apply paths break alias collisions by suffixing the adapter name/file, so a
# repeated id would load the SAME adapter several times and stack its effect past the
# per-adapter weight bound. The UI already blocks duplicates; reject them for API clients
# too so each adapter takes effect at most once.
# Both apply paths break alias collisions by suffixing the adapter name/file, so a repeated id
# would load the SAME adapter several times and stack its effect past the per-adapter weight
# bound. The UI already blocks duplicates; reject them for API clients too.
if value:
seen: set[str] = set()
for spec in value:
@ -2514,8 +2507,8 @@ class DiffusionGenerateRequest(BaseModel):
@field_validator("reference_images")
@classmethod
def _bounded_reference_items(cls, value: Optional[list[str]]) -> Optional[list[str]]:
# Each reference is a base64 image; bound its length like init_image/mask_image so
# several references can't buffer a multi-GB payload.
# Each reference is a base64 image; bound its length like init_image/mask_image so several
# references can't buffer a multi-GB payload.
if value is not None:
for item in value:
if len(item) > 32 * 1024 * 1024:
@ -2525,18 +2518,17 @@ class DiffusionGenerateRequest(BaseModel):
@field_validator("width", "height")
@classmethod
def _multiple_of_16(cls, value: int) -> int:
# Z-Image requires dimensions divisible by 16 (8x VAE downsample + 2x patch).
# Non-multiples crash deep in the pipeline, so reject them here for a clean 422.
# Z-Image requires dimensions divisible by 16 (8x VAE downsample + 2x patch). Non-multiples crash
# deep in the pipeline, so reject them here for a clean 422.
if value % 16 != 0:
raise ValueError("must be a multiple of 16")
return value
@model_validator(mode = "after")
def _batch_seeds_json_safe(self) -> "DiffusionGenerateRequest":
# A batch derives per-image seeds as seed .. seed+batch_size-1. The base seed is capped at
# 2**53-1 to round-trip through the JSON recipe, but a derived top-of-batch seed near the cap
# can exceed it, where the frontend rounds it and a restored recipe replays a different
# image. Reject at the boundary so an API client can't persist an unreplayable seed.
# A batch derives per-image seeds as seed .. seed+batch_size-1. The base seed is capped at 2**53-1
# to round-trip through the JSON recipe, but a derived top-of-batch seed near the cap can exceed
# it, where the frontend rounds it and a restored recipe replays a different image.
if self.seed is not None and self.seed + self.batch_size - 1 > 2**53 - 1:
raise ValueError(
"seed + batch_size - 1 must not exceed 2**53 - 1 so every per-image seed "
@ -2721,10 +2713,10 @@ class DiffusionStatusResponse(BaseModel):
"picker's enabled state). Diffusers only, for families with a ControlNet pipeline; False "
"for the native engine, GGUF-via-diffusers, and torchao fp8/int8 dense.",
)
# Additive: per-Advanced-control provenance {control: {value, source, reason}}. Present only
# on backends that record it; null when nothing is loaded or on older backends. The frontend
# renders an "Auto: X" badge next to each control whose source == "auto". Declared explicitly
# so pydantic's extra='ignore' doesn't drop the resolved record.
# Additive: per-Advanced-control provenance {control: {value, source, reason}}. Present only on
# backends that record it; null when nothing is loaded. The frontend renders an "Auto: X" badge
# next to each control whose source == "auto". Declared explicitly so pydantic's extra='ignore'
# doesn't drop it.
resolved: Optional[Dict[str, DiffusionResolvedControl]] = Field(
None,
description = "Per-control resolved value + provenance (source auto|explicit + reason), "
@ -2762,10 +2754,9 @@ class DiffusionInferenceInfoResponse(BaseModel):
# ── OpenAI-compatible images API (POST /v1/images/generations) ──
#
# Shapes mirror OpenAI's CreateImageRequest / ImagesResponse so off-the-shelf clients work
# unchanged. The loaded image GGUF stands in for the model; GPT-image-only knobs (quality,
# style, background, output_format, ...) are accepted and ignored, like dall-e-2. The size
# string is parsed and `stream` is rejected in the route (where the diffusion backend is in
# reach); everything Pydantic can check declaratively lives here.
# unchanged. The loaded image GGUF stands in for the model; GPT-image-only knobs (quality, style,
# background, output_format, ...) are accepted and ignored, like dall-e-2. The size string is
# parsed and `stream` rejected in the route; everything Pydantic can check declaratively is here.
class ImageGenerationRequest(BaseModel):
@ -2787,8 +2778,8 @@ class ImageGenerationRequest(BaseModel):
"url", description = "Return each image as a URL or a base64-encoded PNG."
)
user: Optional[str] = Field(None, description = "End-user identifier (accepted, unused).")
# gpt-image-only; declared so we can reject it clearly instead of returning JSON to a
# client that asked for an SSE stream.
# gpt-image-only; declared so we can reject it clearly instead of returning JSON to a client that
# asked for an SSE stream.
stream: Optional[bool] = Field(
None, description = "Streaming image generation is not supported; omit or set false."
)
@ -2796,8 +2787,8 @@ class ImageGenerationRequest(BaseModel):
@field_validator("n", "size", "response_format", mode = "before")
@classmethod
def _null_means_default(cls, value, info):
# OpenAI marks these nullable WITH a default, so an explicit null means "use the
# default" -- coalesce it instead of 400-ing a spec-valid body.
# OpenAI marks these nullable WITH a default, so an explicit null means "use the default":
# coalesce it instead of 400-ing a spec-valid body.
if value is None:
return cls.model_fields[info.field_name].default
return value
@ -2929,9 +2920,8 @@ class VideoLoadRequest(BaseModel):
@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.
# 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 form here.
return value.strip().lower() if isinstance(value, str) else value
@ -2942,8 +2932,8 @@ class VideoGenerateRequest(BaseModel):
negative_prompt: Optional[str] = Field(
None, description = "What to avoid (if the model supports it)"
)
# Width/height/num_frames/fps default per loaded family (the backend snaps them to its
# required multiples/lattice), so they are optional here.
# Width/height/num_frames/fps default per loaded family (the backend snaps them to its required
# multiples/lattice), so they are optional here.
width: Optional[int] = Field(
None, ge = 32, le = 2048, description = "Frame width in pixels (family multiple)"
)
@ -2974,9 +2964,8 @@ class VideoGenerateRequest(BaseModel):
"pipeline default it to the main guidance. Ignored by single-DiT families (their pipeline "
"signature has no second guidance kwarg).",
)
# 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 clip. Random seeds are already masked to this range.
# le = 2**53-1: seeds round-trip through JSON gallery recipes, where JavaScript rounds integers
# above Number.MAX_SAFE_INTEGER and a restored recipe would generate a different clip.
seed: Optional[int] = Field(
None, ge = 0, le = 2**53 - 1, description = "Seed for reproducibility (random if omitted)"
)
@ -3127,9 +3116,9 @@ class VideoStatusResponse(BaseModel):
defaults: Optional[VideoGenerationDefaults] = Field(
None, description = "Per-family generation defaults + shape constraints; null when unloaded"
)
# Additive: per-Advanced-control provenance {control: {value, source, reason}}. Same shape
# as the diffusion status; null when nothing is loaded. The frontend renders an "Auto: X"
# badge next to each control whose source == "auto".
# Additive: per-Advanced-control provenance {control: {value, source, reason}}. Same shape as the
# diffusion status; null when nothing is loaded. The frontend renders an "Auto: X" badge next to
# each control whose source == "auto".
resolved: Optional[Dict[str, DiffusionResolvedControl]] = Field(
None,
description = "Per-control resolved value + provenance (source auto|explicit + reason), "

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

@ -4050,9 +4050,8 @@ def _guard_chat_load_against_training(
return
if not llm_active:
# An SDXL LoRA trainer runs in its own subprocess and its VRAM can't be cheaply
# fit-checked here, so refuse the chat load while one is active rather than risk
# OOMing the run. Symmetric with the image-load guard.
# An SDXL LoRA trainer runs in its own subprocess and its VRAM can't be cheaply fit-checked here,
# so refuse the chat load while one is active. Symmetric with the image-load guard.
if _diffusion_training_active():
raise HTTPException(
status_code = 409,
@ -4389,12 +4388,10 @@ async def _load_model_impl(
user_override = request.chat_template_override,
)
# Reclaim the GPU for chat (evicting a resident Images/Video pipeline) only once the
# load is known viable: the already-loaded fast paths below re-assert CHAT ownership
# themselves, and the real handoff is deferred past identifier / gpu_ids / training-memory
# validation so a doomed load (bad id, unsupported gpu_ids on GGUF, training 409) can't
# evict a working image/video model and then error. Mirrors the image/video loaders,
# which validate before acquire_for.
# Reclaim the GPU for chat (evicting a resident Images/Video pipeline) only once the load is known
# viable: the already-loaded fast paths below re-assert CHAT ownership themselves, and the real
# handoff is deferred past identifier / gpu_ids / training-memory validation so a doomed load
# can't evict a working image/video model and then error. Mirrors the image/video loaders.
from core.inference.gpu_arbiter import acquire_for, current_owner, release, CHAT
# ── Already-loaded check: skip reload if the exact model is active ──
@ -4431,11 +4428,10 @@ async def _load_model_impl(
_gguf_audio = getattr(llama_backend, "_audio_type", None)
_gguf_is_audio = getattr(llama_backend, "_is_audio", False)
# Requested GGUF chat model already resident: assert CHAT ownership (no-op when
# held) to correct a drifted arbiter owner. Guaranteed-success path, so evicting
# here is correct -- unless the resident server is a confirmed zero-VRAM one
# (manual zero offload, GPUs hidden from the child), which coexists with an
# image/video pipeline and so must not evict it to re-announce itself.
# Requested GGUF chat model already resident: assert CHAT ownership (no-op when held) to correct
# a drifted arbiter owner. A guaranteed-success path, so evicting here is correct -- unless the
# resident server is a confirmed zero-VRAM one, which coexists with an image/video pipeline and
# so must not evict it to re-announce itself.
if not llama_backend.holds_no_vram:
await asyncio.to_thread(acquire_for, CHAT)
return LoadResponse(
@ -4498,9 +4494,8 @@ async def _load_model_impl(
_sf_flags = _detect_safetensors_features(backend, _chat_template)
_sf_supports_reasoning = _sf_flags["supports_reasoning"]
_sf_reasoning_style = _sf_flags["reasoning_style"]
# Requested chat model already resident: assert CHAT ownership (no-op when held)
# to correct a drifted arbiter owner. Guaranteed-success path, so evicting here
# is correct.
# Requested chat model already resident: assert CHAT ownership (no-op when held) to correct a
# drifted arbiter owner. A guaranteed-success path, so evicting here is correct.
await asyncio.to_thread(acquire_for, CHAT)
return LoadResponse(
status = "already_loaded",
@ -4609,12 +4604,11 @@ async def _load_model_impl(
gpu_memory_mode = request.gpu_memory_mode,
)
# Mark the load and refuse one the download manager already owns, BEFORE the eviction
# below: this 409 leaves nothing loaded, so checking it afterwards destroyed a working
# Images/Video pipeline for a load that could never start. The marker/check order is the
# handshake with the download manager, so both move together. It also has to run after
# pass-through argument inheritance, since a carried --no-mmproj changes the companion
# requirement exactly as it does for the load.
# Mark the load and refuse one the download manager already owns, BEFORE the eviction below: this
# 409 leaves nothing loaded, so checking it afterwards destroyed a working Images/Video pipeline
# for a load that could never start. The marker/check order is the handshake with the download
# manager. It also runs after pass-through argument inheritance, since a carried --no-mmproj
# changes the companion requirement exactly as it does for the load.
if config.is_gguf and config.gguf_hf_repo:
from core.inference.llama_cpp import gguf_load_in_flight
@ -4640,21 +4634,18 @@ async def _load_model_impl(
),
)
# Load now known viable (valid identifier, gpu_ids ok, fits alongside any active
# training): reclaim the GPU for chat, evicting a resident Images/Video pipeline. Doing
# this only here -- not before the validation above -- keeps a doomed load from evicting
# a working image/video model and then erroring. No-op when chat already owns the GPU.
# The in-flight marker is entered UNDER the arbiter lock (the `register` hook), as the
# image and video loads do: a chat load holds no llama-server process until its GGUF has
# downloaded, so a competing Images/Video acquire in that window found nothing to evict
# and both then allocated VRAM at once. With the marker, that evictor cancels this load.
# Load now known viable (valid identifier, gpu_ids ok, fits alongside any active training):
# reclaim the GPU for chat, evicting a resident Images/Video pipeline. Doing this only here keeps
# a doomed load from evicting a working model and then erroring. No-op when chat already owns the
# GPU. The in-flight marker is entered UNDER the arbiter lock, as the image and video loads do: a
# chat load holds no llama-server process until its GGUF downloaded, so a competing acquire in
# that window found nothing to evict and both allocated VRAM at once.
from core.inference.llama_cpp import chat_load_in_flight, zero_vram_chat_load
# ...but only when this load will actually use the GPU, exactly as the image and video
# loaders gate on their resolved device. A manual gpu_layers=0 GGUF load runs on the CPU
# with the GPUs hidden from the child, so taking the arbiter for it would cancel a
# running image/video generation for a model that needs no VRAM, and leave CHAT recorded
# as owner so the next GPU workload pointlessly unloads it.
# ...but only when this load will actually use the GPU, exactly as the image and video loaders
# gate on their resolved device. A manual gpu_layers=0 GGUF load runs on the CPU with the GPUs
# hidden from the child, so taking the arbiter would cancel a running image/video generation for
# a model that needs no VRAM, and leave CHAT recorded as owner.
chat_load_needs_gpu = not (
config.is_gguf
and await asyncio.to_thread(
@ -4673,10 +4664,10 @@ async def _load_model_impl(
lambda: gguf_load_stack.enter_context(chat_load_in_flight()),
)
else:
# The marker still goes up (the download manager's handshake reads it, and it keeps
# this load cancellable). Any stale CHAT claim is dropped AFTER the load, not here:
# this load may still be replacing a GPU-backed chat model, and releasing up front
# would let an image/video load allocate alongside the model not yet unloaded.
# The marker still goes up (the download manager's handshake reads it, and it keeps this load
# cancellable). Any stale CHAT claim is dropped AFTER the load, not here: this load may still be
# replacing a GPU-backed chat model, and releasing up front would let an image/video load
# allocate alongside the model not yet unloaded.
gguf_load_stack.enter_context(chat_load_in_flight())
# ── GGUF path: load via llama-server ──────────────────────
@ -4839,11 +4830,10 @@ async def _load_model_impl(
detail = f"Failed to load GGUF model: {model_log_label if native_grant_backed else config.display_name}",
)
# An Images/Video acquire can land in the gap between the acquire above and
# load_model clearing the cancel event, so its cancellation is lost and this load
# spawns anyway. Ownership survives that gap: whoever took the GPU keeps it, and this
# load undoes itself rather than leaving two models resident on one device. A
# zero-VRAM load never took ownership, so it has nothing to lose and never yields.
# An Images/Video acquire can land in the gap between the acquire above and load_model clearing
# the cancel event, so its cancellation is lost and this load spawns anyway. Ownership survives
# that gap: whoever took the GPU keeps it, and this load undoes itself rather than leaving two
# models resident on one device. A zero-VRAM load never took ownership, so it never yields.
if chat_load_needs_gpu and current_owner() != CHAT:
await asyncio.to_thread(llama_backend.unload_model)
raise HTTPException(
@ -4854,11 +4844,9 @@ async def _load_model_impl(
),
)
if not chat_load_needs_gpu:
# Zero-VRAM load done, and whatever GPU-backed chat model it replaced went with
# it, so drop a now-stale CHAT claim: leaving it would make the next image/video
# load "evict" a server holding nothing. Owner-guarded, so it no-ops when an
# image/video model took the GPU while this was loading -- which is fine, they
# coexist.
# Zero-VRAM load done, and whatever GPU-backed chat model it replaced went with it, so drop a
# now-stale CHAT claim: leaving it would make the next image/video load "evict" a server holding
# nothing. Owner-guarded, so it no-ops when an image/video model took the GPU meanwhile.
await asyncio.to_thread(release, CHAT)
logger.info(
@ -4979,14 +4967,13 @@ async def _load_model_impl(
detail = f"Failed to load model: {model_log_label if native_grant_backed else config.display_name}",
)
# Same guard the GGUF branch runs above: an Images/Video acquire can land in the gap
# between this load's cancellation and its publish, so the eviction is lost and the
# model lands anyway. Ownership survives that gap, so this load undoes itself rather
# than leaving two models resident on one device.
# Same guard the GGUF branch runs above: an Images/Video acquire can land in the gap between this
# load's cancellation and its publish, so the eviction is lost and the model lands anyway.
# Ownership survives that gap, so this load undoes itself rather than leaving two models resident.
if current_owner() != CHAT:
await asyncio.to_thread(backend.unload_model, config.identifier)
# The worker's base CUDA context outlives the model unload, so kill it too --
# that VRAM is exactly what the image/video pipeline just took the GPU for.
# The worker's base CUDA context outlives the model unload, so kill it too -- that VRAM is
# exactly what the image/video pipeline just took the GPU for.
await asyncio.to_thread(backend._shutdown_subprocess, 5.0)
raise HTTPException(
status_code = 409,
@ -15993,9 +15980,9 @@ async def _openai_passthrough_non_streaming_upstream(
# ──────────────────────────────────────────────────────────────────────────
# Diffusion (local text-to-image)
#
# Studio-only routes (studio_router is not mounted under /v1). The diffusion backend
# runs in-process and synchronously, so blocking load/generate/unload calls are
# offloaded with asyncio.to_thread. Single error boundary: backend raises, we map to HTTP.
# Studio-only routes (studio_router is not mounted under /v1). The diffusion backend runs
# in-process and synchronously, so blocking load/generate/unload calls are offloaded with
# asyncio.to_thread. Single error boundary: the backend raises, we map to HTTP.
# ──────────────────────────────────────────────────────────────────────────
@ -16021,9 +16008,9 @@ def _guard_diffusion_load_against_training() -> None:
except Exception as e:
logger.warning("Could not check training state for image-load guard: %s", e)
return
# An SDXL LoRA trainer runs in its own subprocess on the same GPU, so an image load
# must be refused while one is active too, or the pipeline contends with the trainer
# for VRAM. Symmetric with the diffusion-start interlock.
# An SDXL LoRA trainer runs in its own subprocess on the same GPU, so an image load must be
# refused while one is active or the pipeline contends with the trainer for VRAM. Symmetric with
# the diffusion-start interlock.
if not llm_active and not _diffusion_training_active():
return
raise HTTPException(
@ -16055,8 +16042,8 @@ async def diffusion_download_plan(
backend = get_diffusion_backend()
try:
kind = resolve_model_kind(request.gguf_filename, request.model_kind)
# Same bare-single-file-directory reinterpretation as the load route, so the plan
# describes the load that will actually run.
# Same bare-single-file-directory reinterpretation as the load route, so the plan describes the
# load that will actually run.
if kind == "pipeline" and not request.gguf_filename:
sole = await asyncio.to_thread(resolve_local_single_file, request.model_path)
if sole is not None:
@ -16080,11 +16067,10 @@ async def diffusion_download_plan(
hf_token = request.hf_token,
transformer_quant = request.transformer_quant,
speed_mode = request.speed_mode,
# The dense-quant prefetch decision reads the memory policy, the prequant path
# and the adapter selection too (an offload policy never runs the dense build;
# a baked LoRA always does), so the plan has to see the same values the load
# will. Without them it stages the base transformer/ shards for a low-VRAM
# load that never opens them, or omits them for a baked-LoRA load that does.
# The dense-quant prefetch decision reads the memory policy, the prequant path and the adapter
# selection too (an offload policy never runs the dense build; a baked LoRA always does), so the
# plan has to see the same values the load will. Without them it stages the base transformer/
# shards for a low-VRAM load that never opens them, or omits them for a baked-LoRA load that does.
memory_mode = request.memory_mode,
cpu_offload = request.cpu_offload,
transformer_prequant_path = request.transformer_prequant_path,
@ -16115,22 +16101,21 @@ async def load_diffusion_model(
backend = get_diffusion_backend()
try:
# Resolve the load kind once (gguf / single_file / pipeline) so validation,
# engine selection, and the load all agree. A bad explicit kind raises here -> 400.
# Resolve the load kind once (gguf / single_file / pipeline) so validation, engine selection and
# the load all agree. A bad explicit kind raises here, so a 400.
kind = resolve_model_kind(request.gguf_filename, request.model_kind)
# A local On-Device pick can be a bare single-file .safetensors directory (no
# model_index.json): the scanner advertises it as text-to-image, but the picker starts
# it as a pipeline with no filename, so a pipeline load would 400 on the missing
# model_index.json. If the directory holds exactly one checkpoint, reinterpret the pick
# as a single_file load of it (its only loadable shape), so all three paths agree.
# A local On-Device pick can be a bare single-file .safetensors directory (no model_index.json):
# the scanner advertises it as text-to-image, but the picker starts it as a pipeline with no
# filename, so a pipeline load would 400. If the directory holds exactly one checkpoint,
# reinterpret the pick as a single_file load of it, so all three paths agree.
if kind == "pipeline" and not request.gguf_filename:
sole = await asyncio.to_thread(resolve_local_single_file, request.model_path)
if sole is not None:
request.gguf_filename = sole
kind = resolve_model_kind(sole)
# Validate cheaply BEFORE touching the GPU: an unloadable pick (bad family, missing
# local GGUF, non-unsloth non-GGUF repo) must not evict a working chat model and then
# 400. The validated family also drives engine selection below.
# Validate cheaply BEFORE touching the GPU: an unloadable pick (bad family, missing local GGUF,
# non-unsloth non-GGUF repo) must not evict a working chat model and then 400. The validated
# family also drives engine selection below.
fam = await asyncio.to_thread(
backend.validate_load_request,
request.model_path,
@ -16139,26 +16124,24 @@ async def load_diffusion_model(
model_kind = kind,
base_repo = request.base_repo,
)
# Refuse while training is running: a multi-GB diffusion pipeline would
# compete with the training subprocess for VRAM. The chat path does the
# same via _guard_chat_load_against_training; this is its image sibling.
# Refuse while training is running: a multi-GB diffusion pipeline would compete with the training
# subprocess for VRAM. The image sibling of _guard_chat_load_against_training.
_guard_diffusion_load_against_training()
# Pick the engine for this host (diffusers on GPU, native sd.cpp with no GPU),
# installing the sd-cli binary if needed -- all BEFORE evicting chat, so a native
# fallback never strands a half-loaded state. Non-GGUF kinds force diffusers.
# Pick the engine for this host (diffusers on GPU, native sd.cpp with no GPU), installing the
# sd-cli binary if needed, all BEFORE evicting chat so a native fallback never strands a
# half-loaded state. Non-GGUF kinds force diffusers.
engine = await asyncio.to_thread(
select_and_activate_engine, fam, hf_token = request.hf_token, model_kind = kind
)
# Take the GPU from chat only when this load will actually use it, i.e. the resolved
# device is non-CPU. diffusers on an accelerator and a force-native sd.cpp load on
# CUDA/XPU/MPS both resolve to a device; a native sd.cpp load on a pure-CPU host, and a
# CPU-only host falling back to diffusers ON CPU, do not. So gate on the device, not the
# engine name -- else we'd evict a resident chat model for a load that can't use the GPU.
# Take the GPU from chat only when this load will actually use it, i.e. the resolved device is
# non-CPU. diffusers on an accelerator and a force-native sd.cpp load on CUDA/XPU/MPS both
# resolve to a device; a native load on a pure-CPU host, and a CPU-only host falling back to
# diffusers, do not. Gate on the device, not the engine name.
device = await asyncio.to_thread(lambda: resolve_diffusion_device_target().device)
needs_gpu = device != "cpu"
def _start_engine_load():
# Kicks the (slow) load onto a background thread and returns at once (client polls
# Kicks the (slow) load onto a background thread and returns at once (the client polls
# images/load-progress); begin_load itself validates network-free.
return engine.begin_load(
request.model_path,
@ -16181,22 +16164,21 @@ async def load_diffusion_model(
)
def _begin_load():
# Under the router's transition lock, refusing if a competing load switched engines
# since select_and_activate_engine above: begin_load on a deactivated engine leaves a
# resident model that generate / status / unload and the evictor can no longer reach.
# Under the router's transition lock, refusing if a competing load switched engines since
# select_and_activate_engine above: begin_load on a deactivated engine leaves a resident model
# that generate / status / unload and the evictor can no longer reach.
return begin_load_on(engine, _start_engine_load)
if needs_gpu:
# Register the in-flight load UNDER the arbiter lock (not after acquire_for returns):
# otherwise a competing Video/chat acquire in that gap evicts DIFFUSION before the load
# is marked in-flight, finds nothing to cancel, and both loaders allocate VRAM at once.
# Register the in-flight load UNDER the arbiter lock (not after acquire_for returns): otherwise a
# competing Video/chat acquire in that gap evicts DIFFUSION before the load is marked in-flight,
# finds nothing to cancel, and both loaders allocate VRAM at once.
status_dict = await asyncio.to_thread(acquire_for, DIFFUSION, _begin_load)
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
# leaves DIFFUSION still marked as arbiter owner, so a later chat acquire would
# "evict" this CPU model for no reason. Release that stale ownership -- release()
# is owner-guarded, so it's a no-op when diffusion never owned the GPU.
# 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 leaves DIFFUSION still marked as
# owner, so a later chat acquire would "evict" this CPU model for no reason. release() is
# owner-guarded, so it no-ops when diffusion never owned the GPU.
await asyncio.to_thread(release, DIFFUSION)
status_dict = await asyncio.to_thread(_begin_load)
return DiffusionStatusResponse(**annotate_status(status_dict))
@ -16208,9 +16190,8 @@ async def load_diffusion_model(
# Count of finished generations still writing their PNG/gallery records. generate-progress reports
# active while this is > 0 so a reload's mount probe never reads idle between the denoise finishing
# (engine drops _gen) and the image reaching the gallery, which would refresh the gallery before the
# record exists. Mutated only on the event loop (around the persist await), so no lock is needed.
# active while this is above 0, so a reload's mount probe never reads idle between the denoise
# finishing and the image reaching the gallery. Mutated only on the event loop, so no lock.
_diffusion_persist_active = 0
@ -16259,15 +16240,14 @@ async def generate_diffusion_image(
),
)
except ValueError as exc:
# Bad client input (undecodable image/mask, or a workflow the loaded family
# doesn't support) — a 400 with the reason, not a generic 500.
# Bad client input (undecodable image/mask, or a workflow the loaded family doesn't support): a
# 400 with the reason, not a generic 500.
raise HTTPException(status_code = 400, detail = str(exc))
except RuntimeError as exc:
# Only "no model loaded" / user-cancelled are client-state (409); both engines raise
# these two EXACT messages. The native sd.cpp engine also raises RuntimeError for
# execution failures whose text can embed the raw sd-cli tail (local paths / argv) --
# those are 500s returned as a fixed literal, never echoed. Match the sentinels exactly
# (not as substrings) so an sd-cli failure containing "cancelled" can't misroute to 409.
# Only "no model loaded" / user-cancelled are client-state (409); both engines raise these two
# EXACT messages. The native sd.cpp engine also raises RuntimeError for execution failures whose
# text can embed the raw sd-cli tail (local paths / argv), which are 500s returned as a fixed
# literal. Match the sentinels exactly so an sd-cli failure containing "cancelled" can't 409.
msg = str(exc)
if msg in (DIFFUSION_NOT_LOADED_MSG, DIFFUSION_CANCELLED_MSG):
raise HTTPException(status_code = 409, detail = msg)
@ -16277,16 +16257,14 @@ async def generate_diffusion_image(
logger.error("diffusion.generate_failed: %s", exc, exc_info = True)
raise HTTPException(status_code = 500, detail = "Image generation failed.")
# Persist each image with its full recipe embedded. BOTH engines batch with a distinct
# seed per image (diffusers via one torch.Generator per image, native sd.cpp via
# base + index), returned in ``seeds`` so each image is individually reproducible.
# Persist each image with its full recipe embedded. BOTH engines batch with a distinct seed per
# image (diffusers via one torch.Generator per image, native sd.cpp via base + index), returned
# in ``seeds`` so each image is individually reproducible.
created_at = time.time()
per_image_seeds = result.get("seeds")
# A prompts/seeds LIST drives the image count and each image's own seed, so
# ``batch_size`` is only a per-forward cap there and the base seed no longer replays
# image i (seeds=[5, 99] would restore 5 for the 99 image). Persist those outputs as
# single-image recipes keyed on their OWN seed instead, so the gallery's Restore
# reproduces every image and not just the first.
# A prompts/seeds LIST drives the image count and each image's own seed, so ``batch_size`` is only
# a per-forward cap there and the base seed no longer replays image i. Persist those outputs as
# single-image recipes keyed on their OWN seed, so Restore reproduces every image.
list_driven = bool(request.prompts or request.seeds)
def _persist() -> list[dict]:
@ -16301,34 +16279,29 @@ async def generate_diffusion_image(
image_gallery.save(
image,
{
# A prompts-list batch records each image's OWN prompt so its recipe
# replays exactly; otherwise every image shares the single prompt.
# A prompts-list batch records each image's OWN prompt so its recipe replays exactly.
"prompt": (
request.prompts[index]
if request.prompts and index < len(request.prompts)
else request.prompt
),
"negative_prompt": request.negative_prompt,
# Persist the ACTUAL output size, not the request sliders:
# Transform/Inpaint/Edit derive it from the uploaded image, Extend grows
# the canvas, Upscale resizes it, so request.width/height would record the
# wrong dims. For plain txt2img the size equals the sliders anyway.
# Persist the ACTUAL output size, not the request sliders: Transform/Inpaint/Edit derive it from
# the uploaded image, Extend grows the canvas and Upscale resizes it, so request.width/height
# would record the wrong dims.
"width": getattr(image, "width", None) or request.width,
"height": getattr(image, "height", None) or request.height,
"steps": request.steps,
"guidance": request.guidance,
"seed": seed,
# Base seed the batch launched with. The native engine derives per-image
# seeds as base + index, so ``seed`` above is already advanced for index>0;
# restore replays from this base (diffusers shares one seed, so base == seed).
# A list-driven image carries its OWN seed instead (see ``list_driven``).
# Base seed the batch launched with. The native engine derives per-image seeds as base + index,
# so ``seed`` above is already advanced for index>0 and restore replays from this base (diffusers
# shares one seed, so base == seed). A list-driven image carries its OWN seed instead.
"batch_seed": seed if list_driven else result["seed"],
# Position within the batch (shared timestamp), so the export filename
# stays unique.
# Position within the batch (shared timestamp), so the export filename stays unique.
"batch_index": index,
# The batch shares one seed, so reproducing a batch_index>0 image needs
# the original batch_size: persist it so restore can replay. A list-driven
# image needs no replay -- it restores as a single image on its own seed.
# The batch shares one seed, so reproducing a batch_index>0 image needs the original batch_size:
# persist it so restore can replay. A list-driven image restores on its own seed instead.
"batch_size": 1 if list_driven else request.batch_size,
"model": result.get("repo_id"),
"loras": (
@ -16337,8 +16310,8 @@ async def generate_diffusion_image(
"controlnet": (
f"{request.controlnet.id}:{request.controlnet.control_type}:"
f"{request.controlnet.strength:g}"
# strength 0 is disabled and skipped before loading/conditioning,
# so don't claim a ControlNet was applied in the recipe/metadata.
# strength 0 is disabled and skipped before loading/conditioning, so don't claim a ControlNet was
# applied in the recipe/metadata.
if request.controlnet and request.controlnet.strength > 0
else None
),
@ -16350,7 +16323,7 @@ async def generate_diffusion_image(
# Hold generate-progress "active" across the persist so a concurrent reload's mount probe can't
# see idle and refresh the gallery before these records exist. Set synchronously right after the
# engine returned (no await between, so no idle gap), cleared in the finally.
# engine returned (no await between), cleared in the finally.
global _diffusion_persist_active
_diffusion_persist_active += 1
try:
@ -16379,8 +16352,8 @@ async def list_gallery_images(
# Validate inside the pager so offset / limit / has_more all count over the accepted domain. A
# recipe with all keys but a wrong value type passes the presence-only read yet fails
# GalleryImage(**r); dropping it only after slicing let a leading bad record return an empty
# page with has_more=True, stalling infinite scroll at offset 0.
# GalleryImage(**r); dropping it only after slicing let a leading bad record return an empty page
# with has_more=True, stalling infinite scroll at offset 0.
def _valid_gallery_image(record: dict) -> bool:
try:
GalleryImage(**record)
@ -16404,7 +16377,7 @@ async def get_gallery_image_file(
from core.inference import image_gallery
# Ownership-gate the serve like delete/clear: resolve only a Studio-owned PNG (readable recipe),
# so a guessed stem for a hand-dropped foreign PNG the listing hides can't be streamed out.
# so a guessed stem for a hand-dropped foreign PNG can't be streamed out.
path = await asyncio.to_thread(image_gallery.owned_image_path, image_id)
if path is None:
raise HTTPException(status_code = 404, detail = "Image not found.")
@ -16440,12 +16413,11 @@ async def unload_diffusion_model(current_subject: str = Depends(get_current_subj
from core.inference.gpu_arbiter import release_if, DIFFUSION
status_dict = await asyncio.to_thread(get_active_diffusion_engine().unload)
# Drop DIFFUSION ownership only if nothing is resident AND no new load is in flight: a
# concurrent /images/load that re-acquired DIFFUSION while this (slow) unload ran must keep
# ownership, or a later chat load sees no owner, skips eviction, and OOMs the newly resident
# pipeline. An in-flight load has is_loaded False for its whole window, so gate on
# loading_repo_ids() too. The idle check and release must be ATOMIC (release_if): the load's
# register runs under the same lock, so a plain check-then-release could clear the newer claim.
# Drop DIFFUSION ownership only if nothing is resident AND no new load is in flight: a concurrent
# /images/load that re-acquired DIFFUSION while this slow unload ran must keep ownership, or a
# later chat load sees no owner, skips eviction, and OOMs the newly resident pipeline. An
# in-flight load has is_loaded False for its whole window, so gate on loading_repo_ids() too. The
# idle check and release must be ATOMIC (release_if), since the load's register takes the lock.
engine = get_active_diffusion_engine()
await asyncio.to_thread(
release_if,
@ -16482,8 +16454,8 @@ async def diffusion_generate_progress(current_subject: str = Depends(get_current
from core.inference.diffusion_engine_router import get_active_diffusion_engine
progress = get_active_diffusion_engine().generate_progress()
# A finished generation still persisting its gallery record counts as active, so a reload's
# mount probe keeps polling instead of refreshing the gallery before the image lands.
# A finished generation still persisting its gallery record counts as active, so a reload's mount
# probe keeps polling instead of refreshing the gallery before the image lands.
if _diffusion_persist_active > 0 and not progress["active"]:
progress = {**progress, "active": True}
return DiffusionGenerateProgressResponse(**progress)
@ -16493,20 +16465,20 @@ async def diffusion_generate_progress(current_subject: str = Depends(get_current
# OpenAI-compatible images API (POST /v1/images/generations)
#
# The inference router is mounted at both /api/inference and /v1, so this also answers
# /v1/images/generations for off-the-shelf OpenAI clients, mapping CreateImageRequest onto
# the in-process diffusion backend. Studio's Image tab uses the richer /images/generate
# above; this is the spec-shaped surface and the single error boundary mapping backend
# exceptions to OpenAI error envelopes (the global /v1 handler wraps HTTPException detail).
# /v1/images/generations for off-the-shelf OpenAI clients, mapping CreateImageRequest onto the
# in-process diffusion backend. Studio's Image tab uses the richer /images/generate above; this
# is the spec-shaped surface and the single error boundary mapping backend exceptions to OpenAI
# error envelopes.
# ──────────────────────────────────────────────────────────────────────────
# Diffusion dims must land in [256, 2048] on a multiple of 16 (8x VAE downsample x 2x patch);
# the named OpenAI sizes (1024x1024, 1536x1024, 256x256, ...) all satisfy this. Mirrors
# DiffusionGenerateRequest's width/height bounds so both generate paths accept the same geometry.
# Diffusion dims must land in [256, 2048] on a multiple of 16 (8x VAE downsample x 2x patch); the
# named OpenAI sizes all satisfy this. Mirrors DiffusionGenerateRequest's bounds so both generate
# paths accept the same geometry.
_IMAGE_SIZE_RE = _re.compile(r"^(\d{1,5})\s*x\s*(\d{1,5})$")
_IMAGE_DIM_MIN, _IMAGE_DIM_MAX = 256, 2048
# Sanitized 503 detail shared by the pre-check and the unload-race branch, so both
# "no image model" responses stay identical.
# Sanitized 503 detail shared by the pre-check and the unload-race branch, so both "no image
# model" responses stay identical.
_NO_IMAGE_MODEL_MSG = "No image model loaded. Load an image model first."
@ -16569,16 +16541,16 @@ async def openai_image_generations(
)
# Use the active engine (diffusers OR native sd.cpp on a no-GPU host), the same accessor
# /images/generate uses, so a native-engine model isn't wrongly reported unloaded here.
# /images/generate uses, so a native-engine model isn't wrongly reported unloaded.
backend = get_active_diffusion_engine()
status = backend.status()
if not status.get("loaded"):
# Mirror /v1/completions and /v1/embeddings, which 503 when their backend
# isn't loaded; the global handler turns this into the OpenAI envelope.
# Mirror /v1/completions and /v1/embeddings, which 503 when their backend isn't loaded; the
# global handler turns this into the OpenAI envelope.
raise HTTPException(status_code = 503, detail = _NO_IMAGE_MODEL_MSG)
# An edit-only model (Qwen-Image-Edit, FLUX Kontext) needs an input image this API can't
# supply; refuse up front with a 400 rather than let the backend ValueError become a 500.
# An edit-only model (Qwen-Image-Edit, FLUX Kontext) needs an input image this API can't supply;
# refuse up front with a 400 rather than let the backend ValueError become a 500.
workflows = status.get("workflows") or []
if workflows and "txt2img" not in workflows:
raise HTTPException(
@ -16591,8 +16563,8 @@ async def openai_image_generations(
),
)
# Fall back to the resolved base repo so a local-path load (whose repo_id is a
# filesystem path) still gets the right per-model steps/guidance.
# Fall back to the resolved base repo so a local-path load (whose repo_id is a filesystem path)
# still gets the right per-model steps/guidance.
steps, guidance = default_generation_params(status.get("repo_id"), status.get("base_repo"))
try:
result = await asyncio.to_thread(
@ -16605,9 +16577,9 @@ async def openai_image_generations(
batch_size = body.n,
)
except Exception as exc: # noqa: BLE001 (single boundary, sanitized envelope)
# A RuntimeError with the model now unloaded means it was evicted between the readiness
# check and the call (a transient race): 503. Every other failure (CUDA OOM, a diffusers
# shape/device error) is a real 500 whose raw message must not reach the client.
# A RuntimeError with the model now unloaded means it was evicted between the readiness check and
# the call (a transient race): 503. Every other failure (CUDA OOM, a diffusers shape/device
# error) is a real 500 whose raw message must not reach the client.
if isinstance(exc, RuntimeError) and not backend.is_loaded:
raise HTTPException(status_code = 503, detail = _NO_IMAGE_MODEL_MSG)
logger.error("openai_images.generate_failed: %s", exc)
@ -16615,8 +16587,8 @@ async def openai_image_generations(
created = int(time.time())
want_b64 = body.response_format == "b64_json"
# Persist each image with its full recipe, like /images/generate, so response_format=url
# links resolve and the images show up in the gallery.
# Persist each image with its full recipe, like /images/generate, so response_format=url links
# resolve and the images show up in the gallery.
recipe = {
"prompt": body.prompt,
"negative_prompt": None,
@ -16624,15 +16596,15 @@ async def openai_image_generations(
"height": height,
"steps": steps,
"guidance": guidance,
# The batch shares one base seed, so restoring a batch_index>0 sibling needs the
# original batch_size to replay (same as /images/generate); persist it.
# The batch shares one base seed, so restoring a batch_index>0 sibling needs the original
# batch_size to replay (same as /images/generate).
"batch_size": body.n,
"model": result.get("repo_id"),
"created_at": float(created),
}
# The diffusers batch shares one seed; the native sd.cpp batch uses a distinct seed per
# image (returned in ``seeds``), so record each image's own seed like /images/generate,
# or a native batch_index>0 image shows the wrong seed.
# The diffusers batch shares one seed; the native sd.cpp batch uses a distinct seed per image
# (returned in ``seeds``), so record each image's own seed like /images/generate, or a native
# batch_index>0 image shows the wrong seed.
per_image_seeds = result.get("seeds")
def _persist() -> list[ImageGenerationData]:
@ -16643,8 +16615,8 @@ async def openai_image_generations(
if per_image_seeds and index < len(per_image_seeds)
else result["seed"]
)
# batch_seed is the base the native engine derives per-image seeds from (base + index),
# so restore replays from it rather than double-advancing the derived seed above.
# batch_seed is the base the native engine derives per-image seeds from (base + index), so
# restore replays from it rather than double-advancing the derived seed above.
record = image_gallery.save(
image,
{**recipe, "batch_index": index, "seed": seed, "batch_seed": result["seed"]},

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

@ -28,15 +28,15 @@ class CachedModelRepo(BaseModel):
repo_id: str
size_bytes: int
last_modified: Optional[float] = None
# "text-to-image" for cached diffusers image repos; declared here or response_model
# drops it, letting image-only repos pass the chat picker's task gate.
# "text-to-image" for cached diffusers image repos; declared here or response_model drops it,
# letting image-only repos pass the chat picker's task gate.
task: Optional[str] = None
# True when the snapshot is incomplete (cancelled/partial download): the picker must
# not treat it as usable, or an On Device click re-downloads the full GGUF.
# True when the snapshot is incomplete (cancelled/partial download): the picker must not treat it
# as usable, or an On Device click re-downloads the full GGUF.
partial: Optional[bool] = None
# True for a diffusion-tagged repo with NO top-level model_index.json: a single-file
# checkpoint needing from_single_file + a filename. Pickers must not offer it as a
# pipeline load (from_pretrained fails) unless the curated catalog carries its artifact.
# True for a diffusion-tagged repo with NO top-level model_index.json: a single-file checkpoint
# needing from_single_file + a filename. Pickers must not offer it as a pipeline load unless the
# curated catalog carries its artifact.
single_file: Optional[bool] = None
@ -308,10 +308,9 @@ def _scan_models_dir(models_dir: Path, *, limit: int | None = None) -> List[Loca
if not models_dir.exists() or not models_dir.is_dir():
return []
# A scan folder can point directly at a diffusers PIPELINE dir, which _is_model_directory
# rejects; without admitting it the child scan surfaces the component subdirs as bogus models
# and hides the real pipeline. The Images/Video load path loads it, so admit the root as one
# model (task tagging classifies it).
# A scan folder can point directly at a diffusers PIPELINE dir, which _is_model_directory rejects;
# without admitting it the child scan surfaces the component subdirs as bogus models and hides the
# real pipeline. The Images/Video load path loads it, so admit the root as one model.
_is_self_model = _is_model_directory(models_dir) or _local_pipeline_index(models_dir)
if _is_self_model:
@ -342,9 +341,9 @@ def _scan_models_dir(models_dir: Path, *, limit: int | None = None) -> List[Loca
has_config = (child / "config.json").exists() or (
child / "adapter_config.json"
).exists()
# A diffusers PIPELINE folder (weights in component subdirs, only model_index.json at
# the root) is missed by the checks above; the Images/Video load path accepts it, so
# admit it too or it is hidden from the On Device picker.
# A diffusers PIPELINE folder (weights in component subdirs, only model_index.json at the root) is
# missed by the checks above; the Images/Video load path accepts it, so admit it too or it is
# hidden from the On Device picker.
has_pipeline_index = _local_pipeline_index(child)
has_model_files = has_gguf or has_non_gguf_weights or has_config or has_pipeline_index
except OSError:
@ -393,11 +392,9 @@ def _scan_models_dir(models_dir: Path, *, limit: int | None = None) -> List[Loca
# A scan folder can also point directly at a BARE single-file checkpoint dir (one loose
# .safetensors / weight .bin, no config.json and no model_index.json): both root checks above
# reject it, yet the child loop admits exactly that shape via _has_non_gguf_weights and the
# Images/Video load path reinterprets such a directory through resolve_local_single_file. So
# registering the model folder itself returned no On Device row while registering its parent
# worked. Only when nothing else matched, so a models root holding a stray loose weight file
# next to real model subdirs still lists those children instead of collapsing to one row.
# reject it, yet the child loop admits exactly that shape and the Images/Video load path
# reinterprets such a directory through resolve_local_single_file. Only when nothing else
# matched, so a models root holding a stray weight file still lists its real child models.
if not found and (limit is None or limit > 0) and _has_non_gguf_weights(models_dir):
try:
updated_at = models_dir.stat().st_mtime
@ -974,8 +971,8 @@ async def list_local_models(
try:
models = collect_local_models(models_root)
# Tag each model with its task so the Images picker can filter to diffusion
# (GGUF by architecture; local checkpoints by pipeline / family).
# Tag each model with its task so the Images picker can filter to diffusion (GGUF by
# architecture; local checkpoints by pipeline / family).
models = [m.model_copy(update = {"task": _local_model_task(m)}) for m in models]
return LocalModelListResponse(
@ -2581,11 +2578,10 @@ async def delete_finetuned_model(
detail = "Could not verify model load status before deleting",
) from e
# Every guard above is chat-only, and Images / Video hold their own pipelines: a local
# diffusion or video model under the storage root loads by path, so without this rmtree
# would pull the weights (and the companion VAE / text encoders sd.cpp re-reads on every
# generation) out from under a live engine. The cached-model delete route runs the same
# check via hub.services.models.deletion; it matches by repo id, this one by path.
# Every guard above is chat-only, and Images / Video hold their own pipelines: a local diffusion
# or video model under the storage root loads by path, so without this rmtree would pull the
# weights (and the companion VAE / text encoders sd.cpp re-reads every generation) out from under
# a live engine. The cached-model delete route matches by repo id, this one by path.
for label, get_backend in (
("Images", _active_diffusion_backend),
("Video", _active_video_backend),
@ -3202,14 +3198,13 @@ def _repo_gguf_last_modified(repo_info) -> float:
return latest
# GGUF general.architecture values that denote a diffusion (image) model (everything
# else is text); lets the Images picker show only image GGUFs in its On Device list.
# GGUF general.architecture values that denote a diffusion (image) model (everything else is
# text); lets the Images picker show only image GGUFs in its On Device list.
_DIFFUSION_GGUF_ARCHS = frozenset(
{
# ONLY the families the diffusion backend can assemble (see
# diffusion_families._FAMILIES). Other diffusion archs (SD1/2/3, SDXL,
# PixArt, Lumina2, AuraFlow, Wan, HunyuanVideo, ...) would pass this filter
# then 400 in validate_load, so they stay excluded until the backend supports them.
# ONLY the families the diffusion backend can assemble (see diffusion_families._FAMILIES). Other
# diffusion archs (SD1/2/3, SDXL, PixArt, Lumina2, AuraFlow, Wan, HunyuanVideo, ...) would pass
# this filter then 400 in validate_load, so they stay excluded until the backend supports them.
"flux", # flux.1
"flux2", # flux.2-klein
"qwen_image", # qwen-image
@ -3219,10 +3214,9 @@ _DIFFUSION_GGUF_ARCHS = frozenset(
}
)
# Diffusion / image-video GGUF archs the backend can NOT assemble yet (llama.cpp also
# lacks an architecture for them); kept in sync with
# core.inference.llama_cpp.LlamaCppBackend._DIFFUSION_ARCHES minus the loadable set above.
# A dedicated non-loadable task keeps them out of the chat picker (they die with
# Diffusion / image-video GGUF archs the backend can NOT assemble yet (llama.cpp also lacks an
# architecture for them); kept in sync with LlamaCppBackend._DIFFUSION_ARCHES minus the loadable
# set above. A dedicated non-loadable task keeps them out of the chat picker (they die with
# "unknown model architecture") and out of Images (not an IMAGE_GEN_TASK; would 400).
_UNSUPPORTED_DIFFUSION_GGUF_ARCHS = frozenset(
{
@ -3237,9 +3231,8 @@ _UNSUPPORTED_DIFFUSION_GGUF_ARCHS = frozenset(
}
)
# Video GGUF archs the video backend CAN load (LTX-2.x ships as "ltxv"; the Wan
# community GGUFs as "wan"). Tagged text-to-video so they surface in the Video
# picker (VIDEO_GEN_TASKS) and stay out of chat (NON_CHAT_TASKS).
# Video GGUF archs the video backend CAN load (LTX-2.x ships as "ltxv"; the Wan community GGUFs
# as "wan"). Tagged text-to-video so they surface in the Video picker and stay out of chat.
_VIDEO_GGUF_ARCHS = frozenset({"ltxv", "wan"})
_VIDEO_GEN_TASK = "text-to-video"
@ -3263,13 +3256,11 @@ def _arch_to_task(arch: Optional[str], name_hints: tuple[Optional[str], ...] = (
if a in _DIFFUSION_GGUF_ARCHS:
return "text-to-image"
if a in _VIDEO_GGUF_ARCHS:
# Advertise as loadable video only when a VideoFamily resolves. Some archs map
# straight from the arch (ltxv); bare "wan" is ambiguous -- it covers both the
# GGUF-loadable single-DiT TI2V-5B and the dual-expert A14B MoE the loader refuses --
# so when the bare arch doesn't resolve, fall back to repo/file names (each tried
# separately, matching name segments not substrings) like the loader's own
# detect_video_family, surfacing only a non-MoE match. Without a name we can't
# disambiguate, so a bare-arch Wan GGUF stays unsupported rather than 400ing on load.
# Advertise as loadable video only when a VideoFamily resolves. Some archs map straight from the
# arch (ltxv); bare "wan" is ambiguous -- it covers both the GGUF-loadable single-DiT TI2V-5B and
# the dual-expert A14B MoE the loader refuses -- so when the bare arch doesn't resolve, fall back
# to repo/file names (each tried separately, matching name segments) like the loader's own
# detect_video_family, surfacing only a non-MoE match.
from core.inference.video_families import detect_video_family
fam = detect_video_family("", override = a)
@ -3282,8 +3273,8 @@ def _arch_to_task(arch: Optional[str], name_hints: tuple[Optional[str], ...] = (
if fam is not None and not getattr(fam, "is_moe", False):
return _VIDEO_GEN_TASK
return _UNSUPPORTED_DIFFUSION_TASK
# A diffusion arch the backend can't assemble: hide from chat (dies in llama.cpp)
# without surfacing in Images (would 400 in validate_load).
# A diffusion arch the backend can't assemble: hide from chat (dies in llama.cpp) without
# surfacing in Images (would 400 in validate_load).
if a in _UNSUPPORTED_DIFFUSION_GGUF_ARCHS:
return _UNSUPPORTED_DIFFUSION_TASK
return "text-generation"
@ -3349,10 +3340,9 @@ def _local_model_task(model: "LocalModelInfo") -> Optional[str]:
pass
return None
if _local_is_diffusers(model):
# A local diffusers pipeline can be a VIDEO family (LTX / Wan / Hunyuan), not just
# image. Tag it text-to-video so it surfaces in the Video On-Device picker instead of
# Images (which would reject it), mirroring _cached_repo_task. Gated on
# _local_is_diffusers, so only a real pipeline dir or name-matched checkpoint reaches here.
# A local diffusers pipeline can be a VIDEO family (LTX / Wan / Hunyuan), not just image. Tag it
# text-to-video so it surfaces in the Video On-Device picker instead of Images, mirroring
# _cached_repo_task. Gated on _local_is_diffusers, so only a real pipeline dir reaches here.
try:
from core.inference.video import _is_trusted_video_repo
from core.inference.video_families import detect_video_family
@ -3361,10 +3351,9 @@ def _local_model_task(model: "LocalModelInfo") -> Optional[str]:
return _VIDEO_GEN_TASK
except Exception:
pass
# The Images load path rejects a pick with no supported image-family token, 400ing AFTER
# evicting the GPU owner (a bare model_index.json dir is not enough). Tag text-to-image
# only when that same detection succeeds, so the picker never advertises a pipeline the
# load will always reject.
# The Images load path rejects a pick with no supported image-family token, 400ing AFTER evicting
# the GPU owner. Tag text-to-image only when that same detection succeeds, so the picker never
# advertises a pipeline the load will always reject.
try:
from core.inference.diffusion_families import detect_family
for needle in _local_family_needles(model):
@ -3398,9 +3387,8 @@ def _local_is_diffusers(model: "LocalModelInfo") -> bool:
return True
except Exception:
pass
# A single-file VIDEO checkpoint (LTX / Wan / Hunyuan .safetensors, no model_index.json) is
# missed above but loaded as a single_file by the video route, so surface it or the picker
# hides it. Uses _local_family_needles; _local_model_task then routes it to text-to-video.
# A single-file VIDEO checkpoint (LTX / Wan / Hunyuan .safetensors, no model_index.json) is missed
# above but loaded as a single_file by the video route, so surface it or the picker hides it.
try:
from core.inference.video_families import detect_video_family
for needle in _local_family_needles(model):
@ -3534,8 +3522,8 @@ def _repo_pipeline_missing_denoiser(repo_info) -> bool:
except ValueError:
parts = ()
if not parts:
# No snapshot scoping: fall back to the recorded name, which may itself carry
# the component subdir (e.g. 'transformer/model...').
# No snapshot scoping: fall back to the recorded name, which may itself carry the component
# subdir (e.g. 'transformer/model...').
parts = Path(name).parts
if (
len(parts) >= 2
@ -3576,8 +3564,8 @@ def _cached_repo_task(repo_info) -> Optional[str]:
from core.inference.video import _is_trusted_video_repo
from core.inference.video_families import detect_video_family
# Both gates: a detected video family (so image repos don't match) AND the
# load path's trust rule (so an untrusted video repo isn't advertised as loadable).
# Both gates: a detected video family (so image repos don't match) AND the load path's trust rule
# (so an untrusted video repo isn't advertised as loadable).
if detect_video_family(repo_id) is not None and _is_trusted_video_repo(repo_id):
return _VIDEO_GEN_TASK
except Exception:
@ -3629,15 +3617,13 @@ async def list_cached_models(
)
key = repo_id.lower()
existing = seen_lower.get(key)
# A companion-only prefetch (manifest + VAE/text-encoder but no transformer
# shards) is not a loadable pipeline; treat it as partial so the picker does
# not advertise it as on-device.
# A companion-only prefetch (manifest + VAE/text-encoder but no transformer shards) is not a
# loadable pipeline; treat it as partial so the picker does not advertise it as on-device.
is_partial = _cached_repo_partial(
repo_id, Path(repo_info.repo_path)
) or _repo_pipeline_missing_denoiser(repo_info)
# Prefer the most COMPLETE snapshot, then largest. The picker drops partial
# rows, so a partial copy in one cache root must not shadow a smaller complete
# copy in another (size only breaks ties among equal completeness).
# Prefer the most COMPLETE snapshot, then largest: a partial copy in one cache root must not
# shadow a smaller complete copy in another (size only breaks ties among equal completeness).
if existing is None or (not is_partial, total_size) > (
not bool(existing.get("partial")),
existing["size_bytes"],
@ -3649,9 +3635,8 @@ async def list_cached_models(
}
if is_partial:
row["partial"] = True
# Flag diffusion repos with no pipeline index: loadable only via
# from_single_file, so pickers must not offer them as pipeline
# loads unless the catalog carries them.
# Flag diffusion repos with no pipeline index: loadable only via from_single_file, so pickers must
# not offer them as pipeline loads unless the catalog carries them.
if row["task"] is not None and not _repo_has_pipeline_index(repo_info):
row["single_file"] = True
# Keep the newest timestamp across duplicate caches;

3
studio/backend/routes/settings.py
View file

@ -734,8 +734,7 @@ SIDEBAR_MENU_ITEM_DEFAULTS = {
"connections": False,
}
# Navigable sidebar rows the user can pin/reorder; the boolean is each id's
# default pin state, matching the shipped layout.
# Navigable sidebar rows the user can pin/reorder; the boolean is each id's default pin state.
SIDEBAR_NAV_ITEM_DEFAULTS = {
"projects": True,
"hub": True,

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

@ -224,9 +224,9 @@ async def start_training(
error = "Training already active",
)
# A diffusion (SDXL) LoRA job runs in its own subprocess on the same GPU, so an LLM
# start must refuse while one is active, or the two trainers contend for VRAM and both
# fail. Symmetric with the check in start_diffusion_training.
# A diffusion (SDXL) LoRA job runs in its own subprocess on the same GPU, so an LLM start must
# refuse while one is active or the two trainers contend for VRAM. Symmetric with the check in
# start_diffusion_training.
if _diffusion_training_active():
return TrainingJobResponse(
job_id = "",
@ -468,20 +468,18 @@ async def start_training(
logger.warning("Could not shut down export subprocess: %s", e)
try:
# A resident or in-flight Images pipeline also holds GPU memory the run needs
# and can't be cheaply sized, so tear it down unconditionally like the export
# subprocess above (the chat block below fit-checks; diffusion can't). unload()
# is a no-op when nothing is loaded and preempts an in-flight load; release the
# arbiter so it doesn't think the gone pipeline owns the GPU. Must precede the
# A resident or in-flight Images pipeline also holds GPU memory the run needs and can't be cheaply
# sized, so tear it down unconditionally like the export subprocess above (the chat block below
# fit-checks; diffusion can't). unload() no-ops when nothing is loaded and preempts an in-flight
# load; 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_engine_router import (
get_active_diffusion_engine,
)
# 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.
# 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(
@ -493,11 +491,9 @@ async def start_training(
logger.warning("Could not unload diffusion model for training: %s", e)
try:
# A resident or in-flight Video pipeline holds GPU memory the run needs too, and
# loads under the VIDEO arbiter owner the diffusion teardown above never touches.
# Tear it down the same way (unload no-ops when nothing is loaded, preempts an
# in-flight load) and release VIDEO, so a resident video session can't OOM the
# run. Must precede the chat block, which early-returns.
# A resident or in-flight Video pipeline holds GPU memory the run needs too, and loads under the
# VIDEO arbiter owner the diffusion teardown above never touches. Tear it down the same way and
# release VIDEO, so a resident video session can't OOM the run. Must precede the chat block.
from core.inference import gpu_arbiter
from core.inference.video import get_video_backend
@ -540,11 +536,10 @@ async def start_training(
from utils.transformers_version import SidecarSwapInProgress
try:
# Offloaded to a worker thread: the hook's diffusion/video unload() waits on the
# engines' generation locks until an in-flight denoise step hits its cancel callback
# (and the export subprocess teardown can take seconds), which would otherwise block
# the event loop and freeze every concurrent status/cancel/UI request. Overlapping
# starts are serialized by the backend's own start-in-progress guard.
# Offloaded to a worker thread: the hook's diffusion/video unload() waits on the engines'
# generation locks until an in-flight denoise step hits its cancel callback (and the export
# subprocess teardown can take seconds), which would otherwise freeze every concurrent
# status/cancel/UI request. Overlapping starts are serialized by the backend's own guard.
success = await asyncio.to_thread(
backend.start_training,
job_id = job_id,
@ -1149,10 +1144,9 @@ async def stream_training_progress(
# ── Diffusion (SDXL) LoRA training ────────────────────────────────────────────
# A separate, lightweight job path from the LLM endpoints above: diffusion runs are driven
# by DiffusionTrainingService (its own subprocess + event pump), not the LLM TrainingBackend,
# so the two never contend and diffusion never triggers LLM lifecycle (DB run rows, plots,
# transfer-to-chat-inference).
# A separate, lightweight job path from the LLM endpoints above: diffusion runs are driven by
# DiffusionTrainingService (its own subprocess + event pump), not the LLM TrainingBackend, so the
# two never contend and diffusion never triggers LLM lifecycle (DB run rows, plots, transfer).
def _diffusion_training_active() -> bool:
@ -1204,9 +1198,9 @@ def _free_gpu_for_diffusion_training() -> None:
from core.inference import gpu_arbiter
from core.inference.diffusion_engine_router import get_active_diffusion_engine
# The ACTIVE engine, not the diffusers singleton: on a native (sd_cpp) selection the
# diffusers backend reports unloaded while the resident sd-server still holds the GPU,
# so unloading only the singleton is a no-op. Mirrors the LLM training start path.
# The ACTIVE engine, not the diffusers singleton: on a native (sd_cpp) selection the diffusers
# backend reports unloaded while the resident sd-server still holds the GPU, so unloading only
# the singleton is a no-op. Mirrors the LLM training start path.
diffusion = get_active_diffusion_engine()
if diffusion.is_loaded:
logger.info("Unloading resident Images pipeline to free GPU memory for training")
@ -1216,9 +1210,9 @@ def _free_gpu_for_diffusion_training() -> None:
logger.warning("Could not unload Images pipeline for diffusion training: %s", e)
try:
# A resident Video pipeline loads under the VIDEO arbiter owner the Images teardown
# above doesn't free; unload it too (no-op when nothing is loaded) and release VIDEO
# so a resident video session can't OOM the diffusion trainer.
# A resident Video pipeline loads under the VIDEO arbiter owner the Images teardown above
# doesn't free; unload it too (no-op when nothing is loaded) and release VIDEO so a resident
# video session can't OOM the diffusion trainer.
from core.inference import gpu_arbiter
from core.inference.video import get_video_backend
@ -1231,9 +1225,8 @@ def _free_gpu_for_diffusion_training() -> None:
logger.warning("Could not unload Video pipeline for diffusion training: %s", e)
try:
# The SDXL trainer's footprint can't be cheaply sized against a resident chat model,
# so free chat unconditionally (like the LLM path does for an in-flight load) rather
# than risk an OOM.
# The SDXL trainer's footprint can't be cheaply sized against a resident chat model, so free chat
# unconditionally (like the LLM path does for an in-flight load) rather than risk an OOM.
from routes.training_vram import free_chat_models_for_training, summarize_resident_chat
if summarize_resident_chat()["any"]:
freed = free_chat_models_for_training(reason = "diffusion training starting")
@ -1274,8 +1267,8 @@ def _preflight_gated_base(base_model: str, hf_token: Optional[str]) -> None:
f"try again."
),
)
# 404 (e.g. a repo without a root model_index.json) and other codes are not an
# access problem -- let the trainer surface any genuine load error.
# 404 (e.g. a repo without a root model_index.json) and other codes are not an access problem;
# let the trainer surface any genuine load error.
except Exception: # noqa: BLE001 -- network/DNS hiccup must not block a start
return
@ -1295,12 +1288,12 @@ def _resolve_diffusion_data_dir(raw: str) -> Path:
value = str(raw or "").strip()
if value and "\x00" not in value:
p = Path(value)
# Single component and not ".." -> joining under datasets_root() cannot escape it.
# A single component that is not "..", so joining under datasets_root() cannot escape it.
if not p.is_absolute() and len(p.parts) == 1 and p.parts[0] != "..":
direct = datasets_root() / value
# Route a bare name through the same protected resolver the CRUD routes use, so a
# name -> external-directory symlink is rejected here too (is_dir() follows the link).
# Include a broken symlink so it is rejected, not passed to resolve_dataset_path.
# Route a bare name through the same protected resolver the CRUD routes use, so a name to
# external-directory symlink is rejected here too (is_dir() follows the link). A broken symlink
# is included so it is rejected, not passed to resolve_dataset_path.
if direct.is_dir() or direct.is_symlink():
return _resolve_dataset_folder(value)
return resolve_dataset_path(raw)
@ -1317,7 +1310,7 @@ async def start_diffusion_training(
# Under API-key auth, refuse to start training while a request is in flight:
# _free_gpu_for_diffusion_training() below unloads the chat backends, killing the stream.
# Mirrors start_training so a diffusion start can't silently drop an active API request.
# Mirrors start_training.
if via_api_key is True:
from core.inference.llama_keepwarm import other_inference_request_count
if (
@ -1333,8 +1326,8 @@ async def start_diffusion_training(
),
)
# Interlock: refuse while an LLM training run holds the GPU (symmetric with the diffusion
# check in start_training), so the two trainers never contend for VRAM.
# Interlock: refuse while an LLM training run holds the GPU (symmetric with the diffusion check
# in start_training), so the two trainers never contend for VRAM.
try:
if get_training_backend().is_training_active():
raise HTTPException(
@ -1349,9 +1342,9 @@ async def start_diffusion_training(
except Exception: # noqa: BLE001 -- backend import/health issue must not block a start
pass
# Resolve + contain the dataset and output paths BEFORE spawning, so Studio-relative names
# ("uploads/my-images") work and absolute paths stay under a Studio root -- the trainer
# subprocess otherwise resolves them relative to its own cwd.
# Resolve + contain the dataset and output paths BEFORE spawning, so Studio-relative names work
# and absolute paths stay under a Studio root -- the trainer subprocess otherwise resolves them
# relative to its own cwd.
config = body.model_dump()
try:
from utils.paths import resolve_output_dir
@ -1360,9 +1353,9 @@ async def start_diffusion_training(
except ValueError as e:
raise HTTPException(status_code = 400, detail = str(e))
# Validate the config BEFORE freeing resident GPU workloads, so a start then refused (bad
# numbers, non-SDXL base) never tears down the user's chat/Images model. service.start()
# re-runs this cheaply before spawn.
# Validate the config BEFORE freeing resident GPU workloads, so a start then refused (bad numbers,
# non-SDXL base) never tears down the user's chat/Images model. service.start() re-runs this
# cheaply before spawn.
from core.training.diffusion_lora_trainer import _config_from_dict
try:
@ -1370,11 +1363,10 @@ async def start_diffusion_training(
except ValueError as e:
raise HTTPException(status_code = 400, detail = str(e))
# Preflight the requested DiT precision BEFORE freeing GPU residents: the trainer's own
# checks (bf16-capable GPU required; explicit int8 needs a functional torchao) fire only in
# the child, AFTER _free_gpu_for_diffusion_training() evicted the user's model. Fail fast
# (400) so a pre-Ampere GPU (T4 / V100 / RTX 20xx) or stub-torchao host never tears down
# residents for a run that cannot start.
# Preflight the requested DiT precision BEFORE freeing GPU residents: the trainer's own checks
# (bf16-capable GPU required; explicit int8 needs a functional torchao) fire only in the child,
# AFTER _free_gpu_for_diffusion_training() evicted the user's model. Fail fast (400) so a
# pre-Ampere GPU or stub-torchao host never tears down residents for a run that cannot start.
from core.training.diffusion_train_common import training_precision_preflight_error
_precision_reason = training_precision_preflight_error(
@ -1383,9 +1375,8 @@ async def start_diffusion_training(
if _precision_reason:
raise HTTPException(status_code = 400, detail = _precision_reason)
# Run the trainers' trust gate here too (both assert the same predicate before
# from_pretrained), so an untrusted/typoed base 400s BEFORE freeing GPU residents rather
# than tearing down the user's model and failing in the child.
# Run the trainers' trust gate here too (both assert the same predicate before from_pretrained),
# so an untrusted/typoed base 400s BEFORE freeing GPU residents rather than failing in the child.
from core.training.diffusion_train_common import _assert_trusted_base_model
try:
@ -1393,11 +1384,10 @@ async def start_diffusion_training(
except ValueError as e:
raise HTTPException(status_code = 400, detail = str(e))
# Preflight access to a gated base repo with the user's token BEFORE freeing GPU residents,
# so a missing/insufficient token fails fast (400) without tearing down the user's model, and
# never surfaces as a confusing mid-load 401. Offloaded to a worker thread: it does a blocking
# urlopen HEAD (5s timeout) to HF, which would otherwise stall the event loop and every
# concurrent status/progress/cancel request (as the filesystem preflight below also does).
# Preflight access to a gated base repo with the user's token BEFORE freeing GPU residents, so a
# missing/insufficient token fails fast (400) without tearing down the user's model and never
# surfaces as a confusing mid-load 401. Offloaded to a worker thread: it does a blocking urlopen
# HEAD (5s timeout) that would otherwise stall the event loop.
await asyncio.to_thread(
_preflight_gated_base, config.get("base_model", ""), config.get("hf_token")
)
@ -1407,38 +1397,33 @@ async def start_diffusion_training(
service = get_diffusion_training_service()
# Reserve the training slot BEFORE the dataset preflight (not just before freeing residents):
# is_active() otherwise flips true only at service.start(), so during this scan -- which
# decode-probes every image and can take noticeable time on a large folder -- a concurrent
# upload/caption/delete would pass _require_diffusion_dataset_mutable() and mutate the dataset
# the trainer is about to read (training the wrong data, or a missing file mid-step), and a
# concurrent /images/load or /video/load would pass its guard and double-allocate VRAM.
# reserve() is a compare-and-set: a second overlapping /diffusion/start raises RuntimeError
# (-> 409) before touching anything. unreserve() runs in the finally ONLY when THIS request
# reserved, so a rejected second request can't clear the claim, and any preflight failure below
# rolls the reservation back.
# decode-probes every image and can take a while -- a concurrent upload/caption/delete would pass
# _require_diffusion_dataset_mutable() and mutate the dataset the trainer is about to read, and a
# concurrent /images/load or /video/load would double-allocate VRAM. reserve() is a
# compare-and-set, so a second overlapping start 409s before touching anything; unreserve() runs
# in the finally ONLY when THIS request reserved.
reserved = False
try:
service.reserve()
reserved = True
# Preflight the dataset: a missing/empty/uncaptionable data_dir otherwise fails inside the
# spawned trainer AFTER the user's model was evicted. Same discovery the trainer runs, so
# the two cannot disagree.
# Preflight the dataset: a missing/empty/uncaptionable data_dir otherwise fails inside the spawned
# trainer AFTER the user's model was evicted. Same discovery the trainer runs, so the two cannot
# disagree.
try:
await asyncio.to_thread(
_dtc.discover_image_caption_pairs,
config["data_dir"],
instance_prompt = config.get("instance_prompt") or None,
caption_column = config.get("caption_column") or "text",
# Decode-probe every image now (cheap PIL header check) so a corrupt/zero-byte
# upload 400s BEFORE _free_gpu_for_diffusion_training() tears down the user's
# models, rather than crashing the spawned trainer post-eviction.
# Decode-probe every image now (cheap PIL header check) so a corrupt/zero-byte upload 400s BEFORE
# _free_gpu_for_diffusion_training() tears down the user's models.
verify_images = True,
)
except (FileNotFoundError, ValueError) as e:
raise HTTPException(status_code = 400, detail = str(e))
# Free resident GPU workloads (export / Images pipeline / chat) before the trainer loads
# its own pipeline. Offload the blocking teardown (engine unload waits on generation
# locks; export subprocess join can take seconds) to a worker thread so the event loop
# stays free for concurrent status/progress/cancel requests.
# Free resident GPU workloads (export / Images pipeline / chat) before the trainer loads its own
# pipeline. Offload the blocking teardown (engine unload waits on generation locks; export
# subprocess join can take seconds) to a worker thread so the event loop stays responsive.
await asyncio.to_thread(_free_gpu_for_diffusion_training)
job_id = service.start(config)
except ValueError as e:
@ -1457,9 +1442,8 @@ async def start_diffusion_training(
log = logger,
)
finally:
# On success the now-live proc keeps is_active() true; on failure this clears the
# reservation so training isn't left permanently "active". Only the request that reserved
# clears it, so a rejected overlapping start doesn't drop the winner's claim.
# On success the now-live proc keeps is_active() true; on failure this clears the reservation so
# training isn't left permanently "active". Only the request that reserved clears it.
if reserved:
service.unreserve()
return DiffusionTrainingStartResponse(job_id = job_id, status = "running")
@ -1506,9 +1490,8 @@ async def list_diffusion_training_runs(
summaries: list[DiffusionTrainingRunSummary] = []
for r in list_diffusion_runs(limit = limit):
# list_diffusion_runs already skips non-dict / missing-id records, but a wrong-typed
# field (e.g. a non-numeric avg_loss) would still raise here; catch it per record so
# one bad file never breaks the whole Previous runs panel.
# list_diffusion_runs already skips non-dict / missing-id records, but a wrong-typed field would
# still raise here; catch it per record so one bad file never breaks the whole Previous runs panel.
try:
summaries.append(DiffusionTrainingRunSummary(**r))
except ValidationError:
@ -1526,20 +1509,20 @@ async def get_diffusion_training_run(
rec = get_diffusion_run(job_id)
# A valid-JSON file that is not an object (a truncated / hand-edited [] record) makes
# DiffusionTrainingRunDetail(**rec) raise TypeError -- not the ValidationError caught below
# -- and 500 the endpoint. Treat any non-dict record as absent, like the list route.
# DiffusionTrainingRunDetail(**rec) raise TypeError -- not the ValidationError caught below -- and
# 500 the endpoint. Treat any non-dict record as absent, like the list route.
if not isinstance(rec, dict):
raise HTTPException(status_code = 404, detail = "No such training run.")
try:
return DiffusionTrainingRunDetail(**rec)
except ValidationError:
# A malformed on-disk record (hand-edited / older shape) reads as absent rather than
# 500 the endpoint, like the list route skips bad records.
# A malformed on-disk record (hand-edited / older shape) reads as absent rather than 500 the
# endpoint, like the list route skips bad records.
raise HTTPException(status_code = 404, detail = "No such training run.")
# Extensions accepted into an image-training dataset folder: images the trainer reads,
# plus its caption sources (per-image sidecars and metadata/captions jsonl).
# Extensions accepted into an image-training dataset folder: images the trainer reads, plus its
# caption sources (per-image sidecars and metadata/captions jsonl).
_DIFFUSION_DATASET_IMAGE_EXTS = {".png", ".jpg", ".jpeg", ".webp", ".bmp"}
_DIFFUSION_DATASET_TEXT_EXTS = {".txt", ".caption", ".jsonl"}
@ -1573,10 +1556,9 @@ def _resolve_dataset_caption(
def _diffusion_dataset_summary(folder: Path) -> DiffusionDatasetSummary:
# Count an image as captioned only when it resolves to a NON-EMPTY caption via the same
# sidecar > metadata precedence the trainer uses -- an empty tombstone sidecar shadows a
# metadata row and makes the trainer skip the image, so counting it would over-report
# caption_count and mislabel an uncaptioned dataset as captioned.
# Count an image as captioned only when it resolves to a NON-EMPTY caption via the same sidecar
# over metadata precedence the trainer uses: an empty tombstone sidecar shadows a metadata row and
# makes the trainer skip the image, so counting it would mislabel an uncaptioned dataset.
meta_captions = _load_metadata_captions(folder)
images = captions = 0
for f in folder.iterdir():
@ -1602,13 +1584,13 @@ async def diffusion_training_info(current_subject: str = Depends(get_current_sub
root = datasets_root()
found: list[DiffusionDatasetSummary] = []
try:
# Skip hidden dirs: never user datasets, and an in-progress example import stages
# into a dot-prefixed sibling that must not surface as a dataset.
# Skip hidden dirs: never user datasets, and an in-progress example import stages into a
# dot-prefixed sibling that must not surface as a dataset.
children = sorted(
p
for p in root.iterdir()
# Skip symlinked dirs: the CRUD resolver rejects them, so discovery must not
# advertise one as selectable (the read/caption/delete routes would refuse it).
# Skip symlinked dirs: the CRUD resolver rejects them, so discovery must not advertise one as
# selectable (the read/caption/delete routes would refuse it).
if p.is_dir() and not p.is_symlink() and not p.name.startswith(".")
)
except OSError:
@ -1673,7 +1655,7 @@ async def upload_diffusion_dataset(
_require_diffusion_dataset_mutable()
cleaned = _clean_diffusion_dataset_name(name)
# Run the same symlink + root-containment check as the read/caption/delete endpoints before any
# write, so a name -> external-directory symlink can't make the staged upload write outside root.
# write, so a name to external-directory symlink can't make the staged upload write outside root.
folder = _resolve_dataset_folder(name, must_exist = False)
folder.mkdir(parents = True, exist_ok = True)
@ -1681,17 +1663,15 @@ async def upload_diffusion_dataset(
total_bytes = 0
uploaded = 0
allowed = _DIFFUSION_DATASET_IMAGE_EXTS | _DIFFUSION_DATASET_TEXT_EXTS
# Validate every filename up front so a valid image ahead of a bad one isn't left on disk
# when the 400 fires -- make the upload all-or-nothing.
# Validate every filename up front so a valid image ahead of a bad one isn't left on disk when the
# 400 fires; the upload is all-or-nothing.
names: list[str] = []
for f in files:
# Normalise to a safe basename. Path.name doesn't split on a backslash on POSIX, so a
# Windows client sending a backslash path in the multipart filename would be stored
# verbatim; fold backslashes to forward slashes first so the true basename is taken for
# both separators. The read/caption/delete endpoints run the stored name through
# _safe_dataset_image_path (rejects "\\" / ".." / path chars), so a name still holding
# ".." here would list an image the grid can never preview, caption, or delete -- reject
# it now instead of persisting an unmanageable orphan.
# Normalise to a safe basename. Path.name doesn't split on a backslash on POSIX, so a Windows
# client sending a backslash path in the multipart filename would be stored verbatim; fold
# backslashes first so the true basename is taken for both separators. The read/caption/delete
# endpoints run the stored name through _safe_dataset_image_path, so a name still holding ".."
# here would list an image the grid can never preview, caption, or delete.
filename = Path((f.filename or "").replace("\\", "/")).name.strip().replace("\x00", "")
ext = Path(filename).suffix.lower()
if not filename or ".." in filename or ext not in allowed:
@ -1700,13 +1680,11 @@ async def upload_diffusion_dataset(
status_code = 400,
detail = f"Unsupported file '{f.filename}'. Allowed: {exts}",
)
# Reject an EXACT duplicate name within THIS batch (two cat.png from different folders,
# or an API client repeating a part). The same-name exemption below is for SEPARATE
# repeat uploads, a deliberate overwrite of the file on disk; inside one batch the two
# parts are distinct files staged to the same destination on EVERY filesystem, so the
# later tmp.replace(dest) would silently discard the earlier one while `uploaded` counts
# both. Exact match only: a case VARIANT pair (pic.png vs Pic.png) stays exempt per the
# stem guard -- one file / overwrite on case-insensitive filesystems, two on Linux.
# Reject an EXACT duplicate name within THIS batch (two cat.png from different folders, or an API
# client repeating a part). The same-name exemption below is for SEPARATE repeat uploads, a
# deliberate overwrite; inside one batch the two parts are distinct files staged to the same
# destination on EVERY filesystem, so the later replace would silently discard the earlier one.
# Exact match only: a case VARIANT pair stays exempt per the stem guard.
fname_cf = filename.casefold()
if filename in names:
raise HTTPException(
@ -1717,22 +1695,18 @@ async def upload_diffusion_dataset(
"uploading."
),
)
# Reject a second IMAGE sharing this stem but differing by extension (sample.png vs
# sample.jpg): both resolve to the same <stem>.txt sidecar (the kohya/diffusers
# convention the reader, editor, and delete paths use), so keeping both would silently
# share -- and corrupt -- one caption. Check files already on disk (uploads accumulate)
# and earlier images in THIS batch. Re-uploading the exact same name (stem AND extension)
# stays an overwrite; caption/text files are exempt (sample.txt for sample.png is fine).
# Reject a second IMAGE sharing this stem but differing by extension (sample.png vs sample.jpg):
# both resolve to the same <stem>.txt sidecar (the kohya/diffusers convention the reader, editor
# and delete paths use), so keeping both would silently share -- and corrupt -- one caption. Check
# files already on disk and earlier images in THIS batch. Re-uploading the exact same name stays
# an overwrite; caption/text files are exempt.
if ext in _DIFFUSION_DATASET_IMAGE_EXTS:
stem = Path(filename).stem
# Compare stems (and the same-name guard) case-insensitively: on case-insensitive
# filesystems (Windows/macOS) two images whose stems differ only by case (sample.png
# vs Sample.jpg) resolve to the SAME <stem>.txt sidecar, so a case-sensitive check
# would let both share -- and corrupt -- one caption. A same-name case variant is
# exempt ONLY when its stem also differs in case (sample.png vs Sample.png): one file /
# overwrite on case-insensitive filesystems, SEPARATE sidecars on Linux. An
# EXTENSION-case variant (cat.PNG vs cat.png) has equal stems, so on Linux both land
# and resolve to ONE cat.txt -- the collision this guard exists for -- and is rejected.
# Compare stems (and the same-name guard) case-insensitively: on case-insensitive filesystems two
# images whose stems differ only by case resolve to the SAME <stem>.txt sidecar, so a
# case-sensitive check would let both corrupt one caption. A same-name case variant is exempt ONLY
# when its stem also differs in case (one file on case-insensitive filesystems, separate sidecars
# on Linux). An EXTENSION-case variant (cat.PNG vs cat.png) has equal stems, so it is rejected.
stem_cf = stem.casefold()
def _shares_sidecar(other_name: str) -> bool:
@ -1743,8 +1717,8 @@ async def upload_diffusion_dataset(
or other.stem.casefold() != stem_cf
):
return False
# A casefold-equal full name is exempt unless the stems match EXACTLY
# (extension-case variants collide on one sidecar on case-sensitive FS).
# A casefold-equal full name is exempt unless the stems match EXACTLY (extension-case variants
# collide on one sidecar on case-sensitive filesystems).
return other.stem == stem or other_name.casefold() != fname_cf
clash = next(
@ -1763,16 +1737,16 @@ async def upload_diffusion_dataset(
),
)
names.append(filename)
# Stage each file to a temp name and move it into place only once the whole batch is written,
# so a mid-batch failure (size limit, disk error, disconnect) leaves the dataset untouched --
# including any pre-existing same-name file a direct write would have truncated.
# Stage each file to a temp name and move it into place only once the whole batch is written, so a
# mid-batch failure (size limit, disk error, disconnect) leaves the dataset untouched, including
# any pre-existing same-name file a direct write would have truncated.
staged: list[tuple[Path, Path]] = [] # (temp, final)
committed = False
try:
for f, filename in zip(files, names):
dest = folder / filename
# A filename-independent temp name so a long (but valid, <= NAME_MAX) filename
# can't overflow NAME_MAX once the staging suffix is added.
# A filename-independent temp name so a long (but valid) filename can't overflow NAME_MAX once the
# staging suffix is added.
tmp = folder / f".upload-{_uuid.uuid4().hex}.part"
staged.append((tmp, dest))
with open(tmp, "wb") as out:
@ -1788,22 +1762,21 @@ async def upload_diffusion_dataset(
),
)
out.write(chunk)
# Reject a decompression bomb before commit: a small compressible PNG can pass the byte
# limit yet decode to huge pixels and OOM the trainer's latent cache, so bound each
# image's dimensions from the header (mirrors diffusion._decode_b64_image).
# Reject a decompression bomb before commit: a small compressible PNG can pass the byte limit yet
# decode to huge pixels and OOM the trainer's latent cache, so bound each image's dimensions from
# the header (mirrors diffusion._decode_b64_image).
if Path(filename).suffix.lower() in _DIFFUSION_DATASET_IMAGE_EXTS:
_validate_uploaded_training_image(tmp, filename)
uploaded += 1
# Re-check the interlock immediately before the commit: the entry guard only saw the
# pre-upload state, so a /diffusion/start could have reserved the training slot while we
# were streaming (reserve() flips is_active before the run launches). Committing now would
# move images/captions underneath the trainer despite the guard; a 409 here leaves the
# staged temps to be cleaned by the finally below (committed stays False).
# Re-check the interlock immediately before the commit: the entry guard only saw the pre-upload
# state, so a /diffusion/start could have reserved the training slot while we were streaming.
# Committing now would move images/captions underneath the trainer; a 409 here leaves the staged
# temps to the finally below.
_require_diffusion_dataset_mutable()
# Commit every staged file as one transaction. A plain replace loop is not atomic across
# files: a mid-loop tmp.replace(dest) failure leaves earlier destinations already overwritten
# while the request errors. Back up each pre-existing destination first, then on any failure
# drop the versions this request installed and restore every displaced original.
# Commit every staged file as one transaction. A plain replace loop is not atomic across files: a
# mid-loop failure leaves earlier destinations already overwritten while the request errors. Back
# up each pre-existing destination first, then on any failure drop the versions this request
# installed and restore every displaced original.
backups: list[tuple[Path, Optional[Path]]] = [] # (dest, backup path or None)
installed: list[Path] = []
try:
@ -1856,8 +1829,8 @@ async def upload_diffusion_dataset(
# ── Dataset labeling (per-image caption editing) + one-click example imports ──
# Thumbnails live in a hidden subdir so they never appear in dataset listings or the
# trainer's image discovery (both scan only top-level files).
# Thumbnails live in a hidden subdir so they never appear in dataset listings or the trainer's
# image discovery (both scan only top-level files).
_THUMBS_DIRNAME = ".thumbs"
_MAX_CAPTION_CHARS = 2000
@ -1907,9 +1880,9 @@ def _validate_uploaded_training_image(path: Path, original_name: str) -> None:
with Image.open(path) as image:
width, height = image.size
except Image.DecompressionBombError:
# Past Pillow's own hard limit (> 2 x MAX_IMAGE_PIXELS ~ 179 MP) Image.open() raises before
# .size can be read. That error derives straight from Exception (not OSError/ValueError), so
# letting it escape 500s the upload; it is exactly the oversized image this guard rejects.
# Past Pillow's own hard limit (~179 MP) Image.open() raises before .size can be read. That error
# derives straight from Exception (not OSError/ValueError), so letting it escape 500s the upload;
# it is exactly the oversized image this guard rejects.
raise HTTPException(
status_code = 400,
detail = (
@ -1957,8 +1930,8 @@ def _load_metadata_captions(folder: Path) -> dict[str, str]:
meta_path = folder / meta_name
if not meta_path.is_file():
continue
# Tolerate a bad upload (invalid UTF-8, or a line of non-object JSON): skip the record so
# the info / labeling / caption / summary endpoints don't 500.
# Tolerate a bad upload (invalid UTF-8, or a line of non-object JSON): skip the record so the
# info / labeling / caption / summary endpoints don't 500.
try:
lines = meta_path.read_text(encoding = "utf-8").splitlines()
except (OSError, UnicodeError):
@ -1998,14 +1971,13 @@ def _image_record(
source = "sidecar"
except (OSError, UnicodeError):
# Unreadable / invalid UTF-8 sidecar (uploads store text sidecars as raw bytes):
# UnicodeDecodeError is a ValueError, not an OSError, so an OSError-only guard let
# it 500 the whole labeling grid. Read it as no caption, like the info summary's
# _resolve_dataset_caption does.
# UnicodeDecodeError is a ValueError, not an OSError, so an OSError-only guard let it 500 the
# whole labeling grid. Read it as no caption, like the info summary does.
caption = None
break
if caption is None:
# Basename first, then the relative path as written in the jsonl (as_posix so a Windows
# backslash path still matches forward-slash keys) -- discover_image_caption_pairs's order.
# Basename first, then the relative path as written in the jsonl (as_posix so a Windows backslash
# path still matches forward-slash keys): discover_image_caption_pairs's order.
meta = meta_captions.get(image_path.name)
if meta is None:
try:
@ -2080,9 +2052,9 @@ async def get_diffusion_dataset_image(
thumbs_dir = folder / _THUMBS_DIRNAME
thumbs_dir.mkdir(exist_ok = True)
# Key on the full filename (stem + extension), not the stem: two images sharing a stem
# but differing by extension (sample.png / sample.jpg) would otherwise collide on one
# cache file, and an mtime-newer cache for the first would be served for the second.
# Key on the full filename (stem + extension), not the stem: two images sharing a stem but
# differing by extension would otherwise collide on one cache file, and an mtime-newer cache for
# the first would be served for the second.
thumb_path = thumbs_dir / f"{image_path.name}_{size}.jpg"
src_mtime = image_path.stat().st_mtime
if thumb_path.is_file() and thumb_path.stat().st_mtime >= src_mtime:
@ -2131,10 +2103,10 @@ async def set_diffusion_dataset_caption(
sidecar.write_text(caption, encoding = "utf-8")
image_path.with_suffix(".caption").unlink(missing_ok = True)
return _image_record(folder, image_path, _load_metadata_captions(folder))
# Blank must actually clear. Unlinking alone would resurface this image's metadata.jsonl
# / captions.jsonl caption (the fallback), so when one exists write an EMPTY sidecar
# instead: both the reader and the trainer's discovery treat an existing sidecar as
# authoritative even when empty, a tombstone. No metadata caption -> plain cleanup.
# Blank must actually clear. Unlinking alone would resurface this image's metadata.jsonl /
# captions.jsonl caption, so when one exists write an EMPTY sidecar instead: both the reader and
# the trainer's discovery treat an existing sidecar as authoritative even when empty, a tombstone.
# With no metadata caption it is a plain cleanup.
meta = _load_metadata_captions(folder)
try:
rel = image_path.relative_to(folder).as_posix()
@ -2171,10 +2143,9 @@ async def delete_diffusion_dataset_image(
image_path.with_suffix(ext).unlink(missing_ok = True)
thumbs_dir = folder / _THUMBS_DIRNAME
if thumbs_dir.is_dir():
# Thumbs are keyed on the full filename (stem + extension), so match that here too;
# a stem-only glob would strand this image's thumbs or delete a same-stem sibling's.
# Escape the name: a raw glob metacharacter (e.g. "[ab].png") would match siblings'
# thumbs while leaving its own behind.
# Thumbs are keyed on the full filename (stem + extension), so match that here too; a stem-only
# glob would strand this image's thumbs or delete a same-stem sibling's. Escape the name: a raw
# glob metacharacter would match siblings' thumbs while leaving its own behind.
for t in thumbs_dir.glob(f"{_glob.escape(image_path.name)}_*.jpg"):
t.unlink(missing_ok = True)
return {"deleted": image_path.name}
@ -2183,9 +2154,9 @@ async def delete_diffusion_dataset_image(
# Curated, license-labelled example datasets for one-click import. ``loader`` picks the
# materialization strategy: "hf_dataset" streams rows from datasets.load_dataset (image +
# optional caption column); "imagefolder_jsonl" snapshot-downloads a dataset repo whose
# captions live in a *.jsonl (file_name/text) not a standard metadata.jsonl.
# materialization strategy: "hf_dataset" streams rows from datasets.load_dataset (image + optional
# caption column); "imagefolder_jsonl" snapshot-downloads a dataset repo whose captions live in a
# *.jsonl (file_name/text) not a standard metadata.jsonl.
_DATASET_EXAMPLES: list[dict] = [
{
"id": "dreambooth-dog",
@ -2230,8 +2201,8 @@ _DATASET_EXAMPLES: list[dict] = [
"description": "100 butterfly photos. No captions, so use the trigger prompt.",
"license": "CC0",
"image_cap": 100,
# The metadata columns are species names / boilerplate alt-text, not captions, so train
# it as a subject set with the trigger prompt instead.
# The metadata columns are species names / boilerplate alt-text, not captions, so train it as a
# subject set with the trigger prompt instead.
"suggested_trigger": "a photo of a sks butterfly",
"loader": "hf_dataset",
"caption_column": None,
@ -2423,13 +2394,11 @@ async def import_diffusion_dataset_example(
imported = 0
if existing.image_count == 0:
cap = int(entry["image_cap"])
# Materialize into a private staging dir and promote into the dataset folder only
# after the whole import succeeds. A partial materialize (a transient fetch/copy
# error after some images) then leaves only the staging dir, never a half-filled
# dataset -- otherwise the image_count>0 idempotency check above would treat that
# partial as complete on retry (imported=0) and strand a truncated dataset (there is
# no dataset-delete flow). Staged as a hidden same-filesystem sibling so promotion is
# an atomic rename.
# Materialize into a private staging dir and promote into the dataset folder only after the whole
# import succeeds. A partial materialize then leaves only the staging dir, never a half-filled
# dataset -- otherwise the image_count>0 idempotency check above would treat that partial as
# complete on retry and strand a truncated dataset (there is no dataset-delete flow). Staged as a
# hidden same-filesystem sibling so promotion is an atomic rename.
staging = Path(tempfile.mkdtemp(dir = folder.parent, prefix = f".{folder.name}.import-"))
try:
try:
@ -2449,13 +2418,11 @@ async def import_diffusion_dataset_example(
status_code = 502,
detail = f"No images found in '{entry['repo']}'.",
)
# Promote the fully-materialized staging dir as a UNIT. A per-file move loop is
# not atomic: a hard process death (SIGKILL / OOM / power loss) mid-loop would
# leave SOME images, which the image_count>0 idempotency check above would accept
# as complete on retry. The folder was created empty here (runs only when it holds
# no images), so a single same-filesystem rename is atomic. If the folder holds
# unrelated non-image files (rmdir refuses), fall back to a per-file move rather
# than abort -- the common fresh-import path stays atomic.
# Promote the fully-materialized staging dir as a UNIT. A per-file move loop is not atomic: a hard
# process death mid-loop would leave SOME images, which the image_count>0 idempotency check would
# accept as complete on retry. The folder was created empty here, so a single same-filesystem
# rename is atomic. If it holds unrelated non-image files (rmdir refuses), fall back to a per-file
# move rather than abort.
try:
os.rmdir(folder)
except OSError:

49
studio/backend/routes/video.py
View file

@ -60,8 +60,8 @@ def _guard_video_load_against_training() -> None:
diffusion_active = get_diffusion_training_service().is_active()
except Exception: # noqa: BLE001
diffusion_active = False
# An SDXL LoRA trainer runs in its own subprocess on the same GPU, so refuse a video
# load while one is active too (VRAM competition). Symmetric with the image-load interlock.
# An SDXL LoRA trainer runs in its own subprocess on the same GPU, so refuse a video load while
# one is active too. Symmetric with the image-load interlock.
if not llm_active and not diffusion_active:
return
raise HTTPException(
@ -126,12 +126,12 @@ async def load_video_model(
backend = get_video_backend()
try:
# Resolve the load kind once (gguf / single_file / pipeline) so validation and the load
# agree; a bad explicit kind raises here -> 400.
# Resolve the load kind once (gguf / single_file / pipeline) so validation and the load agree; a
# bad explicit kind raises here, so a 400.
kind = resolve_video_model_kind(request.gguf_filename, request.model_kind)
# A local On-Device pick can be a bare single-file .safetensors dir (no model_index.json)
# that the picker starts as a pipeline with no filename, which would 400 on the missing
# index. If the dir holds exactly one checkpoint, load it as a single_file. Mirrors images.
# A local On-Device pick can be a bare single-file .safetensors dir (no model_index.json) that the
# picker starts as a pipeline with no filename, which would 400 on the missing index. If the dir
# holds exactly one checkpoint, load it as a single_file. Mirrors images.
if kind == "pipeline" and not request.gguf_filename:
sole = await asyncio.to_thread(resolve_local_single_file, request.model_path)
if sole is not None:
@ -150,13 +150,13 @@ async def load_video_model(
)
# Refuse while training is running (VRAM competition). Mirrors the image-load guard.
_guard_video_load_against_training()
# Take the GPU from chat only for a non-CPU load; a CPU load never touches GPU memory,
# so key off the device. Release stale VIDEO ownership on a CPU load (owner-guarded no-op).
# Take the GPU from chat only for a non-CPU load; a CPU load never touches GPU memory, so key off
# the device. Release stale VIDEO ownership on a CPU load (owner-guarded no-op).
device = await asyncio.to_thread(lambda: resolve_diffusion_device_target().device)
def _begin_load():
# Kicks the (slow) load onto a background thread and returns at once;
# begin_load itself validates network-free.
# Kicks the (slow) load onto a background thread and returns at once; begin_load itself validates
# network-free.
return backend.begin_load(
request.model_path,
gguf_filename = request.gguf_filename,
@ -174,10 +174,10 @@ async def load_video_model(
)
if device != "cpu":
# Register the in-flight load UNDER the arbiter lock (not after acquire_for returns):
# otherwise a competing Images/chat acquire in that gap evicts VIDEO before the load is
# marked in-flight, finds nothing to cancel, and both loaders allocate VRAM at once.
# Mirrors the images/load handoff.
# Register the in-flight load UNDER the arbiter lock (not after acquire_for returns): otherwise a
# competing Images/chat acquire in that gap evicts VIDEO before the load is marked in-flight,
# finds nothing to cancel, and both loaders allocate VRAM at once. Mirrors the images/load
# handoff.
status_dict = await asyncio.to_thread(acquire_for, VIDEO, _begin_load)
else:
await asyncio.to_thread(release, VIDEO)
@ -227,8 +227,8 @@ async def generate_video(
# Bad client input -- a 400 with the reason, not a generic 500.
raise HTTPException(status_code = 400, detail = str(exc))
except RuntimeError as exc:
# Only the not-loaded / busy sentinels are client-state (409); match exactly so an
# unrelated failure can't misroute and leak its message.
# Only the not-loaded / busy sentinels are client-state (409); match exactly so an unrelated
# failure can't misroute and leak its message.
msg = str(exc)
if msg in (VIDEO_NOT_LOADED_MSG, VIDEO_GENERATION_BUSY_MSG):
raise HTTPException(status_code = 409, detail = msg)
@ -266,8 +266,7 @@ async def unload_video_model(current_subject: str = Depends(get_current_subject)
status_dict = await asyncio.to_thread(backend.unload)
# Drop VIDEO ownership only if nothing is resident AND no new load is in flight: a concurrent
# /video/load that re-acquired VIDEO must keep ownership. The idle check and release must be
# ATOMIC (release_if): the load's register runs under the same lock, so a plain
# check-then-release could clear the newer claim. Mirrors the images route.
# ATOMIC (release_if), since the load's register runs under the same lock. Mirrors images.
await asyncio.to_thread(
release_if,
VIDEO,
@ -289,8 +288,8 @@ async def list_gallery_videos(
# Validate inside the pager so offset / limit / has_more all count over the accepted domain. A
# sidecar that parses as JSON but has a wrong value type passes the read yet fails
# GalleryVideo(**r); dropping it only after slicing let a leading bad record return an empty
# page with has_more=True, stalling infinite scroll at offset 0.
# GalleryVideo(**r); dropping it only after slicing let a leading bad record return an empty page
# with has_more=True, stalling infinite scroll at offset 0.
def _valid_gallery_video(record: dict) -> bool:
try:
GalleryVideo(**record)
@ -313,15 +312,15 @@ async def get_gallery_video_file(
):
from core.inference import video_gallery
# Ownership-gate the serve like delete/clear: resolve only a Studio-owned MP4 (readable
# sidecar), so a guessed stem for a foreign/orphan clip the listing hides can't be streamed out.
# Ownership-gate the serve like delete/clear: resolve only a Studio-owned MP4 (readable sidecar),
# so a guessed stem for a foreign/orphan clip can't be streamed out.
path = await asyncio.to_thread(video_gallery.owned_video_path, video_id)
if path is None:
raise HTTPException(status_code = 404, detail = "Video not found.")
from fastapi.responses import FileResponse
# FileResponse streams from disk and serves range requests (seek without a full fetch).
# Immutable per id, so let the browser cache it.
# FileResponse streams from disk and serves range requests (seek without a full fetch). Immutable
# per id, so let the browser cache it.
return FileResponse(
path,
media_type = "video/mp4",

5
studio/backend/tests/conftest.py
View file

@ -22,9 +22,8 @@ _backend_root = Path(__file__).resolve().parent.parent
if str(_backend_root) not in sys.path:
sys.path.insert(0, str(_backend_root))
# Let the diffusion patch backend lazily import unsloth_zoo on a CPU-only / no-GPU test
# host: unsloth_zoo runs accelerator detection at import and raises without a GPU unless
# this is set (device_type.get_device_type checks torch.cuda first, so it is a no-op on a
# Let the diffusion patch backend lazily import unsloth_zoo on a CPU-only test host: unsloth_zoo
# runs accelerator detection at import and raises without a GPU unless this is set (a no-op on a
# real GPU run). setdefault so an explicit override wins.
os.environ.setdefault("UNSLOTH_ALLOW_CPU", "1")

83
studio/backend/tests/test_cached_gguf_routes.py
View file

@ -527,8 +527,7 @@ def test_list_cached_models_skips_non_suffix_repo_when_gguf_files_exist(monkeypa
def test_list_cached_models_prefers_complete_over_larger_partial(monkeypatch, tmp_path):
# The same repo cached in two roots: a LARGER but PARTIAL copy must not shadow a SMALLER but
# COMPLETE one. The picker drops partial rows, so picking the partial winner by size alone would
# make a usable model vanish from On Device. Completeness wins over size.
# COMPLETE one, or the picker (which drops partial rows) hides a usable model.
complete = _repo(
"Org/Dup",
[_file("model.safetensors", 10_000)],
@ -559,7 +558,7 @@ def test_list_cached_models_prefers_complete_over_larger_partial(monkeypatch, tm
assert len(result["cached"]) == 1
row = result["cached"][0]
assert row["repo_id"] == "Org/Dup"
# The COMPLETE (smaller) copy won: it is not flagged partial and carries its 10_000 size.
# The COMPLETE (smaller) copy won.
assert row.get("partial") is not True
assert row["size_bytes"] == 10_000
@ -713,7 +712,6 @@ def test_list_cached_models_tags_diffusers_pipeline_as_text_to_image(monkeypatch
[
_file("model_index.json", 1_000),
_file("text_encoder/model.safetensors", 9_000),
# A complete pipeline carries its denoiser weights; without them it would be partial.
_file("transformer/diffusion_pytorch_model.safetensors", 9_000),
],
tmp_path / "models--Tongyi-MAI--Z-Image-Turbo",
@ -1137,34 +1135,26 @@ def test_legacy_delete_delegates_to_shared_service(monkeypatch):
def test_arch_to_task_hides_unsupported_diffusion_from_chat():
# Loadable diffusion archs -> the Images-picker task.
assert models_route._arch_to_task("flux") == "text-to-image"
assert models_route._arch_to_task("z_image") == "text-to-image"
assert models_route._arch_to_task("qwen_image") == "text-to-image"
# A real LLM arch stays a chat model; None passes through.
assert models_route._arch_to_task("llama") == "text-generation"
assert models_route._arch_to_task(None) is None
# Known-but-unsupported diffusion archs get a task that is NEITHER chat
# ("text-generation") NOR a loadable image task ("text-to-image"), so the chat
# picker hides them (they'd die in llama.cpp) and the Images picker leaves them
# out (they'd 400 in validate_load).
# Known-but-unsupported diffusion archs get a task that is NEITHER chat nor a loadable image
# task, so the chat picker hides them and the Images picker leaves them out.
for arch in ("sdxl", "sd1", "sd3", "lumina2", "hidream", "cosmos", "hyvid"):
task = models_route._arch_to_task(arch)
assert task == models_route._UNSUPPORTED_DIFFUSION_TASK
assert task not in ("text-generation", "text-to-image")
# A video arch with a REGISTERED VideoFamily surfaces with the Video-picker task (unsloth
# LTX-2.x GGUFs ship general.architecture "ltxv" and ltx-2 is registered).
# A video arch with a REGISTERED VideoFamily surfaces with the Video-picker task.
assert models_route._arch_to_task("ltxv") == models_route._VIDEO_GEN_TASK
assert models_route._arch_to_task("ltxv") not in ("text-generation", "text-to-image")
# A video arch that does not resolve from the bare arch alone ("wan" is ambiguous -- it
# covers both the loadable single-DiT TI2V-5B and the A14B MoE whose single file the loader
# refuses) stays unsupported when no repo/file name is available to disambiguate, rather than
# surfacing a GGUF that might 400 on load.
# A video arch that does not resolve from the bare arch alone ("wan" covers both the loadable
# TI2V-5B and the A14B MoE) stays unsupported when no name is available to disambiguate.
assert models_route._arch_to_task("wan") == models_route._UNSUPPORTED_DIFFUSION_TASK
assert models_route._arch_to_task("wan") not in ("text-generation", "text-to-image")
# With a repo/file name hint, the loadable TI2V-5B Wan GGUF resolves to the Video task (so it
# surfaces in the Video On-Device picker), while the A14B MoE (single file refused by the
# loader) stays in the unsupported bucket -- matching the loader's own name-aware detection.
# With a repo/file name hint, the loadable TI2V-5B Wan GGUF resolves to the Video task while the
# A14B MoE stays unsupported, matching the loader's own name-aware detection.
assert (
models_route._arch_to_task("wan", ("QuantStack/Wan2.2-TI2V-5B-GGUF",))
== models_route._VIDEO_GEN_TASK
@ -1177,8 +1167,8 @@ def test_arch_to_task_hides_unsupported_diffusion_from_chat():
models_route._arch_to_task("wan", ("QuantStack/Wan2.2-T2V-A14B-GGUF",))
== models_route._UNSUPPORTED_DIFFUSION_TASK
)
# Drift guard: every diffusion arch llama.cpp rejects as a chat model must be
# classified here as some non-chat task (image, video, or unsupported).
# Drift guard: every diffusion arch llama.cpp rejects as a chat model must classify here as some
# non-chat task (image, video, or unsupported).
from core.inference.llama_cpp import LlamaCppBackend
classified = (
@ -1228,8 +1218,8 @@ def _idle_diffusion_engine():
def test_delete_cached_refuses_diffusion_loaded_repo(monkeypatch):
# The cached-delete guard refuses deleting a repo the diffusion (Images) backend has loaded,
# mirroring the chat guard, so its GGUF can't be removed from under a live pipeline.
# The cached-delete guard refuses deleting a repo the diffusion (Images) backend has loaded, so
# its GGUF can't be removed from under a live pipeline.
from fastapi import HTTPException
from hub.services.models import deletion
import core.inference.diffusion_engine_router as der
@ -1256,9 +1246,7 @@ def test_delete_cached_refuses_diffusion_loaded_repo(monkeypatch):
def test_delete_cached_refuses_video_loaded_repo(monkeypatch):
# The cached-delete guard must refuse deleting a repo the Video backend has loaded (it
# shares the On-Device GGUF delete UI with chat/Images), so its GGUF can't be removed from
# under a live video pipeline -- the same invariant the sibling guards enforce.
# Same for the Video backend, which shares the On-Device GGUF delete UI with chat/Images.
from fastapi import HTTPException
from hub.services.models import deletion
import core.inference.diffusion_engine_router as der
@ -1284,10 +1272,9 @@ def test_delete_cached_refuses_video_loaded_repo(monkeypatch):
def test_delete_cached_refuses_loaded_native_companion_repo(monkeypatch):
# The native sd.cpp one-shot engine re-reads its companion VAE / text-encoder files from the
# HF cache on every generation, so deleting a companion repo (comfyanonymous/flux_text_encoders)
# while a FLUX GGUF is loaded must be refused. The loaded main repo_id does not match the
# companion, so the guard relies on loaded_repo_ids() to cover the committed companions.
# The native sd.cpp one-shot engine re-reads its companion VAE / text-encoder files every
# generation, so deleting a companion repo while a FLUX GGUF is loaded must be refused. The
# loaded repo_id does not match the companion, so the guard relies on loaded_repo_ids().
from fastapi import HTTPException
from hub.services.models import deletion
import core.inference.diffusion_engine_router as der
@ -1318,9 +1305,8 @@ def test_delete_cached_refuses_loaded_native_companion_repo(monkeypatch):
def test_delete_cached_refuses_repo_a_diffusion_load_is_downloading(monkeypatch):
# status().loaded is still False while a background Images load DOWNLOADS the repo (or its
# companion base), but deleting then would remove blobs from under the in-flight
# download/assembly, so loading_repo_ids() must refuse with the wait-for-it detail.
# status().loaded is still False while a background Images load DOWNLOADS the repo, but deleting
# then would remove blobs from under it, so loading_repo_ids() must refuse.
from fastapi import HTTPException
from hub.services.models import deletion
import core.inference.diffusion_engine_router as der
@ -1347,11 +1333,8 @@ def test_delete_cached_refuses_repo_a_diffusion_load_is_downloading(monkeypatch)
def test_delete_cached_allows_sibling_of_loaded_diffusion_repo(monkeypatch):
# A loaded Images repo must not block deleting a DIFFERENT cached repo that merely shares a
# name prefix. Qwen/Qwen-Image and Qwen/Qwen-Image-2512 are both real catalog artifacts, so
# with Qwen-Image-2512 loaded, deleting the sibling Qwen-Image is a supported operation. The
# guard is `/`-boundary aware, so it refuses only the loaded repo (or a file within it), not
# a prefix sibling.
# A loaded Images repo must not block deleting a DIFFERENT cached repo that merely shares a name
# prefix (Qwen/Qwen-Image vs Qwen/Qwen-Image-2512). The guard is `/`-boundary aware.
from fastapi import HTTPException
from hub.services.models import deletion
import core.inference.diffusion_engine_router as der
@ -1378,8 +1361,7 @@ def test_delete_cached_allows_sibling_of_loaded_diffusion_repo(monkeypatch):
},
)
# The sibling repo clears every guard and reaches the delete -- it is NOT refused with the
# 400 "Unload" that an un-delimited prefix match would have produced.
# The sibling repo clears every guard and reaches the delete.
result = asyncio.run(deletion.delete_cached_model_response("Qwen/Qwen-Image"))
assert result == {"status": "deleted", "repo_id": "Qwen/Qwen-Image"}
@ -1393,10 +1375,9 @@ def test_delete_cached_allows_sibling_of_loaded_diffusion_repo(monkeypatch):
def test_cached_repo_partial_scopes_probe_to_snapshot_dir(monkeypatch):
# The partial probe must be scoped to the snapshot row being listed. Unscoped, the scan
# spans every HF cache root, so a stale .incomplete copy in one root would flag a complete
# copy living in another root as partial and hide the usable model from the picker. Verify
# _cached_repo_partial forwards the snapshot dir (matching the sibling inventory paths).
# The partial probe must be scoped to the snapshot row being listed: unscoped, a stale
# .incomplete copy in one cache root would flag a complete copy in another as partial and hide
# the usable model.
import hub.utils.inventory_scan as scan
calls = []
@ -1414,7 +1395,6 @@ def test_cached_repo_partial_scopes_probe_to_snapshot_dir(monkeypatch):
assert models_route._cached_repo_partial("Org/Repo", snapshot_dir) is False
assert calls == [("model", "Org/Repo", snapshot_dir)]
# When that specific snapshot is partial, the row is flagged.
monkeypatch.setattr(scan, "is_snapshot_partial", lambda *a, **k: True)
assert models_route._cached_repo_partial("Org/Repo", snapshot_dir) is True
@ -1427,10 +1407,9 @@ def test_cached_repo_partial_scopes_probe_to_snapshot_dir(monkeypatch):
def test_repo_has_pipeline_index_requires_root_model_index(tmp_path):
# Only a ROOT model_index.json makes a repo pipeline-loadable: from_pretrained
# reads the repo root, so a nested subdir/model_index.json must NOT clear the
# single_file flag. CachedFileInfo.file_name is the basename, so the helper has
# to scope by file_path/snapshot_path -- a name-only match would claim both.
# Only a ROOT model_index.json makes a repo pipeline-loadable, so a nested subdir one must NOT
# clear the single_file flag. CachedFileInfo.file_name is the basename, so the helper scopes by
# snapshot path -- a name-only match would claim both.
snap = tmp_path / "snapshots" / "abc"
nested = SimpleNamespace(
file_name = "model_index.json",
@ -1454,10 +1433,8 @@ def test_repo_has_pipeline_index_requires_root_model_index(tmp_path):
def test_list_cached_models_flags_single_file_diffusion_repos(monkeypatch, tmp_path):
# A diffusion-tagged repo with NO top-level model_index.json is a single-file
# checkpoint: the task pickers must not offer it as a pipeline load (from_pretrained
# fails on it), so the row carries single_file=True. A full pipeline repo (has
# model_index.json) and a chat repo (task None) carry no flag.
# A diffusion-tagged repo with NO top-level model_index.json is a single-file checkpoint, so it
# carries single_file=True; a full pipeline repo and a chat repo carry no flag.
single = _repo(
"unsloth/Qwen-Image-fp8-single",
[_file("qwen-image-fp8.safetensors", 10_000)],

8
studio/backend/tests/test_diffusion_arch_patches.py
View file

@ -251,8 +251,8 @@ def test_krea2_forward_matches_stock():
num_kv_heads = H // 2,
norm_eps = 1e-6,
).eval()
# Give the zero-init modulation table real values so all six scale/shift/gate
# branches contribute to the output.
# Give the zero-init modulation table real values so all six scale/shift/gate branches
# contribute to the output.
with torch.no_grad():
blk.scale_shift_table.normal_()
# A [text + 2x2 image grid] sequence with the real rotary embed (axes sum to head_dim).
@ -298,8 +298,8 @@ def test_kill_switch(monkeypatch):
def test_body_drift_guard_skips_changed_block(monkeypatch):
# If a resolver's body-check fails (diffusers changed the lines we rewrite), that patch
# is skipped. Force the qwen resolver to see a drifted body.
# A resolver whose body-check fails (diffusers changed the lines we rewrite) is skipped. Force
# the qwen resolver to see a drifted body.
monkeypatch.setattr(ap, "_body_has", lambda fn, *needles: False)
assert ap.install_arch_patches() == 0
assert not ap.is_installed()

88
studio/backend/tests/test_diffusion_attention.py
View file

@ -38,13 +38,13 @@ def test_normalize_defaults_and_aliases():
def test_normalize_rejects_unknown():
with pytest.raises(ValueError):
normalize_attention_backend("bogus")
# dashes are no longer silently rewritten to underscores -> a dashed alias is rejected.
# dashes are no longer silently rewritten to underscores, so a dashed alias is rejected.
with pytest.raises(ValueError):
normalize_attention_backend("flash-3")
def test_sdpa_alias_maps_to_native():
# sdpa is an alias for native -> nothing to set on the dispatcher.
# sdpa is an alias for native, so nothing to set on the dispatcher.
assert select_attention_backend(_target(), "sdpa", speed_active = True) is None
@ -56,15 +56,15 @@ def test_auto_upgrades_to_cudnn_on_nvidia_when_speed_active(monkeypatch):
def test_auto_does_not_pin_cudnn_below_sm80(monkeypatch):
# cuDNN fused SDPA fails at run time on pre-SM80 (T4 SM75 / V100 SM70); auto must stay
# on the native default there rather than pin a backend that crashes on first generation.
# cuDNN fused SDPA fails at run time on pre-SM80 (T4 / V100), so auto must stay native there
# rather than pin a backend that crashes on first generation.
monkeypatch.setattr(att, "_is_cuda_nvidia", lambda target: True)
monkeypatch.setattr(att, "_cuda_capability", lambda: (7, 5)) # Turing T4
assert select_attention_backend(_target(), "auto", speed_active = True) is None
def test_auto_stays_native_when_speed_off(monkeypatch):
# off must stay bit-identical -> no backend change even on NVIDIA.
# off must stay bit-identical, so no backend change even on NVIDIA.
monkeypatch.setattr(att, "_is_cuda_nvidia", lambda target: True)
assert select_attention_backend(_target(), "auto", speed_active = False) is None
@ -84,8 +84,8 @@ def test_explicit_backend_honored_regardless_of_speed(monkeypatch):
def test_explicit_backend_dropped_off_nvidia_cuda(monkeypatch):
# Explicit cuDNN/flash/sage on ROCm / MPS / CPU passes diffusers' set-time check
# and crashes at the first generation, so selection drops to the native default.
# Explicit cuDNN/flash/sage on ROCm / MPS / CPU passes diffusers' set-time check and crashes at
# the first generation, so selection drops to the native default.
monkeypatch.setattr(att, "_is_cuda_nvidia", lambda target: False)
monkeypatch.setattr(att, "_cuda_capability", lambda: (10, 0))
for alias in ("sage", "flash", "flash4", "cudnn"):
@ -93,8 +93,8 @@ def test_explicit_backend_dropped_off_nvidia_cuda(monkeypatch):
def test_aiter_honored_on_rocm(monkeypatch):
# AITER is the AMD ROCm kernel; on a ROCm CUDA target it must be honored, not dropped by
# the NVIDIA-only guard -- it is the one explicit backend that only ever works on ROCm.
# AITER is the AMD ROCm kernel, so on a ROCm CUDA target it must be honored, not dropped by the
# NVIDIA-only guard.
monkeypatch.setattr(att, "_is_cuda_nvidia", lambda target: False) # hip build
assert select_attention_backend(_target(), "aiter", speed_active = False) == "aiter"
@ -108,7 +108,7 @@ def test_aiter_dropped_off_rocm(monkeypatch):
def test_explicit_native_returns_none():
# native is the default -> nothing to set.
# native is the default, so nothing to set.
assert select_attention_backend(_target(), "native", speed_active = True) is None
@ -126,14 +126,14 @@ def test_flash4_dropped_below_blackwell(monkeypatch):
def test_arch_gate_does_not_block_when_capability_unknown(monkeypatch):
# Unknown capability (e.g. no CUDA) must not block -> diffusers' set-time check still guards.
# Unknown capability must not block; diffusers' set-time check still guards.
monkeypatch.setattr(att, "_cuda_capability", lambda: None)
assert select_attention_backend(_target(), "flash4", speed_active = False) == "flash_4_hub"
def test_flash3_dropped_on_blackwell(monkeypatch):
# FlashAttention 3 is a Hopper-SM90 rewrite with no Blackwell kernel: an explicit
# flash3 on a B200 (SM100) must drop to native rather than set fine then crash.
# FlashAttention 3 is a Hopper-SM90 rewrite with no Blackwell kernel, so an explicit flash3 on a
# B200 must drop to native rather than set fine then crash.
monkeypatch.setattr(att, "_cuda_capability", lambda: (10, 0))
assert select_attention_backend(_target(), "flash3", speed_active = False) is None
# FA4 is still honored on Blackwell.
@ -144,8 +144,8 @@ def test_flash3_dropped_on_blackwell(monkeypatch):
def test_explicit_cudnn_dropped_below_sm80(monkeypatch):
# An explicit cuDNN request on pre-Ampere (T4 SM75 / V100 SM70) must drop to native,
# not set fine and crash at first generation -- the same gate the auto path applies.
# An explicit cuDNN request on pre-Ampere must drop to native, not set fine and crash at first
# generation -- the same gate the auto path applies.
monkeypatch.setattr(att, "_cuda_capability", lambda: (7, 5))
assert select_attention_backend(_target(), "cudnn", speed_active = False) is None
# Ampere+ still honors it.
@ -170,7 +170,7 @@ def _pipe(transformer):
def test_apply_none_leaves_native_when_global_already_native(monkeypatch):
# Global already native -> no redundant set call, returns None.
# Global already native, so no redundant set call.
monkeypatch.setattr(att, "_active_attention_backend", lambda: "native")
t = _FakeTransformer()
assert apply_attention_backend(_pipe(t), None) is None
@ -178,8 +178,8 @@ def test_apply_none_leaves_native_when_global_already_native(monkeypatch):
def test_apply_none_restores_native_when_global_polluted(monkeypatch):
# A previous load pinned cuDNN process-wide; a native load must reset it so it can't
# silently inherit cuDNN (the bit-identical/off guarantee).
# A previous load pinned cuDNN process-wide; a native load must reset it so it can't silently
# inherit cuDNN (the bit-identical/off guarantee).
monkeypatch.setattr(att, "_active_attention_backend", lambda: "_native_cudnn")
t = _FakeTransformer()
assert apply_attention_backend(_pipe(t), None) is None
@ -193,9 +193,8 @@ def test_apply_sets_backend():
def test_apply_sets_backend_on_both_dits():
# A dual-DiT family (Ideogram) runs transformer + unconditional_transformer each step, so the
# backend must be set on BOTH; otherwise the second DiT keeps the native default while status
# reports the requested kernel as engaged.
# A dual-DiT family (Ideogram) runs both DiTs each step, so the backend must be set on BOTH, else
# the second keeps native while status reports the requested kernel as engaged.
t1, t2 = _FakeTransformer(), _FakeTransformer()
pipe = types.SimpleNamespace(transformer = t1, unconditional_transformer = t2)
engaged = apply_attention_backend(pipe, "_native_cudnn")
@ -204,7 +203,7 @@ def test_apply_sets_backend_on_both_dits():
def test_apply_falls_back_on_unavailable_kernel(monkeypatch):
# an unavailable kernel must not fail the load -> returns None (diffusers default).
# An unavailable kernel must not fail the load: returns None (diffusers default).
monkeypatch.setattr(att, "_active_attention_backend", lambda: "native")
t = _FakeTransformer(fail = True)
assert apply_attention_backend(_pipe(t), "sage") is None
@ -234,9 +233,8 @@ def test_apply_handles_missing_method():
def test_apply_resets_global_registry_after_success(monkeypatch):
# After a successful per-transformer set, the process-wide registry must be reset to
# native so a later component (unconfigured processors) can't inherit this kernel --
# while the transformer's own backend stays the engaged one.
# After a successful per-transformer set, the process-wide registry must be reset to native so a
# later component can't inherit this kernel, while the transformer keeps the engaged one.
called = {"reset": False}
monkeypatch.setattr(
att, "_reset_global_backend_to_native", lambda logger: called.__setitem__("reset", True)
@ -248,8 +246,8 @@ def test_apply_resets_global_registry_after_success(monkeypatch):
def test_active_attention_backend_reads_tuple_return():
# get_active_backend() returns a (AttentionBackendName, fn) tuple; the helper must read
# the name's .value, not stringify the tuple (which never compares equal to a name).
# get_active_backend() returns a (AttentionBackendName, fn) tuple; the helper must read the
# name's .value, not stringify the tuple (which never compares equal to a name).
pytest.importorskip("diffusers")
from diffusers.models.attention_dispatch import (
AttentionBackendName,
@ -263,13 +261,10 @@ def test_active_attention_backend_reads_tuple_return():
# ── on-demand wheel-only install of optional kernels ─────────────────────────────
@pytest.fixture(autouse = True)
def _no_real_installs(monkeypatch):
# Unit tests must never shell out to pip: the apply path probes installable
# backends (sage/flash*), so hard-disable the gate; install tests re-enable it
# with a stubbed subprocess.
# Unit tests must never shell out to pip: the apply path probes installable backends, so
# hard-disable the gate; install tests re-enable it with a stubbed subprocess.
monkeypatch.setenv("UNSLOTH_DIFFUSION_ATTENTION_INSTALL", "0")
# The install once-per-process memo is module state; clear it so each test starts
# with a fresh "not yet attempted" set (otherwise an earlier test's attempt would
# make a later install a no-op).
# The install once-per-process memo is module state; clear it so each test starts fresh.
att._INSTALL_ATTEMPTED.clear()
@ -321,10 +316,8 @@ def test_install_runs_wheel_only_for_missing_kernel(monkeypatch):
def test_install_uses_no_deps_to_protect_core_deps(monkeypatch):
# A kernel add-on (xformers/flash-attn) pins an exact torch, so a normal install would
# upgrade/replace the running torch/triton. --no-deps installs only the kernel wheel;
# an ABI-incompatible one fails to import and falls back to native rather than clobbering
# the environment's core deps.
# A kernel add-on pins an exact torch, so a normal install would replace the running torch/triton.
# --no-deps installs only the kernel wheel; an ABI-incompatible one just fails to import.
monkeypatch.setenv("UNSLOTH_DIFFUSION_ATTENTION_INSTALL", "auto")
import importlib.util
@ -337,10 +330,9 @@ def test_install_uses_no_deps_to_protect_core_deps(monkeypatch):
def test_failed_install_not_retried_in_same_process(monkeypatch):
# The loader pre-installs the kernel OUTSIDE its locks and then re-resolves the same
# backend under _generate_lock; if the pre-install failed (no wheel / offline) the
# in-lock apply path must NOT re-run pip (a second up-to-600s install holding the load
# lock blocks unload/cancel). The once-per-process memo makes the retry a no-op.
# The loader pre-installs the kernel OUTSIDE its locks and re-resolves under _generate_lock; if
# the pre-install failed the in-lock apply must NOT re-run pip (a second 600s install would block
# unload/cancel). The once-per-process memo makes the retry a no-op.
monkeypatch.setenv("UNSLOTH_DIFFUSION_ATTENTION_INSTALL", "auto")
import importlib.util
import subprocess as sp
@ -360,9 +352,8 @@ def test_failed_install_not_retried_in_same_process(monkeypatch):
def test_install_invalidates_import_caches_on_success(monkeypatch):
# A wheel written to site-packages after the finder cached that directory can be
# missed by the very next import, so a successful install must invalidate the caches
# (otherwise set_attention_backend imports the missing package and falls back).
# A wheel written to site-packages after the finder cached that directory can be missed by the
# very next import, so a successful install must invalidate the caches.
monkeypatch.setenv("UNSLOTH_DIFFUSION_ATTENTION_INSTALL", "auto")
import importlib
import importlib.util
@ -404,8 +395,8 @@ def test_install_never_attempted_for_builtin_backends(monkeypatch):
def test_install_failure_logs_pip_stderr(monkeypatch):
# A CalledProcessError's str() hides the pip reason; the warning must surface the
# captured stderr (decoding bytes) so a fallback to native is diagnosable.
# A CalledProcessError's str() hides the pip reason; the warning must surface the captured stderr
# so a fallback to native is diagnosable.
monkeypatch.setenv("UNSLOTH_DIFFUSION_ATTENTION_INSTALL", "auto")
import importlib.util
import subprocess as sp
@ -433,9 +424,8 @@ def test_install_failure_logs_pip_stderr(monkeypatch):
def test_install_failure_falls_back_to_native(monkeypatch):
# pip failing (no wheel for this platform) must not break the load: the apply
# path proceeds, set_attention_backend raises on the missing package, and the
# dispatcher is restored to native -- same contract as before the hook.
# pip failing (no wheel for this platform) must not break the load: apply proceeds,
# set_attention_backend raises on the missing package, and the dispatcher is restored to native.
monkeypatch.setenv("UNSLOTH_DIFFUSION_ATTENTION_INSTALL", "auto")
import importlib.util
import subprocess as sp

29
studio/backend/tests/test_diffusion_auto_policy.py
View file

@ -47,8 +47,8 @@ def test_family_table_unknown_family_returns_none():
def test_base_repo_override_wins_over_the_family_default():
# flux.2-klein's family default is the 4B base; loading the 9B GGUF passes the 9B
# base repo, whose transformer is more than twice the size.
# flux.2-klein's family default is the 4B base; loading the 9B GGUF passes the 9B base repo,
# whose transformer is more than twice the size.
default = family_bf16_components_gb(_fam("flux.2-klein"))
nine_b = family_bf16_components_gb(
_fam("flux.2-klein"), base_repo = "black-forest-labs/FLUX.2-klein-9B"
@ -69,8 +69,8 @@ def test_estimate_int8_steady_is_roughly_half_bf16():
def test_estimate_prequant_transient_equals_steady():
# A pre-quantized checkpoint loads via the meta device: dense bf16 never lands on
# the GPU, so the build peak IS the quantised size.
# A pre-quantized checkpoint loads via the meta device: dense bf16 never lands on the GPU, so the
# build peak IS the quantised size.
est = estimate_dense_quant(_fam("z-image"), "int8", prequant_available = True)
assert est is not None
assert est.transient_transformer_mib == est.steady_transformer_mib
@ -110,9 +110,8 @@ def _patch_selector(
"resolve_prequant_source",
lambda fam, s, path_override = None, base_repo = None: prequant,
)
# Neutralize the cache-disk gate by default so resolution tests are independent of the
# runner's free space (a small CI disk otherwise drops the candidate). The two disk-gate
# tests re-patch this after calling the helper to exercise the gate explicitly.
# Neutralize the cache-disk gate by default so resolution tests don't depend on the runner's free
# space. The two disk-gate tests re-patch this to exercise it explicitly.
monkeypatch.setattr(ap, "_hf_cache_free_mib", lambda: None)
@ -141,8 +140,8 @@ def test_candidate_none_when_no_scheme_resolves(monkeypatch):
def test_candidate_disk_gate_skips_when_cache_disk_low(monkeypatch):
# The dense artifact may be a multi-GB download; a nearly-full model-cache disk
# drops the candidate (the loader then keeps the GGUF build).
# The dense artifact may be a multi-GB download; a nearly-full model-cache disk drops the
# candidate (the loader then keeps the GGUF build).
import core.inference.diffusion_auto_policy as ap
_patch_selector(monkeypatch, scheme = "int8")
@ -164,7 +163,7 @@ def test_candidate_disk_gate_unprobeable_disk_passes(monkeypatch):
def test_candidate_none_for_an_unlisted_family(monkeypatch):
# No size entry -> no basis to re-plan; the loader keeps today's resident-only gate.
# No size entry means no basis to re-plan; the loader keeps today's resident-only gate.
_patch_selector(monkeypatch)
assert (
resolve_dense_quant_candidate(fam = _fam("not-a-family"), target = object(), requested = "auto")
@ -173,7 +172,7 @@ def test_candidate_none_for_an_unlisted_family(monkeypatch):
def test_candidate_uses_prequant_transient_when_available(monkeypatch):
# A hosted-repo prequant source (kind="repo") is available without a local-path check.
# A hosted-repo prequant source is available without a local-path check.
_patch_selector(monkeypatch, prequant = SimpleNamespace(kind = "repo", location = "org/int8"))
est = resolve_dense_quant_candidate(fam = _fam("z-image"), target = object(), requested = "int8")
assert est is not None and est.prequant is True
@ -186,11 +185,9 @@ def _cuda_target():
def test_quant_candidate_fits_resident_where_gguf_plan_offloads():
# The ordering-fix mechanism, on a 32 GiB consumer card (RTX 5090 class): the user
# picked a LARGE GGUF (the BF16 file), so the file-size plan forces offload -- but
# the dense-quant candidate is far smaller (int8 prequant of z-image: the transient
# IS the quantised size), and re-planning against the candidate keeps everything
# resident. Before the fix the loader never attempted the fast path here.
# The ordering fix on a 32 GiB consumer card: the user picked a LARGE (BF16) GGUF, so the
# file-size plan forces offload -- but the dense-quant candidate is far smaller, and re-planning
# against it keeps everything resident. Before the fix the loader never attempted the fast path.
memory = DeviceMemory("cuda", "cuda", "discrete_vram", 30000, 32768)
z_bf16_gguf_mib = int(12.3 * ap._MIB_PER_GB * 1.05) # BF16 GGUF resident estimate
companions_mib = 2600 # fp8-quantised text encoders + VAE

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

File diff suppressed because it is too large Load diff

190
studio/backend/tests/test_diffusion_base_precision.py
View file

@ -28,15 +28,15 @@ from core.training.diffusion_train_common import (
)
from models.training import DiffusionTrainingStartRequest
# A dense (non-prequant) DiT base and a prequant bnb-4bit base. Both resolve a trainer
# family from their names alone, so normalized() runs without a network call.
# A dense (non-prequant) DiT base and a prequant bnb-4bit base. Both resolve a trainer family from
# their names alone, so normalized() runs without a network call.
_FLUX_DENSE = "black-forest-labs/FLUX.1-dev"
_Z_PREQUANT = "unsloth/Z-Image-Turbo-unsloth-bnb-4bit"
# An SDXL base whose name LOOKS prequant (bnb-4bit): SDXL ignores base_precision, so the
# dense-mode gates must not fire for it even with a dense mode + fp16 compute.
# An SDXL base whose name LOOKS prequant: SDXL ignores base_precision, so the dense-mode gates
# must not fire for it even with a dense mode + fp16 compute.
_SDXL_PREQUANT_NAME = "some/sdxl-model-bnb-4bit"
# A dense Qwen-Image base: its DiT is corrupted by fp8 (activation outliers), so fp8 is
# denied for training the same way the inference path denies it.
# A dense Qwen-Image base: its DiT is corrupted by fp8, so fp8 is denied for training the same way
# the inference path denies it.
_QWEN_DENSE = "Qwen/Qwen-Image"
@ -53,13 +53,13 @@ def test_base_precision_validation():
with pytest.raises(ValueError, match = "base_precision"):
_cfg(base_precision = "banana").normalized()
# A dense mode is case/space-insensitive and stored lowered: " FP8 " on a dense base
# with bf16 compute normalizes cleanly to "fp8".
# A dense mode is case/space-insensitive and stored lowered: " FP8 " on a dense base with bf16
# compute normalizes to "fp8".
norm = _cfg(base_precision = " FP8 ", mixed_precision = "bf16").normalized()
assert norm.base_precision == "fp8"
# A dense mode against a prequant (bnb-4bit) base is refused: the repo already ships a
# 4-bit transformer and cannot serve the dense precisions.
# A dense mode against a prequant (bnb-4bit) base is refused: the repo already ships a 4-bit
# transformer.
with pytest.raises(ValueError, match = "dense base repo"):
_cfg(base_model = _Z_PREQUANT, base_precision = "bf16").normalized()
@ -67,19 +67,18 @@ def test_base_precision_validation():
with pytest.raises(ValueError, match = "bf16 compute"):
_cfg(base_precision = "int8", mixed_precision = "fp16").normalized()
# "auto" is ACCEPTED by normalized() even on a prequant base: the concrete mode is
# resolved at runtime against the live GPU, not at config validation.
# "auto" is ACCEPTED even on a prequant base: the concrete mode is resolved at runtime against
# the live GPU, not at config validation.
assert _cfg(base_model = _Z_PREQUANT, base_precision = "auto").normalized().base_precision == "auto"
def test_base_precision_denies_fp8_for_corrupted_family():
# fp8 corrupts the Qwen-Image DiT (activation outliers exceed fp8's range), so a dense
# Qwen base with base_precision="fp8" is refused up front -- mirroring the inference deny.
# fp8 corrupts the Qwen-Image DiT, so a dense Qwen base with base_precision="fp8" is refused up
# front, mirroring the inference deny.
with pytest.raises(ValueError, match = "fp8"):
_cfg(base_model = _QWEN_DENSE, base_precision = "fp8", mixed_precision = "bf16").normalized()
# The deny is fp8-specific: int8 (per-token, unaffected) and the other dense modes stay
# allowed for the same Qwen base.
# The deny is fp8-specific: int8 and the other dense modes stay allowed for the same Qwen base.
for mode in ("nf4", "bf16", "int8", "auto"):
norm = _cfg(
base_model = _QWEN_DENSE, base_precision = mode, mixed_precision = "bf16"
@ -94,15 +93,13 @@ def test_base_precision_denies_fp8_for_corrupted_family():
def test_family_train_infos_drops_denied_fp8_for_qwen(monkeypatch):
# /info advertises the machine's DiT modes per family, but a family whose DiT the mode
# corrupts must not offer it: with fp8 in the machine list, Qwen-Image drops fp8 while
# FLUX keeps it, so the UI never surfaces a mode normalized() would reject.
# /info advertises the machine's DiT modes per family, but a family whose DiT the mode corrupts
# must not offer it, so the UI never surfaces a mode normalized() would reject.
monkeypatch.setattr(
common, "train_precision_modes", lambda: (["nf4", "bf16", "int8", "fp8", "auto"], "auto")
)
# family_train_infos reads the live GPU via bf16_unsupported_reason; pin it to "bf16 OK" so
# this positive-path assertion is deterministic across GPU types (a non-bf16 CUDA box would
# otherwise empty every DiT family's modes). The empty-on-non-bf16 path is covered separately.
# family_train_infos reads the live GPU via bf16_unsupported_reason; pin it so this assertion is
# deterministic across GPU types. The empty-on-non-bf16 path is covered separately.
monkeypatch.setattr(common, "bf16_unsupported_reason", lambda name: None)
infos = {i["name"]: i for i in common.family_train_infos()}
assert "fp8" not in infos["qwen-image"]["precision_modes"]
@ -111,9 +108,8 @@ def test_family_train_infos_drops_denied_fp8_for_qwen(monkeypatch):
def test_resolve_base_precision_explicit_int8_gates_on_torchao(monkeypatch):
# Explicit int8 has no runtime fallback, so a missing/stub torchao must fail fast here
# rather than load dense with compile disabled. Gate the explicit request the same way
# auto + /info already gate it.
# Explicit int8 has no runtime fallback, so a missing/stub torchao must fail fast here rather than
# load dense with compile disabled -- the same gate auto + /info already apply.
spec = dit._SPECS["flux.1"]
cfg = _cfg(base_precision = "int8")
@ -125,8 +121,8 @@ def test_resolve_base_precision_explicit_int8_gates_on_torchao(monkeypatch):
monkeypatch.setattr(dit, "has_functional_torchao", lambda: True)
assert dit._resolve_base_precision(cfg, spec, "cuda") == "int8"
# The gate is int8-specific: explicit bf16/fp8 pass through regardless of torchao (fp8 has
# its own graceful fallback; bf16 needs no torchao).
# The gate is int8-specific: explicit bf16/fp8 pass through regardless of torchao (fp8 has its own
# fallback; bf16 needs no torchao).
monkeypatch.setattr(dit, "has_functional_torchao", lambda: False)
assert dit._resolve_base_precision(_cfg(base_precision = "bf16"), spec, "cuda") == "bf16"
assert dit._resolve_base_precision(_cfg(base_precision = "fp8"), spec, "cuda") == "fp8"
@ -142,9 +138,8 @@ def test_bf16_unsupported_reason(monkeypatch):
assert bf16_unsupported_reason("sdxl") is None
assert bf16_unsupported_reason("") is None
# A DiT family on a pre-Ampere CUDA GPU -> a clear reason. Pre-Ampere cards EMULATE bf16 and
# report is_bf16_supported() True, so the gate is native compute capability (major >= 8), not
# is_bf16_supported() -- otherwise the emulation case would slip through and evict-then-fail.
# A DiT family on a pre-Ampere CUDA GPU gives a clear reason. Pre-Ampere cards EMULATE bf16 and
# report is_bf16_supported() True, so the gate is native compute capability (major >= 8).
monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
monkeypatch.setattr(
torch.cuda, "is_bf16_supported", lambda *a, **k: True
@ -156,15 +151,14 @@ def test_bf16_unsupported_reason(monkeypatch):
monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *a, **k: (8, 0))
assert bf16_unsupported_reason("qwen-image") is None
# A CPU-only host (fp32 fallback for import/unit tests) -> no reason even for a DiT family.
# A CPU-only host (fp32 fallback for import/unit tests) gives no reason even for a DiT family.
monkeypatch.setattr(torch.cuda, "is_available", lambda: False)
assert bf16_unsupported_reason("z-image") is None
def test_native_bf16_supported_gates_on_capability(monkeypatch):
# Native bf16 is gated by compute capability (Ampere major >= 8), NOT is_bf16_supported(),
# which defaults to counting pre-Ampere emulation. A Turing card that emulates bf16 is
# correctly reported unsupported; Ampere+ is supported; a CPU-only host is always False.
# Native bf16 is gated by compute capability (Ampere major >= 8), NOT is_bf16_supported(), which
# counts pre-Ampere emulation. A Turing card that emulates bf16 is correctly reported unsupported.
import torch
from core.training.diffusion_train_common import native_bf16_supported
@ -183,26 +177,24 @@ def test_native_bf16_supported_gates_on_capability(monkeypatch):
def test_training_precision_preflight_error(monkeypatch):
# The start route calls this BEFORE evicting resident GPU workloads: it folds the bf16-GPU
# requirement together with the explicit-int8 torchao requirement, so both fail fast instead
# of only surfacing in the trainer child after the GPU has already been freed.
# requirement together with the explicit-int8 torchao requirement, so both fail fast instead of
# surfacing in the trainer child after the GPU has been freed.
import torch
from core.training.diffusion_train_common import training_precision_preflight_error
# Present a NATIVE bf16-capable CUDA GPU (Ampere+, cap major >= 8) so the int8 gate (not the
# bf16 gate) is what we exercise. bf16 is gated by capability, not is_bf16_supported() (which
# counts pre-Ampere emulation).
# Present a NATIVE bf16-capable CUDA GPU so the int8 gate, not the bf16 gate, is exercised.
monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
monkeypatch.setattr(torch.cuda, "is_bf16_supported", lambda *a, **k: True)
monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *a, **k: (8, 6))
# The bf16 gate takes precedence: a pre-Ampere GPU (emulated bf16) rejects any DiT precision.
# The bf16 gate takes precedence: a pre-Ampere GPU rejects any DiT precision.
monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *a, **k: (7, 5))
assert "bfloat16" in (training_precision_preflight_error("flux.1", "int8") or "")
monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *a, **k: (8, 6))
# Explicit int8 on a DiT family with a NON-functional torchao -> a clear int8 reason
# (its _int8_quantize_base has no fallback, so the child would otherwise raise post-eviction).
# Explicit int8 on a DiT family with a NON-functional torchao gives a clear int8 reason (its
# _int8_quantize_base has no fallback, so the child would raise post-eviction).
monkeypatch.setattr(common, "has_functional_torchao", lambda: False)
reason = training_precision_preflight_error("qwen-image", "int8")
assert reason is not None and "int8" in reason and "torchao" in reason
@ -211,8 +203,8 @@ def test_training_precision_preflight_error(monkeypatch):
monkeypatch.setattr(common, "has_functional_torchao", lambda: True)
assert training_precision_preflight_error("qwen-image", "int8") is None
# With a broken torchao, only EXPLICIT int8 is gated -- nf4/bf16/auto pass, and the int8
# gate never applies to a non-DiT (SDXL) or unknown family.
# With a broken torchao only EXPLICIT int8 is gated: nf4/bf16/auto pass, and the gate never
# applies to a non-DiT (SDXL) or unknown family.
monkeypatch.setattr(common, "has_functional_torchao", lambda: False)
assert training_precision_preflight_error("flux.1", "nf4") is None
assert training_precision_preflight_error("flux.1", "auto") is None
@ -220,8 +212,7 @@ def test_training_precision_preflight_error(monkeypatch):
assert training_precision_preflight_error("", "int8") is None
# On a host with NO accelerator every DiT precision is rejected up front rather than after
# eviction -- nf4 included, since its 4-bit load goes through bitsandbytes, which needs a GPU.
# SDXL (its own fp32-on-CPU path) still passes.
# eviction -- nf4 included, since its 4-bit load goes through bitsandbytes. SDXL still passes.
monkeypatch.setattr(common, "has_functional_torchao", lambda: True)
monkeypatch.setattr(torch.cuda, "is_available", lambda: False)
monkeypatch.setattr(torch.xpu, "is_available", lambda: False)
@ -242,10 +233,8 @@ def test_training_precision_preflight_error(monkeypatch):
assert training_precision_preflight_error("flux.1", "bf16") is not None
monkeypatch.setattr(torch.xpu, "is_available", lambda: False)
# mxfp8 needs a Blackwell (sm100+) GPU: its MX GEMM has no kernel below sm100 and would raise at
# the first training step, AFTER a full dense-transformer load. On a CUDA GPU that is older than
# Blackwell the preflight rejects mxfp8 UP FRONT (mirroring _resolve_base_precision) so eviction
# is skipped; other dense precisions on the same GPU still pass.
# mxfp8 needs a Blackwell (sm100+) GPU: below sm100 its MX GEMM raises at the first training step,
# AFTER a full dense load, so the preflight rejects it UP FRONT. Other dense precisions still pass.
monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *a, **k: (9, 0))
reason = training_precision_preflight_error("flux.1", "mxfp8")
@ -262,7 +251,7 @@ def test_training_precision_preflight_error(monkeypatch):
def test_family_train_infos_empties_dit_modes_on_non_bf16(monkeypatch):
# On a non-bf16 GPU the start route rejects EVERY DiT family (even nf4), so /info must not
# advertise a DiT precision option that always 400s: the modes empty, the reason surfaces in
# vram_note, compile is off, and the recommendation degrades to nf4. SDXL (non-DiT) is exempt.
# vram_note, compile is off, and the recommendation degrades to nf4. SDXL is exempt.
from core.training.diffusion_train_common import _DIT_TRAIN_FAMILIES, family_train_infos
monkeypatch.setattr(common, "bf16_unsupported_reason", lambda name: "no bfloat16 on this GPU")
@ -280,15 +269,13 @@ def test_family_train_infos_empties_dit_modes_on_non_bf16(monkeypatch):
def test_base_precision_gates_skip_sdxl():
# SDXL ignores base_precision, so the dense-mode gates (prequant base / non-bf16 compute)
# must not fire for it: a prequant-looking SDXL name with base_precision="bf16" does not
# raise, and the mode is still stored lowered.
# SDXL ignores base_precision, so the dense-mode gates must not fire for it: a prequant-looking
# SDXL name with base_precision="bf16" does not raise, and the mode is still stored lowered.
norm = _cfg(base_model = _SDXL_PREQUANT_NAME, base_precision = "bf16").normalized()
assert norm.resolved_family == "sdxl"
assert norm.base_precision == "bf16"
# The non-bf16-compute gate is also skipped for SDXL (fp16 is a valid SDXL mixed
# precision), even with a dense base_precision requested.
# The non-bf16-compute gate is also skipped for SDXL (fp16 is a valid SDXL mixed precision).
norm2 = _cfg(
base_model = "stabilityai/stable-diffusion-xl-base-1.0",
base_precision = "int8",
@ -322,8 +309,8 @@ def test_repo_is_prequantized_cases(repo, expected):
def test_repo_is_prequantized_alias_is_same_object():
# The trainer keeps a module-level alias for callers/tests; it must be the exact same
# function object as the common heuristic (moved there for config validation).
# The trainer keeps a module-level alias for callers/tests; it must be the exact same function
# object as the common heuristic.
assert dit._repo_is_prequantized is repo_is_prequantized
@ -338,21 +325,20 @@ def test_pick_auto_precision_policy_table():
# Missing free-VRAM number -> the safe nf4 mode.
assert p(False, "cuda", None, 23.8, (10, 0), True) == "nf4"
# Plenty of free VRAM -> bf16 regardless of fp8 capability: compiled bf16 measured
# FASTER than torchao float8 at LoRA-training shapes, so fp8 is opt-in only.
# Plenty of free VRAM gives bf16 regardless of fp8 capability: compiled bf16 measured FASTER than
# torchao float8 at LoRA-training shapes, so fp8 is opt-in only.
assert p(False, "cuda", 140, 23.8, (10, 0), True) == "bf16"
assert p(False, "cuda", 140, 23.8, (8, 0), True) == "bf16"
assert p(False, "cuda", 140, 23.8, (10, 0), False) == "bf16"
# Middle band (30 > 23.8 * 1.15 = 27.4, but not > 23.8 * 1.5 = 35.7) -> int8.
assert p(False, "cuda", 30, 23.8, (10, 0), True) == "int8"
# int8 needs torchao at runtime (no fallback), so the int8 band drops to nf4 when
# torchao is not importable while the bf16 band is unaffected.
# int8 needs torchao at runtime (no fallback), so the int8 band drops to nf4 when torchao is not
# importable while the bf16 band is unaffected.
assert p(False, "cuda", 30, 23.8, (10, 0), True, False) == "nf4"
assert p(False, "cuda", 140, 23.8, (10, 0), True, False) == "bf16"
# int8 still materialises the full bf16 transformer before quantize_ shrinks it, so
# free VRAM below the dense-load transient (25 < 27.4) must fall back to nf4 even
# though the QUANTIZED weights would have fit.
# int8 still materialises the full bf16 transformer before quantize_ shrinks it, so free VRAM
# below the dense-load transient must fall back to nf4 even though the quantized weights fit.
assert p(False, "cuda", 25, 23.8, (10, 0), True) == "nf4"
# Too little free VRAM for any dense load -> nf4.
assert p(False, "cuda", 10, 23.8, (10, 0), True) == "nf4"
@ -360,14 +346,14 @@ def test_pick_auto_precision_policy_table():
# ── _resolve_base_precision passthrough ───────────────────────────────────────
def test_resolve_base_precision_passes_explicit_through():
# An explicit mode passes straight through without probing the GPU (normalized() already
# validated it); the spec is only consulted for "auto".
# An explicit mode passes straight through without probing the GPU; the spec is only consulted
# for "auto".
spec = dit._SPECS["flux.1"]
cfg = _cfg(base_precision = "bf16")
assert dit._resolve_base_precision(cfg, spec, "cuda") == "bf16"
# The dense modes are CUDA-only: an explicit request on a GPU-less host fails fast
# (before any model load) instead of silently proceeding; /info never advertised it.
# The dense modes are CUDA-only: an explicit request on a GPU-less host fails fast, before any
# model load.
with pytest.raises(ValueError, match = "CUDA"):
dit._resolve_base_precision(cfg, spec, "cpu")
# nf4 stays a passthrough on any device (the bnb load path owns its own errors).
@ -375,19 +361,16 @@ def test_resolve_base_precision_passes_explicit_through():
def test_resolve_auto_requires_bf16_compute():
# auto may resolve to bf16/int8 which train in bf16 compute, so a non-bf16
# mixed_precision pins auto to the nf4 floor BEFORE any GPU probe (pure, no CUDA
# needed here) -- mirroring the normalized() rule for explicit dense modes.
# auto may resolve to bf16/int8, which train in bf16 compute, so a non-bf16 mixed_precision pins
# auto to the nf4 floor before any GPU probe -- mirroring the normalized() rule.
spec = dit._SPECS["flux.1"]
cfg = _cfg(base_precision = "auto", mixed_precision = "fp16")
assert dit._resolve_base_precision(cfg, spec, "cuda") == "nf4"
def test_resolve_auto_int8_band_gates_on_torchao(monkeypatch):
# The int8 auto band needs a FUNCTIONAL torchao at runtime; when torchao is not
# importable _resolve_base_precision must fall to nf4 instead of picking an int8 that
# would crash in _int8_quantize_base. Drive the probe into the int8 band and toggle the
# functional-torchao probe (shared with train_precision_modes, imported into the trainer).
# The int8 auto band needs a FUNCTIONAL torchao at runtime; without it _resolve_base_precision
# must fall to nf4 instead of picking an int8 that would crash in _int8_quantize_base.
import torch
spec = dit._SPECS["flux.1"] # dense_bf16_gb = 23.8
@ -414,9 +397,8 @@ def test_resolve_auto_int8_band_gates_on_torchao(monkeypatch):
def test_resolve_auto_int8_band_treats_stub_as_absent(monkeypatch):
# Simulate the Windows-ROCm torchao STUB: has_functional_torchao returns False (the
# stub satisfies find_spec but its quantize_ is a no-op), so the int8 band must fall to
# nf4 rather than pick an int8 whose quantization silently does nothing.
# Simulate the Windows-ROCm torchao STUB (find_spec succeeds but quantize_ is a no-op), so the
# int8 band must fall to nf4 rather than pick an int8 that silently does nothing.
import torch
spec = dit._SPECS["flux.1"]
@ -438,10 +420,8 @@ def test_resolve_auto_int8_band_treats_stub_as_absent(monkeypatch):
def test_has_functional_torchao_rejects_stub(monkeypatch):
# has_functional_torchao must reject the Unsloth import stub: even though
# `from torchao.quantization import quantize_` would succeed against the stub, the
# symbols are no-op stub types. Simulate a stub torchao.quantization module carrying the
# stub sentinel and assert the probe returns False.
# has_functional_torchao must reject the Unsloth import stub: `from torchao.quantization import
# quantize_` succeeds against it, but the symbols are no-op stub types.
import importlib
import types
@ -491,29 +471,28 @@ def test_fp8_module_filter():
# ── _should_compile fp8 branch ────────────────────────────────────────────────
def test_should_compile_fp8_branch():
# fp8 is only competitive compiled, so auto arms compile for it on a dense (non-bnb)
# cuda base.
# fp8 is only competitive compiled, so auto arms compile for it on a dense (non-bnb) cuda base.
cfg = _cfg(compile_transformer = "auto")
assert dit._should_compile(cfg, False, "cuda", "fp8") is True
# fp8 forces compile under auto even when the base is (hypothetically) reported as bnb.
assert dit._should_compile(cfg, True, "cuda", "fp8") is True
# An explicit "off" still wins over fp8 -- compile stays off.
# An explicit "off" still wins over fp8: compile stays off.
assert dit._should_compile(_cfg(compile_transformer = "off"), False, "cuda", "fp8") is False
# ── train_precision_modes machine probe ───────────────────────────────────────
def test_train_precision_modes_no_cuda(monkeypatch):
# Patch the torch module attribute the function imports so it observes a CPU-only box:
# no CUDA -> the nf4-only floor with nf4 recommended, and it never raises.
# Patch the torch module attribute the function imports so it observes a CPU-only box: no CUDA
# gives the nf4-only floor with nf4 recommended, and it never raises.
import torch
monkeypatch.setattr(torch.cuda, "is_available", lambda: False)
assert train_precision_modes() == (["nf4"], "nf4")
def test_train_precision_modes_gates_int8_fp8_on_torchao(monkeypatch):
# int8/fp8 are only advertised when torchao is FUNCTIONAL: on a CUDA host WITHOUT a real
# torchao (or with only the Windows-ROCm stub) /info must not offer int8/fp8, since their
# explicit paths import torchao with no fallback. bf16 + auto stay advertised.
# int8/fp8 are only advertised when torchao is FUNCTIONAL: on a CUDA host without a real torchao
# (or with only the Windows-ROCm stub) /info must not offer them, since their explicit paths
# import torchao with no fallback.
import torch
monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
@ -534,10 +513,9 @@ def test_train_precision_modes_gates_int8_fp8_on_torchao(monkeypatch):
def test_train_precision_modes_gates_dense_on_bf16_support(monkeypatch):
# The dense modes (bf16/int8/fp8/auto) all train in bf16 compute, which the DiT trainer
# requires. On a CUDA GPU that cannot do bf16 (T4/V100/RTX 20xx), /info must offer ONLY
# nf4 -- otherwise the UI advertises a start that evicts resident models and then fails the
# trainer's bf16 guard.
# The dense modes all train in bf16 compute, which the DiT trainer requires. On a CUDA GPU that
# cannot do bf16, /info must offer ONLY nf4, else the UI advertises a start that evicts resident
# models and then fails the trainer's bf16 guard.
import torch
monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
@ -551,11 +529,11 @@ def test_train_precision_modes_gates_dense_on_bf16_support(monkeypatch):
# ── family_train_infos precision fields ───────────────────────────────────────
def test_family_train_infos_carries_precision_fields(monkeypatch):
# Pin the machine probe so the DiT families carry a deterministic mode list, while SDXL
# (no precision selector) stays empty regardless of the probe.
# Pin the machine probe so the DiT families carry a deterministic mode list, while SDXL (no
# precision selector) stays empty regardless.
monkeypatch.setattr(common, "train_precision_modes", lambda: (["nf4", "bf16"], "auto"))
# Also pin bf16_unsupported_reason (family_train_infos reads the live GPU through it): "bf16 OK"
# so this positive-path assertion is deterministic across GPU types, not just on CPU-only CI.
# Also pin bf16_unsupported_reason (family_train_infos reads the live GPU through it) so this
# positive-path assertion is deterministic across GPU types.
monkeypatch.setattr(common, "bf16_unsupported_reason", lambda name: None)
infos = {i["name"]: i for i in common.family_train_infos()}
@ -567,8 +545,8 @@ def test_family_train_infos_carries_precision_fields(monkeypatch):
sdxl = infos["sdxl"]
assert sdxl["precision_modes"] == []
assert sdxl["recommended_precision"] == "nf4"
# The SDXL trainer regionally compiles its U-Net blocks too, so compile is advertised
# for every family; only the precision selector stays DiT-only.
# The SDXL trainer regionally compiles its U-Net blocks too, so compile is advertised for every
# family; only the precision selector stays DiT-only.
assert sdxl["supports_compile"] is True
@ -605,11 +583,9 @@ def test_request_model_base_precision():
def test_assert_trusted_base_model_rejects_local_non_pipeline(tmp_path):
# A local base_model dir that is NOT a diffusers pipeline (no model_index.json) is "trusted"
# (any existing path passes the trust check), but the spawned trainer loads it via
# from_pretrained, so it must be rejected in the /diffusion/start preflight BEFORE
# _free_gpu_for_diffusion_training tears down the resident models -- not fail the child after
# the eviction.
# A local base_model dir that is NOT a diffusers pipeline is "trusted" (any existing path passes),
# but the spawned trainer loads it via from_pretrained, so it must be rejected in the
# /diffusion/start preflight BEFORE _free_gpu_for_diffusion_training tears down the residents.
bad = tmp_path / "bare-base"
bad.mkdir()
with pytest.raises(ValueError, match = "model_index.json"):
@ -625,7 +601,7 @@ def test_assert_trusted_base_model_rejects_local_non_pipeline(tmp_path):
def test_dit_accelerator_missing_reason_and_info_hide_train_without_a_gpu(monkeypatch):
# Clicking Start on a GPU-less host evicted the resident Images pipeline, downloaded the text
# encoders, and only then died in the child: diffusers' bitsandbytes quantizer refuses 4-bit
# without CUDA/XPU/MPS. Reject it up front, and stop /info advertising a mode that always 400s.
# without an accelerator. Reject it up front, and stop /info advertising a mode that always 400s.
import torch
from core.training.diffusion_train_common import (

76
studio/backend/tests/test_diffusion_cache.py
View file

@ -144,8 +144,8 @@ def test_explicit_threshold_overrides_quant(monkeypatch):
def test_non_cachemixin_runs_uncached(monkeypatch):
# A transformer without enable_cache (e.g. Z-Image) must NOT install the standalone hook
# -- its pipeline opens no cache_context, so it runs uncached instead of crashing at gen.
# A transformer without enable_cache (e.g. Z-Image) must NOT install the standalone hook: its
# pipeline opens no cache_context, so it runs uncached instead of crashing at generation.
rec: dict = {}
_stub_diffusers(monkeypatch, hook_recorder = rec)
t = _NonCacheMixinTransformer()
@ -154,10 +154,9 @@ def test_non_cachemixin_runs_uncached(monkeypatch):
def test_pipeline_without_cache_context_runs_uncached(monkeypatch):
# A CacheMixin transformer whose PIPELINE never opens a cache_context (Flux Kontext /
# img2img / inpaint / controlnet reuse the CacheMixin FluxTransformer2DModel) must run
# uncached -- otherwise the First-Block-Cache hook raises "No context is set" on the
# first forward, crashing every default generation.
# A CacheMixin transformer whose PIPELINE never opens a cache_context (Flux Kontext / img2img /
# inpaint / controlnet reuse FluxTransformer2DModel) must run uncached, else the First-Block-Cache
# hook raises "No context is set" on the first forward.
_stub_diffusers(monkeypatch)
t = _MixinTransformer()
assert apply_step_cache(_NoCtxPipe(t), mode = "fbcache") is None
@ -172,8 +171,8 @@ def test_incompatible_model_runs_uncached(monkeypatch):
def test_enable_cache_failure_rolls_back_partial_hooks(monkeypatch):
# enable_cache can raise after hooking some blocks; the reported-uncached model
# must not actually run half-cached, so the failure path calls disable_cache.
# enable_cache can raise after hooking some blocks; the reported-uncached model must not actually
# run half-cached, so the failure path calls disable_cache.
_stub_diffusers(monkeypatch)
t = _MixinTransformer(fail = True)
t.disabled = False
@ -201,9 +200,9 @@ def test_missing_transformer_is_none(monkeypatch):
def test_diffusers_unavailable_runs_uncached(monkeypatch):
# no diffusers import -> best-effort returns None, load proceeds uncached. Block the
# hooks module too: the config import falls back to diffusers.hooks, which a REAL
# earlier import in the test session may have left cached in sys.modules.
# No diffusers import: best-effort returns None and the load proceeds uncached. Block the hooks
# module too, since the config import falls back to it and a real earlier import may have left it
# cached in sys.modules.
monkeypatch.setitem(sys.modules, "diffusers", None)
monkeypatch.setitem(sys.modules, "diffusers.hooks", None)
t = _MixinTransformer()
@ -228,9 +227,8 @@ def test_effective_steps_txt2img_is_full_count():
def test_effective_steps_low_strength_shrinks_below_the_bar():
# A 28-step upscale at strength 0.35 denoises int(9.8) = 9 steps (diffusers get_timesteps
# floors the product), which is below FBCACHE_MIN_STEPS -> the auto policy must NOT engage
# FBCache there.
# A 28-step upscale at strength 0.35 denoises int(9.8) = 9 steps (diffusers floors the product),
# below FBCACHE_MIN_STEPS, so the auto policy must NOT engage FBCache.
eff = effective_denoise_steps(28, 0.35)
assert eff == 9
assert eff < FBCACHE_MIN_STEPS
@ -239,33 +237,31 @@ def test_effective_steps_low_strength_shrinks_below_the_bar():
def test_effective_request_strength_uses_pipe_default_when_omitted():
import inspect
# txt2img (no init image) or a pipe without the strength kwarg -> full trajectory (None).
# txt2img or a pipe without the strength kwarg gives the full trajectory (None).
assert effective_request_strength(None, False, True, 0.6) is None
assert effective_request_strength(0.5, True, False, None) is None
# img2img with an explicit strength -> that value.
assert effective_request_strength(0.2, True, True, 0.6) == 0.2
# img2img with an OMITTED strength -> the pipe's own signature default (< 1), so the auto
# policy keys on the real (short) trajectory, not the full step count. This is the fix:
# int(28 * 0.6) = 16 real steps, not 28.
# img2img with an OMITTED strength uses the pipe's own signature default, so the auto policy keys
# on the real (short) trajectory: int(28 * 0.6) = 16 real steps, not 28.
s = effective_request_strength(None, True, True, 0.6)
assert s == 0.6
assert effective_denoise_steps(28, s) == 16
# A non-numeric signature default (inspect.Parameter.empty) falls back to the full count.
# A non-numeric signature default falls back to the full count.
assert effective_request_strength(None, True, True, inspect.Parameter.empty) is None
assert effective_request_strength(None, True, True, None) is None
def test_effective_steps_matches_diffusers_get_timesteps():
# Mirror diffusers exactly: it denoises init_timestep = min(int(num_inference_steps *
# strength), num_inference_steps) steps (the product is floored, not rounded).
# Mirror diffusers exactly: it denoises min(int(steps * strength), steps) steps (floored).
for steps, strength in [(28, 0.35), (28, 0.8), (50, 0.5), (20, 0.99), (30, 0.1)]:
expected = max(1, min(int(steps * strength), steps))
assert effective_denoise_steps(steps, strength) == expected
def test_toggle_stays_off_for_low_strength_workflow(monkeypatch):
# End to end: a 28-step request would engage FBCache, but at strength 0.35 the
# effective ~10 steps keep it uncached.
# End to end: a 28-step request would engage FBCache, but at strength 0.35 the effective ~10 steps
# keep it uncached.
_stub_diffusers(monkeypatch)
t = _ToggleTransformer()
mode = maybe_toggle_step_cache(_pipe(t), steps = effective_denoise_steps(28, 0.35))
@ -294,8 +290,8 @@ def test_normalize_auto_is_a_distinct_state():
def test_apply_treats_stray_auto_as_off(monkeypatch):
# AUTO must be resolved by the loader; if it ever reaches the engage call the
# load runs uncached instead of crashing.
# AUTO must be resolved by the loader; if it ever reaches the engage call the load runs uncached
# instead of crashing.
_stub_diffusers(monkeypatch)
t = _MixinTransformer()
assert apply_step_cache(_pipe(t), mode = "auto") is None
@ -432,8 +428,8 @@ def test_arming_is_idempotent(monkeypatch):
def test_arming_skips_uncompiled_blocks(monkeypatch):
# An eager-tier load has no _compiled_call_impl: the hook must stay untouched
# (compiling the inner would ADD compile where the user chose eager).
# An eager-tier load has no _compiled_call_impl: the hook must stay untouched (compiling the inner
# would ADD compile where the user chose eager).
_stub_torch_compile(monkeypatch)
block, hook, orig = _hooked_block(compiled = False)
assert _compile_hooked_block_inners(_fake_dit([block])) == 0
@ -441,8 +437,8 @@ def test_arming_skips_uncompiled_blocks(monkeypatch):
def test_arming_skips_partial_captured_inner(monkeypatch):
# A stacked hook chain (e.g. group offload) captures a functools.partial, not the
# plain bound method; arming would compile the wrong layer of the chain.
# A stacked hook chain (e.g. group offload) captures a functools.partial, not the plain bound
# method; arming would compile the wrong layer of the chain.
_stub_torch_compile(monkeypatch)
block, hook, orig = _hooked_block(bound = False)
assert _compile_hooked_block_inners(_fake_dit([block])) == 0
@ -450,8 +446,8 @@ def test_arming_skips_partial_captured_inner(monkeypatch):
def test_arming_covers_every_cache_hook_family(monkeypatch):
# FBCache is the image cache today, but the hook-name table already covers the
# MagCache layout too (same fn_ref shape), so a future mode arms for free.
# FBCache is the image cache today, but the hook-name table already covers the MagCache layout
# (same fn_ref shape), so a future mode arms for free.
_stub_torch_compile(monkeypatch)
names = (
"mag_cache_leader_block_hook",
@ -478,8 +474,8 @@ def test_restore_tolerates_fakes_without_modules():
def test_apply_step_cache_arms_compiled_blocks_on_toggle(monkeypatch):
# The generation-time toggle engages the cache AFTER the load already compiled the
# blocks; apply_step_cache must arm the fresh hooks itself.
# The generation-time toggle engages the cache AFTER the load already compiled the blocks, so
# apply_step_cache must arm the fresh hooks itself.
_stub_diffusers(monkeypatch)
_stub_torch_compile(monkeypatch)
block, hook, orig = _hooked_block()
@ -496,8 +492,8 @@ def test_apply_step_cache_arms_compiled_blocks_on_toggle(monkeypatch):
def test_toggle_disable_restores_inners_before_disable(monkeypatch):
# remove_hook splices fn_ref.original_forward back into module.forward, so the
# compiled wrapper must be swapped out BEFORE disable_cache runs.
# remove_hook splices fn_ref.original_forward back into module.forward, so the compiled wrapper
# must be swapped out BEFORE disable_cache runs.
_stub_diffusers(monkeypatch)
order = []
@ -518,8 +514,8 @@ def test_toggle_disable_restores_inners_before_disable(monkeypatch):
def test_enable_failure_restores_inners_before_partial_disable(monkeypatch):
# enable_cache can fail after hooking (and arming) some blocks; the partial-hook
# cleanup must un-arm them before disable_cache splices original_forward back.
# enable_cache can fail after hooking (and arming) some blocks; the partial-hook cleanup must
# un-arm them before disable_cache splices original_forward back.
_stub_diffusers(monkeypatch)
order = []
@ -544,9 +540,9 @@ def test_enable_failure_restores_inners_before_partial_disable(monkeypatch):
def test_enable_invalidates_stale_child_registry_cache(monkeypatch):
# diffusers 0.39 caches the child-registry list on first cache_context use; an
# UNCACHED generation already populates it (empty), so a later toggle-time
# enable_cache would install hooks the context never reaches ("No context is set").
# diffusers 0.39 caches the child-registry list on first cache_context use, and an UNCACHED
# generation already populates it (empty), so a later toggle-time enable_cache would install hooks
# the context never reaches ("No context is set").
_stub_diffusers(monkeypatch)
t = _MixinTransformer()
t._diffusers_hook = types.SimpleNamespace(_child_registries_cache = ["stale"])

15
studio/backend/tests/test_diffusion_compile_cache.py
View file

@ -245,8 +245,8 @@ def test_new_static_shape_redirties_a_hit(monkeypatch, tmp_path, fake_megacache)
assert ctx2.shapes == {(1024, 1024, 1)}
cc.register_shape(ctx2, (1024, 1024, 1), static = True)
assert cc.save(ctx2) is False
# ...but a NEW static shape (its compile just produced new artifacts) does, and the
# rewritten manifest covers both.
# ...but a NEW static shape (its compile just produced new artifacts) does, and the rewritten
# manifest covers both.
cc.register_shape(ctx2, (768, 768, 1), static = True)
assert ctx2.saved is False
assert cc.save(ctx2) is True
@ -255,9 +255,8 @@ def test_new_static_shape_redirties_a_hit(monkeypatch, tmp_path, fake_megacache)
def test_new_batch_size_is_its_own_static_shape(monkeypatch, tmp_path, fake_megacache):
# A static compile produces one artifact PER (w, h, batch): a batched generation at a
# batch size the bundle has not seen (incl. an OOM-backoff half) must re-dirty it, and
# the same (w, h) at the covered batch must not.
# A static compile produces one artifact PER (w, h, batch): a batched generation at an unseen
# batch size (incl. an OOM-backoff half) must re-dirty it, and the covered batch must not.
monkeypatch.setenv(cc._ENV_MODE, "auto")
monkeypatch.delenv(cc._ENV_SAVE, raising = False)
monkeypatch.setenv(cc._ENV_DIR, str(tmp_path))
@ -277,8 +276,8 @@ def test_new_batch_size_is_its_own_static_shape(monkeypatch, tmp_path, fake_mega
def test_gguf_quant_keys_apart_from_dense():
# A GGUF transformer compiles a different graph (the dequant chain) than the dense
# family; the load path fingerprints it quant="gguf" so bundles never cross-hit.
# A GGUF transformer compiles a different graph (the dequant chain) than the dense family; the
# load path fingerprints it quant="gguf" so bundles never cross-hit.
efp = cc.environment_fingerprint()
base = dict(
family = "flux.1",
@ -312,7 +311,7 @@ def test_fingerprint_mismatch_falls_back(monkeypatch, tmp_path, fake_megacache):
ctx = cc.begin(transformer = _transformer(), **_BEGIN_KW)
cc.save(ctx)
# Tamper the manifest's env fingerprint -> exact-match guard must reject the bundle.
# Tamper the manifest's env fingerprint: the exact-match guard must reject the bundle.
manifest = json.loads(ctx.manifest_path.read_text())
manifest["env"]["torch"] = "0.0.0-other"
ctx.manifest_path.write_text(json.dumps(manifest))

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

@ -96,8 +96,8 @@ def test_device_argument_excluded_from_the_key(cache_env):
def test_warm_reuse_across_installs(cache_env):
# A NEW pipe (fresh load) over the same directory hits the persisted entry without
# ever encoding -- the property that lets warm loads keep the text encoder off GPU.
# A NEW pipe (fresh load) over the same directory hits the persisted entry without ever encoding:
# the property that lets warm loads keep the text encoder off GPU.
first = _EncodePipe()
_install(first)
reference = first.encode_prompt("a sloth")
@ -123,9 +123,9 @@ def test_load_fingerprint_keys_apart(cache_env):
def test_companion_base_keys_apart(cache_env):
# A GGUF / single-file checkpoint takes its TEXT ENCODERS from the companion base, so
# the SAME checkpoint reloaded against a different base must re-encode rather than reuse
# the previous base's embeddings (silently different conditioning otherwise).
# A GGUF / single-file checkpoint takes its TEXT ENCODERS from the companion base, so the SAME
# checkpoint reloaded against a different base must re-encode rather than reuse the previous
# base's embeddings.
first = _EncodePipe()
_install(first, repo_id = "org/model-GGUF", base_repo = "base/one")
first.encode_prompt("a sloth")
@ -141,8 +141,8 @@ def test_companion_base_keys_apart(cache_env):
def test_a_local_base_updated_in_place_keys_apart(cache_env, tmp_path):
# A directory path is not a version: editing the text encoder in place must MISS, or the
# run silently conditions on embeddings from the encoder that was there before.
# A directory path is not a version: editing the text encoder in place must MISS, or the run
# silently conditions on embeddings from the encoder that was there before.
base = tmp_path / "base"
(base / "text_encoder").mkdir(parents = True)
weights = base / "text_encoder" / "model.safetensors"
@ -164,8 +164,8 @@ def test_a_local_base_updated_in_place_keys_apart(cache_env, tmp_path):
def test_source_revision_never_raises():
# Best-effort by contract: a missing path, a bare name and junk all resolve to a marker
# instead of blocking the load.
# Best-effort by contract: a missing path, a bare name and junk all resolve to a marker instead of
# blocking the load.
for ref in (None, "", "no/such/repo-xyz", "/does/not/exist", 1234):
assert isinstance(cond_cache._source_revision(ref), str)
@ -196,8 +196,8 @@ class _ListEncodePipe(_EncodePipe):
def test_tensor_list_slots_round_trip(cache_env):
# Z-Image returns list-of-tensors slots; the flatten/unflatten layout must
# reproduce them exactly on a warm hit.
# Z-Image returns list-of-tensors slots; the flatten/unflatten layout must reproduce them exactly
# on a warm hit.
pipe = _ListEncodePipe()
_install(pipe)
cold = pipe.encode_prompt(["a", "bb"])

75
studio/backend/tests/test_diffusion_controlnet.py
View file

@ -41,16 +41,16 @@ def test_resolve_controlnet_catalog_bare_repo_and_unknown():
def test_resolve_controlnet_rejects_filesystem_like_ids():
# The bare-repo fallback must never accept a path-shaped id: from_pretrained
# would treat it as a local directory, bypassing the controlnets_dir() contract.
# The bare-repo fallback must never accept a path-shaped id: from_pretrained would treat it as a
# local directory, bypassing the controlnets_dir() contract.
for bad in ("/tmp/model", "../some/model", "./x/y", "~/x/y", "a/b/c", "C:\\x/y", ".hidden/x"):
with pytest.raises(FileNotFoundError):
dc.resolve_controlnet(bad)
def test_resolve_controlnet_enforces_family_match():
# A curated entry tagged for another family must be rejected before download so it
# never reaches the wrong ControlNet pipeline class.
# A curated entry tagged for another family must be rejected before download so it never reaches
# the wrong ControlNet pipeline class.
with pytest.raises(ValueError, match = "not the"):
dc.resolve_controlnet("qwen-union", family = "flux.1")
# The matching family resolves fine, and no family (unfiltered) is permissive.
@ -59,9 +59,8 @@ def test_resolve_controlnet_enforces_family_match():
def test_resolve_controlnet_repo_id_still_family_gated():
# A curated ControlNet addressed by its full repo id (not its short catalog id) must still
# hit the family gate, not slip through the bare-repo fallback and load through the wrong
# family's ControlNet class.
# A curated ControlNet addressed by its full repo id must still hit the family gate, not slip
# through the bare-repo fallback into the wrong family's ControlNet class.
with pytest.raises(ValueError, match = "is for"):
dc.resolve_controlnet("InstantX/Qwen-Image-ControlNet-Union", family = "flux.1")
r = dc.resolve_controlnet("InstantX/Qwen-Image-ControlNet-Union", family = "qwen-image")
@ -69,9 +68,8 @@ def test_resolve_controlnet_repo_id_still_family_gated():
def test_union_control_mode_maps_only_union_entries():
# Union entries map a known control type to its integer mode; a union model always
# needs a concrete mode, so an unmapped type (passthrough) defaults to 0. A non-union
# id returns None so the caller omits control_mode.
# Union entries map a known control type to its integer mode; a union model always needs a
# concrete mode, so an unmapped type defaults to 0. A non-union id returns None.
assert dc.union_control_mode("flux-union-pro", "canny") == 0
assert dc.union_control_mode("flux-union-pro", "depth") == 2
assert dc.union_control_mode("flux-union-pro", "pose") == 4
@ -80,33 +78,30 @@ def test_union_control_mode_maps_only_union_entries():
def test_union_control_mode_rejects_unknown_type():
# An unknown / typo'd control type (e.g. 'detph') must NOT silently fall back to the canny
# head (0): preprocess_control passes non-canny maps through unchanged, so mode 0 would
# condition a map meant for another mode as canny -- silently wrong. Only passthrough (or an
# empty type) defaults to 0; anything else raises so the route returns a 400.
# An unknown / typo'd control type must NOT silently fall back to the canny head (0):
# preprocess_control passes non-canny maps through unchanged, so mode 0 would condition a map
# meant for another mode as canny. Only passthrough (or empty) defaults to 0; anything else raises.
with pytest.raises(ValueError, match = "Unknown control type"):
dc.union_control_mode("flux-union-pro", "detph")
with pytest.raises(ValueError, match = "Unknown control type"):
dc.union_control_mode("flux-union-pro", "scribble")
# passthrough and empty still default to 0 (the intended no-intrinsic-mode case); a non-union
# entry is unaffected (returns None, never raises).
# passthrough and empty still default to 0; a non-union entry returns None and never raises.
assert dc.union_control_mode("flux-union-pro", "") == 0
assert dc.union_control_mode("some/bare-repo", "detph") is None
def test_union_control_mode_matches_curated_repo_id():
# resolve_controlnet() accepts a curated union model by its bare HF repo id (owner/name),
# loading the same union repo the short catalog id points at. union_control_mode() must then
# recognise that repo id as union too -- otherwise control_mode is dropped and the union
# pipeline runs the wrong/default head (or diffusers raises on the missing mode).
# resolve_controlnet() accepts a curated union model by its bare HF repo id, so union_control_mode()
# must recognise that repo id as union too -- else control_mode is dropped and the union pipeline
# runs the wrong head.
assert dc.union_control_mode("Shakker-Labs/FLUX.1-dev-ControlNet-Union-Pro", "depth") == 2
assert dc.union_control_mode("Shakker-Labs/FLUX.1-dev-ControlNet-Union-Pro", "pose") == 4
assert dc.union_control_mode("Shakker-Labs/FLUX.1-dev-ControlNet-Union-Pro", "passthrough") == 0
assert dc.union_control_mode("InstantX/Qwen-Image-ControlNet-Union", "canny") == 0
# A bare repo id that is NOT a curated union model still returns None (caller omits the kwarg).
assert dc.union_control_mode("some/other-controlnet", "canny") is None
# A typo'd type against a repo-id-matched union still raises (route -> 400), same as the
# short-id path, rather than silently defaulting to the canny head.
# A typo'd type against a repo-id-matched union still raises (route 400), same as the short-id
# path.
with pytest.raises(ValueError, match = "Unknown control type"):
dc.union_control_mode("Shakker-Labs/FLUX.1-dev-ControlNet-Union-Pro", "detph")
@ -126,8 +121,8 @@ def test_resolve_controlnet_local(tmp_path, monkeypatch):
def test_scan_local_skips_config_only_folder(tmp_path, monkeypatch):
# A folder with config.json but no weight/index (interrupted copy) must NOT be
# advertised: it would otherwise fail deep in from_pretrained as a generic 500.
# A folder with config.json but no weight/index (interrupted copy) must NOT be advertised: it
# would fail deep in from_pretrained as a generic 500.
d = tmp_path / "controlnets"
d.mkdir()
incomplete = d / "incomplete-cn"
@ -146,14 +141,14 @@ def test_preprocess_control_passthrough_and_canny():
img = Image.new("RGB", (32, 24), (10, 20, 30))
# passthrough returns the same object.
assert dc.preprocess_control(img, "passthrough") is img
# a flat image has no edges, so the map is all black -- passing the source through would
# condition the ControlNet on its raw luminance, which is not an edge map at all.
# A flat image has no edges, so the map is all black; passing the source through would condition
# the ControlNet on its raw luminance.
import numpy as np
flat = dc.preprocess_control(img, "canny")
assert flat.mode == "RGB" and flat.size == (32, 24)
assert np.asarray(flat).max() == 0
# an image with structure yields an edge map: RGB, same size, some white pixels.
# An image with structure yields an edge map: RGB, same size, some white pixels.
arr = np.zeros((24, 32, 3), np.uint8)
arr[:, 16:, :] = 255 # a hard vertical edge
@ -240,8 +235,8 @@ class _FakeCNModel:
torch_dtype = None,
token = None,
use_safetensors = None,
# cache_dir (and any future loader kwarg) rides through: the real call pins the
# live cache root so a load cannot split across two of them.
# cache_dir (and any future loader kwarg) rides through: the real call pins the live cache root so
# a load cannot split across two of them.
**kwargs,
):
m = cls()
@ -306,15 +301,15 @@ def test_controlnet_pipe_loads_once_and_caches(monkeypatch):
_allow_cn_security(monkeypatch)
b = DiffusionBackend()
st = _state()
# The pipe cache only commits while ``st`` is the CURRENT load (an unload racing
# from_pipe must not repopulate the cache), so mirror the loaded invariant.
# The pipe cache only commits while ``st`` is the CURRENT load (an unload racing from_pipe must
# not repopulate it), so mirror the loaded invariant.
b._state = st
resolved = dc.ResolvedControlNet("flux-union-pro", "repo/id", is_local = False)
p1 = b._controlnet_pipe(st, resolved, threading.Event())
assert isinstance(p1, _FakeCNPipe) and isinstance(p1.controlnet, _FakeCNModel)
assert p1.controlnet.path == "repo/id" and p1.controlnet.device == "cpu"
# A remote (non-local) ControlNet must force safetensors so a pickle can't deserialize even if
# the Hub scan failed open.
# A remote (non-local) ControlNet must force safetensors so a pickle can't deserialize even if the
# Hub scan failed open.
assert p1.controlnet.use_safetensors is True
# cached: same id -> same model + same pipe, no reload.
p2 = b._controlnet_pipe(st, resolved, threading.Event())
@ -323,9 +318,9 @@ def test_controlnet_pipe_loads_once_and_caches(monkeypatch):
def test_controlnet_pipe_blocks_flagged_remote_repo(monkeypatch):
# A bare owner/name ControlNet is accepted by resolve_controlnet without the base
# trust gate, so the load path must run the Hub malware preflight: a flagged remote
# repo must raise BEFORE from_pretrained downloads/deserializes it.
# A bare owner/name ControlNet is accepted by resolve_controlnet without the base trust gate, so
# the load path must run the Hub malware preflight: a flagged remote repo must raise BEFORE
# from_pretrained downloads it.
import threading
import utils.security
@ -367,8 +362,8 @@ def test_controlnet_pipe_blocks_flagged_remote_repo(monkeypatch):
def test_controlnet_pipe_skips_scan_for_local_dir(monkeypatch, tmp_path):
# A local dir the user picked has no Hub scan; the preflight must not block it even
# if the (unused) scan stub would say blocked.
# A local dir the user picked has no Hub scan; the preflight must not block it even if the
# (unused) scan stub would say blocked.
import threading
import utils.security
@ -403,8 +398,8 @@ def test_controlnet_pipe_rejects_family_without_classes():
def test_controlnet_pipe_not_cached_after_unload_race(monkeypatch):
# An unload that lands while from_pipe is assembling must not let the wrapper
# repopulate the cache around the torn-down base pipe.
# An unload that lands while from_pipe is assembling must not let the wrapper repopulate the cache
# around the torn-down base pipe.
import threading
from core.inference.diffusion import DiffusionBackend

93
studio/backend/tests/test_diffusion_dataset_api.py
View file

@ -61,8 +61,8 @@ def test_list_images_caption_precedence(client, ds_root):
_write_png(folder / "a.png")
_write_png(folder / "b.png")
_write_png(folder / "c.png")
# a.png -> sidecar (an explicit edit beats the metadata row), b.png -> metadata-only,
# c.png -> none.
# a.png has a sidecar (an explicit edit beats the metadata row), b.png is metadata-only, c.png has
# none.
(folder / "metadata.jsonl").write_text(
json.dumps({"file_name": "a.png", "text": "from metadata"})
+ "\n"
@ -89,11 +89,10 @@ def test_list_images_caption_precedence(client, ds_root):
def test_list_images_tolerates_invalid_utf8_sidecar(client, ds_root):
# The upload route stores .txt/.caption sidecars as raw bytes (and users hand-drop them), so a
# sidecar can hold non-UTF-8 text. read_text then raises UnicodeDecodeError, which is a
# ValueError -- NOT an OSError -- so an `except OSError` around it 500s the whole labeling grid
# and the user cannot open it to repair the caption. One bad sidecar must read as no caption
# while every other image still lists (the info summary already behaves this way).
# The upload route stores .txt/.caption sidecars as raw bytes, so a sidecar can hold non-UTF-8
# text. read_text then raises UnicodeDecodeError, a ValueError not an OSError, so an `except
# OSError` around it 500s the whole labeling grid. One bad sidecar must read as no caption while
# every other image still lists.
folder = ds_root / "badutf8"
folder.mkdir()
_write_png(folder / "a.png")
@ -167,9 +166,8 @@ def test_put_caption_roundtrip_and_clear(client, ds_root):
def test_put_caption_overrides_metadata_row(client, ds_root):
# Editing a caption for an image that already has a metadata.jsonl row must take
# effect: the sidecar edit wins over the metadata caption in the response (and in
# the data the trainer reads), not the other way round.
# Editing a caption for an image that already has a metadata.jsonl row must take effect: the
# sidecar edit wins over the metadata caption, in the response and in what the trainer reads.
folder = ds_root / "cap"
folder.mkdir()
_write_png(folder / "x.png")
@ -205,8 +203,8 @@ def test_delete_image_cleans_sidecar_and_thumb(client, ds_root):
folder.mkdir()
_write_png(folder / "x.png")
(folder / "x.txt").write_text("cap", encoding = "utf-8")
# Generate a thumbnail so we can assert it is cleaned up too. Thumbs are keyed on
# the full filename (stem + extension) to avoid same-stem collisions across formats.
# Generate a thumbnail so we can assert it is cleaned up too. Thumbs are keyed on the full
# filename to avoid same-stem collisions across formats.
client.get("/api/train/diffusion/dataset/d/image/x.png?thumb=32")
assert list((folder / ".thumbs").glob("x.png_*.jpg"))
@ -218,8 +216,8 @@ def test_delete_image_cleans_sidecar_and_thumb(client, ds_root):
def test_thumb_cache_key_distinguishes_same_stem_extensions(client, ds_root):
# sample.png and sample.jpg share a stem; each must get its OWN thumbnail cache
# file, so the labeling grid never serves one image's thumbnail for the other.
# sample.png and sample.jpg share a stem; each must get its OWN thumbnail cache file, so the grid
# never serves one image's thumbnail for the other.
folder = ds_root / "d"
folder.mkdir()
Image.new("RGB", (8, 8), (10, 20, 30)).save(folder / "sample.png", format = "PNG")
@ -276,8 +274,8 @@ def test_list_dataset_examples(client, ds_root):
def test_list_dataset_examples_large_sets(client, ds_root):
# The two ~100-image sets: butterflies is a subject set (trigger, no caption column),
# nouns is a captioned style set (caption column, no trigger). Both cap at 100.
# The two ~100-image sets: butterflies is a subject set (trigger, no caption column), nouns a
# captioned style set. Both cap at 100.
r = client.get("/api/train/diffusion/dataset-examples")
examples = {e["id"]: e for e in r.json()["examples"]}
butterflies = examples["smithsonian-butterflies"]
@ -401,8 +399,8 @@ def _upload(client, name, files):
def test_upload_rejects_same_stem_different_extension(client, ds_root):
# sample.png and sample.jpg share the stem "sample", so both map to one sample.txt caption
# sidecar; keeping both would silently corrupt captions during training. The second must 400.
# sample.png and sample.jpg share the stem "sample", so both map to one sample.txt sidecar and
# keeping both would silently corrupt captions. The second must 400.
assert _upload(client, "styleset", [("sample.png", _png_bytes())]).status_code == 200
dup = _upload(client, "styleset", [("sample.jpg", _jpg_bytes())])
assert dup.status_code == 400
@ -413,19 +411,17 @@ def test_upload_rejects_same_stem_different_extension(client, ds_root):
def test_upload_same_stem_collision_within_one_batch(client, ds_root):
# The scan must cover files uploaded earlier IN THE SAME batch, not just those already on disk,
# so a single multipart request carrying both sample.png and sample.jpg is rejected too.
# The scan must cover files uploaded earlier IN THE SAME batch, not just those already on disk.
r = _upload(client, "styleset", [("sample.png", _png_bytes()), ("sample.jpg", _jpg_bytes())])
assert r.status_code == 400
assert "Duplicate image name" in r.json()["detail"]
def test_upload_rejects_exact_duplicate_name_within_one_batch(client, ds_root):
# Two parts with the SAME name in ONE multipart batch are distinct files (dragged from
# different folders, or an API client repeating a part); the staged commit would let the
# later tmp.replace(dest) silently discard the earlier one while `uploaded` still counts
# both. The batch must be rejected whole. Re-sending a name in a SEPARATE upload stays a
# deliberate overwrite (test_upload_allows_exact_name_overwrite_and_caption_sidecar).
# Two parts with the SAME name in ONE multipart batch are distinct files, and the staged commit
# would let the later replace silently discard the earlier one while `uploaded` counts both. The
# batch must be rejected whole; re-sending a name in a SEPARATE upload stays a deliberate
# overwrite.
r = _upload(
client,
"styleset",
@ -438,17 +434,15 @@ def test_upload_rejects_exact_duplicate_name_within_one_batch(client, ds_root):
r = _upload(client, "styleset", [("sample.txt", b"a"), ("sample.txt", b"b")])
assert r.status_code == 400
assert "more than once" in r.json()["detail"]
# A STEM case variant pair (Cat.png vs cat.png) stays exempt, matching the stem-guard
# contract: it is one file / an overwrite on case-insensitive filesystems, and on Linux
# the two files write separate sidecars (Cat.txt vs cat.txt) -- no caption collision.
# A STEM case variant pair (Cat.png vs cat.png) stays exempt: it is one file / an overwrite on
# case-insensitive filesystems, and on Linux the two write separate sidecars.
r = _upload(client, "styleset", [("Cat.png", _png_bytes()), ("cat.png", _png_bytes())])
assert r.status_code == 200
def test_upload_rejects_extension_case_variant_sidecar_collision(client, ds_root):
# An EXTENSION-case variant pair (dog.PNG vs dog.png) has exactly equal stems: on a
# case-sensitive filesystem both files land and both resolve to ONE dog.txt caption
# sidecar, silently sharing/corrupting the caption. Must 400 both within one batch and
# An EXTENSION-case variant pair (dog.PNG vs dog.png) has exactly equal stems: on a case-sensitive
# filesystem both land and both resolve to ONE dog.txt sidecar. Must 400 within one batch and
# against a file already on disk.
r = _upload(client, "styleset", [("dog.PNG", _png_bytes()), ("dog.png", _png_bytes())])
assert r.status_code == 400
@ -463,8 +457,8 @@ def test_upload_rejects_extension_case_variant_sidecar_collision(client, ds_root
def test_upload_allows_exact_name_overwrite_and_caption_sidecar(client, ds_root):
# Re-uploading the EXACT same name (stem AND extension) is an allowed overwrite, and a .txt
# caption for the same stem is the intended kohya flow -- neither is a same-stem image collision.
# Re-uploading the EXACT same name is an allowed overwrite, and a .txt caption for the same stem
# is the intended kohya flow: neither is a same-stem image collision.
assert (
_upload(client, "styleset", [("sample.png", _png_bytes((10, 20, 30)))]).status_code == 200
)
@ -479,14 +473,13 @@ def test_upload_allows_exact_name_overwrite_and_caption_sidecar(client, ds_root)
# ── import: promotion is all-or-nothing ──────────────────────────────────────
def test_import_promotion_leaves_no_partial_dataset_on_failure(ds_root, monkeypatch):
# The staging dir is promoted into the dataset folder in one atomic rename. If that rename
# fails (a crash / filesystem error mid-promotion), the folder must be left with NO images
# rather than a half-filled dataset that the image_count>0 idempotency check would then accept
# as complete on retry -- stranding the user with a truncated dataset. Simulate the promotion
# rename failing, assert nothing partial is left, and assert a retry re-imports cleanly.
# The staging dir is promoted into the dataset folder in one atomic rename. If that rename fails,
# the folder must be left with NO images rather than a half-filled dataset the image_count>0
# idempotency check would accept as complete on retry. Simulate the rename failing, assert nothing
# partial is left, and assert a retry re-imports cleanly.
import os
# A client that returns the 500 (as production does) instead of re-raising the server exception.
# A client that returns the 500 (as production does) instead of re-raising.
app = FastAPI()
app.include_router(training_router, prefix = "/api/train")
app.dependency_overrides[get_current_subject] = lambda: "test-user"
@ -497,7 +490,7 @@ def test_import_promotion_leaves_no_partial_dataset_on_failure(ds_root, monkeypa
real_replace = os.replace
def flaky_replace(src, dst, *a, **k):
# Only sabotage the staging -> folder promotion; leave every other rename working.
# Only sabotage the staging to folder promotion; leave every other rename working.
if str(dst) == str(folder):
raise OSError("simulated crash during promotion")
return real_replace(src, dst, *a, **k)
@ -508,7 +501,7 @@ def test_import_promotion_leaves_no_partial_dataset_on_failure(ds_root, monkeypa
json = {"id": "tuxemon", "name": "my-tux"},
)
assert r.status_code == 500
# No half-filled dataset: the folder holds zero images (and no stray staging dir lingers).
# No half-filled dataset: the folder holds zero images and no stray staging dir.
assert list(ds_root.glob("my-tux/*.png")) == []
assert not any(p.name.startswith(".my-tux.import-") for p in ds_root.iterdir())
@ -542,7 +535,7 @@ def test_upload_rolls_back_when_a_later_promotion_fails(ds_root, monkeypatch):
state = {"failed": False}
def flaky_replace(self, target, *a, **k):
# Fail once on the tmp -> b.txt promotion only (not the backup restore), so rollback works.
# Fail once on the tmp to b.txt promotion only (not the backup restore), so rollback works.
if (
not state["failed"]
and str(target).endswith("b.txt")
@ -562,7 +555,7 @@ def test_upload_rolls_back_when_a_later_promotion_fails(ds_root, monkeypatch):
assert r.status_code == 500
monkeypatch.setattr(Path, "replace", real_replace)
# Both originals are intact -- no partial overwrite of the live dataset.
# Both originals are intact: no partial overwrite of the live dataset.
assert (folder / "a.txt").read_bytes() == b"ORIGINAL-A"
assert (folder / "b.txt").read_bytes() == b"ORIGINAL-B"
# No staging or backup artifacts left behind.
@ -572,8 +565,8 @@ def test_upload_rolls_back_when_a_later_promotion_fails(ds_root, monkeypatch):
def test_upload_rechecks_training_state_before_commit(ds_root, monkeypatch):
# A /diffusion/start that reserves the training slot AFTER the upload passed its entry guard but
# BEFORE the commit must not have its dataset mutated: the recheck just before the tmp->dest
# promotion catches the now-active run, 409s, and leaves the on-disk dataset untouched.
# BEFORE the commit must not have its dataset mutated: the recheck just before the promotion
# catches the now-active run, 409s, and leaves the on-disk dataset untouched.
import routes.training as tr
folder = ds_root / "styleset"
@ -583,8 +576,8 @@ def test_upload_rechecks_training_state_before_commit(ds_root, monkeypatch):
calls = {"n": 0}
def fake_active():
# Inactive at the entry guard (call 1), active by the pre-commit recheck (call 2+): the
# training run started while the upload was streaming.
# Inactive at the entry guard (call 1), active by the pre-commit recheck (call 2+): the training
# run started while the upload was streaming.
calls["n"] += 1
return calls["n"] >= 2
@ -672,9 +665,9 @@ def test_delete_image_with_glob_chars_only_removes_own_thumbs(client, ds_root):
def test_import_preserves_unrelated_files_when_folder_not_empty(client, ds_root, monkeypatch):
# If the target folder already holds unrelated NON-image files (so image_count is still 0 and
# the import runs), the atomic rmdir refuses and the code falls back to a per-file move: the
# images are imported AND the pre-existing file is preserved rather than clobbered or lost.
# If the target folder already holds unrelated NON-image files (so image_count is still 0 and the
# import runs), the atomic rmdir refuses and the code falls back to a per-file move: the images
# are imported AND the pre-existing file is preserved.
_install_fake_load_dataset(monkeypatch, n_rows = 3)
folder = ds_root / "my-tux"
folder.mkdir(parents = True)

4
studio/backend/tests/test_diffusion_device.py
View file

@ -121,7 +121,7 @@ def _install(
"""Install the fake torch and either a fake or failing `utils.hardware`."""
monkeypatch.setitem(sys.modules, "torch", torch)
if hardware_fails:
# Force `from utils.hardware import ...` to raise -> torch-probe fallback.
# Force `from utils.hardware import ...` to raise, exercising the torch-probe fallback.
monkeypatch.setitem(sys.modules, "utils.hardware", None)
return
@ -157,7 +157,7 @@ def test_cuda_pre_ampere_fp16(monkeypatch):
torch = _make_torch(cuda_available = True, capability = (7, 5), bf16_supported = True)
_install(monkeypatch, torch, studio_device = "cuda")
t = dd.resolve_diffusion_device_target()
# is_bf16_supported() is True (emulated) but capability < 8 -> fp16.
# is_bf16_supported() is True (emulated) but capability < 8, so fp16.
assert t.dtype == FP16 and t.backend == "cuda"

82
studio/backend/tests/test_diffusion_dit_trainer.py
View file

@ -47,8 +47,7 @@ def test_specs_cover_the_dit_families():
"flux.2-klein",
"flux.2-dev",
}
# FLUX / Qwen share the added-kv attention target set; Z-Image and Krea 2 are
# single-stream.
# FLUX / Qwen share the added-kv attention target set; Z-Image and Krea 2 are single-stream.
assert "add_q_proj" in _SPECS["flux.1"].lora_targets
assert "add_q_proj" in _SPECS["qwen-image"].lora_targets
assert "add_q_proj" not in _SPECS["z-image"].lora_targets
@ -62,12 +61,12 @@ def test_specs_cover_the_dit_families():
def test_flux2_specs_share_targets_and_split_conditioners():
# dev and Klein share the transformer (and so the LoRA target set) but load different
# conditioning pipelines and save through their own pipeline class.
# dev and Klein share the transformer (and so the LoRA target set) but load different conditioning
# pipelines and save through their own pipeline class.
klein, dev = _SPECS["flux.2-klein"], _SPECS["flux.2-dev"]
assert klein.lora_targets == dev.lora_targets == _FLUX2_TARGETS
# The fused single-stream projection is targeted; the plain to_out suffix is not (it
# would also match the double-stream ModuleList container, which peft cannot wrap).
# The fused single-stream projection is targeted; the plain to_out suffix is not (it would also
# match the double-stream ModuleList container, which peft cannot wrap).
assert "to_qkv_mlp_proj" in _FLUX2_TARGETS
assert "to_out.0" in _FLUX2_TARGETS
assert "to_out" not in _FLUX2_TARGETS
@ -79,9 +78,8 @@ def test_flux2_specs_share_targets_and_split_conditioners():
def test_select_lora_targets_uses_family_default_for_generic_config():
# normalized() fills lora_target_modules with the generic DEFAULT_LORA_TARGETS when a
# caller doesn't set it, so that value must resolve to the family's targets (which add
# the DiT-specific projections), not stay stuck on the generic SDXL list.
# normalized() fills lora_target_modules with the generic DEFAULT_LORA_TARGETS when a caller
# doesn't set it, so that value must resolve to the family's targets, not stay on the SDXL list.
assert _select_lora_targets(DEFAULT_LORA_TARGETS, _FLUX_TARGETS) == _FLUX_TARGETS
assert _select_lora_targets(DEFAULT_LORA_TARGETS, _QWEN_TARGETS) == _QWEN_TARGETS
assert _select_lora_targets(DEFAULT_LORA_TARGETS, _ZIMAGE_TARGETS) == _ZIMAGE_TARGETS
@ -129,9 +127,9 @@ def test_zimage_rejects_fp16_before_loading():
def test_flux2_rejects_fp16_before_loading():
# Both FLUX.2 variants resolve from their repo names, and both are bf16-only: an
# explicit fp16 request fails in normalized() itself, before anything loads. Klein's
# base is ungated, so this exercises the precision guard directly (no token in play).
# Both FLUX.2 variants resolve from their repo names and are bf16-only: an explicit fp16 request
# fails in normalized() itself. Klein's base is ungated, so this exercises the precision guard
# directly.
ok = DiffusionLoraConfig(
base_model = "black-forest-labs/FLUX.2-klein-4B", data_dir = "d", output_dir = "o"
).normalized()
@ -165,9 +163,8 @@ def test_flux2_bases_pass_the_trusted_base_gate():
def test_every_train_base_is_deployable_as_an_inference_pipeline():
# "Deploy to Create" reloads the trained-on base (or the family's deploy_base) through
# /images/load as a PIPELINE, which gates non-GGUF loads on _is_trusted_diffusion_repo. Any
# advertised training base that fails that gate makes Deploy 400 for every adapter trained on
# it -- which is what happened to both FLUX.2 families, trusted for training only.
# /images/load as a PIPELINE, gated on _is_trusted_diffusion_repo. Any advertised training base
# failing that gate makes Deploy 400 for every adapter trained on it.
from core.inference.diffusion import _is_trusted_diffusion_repo
from core.inference.diffusion_families import _FAMILIES
for fam in _FAMILIES:
@ -217,10 +214,9 @@ def test_family_train_infos_lists_dit_families():
def test_family_train_infos_sdxl_supports_compile_without_precision_modes(monkeypatch):
# Regional compile now applies to every family (the SDXL trainer compiles its U-Net
# blocks too), but base_precision stays DiT-only, so SDXL advertises no precision modes
# while a DiT family (z-image) keeps its own. Pin the precision list so the assertion
# holds regardless of the test host's GPU capability.
# Regional compile now applies to every family (the SDXL trainer compiles its U-Net blocks too),
# but base_precision stays DiT-only, so SDXL advertises no precision modes while z-image keeps
# its own. Pin the precision list so the assertion holds regardless of the host GPU.
import core.training.diffusion_train_common as dtc
monkeypatch.setattr(dtc, "train_precision_modes", lambda: (["nf4", "bf16", "auto"], "auto"))
@ -247,16 +243,14 @@ def test_mx_module_filter_accepts_dense_block_linear():
def test_mx_module_filter_skips_biased_linear():
# The torchao 0.17 MX training path drops the bias term (its linear override computes
# input @ weight_t only), so an mxfp8'd biased FROZEN linear would silently lose its bias and
# corrupt the base output the LoRA regresses against. Biased linears must stay bf16.
# The torchao 0.17 MX training path drops the bias term, so an mxfp8'd biased FROZEN linear would
# silently lose its bias and corrupt the base output the LoRA regresses against.
assert _mx_module_filter(_linear(3072, 3072, bias = True), "blocks.0.ff.up") is False
def test_resolve_base_precision_explicit_mxfp8_requires_blackwell(monkeypatch):
# An explicit mxfp8 request on a non-Blackwell CUDA GPU must fail fast: its MX GEMM has no
# kernel below sm100 and would otherwise crash at the first training step, after a full dense
# transformer load. /info only advertises mxfp8 on sm100+, so this mirrors that gate.
# An explicit mxfp8 request on a non-Blackwell CUDA GPU must fail fast: its MX GEMM has no kernel
# below sm100 and would otherwise crash at the first training step, after a full dense load.
import torch
monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *a, **k: (8, 9))
@ -274,8 +268,7 @@ def test_resolve_base_precision_explicit_mxfp8_ok_on_blackwell(monkeypatch):
def test_mx_module_filter_skips_lora_and_proj_out():
# LoRA-owned modules (adapters stay high precision) and the output projection are
# excluded, mirroring the fp8 filter's guards.
# LoRA-owned modules and the output projection are excluded, mirroring the fp8 filter.
lin = _linear(3072, 3072)
assert _mx_module_filter(lin, "blocks.0.attn.to_q.lora_A.default") is False
assert _mx_module_filter(lin, "proj_out") is False
@ -283,7 +276,7 @@ def test_mx_module_filter_skips_lora_and_proj_out():
def test_mx_module_filter_rejects_non_block_aligned_dims():
# MX block scaling tiles 32-wide, so a dim not divisible by 32 (3000) is rejected.
# MX block scaling tiles 32-wide, so a dim not divisible by 32 is rejected.
assert _mx_module_filter(_linear(3000, 3072), "blocks.0.ff.up") is False
@ -295,8 +288,8 @@ def test_mx_module_filter_rejects_non_linear():
def test_should_compile_auto_mxfp8_on_cuda():
# auto compiles the dense speed modes on cuda; int8 stays eager (torchao subclass);
# an explicit "off" wins over the mode.
# auto compiles the dense speed modes on cuda; int8 stays eager (torchao subclass); an explicit
# "off" wins over the mode.
cfg = DiffusionLoraConfig(base_model = "b", data_dir = "d", output_dir = "o")
assert _should_compile(cfg, False, "cuda", base_precision = "mxfp8") is True
assert _should_compile(cfg, False, "cuda", base_precision = "int8") is False
@ -307,9 +300,8 @@ def test_should_compile_auto_mxfp8_on_cuda():
def test_apply_mxfp8_training_failure_falls_back_with_warning(monkeypatch):
# An unavailable torchao MX path must never be fatal: force both API revisions'
# imports to raise, then assert the helper returns False and emits exactly one
# warning naming mxfp8.
# An unavailable torchao MX path must never be fatal: force both API revisions' imports to raise,
# then assert the helper returns False with exactly one warning naming mxfp8.
monkeypatch.setitem(sys.modules, "torchao.prototype.mx_formats", None)
monkeypatch.setitem(sys.modules, "torchao.prototype.moe_training.config", None)
events = []
@ -321,10 +313,9 @@ def test_apply_mxfp8_training_failure_falls_back_with_warning(monkeypatch):
def test_mxfp8_training_config_falls_back_to_the_torchao_0_17_api(monkeypatch):
# torchao 0.17 removed prototype.mx_formats.MXLinearConfig in favour of the
# MXFP8TrainingOpConfig recipe API; the config helper must fall back to it so the
# advertised mxfp8 mode keeps engaging on those installs instead of silently
# training dense bf16.
# torchao 0.17 removed prototype.mx_formats.MXLinearConfig in favour of the MXFP8TrainingOpConfig
# recipe API, so the config helper must fall back to it or the advertised mxfp8 mode silently
# trains dense bf16.
from types import SimpleNamespace
from core.training.diffusion_dit_trainer import _mxfp8_training_config
@ -351,13 +342,10 @@ def test_mxfp8_training_config_falls_back_to_the_torchao_0_17_api(monkeypatch):
def _patch_capability(monkeypatch, capability):
# Drive train_precision_modes' GPU probe: pretend CUDA is present at the given tensor
# core capability (fp8 needs sm89+, mxfp8 needs sm100+). The torchao probe is stubbed
# functional so these tests exercise the CAPABILITY gate on hosts without torchao
# (the CPU-only CI runner does not install it). is_bf16_supported must be stubbed True
# too: the dense modes gate on it, and an Ada/Blackwell GPU is by definition bf16-capable,
# so without this the modes collapse to nf4 on a CPU runner where the real probe is False
# (the test otherwise only passes on a bf16 GPU host).
# Drive train_precision_modes' GPU probe: pretend CUDA is present at the given capability (fp8
# needs sm89+, mxfp8 sm100+). torchao is stubbed functional so these exercise the CAPABILITY gate
# on hosts without it, and is_bf16_supported is stubbed True (an Ada/Blackwell GPU is bf16-capable
# by definition, and the dense modes gate on it).
import torch
import core.training.diffusion_train_common as dtc
@ -394,9 +382,9 @@ def test_train_precision_modes_newer_blackwell_has_mxfp8(monkeypatch):
def test_train_precision_modes_pre_ampere_is_nf4_only(monkeypatch):
# A pre-Ampere GPU EMULATES bf16 (is_bf16_supported() True) but has no native bf16 tensor
# cores; the DiT trainer requires native bf16, so /info must offer nf4 only. Otherwise it
# advertises a start that evicts resident models and then fails the trainer's bf16 guard.
# A pre-Ampere GPU EMULATES bf16 but has no native bf16 tensor cores, and the DiT trainer requires
# native bf16, so /info must offer nf4 only. Otherwise it advertises a start that evicts resident
# models and then fails the trainer's bf16 guard.
import torch
monkeypatch.setattr(torch.cuda, "is_available", lambda: True)

28
studio/backend/tests/test_diffusion_dit_trainer_flow_shift.py
View file

@ -26,9 +26,8 @@ QWEN_SHIFT_TERMINAL = 0.02
def _qwen_scheduler():
# The Qwen/Qwen-Image scheduler config: shift=1.0 is SKIPPED at init because
# use_dynamic_shifting is true, base_shift = max_shift = log 3 (constant inference mu),
# exponential time shift, terminal stretch to 0.02.
# The Qwen/Qwen-Image scheduler config: shift=1.0 is SKIPPED at init because use_dynamic_shifting
# is true, base_shift = max_shift = log 3, exponential time shift, terminal stretch to 0.02.
diffusers = pytest.importorskip("diffusers")
FlowMatchEulerDiscreteScheduler = diffusers.FlowMatchEulerDiscreteScheduler
return FlowMatchEulerDiscreteScheduler(
@ -91,9 +90,8 @@ def test_flow_shift_explicit_values_and_validation():
DiffusionLoraConfig(
base_model = "b", data_dir = "d", output_dir = "o", flow_shift = "bogus"
).normalized()
# Non-finite must be rejected too: JSON accepts 1e309, which floats to inf, and
# inf <= 0 is False (NaN fails every comparison), so a positivity-only guard passed
# them through to the sigma table as NaN and the run saved a corrupted adapter.
# Non-finite must be rejected too: JSON accepts 1e309, which floats to inf, and a positivity-only
# guard passed it through to the sigma table as NaN, saving a corrupted adapter.
for bad in (float("inf"), float("-inf"), float("nan"), 1e309):
with pytest.raises(ValueError, match = "flow_shift"):
DiffusionLoraConfig(
@ -149,14 +147,14 @@ def test_auto_table_matches_the_exact_qwen_transform():
sched = _qwen_scheduler()
table = _training_sigma_table(sched, "auto")
base = sched.sigmas
# Exponential shift at mu = log 3 with sigma exponent 1 is exp(mu)/(exp(mu) + 1/u - 1)
# = 3u/(1 + 2u), then the terminal stretch maps the schedule's last sigma to 0.02.
# Exponential shift at mu = log 3 with sigma exponent 1 is exp(mu)/(exp(mu) + 1/u - 1) =
# 3u/(1 + 2u), then the terminal stretch maps the schedule's last sigma to 0.02.
shifted = 3.0 * base / (1.0 + 2.0 * base)
scale = (1.0 - shifted[-1]) / (1.0 - QWEN_SHIFT_TERMINAL)
expected = 1.0 - (1.0 - shifted) / scale
assert torch.allclose(table, expected, atol = 1e-6)
# Fixed-point spot checks: sigma 1.0 stays 1.0, the terminal sigma lands on 0.02, and
# the midpoint u = 0.5 rises to ~0.754 (3u/(1+2u) = 0.75 before the stretch).
# Fixed-point spot checks: sigma 1.0 stays 1.0, the terminal sigma lands on 0.02, and the midpoint
# u = 0.5 rises to ~0.754 (3u/(1+2u) = 0.75 before the stretch).
assert abs(float(table[0]) - 1.0) < 1e-6
assert abs(float(table[-1]) - QWEN_SHIFT_TERMINAL) < 1e-6
assert abs(float(table[499]) - 0.75427) < 1e-3
@ -176,9 +174,9 @@ def test_numeric_table_applies_the_linear_shift():
def test_identity_and_static_families_are_untouched():
# flow_shift 1.0 must return the scheduler's own table object (no numeric drift for
# FLUX / Z-Image / Krea 2), and "auto" on a static-shift scheduler is a no-op too:
# its init already baked the shift into sigmas.
# flow_shift 1.0 must return the scheduler's own table object (no numeric drift for FLUX / Z-Image
# / Krea 2), and "auto" on a static-shift scheduler is a no-op too: its init already baked the
# shift into sigmas.
sched = _qwen_scheduler()
assert _training_sigma_table(sched, 1.0) is sched.sigmas
static = _flux_static_scheduler()
@ -195,8 +193,8 @@ def test_sampled_sigma_distribution_shifts_under_auto():
_, idx = _sample_timesteps(sched, 4096, "cpu")
base = _gather_sigmas(sched.sigmas, idx, "cpu", torch.float32, 1)
shifted = _gather_sigmas(auto_table, idx, "cpu", torch.float32, 1)
# Unshifted logit-normal draws center at 0.5; the mu = log 3 shift + terminal stretch
# pushes the mass toward high noise (mean ~0.72). Shift raises EVERY sample.
# Unshifted logit-normal draws center at 0.5; the mu = log 3 shift + terminal stretch pushes the
# mass toward high noise (mean ~0.72). Shift raises EVERY sample.
assert abs(float(base.mean()) - 0.5) < 0.03
assert float(shifted.mean()) > 0.68
assert bool((shifted >= base - 1e-6).all())

4
studio/backend/tests/test_diffusion_eager_patches.py
View file

@ -94,8 +94,8 @@ def test_patched_matches_original(cls, device, dtype):
with torch.inference_mode():
got = _first(_call(cls, m, args))
# The fused ops are FMA-based (addcmul) / fused (F.rms_norm): within ~1 ULP of the
# stock mul+add (and more accurate, single rounding), NOT bit-identical in fp32.
# The fused ops are FMA-based (addcmul) / fused (F.rms_norm): within ~1 ULP of the stock mul+add
# (and more accurate, single rounding), NOT bit-identical in fp32.
atol, rtol = (1e-5, 1e-4) if dtype == torch.float32 else (8e-3, 8e-3)
torch.testing.assert_close(got, ref, atol = atol, rtol = rtol)

46
studio/backend/tests/test_diffusion_engine_router.py
View file

@ -25,15 +25,15 @@ _ENVS = (
def _clean_env_and_state(monkeypatch):
for e in _ENVS:
monkeypatch.delenv(e, raising = False)
# A light status-capable stub so neither selection nor active_status() imports the
# heavy diffusers/sd.cpp backends; the active engine NAME comes from module state.
# A light status-capable stub so neither selection nor active_status() imports the heavy
# diffusers/sd.cpp backends; the active engine NAME comes from module state.
monkeypatch.setattr(
r,
"get_active_diffusion_engine",
lambda: SimpleNamespace(status = lambda: {"loaded": False, "repo_id": None}),
)
# Default: no resident sd-server (so existing tests exercise the sd-cli path only) and
# a stubbed runnability probe, so neither reaches the real install/exec path.
# Default: no resident sd-server (so existing tests exercise the sd-cli path) and a stubbed
# runnability probe, so neither reaches the real install/exec path.
monkeypatch.setattr(r, "ensure_sd_server_binary", lambda **_: None)
monkeypatch.setattr(r, "_server_binary_runnable", lambda *_a, **_k: True)
yield
@ -75,8 +75,8 @@ def test_cpu_with_binary_and_supported_family_picks_sd_cpp(monkeypatch):
def test_cpu_with_only_sd_server_picks_sd_cpp(monkeypatch):
# An sd-server-only install (no runnable sd-cli) must still route to native: the
# backend prefers the resident server, so a runnable sd-server is native availability.
# An sd-server-only install (no runnable sd-cli) must still route to native: the backend prefers
# the resident server, so a runnable sd-server is native availability.
_set_device(monkeypatch, "cpu")
_set_binary(monkeypatch, None) # no sd-cli
monkeypatch.setattr(r, "SdCppEngine", lambda **_: SimpleNamespace(version = lambda: None))
@ -85,8 +85,8 @@ def test_cpu_with_only_sd_server_picks_sd_cpp(monkeypatch):
def test_present_but_not_runnable_binary_falls_back(monkeypatch):
# A binary that exists but cannot run (version() -> None) must fall back to
# diffusers at selection, not commit native and fail inside the load.
# A binary that exists but cannot run must fall back to diffusers at selection, not commit native
# and fail inside the load.
_set_device(monkeypatch, "cpu")
_set_binary(monkeypatch, "/usr/bin/sd-cli")
monkeypatch.setattr(r, "SdCppEngine", lambda **_: SimpleNamespace(version = lambda: None))
@ -168,8 +168,8 @@ def test_install_accelerator_maps_backend(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.
# Forcing sd_cpp on a ROCm host with no binary must install the ROCm build, not the default CPU
# one, else the forced-native generation silently runs on CPU.
_set_device(monkeypatch, "rocm")
_set_runnable(monkeypatch)
seen = {}
@ -200,12 +200,10 @@ def test_active_status_injects_engine_and_reason(monkeypatch):
def test_switch_unloads_old_engine_before_publishing_new(monkeypatch):
# The arbiter's diffusion evictor unloads get_active_diffusion_engine(); if the router
# published the new (empty) engine BEFORE the old one finished unloading, a concurrent
# chat/video acquire_for could evict the empty engine and take the GPU while the old
# model was still resident (two large models briefly co-resident -> OOM). Assert the
# OLD engine stays the published active target until its unload() completes, then the
# new engine is published.
# The arbiter's diffusion evictor unloads get_active_diffusion_engine(); if the router published
# the new (empty) engine BEFORE the old one finished unloading, a concurrent acquire could evict
# the empty engine and take the GPU while the old model was still resident. Assert the OLD engine
# stays the published target until its unload() completes.
seen = {}
def _fake_engine():
@ -222,8 +220,8 @@ def test_switch_unloads_old_engine_before_publishing_new(monkeypatch):
def test_no_switch_keeps_engine_and_refreshes_reason(monkeypatch):
# When the engine does not change, _activate must not spuriously unload anything and must
# still refresh the recorded fallback reason (the diffusers-only steady state).
# When the engine does not change, _activate must not spuriously unload anything and must still
# refresh the recorded fallback reason.
calls = {"unload": 0}
def _fake_engine():
@ -286,10 +284,10 @@ def test_activate_serializes_switch_and_concurrent_query(monkeypatch):
def test_begin_load_on_refuses_an_engine_that_was_switched_away(monkeypatch):
# A load selects its engine, then yields (device probe, arbiter acquire) before registering.
# A second load choosing the OTHER engine transitions in that gap and unloads the still-idle
# engine this one captured, so registering there would leave a model that generate / status /
# unload and the arbiter's evictor can no longer reach (they all resolve the ACTIVE engine).
# A load selects its engine, then yields (device probe, arbiter acquire) before registering. A
# second load choosing the OTHER engine transitions in that gap and unloads the still-idle engine
# this one captured, so registering there would leave a model that generate / status / unload and
# the evictor can no longer reach (they all resolve the ACTIVE engine).
diffusers = SimpleNamespace(name = "diffusers")
sd_cpp = SimpleNamespace(name = "sd_cpp")
active = {"engine": diffusers}
@ -307,8 +305,8 @@ def test_begin_load_on_refuses_an_engine_that_was_switched_away(monkeypatch):
def test_begin_load_on_holds_the_transition_lock_while_registering(monkeypatch):
# The check and the registration must be one operation: taken under the same lock a switch
# takes, so no _activate can slip between them.
# The check and the registration must be one operation, taken under the same lock a switch takes,
# so no _activate can slip between them.
import threading
engine = SimpleNamespace(name = "diffusers")

4
studio/backend/tests/test_diffusion_gguf_compile.py
View file

@ -20,8 +20,8 @@ from core.inference import diffusion_gguf_compile as gc # noqa: E402
@pytest.fixture(autouse = True)
def _clean():
# Always start and end from a clean, unpatched state so tests do not leak the
# process-wide patch into each other.
# Always start and end from a clean, unpatched state so tests do not leak the process-wide patch
# into each other.
gc.uninstall_all()
yield
gc.uninstall_all()

8
studio/backend/tests/test_diffusion_inference_info.py
View file

@ -41,8 +41,8 @@ def test_component_sizes_match_the_table():
def test_quantised_estimate_is_below_bf16():
# A quantised transformer is smaller than bf16, so its resident estimate must be too
# (the companions are shared, and every steady factor is < 1).
# A quantised transformer is smaller than bf16, so its resident estimate must be too (the
# companions are shared, and every steady factor is below 1).
for info in family_inference_infos():
estimated = info["estimated_resident_gb"]
for scheme in _QUANT_STEADY_FACTOR:
@ -50,8 +50,8 @@ def test_quantised_estimate_is_below_bf16():
def test_nvfp4_is_below_int8():
# nvfp4 packs two params per byte (~0.33x) vs int8's one byte per param (~0.55x), so
# nvfp4's estimate is the smaller of the two on every family.
# nvfp4 packs two params per byte vs int8's one, so nvfp4's estimate is the smaller on every
# family.
for info in family_inference_infos():
estimated = info["estimated_resident_gb"]
assert estimated["nvfp4"] < estimated["int8"], info["family"]

43
studio/backend/tests/test_diffusion_krea2.py
View file

@ -44,7 +44,7 @@ def test_remap_rope_parameters_copies_5x_values():
def test_remap_rope_parameters_noop_on_5x_runtime_or_plain_4x_config():
# rope_scaling already parsed (5.x runtime exposing the alias): untouched.
# rope_scaling already parsed (a 5.x runtime exposing the alias): untouched.
parsed = {"rope_type": "default", "mrope_section": [1, 2, 3]}
cfg = SimpleNamespace(rope_scaling = parsed, rope_theta = 7.0, rope_parameters = {"x": 1})
remap_rope_parameters(cfg)
@ -110,8 +110,8 @@ def test_load_krea2_pipeline_threads_init_config(monkeypatch, tmp_path):
pipe = load_krea2_pipeline(str(tmp_path), "bf16")
# Turbo's fixed-mu schedule rides on is_distilled; dropping any of these would
# silently degrade generations, so the ctor kwargs are asserted exactly.
# Turbo's fixed-mu schedule rides on is_distilled; dropping any of these would silently degrade
# generations, so the ctor kwargs are asserted exactly.
assert captured["pipeline"]["is_distilled"] is True
assert captured["pipeline"]["patch_size"] == 2
assert captured["pipeline"]["text_encoder_select_layers"] == [2, 5, 8]
@ -126,8 +126,8 @@ def test_load_krea2_pipeline_threads_init_config(monkeypatch, tmp_path):
def test_load_krea2_pipeline_requires_krea_capable_diffusers(monkeypatch):
# On diffusers < 0.39 (no Krea2Pipeline) the loader must fail fast with the upgrade
# hint instead of dying with a bare AttributeError mid-load.
# On diffusers < 0.39 (no Krea2Pipeline) the loader must fail fast with the upgrade hint instead
# of dying with a bare AttributeError mid-load.
import pytest
fake = SimpleNamespace(__version__ = "0.38.0")
@ -147,19 +147,18 @@ def test_krea2_family_wiring():
fam = detect_family("krea/Krea-2-Turbo")
assert fam is not None and fam.name == KREA2_FAMILY_NAME
# Both vendor repos are non-GGUF allowlisted (Turbo for inference, Raw for training);
# no sd.cpp mapping -> diffusers fallback.
# Both vendor repos are non-GGUF allowlisted (Turbo for inference, Raw for training); no sd.cpp
# mapping, so diffusers fallback.
assert _is_trusted_diffusion_repo("krea/Krea-2-Turbo")
assert _is_trusted_diffusion_repo("krea/Krea-2-Raw")
assert not family_sd_cpp_supported(fam)
# Krea2TimestepEmbedding runs at M = batch; int8 (torch._int_mm, M > 16) must skip it.
# Krea2TimestepEmbedding runs at M = batch; int8 (torch._int_mm, M above 16) must skip it.
assert "time_embed" in exclude_tokens_for_scheme(TQ_INT8)
# Adapters train on Raw but run on Turbo, so the family carries a deploy override.
assert fam.deploy_base_repo == "krea/Krea-2-Turbo"
# The OpenAI /v1/images/generations route reads (steps, guidance) from this table; Krea
# Turbo is distilled (8 steps, no CFG), matching the Create UI seed instead of the
# generic (9, 0.0) fallback. Raw is the undistilled base (also inference-loadable) and runs
# its full 52-step / CFG 3.5 recipe, so its more specific key must win over the "krea" one.
# The OpenAI /v1/images/generations route reads (steps, guidance) from this table; Krea Turbo is
# distilled (8 steps, no CFG), matching the Create UI seed. Raw is the undistilled base and runs
# its full 52-step / CFG 3.5 recipe, so its more specific key must win over "krea".
assert default_generation_params("krea/Krea-2-Turbo") == (8, 0.0)
assert default_generation_params("krea/Krea-2-Raw") == (52, 3.5)
@ -185,13 +184,13 @@ def test_krea2_training_registry():
"resolution": 512,
}
info = {i["name"]: i for i in family_train_infos()}["krea-2"]
# Krea's guidance: train LoRAs on the undistilled Raw model, run them on Turbo, so
# Raw leads the training bases while Turbo stays available.
# Krea's guidance: train LoRAs on the undistilled Raw model and run them on Turbo, so Raw leads
# the training bases while Turbo stays available.
assert info["default_base"] == "krea/Krea-2-Raw"
assert info["base_repos"] == ["krea/Krea-2-Raw", "krea/Krea-2-Turbo"]
assert info["supports_compile"] is True
# Deploy previews the adapter on Turbo, not the Raw checkpoint it trained on, so the UI
# loads the distilled inference recipe; other families leave this None.
# Deploy previews the adapter on Turbo, not the Raw checkpoint it trained on, so the UI loads the
# distilled inference recipe; other families leave this None.
assert info["deploy_base"] == "krea/Krea-2-Turbo"
assert {i["name"]: i for i in family_train_infos()}["flux.1"]["deploy_base"] is None
@ -208,8 +207,8 @@ def test_krea2_spec_registered_with_authors_targets():
def test_krea2_collate_and_forward_roundtrip():
# spec.forward imports Krea2Pipeline (prepare_position_ids), so this needs a real
# diffusers install; CI hosts run the backend suite without one.
# spec.forward imports Krea2Pipeline (prepare_position_ids), so this needs a real diffusers
# install; CI hosts run the backend suite without one.
pytest.importorskip("diffusers")
import torch
from core.training.diffusion_dit_trainer import _SPECS
@ -229,8 +228,8 @@ def test_krea2_collate_and_forward_roundtrip():
class _FakeTransformer:
def __call__(self, **kwargs):
captured.update(kwargs)
# Echo the packed sequence: unpack(pack(x)) == x proves the inlined
# packing mirrors Krea2Pipeline exactly (they are mutual inverses).
# Echo the packed sequence: unpack(pack(x)) == x proves the inlined packing mirrors Krea2Pipeline
# exactly (they are mutual inverses).
return (kwargs["hidden_states"],)
noisy = torch.randn(2, 16, 1, 8, 8)
@ -239,8 +238,8 @@ def test_krea2_collate_and_forward_roundtrip():
_FakeTransformer(), noisy, timesteps, None, (pe_b, mask_b), None, "cpu", torch.float32
)
assert torch.equal(pred, noisy)
# [B, (H/2)*(W/2), C*4] patches, one shared [(txt+img), 3] position grid, and the
# [0, 1] timestep convention.
# [B, (H/2)*(W/2), C*4] patches, one shared [(txt+img), 3] position grid, and the [0, 1] timestep
# convention.
assert captured["hidden_states"].shape == (2, 16, 64)
assert captured["position_ids"].shape == (8 + 16, 3)
assert torch.allclose(captured["timestep"], torch.tensor([0.25, 0.75]))

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

@ -23,8 +23,8 @@ def test_sanitize_alias_strips_path_ext_and_unsafe_chars():
assert dl.sanitize_alias("owner/repo-name") == "repo-name"
assert dl.sanitize_alias("weird:<>chars.gguf") == "weird_chars"
assert dl.sanitize_alias("") == "lora"
# Internal dots (version tags like "V1.0") must be replaced: the alias becomes a
# diffusers PEFT adapter name and PEFT rejects "." in module/adapter names.
# Internal dots (version tags like "V1.0") must be replaced: the alias becomes a diffusers PEFT
# adapter name and PEFT rejects "." in adapter names.
assert (
dl.sanitize_alias("Qwen-Image-2512-Lightning-8steps-V1.0-bf16")
== "Qwen-Image-2512-Lightning-8steps-V1_0-bf16"
@ -42,16 +42,15 @@ def test_inject_prompt_tags_appends_with_spacing():
def test_inject_prompt_tags_validated_weight_overrides_user_typed():
r = dl.ResolvedLora("id", "style", "/p", "safetensors", 0.8)
# A user-typed tag for a SELECTED adapter is replaced by the backend-validated weight
# (so the recorded/validated 0-2 weight wins over whatever was typed), not duplicated.
# A user-typed tag for a SELECTED adapter is replaced by the backend-validated weight, not
# duplicated.
assert dl.inject_prompt_tags("a cat <lora:style:1>", [r]) == "a cat <lora:style:0.8>"
def test_inject_prompt_tags_strips_unselected_user_tags():
r = dl.ResolvedLora("id", "style", "/p", "safetensors", 0.8)
# A user tag for an alias that is NOT selected is stripped: only selected adapters are
# materialized in the managed --lora-model-dir, so sd-cli would drop the dead tag anyway;
# removing it keeps the prompt clean and unambiguous.
# materialized in the managed --lora-model-dir, so sd-cli would drop the dead tag anyway.
out = dl.inject_prompt_tags("a cat <lora:other:0.5>", [r])
assert "<lora:other:0.5>" not in out
assert out == "a cat <lora:style:0.8>"
@ -86,9 +85,9 @@ def test_supports_lora_matrix():
assert dl.supports_lora(
engine = "diffusers", family = "flux.1", model_kind = "single_file", transformer_quant = "int8"
)
# The quant fast path keeps the PICKER kind ("gguf") while the effective transformer is a
# dense torchao build, so the quant check must decide BEFORE the gguf-kind check; and the
# bake precedes compilation by construction, so compiled does not gate quant builds.
# The quant fast path keeps the PICKER kind ("gguf") while the effective transformer is a dense
# torchao build, so the quant check must decide BEFORE the gguf-kind check; the bake precedes
# compilation by construction, so compiled does not gate quant builds.
assert dl.supports_lora(
engine = "diffusers",
family = "z-image",
@ -105,8 +104,8 @@ def test_supports_lora_matrix():
assert not dl.supports_lora(
engine = "diffusers", family = "flux.1", model_kind = "gguf", transformer_quant = None
)
# A torch.compile'd diffusers transformer (Speed=default/max) can't take a non-hotswap
# adapter: diffusers needs the adapter loaded before compilation.
# A torch.compile'd diffusers transformer can't take a non-hotswap adapter: diffusers needs the
# adapter loaded before compilation.
assert not dl.supports_lora(
engine = "diffusers",
family = "flux.1",
@ -125,9 +124,8 @@ def test_supports_lora_matrix():
def test_resolve_specs_maps_cancelled_to_diffusion_sentinel(tmp_path, monkeypatch):
# A Hub download cancelled mid-flight raises RuntimeError("Cancelled"); resolve_specs
# must convert it to the diffusion cancellation sentinel so the route maps it to 409,
# not a generic 500 server-error toast.
# A Hub download cancelled mid-flight raises RuntimeError("Cancelled"); resolve_specs must convert
# it to the diffusion cancellation sentinel so the route maps it to 409, not a generic 500.
def _boom(spec_id, weight, **kw):
raise RuntimeError("Cancelled")
@ -192,8 +190,8 @@ def test_resolve_one_local_and_unknown(tmp_path, monkeypatch):
def test_resolve_one_rejects_cross_family_catalog_entry(tmp_path, monkeypatch):
# A family-tagged adapter must be rejected in the resolver (not just the UI picker) when the
# loaded model is a different family, so a direct API client cannot apply a mismatched LoRA.
# A family-tagged adapter must be rejected in the resolver (not just the UI picker) when the loaded
# model is a different family, so a direct API client cannot apply a mismatched LoRA.
d = tmp_path / "loras"
d.mkdir()
(d / "krea-style.safetensors").write_bytes(b"x")
@ -219,8 +217,8 @@ def test_resolve_specs_drops_zero_weight(tmp_path, monkeypatch):
def test_resolve_specs_maps_unknown_id_to_valueerror(tmp_path, monkeypatch):
# An unknown / stale id raises FileNotFoundError in resolve_one; resolve_specs must
# surface it as ValueError so the route returns 400, not a generic 500.
# An unknown / stale id raises FileNotFoundError in resolve_one; resolve_specs must surface it as
# ValueError so the route returns 400, not a generic 500.
d = tmp_path / "loras"
d.mkdir()
monkeypatch.setattr(dl, "loras_dir", lambda: d)
@ -229,15 +227,14 @@ def test_resolve_specs_maps_unknown_id_to_valueerror(tmp_path, monkeypatch):
def test_resolve_specs_maps_hub_error_to_valueerror(tmp_path, monkeypatch):
# A mistyped Hub repo id makes the Hub resolution raise a huggingface_hub client error
# (RepositoryNotFoundError, an HfHubHTTPError). resolve_specs must surface it as
# ValueError so the route returns 400, not a generic 500. The Hub message embeds the
# request URL, which must be scrubbed out of the client-facing 400.
# A mistyped Hub repo id makes the Hub resolution raise an HfHubHTTPError. resolve_specs must
# surface it as ValueError so the route returns 400, and the Hub message embeds the request URL,
# which must be scrubbed from the client-facing 400.
from huggingface_hub.errors import RepositoryNotFoundError
def _boom(spec_id, weight, **kw):
# response is optional in huggingface_hub 0.x but required in 1.x; both only
# read .headers / .request, so a stub keeps the test working on either.
# response is optional in huggingface_hub 0.x but required in 1.x; both only read .headers /
# .request, so a stub works on either.
raise RepositoryNotFoundError(
"404 Client Error. Repository Not Found for url: "
"https://huggingface.co/api/models/nope/nope (Request ID: abc)",
@ -252,8 +249,8 @@ def test_resolve_specs_maps_hub_error_to_valueerror(tmp_path, monkeypatch):
def test_scan_local_disambiguates_identical_stems(tmp_path, monkeypatch):
# foo.safetensors and foo.gguf must get distinct ids so each is addressable; a
# unique stem keeps its clean stem id.
# foo.safetensors and foo.gguf must get distinct ids so each is addressable; a unique stem keeps
# its clean stem id.
d = tmp_path / "loras"
d.mkdir()
(d / "foo.safetensors").write_bytes(b"x")
@ -268,8 +265,8 @@ def test_scan_local_disambiguates_identical_stems(tmp_path, monkeypatch):
def test_resolve_one_rejects_traversal_weight_name(tmp_path, monkeypatch):
# A client-supplied weight file with traversal / absolute path is rejected before it
# can reach the downloader (it must stay a plain filename inside the repo).
# A client-supplied weight file with traversal / absolute path is rejected before it can reach the
# downloader.
monkeypatch.setattr(dl, "loras_dir", lambda: tmp_path)
for bad in ("owner/name:../secret.safetensors", "owner/name:/etc/x.safetensors"):
with pytest.raises(ValueError):
@ -295,8 +292,8 @@ def test_lora_spec_and_request_validation():
LoraSpec(id = "a", weight = -0.1)
# default weight
assert LoraSpec(id = "a").weight == 1.0
# duplicate ids are rejected: repeating an id would load the same adapter as several
# distinct suffixed adapters and stack its effect past the per-adapter weight bound.
# Duplicate ids are rejected: repeating an id would load the same adapter as several distinct
# suffixed adapters and stack its effect past the per-adapter weight bound.
with pytest.raises(Exception):
DiffusionGenerateRequest(
prompt = "x", loras = [{"id": "a", "weight": 0.5}, {"id": "a", "weight": 1.0}]
@ -409,10 +406,9 @@ def test_diffusers_apply_clears_when_empty(monkeypatch):
def test_diffusers_apply_rejects_unsupported_quant():
# int8/fp8 pipes bake adapters at load time; a bake-less quant pipe cannot take one at
# generation time (frozen topology) and must direct the client to reload with the
# selection. nvfp4/mxfp8 are never baked, so the same reload error is unreachable there
# via the API (supports_lora blocks the load), but the backend path is shared.
# int8/fp8 pipes bake adapters at load time; a bake-less quant pipe cannot take one at generation
# time (frozen topology) and must direct the client to reload. nvfp4/mxfp8 are never baked, so the
# error is unreachable there via the API, but the backend path is shared.
import threading
pipe = _FakePipe()
@ -426,8 +422,8 @@ def test_diffusers_apply_rejects_unsupported_quant():
def test_diffusers_apply_rejects_gguf_adapter(monkeypatch):
# A .gguf adapter (discoverable in the shared catalog) cannot load on the diffusers
# engine; it must be rejected as a clean 400 before touching the pipe.
# A .gguf adapter (discoverable in the shared catalog) cannot load on the diffusers engine; it
# must be rejected as a clean 400 before touching the pipe.
import threading
monkeypatch.setattr(
@ -461,8 +457,8 @@ def test_scan_local_reads_family_sidecar(tmp_path, monkeypatch):
assert by_id["plain"].families == ()
assert by_id["plain"].weight_default == 1.0
# Family filter: the sdxl-tagged adapter is kept for sdxl and hidden for flux.1;
# the untagged one is always shown (unknown compatibility).
# Family filter: the sdxl-tagged adapter is kept for sdxl and hidden for flux.1; the untagged one
# is always shown (unknown compatibility).
sdxl_ids = {e.id for e in dl.list_loras(family = "sdxl")}
flux_ids = {e.id for e in dl.list_loras(family = "flux.1")}
assert "trained" in sdxl_ids and "plain" in sdxl_ids

91
studio/backend/tests/test_diffusion_lora_trainer.py
View file

@ -46,8 +46,8 @@ def test_discover_prefers_sidecar_then_metadata_then_instance(tmp_path):
def test_discover_sidecar_overrides_metadata_row(tmp_path):
# A per-image sidecar is the user's explicit edit and must win over a metadata row
# for the same image (the labeling grid writes sidecars).
# A per-image sidecar is the user's explicit edit and must win over a metadata row for the same
# image (the labeling grid writes sidecars).
_touch(tmp_path / "a.png")
(tmp_path / "metadata.jsonl").write_text(
json.dumps({"file_name": "a.png", "text": "from metadata"}) + "\n", encoding = "utf-8"
@ -58,8 +58,8 @@ def test_discover_sidecar_overrides_metadata_row(tmp_path):
def test_discover_empty_sidecar_suppresses_metadata_but_uses_instance_prompt(tmp_path):
# An empty sidecar tombstone must suppress the metadata caption yet leave the image uncaptioned
# so the dreambooth instance_prompt still applies (not drop the image).
# An empty sidecar tombstone must suppress the metadata caption yet leave the image uncaptioned so
# the dreambooth instance_prompt still applies (not drop the image).
_touch(tmp_path / "cat.png")
(tmp_path / "metadata.jsonl").write_text(
json.dumps({"file_name": "cat.png", "text": "old metadata caption"}) + "\n",
@ -97,8 +97,7 @@ def test_discover_reads_invalid_utf8_sidecar_as_tombstone(tmp_path):
def test_discover_null_metadata_caption_is_not_the_string_none(tmp_path):
# str(None) stored "None" as a real caption; a null row must fall through to the
# instance prompt instead.
# str(None) stored "None" as a real caption; a null row must fall through to the instance prompt.
_touch(tmp_path / "cat.png")
(tmp_path / "metadata.jsonl").write_text(
json.dumps({"file_name": "cat.png", "text": None}) + "\n", encoding = "utf-8"
@ -124,8 +123,8 @@ def test_discover_captions_jsonl_and_image_key(tmp_path):
def test_discover_tolerates_non_object_and_invalid_utf8_jsonl(tmp_path):
# A metadata.jsonl line that is valid JSON but not an object ([]/null/string/number) or malformed
# must be skipped per-line rather than crash the trainer in .get(); a valid row still resolves.
# A metadata.jsonl line that is valid JSON but not an object, or malformed, must be skipped
# per-line rather than crash the trainer in .get(); a valid row still resolves.
_touch(tmp_path / "x.png")
(tmp_path / "metadata.jsonl").write_text(
'[]\nnull\n"str"\n123\n{not json\n'
@ -134,8 +133,8 @@ def test_discover_tolerates_non_object_and_invalid_utf8_jsonl(tmp_path):
encoding = "utf-8",
)
assert discover_image_caption_pairs(tmp_path) == [(str(tmp_path / "x.png"), "hi")]
# Invalid UTF-8 in the metadata file must not raise; the file is skipped (image falls back to
# the instance prompt).
# Invalid UTF-8 in the metadata file must not raise; the file is skipped (image falls back to the
# instance prompt).
_touch(tmp_path / "y.png")
(tmp_path / "captions.jsonl").write_bytes(b"\xff\xfe not utf-8\n")
pairs = dict(discover_image_caption_pairs(tmp_path, instance_prompt = "fallback"))
@ -151,10 +150,9 @@ def test_discover_custom_caption_column(tmp_path):
def test_discover_verify_images_rejects_undecodable(tmp_path):
# verify_images (opt-in, enabled by the start route) rejects a corrupt / zero-byte image with
# a clear ValueError -> 400 BEFORE the route frees the resident GPU models, instead of letting
# the spawned trainer crash in PIL after the teardown. The trainers leave it off (default),
# since they decode every image anyway.
# verify_images (opt-in, enabled by the start route) rejects a corrupt / zero-byte image with a
# clear ValueError -> 400 BEFORE the route frees the resident GPU models, instead of letting the
# spawned trainer crash in PIL after teardown. The trainers leave it off (they decode anyway).
from PIL import Image
good = tmp_path / "good.png"
@ -165,7 +163,7 @@ def test_discover_verify_images_rejects_undecodable(tmp_path):
bad.write_bytes(b"")
(tmp_path / "bad.txt").write_text("broken", encoding = "utf-8")
# Default (verify off): the bad file is accepted (filename-only), matching trainer behavior.
# Default (verify off): the bad file is accepted, matching trainer behavior.
pairs = dict(discover_image_caption_pairs(tmp_path))
assert str(bad) in pairs and str(good) in pairs
@ -256,9 +254,9 @@ def test_config_normalized_lists_mxfp8_in_invalid_mode_error():
def test_config_normalized_krea2_requires_bf16_compute():
# krea-2 (like qwen-image / z-image) has fp32 RoPE/embedder internals that overflow fp16,
# so its DiT trains in bf16 only; fp16 must be refused up front by the route preflight,
# before it reserves training and evicts resident GPU models and the child trainer raises.
# krea-2 (like qwen-image / z-image) has fp32 RoPE/embedder internals that overflow fp16, so its
# DiT trains in bf16 only; fp16 must be refused by the route preflight, before it reserves
# training and evicts resident GPU models.
with pytest.raises(ValueError, match = "bf16"):
DiffusionLoraConfig(
base_model = "b",
@ -270,10 +268,9 @@ def test_config_normalized_krea2_requires_bf16_compute():
def test_force_bf16_families_matches_trainer_specs():
# The route-level bf16-only preflight set must list exactly the DiT families whose trainer
# spec sets force_bf16. If a force_bf16 family is missing from the set (as krea-2 was), an
# fp16 start passes the route preflight, reserves training + evicts resident models, and
# only the child trainer raises -- the evict-then-fail the preflight exists to prevent.
# The route-level bf16-only preflight set must list exactly the DiT families whose trainer spec
# sets force_bf16. A missing family (as krea-2 was) lets an fp16 start pass the preflight, reserve
# training and evict resident models, with only the child trainer raising.
from core.training.diffusion_dit_trainer import _SPECS
from core.training.diffusion_train_common import _FORCE_BF16_FAMILIES
assert _FORCE_BF16_FAMILIES == {fam for fam, spec in _SPECS.items() if spec.force_bf16}
@ -291,12 +288,12 @@ def test_resolve_train_steps_uses_train_steps_when_epochs_disabled():
def test_resolve_train_steps_epochs_ceil_over_batch_and_grad_accum():
# One epoch = ceil(N / (batch x grad_accum)) optimizer steps; num_epochs multiplies it.
# 10 images, batch 4, grad_accum 1 -> ceil(10/4)=3 steps/epoch.
# One epoch = ceil(N / (batch x grad_accum)) optimizer steps; num_epochs multiplies it. 10 images,
# batch 4, grad_accum 1 gives 3 steps/epoch.
assert resolve_train_steps(_cfg(num_epochs = 1, train_batch_size = 4), 10) == 3
assert resolve_train_steps(_cfg(num_epochs = 5, train_batch_size = 4), 10) == 15
# grad_accum widens the effective batch: 100 images, batch 2, grad_accum 3 -> per_step=6,
# ceil(100/6)=17 steps/epoch, 2 epochs -> 34.
# grad_accum widens the effective batch: 100 images, batch 2, grad_accum 3 gives per_step=6,
# 17 steps/epoch, 2 epochs = 34.
cfg = _cfg(num_epochs = 2, train_batch_size = 2, gradient_accumulation_steps = 3)
assert resolve_train_steps(cfg, 100) == 34
# An exact multiple does not round up: 8 images / batch 4 -> 2 steps/epoch.
@ -309,8 +306,8 @@ def test_resolve_train_steps_single_image_dataset():
def test_resolve_train_steps_caps_at_100000():
# The run length is capped at 100000 even for absurd epoch counts (matches the request
# model's train_steps ceiling), so a huge epochs x dataset never overflows the loop.
# The run length is capped at 100000 even for absurd epoch counts (matching the request model's
# ceiling), so a huge epochs x dataset never overflows the loop.
cfg = _cfg(num_epochs = 1000, train_batch_size = 1)
assert resolve_train_steps(cfg, 10_000) == 100000
@ -334,9 +331,9 @@ def test_config_from_dict_threads_num_epochs():
def test_normalized_rejects_piecewise_constant():
# piecewise_constant needs a step_rules string the trainers never supply, so get_scheduler()
# would crash in the trainer subprocess AFTER the resident GPU workloads are freed. It must be
# rejected up front (a clean ValueError -> 400), not accepted like the other schedulers.
# piecewise_constant needs a step_rules string the trainers never supply, so get_scheduler() would
# crash in the trainer subprocess AFTER the resident GPU workloads are freed. It must be rejected
# up front.
with pytest.raises(ValueError, match = "lr_scheduler"):
DiffusionLoraConfig(
base_model = "b", data_dir = "d", output_dir = "o", lr_scheduler = "piecewise_constant"
@ -344,7 +341,7 @@ def test_normalized_rejects_piecewise_constant():
def test_normalized_accepts_supported_schedulers():
# Every scheduler in the allow-list runs with only warmup/training steps (no extra required arg).
# Every scheduler in the allow-list runs with only warmup/training steps.
for sched in (
"linear",
"cosine",
@ -360,10 +357,8 @@ def test_normalized_accepts_supported_schedulers():
def test_api_scheduler_enum_never_advertises_a_rejected_scheduler():
# The request-model enum must not offer a scheduler that normalized() rejects: a client that
# picks it straight from the schema would get a 400. Every option the API advertises must be in
# the validation allow-list (this guards against the enum and allow-list drifting apart again,
# e.g. piecewise_constant left in one but removed from the other).
# The request-model enum must not offer a scheduler that normalized() rejects: every option the
# API advertises must be in the validation allow-list, so the two cannot drift apart again.
import typing
from core.training.diffusion_train_common import _LR_SCHEDULERS
@ -401,7 +396,7 @@ def test_config_rejects_zero_lora_alpha():
def test_config_rejects_nonpositive_snr_gamma():
# gamma <= 0 zeroes/inverts the min-SNR weight; None is the documented disable.
# A gamma at or below 0 zeroes/inverts the min-SNR weight; None is the documented disable.
with pytest.raises(ValueError, match = "snr_gamma"):
DiffusionLoraConfig(base_model = "b", data_dir = "d", output_dir = "o", snr_gamma = 0).normalized()
cfg = DiffusionLoraConfig(
@ -480,8 +475,8 @@ def test_config_rejects_nonpositive_learning_rate():
def test_config_rejects_untrainable_base_models():
# GGUF checkpoints and families without a trainer (Kontext editing, SD3) must fail at
# normalise time (an instant 400 via the API), not minutes later inside from_pretrained.
# GGUF checkpoints and families without a trainer must fail at normalise time (an instant 400),
# not minutes later inside from_pretrained.
for bad in (
"unsloth/FLUX.1-dev-GGUF",
"z-image-turbo-Q4_K_M.gguf",
@ -505,8 +500,8 @@ def test_config_resolves_dit_families():
def test_config_accepts_sdxl_and_unknown_base_models():
# SDXL names and unclassifiable custom names/paths must pass the guard (a wrong
# custom pick still fails cleanly in from_pretrained).
# SDXL names and unclassifiable custom names/paths must pass the guard (a wrong custom pick still
# fails cleanly in from_pretrained).
for ok in (
"stabilityai/stable-diffusion-xl-base-1.0",
"stabilityai/sdxl-turbo",
@ -610,8 +605,8 @@ def test_publish_writes_metadata_sidecar(tmp_path, monkeypatch):
def test_publish_does_not_clobber_same_name_adapter(tmp_path, monkeypatch):
# A retrain with the same adapter name must not overwrite a prior mirror: the second
# publish lands under a numeric suffix (my-style -> my-style-2), sidecar alongside it.
# A retrain with the same adapter name must not overwrite a prior mirror: the second publish lands
# under a numeric suffix, sidecar alongside it.
from pathlib import Path
from core.inference import diffusion_lora
@ -644,7 +639,7 @@ def test_publish_does_not_clobber_same_name_adapter(tmp_path, monkeypatch):
def test_config_rejects_bad_lr_scheduler():
# A typo'd scheduler ('constnat') must fail at normalize time, not later in the subprocess.
# A typo'd scheduler must fail at normalize time, not later in the subprocess.
with pytest.raises(ValueError, match = "lr_scheduler"):
DiffusionLoraConfig(
base_model = "b", data_dir = "d", output_dir = "o", lr_scheduler = "constnat"
@ -657,8 +652,8 @@ def test_config_rejects_bad_lr_scheduler():
def test_config_rejects_fp16_on_bf16_only_family():
# qwen-image / z-image are bf16-only: an fp16 request must be rejected before spawn,
# in normalized(), not only by the subprocess-side guard.
# qwen-image / z-image are bf16-only: an fp16 request must be rejected before spawn, in
# normalized(), not only by the subprocess-side guard.
for base in ("Tongyi-MAI/Z-Image-Turbo", "unsloth/Qwen-Image-2512-unsloth-bnb-4bit"):
with pytest.raises(ValueError, match = "bf16"):
DiffusionLoraConfig(
@ -675,8 +670,8 @@ def test_config_rejects_fp16_on_bf16_only_family():
def test_gguf_substring_does_not_reject_local_diffusers_dir(tmp_path):
# A local diffusers directory whose path merely contains 'gguf' is a valid training base
# (it carries model_index.json, not GGUF weights); the broad substring must not reject it.
# A local diffusers directory whose path merely contains 'gguf' is a valid training base (it
# carries model_index.json, not GGUF weights); the broad substring must not reject it.
from core.training.diffusion_train_common import resolve_trainable_family
local = tmp_path / "my-gguf-experiments" / "sdxl-finetune"

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

@ -70,11 +70,10 @@ def test_normalize_memory_mode_accepts_and_rejects():
def test_estimate_gguf_resident_mib_matches_packed_size():
# GGUF weights stay packed (uint8) on-device; diffusers dequantises per-matmul
# transiently, so the resident footprint ~= the on-disk size regardless of quant
# level (measured on Z-Image-Turbo: Q2_K 3.64->3.68 GiB, Q8_0 7.22->7.25 GiB). A
# small margin covers allocator overhead. The prior per-quant expansion over-
# estimated (Q2 ~7.6x) and forced needless offload on a roomy card.
# GGUF weights stay packed (uint8) on-device and diffusers dequantises per-matmul transiently, so
# the resident footprint is about the on-disk size regardless of quant level (measured on
# Z-Image-Turbo). A small margin covers allocator overhead; the prior per-quant expansion
# over-estimated and forced needless offload on a roomy card.
assert estimate_gguf_resident_mib(1000) == 1050
assert estimate_gguf_resident_mib(7220) == 7581
assert estimate_gguf_resident_mib(None) is None
@ -132,7 +131,7 @@ def test_unified_cuda_skips_offload_even_if_offload_capable():
def test_auto_resident_when_roomy():
# 80 GB card, ~16 GB model: fits with headroom -> stay resident (bit-identical).
# 80 GB card, ~16 GB model: fits with headroom, so stay resident (bit-identical).
plan = plan_diffusion_memory(
target = _target(),
device_memory = _discrete(80000),
@ -144,9 +143,8 @@ def test_auto_resident_when_roomy():
def test_auto_model_offload_on_tight_fit():
# 24 GB free -> reserve max(2048, 2400)=2400 -> budget 21600, 0.85*budget=18360.
# required = 16000+4000+1000 = 21000: over 0.85*budget but still under budget
# -> whole-module offload.
# 24 GB free -> reserve 2400 -> budget 21600, 0.85*budget = 18360. required = 21000: over
# 0.85*budget but under budget, so whole-module offload.
plan = plan_diffusion_memory(
target = _target(),
device_memory = _discrete(24000, 24000),
@ -159,8 +157,8 @@ def test_auto_model_offload_on_tight_fit():
def test_auto_group_offload_when_transformer_overflows_but_companions_fit():
# Big transformer pushes the resident total over budget, but the companions
# (text encoder + VAE) still fit -> stream the transformer (fast, moderate cut).
# A big transformer pushes the resident total over budget, but the companions still fit, so stream
# the transformer (fast, moderate cut).
plan = plan_diffusion_memory(
target = _target(),
device_memory = _discrete(8000, 8000),
@ -170,8 +168,8 @@ def test_auto_group_offload_when_transformer_overflows_but_companions_fit():
base_overhead_mib = 1000,
)
assert plan.offload_policy == OFFLOAD_GROUP
# Group keeps the VAE resident, so it uses exact slicing but NOT lossy tiling
# -> balanced stays bit-identical while still capping the offload footprint.
# Group keeps the VAE resident, so it uses exact slicing but NOT lossy tiling: balanced stays
# bit-identical while still capping the offload footprint.
assert plan.vae_slicing is True and plan.vae_tiling is False
@ -188,7 +186,7 @@ def test_auto_model_offload_when_companions_exceed_budget():
def test_auto_model_offload_when_companion_size_unknown():
# Without a companion estimate the planner can't prove group fits -> safest cut.
# Without a companion estimate the planner can't prove group fits, so it takes the safest cut.
plan = plan_diffusion_memory(
target = _target(),
device_memory = _discrete(8000, 8000),
@ -258,7 +256,7 @@ def test_fast_falls_back_to_model_offload_when_it_does_not_fit():
def test_explicit_cpu_offload_overrides_resident_auto_choice():
# Roomy GPU -> auto would stay resident, but cpu_offload=True forces offload.
# A roomy GPU would stay resident under auto, but cpu_offload=True forces offload.
plan = plan_diffusion_memory(
target = _target(),
device_memory = _discrete(80000),
@ -271,8 +269,8 @@ def test_explicit_cpu_offload_overrides_resident_auto_choice():
def test_explicit_memory_mode_wins_over_legacy_cpu_offload():
# The API documents memory_mode as overriding cpu_offload when set: fast +
# the legacy flag must stay resident, not silently downgrade to offload.
# The API documents memory_mode as overriding cpu_offload when set: fast + the legacy flag must
# stay resident, not silently downgrade to offload.
plan = plan_diffusion_memory(
target = _target(),
device_memory = _discrete(80000),
@ -407,9 +405,8 @@ def test_apply_model_offload_engages_offload_and_tiling():
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.
# enable_model_cpu_offload defaults to CUDA in diffusers, so on a non-CUDA accelerator (e.g. Intel
# XPU) the target device must be forwarded or diffusers offloads to the wrong backend.
pipe = _RecordingPipe()
apply_memory_plan(pipe, _plan(OFFLOAD_MODEL, tiling = False), device = "xpu")
assert pipe.offload_device == "xpu"
@ -441,8 +438,8 @@ def test_apply_vae_tiling_falls_back_to_vae_submodule():
def test_apply_group_falls_back_to_model_without_transformer():
# The recording pipe has no .transformer, so group offload can't engage and the
# applier falls back to whole-module offload, reporting the real policy.
# The recording pipe has no .transformer, so group offload can't engage and the applier falls back
# to whole-module offload, reporting the real policy.
pipe = _RecordingPipe()
effective, _ = apply_memory_plan(pipe, _plan(OFFLOAD_GROUP, tiling = True), device = "cuda")
assert effective == OFFLOAD_MODEL and "model_offload" in pipe.calls
@ -468,10 +465,9 @@ def _install_fake_torch_and_hooks(monkeypatch, apply_group_offloading):
def test_apply_group_partial_hooks_propagates_not_crash_fallback(monkeypatch):
# A dual-DiT pipe whose second transformer fails group offload AFTER the first installed
# hooks is left in a partial group-offload state that enable_model_cpu_offload rejects. The
# applier must PROPAGATE the failure (load fails with the real cause) instead of returning
# False and letting the caller's whole-module fallback crash on the partially-hooked pipe.
# A dual-DiT pipe whose second transformer fails group offload AFTER the first installed hooks is
# left in a partial state that enable_model_cpu_offload rejects. The applier must PROPAGATE the
# failure instead of returning False and letting the caller's fallback crash.
import core.inference.diffusion_memory as mem
calls = {"n": 0}
@ -494,8 +490,8 @@ def test_apply_group_partial_hooks_propagates_not_crash_fallback(monkeypatch):
def test_apply_group_single_transformer_failure_falls_back(monkeypatch):
# A single-DiT pipe whose group offload fails with NO hooks installed must still return
# False so the caller falls back cleanly to whole-module offload.
# A single-DiT pipe whose group offload fails with NO hooks installed must still return False so
# the caller falls back cleanly to whole-module offload.
import core.inference.diffusion_memory as mem
def _apply(module, **kw):
@ -511,9 +507,8 @@ def test_apply_group_single_transformer_failure_falls_back(monkeypatch):
def test_apply_group_fallback_enables_vae_tiling():
# A balanced/group plan keeps the VAE resident (tiling off); when group offload can't
# engage and we drop to whole-module offload, the applier must turn VAE tiling ON to
# cap the decode-time spike on what is now a low-VRAM path.
# A balanced/group plan keeps the VAE resident (tiling off); when group offload can't engage and
# we drop to whole-module offload, the applier must turn VAE tiling ON to cap the decode spike.
plan = _plan(OFFLOAD_GROUP, tiling = True)
assert plan.vae_tiling is False # group plan leaves tiling off by design
pipe = _RecordingPipe() # no .transformer -> group offload falls back to model
@ -533,8 +528,8 @@ def test_apply_sequential_offload():
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.
# 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, device = None):
raise RuntimeError("sequential offload not supported for this transformer")
@ -565,9 +560,8 @@ def test_apply_tolerates_pipe_without_vae_savers():
def test_settled_snapshot_takes_max_free_over_reads(monkeypatch):
# A transient foreign allocation can only SHRINK free, so the settled snapshot must
# reject a transient undercount (60 GB free on an idle 183 GB card) by keeping the max
# free across the retry reads. Measured incident: FLUX.2-dev int8 cold load.
# A transient foreign allocation can only SHRINK free, so the settled snapshot must reject a
# transient undercount (60 GB free on an idle 183 GB card) by keeping the max free across reads.
from core.inference import diffusion_memory as dm
reads = [
@ -613,8 +607,8 @@ def test_settled_snapshot_passthrough_off_cuda(monkeypatch):
def test_plan_fits_total_capacity():
# True exactly when required fits (total - reserve) * 0.85: the decline can then only
# stem from the instantaneous free reading, so a settled retry is worthwhile.
# True exactly when required fits (total - reserve) * 0.85: the decline can then only stem from
# the instantaneous free reading, so a settled retry is worthwhile.
from core.inference.diffusion_memory import plan_fits_total_capacity
def plan(
@ -627,9 +621,9 @@ def test_plan_fits_total_capacity():
device_memory = DeviceMemory("cuda", "cuda", kind, free_mib = 1, total_mib = total),
)
# FLUX.2-dev int8 incident numbers: 90,228 required on a 183,359 MiB card -> fits.
# FLUX.2-dev int8 incident numbers: 90,228 required on a 183,359 MiB card, so it fits.
assert plan_fits_total_capacity(plan(90_228, 183_359)) is True
# Larger than the capacity margin (0.85 * (183,359 - 18,335) = 140,270) -> no retry.
# Larger than the capacity margin (0.85 * (183,359 - 18,335) = 140,270), so no retry.
assert plan_fits_total_capacity(plan(150_000, 183_359)) is False
# Unknown sizes keep today's behaviour (no retry).
assert plan_fits_total_capacity(plan(None, 183_359)) is False

96
studio/backend/tests/test_diffusion_more_families.py
View file

@ -42,8 +42,8 @@ def test_detect_family_ideogram4_override():
def test_ideogram4_repos_are_trusted_non_gguf():
# The three official vendor pipelines load via from_pretrained, which is gated
# to the unsloth org + the explicit allowlist.
# The three official vendor pipelines load via from_pretrained, gated to the unsloth org + the
# explicit allowlist.
for rid in (
"ideogram-ai/ideogram-4-fp8",
"ideogram-ai/ideogram-4-nf4",
@ -64,21 +64,21 @@ def test_ideogram4_repos_are_trusted_non_gguf():
],
)
def test_detect_family_flux1_krea_dev(repo_id):
# Krea's FLUX.1-dev finetune keeps the exact dev layout, so it must resolve to the
# existing flux.1 family (FluxPipeline), never to krea-2 (a different arch).
# Krea's FLUX.1-dev finetune keeps the exact dev layout, so it must resolve to the existing flux.1
# family, never to krea-2 (a different arch).
fam = detect_family(repo_id)
assert fam is not None and fam.name == "flux.1"
assert fam.pipeline_class == "FluxPipeline"
def test_flux1_krea_dev_is_trusted_non_gguf():
# The gated official pipeline loads via from_pretrained -> needs the allowlist.
# The gated official pipeline loads via from_pretrained, so it needs the allowlist.
assert _is_trusted_diffusion_repo("black-forest-labs/FLUX.1-Krea-dev")
def test_flux1_krea_dev_generation_defaults():
# Model-card recipe: 28 steps at guidance 4.5. The generic "krea" key (Krea-2-Turbo's
# 8-step no-CFG shape) must NOT swallow it, and the krea-2 defaults must stay intact.
# Model-card recipe: 28 steps at guidance 4.5. The generic "krea" key (Turbo's 8-step no-CFG
# shape) must NOT swallow it, and the krea-2 defaults must stay intact.
assert default_generation_params("black-forest-labs/FLUX.1-Krea-dev") == (28, 4.5)
assert default_generation_params("QuantStack/FLUX.1-Krea-dev-GGUF") == (28, 4.5)
assert default_generation_params("krea/Krea-2-Turbo") == (8, 0.0)
@ -107,8 +107,8 @@ def test_detect_family_lumina2_repos(repo_id):
def test_detect_family_lumina2_override_and_next_rejected():
assert detect_family("x", override = "lumina-2").name == "lumina-2"
assert detect_family("x", override = "lumina2").name == "lumina-2"
# Lumina-Next is a DIFFERENT arch (LuminaText2ImgPipeline): it must stay unknown
# instead of resolving here and crashing mid-load.
# Lumina-Next is a DIFFERENT arch (LuminaText2ImgPipeline): it must stay unknown instead of
# resolving here and crashing mid-load.
assert detect_family("Alpha-VLLM/Lumina-Next-SFT-diffusers") is None
@ -119,8 +119,8 @@ def test_lumina2_is_trusted_non_gguf():
def test_lumina2_generation_defaults():
# Model-card recipe: 50 steps at guidance 4.0 (cfg_trunc_ratio is added by the
# backend generate call itself, not the defaults table).
# Model-card recipe: 50 steps at guidance 4.0 (cfg_trunc_ratio is added by the backend generate
# call itself, not the defaults table).
assert default_generation_params("Alpha-VLLM/Lumina-Image-2.0") == (50, 4.0)
@ -149,8 +149,8 @@ def test_lumina2_bf16_component_table_present():
[
"hunyuanvideo-community/HunyuanImage-2.1-Diffusers",
"QuantStack/HunyuanImage-2.1-GGUF",
# A local GGUF pick where the family keyword lives in the filename (QuantStack's
# actual naming drops the dash: the hunyuanimage2.1 alias covers it).
# A local GGUF pick where the family keyword lives in the filename (QuantStack's naming drops the
# dash, which the hunyuanimage2.1 alias covers).
"QuantStack/HunyuanImage-2.1-GGUF/HunyuanImage2.1-Q4_K_M.gguf",
],
)
@ -169,8 +169,8 @@ def test_detect_family_hunyuanimage21_repos(repo_id):
def test_detect_family_hunyuanimage21_override_and_30_still_excluded():
assert detect_family("x", override = "hunyuanimage-2.1").name == "hunyuanimage-2.1"
assert detect_family("x", override = "hunyuanimage2.1").name == "hunyuanimage-2.1"
# The HunyuanImage-3.0 structured exclusion must survive the 2.1 family: 3.0 has
# no diffusers pipeline and must stay unknown with its stated reason.
# The HunyuanImage-3.0 structured exclusion must survive the 2.1 family: 3.0 has no diffusers
# pipeline and must stay unknown with its stated reason.
assert detect_family("tencent/HunyuanImage-3.0") is None
assert excluded_model_reason("tencent/HunyuanImage-3.0") is not None
assert excluded_model_reason("hunyuanvideo-community/HunyuanImage-2.1-Diffusers") is None
@ -183,8 +183,8 @@ def test_hunyuanimage21_is_trusted_non_gguf():
def test_hunyuanimage21_generation_defaults():
# Card recipe: 50 steps; guidance feeds the call's distilled_guidance_scale (3.25
# default), while classifier-free guidance runs inside the repo's guider components.
# Card recipe: 50 steps; guidance feeds the call's distilled_guidance_scale, while classifier-free
# guidance runs inside the repo's guider components.
assert default_generation_params("hunyuanvideo-community/HunyuanImage-2.1-Diffusers") == (
50,
3.25,
@ -233,8 +233,8 @@ def test_detect_family_hidream_repos(repo_id):
def test_hidream_override_and_trust():
assert detect_family("x", override = "hidream-i1").name == "hidream-i1"
assert detect_family("x", override = "hidream").name == "hidream-i1"
# The three official repos load via from_pretrained -> allowlisted; the Llama TE4
# comes from the unsloth mirror, which the org prefix already trusts.
# The three official repos load via from_pretrained, so they are allowlisted; the Llama TE4 comes
# from the unsloth mirror, which the org prefix already trusts.
for rid in (
"HiDream-ai/HiDream-I1-Full",
"HiDream-ai/HiDream-I1-Dev",
@ -246,9 +246,8 @@ def test_hidream_override_and_trust():
def test_hidream_generation_defaults():
# Upstream inference.py: Full 50 steps / guidance 5; Dev and Fast are distilled and
# run guidance-free at 28 / 16 steps. The specific keys must beat the generic
# "hidream" (which also appears in the owner segment of every variant id).
# Upstream inference.py: Full 50 steps / guidance 5; Dev and Fast are distilled and run
# guidance-free at 28 / 16 steps. The specific keys must beat the generic "hidream".
assert default_generation_params("HiDream-ai/HiDream-I1-Full") == (50, 5.0)
assert default_generation_params("HiDream-ai/HiDream-I1-Dev") == (28, 0.0)
assert default_generation_params("HiDream-ai/HiDream-I1-Fast") == (16, 0.0)
@ -259,25 +258,23 @@ def test_hidream_bf16_component_table_present():
sizes = family_bf16_components_gb(fam)
assert sizes is not None
transformer_gb, encoders_gb, vae_gb = sizes
# 17B MoE DiT 34.2 GB; TEs = CLIP-L 0.5 + CLIP-G 2.8 + T5-XXL 9.5 from the repo plus
# the ~16 GB Llama TE4 assembled from the mirror -> ~28.8 GB.
# 17B MoE DiT 34.2 GB; TEs are CLIP-L 0.5 + CLIP-G 2.8 + T5-XXL 9.5 from the repo plus the ~16 GB
# Llama TE4 from the mirror, so ~28.8 GB.
assert 32.0 <= transformer_gb <= 37.0
assert 26.0 <= encoders_gb <= 32.0
assert vae_gb <= 0.5
def test_ideogram4_generation_defaults():
# Model-card settings: 48 steps, guidance 7 (the backend keeps the pipeline's
# recommended tapered schedule when the request matches exactly).
# Model-card settings: 48 steps, guidance 7 (the backend keeps the pipeline's recommended tapered
# schedule when the request matches exactly).
assert default_generation_params("ideogram-ai/ideogram-4-fp8") == (48, 7.0)
def test_ideogram4_bf16_reservation_table_present():
# The memory planner reserves this bf16 footprint for a narrow (fp8) ideogram-4 base even
# when the blob-cache estimate is absent (empty cache / a best-effort download probe that
# swallowed a transient HF error), so the ~54 GB pipeline never plans a resident placement
# it cannot fit. If this constant table ever went None, that fp8 OOM safeguard would
# silently disable, so pin that it is present and sums to the expected ~54 GB.
# The memory planner reserves this bf16 footprint for a narrow (fp8) ideogram-4 base even when the
# blob-cache estimate is absent, so the ~54 GB pipeline never plans a resident placement it cannot
# fit. If the table went None that safeguard would silently disable, so pin its presence and sum.
fam = detect_family("ideogram-ai/ideogram-4-fp8")
table = family_bf16_components_gb(fam, fam.base_repo)
assert table is not None
@ -289,15 +286,15 @@ def test_ideogram4_memory_table_counts_both_dits():
components = family_bf16_components_gb(fam)
assert components is not None
transformer_gb, text_encoders_gb, _vae_gb = components
# Two ~9.3B DiTs (conditional + unconditional) at bf16: well above one DiT's
# ~18.6 GB. A single-DiT entry here would let auto planning under-reserve and OOM.
# Two ~9.3B DiTs (conditional + unconditional) at bf16: well above one DiT's ~18.6 GB. A
# single-DiT entry would let auto planning under-reserve and OOM.
assert transformer_gb > 30.0
assert text_encoders_gb > 5.0
def test_hidream_prequant_wiring():
# Hosted int8/fp8 checkpoints (28/28 per-case gate pairs per scheme; int8 verified
# bit-identical to on-the-fly quantize) serve the family default base.
# Hosted int8/fp8 checkpoints (28/28 per-case gate pairs per scheme; int8 verified bit-identical
# to on-the-fly quantize) serve the family default base.
from core.inference.diffusion_families import family_prequant_repo
fam = detect_family("HiDream-ai/HiDream-I1-Full")
for scheme in ("int8", "fp8"):
@ -305,10 +302,9 @@ def test_hidream_prequant_wiring():
def test_hidream_quant_schemes_not_denied_and_no_extra_excludes():
# Measured on a B200 (outputs/hidream_smoke): int8 and fp8 both engage and render
# cleanly, including a 2-3 token prompt on int8 -- the routed MoE expert Linears
# only ever see the concatenated image+text token stream (M >> 16), so the
# torch._int_mm minimum never binds and no family exclude tokens are needed.
# Measured on a B200: int8 and fp8 both engage and render cleanly, including a 2-3 token prompt on
# int8 -- the routed MoE expert Linears only ever see the concatenated image+text stream (M well
# above 16), so the torch._int_mm minimum never binds and no family exclude tokens are needed.
from core.inference.diffusion_transformer_quant import (
_FAMILY_SCHEME_DENY,
_INT8_EXCLUDE_NAME_TOKENS,
@ -363,10 +359,9 @@ def test_list_loras_family_filter_gates_krea_entries():
# ── ideogram-4 fp8 transformer remap ─────────────────────────────────────────
def test_convert_fp8_state_dict_dequantizes_and_splits_qkv():
# The vendor fp8 transformer stores fused attention.qkv (Q/K/V rows stacked) +
# attention.o, each with a per-output-channel weight_scale; diffusers expects split
# to_q/to_k/to_v/to_out.0 with the scale already applied. The converter must undo
# both, or every attention weight loads wrong (garbage) and on meta (a load crash).
# The vendor fp8 transformer stores fused attention.qkv (Q/K/V rows stacked) + attention.o, each
# with a per-output-channel weight_scale; diffusers expects split to_q/to_k/to_v/to_out.0 with the
# scale already applied. The converter must undo both, or every attention weight loads wrong.
torch = pytest.importorskip("torch")
from core.inference.diffusion_ideogram4 import _convert_fp8_state_dict
@ -394,8 +389,8 @@ def test_convert_fp8_state_dict_dequantizes_and_splits_qkv():
}
out = _convert_fp8_state_dict(raw, hidden, torch.bfloat16)
# Every converted tensor is cast to the requested compute dtype (the load_state_dict
# copy would silently up/down-cast otherwise).
# Every converted tensor is cast to the requested compute dtype (the load_state_dict copy would
# silently up/down-cast otherwise).
assert all(t.dtype == torch.bfloat16 for t in out.values())
# Re-run in float32 for the exact value checks below (bf16 loses precision).
out = _convert_fp8_state_dict(raw, hidden, torch.float32)
@ -417,10 +412,9 @@ def test_convert_fp8_state_dict_dequantizes_and_splits_qkv():
def test_ideogram4_repo_is_fp8_detects_local_layout(tmp_path):
# A local mirror of the fp8 base never string-matches base_repo, so memory planning
# relies on this shard-header probe to reserve the bf16 footprint. The fp8 layout is
# marked by a companion ``*.weight_scale``; the bnb-4bit (nf4) mirror carries none and
# must read as not-fp8 so it stays (correctly) planned against its compressed bytes.
# A local mirror of the fp8 base never string-matches base_repo, so memory planning relies on this
# shard-header probe to reserve the bf16 footprint. The fp8 layout is marked by a companion
# ``*.weight_scale``; the bnb-4bit mirror carries none and must read as not-fp8.
torch = pytest.importorskip("torch")
st = pytest.importorskip("safetensors.torch")
@ -451,8 +445,8 @@ def test_ideogram4_repo_is_fp8_detects_local_layout(tmp_path):
def test_create_causal_mask_patch_is_self_disabling_and_idempotent():
# The patch adapts the pipeline's inputs_embeds kwarg to the installed transformers
# create_causal_mask signature; on a matching signature it must forward unchanged,
# and a second apply must not double-wrap.
# create_causal_mask signature; on a matching signature it forwards unchanged, and a second apply
# must not double-wrap.
pytest.importorskip("torch")
pytest.importorskip("diffusers")

51
studio/backend/tests/test_diffusion_precision.py
View file

@ -49,9 +49,8 @@ def _stub_torch(
torch.float16 = "float16"
if with_fp8:
torch.float8_e4m3fn = "float8_e4m3fn"
# _cast_fp8 skips nn.Embedding tables (skip_modules_classes) to keep prompt
# tokens full precision, and _keep_bf16_block_fqns walks for nn.ModuleList block
# stacks, so the stub torch must expose both.
# _cast_fp8 skips nn.Embedding tables to keep prompt tokens full precision, and
# _keep_bf16_block_fqns walks for nn.ModuleList block stacks, so the stub torch must expose both.
torch.nn = types.SimpleNamespace(
Embedding = type("Embedding", (), {}),
ModuleList = type("ModuleList", (list,), {}),
@ -69,7 +68,7 @@ def _stub_casters(monkeypatch, recorder):
hooks.apply_layerwise_casting = lambda module, **kw: recorder.append(("fp8", module))
monkeypatch.setitem(sys.modules, "diffusers.hooks", hooks)
monkeypatch.setitem(sys.modules, "diffusers.hooks.layerwise_casting", casting)
# torchao nvfp4 -- quantize_ now receives the vision-tower exclusion filter_fn; accept + ignore.
# torchao nvfp4: quantize_ now receives the vision-tower exclusion filter_fn; accept + ignore.
tq = types.ModuleType("torchao.quantization")
tq.quantize_ = lambda module, config, filter_fn = None: recorder.append(("nvfp4", module))
mx = types.ModuleType("torchao.prototype.mx_formats")
@ -206,8 +205,8 @@ def test_quantize_tolerates_caster_failure(monkeypatch):
def test_quantize_int8_uses_family_keep_bf16_schedule(monkeypatch):
# int8 for a family with a measured schedule routes to the selective caster with
# that family's (skip_first, skip_last); qwen-image keeps first+last 6 blocks bf16.
# int8 for a family with a measured schedule routes to the selective caster with that family's
# (skip_first, skip_last); qwen-image keeps first+last 6 blocks bf16.
_stub_torch(monkeypatch, cc = (10, 0))
calls: list = []
monkeypatch.setattr(
@ -221,8 +220,8 @@ def test_quantize_int8_uses_family_keep_bf16_schedule(monkeypatch):
def test_quantize_int8_unknown_family_falls_back_to_fp8(monkeypatch):
# A family without an int8 keep-bf16 schedule falls back to layerwise fp8 (logged),
# never silently running full int8 that would degrade the encoder.
# A family without an int8 keep-bf16 schedule falls back to layerwise fp8 (logged), never silently
# running full int8 that would degrade the encoder.
_stub_torch(monkeypatch, cc = (10, 0))
int8_calls: list = []
fp8_calls: list = []
@ -236,8 +235,8 @@ def test_quantize_int8_unknown_family_falls_back_to_fp8(monkeypatch):
def test_quantize_fp8_dynamic_uses_compute_caster(monkeypatch):
# fp8_dynamic routes to the torchao per-row compute caster (not the layerwise one)
# and needs no per-family schedule.
# fp8_dynamic routes to the torchao per-row compute caster (not the layerwise one) and needs no
# per-family schedule.
_stub_torch(monkeypatch, cc = (9, 0))
calls: list = []
monkeypatch.setattr(dp, "_cast_fp8_dynamic", lambda enc, tgt: calls.append(enc))
@ -257,10 +256,9 @@ def test_quantize_int8_unsupported_hw_is_noop(monkeypatch):
def test_quantize_te_skips_torchao_modes_under_offload(monkeypatch):
# The torchao modes (int8-with-schedule / fp8_dynamic / nvfp4) produce tensor subclasses that
# reject Module.to(), which an offload hook uses, so they must be skipped under offload (the DiT
# path skips torchao quant for the same reason). Hardware supports every mode here, so a None
# result proves the offload skip, not a capability gate; the casters fail if wrongly invoked.
# The torchao modes produce tensor subclasses that reject Module.to(), which an offload hook uses,
# so they must be skipped under offload. Hardware supports every mode here, so a None result
# proves the offload skip, not a capability gate; the casters fail if wrongly invoked.
_stub_torch(monkeypatch, cc = (10, 0))
monkeypatch.setattr(
dp, "_cast_fp8_dynamic", lambda *a: pytest.fail("torchao caster must not run")
@ -292,8 +290,8 @@ def test_keep_bf16_block_fqns_selects_first_and_last(monkeypatch):
torch = _stub_torch(monkeypatch)
module_list = torch.nn.ModuleList
layers = module_list([object() for _ in range(10)])
# A short stack (<= skip_first + skip_last) contributes nothing (keeping it all would
# leave no interior to quantise).
# A short stack (at most skip_first + skip_last) contributes nothing, since keeping it all would
# leave no interior to quantise.
short = module_list([object() for _ in range(4)])
enc = types.SimpleNamespace()
enc.named_modules = lambda: [("", enc), ("model.layers", layers), ("aux.blocks", short)]
@ -349,8 +347,8 @@ def _stub_transformer_quant(monkeypatch, captured):
def test_int8_filter_keeps_blocks_and_towers_dense(monkeypatch):
# The real selective closure: interior Linears quantise, but the kept first blocks,
# the vision tower, lm_head, and the encoder's fp32-kept modules (T5 "wo") stay bf16.
# The real selective closure: interior Linears quantise, but the kept first blocks, the vision
# tower, lm_head, and the encoder's fp32-kept modules (T5 "wo") stay bf16.
torch = _stub_torch(monkeypatch)
captured: dict = {}
_stub_transformer_quant(monkeypatch, captured)
@ -373,10 +371,9 @@ def test_int8_filter_keeps_blocks_and_towers_dense(monkeypatch):
def test_nvfp4_filter_keeps_vision_tower_dense(monkeypatch):
# Weight-only NVFP4 on a text encoder must exclude the VLM vision tower / lm_head / T5 "wo"
# like the int8 / fp8 torchao TE modes -- 4-bit-ing a qwen-image(-edit) Qwen2.5-VL image tower
# degrades the edit/image conditioning the sibling schemes deliberately protect. Before the fix
# _cast_nvfp4 quantised every nn.Linear (no filter_fn), so the tower was silently 4-bit.
# Weight-only NVFP4 on a text encoder must exclude the VLM vision tower / lm_head / T5 "wo" like
# the int8 / fp8 torchao TE modes, since 4-bit-ing a Qwen2.5-VL image tower degrades the edit
# conditioning. Before the fix _cast_nvfp4 quantised every nn.Linear (no filter_fn).
_stub_torch(monkeypatch)
captured: dict = {}
_stub_transformer_quant(monkeypatch, captured)
@ -434,9 +431,9 @@ class _FakeWeight:
def test_weight_zero_output_row_detection():
# A dead output row NaNs torchao's per-row fp8 (scale 0 -> 0/0); SDXL's
# text_encoder_2 (OpenCLIP bigG) really ships one in layers.2.self_attn.out_proj --
# measured: every fp8_dynamic SDXL render was black until the row is kept dense.
# A dead output row NaNs torchao's per-row fp8 (scale 0 -> 0/0); SDXL's text_encoder_2 really
# ships one in layers.2.self_attn.out_proj -- measured: every fp8_dynamic SDXL render was black
# until the row is kept dense.
zero_row = types.SimpleNamespace(weight = _FakeWeight([[0.1, 0.2], [0.0, 0.0]]))
dense = types.SimpleNamespace(weight = _FakeWeight([[0.1, 0.2], [0.3, 0.0]]))
assert dp._weight_has_zero_output_row(zero_row) is True
@ -457,8 +454,8 @@ def test_weight_zero_output_row_detection():
def test_fp8_dynamic_filter_skips_zero_row_linear(monkeypatch):
# The fp8_dynamic caster must leave a zero-output-row Linear dense while the rest
# of the encoder still quantises (a family-wide deny would forfeit the whole win).
# The fp8_dynamic caster must leave a zero-output-row Linear dense while the rest of the encoder
# still quantises (a family-wide deny would forfeit the whole win).
_stub_torch(monkeypatch)
captured: dict = {}
_stub_transformer_quant(monkeypatch, captured)

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

@ -76,8 +76,8 @@ def test_resolve_variant_base_picks_variant_repo():
def test_resolve_variant_base_falls_back_to_default():
# An unknown variant base (or no base at all) keeps the family default entry: the
# loader's base_model_id validation then refuses it and dense-quantises, as before.
# An unknown variant base (or none at all) keeps the family default entry: the loader's
# base_model_id validation then refuses it and dense-quantises, as before.
fam = _fam(
prequant_repos = (("int8", "org/default-fp8"),),
prequant_variant_repos = (("org/model-dev", "int8", "org/dev-fp8"),),
@ -118,9 +118,9 @@ def test_resolve_nothing_configured_is_none():
def test_local_prequant_path_ready(tmp_path, monkeypatch):
# The auto-policy planner budgets the small prequant plan only when a request-supplied
# path would actually load: present AND inside an allowlisted root. Missing or not
# allowlisted -> not ready, else the loader refuses it and rebuilds dense after evict.
# The auto-policy planner budgets the small prequant plan only when a request-supplied path would
# actually load: present AND inside an allowlisted root. Otherwise the loader refuses it and
# rebuilds dense after evicting.
import os
ckpt = tmp_path / "model.pt"
@ -135,10 +135,9 @@ def test_local_prequant_path_ready(tmp_path, monkeypatch):
# ── usable_prequant_source ───────────────────────────────────────────────────────
def test_usable_source_missing_path_is_none(tmp_path, monkeypatch):
# An allowlisted but ABSENT request-supplied path must not count as a prequant
# source: load_prequantized_transformer would find no file and fall back to the
# dense bf16 build after the resident pipeline was already evicted, so the memory
# planner must run the dense fit checks up front instead.
# An allowlisted but ABSENT request-supplied path must not count as a prequant source:
# load_prequantized_transformer would find no file and fall back to the dense bf16 build after the
# resident pipeline was already evicted, so the planner must run the dense fit checks up front.
import os
monkeypatch.setattr(pq, "_allowed_prequant_roots", lambda: [os.path.realpath(str(tmp_path))])
@ -148,9 +147,8 @@ def test_usable_source_missing_path_is_none(tmp_path, monkeypatch):
def test_usable_source_disallowed_path_is_none(tmp_path, monkeypatch):
# A path OUTSIDE the UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH allowlist (including the
# default empty allowlist) is refused by the loader, so it must resolve to None
# here even when the file exists.
# A path OUTSIDE the UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH allowlist (including the default empty one)
# is refused by the loader, so it must resolve to None here even when the file exists.
ckpt = tmp_path / "model.pt"
ckpt.write_bytes(b"x")
monkeypatch.setattr(pq, "_allowed_prequant_roots", lambda: [])
@ -159,8 +157,8 @@ def test_usable_source_disallowed_path_is_none(tmp_path, monkeypatch):
def test_usable_source_allowed_present_path_wins(tmp_path, monkeypatch):
# Allowlisted AND present: the override is usable and takes priority over the
# hosted repo, exactly like resolve_prequant_source.
# Allowlisted AND present: the override is usable and takes priority over the hosted repo, exactly
# like resolve_prequant_source.
import os
ckpt = tmp_path / "model.pt"
@ -277,9 +275,8 @@ def _load(
):
_FakeTransformer.calls = {}
_stub_torch_accelerate(monkeypatch, ckpt, load_raises = load_raises)
# The local-path branch is opt-in via a directory ALLOWLIST (it unpickles an arbitrary
# file); these tests exercise the load mechanics, so allowlist tmp_path (where ckpt.pt
# lives) unless a test is checking the gate.
# The local-path branch is opt-in via a directory ALLOWLIST (it unpickles an arbitrary file);
# these tests exercise the load mechanics, so allowlist tmp_path unless a test checks the gate.
if allow_local:
monkeypatch.setenv(pq.ALLOW_LOCAL_PREQUANT_PATH_ENV, str(tmp_path))
else:
@ -314,8 +311,8 @@ def test_load_meta_init_and_assign(monkeypatch, tmp_path):
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.
# 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
@ -345,8 +342,8 @@ def test_load_base_mismatch_is_none(monkeypatch, tmp_path):
def test_load_fp8_stale_per_tensor_is_rejected(monkeypatch, tmp_path):
# A pre-fix fp8 checkpoint has no fp8_granularity (old per-tensor layout); it must be
# rejected so the loader rebuilds instead of reproducing the noise failure.
# A pre-fix fp8 checkpoint has no fp8_granularity (old per-tensor layout); it must be rejected so
# the loader rebuilds instead of reproducing the noise failure.
stale = _good_ckpt(scheme = "fp8")
del stale["metadata"]["fp8_granularity"]
assert _load(monkeypatch, tmp_path, stale, scheme = "fp8") is None
@ -362,16 +359,16 @@ def test_load_int8_ignores_fp8_granularity(monkeypatch, tmp_path):
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.
# A checkpoint whose keys happen to match a different base can load strict=True and then render
# from the wrong weights, so one requested with a base but recording none 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.
# 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
@ -391,8 +388,8 @@ def test_load_fast_accum_auto_ignores_baked(monkeypatch, tmp_path):
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.
# An int8 checkpoint recording a stale exclusion set (which 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
@ -407,11 +404,10 @@ def test_load_exclude_tokens_match_ok(monkeypatch, tmp_path):
def test_load_exclude_tokens_need_the_recorded_family(monkeypatch, tmp_path):
# int8 carries PER-FAMILY exclusions (Qwen's unpadded text stream runs at M = prompt tokens,
# under _int_mm's M > 16 floor). An offline artifact that recorded the family but built its
# exclusion set with family=None baked those linears as int8, so the loader must reject it --
# and accept only the family-aware set. Pins the offline builder
# (scripts/build_prequant_checkpoint.py) to exclude_tokens_for_scheme(scheme, fam.name).
# int8 carries PER-FAMILY exclusions (Qwen's unpadded text stream runs at M = prompt tokens, under
# _int_mm's M floor of 16). An artifact that recorded the family but built its exclusion set with
# family=None baked those linears as int8, so the loader must reject it and accept only the
# family-aware set. Pins the offline builder to exclude_tokens_for_scheme(scheme, fam.name).
from core.inference.diffusion_transformer_quant import exclude_tokens_for_scheme
for family in ("qwen-image", "qwen-image-edit"):
family_less = _good_ckpt(scheme = "int8")
@ -428,8 +424,8 @@ def test_load_exclude_tokens_need_the_recorded_family(monkeypatch, tmp_path):
def test_load_require_bf16_mismatch_is_none(monkeypatch, tmp_path):
# An fp8 (scaled_mm) checkpoint built WITHOUT the bf16 gate quantised a different layer set
# than the runtime filter now produces, so it must be rejected rather than loaded.
# An fp8 (scaled_mm) checkpoint built WITHOUT the bf16 gate quantised a different layer set than
# the runtime filter now produces, so it must be rejected.
ckpt = _good_ckpt(scheme = "fp8")
ckpt["metadata"]["require_bf16"] = False
assert _load(monkeypatch, tmp_path, ckpt, scheme = "fp8") is None
@ -442,32 +438,32 @@ def test_load_require_bf16_match_ok(monkeypatch, tmp_path):
def test_load_require_bf16_int8_true_is_none(monkeypatch, tmp_path):
# int8 (torch._int_mm) tolerates non-bf16 weights, so it never sets the gate; a checkpoint
# claiming it did contradicts the runtime filter and must be rejected.
# int8 (torch._int_mm) tolerates non-bf16 weights, so it never sets the gate; a checkpoint claiming
# it did contradicts the runtime filter and must be rejected.
ckpt = _good_ckpt(scheme = "int8")
ckpt["metadata"]["require_bf16"] = True
assert _load(monkeypatch, tmp_path, ckpt, scheme = "int8") is None
def test_load_require_bf16_nvfp4_false_ok(monkeypatch, tmp_path):
# nvfp4 quantises fp32 weights fine, so the runtime filter does NOT set the bf16 gate; a
# checkpoint built the same way (require_bf16=False) matches and loads.
# nvfp4 quantises fp32 weights fine, so the runtime filter does NOT set the bf16 gate; a checkpoint
# built the same way matches and loads.
ckpt = _good_ckpt(scheme = "nvfp4")
ckpt["metadata"]["require_bf16"] = False
assert _load(monkeypatch, tmp_path, ckpt, scheme = "nvfp4") is not None
def test_load_require_bf16_nvfp4_true_is_none(monkeypatch, tmp_path):
# An nvfp4 checkpoint claiming the bf16 gate contradicts the runtime filter (nvfp4 is not gated),
# so it quantised a different layer set and must be rejected.
# An nvfp4 checkpoint claiming the bf16 gate contradicts the runtime filter, so it quantised a
# different layer set and must be rejected.
ckpt = _good_ckpt(scheme = "nvfp4")
ckpt["metadata"]["require_bf16"] = True
assert _load(monkeypatch, tmp_path, ckpt, scheme = "nvfp4") is 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.
# 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"
@ -540,8 +536,8 @@ def test_load_repo_source_allowed_without_optin(monkeypatch, tmp_path):
def test_load_repo_source_falls_back_to_legacy_filename(monkeypatch, tmp_path):
# A repo still carrying the legacy transformer_<scheme>.pt name serves the download after
# the model-name filename 404s; both names are requested in order.
# A repo still carrying the legacy transformer_<scheme>.pt name serves the download after the
# model-name filename 404s; both names are requested in order.
_FakeTransformer.calls = {}
_stub_torch_accelerate(monkeypatch, _good_ckpt())
monkeypatch.delenv(pq.ALLOW_LOCAL_PREQUANT_PATH_ENV, raising = False)
@ -593,8 +589,8 @@ def test_load_repo_source_falls_back_to_legacy_filename(monkeypatch, tmp_path):
def test_load_local_path_outside_allowlist_refused(monkeypatch, tmp_path):
# Even with the opt-in set, a path OUTSIDE every allowlisted directory must not be
# unpickled: enabling one trusted dir is not a wildcard for arbitrary request paths.
# Even with the opt-in set, a path OUTSIDE every allowlisted directory must not be unpickled:
# enabling one trusted dir is not a wildcard for arbitrary request paths.
called = {"load": False}
def _explode(*a, **k):
@ -627,8 +623,8 @@ def test_load_local_path_outside_allowlist_refused(monkeypatch, tmp_path):
def test_load_min_features_mismatch_is_none(monkeypatch, tmp_path):
# A checkpoint built with a different --min-features quantises a different Linear set,
# so it must be rejected when the runtime threshold is supplied.
# A checkpoint built with a different --min-features quantises a different Linear set, so it must
# be rejected when the runtime threshold is supplied.
ckpt = _good_ckpt()
ckpt["metadata"]["min_features"] = 256 # built with 256, runtime asks for 512
_FakeTransformer.calls = {}

177
studio/backend/tests/test_diffusion_routes.py
View file

@ -24,8 +24,8 @@ from routes.inference import studio_router
class _FakeBackend:
def __init__(self) -> None:
self.loaded = False
# Repo ids of in-flight (not yet committed) loads; empty tuple = none. The unload
# route reads this to keep DIFFUSION ownership while a concurrent load is still loading.
# Repo ids of in-flight (not yet committed) loads; empty tuple = none. The unload route reads this
# to keep DIFFUSION ownership while a concurrent load is still loading.
self.loading: tuple = ()
@property
@ -44,8 +44,8 @@ class _FakeBackend:
model_kind = None,
base_repo = None,
):
# Mirror the real backend's cheap validation so the route's
# validate-before-evict ordering is exercised.
# Mirror the real backend's cheap validation so the route's validate-before-evict ordering is
# exercised.
from core.inference.diffusion import resolve_model_kind
from core.inference.diffusion_families import detect_family
@ -57,8 +57,8 @@ class _FakeBackend:
raise ValueError(
f"Non-GGUF diffusion loads are restricted to unsloth/* repos; got '{model_path}'."
)
# A client-supplied base_repo clears the same trust bar (mirrors the real backend's
# gate), so the route's validate-before-evict rejects an untrusted companion base.
# A client-supplied base_repo clears the same trust bar as the real backend, so the route's
# validate-before-evict rejects an untrusted companion base.
if base_repo and base_repo.strip() and not base_repo.lower().startswith("unsloth/"):
raise ValueError(
f"base_repo is restricted to unsloth/* repos (or a local path); got '{base_repo}'."
@ -106,8 +106,8 @@ class _FakeBackend:
if not self.loaded:
raise RuntimeError("No diffusion model is loaded.")
if prompts is not None or seeds is not None:
# List-driven batch: the LIST sets the image count and each image's own seed
# (batch_size is only a per-forward cap), exactly as the real engine reports it.
# List-driven batch: the LIST sets the image count and each image's own seed (batch_size is only a
# per-forward cap), exactly as the real engine reports it.
base = seeds[0] if seeds else (seed if seed is not None else 4242)
count = len(prompts) if prompts is not None else len(seeds)
per_image = seeds if seeds is not None else [base + i for i in range(count)]
@ -117,8 +117,8 @@ class _FakeBackend:
"seeds": list(per_image),
"repo_id": "x/z-image",
}
# The real backend returns the PIL images; the route persists them. The
# fake returns sentinels since image_gallery is stubbed in the fixture.
# The real backend returns the PIL images and the route persists them; the fake returns sentinels
# since image_gallery is stubbed in the fixture.
return {
"images": [object() for _ in range(batch_size)],
"seed": seed if seed is not None else 4242,
@ -153,14 +153,13 @@ def _unloaded_status():
def client(monkeypatch, tmp_path):
backend = _FakeBackend()
monkeypatch.setattr(diffusion_module, "get_diffusion_backend", lambda: backend)
# Neutralise the engine router so the routes deterministically drive this fake
# (diffusers) backend regardless of the host's real device, and never attempt a
# native sd.cpp install/download. The router's selection logic is covered in
# test_diffusion_engine_router.py; one route-level sd_cpp test lives below.
# Neutralise the engine router so the routes deterministically drive this fake (diffusers) backend
# regardless of the host's real device, and never attempt a native sd.cpp install. The selection
# logic is covered in test_diffusion_engine_router.py.
import core.inference.diffusion_engine_router as engine_router
# Delegate to whatever get_diffusion_backend currently returns, so per-test
# re-patches of the backend still flow through the routes.
# Delegate to whatever get_diffusion_backend currently returns, so per-test re-patches of the
# backend still flow through the routes.
monkeypatch.setattr(
engine_router,
"select_and_activate_engine",
@ -173,14 +172,14 @@ def client(monkeypatch, tmp_path):
)
monkeypatch.setattr(engine_router, "_active_engine_name", "diffusers")
monkeypatch.setattr(engine_router, "_fallback_reason", None)
# Isolate from the real GPU arbiter: reset ownership and stub the evictors so
# the load route's acquire_for() never touches live backend singletons.
# Isolate from the real GPU arbiter: reset ownership and stub the evictors so the load route's
# acquire_for() never touches live backend singletons.
monkeypatch.setattr(gpu_arbiter, "_owner", None)
monkeypatch.setitem(gpu_arbiter._EVICTORS, gpu_arbiter.CHAT, lambda: None)
monkeypatch.setitem(gpu_arbiter._EVICTORS, gpu_arbiter.DIFFUSION, lambda: None)
# In-memory gallery backed by tmp files, so routes exercise persistence wiring
# without PIL/real disk under studio_root.
# In-memory gallery backed by tmp files, so routes exercise persistence wiring without PIL/real
# disk under studio_root.
store: dict[str, dict] = {}
def _save(image, meta):
@ -215,8 +214,8 @@ def client(monkeypatch, tmp_path):
"image_path",
lambda i: (tmp_path / f"{i}.png") if i in store else None,
)
# The serve route resolves through owned_image_path (ownership-gated); the fake store only ever
# holds owned records, so a stem not in it is treated as foreign and refused, like the real guard.
# The serve route resolves through owned_image_path; the fake store only holds owned records, so a
# stem not in it is treated as foreign and refused, like the real guard.
monkeypatch.setattr(
gallery_module,
"owned_image_path",
@ -268,15 +267,14 @@ def test_load_generate_status_unload_roundtrip(client):
def test_gallery_serve_refuses_unowned_id(client):
# The serve route resolves through the ownership guard, so a guessed stem for a PNG the gallery
# does not own (no record) is a 404, not a stream of foreign bytes.
# does not own is a 404, not a stream of foreign bytes.
assert client.get("/api/inference/images/gallery/family-photo/file").status_code == 404
def test_generate_holds_progress_active_during_persist(client, monkeypatch):
# generate-progress must stay active while a finished generation is still writing its gallery
# record, so a concurrent reload's mount probe keeps polling instead of refreshing the gallery
# before the image lands. Probe the persist counter from inside the save call, and confirm it
# is cleared afterwards.
# early. Probe the persist counter from inside the save call, and confirm it clears afterwards.
import core.inference.image_gallery as gallery_module
import routes.inference as inf
@ -310,9 +308,9 @@ def test_generate_holds_progress_active_during_persist(client, monkeypatch):
def test_load_rejects_untrusted_base_repo(client):
# A trusted GGUF model_path paired with an untrusted remote base_repo is rejected at the
# route (validate runs before the GPU handoff), so an authenticated client cannot make the
# server fetch and deserialize an arbitrary companion repo, and no model is loaded/evicted.
# A trusted GGUF model_path paired with an untrusted remote base_repo is rejected at the route
# (validate runs before the GPU handoff), so an authenticated client cannot make the server fetch
# and deserialize an arbitrary companion repo.
r = client.post(
"/api/inference/images/load",
json = {
@ -327,9 +325,9 @@ def test_load_rejects_untrusted_base_repo(client):
def test_unload_keeps_ownership_when_a_model_is_still_resident(client, monkeypatch):
# The unload route must drop DIFFUSION ownership only when nothing is resident. If a
# concurrent load re-established residency while the (slow) unload ran, releasing would
# clear the newer load's claim and a later chat load would skip eviction and OOM.
# The unload route must drop DIFFUSION ownership only when nothing is resident. If a concurrent
# load re-established residency while the slow unload ran, releasing would clear the newer claim
# and a later chat load would skip eviction and OOM.
backend = diffusion_module.get_diffusion_backend()
gpu_arbiter._owner = gpu_arbiter.DIFFUSION
@ -349,11 +347,9 @@ def test_unload_keeps_ownership_when_a_model_is_still_resident(client, monkeypat
def test_unload_keeps_ownership_when_a_load_is_in_flight(client, monkeypatch):
# A concurrent /images/load re-acquires DIFFUSION and starts a background load, so the
# engine is NOT is_loaded yet (the pipeline commits later) but a load IS in flight. The
# unload route must keep ownership on the in-flight state alone, or a later chat load would
# see no owner, skip eviction, and OOM against the newly resident pipeline. is_loaded stays
# False the whole download/finalize window, so the loaded-only check is insufficient here.
# A concurrent /images/load re-acquires DIFFUSION and starts a background load, so the engine is
# NOT is_loaded yet but a load IS in flight. The unload route must keep ownership on the in-flight
# state alone, since is_loaded stays False for the whole download/finalize window.
backend = diffusion_module.get_diffusion_backend()
gpu_arbiter._owner = gpu_arbiter.DIFFUSION
@ -384,9 +380,9 @@ def test_generate_batch_size_persists_each_image(client):
def test_generate_seed_list_records_replay_from_each_own_seed(client):
# A seeds LIST sets each image's own seed, so the recipe must NOT claim the base seed +
# the request's batch_size: restore prefers batch_seed (frontend restoreSettings does
# `batch_seed ?? seed`), which would regenerate seed 5 for the seed-99 image.
# A seeds LIST sets each image's own seed, so the recipe must NOT claim the base seed + the
# request's batch_size: restore prefers batch_seed (`batch_seed ?? seed`), which would regenerate
# seed 5 for the seed-99 image.
client.post(
"/api/inference/images/load", json = {"model_path": "x/z-image", "gguf_filename": "q.gguf"}
)
@ -418,8 +414,8 @@ def test_generate_prompt_list_records_each_prompt_and_seed(client):
def test_generate_legacy_batch_still_records_the_base_seed_and_size(client):
# The batch_size path is unchanged: those images DO share one base seed, so restore
# must replay the whole batch (base seed + batch_size), not the derived per-image seed.
# The batch_size path is unchanged: those images DO share one base seed, so restore must replay
# the whole batch, not the derived per-image seed.
client.post(
"/api/inference/images/load", json = {"model_path": "x/z-image", "gguf_filename": "q.gguf"}
)
@ -434,10 +430,9 @@ def test_generate_legacy_batch_still_records_the_base_seed_and_size(client):
def test_generate_request_rejects_zero_denoise_strength():
# strength 0 does NOT keep the source: every diffusers img2img/inpaint pipeline derives
# t_start = steps - int(steps * strength), so 0 leaves zero denoising steps (FLUX/Qwen/
# Z-Image raise "the number of pipeline steps is 0 which is < 1"; SDXL img2img has no
# such guard and crashes on empty latents). Reject it as a 422 up front.
# strength 0 does NOT keep the source: every diffusers img2img/inpaint pipeline derives t_start =
# steps - int(steps * strength), so 0 leaves zero denoising steps (FLUX/Qwen/Z-Image raise, SDXL
# img2img crashes on empty latents). Reject it as a 422 up front.
import pydantic
from models.inference import DiffusionGenerateRequest
@ -464,8 +459,8 @@ def test_generate_rejects_non_multiple_of_16(client):
client.post(
"/api/inference/images/load", json = {"model_path": "x/z-image", "gguf_filename": "q.gguf"}
)
# Odd, and a multiple of 8 that isn't a multiple of 16: both rejected, since
# Z-Image requires dimensions divisible by 16.
# Odd, and a multiple of 8 that isn't a multiple of 16: both rejected, since Z-Image requires
# dimensions divisible by 16.
for bad in (1001, 1000):
resp = client.post("/api/inference/images/generate", json = {"prompt": "p", "width": bad})
assert resp.status_code == 422, bad
@ -478,8 +473,8 @@ def test_generate_rejects_batch_seed_past_json_safe_range(client):
client.post(
"/api/inference/images/load", json = {"model_path": "x/z-image", "gguf_filename": "q.gguf"}
)
# A seed at the cap with a batch derives per-image seeds (seed+1 ...) past the JSON-safe range,
# so the request is rejected.
# A seed at the cap with a batch derives per-image seeds past the JSON-safe range, so the request
# is rejected.
over = client.post(
"/api/inference/images/generate",
json = {"prompt": "p", "seed": 2**53 - 1, "batch_size": 2},
@ -494,16 +489,16 @@ def test_generate_rejects_batch_seed_past_json_safe_range(client):
def test_non_gguf_load_restricted_to_unsloth(client):
# gguf_filename is optional now; with none, the load is a full-pipeline kind, which
# is gated to unsloth/* repos. A non-unsloth repo (no filename) is rejected -> 400.
# gguf_filename is optional now; with none the load is a full-pipeline kind, gated to unsloth/*
# repos, so a non-unsloth repo is rejected with a 400.
resp = client.post("/api/inference/images/load", json = {"model_path": "x/z-image"})
assert resp.status_code == 400
assert "unsloth" in resp.json()["detail"].lower()
def test_pipeline_load_allowed_for_unsloth_repo(client):
# An unsloth/* repo with no filename loads as a full diffusers pipeline (kind auto
# = pipeline); the route forwards model_kind="pipeline" to begin_load.
# An unsloth/* repo with no filename loads as a full diffusers pipeline, so the route forwards
# model_kind="pipeline" to begin_load.
resp = client.post(
"/api/inference/images/load", json = {"model_path": "unsloth/Z-Image-Turbo-unsloth-bnb-4bit"}
)
@ -519,8 +514,8 @@ def test_generate_without_load_returns_409(client):
def test_generate_pipeline_error_returns_sanitized_500(client, monkeypatch):
# A loaded model that fails mid-pipeline (CUDA OOM, a RuntimeError) is a server
# failure: 500 with a generic message, not a 409 echoing the raw exception.
# A loaded model that fails mid-pipeline (CUDA OOM, a RuntimeError) is a server failure: 500 with
# a generic message, not a 409 echoing the raw exception.
backend = diffusion_module.get_diffusion_backend()
backend.loaded = True
@ -535,9 +530,8 @@ def test_generate_pipeline_error_returns_sanitized_500(client, monkeypatch):
def test_generate_execution_error_with_cancelled_substring_is_sanitized_500(client, monkeypatch):
# A native sd-cli execution failure whose raw tail merely CONTAINS "cancelled"
# must stay a sanitized 500, not misroute to 409 and echo that output (path/arg
# leak). Regression: the handler matched "cancelled" as a substring.
# A native sd-cli execution failure whose raw tail merely CONTAINS "cancelled" must stay a
# sanitized 500, not misroute to 409 and echo that output (path/arg leak).
backend = diffusion_module.get_diffusion_backend()
backend.loaded = True
@ -570,8 +564,8 @@ def test_load_unknown_family_returns_400(client, monkeypatch):
raise ValueError("'x/y' isn't a supported image-generation model. Supported: Z-Image.")
backend = _FakeBackend()
# Validation runs in the pre-flight (before the GPU is taken), so that is
# where an unsupported model is rejected now.
# Validation runs in the pre-flight (before the GPU is taken), so that is where an unsupported
# model is rejected now.
backend.validate_load_request = _raise
monkeypatch.setattr(diffusion_module, "get_diffusion_backend", lambda: backend)
resp = client.post(
@ -582,8 +576,8 @@ def test_load_unknown_family_returns_400(client, monkeypatch):
def test_load_validation_failure_does_not_evict_chat(client, monkeypatch):
# A rejected image-model pick must not tear down the user's loaded chat model:
# validation runs before acquire_for, so chat keeps the GPU on a 400.
# A rejected image-model pick must not tear down the user's loaded chat model: validation runs
# before acquire_for, so chat keeps the GPU on a 400.
monkeypatch.setattr(gpu_arbiter, "_owner", gpu_arbiter.CHAT)
evicted = []
monkeypatch.setitem(gpu_arbiter._EVICTORS, gpu_arbiter.CHAT, lambda: evicted.append(True))
@ -604,8 +598,8 @@ def test_load_validation_failure_does_not_evict_chat(client, monkeypatch):
def test_load_refused_during_training_does_not_evict_chat(client, monkeypatch):
# An image load while training is active is refused (409) before the GPU is
# taken, so the training run and the loaded chat model are both untouched.
# An image load while training is active is refused (409) before the GPU is taken, so the training
# run and the loaded chat model are both untouched.
import core.training as core_training
monkeypatch.setattr(gpu_arbiter, "_owner", gpu_arbiter.CHAT)
@ -649,8 +643,8 @@ def test_routes_require_auth():
def test_invalid_family_returns_400_without_evicting_chat(client):
# An undetectable family fails validation BEFORE the GPU handoff, so the
# arbiter is never acquired and a loaded chat model would not be evicted.
# An undetectable family fails validation BEFORE the GPU handoff, so the arbiter is never acquired
# and a loaded chat model would not be evicted.
resp = client.post(
"/api/inference/images/load", json = {"model_path": "x/y", "gguf_filename": "q.gguf"}
)
@ -752,9 +746,8 @@ def test_invalid_attention_backend_returns_422(client):
def test_prequant_path_doc_describes_allowlist_not_toggle():
# The field help must match the code: UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH is a
# directory allowlist, not a =1 toggle (diffusion_prequant._allowed_prequant_roots
# drops bare on/off tokens), so operators following the doc don't get every
# The field help must match the code: UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH is a directory allowlist,
# not a =1 toggle (bare on/off tokens are dropped), so operators following the doc don't get every
# request silently refused.
from models.inference import DiffusionLoadRequest
@ -830,8 +823,7 @@ def test_load_routes_to_sd_cpp_on_cpu(monkeypatch, tmp_path):
lambda: SimpleNamespace(backend = "cpu", device = "cpu"),
)
monkeypatch.setattr(engine_router, "ensure_sd_cpp_binary", lambda **_: "/x/sd-cli")
# The router now probes runnability before committing to native; treat the stub
# binary as executable.
# The router probes runnability before committing to native; treat the stub binary as executable.
monkeypatch.setattr(
engine_router, "SdCppEngine", lambda **_: SimpleNamespace(version = lambda: "sd-cli v0")
)
@ -862,8 +854,8 @@ def test_load_routes_to_sd_cpp_on_cpu(monkeypatch, tmp_path):
def test_invalid_transformer_quant_returns_422_without_eviction(client):
# An unsupported transformer_quant is rejected by the request schema (Literal), so
# the GPU is never acquired and no chat model is evicted.
# An unsupported transformer_quant is rejected by the request schema (Literal), so the GPU is
# never acquired and no chat model is evicted.
resp = client.post(
"/api/inference/images/load",
json = {"model_path": "x/z-image", "gguf_filename": "q.gguf", "transformer_quant": "int2"},
@ -873,8 +865,8 @@ def test_invalid_transformer_quant_returns_422_without_eviction(client):
def test_invalid_memory_mode_returns_422_without_eviction(client):
# An unsupported memory_mode is rejected by the request schema (Literal), so the
# GPU is never acquired and no chat model is evicted.
# An unsupported memory_mode is rejected by the request schema (Literal), so the GPU is never
# acquired and no chat model is evicted.
resp = client.post(
"/api/inference/images/load",
json = {"model_path": "x/z-image", "gguf_filename": "q.gguf", "memory_mode": "ultra"},
@ -890,8 +882,8 @@ def test_in_progress_returns_409_after_validation_passes(client, monkeypatch):
backend = _FakeBackend()
backend.begin_load = _busy
monkeypatch.setattr(diffusion_module, "get_diffusion_backend", lambda: backend)
# Pin the resolved device to cuda: the route only takes the arbiter for non-CPU
# loads, so on a CPU-only host the ownership assert below would never hold.
# Pin the resolved device to cuda: the route only takes the arbiter for non-CPU loads, so on a
# CPU-only host the ownership assert below would never hold.
import types as _types
import core.inference.diffusion_device as devmod
@ -934,8 +926,8 @@ def _force_engine(monkeypatch, backend, *, engine_name, device):
def test_cpu_native_load_skips_gpu_arbiter(client, monkeypatch):
# A native sd.cpp load on a pure-CPU host never touches the GPU, so the route must NOT
# evict the resident chat model -- the arbiter handoff is skipped.
# A native sd.cpp load on a pure-CPU host never touches the GPU, so the route must NOT evict the
# resident chat model: the arbiter handoff is skipped.
from core.inference.sd_cpp_engine import ENGINE_SD_CPP
backend = diffusion_module.get_diffusion_backend()
@ -948,8 +940,8 @@ def test_cpu_native_load_skips_gpu_arbiter(client, monkeypatch):
def test_gpu_native_load_takes_arbiter(client, monkeypatch):
# A force-native sd.cpp load on a GPU box DOES use the GPU, so the arbiter is acquired
# (same as the always-GPU diffusers path).
# A force-native sd.cpp load on a GPU box DOES use the GPU, so the arbiter is acquired, like the
# always-GPU diffusers path.
from core.inference.sd_cpp_engine import ENGINE_SD_CPP
backend = diffusion_module.get_diffusion_backend()
@ -962,8 +954,8 @@ def test_gpu_native_load_takes_arbiter(client, monkeypatch):
def test_images_info_lists_every_family(client):
# The pure info endpoint is hardware-independent (no load required): it returns one
# entry per auto-policy family, each with the quant estimates the UI shows.
# The pure info endpoint is hardware-independent (no load required): one entry per auto-policy
# family, each with the quant estimates the UI shows.
from core.inference.diffusion_auto_policy import _FAMILY_BF16_GB
resp = client.get("/api/inference/images/info")
@ -972,14 +964,14 @@ def test_images_info_lists_every_family(client):
assert {f["family"] for f in families} == set(_FAMILY_BF16_GB)
sample = families[0]
est = sample["estimated_resident_gb"]
# Quantised estimates undercut bf16, and nvfp4 undercuts int8 (matches the pure helper).
# Quantised estimates undercut bf16, and nvfp4 undercuts int8 (matching the pure helper).
assert est["int8"] < est["bf16"]
assert est["nvfp4"] < est["int8"]
def test_status_passes_through_resolved(client, monkeypatch):
# The additive `resolved` provenance record round-trips through the status route so the
# frontend can render the "Auto: X" badges.
# The additive `resolved` provenance record round-trips through the status route so the frontend
# can render the "Auto: X" badges.
backend = diffusion_module.get_diffusion_backend()
resolved = {
"speed_mode": {"value": "eager", "source": "auto", "reason": "per-kind default"},
@ -998,18 +990,17 @@ def test_status_passes_through_resolved(client, monkeypatch):
def test_status_resolved_defaults_to_null(client):
# A backend status without a `resolved` key leaves the additive field null (older
# backends and the unloaded state).
# A backend status without a `resolved` key leaves the additive field null (older backends and the
# unloaded state).
body = client.get("/api/inference/images/status").json()
assert body["resolved"] is None
def test_download_plan_forwards_the_load_time_controls(client, monkeypatch):
# The plan drives the staged download, so it has to be computed from the SAME
# configuration the load will run with. The dense-quant prefetch decision reads the
# memory policy, the prequant path and the adapter selection as well as speed/quant:
# dropping them stages the base transformer/ shards for a low-VRAM load that never
# opens them, or omits them for a baked-LoRA load that does.
# The plan drives the staged download, so it must be computed from the SAME configuration the load
# will run with. The dense-quant prefetch decision reads the memory policy, the prequant path and
# the adapter selection as well as speed/quant: dropping them stages the base transformer/ shards
# for a low-VRAM load that never opens them, or omits them for a baked-LoRA load that does.
backend = diffusion_module.get_diffusion_backend()
seen: dict = {}

51
studio/backend/tests/test_diffusion_sdxl.py
View file

@ -54,8 +54,8 @@ def test_sdxl_detection_by_repo_and_override():
def test_dit_families_keep_transformer_denoiser():
# The generalisation must not change existing DiT families: they stay on
# pipe.transformer and their single file is transformer-only.
# The generalisation must not change existing DiT families: they stay on pipe.transformer and
# their single file is transformer-only.
for rid in ("unsloth/FLUX.1-schnell-GGUF", "unsloth/Qwen-Image-GGUF", "unsloth/Z-Image-GGUF"):
fam = detect_family(rid)
assert fam.denoiser_attr == "transformer"
@ -63,8 +63,8 @@ def test_dit_families_keep_transformer_denoiser():
def test_sdxl_has_no_native_sd_cpp_mapping():
# No single-file VAE/TE mapping yet, so the no-GPU route falls back to diffusers
# rather than trying to drive sd-cli.
# No single-file VAE/TE mapping yet, so the no-GPU route falls back to diffusers rather than
# trying to drive sd-cli.
assert family_sd_cpp_supported(detect_family("stabilityai/sdxl-turbo")) is False
@ -72,9 +72,8 @@ def test_sdxl_base_repos_are_trusted_non_gguf():
# Official safetensors-only base repos are allowlisted so their catalog entries load.
assert _is_trusted_diffusion_repo("stabilityai/stable-diffusion-xl-base-1.0")
assert _is_trusted_diffusion_repo("stabilityai/sdxl-turbo")
# The refiner is img2img-only and is intentionally NOT allowlisted (see
# test_sdxl_refiner_not_trusted).
# Case-insensitive match.
# The refiner is img2img-only and intentionally NOT allowlisted (see
# test_sdxl_refiner_not_trusted). Case-insensitive match.
assert _is_trusted_diffusion_repo("StabilityAI/SDXL-Turbo")
# A random repo (even one that detects as SDXL) is NOT trusted for a non-GGUF load.
assert not _is_trusted_diffusion_repo("randomorg/my-sdxl-merge")
@ -82,8 +81,8 @@ def test_sdxl_base_repos_are_trusted_non_gguf():
def test_sdxl_model_kind_resolution():
# A full-pipeline load (no single-file name) is "pipeline"; a single .safetensors
# is "single_file" (handled by the whole-pipeline branch for SDXL).
# A full-pipeline load (no single-file name) is "pipeline"; a single .safetensors is
# "single_file" (handled by the whole-pipeline branch for SDXL).
assert resolve_model_kind(None) == "pipeline"
assert resolve_model_kind("sdxl.safetensors") == "single_file"
@ -102,9 +101,9 @@ class _FakeVae:
def test_align_vae_dtype_uses_unet_denoiser():
# For SDXL the denoiser lives at pipe.unet; _align_vae_dtype must read it (a pipe
# with only .unet and no .transformer) and cast the VAE to the U-Net's dtype. The
# dtype is read from a parameter (denoiser has no .dtype), so use a _FakeVae denoiser.
# For SDXL the denoiser lives at pipe.unet; _align_vae_dtype must read it (a pipe with only .unet)
# and cast the VAE to the U-Net's dtype. The dtype comes from a parameter, so use a _FakeVae
# denoiser.
import torch
vae = _FakeVae(dtype = torch.float32)
@ -130,10 +129,9 @@ def test_align_vae_dtype_transformer_default_unchanged():
def test_align_vae_dtype_skips_gguf_packed_uint8_params():
# A GGUF-quantized transformer's leading parameters are packed uint8 storage; the
# dtype probe must skip them and use the first FLOATING dtype, or nn.Module.to()
# rejects the integer dtype and an Edit/img2img call 500s (regression: Qwen-Image-
# Edit GGUF). All-integer params (no floating dtype at all) must be a clean no-op.
# A GGUF-quantized transformer's leading parameters are packed uint8 storage, so the dtype probe
# must skip them and use the first FLOATING dtype, or nn.Module.to() rejects the integer dtype and
# an Edit/img2img call 500s. All-integer params must be a clean no-op.
import torch
class _GgufDenoiser:
@ -156,7 +154,7 @@ def test_align_vae_dtype_skips_gguf_packed_uint8_params():
def test_sdxl_lora_supported_on_diffusers():
# SDXL is bf16/bnb-4bit on diffusers -> LoRA is allowed (unlike GGUF-via-diffusers).
# SDXL is bf16/bnb-4bit on diffusers, so LoRA is allowed (unlike GGUF-via-diffusers).
assert diffusion_lora.supports_lora(
engine = "diffusers", family = "sdxl", model_kind = "pipeline", transformer_quant = None
)
@ -166,10 +164,9 @@ def test_sdxl_lora_supported_on_diffusers():
def test_pipeline_prefetch_skips_non_torch_artifacts():
# The SDXL Base repo ships fp16 variants, ONNX, OpenVINO and Flax exports next to
# the default safetensors; from_pretrained (no variant kwarg) loads only the
# default torch weights, so the prefetch filter must skip everything else or a
# catalog load pulls tens of GB of unused artifacts.
# The SDXL Base repo ships fp16 variants, ONNX, OpenVINO and Flax exports next to the default
# safetensors; from_pretrained loads only the default torch weights, so the prefetch filter must
# skip everything else or a catalog load pulls tens of GB of unused artifacts.
from core.inference.diffusion import _pipeline_file_downloaded as keep
assert keep("model_index.json")
@ -186,8 +183,8 @@ def test_pipeline_prefetch_skips_non_torch_artifacts():
def test_sdxl_refiner_not_trusted():
# The refiner is an img2img-only pipeline; the sdxl family loads every repo as the
# base txt2img pipeline, so the refiner must NOT be allowlisted for a non-GGUF load.
# The refiner is an img2img-only pipeline, and the sdxl family loads every repo as the base
# txt2img pipeline, so it must NOT be allowlisted for a non-GGUF load.
assert not _is_trusted_diffusion_repo("stabilityai/stable-diffusion-xl-refiner-1.0")
# The base and turbo remain trusted.
assert _is_trusted_diffusion_repo("stabilityai/stable-diffusion-xl-base-1.0")
@ -195,8 +192,8 @@ def test_sdxl_refiner_not_trusted():
def test_sdxl_gguf_load_rejected_up_front():
# SDXL has no transformer-only GGUF variant (its single file is the whole pipeline),
# so a GGUF request must fail cheap validation before the GPU handoff.
# SDXL has no transformer-only GGUF variant (its single file is the whole pipeline), so a GGUF
# request must fail cheap validation before the GPU handoff.
backend = DiffusionBackend()
with pytest.raises(ValueError, match = "no GGUF"):
backend.validate_load_request(
@ -205,8 +202,8 @@ def test_sdxl_gguf_load_rejected_up_front():
def test_base_config_filter_skips_weights():
# For a whole-pipeline single file, the base repo supplies only config/tokenizer, not
# its (unused) weight tensors.
# For a whole-pipeline single file, the base repo supplies only config/tokenizer, not its (unused)
# weight tensors.
from core.inference.diffusion import _base_config_file_downloaded as keep
assert keep("model_index.json")

96
studio/backend/tests/test_diffusion_speed.py
View file

@ -93,8 +93,8 @@ def test_resolve_speed_mode_gguf_auto_default():
assert resolve_speed_mode("off", is_gguf = True) == SPEED_OFF
assert resolve_speed_mode("max", is_gguf = True) == SPEED_MAX
assert resolve_speed_mode("max", is_gguf = False) == SPEED_MAX
# The video backend passes a dense default of `default` (clips amortise the
# compile within one run); it must not affect GGUF or explicit values.
# The video backend passes a dense default of `default` (clips amortise the compile within one
# run); it must not affect GGUF or explicit values.
assert resolve_speed_mode(None, is_gguf = False, dense_default = SPEED_DEFAULT) == SPEED_DEFAULT
assert resolve_speed_mode("off", is_gguf = False, dense_default = SPEED_DEFAULT) == SPEED_OFF
@ -106,7 +106,7 @@ def test_compile_eligible_requires_bf16_cuda_friendly(monkeypatch):
_stub_torch(monkeypatch)
# The happy path: bf16, CUDA, compile-friendly family.
assert compile_eligible(_target(), is_gguf = False, family = _family()) is True
# GGUF is now compile-eligible too (measured ~2.3x, PSNR ~37 dB vs eager).
# GGUF is compile-eligible too (measured ~2.3x, PSNR ~37 dB vs eager).
assert compile_eligible(_target(), is_gguf = True, family = _family()) is True
# fp16 (non-bf16) is excluded.
assert compile_eligible(_target(dtype = "float16"), is_gguf = False, family = _family()) is False
@ -139,8 +139,8 @@ def test_restore_backend_flags_tolerates_none():
def test_snapshot_partial_when_some_backends_missing(monkeypatch):
# A build/platform without cuda.matmul (e.g. CPU/MPS) must still snapshot + restore the
# flags it does have, rather than skipping the whole snapshot on one missing attribute.
# A build/platform without cuda.matmul (e.g. CPU/MPS) must still snapshot + restore the flags it
# does have, rather than skipping the whole snapshot on one missing attribute.
torch = types.ModuleType("torch")
torch.backends = types.SimpleNamespace(
cuda = types.SimpleNamespace(), # no .matmul
@ -231,14 +231,13 @@ def test_speed_off_applies_nothing(monkeypatch):
"fp16_accum": False,
}
assert pipe.vae.mem_format is None and pipe.compiled is False
# off must not touch any process-wide flag (bit-identical reference path).
# off must not touch any process-wide flag (the bit-identical reference path).
assert torch.backends.cudnn.benchmark is False
def test_speed_compiles_both_dits_for_dual_dit_family(monkeypatch):
# A dual-DiT family (Ideogram: transformer + unconditional_transformer) runs BOTH DiTs each
# denoise step, so the regional block compile must engage on both, not just the first --
# otherwise the second DiT runs eager while status reports compile as engaged.
# A dual-DiT family runs BOTH DiTs each denoise step, so the regional block compile must engage on
# both -- otherwise the second runs eager while status reports compile as engaged.
_stub_torch(monkeypatch)
_stub_gguf_accel(monkeypatch)
pipe = _Pipe(with_compile = True, with_second_dit = True)
@ -250,8 +249,8 @@ def test_speed_compiles_both_dits_for_dual_dit_family(monkeypatch):
def test_speed_default_dense_falls_back_to_regional_compile(monkeypatch):
# A DENSE model has no GGUF dequant to compile, so `default` falls back to the
# regional block compile (its only compile lever) -- and no GGUF accelerators.
# A DENSE model has no GGUF dequant to compile, so `default` falls back to the regional block
# compile (its only compile lever), and no GGUF accelerators.
torch = _stub_torch(monkeypatch)
called = _stub_gguf_accel(monkeypatch)
pipe = _Pipe(with_compile = True)
@ -260,8 +259,8 @@ def test_speed_default_dense_falls_back_to_regional_compile(monkeypatch):
)
assert applied["channels_last"] is True and pipe.vae.mem_format == torch.channels_last
assert applied["compiled"] is True and pipe.compiled is True
# default compiles with dynamic=True and no autotune mode (fast cold start,
# resolution-robust, sidesteps the CUDA-graph crash).
# default compiles with dynamic=True and no autotune mode (fast cold start, resolution-robust,
# sidesteps the CUDA-graph crash).
assert pipe.compile_kwargs == {"fullgraph": True, "dynamic": True}
# default also autotunes the VAE convs but does NOT flip TF32 or fuse QKV.
assert applied["cudnn_benchmark"] is True and torch.backends.cudnn.benchmark is True
@ -272,10 +271,9 @@ def test_speed_default_dense_falls_back_to_regional_compile(monkeypatch):
def test_offload_active_drops_fullgraph(monkeypatch):
# Group/model/sequential offload installs a torch.compiler.disable'd onload hook;
# compiling with fullgraph=True then crashes at the first denoise step. Same reason
# as an active step cache -> fullgraph must drop to False when offload is planned.
# (Dense model: on this branch GGUF `default` takes the compiled-dequant path.)
# Group/model/sequential offload installs a torch.compiler.disable'd onload hook, so compiling with
# fullgraph=True crashes at the first denoise step -- same reason as an active step cache, so
# fullgraph must drop to False when offload is planned.
_stub_torch(monkeypatch)
pipe = _Pipe(with_compile = True)
applied = apply_speed_optims(
@ -291,8 +289,8 @@ def test_offload_active_drops_fullgraph(monkeypatch):
def test_speed_default_gguf_compiles_only_dequant(monkeypatch):
# GGUF `default` is the LIGHT path: compile ONLY the dequant op chain, NOT the
# regional block compile.
# GGUF `default` is the LIGHT path: compile ONLY the dequant op chain, NOT the regional block
# compile.
_stub_torch(monkeypatch)
called = _stub_gguf_accel(monkeypatch)
pipe = _Pipe(with_compile = True)
@ -307,9 +305,8 @@ def test_speed_default_gguf_compiles_only_dequant(monkeypatch):
def test_speed_eager_gguf_installs_no_accelerator(monkeypatch):
# eager = lossless-but-no-compile: neither the compiled dequant nor the regional
# block compile run; only the process-wide lossless levers (channels_last, cudnn)
# and the shared/per-arch eager monkey-patches (installed elsewhere) engage.
# eager = lossless-but-no-compile: neither the compiled dequant nor the regional block compile
# run; only the process-wide lossless levers and the eager monkey-patches engage.
_stub_torch(monkeypatch)
called = _stub_gguf_accel(monkeypatch)
pipe = _Pipe(with_compile = True)
@ -322,8 +319,8 @@ def test_speed_eager_gguf_installs_no_accelerator(monkeypatch):
def test_speed_max_gguf_regional_compile_not_dequant(monkeypatch):
# GGUF `max` = the FULL regional block compile (which fuses the dequant inline), so
# the standalone compiled dequant is deliberately OFF.
# GGUF `max` is the FULL regional block compile (which fuses the dequant inline), so the
# standalone compiled dequant is deliberately OFF.
_stub_torch(monkeypatch)
called = _stub_gguf_accel(monkeypatch)
pipe = _Pipe(with_compile = True, with_fuse = True)
@ -395,10 +392,9 @@ class _UNetPipe:
def test_unet_whole_compile_default_tier(monkeypatch):
# SDXL's UNet has no _repeated_blocks, so `default` falls back to a whole-module
# STATIC compile (measured 1.61x at LPIPS 0.034 on SDXL): fullgraph on, dynamic OFF.
# The U-Net recipe also fuses QKV (36.3 vs 39.3 ms/step) and compiles the VAE decode
# (4.98 -> 4.25 s over 4 images) on the same tier.
# SDXL's UNet has no _repeated_blocks, so `default` falls back to a whole-module STATIC compile
# (measured 1.61x at LPIPS 0.034): fullgraph on, dynamic OFF. The U-Net recipe also fuses QKV and
# compiles the VAE decode on the same tier.
_stub_torch(monkeypatch)
pipe = _UNetPipe()
applied = apply_speed_optims(
@ -411,9 +407,8 @@ def test_unet_whole_compile_default_tier(monkeypatch):
def test_dit_default_tier_keeps_fuse_and_vae_decode_off(monkeypatch):
# The DiT default tier is unchanged: fused QKV measured exactly neutral there
# (Qwen-Image 6.53 vs 6.52 s) so it stays max-only, and the VAE decode is a few
# percent of a DiT generation so it stays eager.
# The DiT default tier is unchanged: fused QKV measured exactly neutral there so it stays
# max-only, and the VAE decode is a few percent of a DiT generation so it stays eager.
_stub_torch(monkeypatch)
pipe = _Pipe(with_compile = True, with_fuse = True)
applied = apply_speed_optims(
@ -455,8 +450,8 @@ def test_unet_whole_compile_max_tier_mode(monkeypatch):
def test_unet_whole_compile_gated_by_class_name(monkeypatch):
# An unlisted U-Net class (unmeasured architecture) stays eager rather than paying
# an unvalidated whole-module compile.
# An unlisted U-Net class (unmeasured architecture) stays eager rather than paying an unvalidated
# whole-module compile.
_stub_torch(monkeypatch)
pipe = _UNetPipe(unet = _SomeOtherUNet())
applied = apply_speed_optims(
@ -579,9 +574,8 @@ def test_fp16_accum_respects_kill_switch(monkeypatch):
@pytest.mark.parametrize("value", ["TRUE", "Yes", "On", " true "])
def test_fp16_accum_kill_switch_is_case_insensitive(monkeypatch, value):
# The documented safety escape hatch must honor the common boolean spellings, not only
# lowercase "1"/"true"/"yes": an operator setting UNSLOTH_DISABLE_FP16_ACCUM=TRUE to stop
# fp16-accumulation drift would otherwise be silently ignored.
# The documented safety escape hatch must honor the common boolean spellings, not only lowercase
# "1"/"true"/"yes": an operator setting UNSLOTH_DISABLE_FP16_ACCUM=TRUE would otherwise be ignored.
_stub_torch_fp16_accum(monkeypatch, consumer = True)
_stub_gguf_accel(monkeypatch)
monkeypatch.setenv("UNSLOTH_DISABLE_FP16_ACCUM", value)
@ -623,8 +617,8 @@ def test_fp16_accum_not_touched_off_cuda(monkeypatch):
def test_fp16_accum_denied_on_fp16_dtype_below_max(monkeypatch):
# fp16 compute is where the accumulator width actually changes results (measured
# same-seed drift, mean 2-5%): the quality-neutral tiers must refuse it.
# fp16 compute is where the accumulator width actually changes results (measured same-seed drift,
# mean 2-5%), so the quality-neutral tiers must refuse it.
torch = _stub_torch_fp16_accum(monkeypatch, consumer = True)
_stub_gguf_accel(monkeypatch)
for mode in ("eager", "default"):
@ -640,8 +634,8 @@ def test_fp16_accum_denied_on_fp16_dtype_below_max(monkeypatch):
def test_fp16_accum_allowed_on_fp16_dtype_under_max(monkeypatch):
# max already trades exactness for speed (conv algos, max-autotune), so the 2x
# fp16 accumulate joins that tier for fp16 pipelines.
# max already trades exactness for speed (conv algos, max-autotune), so the 2x fp16 accumulate
# joins that tier for fp16 pipelines.
torch = _stub_torch_fp16_accum(monkeypatch, consumer = True)
_stub_gguf_accel(monkeypatch)
applied = apply_speed_optims(
@ -673,10 +667,9 @@ def _stub_inductor_config(
def test_regional_compile_enables_emulate_precision_casts(monkeypatch):
# Inductor's fused pointwise kernels keep intermediates in fp32 where eager rounds
# to bf16 between ops; over a multi-step denoise that compounds to a visible drift.
# emulate_precision_casts restores eager's rounding at zero measured speed cost, so
# the regional compile path must switch it on.
# Inductor's fused pointwise kernels keep intermediates in fp32 where eager rounds to bf16 between
# ops, which compounds over a multi-step denoise. emulate_precision_casts restores eager's
# rounding at zero measured cost, so the regional compile path must switch it on.
torch = _stub_torch(monkeypatch)
_stub_gguf_accel(monkeypatch)
cfg = _stub_inductor_config(monkeypatch, torch, emulate = False)
@ -689,8 +682,8 @@ def test_regional_compile_enables_emulate_precision_casts(monkeypatch):
def test_snapshot_restores_emulate_precision_casts(monkeypatch):
# The flag is process-global, so the unload path must restore the pre-load value
# exactly like the TF32 / cudnn.benchmark globals.
# The flag is process-global, so the unload path must restore the pre-load value exactly like the
# TF32 / cudnn.benchmark globals.
torch = _stub_torch(monkeypatch)
cfg = _stub_inductor_config(monkeypatch, torch, emulate = False)
snap = snapshot_backend_flags()
@ -701,8 +694,8 @@ def test_snapshot_restores_emulate_precision_casts(monkeypatch):
def test_missing_inductor_config_is_tolerated(monkeypatch):
# A build without torch._inductor (or with the flag renamed) must neither break the
# snapshot nor the compile path.
# A build without torch._inductor (or with the flag renamed) must break neither the snapshot nor
# the compile path.
_stub_torch(monkeypatch) # the stub torch has no _inductor attribute
_stub_gguf_accel(monkeypatch)
snap = snapshot_backend_flags()
@ -715,10 +708,9 @@ def test_missing_inductor_config_is_tolerated(monkeypatch):
def test_regional_compile_arms_cache_hook_inners(monkeypatch):
# The production load order engages the step cache BEFORE compile, so the regional
# compile pass must re-arm the already-installed cache hooks with compiled inner
# forwards (otherwise every computed step runs eager under the hook's
# torch.compiler.disable and forfeits the regional compile).
# The production load order engages the step cache BEFORE compile, so the regional compile pass
# must re-arm the already-installed cache hooks with compiled inner forwards, else every computed
# step runs eager under the hook's torch.compiler.disable.
_stub_torch(monkeypatch)
_stub_gguf_accel(monkeypatch)
from core.inference import diffusion_cache as dc_mod

28
studio/backend/tests/test_diffusion_te_prequant.py
View file

@ -60,8 +60,8 @@ def test_family_repo_by_scheme_and_component():
assert (
family_te_prequant_repo(_fam(te_prequant_repos = (("bad",),)), "fp8", "text_encoder") is None
)
# Families without the field resolve to None (both dataclasses default it, but a fake
# or an older family object must not break).
# Families without the field resolve to None (both dataclasses default it, but a fake or an older
# family object must not break).
assert family_te_prequant_repo(types.SimpleNamespace(name = "x"), "fp8", "text_encoder") is None
@ -317,8 +317,8 @@ def test_family_dataclasses_declare_te_prequant_field():
assert DiffusionFamily.__dataclass_fields__["te_prequant_repos"].default_factory is tuple
assert VideoFamily.__dataclass_fields__["te_prequant_repos"].default_factory is tuple
# Families without a hosted TE checkpoint keep the empty default (sdxl's CLIPs
# stay dense; flux.1 now hosts its T5 and is asserted below).
# Families without a hosted TE checkpoint keep the empty default (sdxl's CLIPs stay dense; flux.1
# hosts its T5 and is asserted below).
fam = detect_family("stabilityai/stable-diffusion-xl-base-1.0")
assert fam.te_prequant_repos == ()
@ -350,8 +350,8 @@ def test_hosted_te_prequant_entries():
te_prequant_repo_filename("unsloth/LTX-2-FP8", "text_encoder", "fp8")
== "LTX-2-text_encoder-FP8.pt"
)
# HiDream's heavyweight is TE4 (Llama-3.1-8B), engaged via hidream_te4_kwargs because
# the generic quantize_text_encoders pass only covers text_encoder.._3.
# HiDream's heavyweight is TE4 (Llama-3.1-8B), engaged via hidream_te4_kwargs because the generic
# quantize_text_encoders pass only covers text_encoder.._3.
assert detect_family("HiDream-ai/HiDream-I1-Full").te_prequant_repos == (
("fp8", "text_encoder_4", "unsloth/HiDream-I1-Full-FP8"),
)
@ -359,8 +359,8 @@ def test_hosted_te_prequant_entries():
te_prequant_repo_filename("unsloth/HiDream-I1-Full-FP8", "text_encoder_4", "fp8")
== "HiDream-I1-Full-text_encoder_4-FP8.pt"
)
# Round 2: T5-XXL for every flux.1 base (byte-identical weights, one artifact),
# Gemma2-2B, Qwen3-4B, Qwen3-VL-4B, and hunyuanimage reusing the Qwen-Image artifact.
# Round 2: T5-XXL for every flux.1 base (byte-identical weights, one artifact), Gemma2-2B,
# Qwen3-4B, Qwen3-VL-4B, and hunyuanimage reusing the Qwen-Image artifact.
assert detect_family("black-forest-labs/FLUX.1-schnell").te_prequant_repos == (
("fp8", "text_encoder_2", "unsloth/FLUX.1-schnell-FP8"),
)
@ -380,8 +380,8 @@ def test_hosted_te_prequant_entries():
assert detect_family("hunyuanvideo-community/HunyuanImage-2.1-Diffusers").te_prequant_repos == (
("fp8", "text_encoder", "unsloth/Qwen-Image-FP8"),
)
# flux.2-klein-4B hosts NO TE entry: its Qwen3-4B retrained layer 35's MLP, so the
# z-image artifact must not serve it (verified tensor diff, maxdiff 0.86).
# flux.2-klein-4B hosts NO TE entry: its Qwen3-4B retrained layer 35's MLP, so the z-image artifact
# must not serve it (verified tensor diff, maxdiff 0.86).
assert detect_family("black-forest-labs/FLUX.2-klein-4B").te_prequant_repos == ()
@ -557,11 +557,11 @@ def test_cast_fp8_is_idempotent_on_precast_encoder():
enc = torch.nn.Sequential(torch.nn.Linear(64, 64), torch.nn.LayerNorm(64))
_cast_fp8(enc, target)
assert enc[0].weight.dtype == torch.float8_e4m3fn
# Module.dtype must report the COMPUTE dtype: pipelines derive tensor dtypes from it
# (Flux2 feeds it to randn_tensor, which has no fp8 kernel).
# Module.dtype must report the COMPUTE dtype: pipelines derive tensor dtypes from it (Flux2 feeds
# it to randn_tensor, which has no fp8 kernel).
assert enc.dtype == torch.bfloat16
# EXACT class identity: a dynamic-subclass swap here broke transformers' kwargs-based
# output recording (Qwen3VLModel returned hidden_states=None; krea-2 crashed at encode).
# EXACT class identity: a dynamic-subclass swap here broke transformers' kwargs-based output
# recording (Qwen3VLModel returned hidden_states=None; krea-2 crashed at encode).
assert type(enc) is torch.nn.Sequential
# An uncast sibling of the same (now property-patched) class keeps original behaviour.
sibling = torch.nn.Sequential(torch.nn.Linear(8, 8))

45
studio/backend/tests/test_diffusion_train_perf.py
View file

@ -46,8 +46,8 @@ from core.training.diffusion_training_service import DiffusionTrainingService
from models.training import DiffusionTrainingStartRequest, DiffusionTrainingStopRequest
from routes.training import router as training_router
# A trainable SDXL base so DiffusionLoraConfig.normalized() resolves a family without a
# network call (resolve_trainable_family is pure name matching for this repo).
# A trainable SDXL base so DiffusionLoraConfig.normalized() resolves a family without a network
# call (resolve_trainable_family is pure name matching for this repo).
_SDXL = "stabilityai/stable-diffusion-xl-base-1.0"
@ -57,7 +57,7 @@ def _cfg(**kw) -> DiffusionLoraConfig:
# ── _plan_cache_variants (pure, seed-deterministic) ───────────────────────────
def test_plan_cache_variants_deterministic_and_deduped():
# Same seed -> byte-identical plan (its own rng stream, so it is fully reproducible).
# Same seed gives a byte-identical plan (its own rng stream, so it is fully reproducible).
p1 = _plan_cache_variants(3, 4, center_crop = False, random_flip = True, seed = 123)
p2 = _plan_cache_variants(3, 4, center_crop = False, random_flip = True, seed = 123)
assert p1 == p2
@ -67,8 +67,8 @@ def test_plan_cache_variants_deterministic_and_deduped():
p_one = _plan_cache_variants(3, 1, center_crop = False, random_flip = True, seed = 7)
assert [len(v) for v in p_one] == [1, 1, 1]
# A center crop with no flip collapses to a single distinct variant no matter how many
# draws are requested, and that variant is the fixed (0.5, 0.5, False) center.
# A center crop with no flip collapses to a single distinct variant no matter how many draws are
# requested, and that variant is the fixed (0.5, 0.5, False) center.
p_cc = _plan_cache_variants(2, 8, center_crop = True, random_flip = False, seed = 7)
assert [len(v) for v in p_cc] == [1, 1]
assert p_cc[0][0] == (0.5, 0.5, False)
@ -100,8 +100,8 @@ def test_flux_collate_shapes():
def test_qwen_collate_pads_and_masks():
dim = 8
# A short (mask=None) and a long (mask=ones) entry -> pad to the batch max and build the
# validity mask, with the padded tail of the short sample masked out.
# A short (mask=None) and a long (mask=ones) entry pad to the batch max and build the validity
# mask, with the padded tail of the short sample masked out.
short = (torch.randn(1, 5, dim), None)
long = (torch.randn(1, 9, dim), torch.ones(1, 9, dtype = torch.int64))
pe, mask = _qwen_collate([short, long], "cpu", torch.float32)
@ -123,8 +123,8 @@ def test_qwen_collate_pads_and_masks():
def test_zimage_collate_list():
# Z-Image uses list I/O: the batch is one tuple carrying a list of per-sample tensors, each
# cast to the requested dtype.
# Z-Image uses list I/O: the batch is one tuple carrying a list of per-sample tensors, each cast
# to the requested dtype.
entries = [(torch.randn(7, 2560),), (torch.randn(9, 2560),)]
out = _zimage_collate(entries, "cpu", torch.float32)
assert isinstance(out, tuple) and len(out) == 1
@ -135,8 +135,8 @@ def test_zimage_collate_list():
# ── index-based sigma gather ──────────────────────────────────────────────────
def test_gather_sigmas_matches_search_based_gather():
# CI installs the backend test deps without diffusers; the scheduler math is what we
# are checking, so skip rather than fail there.
# CI installs the backend test deps without diffusers, and the scheduler math is what we check,
# so skip rather than fail there.
pytest.importorskip("diffusers")
from diffusers import FlowMatchEulerDiscreteScheduler
@ -192,8 +192,8 @@ def test_config_validates_new_fields():
assert cfg.enable_tf32 is False
assert cfg.cache_latents is False
# String flags from the generic Studio dict path are coerced: "false" is otherwise a
# non-empty (truthy) string, so an opt-out would silently no-op.
# String flags from the generic Studio dict path are coerced: "false" is otherwise a non-empty
# (truthy) string, so an opt-out would silently no-op.
cfg = _config_from_dict(
{
"base_model": _SDXL,
@ -349,18 +349,17 @@ def test_request_models_new_fields():
# ── perf flags round-trip on cpu ──────────────────────────────────────────────
def test_perf_flags_cpu_roundtrip():
# On a cpu device (or a torch build without cuda), applying the perf flags is a no-op
# snapshot path and restoring it must not raise.
# On a cpu device (or a torch build without cuda), applying the perf flags is a no-op snapshot
# path and restoring it must not raise.
snap = _apply_perf_flags(_cfg(), "cpu")
assert isinstance(snap, dict)
_restore_perf_flags(snap) # no exception
def test_perf_flags_tf32_off_clears_flags():
# enable_tf32=False is the strict-fp32 A/B mode: it must actively clear the TF32 flags
# (cudnn TF32 defaults ON in torch) rather than inherit ambient state, and restore must
# put the ambient values back. The flag attributes are plain Python state, present and
# settable on CPU-only torch builds, so this runs without a GPU.
# enable_tf32=False is the strict-fp32 A/B mode: it must actively clear the TF32 flags (cudnn TF32
# defaults ON in torch) rather than inherit ambient state, and restore must put the ambient values
# back. The flag attributes are plain Python state, so this runs without a GPU.
import torch
before = (
@ -396,8 +395,8 @@ class _FakeEncoded:
class _FakeVae:
# Minimal VAE stand-in: encode() returns a posterior of the requested latent shape so the
# builder measures a real per-variant byte size without a model load or image files.
# Minimal VAE stand-in: encode() returns a posterior of the requested latent shape so the builder
# measures a real per-variant byte size without a model load or image files.
def __init__(self, shape):
self._shape = shape
@ -453,8 +452,8 @@ def test_sdxl_cache_built_under_budget(monkeypatch):
def test_sdxl_cache_gated_over_budget(monkeypatch):
# A budget below one variant forces the gate on the first encode: the sentinel is returned
# so the caller keeps the VAE resident and encodes per step.
# A budget below one variant forces the gate on the first encode: the sentinel is returned so the
# caller keeps the VAE resident and encodes per step.
monkeypatch.delenv("UNSLOTH_DIFFUSION_FORCE_LATENT_CACHE", raising = False)
monkeypatch.setattr(train_common, "_LATENT_CACHE_BUDGET_BYTES", 8)
cache = _build_fake_sdxl_cache(monkeypatch, num_images = 3, latent_shape = (1, 4, 8, 8))

94
studio/backend/tests/test_diffusion_transformer_quant.py
View file

@ -49,8 +49,8 @@ def _stub_torch(
torch.cuda = types.SimpleNamespace(
is_available = lambda: cuda_available,
get_device_capability = lambda *a: cc,
# data-center name by default so the ladder tests get the data-center order;
# consumer tests pass a GeForce name (or monkeypatch _is_consumer_gpu).
# A data-center name by default so the ladder tests get the data-center order; consumer tests pass
# a GeForce name (or monkeypatch _is_consumer_gpu).
get_device_name = lambda *a: device_name,
)
monkeypatch.setitem(sys.modules, "torch", torch)
@ -92,9 +92,8 @@ def _allow(monkeypatch, allowed):
def test_auto_blackwell_prefers_fp8_then_falls_back(monkeypatch):
_stub_torch(monkeypatch, cc = (10, 0))
# Even with every scheme available, auto picks fp8 on Blackwell: measured on a B200
# (torch 2.11 + torchao CUTLASS FP4), fp8 is both faster and more accurate than nvfp4
# for the DiT's shapes -- nvfp4's FP4 GEMM only wins on very large GEMMs, not here.
# Even with every scheme available, auto picks fp8 on Blackwell: measured on a B200, fp8 is both
# faster and more accurate than nvfp4 for the DiT's shapes.
_allow(monkeypatch, {TQ_NVFP4, TQ_MXFP8, TQ_FP8, TQ_INT8})
assert select_transformer_quant_scheme(_target(), "auto") == TQ_FP8
# fp8 unavailable: nvfp4 is the next pick (above mxfp8 / int8).
@ -109,12 +108,12 @@ def test_auto_blackwell_prefers_fp8_then_falls_back(monkeypatch):
def test_auto_consumer_blackwell_prefers_int8(monkeypatch):
# Consumer Blackwell (RTX 50xx): fp8 FP32-accumulate is throughput-halved while int8 is
# full-rate, so auto prefers int8 even though fp8 is available (the data-center default).
# Consumer Blackwell (RTX 50xx): fp8 FP32-accumulate is throughput-halved while int8 is full-rate,
# so auto prefers int8 even though fp8 is available.
_stub_torch(monkeypatch, cc = (10, 0), device_name = "NVIDIA GeForce RTX 5090")
_allow(monkeypatch, {TQ_NVFP4, TQ_MXFP8, TQ_FP8, TQ_INT8})
assert select_transformer_quant_scheme(_target(), "auto") == TQ_INT8
# int8 unavailable -> falls back to the rest of the tier (fp8 next).
# int8 unavailable, so it falls back to the rest of the tier (fp8 next).
_allow(monkeypatch, {TQ_NVFP4, TQ_MXFP8, TQ_FP8})
assert select_transformer_quant_scheme(_target(), "auto") == TQ_FP8
@ -127,15 +126,15 @@ def test_auto_consumer_ada_prefers_int8(monkeypatch):
def test_auto_workstation_unknown_prefers_int8(monkeypatch):
# Unknown / workstation name -> treated as consumer (the safe default) -> int8 first.
# An unknown / workstation name is treated as consumer (the safe default), so int8 first.
_stub_torch(monkeypatch, cc = (8, 9), device_name = "NVIDIA RTX A5000")
_allow(monkeypatch, {TQ_FP8, TQ_INT8})
assert select_transformer_quant_scheme(_target(), "auto") == TQ_INT8
def test_auto_professional_rtx_prefers_fp8(monkeypatch):
# Professional parts (RTX PRO 6000 Blackwell, RTX 6000 Ada) are classified datacenter
# by the rest of the backend, so auto keeps fp8 first (not int8) -- matching llama_cpp.
# Professional parts (RTX PRO 6000 Blackwell, RTX 6000 Ada) are classified datacenter by the rest
# of the backend, so auto keeps fp8 first, matching llama_cpp.
for device_name, cc in (
("NVIDIA RTX PRO 6000 Blackwell Server Edition", (10, 0)),
("NVIDIA RTX 6000 Ada Generation", (8, 9)),
@ -146,7 +145,7 @@ def test_auto_professional_rtx_prefers_fp8(monkeypatch):
def test_auto_ada_hopper_prefers_fp8(monkeypatch):
# Data-center Ada (L40S) / Hopper (H100): not nerfed -> fp8 first.
# Data-center Ada (L40S) / Hopper (H100) are not nerfed, so fp8 comes first.
_stub_torch(monkeypatch, cc = (8, 9), device_name = "NVIDIA L40S")
_allow(monkeypatch, {TQ_NVFP4, TQ_MXFP8, TQ_FP8, TQ_INT8})
assert select_transformer_quant_scheme(_target(), "auto") == TQ_FP8
@ -172,7 +171,7 @@ def test_explicit_scheme_honored_or_none(monkeypatch):
_stub_torch(monkeypatch, cc = (8, 0))
_allow(monkeypatch, {TQ_INT8})
assert select_transformer_quant_scheme(_target(), "int8") == TQ_INT8
# Explicit unsupported scheme is NOT silently downgraded -> None (-> GGUF fallback).
# An explicit unsupported scheme is NOT silently downgraded: None, so the GGUF fallback.
assert select_transformer_quant_scheme(_target(), "fp8") is None
assert select_transformer_quant_scheme(_target(), "nvfp4") is None
@ -188,11 +187,11 @@ def test_select_none_when_disabled_or_non_cuda(monkeypatch):
def test_scheme_supported_shortcircuits(monkeypatch):
# No CUDA -> False without running the smoke probe.
# No CUDA gives False without running the smoke probe.
_stub_torch(monkeypatch, cuda_available = False)
monkeypatch.setattr(tq, "_smoke_probe", lambda *a: pytest.fail("probe should not run"))
assert tq._scheme_supported(TQ_INT8, "cuda") is False
# fp8 requested but the fp8 dtype is missing -> False before the probe.
# fp8 requested but the fp8 dtype is missing gives False before the probe.
_stub_torch(monkeypatch, with_fp8 = False)
monkeypatch.setattr(tq, "_smoke_probe", lambda *a: pytest.fail("probe should not run"))
assert tq._scheme_supported(TQ_FP8, "cuda") is False
@ -230,14 +229,14 @@ def test_smoke_probe_caches_and_tolerates_failure(monkeypatch):
tqz.Int8DynamicActivationInt8WeightConfig = lambda: "int8cfg"
tqz.Float8DynamicActivationFloat8WeightConfig = lambda: "fp8cfg"
monkeypatch.setitem(sys.modules, "torchao.quantization", tqz)
# _Lin is callable? No -> the forward lin(x) would fail. Make instances callable.
# _Lin must be callable or the forward lin(x) would fail, so make instances callable.
_Lin.__call__ = lambda self, x: x
assert tq._smoke_probe(TQ_INT8, "cuda") is True
assert tq._smoke_probe(TQ_INT8, "cuda") is True # cached, no second quantize_
assert calls["n"] == 1
# A scheme whose quantize_ raises -> probe False (and cached).
# A scheme whose quantize_ raises probes False (and is cached).
tq._SMOKE_CACHE.clear()
def _quantize_boom(
@ -297,8 +296,8 @@ def test_is_consumer_gpu_false_for_datacenter(monkeypatch, name):
def test_is_consumer_gpu_defaults_true_on_probe_failure(monkeypatch):
# No torch / no device name available -> assume consumer (safe: fast accum is free
# on data center and a win on consumer).
# No torch / no device name available assumes consumer (safe: fast accum is free on data center
# and a win on consumer).
torch = types.ModuleType("torch")
torch.cuda = types.SimpleNamespace() # no get_device_name
monkeypatch.setitem(sys.modules, "torch", torch)
@ -326,9 +325,8 @@ def test_make_filter_fn(monkeypatch):
def test_require_bf16_schemes_excludes_nvfp4():
# fp8 and mxfp8 assert a bf16 weight (torchao 0.17 / B200: "PerRow quantization only works for
# bfloat16 ..." and "Only supporting bf16 out dtype ..."), so they gate on it; nvfp4 quantises an
# fp32 weight fine, so it is NOT gated (leaving its large fp32 projections quantised, not dense).
# fp8 and mxfp8 assert a bf16 weight (torchao 0.17 / B200), so they gate on it; nvfp4 quantises an
# fp32 weight fine, so it is NOT gated and keeps its large fp32 projections quantised.
from core.inference.diffusion_transformer_quant import (
_REQUIRE_BF16_SCHEMES,
TQ_FP8,
@ -344,9 +342,9 @@ def test_require_bf16_schemes_excludes_nvfp4():
def test_make_filter_fn_require_bf16_skips_non_bf16(monkeypatch):
# fp8 / mxfp8 assert a bf16 weight, so require_bf16 must skip a fp32 Linear (which Wan / Hunyuan
# video DiTs keep) while keeping the bf16 ones -- otherwise a single fp32 layer raises inside
# quantize_ and no-ops the whole pass. int8 and nvfp4 leave it off (they quantise fp32 fine).
# fp8 / mxfp8 assert a bf16 weight, so require_bf16 must skip an fp32 Linear (which Wan / Hunyuan
# video DiTs keep) while keeping the bf16 ones, else one fp32 layer raises inside quantize_ and
# no-ops the whole pass. int8 and nvfp4 leave it off.
torch = types.ModuleType("torch")
torch.bfloat16, torch.float32 = "bf16", "fp32"
@ -367,9 +365,9 @@ def test_make_filter_fn_require_bf16_skips_non_bf16(monkeypatch):
def test_make_filter_fn_int8_excludes_modulation_and_embedders(monkeypatch):
# The int8 path skips the large M=1 AdaLN modulation / conditioning-embedder projections
# (they crash torch._int_mm's M>16), while keeping the attention / FFN compute layers and
# the sequence embedders. fp8 (no exclusion) keeps everything.
# The int8 path skips the large M=1 AdaLN modulation / conditioning-embedder projections (they
# crash torch._int_mm's M floor of 16) while keeping the attention / FFN compute layers and the
# sequence embedders. fp8 (no exclusion) keeps everything.
from core.inference.diffusion_transformer_quant import _INT8_EXCLUDE_NAME_TOKENS
class _Lin:
@ -408,16 +406,16 @@ def test_make_filter_fn_int8_excludes_modulation_and_embedders(monkeypatch):
assert keep(big(), fqn) is True, fqn
# Without the exclusion (fp8 path), the modulation layer is kept.
assert make_filter_fn(512)(big(), "transformer_blocks.0.norm1.linear") is True
# A None / empty fqn must not crash the exclusion check (defensive against the callback
# passing no name); with no name nothing matches the exclusion tokens -> kept.
# A None / empty fqn must not crash the exclusion check; with no name nothing matches, so it is
# kept.
assert keep(big(), None) is True
assert keep(big(), "") is True
def test_exclude_tokens_for_scheme_shared_by_runtime_and_builder():
# The runtime quantiser and the offline prequant builder must apply the SAME int8
# exclusion, or an int8 prequant artifact quantises the M=1 modulation/embedder linears
# and reintroduces the torch._int_mm crash. int8 gets the exclusion; others get none.
# The runtime quantiser and the offline prequant builder must apply the SAME int8 exclusion, or an
# int8 prequant artifact quantises the M=1 modulation/embedder linears and reintroduces the
# torch._int_mm crash.
from core.inference.diffusion_transformer_quant import (
_INT8_EXCLUDE_NAME_TOKENS,
exclude_tokens_for_scheme,
@ -428,10 +426,9 @@ def test_exclude_tokens_for_scheme_shared_by_runtime_and_builder():
def test_exclude_tokens_for_scheme():
# The shared scheme->exclusion decision used by BOTH the runtime quantise path and the offline
# prequant-checkpoint builder, so an int8 checkpoint built ahead of time skips exactly the
# layers the runtime path skips (offline == runtime). int8 excludes the M=1 modulation /
# embedder tokens; every scaled_mm scheme excludes nothing.
# The shared scheme-to-exclusion decision used by BOTH the runtime quantise path and the offline
# prequant-checkpoint builder, so an int8 checkpoint skips exactly the layers the runtime path
# skips. int8 excludes the M=1 modulation / embedder tokens; every scaled_mm scheme excludes none.
from core.inference.diffusion_transformer_quant import (
_INT8_EXCLUDE_NAME_TOKENS,
exclude_tokens_for_scheme,
@ -444,10 +441,9 @@ def test_exclude_tokens_for_scheme():
def test_exclude_tokens_for_scheme_family():
# Qwen-Image never pads its text stream (unlike FLUX's 512-token T5), so a short prompt
# runs the text-stream linears at M <= 16 and torch._int_mm raises ("size(0) needs to be
# greater than 16"); they stay bf16 while the M ~ 4k image stream keeps int8 coverage.
# Unknown families keep the family-independent behaviour.
# Qwen-Image never pads its text stream (unlike FLUX's 512-token T5), so a short prompt runs the
# text-stream linears at M under 16 and torch._int_mm raises; they stay bf16 while the M ~ 4k image
# stream keeps int8 coverage. Unknown families keep the family-independent behaviour.
from core.inference.diffusion_transformer_quant import (
_INT8_EXCLUDE_NAME_TOKENS,
_QWENIMAGE_INT8_EXCLUDES,
@ -563,7 +559,7 @@ def test_quantize_transformer_tolerates_failure(monkeypatch):
tqz.quantize_ = _boom
monkeypatch.setitem(sys.modules, "torchao.quantization", tqz)
pipe = types.SimpleNamespace(transformer = types.SimpleNamespace())
# A quantise failure returns None (caller falls back to GGUF), never raises.
# A quantise failure returns None (the caller falls back to GGUF), never raises.
assert quantize_transformer(pipe, _target(), mode = "int8") is None
@ -571,9 +567,8 @@ def test_quantize_transformer_tolerates_failure(monkeypatch):
def test_family_deny_auto_skips_fp8_for_qwen(monkeypatch):
# B200 with every scheme available: auto must NOT pick fp8 / nvfp4 / mxfp8 for the
# Qwen DiT (per-row fp8 renders black frames on it; see _FAMILY_SCHEME_DENY) and
# falls through the ladder to int8, which measures excellent on Qwen.
# B200 with every scheme available: auto must NOT pick fp8 / nvfp4 / mxfp8 for the Qwen DiT
# (per-row fp8 renders black frames on it) and falls through the ladder to int8.
_stub_torch(monkeypatch, cc = (10, 0))
_allow(monkeypatch, {TQ_FP8, TQ_NVFP4, TQ_MXFP8, TQ_INT8})
assert select_transformer_quant_scheme(_target(), "auto", family = "qwen-image") == TQ_INT8
@ -581,9 +576,8 @@ def test_family_deny_auto_skips_fp8_for_qwen(monkeypatch):
def test_family_deny_refuses_explicit_fp8_for_qwen(monkeypatch):
# An explicit fp8 request on qwen-image returns None (same contract as an
# unsupported scheme: the caller builds the GGUF pipeline instead). int8 stays
# honored on qwen, and fp8 stays honored on families outside the deny table.
# An explicit fp8 request on qwen-image returns None (same contract as an unsupported scheme).
# int8 stays honored on qwen, and fp8 stays honored outside the deny table.
_stub_torch(monkeypatch, cc = (10, 0))
_allow(monkeypatch, {TQ_FP8, TQ_INT8})
assert select_transformer_quant_scheme(_target(), "fp8", family = "qwen-image") is None
@ -600,8 +594,8 @@ def test_family_deny_no_family_keeps_ladder(monkeypatch):
def test_quantize_transformer_threads_family(monkeypatch):
# quantize_transformer passes the family down to the selector, so a denied
# (family, scheme) pair never reaches torchao.
# quantize_transformer passes the family down to the selector, so a denied (family, scheme) pair
# never reaches torchao.
_stub_torch(monkeypatch, cc = (10, 0))
_allow(monkeypatch, {TQ_FP8, TQ_INT8})
pipe = types.SimpleNamespace(transformer = types.SimpleNamespace())

17
studio/backend/tests/test_gguf_load_cache_reuse.py
View file

@ -629,8 +629,8 @@ class TestLoadHubDownloadExclusion:
assert not hf_gguf_load_in_flight("")
def test_chat_load_marker_is_repo_agnostic_and_nests(self):
# The GPU arbiter needs to know a chat load exists before llama-server is spawned, for
# local paths and safetensors too, so this marker carries no repo key.
# The GPU arbiter needs to know a chat load exists before llama-server is spawned, for local paths
# and safetensors too, so this marker carries no repo key.
from core.inference.llama_cpp import chat_load_active, chat_load_in_flight
assert not chat_load_active()
@ -831,13 +831,12 @@ class TestLoadHubDownloadExclusion:
impl = source[source.index("async def _load_model_impl") :]
# One chain, in this order:
# - _resolve_inherited_extra_args first: the inherited value (e.g. a carried
# --no-mmproj) shapes the guard's require_mmproj.
# - the gguf_load_in_flight marker before the hub-download guard: that pair is the
# handshake that keeps a load and the download manager off the same files.
# - both before the CHAT handoff: the guard's 409 loads nothing, so checking it after
# the handoff destroyed a resident Images/Video pipeline for a load that could never
# start. The handoff registers its own marker under the arbiter lock.
# - _resolve_inherited_extra_args first: the inherited value (e.g. a carried --no-mmproj) shapes
# the guard's require_mmproj.
# - the gguf_load_in_flight marker before the hub-download guard: that pair is the handshake that
# keeps a load and the download manager off the same files.
# - both before the CHAT handoff: the guard's 409 loads nothing, so checking it after the handoff
# destroyed a resident Images/Video pipeline for a load that could never start.
# - the resident unload last.
# Anchored on call forms so each assertion pins a call site, not a definition.
assert (

43
studio/backend/tests/test_gpu_arbiter.py
View file

@ -71,9 +71,8 @@ def test_unknown_owner_raises(calls):
def test_evict_chat_unloads_a_still_loading_chat_backend(monkeypatch):
# A chat model still starting up is is_active (process exists) but not yet
# is_loaded (healthy). Eviction must still unload it, or the load would keep
# allocating VRAM after the GPU was handed to diffusion.
# A chat model still starting up is is_active (process exists) but not yet is_loaded (healthy).
# Eviction must still unload it, or the load keeps allocating VRAM after the GPU was handed over.
import core.inference as core_inference
import routes.inference as routes_inference
@ -118,7 +117,7 @@ def test_release_if_drops_only_when_predicate_true(calls):
def test_release_if_by_non_owner_is_noop(calls):
arb.acquire_for(arb.CHAT)
# Predicate is never even consulted for a non-owner; ownership is untouched.
# The predicate is never consulted for a non-owner; ownership is untouched.
consulted: list[bool] = []
assert arb.release_if(arb.DIFFUSION, lambda: consulted.append(True) or True) is False
assert consulted == []
@ -126,8 +125,8 @@ def test_release_if_by_non_owner_is_noop(calls):
def test_release_if_predicate_sees_a_reregistered_same_owner_load(calls):
# The race release_if closes: a slow unload's predicate reports a load now in flight (a
# re-registered same-owner load), so ownership must stay with DIFFUSION.
# The race release_if closes: a slow unload's predicate reports a load now in flight, so ownership
# must stay with DIFFUSION.
arb.acquire_for(arb.DIFFUSION)
loading = {"in_flight": True}
assert arb.release_if(arb.DIFFUSION, lambda: not loading["in_flight"]) is False
@ -135,8 +134,8 @@ def test_release_if_predicate_sees_a_reregistered_same_owner_load(calls):
def test_register_runs_under_ownership_and_returns_result(calls):
# A register callback runs after ownership transfers (owner already set) and its
# return value is forwarded -- the route uses this to register the in-flight load.
# A register callback runs after ownership transfers and its return value is forwarded; the route
# uses this to register the in-flight load.
seen_owner: list = []
def register():
@ -150,8 +149,8 @@ def test_register_runs_under_ownership_and_returns_result(calls):
def test_register_failure_leaves_ownership_in_place(calls):
# A failing register (e.g. begin_load reporting a load already in progress) propagates
# but must not drop ownership -- the prior handoff (chat already evicted) stands.
# A failing register (e.g. begin_load reporting a load already in progress) propagates but must not
# drop ownership: the prior handoff stands.
arb.acquire_for(arb.CHAT)
def register():
@ -164,8 +163,8 @@ def test_register_failure_leaves_ownership_in_place(calls):
def test_competing_acquire_blocks_until_register_completes(monkeypatch):
# While DIFFUSION registers its load, a competing VIDEO acquire must block (not evict) until
# the load is in-flight; holding the lock across register makes eviction never race it.
# While DIFFUSION registers its load, a competing VIDEO acquire must block (not evict) until the
# load is in-flight; holding the lock across register makes eviction never race it.
import threading
import time
@ -207,9 +206,9 @@ def test_competing_acquire_blocks_until_register_completes(monkeypatch):
def test_evict_chat_cancels_a_chat_load_that_has_not_spawned_yet(monkeypatch):
# An HF chat load has no llama-server process until its GGUF finished downloading, which is
# minutes. Gating only on is_active let the evictor find nothing to cancel, grant the GPU to
# the image/video load, and the chat load then spawned onto the same device: two big models
# allocating at once. The in-flight marker is what makes that load cancellable.
# minutes. Gating only on is_active let the evictor find nothing to cancel, grant the GPU to the
# image/video load, and the chat load then spawned onto the same device. The in-flight marker is
# what makes that load cancellable.
import core.inference as core_inference
import routes.inference as routes_inference
from core.inference.llama_cpp import chat_load_in_flight
@ -253,9 +252,9 @@ def test_evict_chat_cancels_a_chat_load_that_has_not_spawned_yet(monkeypatch):
def test_evict_chat_cancels_an_in_flight_safetensors_load(monkeypatch):
# The orchestrator publishes active_model_name only once its worker reports success, so an
# in-flight safetensors load is visible ONLY as an entry in loading_models. Gating the
# cancellation on active_model_name let that worker finish after ownership transferred and
# allocate the model alongside the image/video pipeline.
# in-flight safetensors load is visible ONLY in loading_models. Gating the cancellation on
# active_model_name let that worker finish after ownership transferred and allocate alongside the
# image/video pipeline.
import core.inference as core_inference
import routes.inference as routes_inference
@ -294,8 +293,8 @@ def test_evict_chat_cancels_an_in_flight_safetensors_load(monkeypatch):
def test_evict_chat_cancels_every_pending_load_over_a_live_snapshot(monkeypatch):
# cancel_load discards the marker it cancels, so iterate a snapshot rather than the live
# set (mutating during iteration raises) and cancel each pending entry.
# cancel_load discards the marker it cancels, so iterate a snapshot rather than the live set
# (mutating during iteration raises).
import core.inference as core_inference
import routes.inference as routes_inference
@ -337,8 +336,8 @@ def test_evict_chat_cancels_every_pending_load_over_a_live_snapshot(monkeypatch)
def test_the_safetensors_load_yields_a_gpu_it_lost_while_loading():
# Mirror of the GGUF branch's guard: an Images/Video acquire can land in the gap between the
# eviction and the load's publish, so the load has to undo itself instead of leaving two
# models resident. Without it only the GGUF branch was safe.
# eviction and the load's publish, so the load has to undo itself instead of leaving two models
# resident.
from pathlib import Path
route_src = (Path(__file__).resolve().parent.parent / "routes" / "inference.py").read_text(

21
studio/backend/tests/test_gpu_memory_mode.py
View file

@ -1166,9 +1166,9 @@ def not_vulkan(monkeypatch):
def test_zero_vram_chat_load_only_for_a_deliberate_cpu_only_offload(not_vulkan):
# Manual + gpu_layers=0 is the one shape that launches with the GPUs hidden from the child,
# so it is the one shape allowed to skip the GPU arbiter. Auto (or any pinned layer count)
# puts the model on the GPU and must still evict an image/video pipeline.
# Manual + gpu_layers=0 is the one shape that launches with the GPUs hidden from the child, so it
# is the one shape allowed to skip the GPU arbiter. Auto (or any pinned layer count) puts the
# model on the GPU and must still evict an image/video pipeline.
zero = llama_cpp_module.zero_vram_chat_load
assert zero("manual", 0) is True
assert zero("auto", 0) is False
@ -1178,8 +1178,8 @@ def test_zero_vram_chat_load_only_for_a_deliberate_cpu_only_offload(not_vulkan):
def test_zero_vram_chat_load_refuses_every_gpu_companion(not_vulkan):
# The launch-time mask keeps the GPUs visible for a device pin, tensor mode, an mmproj or a
# drafter, so those loads DO hold VRAM. --mmproj and --model-draft are added by the backend
# rather than carried in the extras, so their intent arrives as flags.
# drafter, so those loads DO hold VRAM. --mmproj and --model-draft are added by the backend, so
# their intent arrives as flags.
zero = llama_cpp_module.zero_vram_chat_load
assert zero("manual", 0, ["--device", "CUDA0"]) is False
assert zero("manual", 0, ["-dev", "CUDA0"]) is False
@ -1200,7 +1200,7 @@ def test_zero_vram_chat_load_is_skipped_on_vulkan(monkeypatch):
def test_holds_no_vram_needs_the_launch_to_have_confirmed_it():
# The property reports the LAUNCHED server, so it follows _gpu_offload_active: True (something
# still reached the GPU) and None (no GPU detected at all) both keep the normal arbiter path.
# reached the GPU) and None (no GPU detected) both keep the normal arbiter path.
backend = LlamaCppBackend()
backend._gpu_memory_mode = "manual"
backend._gpu_layers = 0
@ -1220,17 +1220,16 @@ def test_holds_no_vram_needs_the_launch_to_have_confirmed_it():
def test_a_cpu_only_chat_load_does_not_take_the_gpu_arbiter():
# A load that needs no VRAM must not evict a resident Images/Video pipeline (nor leave CHAT
# recorded as owner, which would make the next GPU workload unload it for nothing). The
# in-flight marker still goes up, and the post-load ownership recheck -- which would 409 a
# load that never acquired -- is gated on the same flag.
# recorded as owner, which would make the next GPU workload unload it for nothing). The in-flight
# marker still goes up, and the post-load ownership recheck is gated on the same flag.
route_src = (Path(_BACKEND_DIR) / "routes" / "inference.py").read_text(encoding = "utf-8")
load_impl = route_src[route_src.index("async def _load_model_impl") :]
assert "chat_load_needs_gpu = not (" in load_impl
gate = load_impl.index("chat_load_needs_gpu = not (")
acquire = load_impl.index("if chat_load_needs_gpu:", gate)
# The stale CHAT claim is dropped only AFTER the load: this load may still be replacing a
# GPU-backed chat model, and releasing before it would let an image/video load allocate
# alongside the model not yet unloaded.
# GPU-backed chat model, and releasing before it would let an image/video load allocate alongside
# the model not yet unloaded.
release = load_impl.index("await asyncio.to_thread(release, CHAT)", acquire)
assert load_impl.index("if not chat_load_needs_gpu:", acquire) < release
assert load_impl.index("success = await load_with_tensor_fallback(", acquire) < release

23
studio/backend/tests/test_gpu_selection.py
View file

@ -1111,9 +1111,9 @@ class TestRouteErrors(unittest.TestCase):
self.assertNotIn("not supported", exc_info.exception.detail.lower())
def test_inference_route_defers_gpu_handoff_until_after_validation(self):
# A doomed chat load (GGUF + gpu_ids -> 400) must NOT reclaim the CHAT arbiter owner
# first: the handoff is deferred past validation, so a resident Images/Video pipeline is
# never evicted for a load that then errors.
# A doomed chat load (GGUF + gpu_ids -> 400) must NOT reclaim the CHAT arbiter owner first: the
# handoff is deferred past validation, so a resident Images/Video pipeline is never evicted for a
# load that then errors.
import core.inference.gpu_arbiter as arb
inference_route = _load_route_module(
@ -1137,9 +1137,8 @@ class TestRouteErrors(unittest.TestCase):
has_audio_input = False,
)
acquired = []
# Make [0, 1] invalid on any host (a duplicate id is rejected everywhere, GPUs or not):
# the point of the test is the ORDER -- validation before the handoff -- not this
# machine's device count.
# Make [0, 1] invalid on any host (a duplicate id is rejected everywhere): the point is the ORDER,
# validation before the handoff, not this machine's device count.
request.gpu_ids = [0, 0]
with (
patch.object(
@ -1169,9 +1168,9 @@ class TestRouteErrors(unittest.TestCase):
self.assertEqual(acquired, []) # no CHAT handoff before the doomed load errored
def test_inference_route_checks_hub_download_conflict_before_the_handoff(self):
# A GGUF the download manager is fetching 409s and loads nothing, so that check has to
# run BEFORE the CHAT handoff: afterwards, it destroyed the resident Images/Video
# pipeline for a load that could never start.
# A GGUF the download manager is fetching 409s and loads nothing, so that check has to run BEFORE
# the CHAT handoff: afterwards, it destroyed the resident Images/Video pipeline for a load that
# could never start.
import core.inference.gpu_arbiter as arb
import core.inference.llama_cpp as llama_cpp
@ -1225,9 +1224,9 @@ class TestRouteErrors(unittest.TestCase):
self.assertEqual(acquired, []) # nothing evicted for a load that cannot start
def test_inference_route_marks_the_chat_load_under_the_arbiter_lock(self):
# A chat load holds no llama-server process until its GGUF downloaded, so the arbiter is
# told about it through acquire_for's `register` hook (which runs under the arbiter lock).
# Passing no register left a competing Images/Video acquire with nothing to cancel.
# A chat load holds no llama-server process until its GGUF downloaded, so the arbiter is told about
# it through acquire_for's `register` hook (which runs under the arbiter lock). Passing no register
# left a competing Images/Video acquire with nothing to cancel.
import core.inference.gpu_arbiter as arb
import core.inference.llama_cpp as llama_cpp

14
studio/backend/tests/test_hf_cache_settings.py
View file

@ -291,14 +291,14 @@ def test_inactive_cache_model_loads_from_snapshot_path(tmp_path):
def test_diffusion_cache_root_follows_a_live_switch(settings_store, tmp_path):
# The image/video backends used huggingface_hub's import-time HF_HUB_CACHE constant,
# which set_hf_cache_home does not update. The download then wrote to the new root
# while progress counted the old one, and a load could split across both.
# The image/video backends used huggingface_hub's import-time HF_HUB_CACHE constant, which
# set_hf_cache_home does not update. The download then wrote to the new root while progress counted
# the old one, and a load could split across both.
import core.inference.diffusion as diffusion
moved = tmp_path / "external-c" / "huggingface"
# Write the setting straight into the store: set_hf_cache_home's folder validation is
# not what is under test, and it rejects the pytest tmp root on macOS.
# Write the setting straight into the store: set_hf_cache_home's folder validation is not under
# test, and it rejects the pytest tmp root on macOS.
settings_store[hf_cache_settings.CACHE_HOME_SETTING_KEY] = str(moved)
assert diffusion.hub_cache_dir() == str(moved / "hub")
@ -308,8 +308,8 @@ def test_diffusion_cache_root_follows_a_live_switch(settings_store, tmp_path):
def test_diffusion_loader_calls_pin_the_cache_dir():
# Every from_pretrained / from_single_file must carry cache_dir, else diffusers
# resolves it through the stale constant and half a load lands in the old root.
# Every from_pretrained / from_single_file must carry cache_dir, else diffusers resolves it through
# the stale constant and half a load lands in the old root.
for rel in ("core/inference/diffusion.py", "core/inference/video.py"):
source = (Path(_BACKEND_DIR) / rel).read_text(encoding = "utf-8")
for call in ("from_pretrained(", "from_single_file("):

15
studio/backend/tests/test_image_gallery.py
View file

@ -129,9 +129,8 @@ def test_image_path_rejects_unsafe_ids():
def test_owned_image_path_serves_only_owned_pngs():
# A hand-dropped foreign PNG resolves via image_path (safe stem, on disk) but must NOT be
# served: owned_image_path applies the same recipe check as delete/clear, so the serve route
# can't stream a file the listing hides.
# A hand-dropped foreign PNG resolves via image_path (safe stem, on disk) but must NOT be served:
# owned_image_path applies the same recipe check as delete/clear.
foreign = gallery.gallery_dir() / "family-photo.png"
_img().save(foreign, format = "PNG")
assert gallery.image_path("family-photo") is not None # resolvable...
@ -154,21 +153,21 @@ def test_list_skips_foreign_pngs(tmp_path):
def test_foreign_png_in_window_does_not_drop_valid_images():
# A foreign PNG sorting INTO the requested page must not consume a window slot and
# drop a valid image that sorts after it: paging is over readable records, not files.
# A foreign PNG sorting INTO the requested page must not consume a window slot and drop a valid
# image that sorts after it: paging is over readable records, not files.
_save_with_mtime("p2", 100.0)
foreign = gallery.gallery_dir() / "zzz_foreign.png"
_img().save(foreign, format = "PNG") # newest by mtime (set below), sorts first
os.utime(foreign, (300.0, 300.0))
_save_with_mtime("p1", 200.0)
# First page of 2 must still return both real images, not [p1] (foreign eating a slot).
# First page of 2 must still return both real images, not [p1].
page1 = gallery.list_images(limit = 2, offset = 0)
assert [r["prompt"] for r in page1] == ["p1", "p2"]
def test_list_skips_recipe_missing_required_fields(tmp_path):
# A PNG carrying our chunk but an incomplete/older-schema recipe (no seed etc.)
# must be skipped, not crash the whole listing when the route builds GalleryImage.
# A PNG carrying our chunk but an incomplete/older-schema recipe must be skipped, not crash the
# whole listing when the route builds GalleryImage.
import json
from PIL.PngImagePlugin import PngInfo

4
studio/backend/tests/test_inference_model_validation.py
View file

@ -234,8 +234,8 @@ def _diff_load(**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.
# 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"

Some files were not shown because too many files have changed in this diff Show more