From 36df317293f4e63ab64b5301394560e47ef06165 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 26 Jul 2026 20:30:46 +0000 Subject: [PATCH] Trim the comments across the diffusion backend 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). --- scripts/build_prequant_checkpoint.py | 34 +- scripts/build_te_prequant_checkpoint.py | 11 +- scripts/diffusion_bench.py | 12 +- scripts/diffusion_quality.py | 14 +- scripts/fbcache_flux_probe.py | 5 +- scripts/fp8_overflow_check.py | 5 +- scripts/image_speedmem_bench.py | 6 +- scripts/int8_linear_probe.py | 4 +- scripts/nvfp4_t211_probe.py | 4 +- scripts/perf_levers_probe.py | 15 +- scripts/verify_prequant_backend.py | 15 +- scripts/video_quality.py | 17 +- studio/backend/core/inference/diffusion.py | 757 +++++++----------- .../core/inference/diffusion_arch_patches.py | 7 +- .../core/inference/diffusion_attention.py | 52 +- .../core/inference/diffusion_auto_policy.py | 41 +- .../core/inference/diffusion_batched.py | 8 +- .../backend/core/inference/diffusion_cache.py | 34 +- .../core/inference/diffusion_compile_cache.py | 16 +- .../core/inference/diffusion_cond_cache.py | 25 +- .../core/inference/diffusion_controlnet.py | 24 +- .../core/inference/diffusion_device.py | 18 +- .../core/inference/diffusion_eager_patches.py | 10 +- .../core/inference/diffusion_engine_router.py | 32 +- .../core/inference/diffusion_families.py | 240 +++--- .../core/inference/diffusion_hidream.py | 8 +- .../core/inference/diffusion_ideogram4.py | 26 +- .../backend/core/inference/diffusion_krea2.py | 9 +- .../backend/core/inference/diffusion_lora.py | 30 +- .../core/inference/diffusion_memory.py | 59 +- .../core/inference/diffusion_precision.py | 70 +- .../core/inference/diffusion_prequant.py | 61 +- .../backend/core/inference/diffusion_speed.py | 82 +- .../core/inference/diffusion_te_prequant.py | 65 +- .../inference/diffusion_transformer_quant.py | 127 ++- studio/backend/core/inference/gpu_arbiter.py | 26 +- .../backend/core/inference/image_gallery.py | 13 +- studio/backend/core/inference/llama_cpp.py | 8 +- studio/backend/core/inference/sd_cpp_args.py | 18 +- .../backend/core/inference/sd_cpp_backend.py | 121 ++- .../backend/core/inference/sd_cpp_engine.py | 21 +- .../backend/core/inference/sd_cpp_server.py | 43 +- studio/backend/core/inference/video.py | 346 ++++---- .../backend/core/inference/video_families.py | 90 +-- .../backend/core/inference/video_gallery.py | 21 +- studio/backend/core/inference/video_ltx2.py | 17 +- .../core/training/diffusion_lora_trainer.py | 51 +- .../core/training/diffusion_train_common.py | 269 +++---- studio/backend/core/training/training.py | 23 +- .../hub/services/download_lifecycle.py | 11 +- .../backend/hub/services/models/deletion.py | 13 +- .../backend/hub/services/models/downloads.py | 5 +- studio/backend/hub/utils/download_registry.py | 7 +- studio/backend/hub/workers/hf_download.py | 18 +- studio/backend/main.py | 29 +- studio/backend/models/inference.py | 105 ++- studio/backend/routes/inference.py | 326 ++++---- studio/backend/routes/models.py | 129 ++- studio/backend/routes/settings.py | 3 +- studio/backend/routes/training.py | 351 ++++---- studio/backend/routes/video.py | 49 +- studio/backend/tests/conftest.py | 5 +- .../backend/tests/test_cached_gguf_routes.py | 83 +- .../tests/test_diffusion_arch_patches.py | 8 +- .../backend/tests/test_diffusion_attention.py | 88 +- .../tests/test_diffusion_auto_policy.py | 29 +- .../backend/tests/test_diffusion_backend.py | 612 +++++++------- .../tests/test_diffusion_base_precision.py | 190 ++--- studio/backend/tests/test_diffusion_cache.py | 76 +- .../tests/test_diffusion_compile_cache.py | 15 +- .../tests/test_diffusion_cond_cache.py | 22 +- .../tests/test_diffusion_controlnet.py | 75 +- .../tests/test_diffusion_dataset_api.py | 93 +-- studio/backend/tests/test_diffusion_device.py | 4 +- .../tests/test_diffusion_dit_trainer.py | 82 +- .../test_diffusion_dit_trainer_flow_shift.py | 28 +- .../tests/test_diffusion_eager_patches.py | 4 +- .../tests/test_diffusion_engine_router.py | 46 +- .../tests/test_diffusion_gguf_compile.py | 4 +- .../tests/test_diffusion_inference_info.py | 8 +- studio/backend/tests/test_diffusion_krea2.py | 43 +- studio/backend/tests/test_diffusion_lora.py | 72 +- .../tests/test_diffusion_lora_trainer.py | 91 +-- studio/backend/tests/test_diffusion_memory.py | 74 +- .../tests/test_diffusion_more_families.py | 96 ++- .../backend/tests/test_diffusion_precision.py | 51 +- .../backend/tests/test_diffusion_prequant.py | 92 +-- studio/backend/tests/test_diffusion_routes.py | 177 ++-- studio/backend/tests/test_diffusion_sdxl.py | 51 +- studio/backend/tests/test_diffusion_speed.py | 96 +-- .../tests/test_diffusion_te_prequant.py | 28 +- .../tests/test_diffusion_train_perf.py | 45 +- .../tests/test_diffusion_transformer_quant.py | 94 +-- .../tests/test_gguf_load_cache_reuse.py | 17 +- studio/backend/tests/test_gpu_arbiter.py | 43 +- studio/backend/tests/test_gpu_memory_mode.py | 21 +- studio/backend/tests/test_gpu_selection.py | 23 +- .../backend/tests/test_hf_cache_settings.py | 14 +- studio/backend/tests/test_image_gallery.py | 15 +- .../tests/test_inference_model_validation.py | 4 +- .../backend/tests/test_local_model_format.py | 49 +- studio/backend/tests/test_middleware.py | 40 +- .../test_openai_images_generations_route.py | 45 +- .../tests/test_personalization_settings.py | 3 +- .../backend/tests/test_scoped_download_job.py | 33 +- studio/backend/tests/test_sd_cpp_args.py | 19 +- studio/backend/tests/test_sd_cpp_backend.py | 84 +- studio/backend/tests/test_sd_cpp_engine.py | 23 +- studio/backend/tests/test_sd_cpp_install.py | 59 +- studio/backend/tests/test_sd_cpp_server.py | 20 +- .../tests/test_training_start_offload.py | 14 +- studio/backend/tests/test_video_backend.py | 341 ++++---- studio/backend/tests/test_video_families.py | 71 +- 113 files changed, 3389 insertions(+), 4087 deletions(-) diff --git a/scripts/build_prequant_checkpoint.py b/scripts/build_prequant_checkpoint.py index f511867c5e..b5c78b52ba 100644 --- a/scripts/build_prequant_checkpoint.py +++ b/scripts/build_prequant_checkpoint.py @@ -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 = { diff --git a/scripts/build_te_prequant_checkpoint.py b/scripts/build_te_prequant_checkpoint.py index 94a7b713b1..3c15ee5c80 100644 --- a/scripts/build_te_prequant_checkpoint.py +++ b/scripts/build_te_prequant_checkpoint.py @@ -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__), } diff --git a/scripts/diffusion_bench.py b/scripts/diffusion_bench.py index f27ad17aeb..96cb6d4a94 100644 --- a/scripts/diffusion_bench.py +++ b/scripts/diffusion_bench.py @@ -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" diff --git a/scripts/diffusion_quality.py b/scripts/diffusion_quality.py index ca121304d3..305aab3d59 100644 --- a/scripts/diffusion_quality.py +++ b/scripts/diffusion_quality.py @@ -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): diff --git a/scripts/fbcache_flux_probe.py b/scripts/fbcache_flux_probe.py index d323e837aa..588272935d 100644 --- a/scripts/fbcache_flux_probe.py +++ b/scripts/fbcache_flux_probe.py @@ -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) diff --git a/scripts/fp8_overflow_check.py b/scripts/fp8_overflow_check.py index b36cc6e93b..7d96a78fa3 100644 --- a/scripts/fp8_overflow_check.py +++ b/scripts/fp8_overflow_check.py @@ -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) diff --git a/scripts/image_speedmem_bench.py b/scripts/image_speedmem_bench.py index d98707831f..273925a2ed 100644 --- a/scripts/image_speedmem_bench.py +++ b/scripts/image_speedmem_bench.py @@ -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() diff --git a/scripts/int8_linear_probe.py b/scripts/int8_linear_probe.py index 5cef1b68e7..4a30af0d0f 100644 --- a/scripts/int8_linear_probe.py +++ b/scripts/int8_linear_probe.py @@ -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(".")) diff --git a/scripts/nvfp4_t211_probe.py b/scripts/nvfp4_t211_probe.py index de697165ea..5a74ccad40 100644 --- a/scripts/nvfp4_t211_probe.py +++ b/scripts/nvfp4_t211_probe.py @@ -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, diff --git a/scripts/perf_levers_probe.py b/scripts/perf_levers_probe.py index 16f8d826c5..8c00216d8e 100644 --- a/scripts/perf_levers_probe.py +++ b/scripts/perf_levers_probe.py @@ -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 diff --git a/scripts/verify_prequant_backend.py b/scripts/verify_prequant_backend.py index d0c2518c89..dbb7946e08 100644 --- a/scripts/verify_prequant_backend.py +++ b/scripts/verify_prequant_backend.py @@ -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) diff --git a/scripts/video_quality.py b/scripts/video_quality.py index 40d3732f2c..a620aae81f 100644 --- a/scripts/video_quality.py +++ b/scripts/video_quality.py @@ -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 diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 44b2f842aa..215a05a4c6 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -181,11 +181,8 @@ def resolve_local_single_file(model_path: str) -> Optional[str]: root = Path(model_path).expanduser() if not root.is_dir() or (root / "model_index.json").is_file(): return None - # A PEFT LoRA adapter folder (adapter_config.json + adapter_model.safetensors) is not a - # base checkpoint: from_single_file would fail on the adapter weights AFTER the route - # evicted the resident GPU model. Skip it so the pick stays a pipeline load and 400s in - # validation, before the GPU handoff. Also drop a bare adapter_model.safetensors so a - # config-less adapter export is never reinterpreted as the sole checkpoint. + # A PEFT adapter folder is not a base checkpoint: from_single_file would fail on it AFTER the + # route evicted the GPU model. Skip it (and a bare adapter_model.safetensors) so validation 400s. if (root / "adapter_config.json").is_file(): return None checkpoints = [ @@ -220,13 +217,11 @@ def _decode_b64_image(data: str, *, mode: str = "RGB") -> Any: blob = base64.b64decode(raw, validate = False) except (binascii.Error, ValueError) as exc: raise ValueError(f"Invalid base64 image data: {exc}") from exc - # Bound the decoded size (this is the single decode path for every image-conditioned - # workflow). 4096px covers txt2img's 2048 max, upscales, and outpaint canvases; larger 400s. + # Bound the decoded size: 4096px covers txt2img's 2048 max, upscales and outpaint canvases. max_side = 4096 try: img = Image.open(io.BytesIO(blob)) - # Reject an over-limit image from the header BEFORE img.load() decompresses pixels, so a - # crafted small-payload/huge-dimension file can't spike memory first. + # Reject from the header before img.load(), so a huge-dimension file can't spike memory first. w, h = img.size if w > max_side or h > max_side: raise ValueError(f"Image is too large ({w}x{h}); maximum is {max_side}px per side.") @@ -292,52 +287,41 @@ def _compile_shape_dims(workflow: str, init_pil: Any, width: int, height: int) - return int(iw), int(ih) -# Official base repos that may load as a full (non-GGUF) pipeline despite not being under -# unsloth/. Safetensors-only, no pickle/remote code; exact-match lowercased (typo-squat safe). -# Extend deliberately; never add pickled weights or remote code. The SDXL refiner is -# intentionally NOT here (img2img-only; this backend loads every sdxl repo as base txt2img). +# Official non-unsloth base repos that may load as a full pipeline. Safetensors-only, no pickle +# or remote code; exact lowercased match. The SDXL refiner is intentionally absent (img2img-only). _TRUSTED_NON_GGUF_REPOS = frozenset( { "stabilityai/stable-diffusion-xl-base-1.0", "stabilityai/sdxl-turbo", - # Vendor safetensors-only bases: LoRA TRAINING bases + the BF16 artifact behind each - # catalog group. FLUX.1 repos are Hub-gated (need the user's token); Qwen/Z-Image are open. + # Vendor safetensors-only bases: LoRA training bases + the BF16 artifact per catalog group. + # FLUX.1 repos are Hub-gated (need the user's token); Qwen/Z-Image are open. "black-forest-labs/flux.1-dev", "black-forest-labs/flux.1-schnell", "black-forest-labs/flux.1-kontext-dev", - # Krea's guidance-distilled FLUX.1-dev finetune: same arch/layout as dev (FluxPipeline, - # CLIP+T5+ae), gated like dev. Detected as the flux.1 family via the "flux.1" token. + # Krea: guidance-distilled FLUX.1-dev finetune, same arch and gating; detected via "flux.1". "black-forest-labs/flux.1-krea-dev", - # FLUX.2: the LoRA TRAINING bases for the flux.2-dev / flux.2-klein families (dev is - # Hub-gated, klein-4B is open). Both ship a safetensors-only diffusers layout - # (model_index.json + transformer/text_encoder/vae, no pickled weights, no remote code), - # so a pipeline load is the same shape as FLUX.1-dev's. Trusted here too because the Train - # tab's "Deploy to Create" reloads the trained-on base as a pipeline: without these the - # deploy of every FLUX.2 adapter 400s on this gate. + # FLUX.2 LoRA training bases (dev gated, klein-4B open), safetensors-only diffusers layout. + # Trusted because the Train tab's "Deploy to Create" reloads the trained-on base as a pipeline. "black-forest-labs/flux.2-dev", "black-forest-labs/flux.2-klein-4b", "tongyi-mai/z-image-turbo", "qwen/qwen-image", "qwen/qwen-image-2512", "qwen/qwen-image-edit-2511", - # Krea 2: assembled per-component (diffusion_krea2.py). Turbo = inference; Raw = the - # undistilled base to train LoRAs on (train on Raw, run adapters on Turbo). + # Krea 2: assembled per-component. Turbo = inference; Raw = the base to train LoRAs on. "krea/krea-2-turbo", "krea/krea-2-raw", - # Lumina Image 2.0: standard diffusers layout (Gemma2-2B encoder), safetensors-only, - # loads through the generic from_pretrained pipeline path. + # Lumina Image 2.0: standard diffusers layout, loads via the generic from_pretrained path. "alpha-vllm/lumina-image-2.0", - # HunyuanImage 2.1: the community diffusers mirror (open, tencent-hunyuan-community - # license), safetensors-only, including the diffusers-native guider components. + # HunyuanImage 2.1: open community diffusers mirror, safetensors-only. "hunyuanvideo-community/hunyuanimage-2.1-diffusers", - # HiDream-I1: open MIT-weights repos, all three variants one family. The Llama TE the - # model_index names comes from the unsloth mirror (diffusion_hidream.py), which the - # unsloth/ org prefix already trusts. + # HiDream-I1: open MIT weights, all three variants one family. The Llama TE comes from the + # unsloth mirror (diffusion_hidream.py), already trusted by the unsloth/ prefix. "hidream-ai/hidream-i1-full", "hidream-ai/hidream-i1-dev", "hidream-ai/hidream-i1-fast", - # Ideogram 4: no bf16 ships. -fp8 stores the two DiTs as raw float8 (the family base); - # the two nf4 repos are identical bnb-4bit exports (both listed so either id loads). + # Ideogram 4: no bf16 ships. -fp8 stores both DiTs as raw float8 (the family base); the two + # nf4 repos are identical bnb-4bit exports (both listed so either id loads). "ideogram-ai/ideogram-4-fp8", "ideogram-ai/ideogram-4-nf4", "ideogram-ai/ideogram-4-nf4-diffusers", @@ -410,14 +394,12 @@ class _LoadState: offload_policy: str = OFFLOAD_NONE vae_tiling: bool = False memory_mode: str = "auto" - # Resolved load kind ("gguf"|"single_file"|"pipeline"); surfaced so the UI can gate - # GGUF-only controls (the dense transformer_quant fast path engages only on gguf). + # Resolved load kind ("gguf"|"single_file"|"pipeline"); lets the UI gate GGUF-only controls. kind: str = "gguf" # The opt-in speed profile. speed_mode: str = SPEED_OFF speed_optims: tuple = () - # Process-wide torch backend flags (TF32 / cudnn.benchmark) captured before the speed - # layer mutated them; restored on unload so a later `off` load isn't contaminated. + # Torch backend flags (TF32 / cudnn.benchmark) captured before the speed layer mutated them. backend_flags_before: Optional[dict] = None # Text-encoder quant engaged: "fp8" | "nvfp4" | None. text_encoder_quant: Optional[str] = None @@ -425,21 +407,18 @@ class _LoadState: transformer_quant: Optional[str] = None # Attention backend engaged via the diffusers dispatcher, or None for default SDPA. attention_backend: Optional[str] = None - # The caller's ORIGINAL attention request, carried so the deferred-speed engagement - # re-runs the SAME selection instead of forcing the auto cuDNN upgrade over an explicit pin. + # The caller's original attention request, so deferred engagement re-runs the same selection. attention_request: Optional[str] = None # Step cache engaged ("fbcache") or None. Opt-in, for many-step models. transformer_cache: Optional[str] = None - # AUTO on a cache-capable transformer: generate() re-checks the step count and toggles - # FBCache across FBCACHE_MIN_STEPS. An explicit request (off / fbcache) is never toggled. + # AUTO: generate() toggles FBCache across FBCACHE_MIN_STEPS; an explicit request never toggles. cache_auto: bool = False # Inputs the generation-time toggle re-applies (quantised threshold + override). cache_quant_active: bool = False cache_threshold: Optional[float] = None # Shared eager patches installed for this load (any non-off tier); uninstalled on unload. eager_patched: bool = False - # Deferred speed auto: the load stays eager/bit-identical; generate() engages the `default` - # compile profile at the 3rd generation this session. Cleared once engaged (or failed). + # Deferred speed auto: the load stays eager; generate() engages `default` at the 3rd generation. speed_deferred: bool = False # Successful generations on this load; drives the deferred engagement above. generation_count: int = 0 @@ -467,8 +446,7 @@ class _GenState: total_steps: int step: int = 0 - # Set when the first step finishes; the ETA rate is measured from there so the - # slower first step (warmup) doesn't skew it. + # Set when the first step finishes, so the slower warmup step doesn't skew the ETA rate. first_step_at: float = 0.0 # Computed once per step (in the callback) so it's stable between polls. eta_seconds: Optional[float] = None @@ -539,28 +517,23 @@ class DiffusionBackend: """Holds at most one loaded diffusers pipeline. All mutations are serialised.""" def __init__(self) -> None: - # _lock serialises the small state mutations; status()/load_progress()/ - # generate_progress() read lock-free so polling never blocks a slow load. + # _lock serialises the small state mutations; the status/progress readers stay lock-free. self._lock = threading.Lock() # _generate_lock serialises generations and is the ONLY lock the denoise holds. self._generate_lock = threading.Lock() self._state: Optional[_LoadState] = None self._loading: Optional[_LoadingState] = None - # Bumped on every begin_load/unload so a superseded/cancelled worker neither - # commits its pipeline nor stamps progress onto the current load. + # Bumped on begin_load/unload so a superseded worker neither commits nor stamps progress. self._load_token = 0 # Set by unload() to abort an in-flight (lock-free) download so an eviction preempts it. self._cancel_event = threading.Event() - # Cancel Event of the in-flight generation (or None), set under _lock to abort THAT - # denoise. Per-generation so a cancel can't be lost to a racing generate nor leak. + # Cancel Event of the in-flight generation; per-generation so a cancel can't be lost or leak. self._active_generate_cancel: Optional[threading.Event] = None # Written by the callback, read lock-free by generate_progress(). self._gen: Optional[_GenState] = None - # Image-conditioned workflow pipes (img2img/inpaint) built via from_pipe around the - # loaded pipe (shared modules, no extra VRAM). Keyed by class name; cleared on unload. + # img2img/inpaint pipes built via from_pipe (shared modules, no extra VRAM); cleared on unload. self._aux_pipes: dict[str, Any] = {} - # Loaded ControlNet models (id -> module) and their from_pipe pipelines - # ((class, cn_id) -> pipe), reusing resident base modules; cleared on unload. + # Loaded ControlNets and their from_pipe pipelines, reusing resident modules; cleared on unload. self._cn_models: dict[str, Any] = {} self._cn_pipes: dict[tuple[str, str], Any] = {} @@ -623,19 +596,16 @@ class DiffusionBackend: if speed is not None and str(speed).strip().lower() == SPEED_OFF: return False try: - # A definite-offload policy forces load_pipeline onto offload, so the dense build - # never runs and never touches the base transformer/ shards. Widening the prefetch - # would only waste a multi-GB pull -- and a disk-full here has NO GGUF fallback - # (unlike an in-load_pipeline dense failure). Mirror those offload gates. + # A definite-offload policy forces offload, so the dense build never runs; widening would waste + # a multi-GB pull, and a disk-full here has no GGUF fallback. Mirror those offload gates. mm = normalize_memory_mode(kwargs.get("memory_mode")) if mm in (MEMORY_MODE_BALANCED, MEMORY_MODE_LOW_VRAM): return False if mm is None and kwargs.get("cpu_offload"): return False target = self._resolve_device_target(fam) - # Only widen when the loader would actually take the dense path: resolve the SAME - # candidate load_pipeline re-plans against (which also checks the cache has disk room). - # When disk/scheme/support/a prequant rule it out, don't eagerly pull the shards. + # Only widen when the loader would take the dense path: resolve the same candidate load_pipeline + # re-plans against. When disk/scheme/support/a prequant rule it out, don't pull the shards. candidate = resolve_dense_quant_candidate( fam = fam, target = target, @@ -645,15 +615,12 @@ class DiffusionBackend: force_dense = bool(kwargs.get("loras")), logger = None, ) - # A prequant candidate loads a small checkpoint, not the dense transformer/ shards, - # so widening for it defeats the savings and can disk-full the fallback-less begin_load pull. + # A prequant loads a small checkpoint, so widening defeats the savings and can disk-full. if candidate is None or candidate.prequant: return False - # Capacity gate: on a device that cannot hold even the candidate's post-quant - # resident set, load_pipeline's re-plan is CERTAIN to decline the dense path, so - # widening would fetch the multi-GB base transformer/ shards (47 GB on Qwen-Image) - # only to run the GGUF as-is. Mirror plan_fits_total_capacity's bar against TOTAL - # capacity (not instantaneous free, which a transient tenant could undercount). + # Capacity gate: if the device cannot hold the candidate's post-quant resident set, load_pipeline + # is certain to decline the dense path, so mirror plan_fits_total_capacity against TOTAL capacity + # (not instantaneous free, which a transient tenant could undercount). from .diffusion_memory import ( _reserve_mib, snapshot_device_memory, @@ -725,9 +692,8 @@ class DiffusionBackend: kind = resolve_model_kind(gguf_filename, model_kind) fam = detect_family_for_pick(repo_id, gguf_filename, family_override) if fam is None: - # A deliberately-excluded model gets its stated reason, not the generic - # unknown-family message (which reads like a detection gap and invites a - # family_override retry that would fail deeper and less clearly). + # A deliberately-excluded model gets its stated reason, not the unknown-family message (which + # reads like a detection gap and invites a family_override retry that fails deeper). excluded = excluded_model_reason(repo_id) if excluded: raise ValueError(f"'{repo_id}' cannot be loaded: {excluded}") @@ -737,15 +703,13 @@ class DiffusionBackend: f"pass family_override with that family name. (Video models and image models " f"whose diffusers transformer has no single-file loader are not supported.)" ) - # Families whose single file IS the whole pipeline (SDXL) have no transformer-only - # GGUF path; reject GGUF here, before the route evicts the current model. + # Families whose single file IS the whole pipeline have no GGUF path; reject before eviction. if kind == "gguf" and fam.single_file_is_pipeline: raise ValueError( f"'{fam.name}' checkpoints are whole-pipeline single files and have no GGUF " f"transformer variant; load the .safetensors pipeline instead of a GGUF." ) - # A multi-denoiser family (Ideogram 4's dual DiTs) has no transformer-only path; - # a single-file/GGUF load would miss its second DiT. Reject here, before eviction. + # A multi-denoiser family (Ideogram 4) has no transformer-only path; reject before eviction. if kind in ("gguf", "single_file") and fam.pipeline_only: raise ValueError( f"'{fam.name}' loads only as a full diffusers pipeline (it assembles " @@ -758,18 +722,16 @@ class DiffusionBackend: f"Non-GGUF diffusion loads are restricted to unsloth/* repos (or a local " f"path); got '{repo_id}'. Pass a gguf_filename to load a GGUF instead." ) - # A companion base repo also loads via from_pretrained, so it must clear the same - # trust bar (else a GGUF pick could smuggle in an arbitrary remote base). Gate here. + # The companion base repo also loads via from_pretrained, so it must clear the same trust bar. if base_repo and base_repo.strip() and not _is_trusted_diffusion_repo(base_repo): raise ValueError( f"base_repo is restricted to unsloth/* repos (or a local path); got " f"'{base_repo}'." ) - # A local base_repo loads as a full pipeline (needs model_index.json); reject a - # non-pipeline local base here, before eviction. + # A local base_repo loads as a full pipeline; reject a non-pipeline one before eviction. _assert_local_base_is_pipeline(base_repo) - # Reject a bad LOCAL pick now so the route never evicts chat for an unloadable request. - # A path-shaped repo_id is meant to be on disk; a bare "org/name" is a remote HF repo. + # Reject a bad LOCAL pick before the route evicts chat. A path-shaped repo_id must be on disk; + # a bare "org/name" is a remote HF repo. local_root = Path(repo_id).expanduser() # Path-shaped: "."/".." prefix, a backslash (never in "org/name"), or an absolute path. path_shaped = ( @@ -778,15 +740,13 @@ class DiffusionBackend: if kind in ("gguf", "single_file"): if not gguf_filename: raise ValueError(f"a single-file checkpoint name is required for a '{kind}' load.") - # Fail a kind/extension mismatch here, before the handoff: gguf needs .gguf, - # single_file must not be handed a .gguf. + # Fail a kind/extension mismatch before the handoff: gguf needs .gguf, single_file must not. is_gguf_name = gguf_filename.lower().endswith(".gguf") if kind == "gguf" and not is_gguf_name: raise ValueError("a 'gguf' load requires a .gguf checkpoint name.") if kind == "single_file" and is_gguf_name: raise ValueError("a .gguf checkpoint needs model_kind 'gguf', not 'single_file'.") - # A single-file load must name an actual .safetensors checkpoint (else it evicts - # chat and only fails in the background from_single_file). + # A single-file load must name a real .safetensors, else it evicts chat then fails in background. if kind == "single_file" and not gguf_filename.lower().endswith(".safetensors"): raise ValueError( f"'{gguf_filename}' is not a loadable single-file checkpoint " @@ -809,8 +769,7 @@ class DiffusionBackend: elif path_shaped: raise FileNotFoundError(f"Local model path does not exist: {repo_id}") elif repo_id.upper().endswith("-GGUF"): - # A remote "*-GGUF" id is a GGUF repo, not a pipeline; loading it as a pipeline - # would evict chat then fail on the missing model_index.json. Reject here. + # A remote "*-GGUF" id is not a pipeline; reject here instead of evicting chat then failing. raise ValueError( f"'{repo_id}' is a single-file GGUF repo; load it with model_kind 'gguf' " f"and a .gguf filename, not as a full pipeline." @@ -843,8 +802,7 @@ class DiffusionBackend: """Validate, then run the (slow) load on a daemon thread. Returns at once.""" # A blank token must mean "anonymous", not an empty credential the Hub 401s. hf_token = (hf_token.strip() if isinstance(hf_token, str) else hf_token) or None - # base_repo is already gated at the route's pre-eviction validate; this re-validation - # is a cheap-fail guard for the resolved repo/family and does not re-gate base_repo. + # base_repo is gated at the route pre-eviction; this only cheap-fails the resolved repo/family. fam = self.validate_load_request( repo_id, gguf_filename = gguf_filename, @@ -892,8 +850,7 @@ class DiffusionBackend: def _run_load(self, **kwargs: Any) -> None: token = kwargs.get("_load_token") try: - # Resolve the base repo and estimate sizes on this thread (both network calls) so - # begin_load returns instantly. Only writer of _loading's fields here. + # Resolve the base repo and estimate sizes here (both network) so begin_load returns instantly. fam = detect_family_for_pick( kwargs["repo_id"], kwargs.get("gguf_filename"), kwargs.get("family_override") ) @@ -913,8 +870,7 @@ class DiffusionBackend: kwargs.get("hf_token"), kind = kind, single_file_is_pipeline = bool(fam and fam.single_file_is_pipeline), - # The dense-quant path otherwise pulls the base transformer/ shards inside the - # locked finalize (unpreemptable); when it can run, pull them in the prefetch here. + # Pull the base shards here rather than inside the locked, unpreemptable finalize. include_transformer = kind == "gguf" and self._dense_quant_prefetch_needed(fam, kwargs), ) @@ -941,8 +897,7 @@ class DiffusionBackend: if self._load_token != token: return logger.error("diffusion.load_failed: %s", exc) - # Free the debris of a failed construction (uncommitted _state, so nothing else - # reclaims the VRAM). Guarded so a sticky CUDA error can't skip stamping the real error. + # Free the debris of a failed construction; guarded so a sticky CUDA error still stamps the error. try: clear_gpu_cache() except Exception: # noqa: BLE001 @@ -962,8 +917,7 @@ class DiffusionBackend: if loading is None: return _progress("ready" if self._state is not None else None) - # Sum checkpoint + companion base cache; for a full-pipeline load base IS the repo, - # so count it once (else the bar double-counts). + # Sum checkpoint + companion cache; for a full-pipeline load base IS the repo, so count once. downloaded = self._cache_bytes(loading.repo_id) if loading.base_repo and loading.base_repo != loading.repo_id: downloaded += self._cache_bytes(loading.base_repo) @@ -1033,8 +987,8 @@ class DiffusionBackend: if sizes_out is not None: sizes_out[repo_id] = total return total, base_files - # Skip the Hub size lookup for a LOCAL gguf path: model_info would raise on a - # filesystem path and (caught below) skip the base lookup, forcing a synchronous companion pull. + # Skip the Hub size lookup for a LOCAL gguf path: model_info raises on a filesystem path and + # (caught below) would force a synchronous companion pull. if gguf_filename and not Path(repo_id).expanduser().exists(): info = api.model_info(repo_id, files_metadata = True, token = hf_token) gguf_bytes = sum(s.size or 0 for s in info.siblings if s.rfilename == gguf_filename) @@ -1252,23 +1206,19 @@ class DiffusionBackend: transformer_cache: Optional[str] = None, transformer_cache_threshold: Optional[float] = None, model_kind: Optional[str] = None, - # LoRA adapters to BAKE into a torchao int8/fp8 build (attached on the dense - # transformer before quantize_ + compile). Ignored by every other load kind: bf16 / - # bnb loads take adapters at generation time, GGUF-as-is has no dense transformer. + # LoRA adapters to BAKE into a torchao int8/fp8 build. Ignored by other kinds: bf16 / bnb take + # adapters at generation time, GGUF-as-is has no dense transformer. loras: Optional[list[tuple[str, float]]] = None, _load_token: Optional[int] = None, _base_local_dir: Optional[str] = None, ) -> dict[str, Any]: - # A blank/whitespace token must degrade to anonymous, not be passed as a credential - # the Hub client can error on. Normalize once for every branch below. + # A blank token must degrade to anonymous, not be passed as a credential. Normalize once. hf_token = hf_token.strip() if isinstance(hf_token, str) else hf_token hf_token = hf_token or None - # Validate first (cheap, no torch/diffusers) so a bad family fails even in a no-diffusers - # runtime. Re-sanitize the token (direct callers bypass begin_load). + # Validate first (no torch/diffusers) so a bad family fails even in a no-diffusers runtime. hf_token = (hf_token.strip() if isinstance(hf_token, str) else hf_token) or None - # base_repo is gated at the route before eviction; this re-validation cheap-fails the - # resolved repo/family and does not re-gate an already-validated base_repo. + # base_repo is gated at the route pre-eviction; this only cheap-fails the resolved repo/family. fam = self.validate_load_request( repo_id, gguf_filename = gguf_filename, @@ -1276,8 +1226,8 @@ class DiffusionBackend: model_kind = model_kind, ) kind = resolve_model_kind(gguf_filename, model_kind) - # Validate every mode string that can raise NOW, before this load evicts the previous - # pipeline. Validate-only for transformer_quant (keep the unset/auto vs explicit-off tri-state). + # Validate every mode string that can raise BEFORE this load evicts the previous pipeline. + # Validate-only for transformer_quant (keeps the unset/auto vs explicit-off tri-state). normalize_transformer_quant(transformer_quant) normalize_speed_mode(speed_mode) normalize_attention_backend(attention_backend) @@ -1292,10 +1242,8 @@ class DiffusionBackend: import diffusers - # Pre-install the optional attention kernel BEFORE the load locks: the wheel-only pip - # install can run up to 600s, and doing it under the lock would block unload/cancel that - # whole window. Only an explicit backend pulls a package (auto uses cuDNN/native from - # torch). Best-effort; the authoritative resolve + set under the lock is then a no-op. + # Pre-install the optional attention kernel before the load locks: the pip install can take 600s + # and would block unload/cancel that whole window. Best-effort; the locked resolve is a no-op. try: preinstall_backend = select_attention_backend( target, attention_backend, speed_active = True @@ -1305,11 +1253,9 @@ class DiffusionBackend: except Exception: # noqa: BLE001 — the locked path re-resolves and validates pass - # Signal an in-flight denoise to abort, then take _generate_lock to WAIT for it to exit - # before allocating the replacement (a load claims VRAM, so it must not overlap a live pipe). + # Abort an in-flight denoise and wait for it to exit: a load claims VRAM, so it must not overlap. with self._lock: - # Bail BEFORE signalling a cancel if this load was already superseded, else a stale - # worker would abort an unrelated live generation from the CURRENT model. + # Bail before signalling if this load was superseded, else a stale worker aborts a live one. if _load_token is not None and _load_token != self._load_token: raise RuntimeError("Diffusion load was cancelled.") if self._active_generate_cancel is not None: @@ -1332,8 +1278,7 @@ class DiffusionBackend: transformer_cls = getattr(diffusers, fam.transformer_class) pipeline_cls = getattr(diffusers, fam.pipeline_class) - # Decide placement up front (weights still on CPU, so free VRAM is the budget). - # Budgets the GGUF file; the dense-quant fast path is preflighted separately below. + # Decide placement up front (weights still on CPU). Budgets the GGUF; dense is preflighted below. plan = self._plan_memory( target, single_file_path, @@ -1345,33 +1290,28 @@ class DiffusionBackend: repo_id = repo_id, ) - # Dtype tri-state: unset/"auto" -> hardware ladder picks a quantised build (int8 - # min, fp8 on datacenter silicon) over GGUF-as-is; "none"/"off" pins GGUF-as-is; - # an explicit scheme pins it. An overwritten "auto" still records source=auto. + # Dtype tri-state: unset/"auto" -> hardware ladder picks a quantised build; "none"/"off" pins + # GGUF-as-is; an explicit scheme pins it. An overwritten "auto" still records source=auto. if transformer_quant is None or str(transformer_quant).strip().lower() in ( "", "auto", ): - # An explicit Speed="off" load must stay GGUF-as-is: auto-quant would engage - # int8/fp8 + compile and break the bit-exact request. "off" -> None (GGUF-as-is). + # An explicit Speed="off" stays GGUF-as-is: auto-quant would break the bit-exact request. speed_off = ( speed_mode is not None and str(speed_mode).strip().lower() == SPEED_OFF ) transformer_quant = "off" if speed_off else TQ_AUTO - # Default-on fast path: load the DENSE bf16 transformer and torchao-quantise it - # (int8/fp8/fp4 tensor cores), which beats GGUF's per-matmul dequant on speed AND - # quality at a higher-memory dense load. CUDA + bf16 + resident fit; ANY failure - # falls back to the GGUF build. GGUF kind only (it has the dense bf16 to materialise). + # Default-on fast path: load the dense bf16 transformer and torchao-quantise it, which beats + # GGUF's per-matmul dequant on speed AND quality. CUDA + bf16 + resident fit; any failure falls + # back to the GGUF build. GGUF kind only (it has the dense bf16 to materialise). pipe = None transformer_quant_engaged = None quant_plan = None - # The GGUF-size `plan` can mis-budget the fast path two ways, so preflight the real - # footprint BEFORE eviction; both branches need the base repo + a resolved scheme. + # The GGUF-size plan can mis-budget the fast path, so preflight the real footprint pre-eviction. dense_declined = False - # False when the memory plan only holds a PREQUANT-sized build: if the prequant - # load then fails, the loader must raise to GGUF instead of materialising the - # dense bf16 transformer the plan never budgeted for. + # False when the plan only holds a PREQUANT-sized build: if the prequant load fails, raise to + # GGUF instead of materialising the dense bf16 transformer the plan never budgeted for. dense_fallback_allowed = True if ( kind == "gguf" @@ -1379,17 +1319,15 @@ class DiffusionBackend: and dense_transformer_supported(target) ): if plan.offload_policy != OFFLOAD_NONE: - # The GGUF plan picked offload, but the quantised artifact is smaller - # (int8/fp8 ~half bf16; a prequant never materialises dense). Re-plan - # against the candidate's estimate -- a resident quant build beats an offloaded GGUF. + # The GGUF plan picked offload, but the quantised artifact is smaller, so re-plan against the + # candidate's estimate: a resident quant build beats an offloaded GGUF. candidate = resolve_dense_quant_candidate( fam = fam, target = target, requested = transformer_quant, base_repo = base, prequant_path = transformer_prequant_path, - # A LoRA bake skips the prequant shortcut, so size the candidate - # for the dense build it will actually run. + # A LoRA bake skips the prequant shortcut, so size for the dense build it will run. force_dense = bool(loras), logger = logger, ) @@ -1408,26 +1346,21 @@ class DiffusionBackend: transformer_resident_override_mib = ( candidate.transient_transformer_mib ), - # Pass the auto-policy's companion estimate so the prefetched base - # transformer/ shards in the cache aren't double-counted. + # Pass the companion estimate so prefetched base shards aren't double-counted. companion_override_mib = candidate.companions_mib, ) replanned = _replan_candidate() if ( replanned.offload_policy != OFFLOAD_NONE - # Explicit balanced/low_vram picks offload BY MODE; a fresh - # snapshot cannot change that, so don't waste a retry. + # Explicit balanced/low_vram picks offload BY MODE, so a fresh snapshot cannot change it. and normalize_memory_mode(memory_mode) not in (MEMORY_MODE_BALANCED, MEMORY_MODE_LOW_VRAM) and plan_fits_total_capacity(replanned) ): - # The candidate fits TOTAL device capacity with the standard - # reserve + resident margin, yet the instantaneous free reading - # said no: a transient foreign allocation (measured on B200: - # ~100 GB held for under a minute on an idle card) must not - # force the GGUF fallback. Re-snapshot (settled) and replan - # once before declining. + # The candidate fits TOTAL capacity with the standard reserve yet the instantaneous free + # reading said no: a transient foreign allocation (~100 GB for under a minute on an idle + # B200) must not force the GGUF fallback. Re-snapshot (settled) and replan once. replanned = _replan_candidate() if replanned.offload_policy != OFFLOAD_NONE: logger.info( @@ -1441,28 +1374,22 @@ class DiffusionBackend: ) if replanned.offload_policy == OFFLOAD_NONE: quant_plan = replanned - # The GGUF plan already declined resident; a prequant-sized - # replan says nothing about the (larger) dense transformer. + # The GGUF plan declined resident; a prequant-sized replan says nothing about the dense build. if candidate.prequant: dense_fallback_allowed = False else: - # The GGUF fits resident, but this path first materialises the base's dense - # bf16 transformer (bigger), so re-check the fit against THAT -- a card that - # fits the GGUF but not the dense must skip the fast path up front, not OOM - # after eviction. A prequant loads a small file (no dense), so skip the - # re-check there. _dense_transformer_resident_bytes returns 0 if shards are absent. + # This path first materialises the base's dense bf16 transformer (bigger than the GGUF), so + # re-check the fit against that, rather than OOMing after eviction. A prequant loads a small + # file, so skip the re-check there (_dense_transformer_resident_bytes is 0 without shards). scheme = select_transformer_quant_scheme( target, transformer_quant, # normalized above family = getattr(fam, "name", None), ) - # usable_prequant_source (not resolve_): a missing/non-allowlisted local - # path must NOT count as prequant here, or it skips the dense-fit re-check - # and OOMs materialising the dense transformer after eviction. + # usable_prequant_source (not resolve_): a missing/non-allowlisted local path must not count + # here, or it skips the dense-fit re-check and OOMs materialising the dense transformer. prequant = ( - # A LoRA bake skips the prequant shortcut (adapters attach on the - # dense transformer), so the dense-fit re-check must gate the fast - # path exactly as if no prequant source existed. + # A LoRA bake skips the prequant shortcut, so gate the fast path as if no prequant existed. None if loras else usable_prequant_source( @@ -1491,10 +1418,8 @@ class DiffusionBackend: ) if dense_plan.offload_policy != OFFLOAD_NONE: dense_fallback_allowed = False - # Without a prequant source the dense build is the ONLY path, - # so a dense misfit skips the fast path entirely (as before); with - # one, the small prequant load proceeds and only the dense - # fallback is forbidden. + # Without a prequant source the dense build is the only path, so a dense misfit skips the fast + # path; with one, the small prequant load proceeds and only the dense fallback is forbidden. if prequant is None: dense_declined = True if ( @@ -1528,8 +1453,7 @@ class DiffusionBackend: ) pipe = None transformer_quant_engaged = None - # Drop the exception BEFORE clearing the cache: its traceback keeps the - # partially-built dense transformer/pipe alive, blocking VRAM reclaim. + # Drop the exception before clearing the cache: its traceback pins the dense transformer's VRAM. del exc # Guarded: a sticky CUDA error can raise; the fallback must reach the GGUF build. try: @@ -1546,11 +1470,8 @@ class DiffusionBackend: and normalize_transformer_quant(transformer_quant) is not None and any(w != 0 for (_lid, w) in (loras or ())) ): - # The client asked for adapters BAKED into a quantized build, and that - # build was declined (memory plan) or failed. The GGUF fallback cannot - # carry the adapters, so completing it would return HTTP success while - # silently generating without the requested LoRAs -- wrong output with - # no signal. Fail the load with the recovery options instead. + # Adapters were requested BAKED but that build was declined or failed. The GGUF fallback cannot + # carry them, so fail loudly instead of silently generating without the requested LoRAs. raise RuntimeError( "The requested LoRA adapters could not be applied: baking adapters " "requires the quantized (int8/fp8) transformer build, which was " @@ -1561,12 +1482,10 @@ class DiffusionBackend: if pipe is None: if kind == "pipeline": - # Full diffusers repo: from_pretrained pulls every component and re-applies - # any embedded quantization_config (e.g. bnb-4bit). + # Full diffusers repo: from_pretrained pulls every component and re-applies embedded quant config. if fam.name == KREA2_FAMILY_NAME: - # krea ships transformers-5.x configs the 4.x line can't parse; assemble - # per-component (see diffusion_krea2.py). The constructor path never - # sees pipe_kwargs, so the pre-cast TE is handed in directly. + # krea ships transformers-5.x configs the 4.x line can't parse; assemble per-component (see + # diffusion_krea2.py). The constructor path never sees pipe_kwargs, so pass the pre-cast TE. pipe = load_krea2_pipeline( repo_id, dtype, @@ -1582,8 +1501,7 @@ class DiffusionBackend: ).get("text_encoder"), ) elif fam.name == IDEOGRAM4_FAMILY_NAME: - # ideogram ships the same transformers-5.x Qwen stack as krea; assemble - # per-component too (see diffusion_ideogram4.py). + # ideogram ships the same transformers-5.x Qwen stack as krea; assemble per-component too. pipe = load_ideogram4_pipeline(repo_id, dtype, hf_token = hf_token) else: pipe_kwargs: dict[str, Any] = { @@ -1593,8 +1511,7 @@ class DiffusionBackend: if hf_token: pipe_kwargs["token"] = hf_token if fam.name == HIDREAM_FAMILY_NAME: - # The repo names a Llama text_encoder_4 it does not ship; - # supply it from the open mirror (diffusion_hidream.py). + # The repo names a Llama text_encoder_4 it does not ship; supply it from the open mirror. pipe_kwargs.update( hidream_te4_kwargs( dtype, @@ -1604,9 +1521,8 @@ class DiffusionBackend: target = target, ) ) - # A hosted pre-cast fp8 text encoder (when the family ships one and - # the runtime cast would engage) skips the dense TE download; the - # later quantize_text_encoders re-applies the cast idempotently. + # A hosted pre-cast fp8 text encoder skips the dense TE download; quantize_text_encoders + # re-applies the cast idempotently. pipe_kwargs.update( te_prequant_pipe_kwargs( fam, @@ -1618,15 +1534,13 @@ class DiffusionBackend: logger = logger, ) ) - # The prefetched snapshot dir keeps from_pretrained off the hub (its - # sweep re-pulls files the scoped prefetch skipped: 24 GB per FLUX.1). + # The prefetched snapshot dir keeps from_pretrained off the hub (24 GB per FLUX.1 otherwise). pipe = pipeline_cls.from_pretrained( _base_local_dir or repo_id, **pipe_kwargs ) elif kind == "single_file" and fam.single_file_is_pipeline: - # A single-file SDXL-style checkpoint is the WHOLE pipeline, so load it - # through the pipeline class; ``config`` points at the base repo so diffusers - # builds the correct structure around the single-file weights. + # A single-file SDXL-style checkpoint is the WHOLE pipeline, so load it through the pipeline + # class; ``config`` points at the base repo so diffusers builds the right structure. sf_pipe_kwargs: dict[str, Any] = { "torch_dtype": dtype, "config": base, @@ -1650,8 +1564,7 @@ class DiffusionBackend: sf_kwargs["quantization_config"] = diffusers.GGUFQuantizationConfig( compute_dtype = dtype ) - # sd.cpp-converted GGUFs prefix tensors with model.diffusion_model.; - # the FLUX.2 / Qwen-Image converters crash or meta-strand on it. + # sd.cpp GGUFs prefix tensors with model.diffusion_model.; the FLUX.2 / Qwen converters choke. _install_gguf_prefix_strip(transformer_cls, logger) # A safetensors single-file (fp8) carries its own dtype: no GGUF dequant config. transformer = transformer_cls.from_single_file( @@ -1694,8 +1607,8 @@ class DiffusionBackend: target = target, ) ) - # Same pre-cast TE injection as the full-pipeline branch: the GGUF - # supplies the transformer, so the companion TE is the big download. + # Same pre-cast TE injection as above: the GGUF supplies the transformer, so the TE is the + # big download. pipe_kwargs.update( te_prequant_pipe_kwargs( fam, @@ -1711,31 +1624,29 @@ class DiffusionBackend: _base_local_dir or base, **pipe_kwargs ) - # Effective speed: GGUF defaults to near-lossless `default` (compile ~2.2x, below - # the quant noise floor); dense stays bit-identical `off`. Explicit is honored. + # Effective speed: GGUF defaults to near-lossless `default` (~2.2x, below the quant noise floor); + # dense stays bit-identical `off`. Explicit is honored. effective_speed = resolve_speed_mode(speed_mode, is_gguf = kind == "gguf") - # A torchao-quantized dense transformer must be compiled (eager is ~30x slower and - # loses to GGUF), so force at least `default` when quant engaged. + # A torchao-quantized dense transformer must be compiled (eager is ~30x slower, losing to GGUF). if transformer_quant_engaged is not None and effective_speed == SPEED_OFF: logger.info( "diffusion.transformer_quant: forcing speed_mode=default " "(quantized transformer must be compiled; eager is ~30x slower)" ) effective_speed = SPEED_DEFAULT - # Deferred speed auto for dense models: stay eager (a one-off image shouldn't pay - # the 25-60s compile), but generate() engages `default` on the 3rd image, where - # repeated use amortises it. Only when speed was unset, nothing forced compile, and this device can compile. + # Deferred speed auto for dense models: stay eager, but generate() engages `default` on the 3rd + # image where repeated use amortises the compile. Only when speed was unset and nothing forced it. speed_deferred = ( speed_mode is None and effective_speed == SPEED_OFF and transformer_quant_engaged is None and compile_eligible(target, is_gguf = False, family = fam) ) - # Speed optims run BEFORE placement (channels_last/compile precede offload). - # Snapshot the global backend flags (TF32/cudnn.benchmark) first for unload restore. + # Speed optims run BEFORE placement (channels_last/compile precede offload). Snapshot the global + # backend flags first for unload restore. backend_flags_before = snapshot_backend_flags() - # Pick the attention kernel BEFORE compile. auto upgrades to cuDNN fused attention - # on NVIDIA when a speed profile is active (~1.18x); explicit is honored. + # Pick the attention kernel BEFORE compile. auto upgrades to cuDNN fused attention on NVIDIA + # when a speed profile is active (~1.18x); explicit is honored. attention_engaged = apply_attention_backend( pipe, select_attention_backend( @@ -1743,10 +1654,9 @@ class DiffusionBackend: ), logger = logger, ) - # Step caching (First-Block-Cache), also before compile: reuses the transformer - # tail across steps (~1.4x on Flux at LPIPS ~0.08); when engaged, compile drops - # fullgraph (graph break). Tri-state: unset/"auto" -> step-count policy decides - # (engage when the DEFAULT schedule reaches FBCACHE_MIN_STEPS); "off"/"fbcache" pinned. + # Step caching (First-Block-Cache), also before compile: reuses the transformer tail across + # steps (~1.4x on Flux at LPIPS ~0.08); when engaged, compile drops fullgraph. Tri-state: + # unset/"auto" -> step-count policy decides (FBCACHE_MIN_STEPS); "off"/"fbcache" pinned. cache_request = normalize_transformer_cache(transformer_cache) cache_auto = transformer_cache is None or cache_request == TC_AUTO cache_quant_active = transformer_quant_engaged is not None or bool(gguf_filename) @@ -1764,8 +1674,7 @@ class DiffusionBackend: quant_active = cache_quant_active, logger = logger, ) - # An auto decision can flip at generation time, but only on a cache-capable - # transformer (a non-CacheMixin one keeps fullgraph). + # An auto decision can flip at generation time, but only on a cache-capable transformer. cache_may_toggle = cache_auto and callable( getattr(getattr(pipe, "transformer", None), "enable_cache", None) ) @@ -1784,12 +1693,10 @@ class DiffusionBackend: ) else: cache_reason = "requested" - # Everything from here to the _LoadState commit mutates PROCESS-WIDE state (class - # patches, TORCHINDUCTOR_CACHE_DIR, backend flags). _unload_locked reverses it via - # _state, so a pre-commit failure would leak it; the try/finally below restores on failure. - # gguf_transformer: the GGUF-specific compiled dequant applies only when the GGUF - # was actually loaded. On the dense fast path gguf_filename is still set (fallback) - # but pipe.transformer is dense (needs REGIONAL block compile), so treat it non-GGUF. + # Everything from here to the _LoadState commit mutates PROCESS-WIDE state (class patches, + # TORCHINDUCTOR_CACHE_DIR, backend flags); the try/finally below restores it on failure. + # gguf_transformer: on the dense fast path gguf_filename is still set (fallback) but + # pipe.transformer is dense (needs REGIONAL block compile), so treat it as non-GGUF. gguf_transformer = kind == "gguf" and transformer_quant_engaged is None eager_patched = False @@ -1808,17 +1715,15 @@ class DiffusionBackend: try: if effective_speed != SPEED_OFF: install_compile_safe_patches() - # Per-arch compile-safe fusions (qwen _modulate / z-image residual, etc.); - # neutral under compile, tracked by the same eager_patched flag. + # Per-arch compile-safe fusions; neutral under compile, tracked by the same eager_patched flag. install_arch_patches() eager_patched = True else: uninstall_patches() uninstall_arch_patches() - # Pre-warmed torch.compile cache: point inductor at a per-fingerprint dir and - # load a matching bundle before the first compiled forward, so the 25-58s - # compile is paid once and reused. A miss is silent -> local compile. + # Pre-warmed torch.compile cache: point inductor at a per-fingerprint dir and load a matching + # bundle before the first compiled forward, so the 25-58s compile is paid once. A miss is silent. if effective_speed in (SPEED_DEFAULT, SPEED_MAX) and compile_eligible( target, is_gguf = gguf_transformer, family = fam ): @@ -1828,16 +1733,13 @@ class DiffusionBackend: transformer = getattr(pipe, "transformer", None) or getattr(pipe, "unet", None), dtype = getattr(target, "dtype", None), - # A GGUF transformer compiles a DIFFERENT graph (the dequant op - # chain under `default`, dequant+block fusion under `max`) than a - # dense load of the same family, so it must key its own bundles - # instead of sharing (and cross-hitting) the dense fingerprint. + # A GGUF transformer compiles a DIFFERENT graph than a dense load of the same family, so it + # must key its own bundles instead of cross-hitting the dense fingerprint. quant = "gguf" if gguf_transformer else transformer_quant_engaged, attention_backend = attention_engaged, compile_kwargs = { - # Mirrors apply_speed_optims' fullgraph decision: an active or - # still-toggleable step cache OR a planned offload graph-breaks, - # so the cached bundle must key on the same fullgraph setting. + # Mirrors apply_speed_optims' fullgraph decision: an active or still-toggleable step cache OR + # a planned offload graph-breaks, so the bundle keys on the same fullgraph setting. "fullgraph": cache_engaged is None and not cache_may_toggle and plan.offload_policy == OFFLOAD_NONE, @@ -1861,15 +1763,15 @@ class DiffusionBackend: logger = logger, ) if transformer_quant_engaged is not None and not speed_applied.get("compiled"): - # Compile couldn't engage: the quantized transformer runs eager, far slower - # than the GGUF it replaced. Surface it loudly. + # Compile couldn't engage: the quantized transformer runs eager, far slower than the GGUF it + # replaced. Surface it loudly. logger.warning( "diffusion.transformer_quant: %s engaged but the transformer is NOT " "compiled; eager torchao quant is ~30x slower than GGUF here", transformer_quant_engaged, ) - # Quantise the dense companion text encoder(s) (opt-in), before placement so - # offload moves the smaller weights. Family drives int8's keep-bf16 schedule. + # Quantise the dense companion text encoder(s) before placement so offload moves the smaller + # weights. Family drives int8's keep-bf16 schedule. te_quant = quantize_text_encoders( pipe, target, @@ -1879,16 +1781,11 @@ class DiffusionBackend: logger = logger, ) - # Inference-side persistent conditioning cache (opt-in via - # UNSLOTH_DIFFUSION_COND_CACHE_DIR, the inference sibling of the - # trainers' cond_cache_dir): repeated prompts skip the text-encoder - # forward entirely, so warm prompts never onload the multi-GB encoders - # under an offload policy. Installed AFTER the TE quant above so the - # cache key reflects the encoders that actually run. ``base`` keys the - # companion repo the TEXT ENCODERS came from (a GGUF/single-file load - # takes them from the base, so the same checkpoint against a different - # base must not reuse the previous base's embeddings). Instance-level; - # dies with the pipe on unload. + # Inference-side persistent conditioning cache (opt-in via UNSLOTH_DIFFUSION_COND_CACHE_DIR): + # repeated prompts skip the text-encoder forward, so warm prompts never onload the multi-GB + # encoders under offload. Installed AFTER the TE quant so the key reflects the encoders that + # run. ``base`` keys the companion repo the text encoders came from, so the same checkpoint + # against a different base can't reuse its embeddings. Instance-level; dies with the pipe. cond_cache.install( pipe, family = fam.name, @@ -1899,14 +1796,14 @@ class DiffusionBackend: logger = logger, ) - # Apply the planned placement; apply_memory_plan returns the (policy, tiling) - # ACTUALLY engaged so status stays honest. Idempotent for the `none` policy. + # Apply the planned placement; apply_memory_plan returns the policy/tiling ACTUALLY engaged so + # status stays honest. Idempotent for the `none` policy. effective_policy, effective_tiling = apply_memory_plan( pipe, plan, device = device, logger = logger ) - # Per-control provenance for status. cpu_offload=False is the unset default, - # so only True is an explicit request. + # Per-control provenance for status. cpu_offload=False is the unset default, so only True is + # an explicit request. resolved = build_resolved_record( { "speed_mode": ( @@ -2062,13 +1959,12 @@ class DiffusionBackend: # 1. Pre-quantized checkpoint, when one is configured for the resolved scheme. scheme = select_transformer_quant_scheme(target, mode, family = getattr(fam, "name", None)) if scheme is None: - # Bail BEFORE the multi-GB dense download: an unsupported scheme (fp8 on Ampere, - # nvfp4 off Blackwell) would otherwise materialise the transformer only to fail at - # quantize, after eviction. load_pipeline catches this and builds the GGUF pipeline. + # Bail BEFORE the multi-GB dense download: an unsupported scheme (fp8 on Ampere, nvfp4 off + # Blackwell) would otherwise materialise the transformer only to fail at quantize, after + # eviction. load_pipeline catches this and builds the GGUF pipeline. raise RuntimeError("transformer quant unsupported for this device/scheme") if fam is not None and not lora_specs: - # A LoRA bake needs the DENSE transformer (adapters attach before quantize_), so - # the prequant shortcut is skipped when adapters were requested. + # A LoRA bake needs the DENSE transformer (adapters attach before quantize_), so skip prequant. source = resolve_prequant_source( fam, scheme, path_override = prequant_path, base_repo = base ) @@ -2083,8 +1979,8 @@ class DiffusionBackend: scheme = scheme, # Reject a checkpoint with a different Linear filter so prequant matches runtime-quant. min_features = DEFAULT_MIN_LINEAR_FEATURES, - # Only enforced when the caller forces fp8 fast-accum; a checkpoint that baked - # the other choice falls to the dense path instead of using the baked kernels. + # Only enforced when the caller forces fp8 fast-accum; a checkpoint that baked the other + # choice falls to the dense path instead of using the baked kernels. fast_accum = fast_accum, logger = logger, ) @@ -2105,8 +2001,8 @@ class DiffusionBackend: # 2. Fallback: materialise the dense bf16 transformer and quantise it on-device. if not allow_dense_fallback: - # The memory plan only budgeted the prequant-sized build; materialising the dense - # bf16 transformer here would exceed it after eviction. Raise to the GGUF build. + # The memory plan only budgeted the prequant-sized build, so materialising the dense bf16 + # transformer would exceed it after eviction. Raise to the GGUF build. raise RuntimeError( "prequant checkpoint unavailable and the dense transformer does not fit resident" ) @@ -2130,10 +2026,9 @@ class DiffusionBackend: target = target, ) if lora_specs: - # Bake the adapters BEFORE quantize_: peft injects its wrappers on the dense - # Linears (the post-quant torchao dispatch would TypeError on a manually - # quantized module), then quantize_ converts only each wrapper's frozen - # base_layer while the "lora_" side path stays high precision. + # Bake the adapters BEFORE quantize_: peft injects its wrappers on the dense Linears (the + # post-quant torchao dispatch would TypeError), then quantize_ converts only each wrapper's + # frozen base_layer while the "lora_" side path stays high precision. baked = self._resolve_lora_set( [(i, w) for (i, w) in lora_specs if w != 0], family = getattr(fam, "name", None), @@ -2180,9 +2075,8 @@ class DiffusionBackend: """Assemble the diffusers pipeline around ``transformer`` and place it on ``device`` (a no-op for an already-placed pre-quantized transformer; it moves the companions).""" if getattr(fam, "name", None) == KREA2_FAMILY_NAME: - # krea ships transformers-5.x configs and no top-level tokenizer files, so - # Pipeline.from_pretrained dies in the tokenizer (vocab_file = None); assemble - # per-component like every other krea load path (see diffusion_krea2.py). + # krea ships transformers-5.x configs and no top-level tokenizer files, so from_pretrained dies + # in the tokenizer; assemble per-component like every other krea path (diffusion_krea2.py). krea_te = None if target is not None: krea_te = te_prequant_pipe_kwargs( @@ -2211,8 +2105,7 @@ class DiffusionBackend: if hf_token: pipe_kwargs["token"] = hf_token if getattr(fam, "name", None) == HIDREAM_FAMILY_NAME: - # The repo ships no Llama text_encoder_4; assemble it from the open mirror - # (diffusion_hidream.py) exactly like the full-pipeline load branch. + # The repo ships no Llama text_encoder_4; assemble it from the open mirror, as above. pipe_kwargs.update( hidream_te4_kwargs( dtype, @@ -2222,8 +2115,8 @@ class DiffusionBackend: target = target, ) ) - # Same pre-cast TE injection as the full-pipeline and GGUF branches: the dense - # fast path supplies only the transformer, so the companion TE is the big download. + # Same pre-cast TE injection as the other branches: the dense fast path supplies only the + # transformer, so the companion TE is the big download. if target is not None: pipe_kwargs.update( te_prequant_pipe_kwargs( @@ -2271,13 +2164,12 @@ class DiffusionBackend: re-plan, so the base repo's PREFETCHED transformer/ shards -- which land in the same blob cache _companion_cache_bytes sums -- are not counted as companions on top of transformer_resident_override_mib (a double-count of the transformer).""" - # Settled (max-over-reads) on cuda: a transient foreign allocation at the wrong instant - # otherwise makes an empty card look full and silently declines the resident/quant fast - # path (see settled_snapshot_device_memory). + # Settled (max-over-reads) on cuda: a transient foreign allocation otherwise makes an empty + # card look full and silently declines the resident/quant fast path. device_memory = settled_snapshot_device_memory(target) if kind == "pipeline": - # The whole repo is one cached download; cached bytes are the resident estimate - # (bnb-4bit/fp8 stay compressed). A LOCAL path isn't cached, so sum its on-disk weights. + # The whole repo is one cached download, so cached bytes are the resident estimate (bnb-4bit/ + # fp8 stay compressed). A LOCAL path isn't cached, so sum its on-disk weights. local_repo = Path(repo_id).expanduser() if repo_id else None if local_repo is not None and local_repo.is_dir(): cached = self._local_dir_weight_bytes(local_repo, exclude_transformer = False) @@ -2285,9 +2177,9 @@ class DiffusionBackend: cached = self._cache_bytes(repo_id) if repo_id else 0 cached_mib = int(cached // (1024 * 1024)) if cached else None model_dense_mib = estimate_safetensors_dense_mib(cached_mib) - # A repo can store weights NARROWER than the loaded dtype: ideogram-4's base ships its - # two DiTs as raw float8, so cached bytes undershoot the bf16 footprint ~2x and auto - # would OOM. When the size table knows the bf16 total for THIS repo, plan against the larger. + # A repo can store weights NARROWER than the loaded dtype: ideogram-4 ships its DiTs as raw + # float8, so cached bytes undershoot the bf16 footprint ~2x and auto would OOM. When the size + # table knows the bf16 total for THIS repo, plan against the larger. is_narrow_base = bool(repo_id) and repo_id.strip().lower() == fam.base_repo.lower() if ( not is_narrow_base @@ -2295,15 +2187,14 @@ class DiffusionBackend: and local_repo is not None and local_repo.is_dir() ): - # A local fp8 mirror never string-matches base_repo, so detect fp8 from the shard - # headers and reserve the bf16 footprint (a local nf4 mirror stays compressed). + # A local fp8 mirror never string-matches base_repo, so detect fp8 from the shard headers and + # reserve the bf16 footprint (a local nf4 mirror stays compressed). is_narrow_base = ideogram4_repo_is_fp8(repo_id) if is_narrow_base: table = family_bf16_components_gb(fam, fam.base_repo) if table is not None: - # Reserve the bf16 footprint from this network-free constant even when the - # cache estimate is absent; else model_dense_mib stays None ("size unknown -> - # resident") and the ~54 GB fp8 pipeline OOMs a card offload would have fit. + # Reserve the bf16 footprint from this network-free constant even without a cache estimate, + # else model_dense_mib stays None ("size unknown -> resident") and the ~54 GB pipeline OOMs. table_mib = int(sum(table) * (1000.0**3) / (1024.0 * 1024.0)) model_dense_mib = ( table_mib if model_dense_mib is None else max(model_dense_mib, table_mib) @@ -2311,12 +2202,12 @@ class DiffusionBackend: companion_mib = None else: if transformer_resident_override_mib is not None: - # Planning for the dense-quant candidate (not the file on disk): the auto-policy's - # estimate replaces the file-size derivation; companions stay measured from cache. + # Planning for the dense-quant candidate: the auto-policy's estimate replaces the file-size + # derivation; companions stay measured from cache. transformer_resident = transformer_resident_override_mib elif kind == "single_file": - # An fp8 checkpoint upcasts to bf16 on load (~2x resident); detect from the - # basename. Excludes the SDXL (single_file_is_pipeline) case (already bf16). + # An fp8 checkpoint upcasts to bf16 on load (~2x resident); detect from the basename. Excludes + # the SDXL (single_file_is_pipeline) case, already bf16. fp8_upcast = not getattr(fam, "single_file_is_pipeline", False) and ( "fp8" in Path(single_file_path).name.lower() if single_file_path else False ) @@ -2325,12 +2216,11 @@ class DiffusionBackend: ) else: transformer_resident = estimate_gguf_resident_mib(file_size_mib(single_file_path)) - # Companions (VAE + text encoders) load near on-disk size; sum the base-repo cache, - # or a LOCAL base's on-disk weights (the blob cache is empty for a local path). + # Companions (VAE + text encoders) load near on-disk size; sum the base-repo cache, or a LOCAL + # base's on-disk weights (the blob cache is empty for a local path). if companion_override_mib is not None: - # Re-planning the dense candidate: the prefetched transformer/ shards land in the - # SAME cache _companion_cache_bytes sums, so use the auto-policy's companion estimate - # instead of double-counting the transformer. + # Re-planning the dense candidate: the prefetched transformer/ shards land in the SAME cache + # _companion_cache_bytes sums, so use the auto-policy's estimate instead of double-counting. companion_mib = companion_override_mib else: companion = self._companion_cache_bytes(base) @@ -2338,8 +2228,8 @@ class DiffusionBackend: model_dense_mib = None if transformer_resident is not None: model_dense_mib = transformer_resident + (companion_mib or 0) - # Feed the variant hint (basename + repo) so estimate_image_runtime_mib sees distilled - # markers ("turbo"/"schnell") normalized out of fam.name (distilled needs ~15% less headroom). + # Feed the variant hint so estimate_image_runtime_mib sees distilled markers normalized out of + # fam.name (distilled needs ~15% less headroom). variant_hint = " ".join( p for p in ( @@ -2375,14 +2265,13 @@ class DiffusionBackend: return cached import diffusers - # torch_dtype=None is load-bearing: from_pipe otherwise recasts EVERY component to fp32, - # which upcasts the reused bf16 modules and hard-crashes the dense-quant path (a - # torchao+compiled transformer's tensor-subclass weights can't swap_tensors). None reuses - # the resident modules at their loaded dtype (the point of from_pipe). + # torch_dtype=None is load-bearing: from_pipe otherwise recasts EVERY component to fp32, which + # hard-crashes the dense-quant path (torchao tensor-subclass weights can't swap_tensors). None + # reuses the resident modules at their loaded dtype, the point of from_pipe. pipe = getattr(diffusers, class_name).from_pipe(state.pipe, torch_dtype = None) # Publish to the shared aux cache only if THIS load is still current: from_pipe runs under - # _generate_lock but NOT _lock, so an unload can null _state while it builds; caching then - # would hand a wrapper over stale modules to a later load. + # _generate_lock but NOT _lock, so an unload can null _state while it builds, and caching would + # hand a wrapper over stale modules to a later load. with self._lock: if self._state is state: self._aux_pipes[class_name] = pipe @@ -2406,30 +2295,27 @@ class DiffusionBackend: if cn_model is None: if cancel.is_set(): raise RuntimeError(DIFFUSION_CANCELLED_MSG) - # resolve_controlnet accepts a bare owner/name without the base trust gate, and - # from_pretrained deserializes it (a malicious pickle would execute), so run the same - # Hub malware preflight the chat/export loaders use. A local dir is exempt (fail-open). - # The preflight fails OPEN when the Hub scan is unavailable, so a remote repo also forces - # safetensors below, closing the pickle RCE vector even when the scan could not run. + # resolve_controlnet accepts a bare owner/name without the base trust gate, and from_pretrained + # would execute a malicious pickle, so run the same Hub malware preflight the chat/export + # loaders use (a local dir is exempt). The preflight fails OPEN, so a remote repo also forces + # safetensors below, closing the pickle RCE vector when the scan could not run. remote_cn = not getattr(resolved_cn, "is_local", False) if remote_cn: from utils.security import evaluate_file_security _cn_fs = evaluate_file_security(resolved_cn.path, hf_token = state.hf_token or None) if _cn_fs.blocked: raise ValueError(_cn_fs.reason) - # Keep at most one ControlNet resident: evict the previous module + wrapper first, - # or swapping ControlNets within a load accumulates until OOM. + # Keep at most one ControlNet resident, else swapping ControlNets accumulates until OOM. if self._cn_models or self._cn_pipes: self._cn_models.clear() self._cn_pipes.clear() clear_gpu_cache() import torch - # state.dtype is the display string ("bfloat16"), not a torch.dtype; pass the real - # dtype so diffusers loads at the base compute dtype, not float32 (extra VRAM). + # state.dtype is the display string ("bfloat16"), so pass the real dtype and avoid a float32 load. cn_dtype = getattr(torch, str(state.dtype).replace("torch.", ""), None) - # Force safetensors for an untrusted remote repo: if the Hub scan failed open above, an - # embedded pickle would still deserialize on load. A local dir the user chose is exempt. + # Force safetensors for an untrusted remote repo: if the Hub scan failed open above, an embedded + # pickle would still deserialize. A local dir the user chose is exempt. cn_from_pretrained_kwargs: dict[str, Any] = {"cache_dir": hub_cache_dir()} if remote_cn: cn_from_pretrained_kwargs["use_safetensors"] = True @@ -2440,12 +2326,12 @@ class DiffusionBackend: **cn_from_pretrained_kwargs, ) if cancel.is_set(): - # An unload raced the blocking download; bail BEFORE placement so we don't - # allocate onto a GPU _unload_locked() just freed. + # An unload raced the blocking download; bail BEFORE placement so we don't allocate onto a GPU + # _unload_locked() just freed. del cn_model raise RuntimeError(DIFFUSION_CANCELLED_MSG) - # Placement follows the base's offload policy: a resident base places it resident, an - # offloaded base streams it via group offloading. Best-effort; failure -> resident. + # Placement follows the base's offload policy: resident base places it resident, offloaded base + # streams it via group offloading. Best-effort; failure -> resident. if getattr(state, "offload_policy", OFFLOAD_NONE) != OFFLOAD_NONE and ( _offload_controlnet_module(cn_model, state.device, logger) ): @@ -2464,8 +2350,7 @@ class DiffusionBackend: state.pipe, controlnet = cn_model, torch_dtype = None ) with self._lock: - # Same race as the model cache: an unload may have cleared _cn_pipes while - # from_pipe ran; caching now would pin a pipeline over the unloaded base. + # Same race as the model cache: an unload may have cleared _cn_pipes while from_pipe ran. if cancel.is_set() or self._state is not state: del pipe raise RuntimeError(DIFFUSION_CANCELLED_MSG) @@ -2487,8 +2372,8 @@ class DiffusionBackend: if denoiser is None or vae is None: return try: - # Read the dtype from the parameters (a plain/compiled nn.Module may hide .dtype). - # Take the first FLOATING dtype (a GGUF transformer's leading params are packed uint8). + # Read the dtype from the parameters (a compiled nn.Module may hide .dtype). Take the first + # FLOATING dtype (a GGUF transformer's leading params are packed uint8). target_dtype = next( (p.dtype for p in denoiser.parameters() if p.dtype.is_floating_point), None, @@ -2714,10 +2599,8 @@ class DiffusionBackend: install_compile_safe_patches() install_arch_patches() object.__setattr__(state, "eager_patched", True) - # Re-run the load-time selection with the caller's ORIGINAL request (not a bare - # None): an explicit backend must survive the deferred upgrade. auto still upgrades - # to cuDNN here (speed_active=True), but an explicit "native"/"sage"/"flash" is - # honored verbatim rather than silently replaced by the auto cuDNN choice. + # Re-run the load-time selection with the caller's ORIGINAL request: auto still upgrades to + # cuDNN here (speed_active=True), but an explicit backend is honored verbatim. attention_engaged = apply_attention_backend( state.pipe, select_attention_backend(target, state.attention_request, speed_active = True), @@ -2736,8 +2619,8 @@ class DiffusionBackend: quant = "gguf" if gguf_transformer else state.transformer_quant, attention_backend = attention_engaged, compile_kwargs = { - # Mirrors the load-time fullgraph decision: an engaged or - # still-toggleable step cache OR an offload graph-breaks. + # Mirrors the load-time fullgraph decision: an engaged or still-toggleable step cache OR an + # offload graph-breaks. "fullgraph": state.transformer_cache is None and not state.cache_auto and state.offload_policy == OFFLOAD_NONE, @@ -2791,12 +2674,10 @@ class DiffusionBackend: guidance: float = 0.0, seed: Optional[int] = None, batch_size: int = 1, - # Batched multi-image generation (see diffusion_batched.py): ``prompts`` renders one - # image per prompt (txt2img only); ``seeds`` renders one image per seed. Each image - # gets its OWN torch.Generator: same-seed repeats at the same batch shape are - # bit-identical, and a solo regeneration with the recorded seed matches its batched - # rendition up to batch-size-dependent kernel numerics. ``batch_size`` alone derives - # per-image seeds as base..base+batch_size-1 (matching the native engine). + # Batched multi-image generation (diffusion_batched.py): ``prompts`` renders one image per + # prompt (txt2img only); ``seeds`` one per seed. Each image gets its OWN torch.Generator, so + # same-seed repeats at the same batch shape are bit-identical and a solo regeneration matches + # up to batch-size kernel numerics. ``batch_size`` alone derives seeds base..base+n-1. prompts: Optional[list[str]] = None, seeds: Optional[list[int]] = None, # Image-conditioned (base64/data-URL): init alone = img2img; init + mask = inpaint. @@ -2816,8 +2697,8 @@ class DiffusionBackend: import torch from PIL import Image - # Per-generation cancel Event that unload()/a superseding load set (registered under - # _lock below) to abort just this denoise. _generate_lock is the only lock the denoise holds. + # Per-generation cancel Event that unload()/a superseding load set (registered under _lock) to + # abort just this denoise. _generate_lock is the only lock the denoise holds. cancel = threading.Event() with self._generate_lock: with self._lock: @@ -2826,18 +2707,15 @@ class DiffusionBackend: raise RuntimeError(DIFFUSION_NOT_LOADED_MSG) # Register under _lock so unload()/a load can signal THIS generation. self._active_generate_cancel = cancel - # Publish an active (step 0) state now, before the slow pre-denoise setup (deferred - # compile, LoRA resolution, ControlNet build), so a reload's mount probe doesn't read - # idle while this generation holds _generate_lock and let a second generate queue - # behind it. The per-step callback swaps in its own _GenState at denoise start. + # Publish an active (step 0) state before the slow pre-denoise setup (deferred compile, LoRA + # resolution, ControlNet build), so a reload's mount probe doesn't read idle while this + # generation holds _generate_lock. The per-step callback swaps in its own _GenState. self._gen = _GenState(total_steps = steps) try: - # The local `state` ref keeps the pipe alive even if unload() nulls _state. - # Resolve the per-image (prompt, seed) jobs up front: N prompts, one prompt x - # N seeds, or the legacy single prompt whose batch derives per-image seeds as - # base..base+batch_size-1 (matching the native sd.cpp engine, so a gallery - # recipe replays any batch member alone). A fresh base seed is drawn in JS's - # safe-integer range (< 2**53) so it round-trips through JSON. + # The local `state` ref keeps the pipe alive even if unload() nulls _state. Resolve the + # per-image (prompt, seed) jobs up front: N prompts, one prompt x N seeds, or the legacy single + # prompt deriving base..base+batch_size-1 (matching the native sd.cpp engine, so a gallery + # recipe replays any batch member alone). A fresh base seed stays < 2**53 for JSON. jobs, seed = resolve_batch_jobs( prompt = prompt, prompts = prompts, @@ -2847,14 +2725,12 @@ class DiffusionBackend: draw_seed = torch.Generator(device = state.device).seed, ) - # Deferred speed auto: engage the compile profile on the 3rd image, before the LoRA/ - # workflow wiring (load-time ordering). Best-effort; a failure stays eager and never retries. - # NOT when a LoRA is requested: a compiled transformer rejects LoRA, so compiling - # here would permanently break every LoRA generation on this load. + # Deferred speed auto: engage the compile profile on the 3rd image, before the LoRA/workflow + # wiring. Best-effort; a failure stays eager and never retries. NOT when a LoRA is requested: + # a compiled transformer rejects LoRA, so this would break every LoRA generation on this load. lora_requested = any(w != 0 for (_id, w) in (loras or [])) - # Also stay eager while a PRIOR generation's adapters are attached: _apply_loras runs - # AFTER the engage below, so compiling would bake the adapter in and the swallowed - # unload_lora_weights() would leave it active forever. Defer until a later gen. + # Also stay eager while a PRIOR generation's adapters are attached: _apply_loras runs AFTER the + # engage below, so compiling would bake the adapter in permanently. Defer until a later gen. loras_attached = bool(getattr(state.pipe, "_unsloth_loras", ())) if ( state.speed_deferred @@ -2873,15 +2749,15 @@ class DiffusionBackend: # Apply/adjust LoRA before picking the workflow pipe; from_pipe pipes share the transformer. self._apply_loras(state, loras, cancel) - # Select the workflow pipeline: txt2img uses the loaded pipe; img2img/inpaint reuse - # its modules via from_pipe; an edit model's own pipe is already the edit pipeline. + # Select the workflow pipeline: txt2img uses the loaded pipe; img2img/inpaint reuse its modules + # via from_pipe; an edit model's own pipe is already the edit pipeline. pipe = state.pipe init_pil = mask_pil = None control_pil = None cn_scale = cn_gstart = cn_gend = cn_mode = None ref_extra: list = [] - # Validate dependencies up front: mask/upscale/reference need an input image, and - # reference needs a supporting family (else the combo silently falls back to txt2img). + # Validate dependencies up front: mask/upscale/reference need an input image, and reference + # needs a supporting family (else the combo silently falls back to txt2img). if init_image is None: if mask_image is not None: raise ValueError("mask_image requires an input image (init_image).") @@ -2895,8 +2771,8 @@ class DiffusionBackend: "model family." ) if getattr(state.family, "edit", False): - # Instruction editing: the loaded pipe IS the edit pipeline; always needs an - # input image, prompt is the instruction. No mask, no from_pipe. + # Instruction editing: the loaded pipe IS the edit pipeline and always needs an input image; + # the prompt is the instruction. No mask, no from_pipe. if init_image is None: raise ValueError( f"{state.family.name} is an image-editing model: provide an input image." @@ -2915,14 +2791,14 @@ class DiffusionBackend: init_pil = _decode_b64_image(init_image, mode = "RGB") mask_pil = _decode_b64_image(mask_image, mode = "L") elif init_image is not None and upscale is not None and upscale > 1.0: - # Upscale (hires fix): enlarge with Lanczos, then re-run img2img at low strength - # to add detail without redrawing. Shares the img2img pipeline via from_pipe. + # Upscale (hires fix): enlarge with Lanczos, then re-run img2img at low strength to add detail + # without redrawing. Shares the img2img pipeline via from_pipe. workflow = "upscale" pipe = self._workflow_pipe(state, state.family.img2img_pipeline_class, workflow) init_pil = _decode_b64_image(init_image, mode = "RGB") iw, ih = init_pil.size - # Cap the factor, then the absolute output (longest side 2048, txt2img's max) - # to avoid an OOM-scale latent; round to a multiple of 16 (VAE downsample + patch). + # Cap the factor, then the absolute output (longest side 2048) to avoid an OOM-scale latent; + # round to a multiple of 16 (VAE downsample + patch). factor = max(1.0, min(float(upscale), 4.0)) tw_f, th_f = iw * factor, ih * factor max_side = 2048 @@ -2940,9 +2816,8 @@ class DiffusionBackend: if strength is None: strength = 0.35 # hires-fix default: preserve content, add detail elif getattr(state.family, "reference", False) and init_image is not None: - # FLUX.2 reference conditioning: the loaded pipe takes the reference via `image` - # and generates at the REQUESTED size. No from_pipe, no strength; output size - # from the sliders. After inpaint/upscale so a mask/upscale on a reference family routes right. + # FLUX.2 reference conditioning: the loaded pipe takes the reference via `image` and generates + # at the REQUESTED size (no from_pipe, no strength). After inpaint/upscale so those route right. workflow = "reference" init_pil = _decode_b64_image(init_image, mode = "RGB") # Additional references (FLUX.2 combines a list); capped to bound VRAM. @@ -2956,8 +2831,8 @@ class DiffusionBackend: else: workflow = "txt2img" - # ControlNet (diffusers): txt2img only (not img2img/inpaint/edit). Builds the - # family's CN pipeline around resident modules and passes a control map. + # ControlNet (diffusers): txt2img only. Builds the family's CN pipeline around resident modules + # and passes a control map. if controlnet is not None: from core.inference import diffusion_controlnet cn_id, cn_image_b64, cn_type, cn_strength, cn_gs, cn_ge = controlnet @@ -2984,8 +2859,8 @@ class DiffusionBackend: "diffusers engine (needs a bf16 or bnb-4bit load of a family with a " "ControlNet pipeline; not GGUF-via-diffusers or torchao fp8/int8)." ) - # Decode + preprocess the control image FIRST so a bad image 400s before - # any CN download/build. Control map at the OUTPUT size to align with latents. + # Decode + preprocess the control image FIRST so a bad image 400s before any CN download/build. + # Control map at the OUTPUT size to align with latents. src = _decode_b64_image(cn_image_b64, mode = "RGB") control_pil = diffusion_controlnet.preprocess_control(src, cn_type).resize( (width, height), Image.LANCZOS @@ -3002,21 +2877,20 @@ class DiffusionBackend: cn_scale, cn_gstart, cn_gend = cn_strength, cn_gs, cn_ge # Flux Union CN selects its head by an integer control_mode; map the type. cn_mode = diffusion_controlnet.union_control_mode(cn_id, cn_type) - # A prompt LIST batches plain text-to-image only: the image-conditioned and - # ControlNet workflows are tuned around one conditioning image per call, and a - # silent broadcast would pair every prompt with the same image. Multiple SEEDS - # of one prompt remain valid everywhere (per-image generators only). + # A prompt LIST batches plain text-to-image only: the image-conditioned and ControlNet workflows + # take one conditioning image per call, and a silent broadcast would pair every prompt with the + # same image. Multiple SEEDS of one prompt remain valid everywhere. if uniform_prompt(jobs) is None and workflow != "txt2img": raise ValueError( "A prompts list is supported for plain text-to-image only; the " f"{workflow} workflow takes one prompt per call (seed lists still work)." ) - # Snap odd-sized inputs to a multiple of 16 for workflows whose OUTPUT size comes - # from the input image (img2img/inpaint/edit); txt2img/reference use the slider, - # upscale already produced a /16 target. Mask is matched to the snapped image. + # Snap odd-sized inputs to a multiple of 16 for workflows whose OUTPUT size comes from the input + # image (img2img/inpaint/edit); txt2img/reference use the slider and upscale is already /16. + # Mask is matched to the snapped image. if init_pil is not None and workflow in ("img2img", "inpaint", "edit"): - # img2img/inpaint take output size from the upload, so bound the longest side to - # 2048 first (else a phone photo drives an OOM-scale latent). edit resizes internally. + # img2img/inpaint take output size from the upload, so bound the longest side to 2048 first + # (else a phone photo drives an OOM-scale latent). edit resizes internally. if workflow in ("img2img", "inpaint"): init_pil = _clamp_max_side(init_pil, 2048) init_pil = _snap_to_multiple(init_pil, 16) @@ -3031,25 +2905,23 @@ class DiffusionBackend: call_params = inspect.signature(pipe.__call__).parameters kwargs: dict[str, Any] = { - # prompt / generator / num_images_per_prompt are set per chunk below: - # each per-forward chunk carries its own prompt slice and one - # torch.Generator PER IMAGE (individual seed reproducibility). + # prompt / generator / num_images_per_prompt are set per chunk below: each chunk carries its own + # prompt slice and one torch.Generator PER IMAGE (individual seed reproducibility). "num_inference_steps": steps, # Most pipelines use "guidance_scale"; Qwen-Image uses "true_cfg_scale". state.family.cfg_kwarg: guidance, } if state.family.name == IDEOGRAM4_FAMILY_NAME: - # Ideogram 4 drives CFG via EITHER a constant guidance_scale OR a per-step - # guidance_schedule (check_inputs rejects both). At the advertised defaults drop - # the constant so the recommended 48-step taper engages; else null the schedule. + # Ideogram 4 drives CFG via EITHER a constant guidance_scale OR a per-step guidance_schedule + # (check_inputs rejects both). At the advertised defaults drop the constant so the recommended + # 48-step taper engages; else null the schedule. if steps == 48 and abs(float(guidance) - 7.0) < 1e-6: kwargs.pop(state.family.cfg_kwarg, None) else: kwargs["guidance_schedule"] = None if state.family.name == LUMINA2_FAMILY_NAME and "cfg_trunc_ratio" in call_params: - # Lumina 2's card recipe runs the CFG double-forward only over the FIRST - # quarter of the trajectory (cfg_trunc_ratio=0.25); the pipeline default (1.0) - # applies it everywhere, visibly oversaturating output. Constant card value. + # Lumina 2's card recipe runs the CFG double-forward over only the first quarter of the + # trajectory (cfg_trunc_ratio=0.25); the pipeline default (1.0) visibly oversaturates output. kwargs["cfg_trunc_ratio"] = 0.25 if init_pil is not None: # Reference passes the whole list (FLUX.2 combines); others take the single image. @@ -3058,9 +2930,9 @@ class DiffusionBackend: kwargs["mask_image"] = mask_pil if strength is not None and "strength" in call_params: kwargs["strength"] = strength - # width/height: txt2img uses the slider; image-conditioned pipes must use the INPUT - # IMAGE's own size (a differing slider mismatches the latents). Many img2img/inpaint - # pipes drop them entirely, so pass only when accepted, derived from the image. + # width/height: txt2img uses the slider; image-conditioned pipes must use the INPUT IMAGE's own + # size (a differing slider mismatches the latents). Many such pipes drop them entirely, so pass + # only when accepted, derived from the image. if workflow in ("txt2img", "reference", "controlnet"): # These generate at the REQUESTED size (reference/control image resized to match). kwargs["width"] = width @@ -3074,8 +2946,8 @@ class DiffusionBackend: if negative_prompt and "negative_prompt" in call_params: kwargs["negative_prompt"] = negative_prompt if workflow == "controlnet" and control_pil is not None: - # CN pipeline takes the control map + scale; guidance start/end bound its step - # range. Every kwarg is signature-gated so a family that omits one still runs. + # CN pipeline takes the control map + scale; guidance start/end bound its step range. Every + # kwarg is signature-gated so a family that omits one still runs. if "control_image" in call_params: kwargs["control_image"] = control_pil elif "image" in call_params: # some CN pipelines name it "image" @@ -3090,14 +2962,13 @@ class DiffusionBackend: if "control_mode" in call_params and cn_mode is not None: kwargs["control_mode"] = cn_mode - # Per-forward chunks: the whole job list in ONE forward by default (the - # measured sweet spot: batch 32 on 4-step models, ~10-22x over serial - # engines), bounded by an explicit batch_size cap; OOM backoff below - # halves a failed chunk instead of failing the call. + # Per-forward chunks: the whole job list in ONE forward by default (the measured sweet spot: + # batch 32 on 4-step models, ~10-22x over serial engines), bounded by an explicit batch_size + # cap; the OOM backoff below halves a failed chunk instead of failing the call. chunks = chunk_jobs(jobs, batch_size) gen = _GenState(total_steps = steps * len(chunks)) - # Steps completed by FINISHED chunks, so the progress bar spans the whole - # multi-chunk call (mutable cell: _on_step closes over it). + # Steps completed by FINISHED chunks, so the bar spans the whole multi-chunk call (mutable + # cell: _on_step closes over it). steps_done = [0] def _on_step(pipe, step_index, timestep, callback_kwargs): @@ -3117,13 +2988,12 @@ class DiffusionBackend: if "callback_on_step_end" in call_params: kwargs["callback_on_step_end"] = _on_step - # Re-check an AUTO cache decision against the ACTUAL step count (a 28-step request - # gains FBCache, a few-step turbo drops it); explicit choices never toggle. + # Re-check an AUTO cache decision against the ACTUAL step count (28 steps gains FBCache, a + # few-step turbo drops it); explicit choices never toggle. if state.cache_auto: - # Key on the EFFECTIVE denoise steps: an img2img/upscale request at strength < 1 - # only denoises a fraction of `steps`, so folding in `strength` (only when - # actually applied) keeps FBCache off the short trajectory. When strength is - # omitted the pipe's own default (< 1) still applies, so key on that too. + # Key on the EFFECTIVE denoise steps: an img2img/upscale request at strength < 1 only denoises + # a fraction of `steps`, so fold in `strength` (only when applied, else the pipe's own default) + # to keep FBCache off the short trajectory. strength_applied = effective_request_strength( strength, init_pil is not None, @@ -3139,8 +3009,8 @@ class DiffusionBackend: logger = logger, ) if toggled != state.transformer_cache: - # _LoadState is frozen; the one deliberate in-place update, tracking the - # pipe-level toggle so status() is truthful. + # _LoadState is frozen; the one deliberate in-place update, tracking the pipe-level toggle so + # status() is truthful. object.__setattr__(state, "transformer_cache", toggled) entry = (state.resolved or {}).get("transformer_cache") if isinstance(entry, dict): @@ -3163,8 +3033,8 @@ class DiffusionBackend: torch.Generator(device = state.device).manual_seed(s) for _, s in chunk ] if len(jobs) == 1: - # Single image: scalar prompt + generator, exactly the pre-batching - # call shape (the regression harness checks this path bit-identical). + # Single image: scalar prompt + generator, exactly the pre-batching call shape (the regression + # harness checks this path bit-identical). chunk_kwargs["prompt"] = shared chunk_kwargs["generator"] = generators[0] chunk_kwargs["num_images_per_prompt"] = 1 @@ -3174,13 +3044,10 @@ class DiffusionBackend: chunk_kwargs["generator"] = generators chunk_kwargs["num_images_per_prompt"] = len(chunk) else: - # Distinct prompts: one image per prompt in a single forward. The - # negative prompt must be broadcast to match: a pipeline either - # asserts on the length (Z-Image) or encodes a batch-1 negative - # against batch-N latents and fails in the transformer's txt/img - # concat (Qwen-Image, Krea 2, FLUX true-CFG). Only the pipes that - # already expand a scalar themselves (SDXL, Lumina 2, HiDream) are - # unaffected, so broadcast here for all of them. + # Distinct prompts: one image per prompt in a single forward. The negative prompt must be + # broadcast to match, else a pipeline asserts on the length (Z-Image) or fails in the + # transformer's txt/img concat (Qwen-Image, Krea 2, FLUX true-CFG). The pipes that already + # expand a scalar themselves are unaffected, so broadcast for all of them. chunk_kwargs["prompt"] = [p for p, _ in chunk] chunk_kwargs["generator"] = generators chunk_kwargs["num_images_per_prompt"] = 1 @@ -3188,12 +3055,9 @@ class DiffusionBackend: chunk_kwargs["negative_prompt"] = [ chunk_kwargs["negative_prompt"] ] * len(chunk) - # Start every forward from a clean step cache. diffusers resets the - # transformer's FBCache state at the END of a successful __call__ - # (maybe_free_model_hooks), but a call that RAISED -- the OOM below, a - # cancelled denoise, a failed prior generation -- leaves the residual of - # its own batch behind, and the next (differently sized) forward then - # compares against it: "size of tensor a (16) must match ... b (32)". + # Start every forward from a clean step cache: diffusers only resets FBCache state at the END + # of a SUCCESSFUL __call__, so a raised call (OOM, cancel) leaves its residual behind and the + # next differently sized forward compares against it ("size of tensor a (16) must match b (32)"). if state.transformer_cache: self._reset_step_cache(state.pipe) try: @@ -3203,8 +3067,8 @@ class DiffusionBackend: except Exception as exc: # noqa: BLE001 — reraised unless a splittable OOM if len(chunk) < 2 or not is_oom_error(exc): raise - # OOM backoff: halve the failed chunk and retry; finished chunks keep - # their images and per-image seeds keep every retry reproducible. + # OOM backoff: halve the failed chunk and retry; finished chunks keep their images and per-image + # seeds keep every retry reproducible. empty_cache = getattr(getattr(torch, "cuda", None), "empty_cache", None) if callable(empty_cache): empty_cache() @@ -3218,28 +3082,22 @@ class DiffusionBackend: len(second_half), ) continue - # A cancelled denoise returns a partial/garbage image; don't persist it - # (nor run any remaining chunks). + # A cancelled denoise returns a partial image; don't persist it, nor run remaining chunks. if cancel.is_set(): raise RuntimeError(DIFFUSION_CANCELLED_MSG) images.extend(out) per_image_seeds.extend(s for _, s in chunk) chunk_shapes.append(len(chunk)) steps_done[0] += steps - # Keep progress ACTIVE through the post-denoise work below (the compile-cache save can - # take a moment) -- don't null _gen here. The route persists the image AFTER this - # returns, so a reload's mount probe that read idle now would refresh the gallery - # before the result exists. The outer finally clears _gen on every exit (return/raise). - # Persist the warm torch.compile bundle after the first compiled generation. A - # STATIC compile makes new artifacts per (w,h,batch), so register this shape first - # (an uncovered shape re-dirties the context and the save rewrites the bundle). - # Idempotent + best-effort. + # Keep progress ACTIVE through the post-denoise work below (the compile-cache save can take a + # moment): the route persists the image AFTER this returns, so a mount probe reading idle now + # would refresh the gallery before the result exists. The outer finally clears _gen on any exit. + # Persist the warm torch.compile bundle after the first compiled generation. A STATIC compile + # makes new artifacts per (w,h,batch), so register this shape first. Idempotent, best-effort. try: - # Register the dims the forward ACTUALLY compiled with (image-conditioned - # workflows run at the input image's size, not the slider; see - # _compile_shape_dims). A static compile makes one artifact per BATCH size - # too, so every distinct chunk size this call ran (OOM backoff can produce - # several) registers its own (w, h, batch) shape. + # Register the dims the forward ACTUALLY compiled with (image-conditioned workflows run at the + # input image's size; see _compile_shape_dims). A static compile makes one artifact per BATCH + # size too, so every distinct chunk size this call ran registers its own (w, h, batch). reg_width, reg_height = _compile_shape_dims(workflow, init_pil, width, height) static_shapes = "compiled" in ( state.speed_optims or () @@ -3255,9 +3113,8 @@ class DiffusionBackend: pass # Count the finished generation (drives deferred speed); a batch is one generation. object.__setattr__(state, "generation_count", state.generation_count + 1) - # Return the PIL images (unencoded); the route embeds recipes and persists them. - # ``seeds`` records each image's OWN seed (the native engine already reports - # these), so every batch member replays individually from its recipe. + # Return the PIL images (unencoded); the route embeds recipes and persists them. ``seeds`` + # records each image's OWN seed, so every batch member replays individually from its recipe. return { "images": list(images), "seed": int(seed), @@ -3269,9 +3126,8 @@ class DiffusionBackend: with self._lock: if self._active_generate_cancel is cancel: self._active_generate_cancel = None - # Sole clear of the published progress state, on every exit (return, cancel, or a - # setup/denoise error), so it stays active through post-denoise work but a crashed - # generation never leaves the UI stuck "active". Safe under _generate_lock. + # Sole clear of the published progress state, on every exit, so it stays active through + # post-denoise work but a crashed generation never leaves the UI stuck "active". self._gen = None def generate_progress(self) -> dict[str, Any]: @@ -3304,9 +3160,8 @@ class DiffusionBackend: self._load_token += 1 self._loading = None # Wait for the signalled denoise to exit BEFORE tearing down: _unload_locked uninstalls - # process-wide state (attention patches, GGUF compile hooks, backend flags, compile cache) - # the denoise still depends on but its pipe ref doesn't pin. The denoise holds _generate_lock - # for its whole body, so acquiring it here blocks until it has finished. Mirrors begin_load. + # process-wide state (attention patches, GGUF compile hooks, backend flags, compile cache) the + # denoise still depends on. It holds _generate_lock for its whole body. Mirrors begin_load. with self._generate_lock: with self._lock: self._unload_locked() @@ -3316,8 +3171,8 @@ class DiffusionBackend: state = self._state if state is None: return - # Restore the process-wide backend flags this load flipped, so the next `off` load is - # bit-identical. compile_cache.restore + gguf_compile.uninstall_all likewise; all idempotent. + # Restore the process-wide backend flags this load flipped so the next `off` load is + # bit-identical; compile_cache.restore + gguf_compile.uninstall_all likewise. All idempotent. restore_backend_flags(state.backend_flags_before) compile_cache.restore(state.compile_cache_ctx) gguf_compile.uninstall_all() @@ -3328,9 +3183,8 @@ class DiffusionBackend: uninstall_patches() uninstall_arch_patches() - # Deliberately NOT unload_lora_weights() here: the whole pipe is dropped below, freeing any - # LoRA adapters with it. Both callers hold _generate_lock across this teardown, so no denoise - # is in flight. + # Deliberately NOT unload_lora_weights(): the whole pipe is dropped below, freeing any adapters + # with it. Both callers hold _generate_lock, so no denoise is in flight. # Drop the workflow pipes so they don't pin the freed pipeline's modules past unload. self._aux_pipes.clear() # Drop any ControlNet models + pipelines so the freed load carries no extra modules. @@ -3514,9 +3368,8 @@ def _base_file_downloaded(rfilename: str, *, include_transformer: bool = False) return not rfilename.startswith("assets/") -# Weight file extensions the base repo need NOT supply when the single file is the whole -# pipeline (SDXL): from_single_file(config=base) reads only the base repo's structure -# (config/tokenizer/scheduler) and takes the weights from the single file. +# Weight extensions the base repo need NOT supply when the single file is the whole pipeline +# (SDXL): from_single_file(config=base) reads only the base's structure, not its weights. _BASE_WEIGHT_EXTS = ( ".safetensors", ".bin", diff --git a/studio/backend/core/inference/diffusion_arch_patches.py b/studio/backend/core/inference/diffusion_arch_patches.py index 79de1f968e..7b3b1530d4 100644 --- a/studio/backend/core/inference/diffusion_arch_patches.py +++ b/studio/backend/core/inference/diffusion_arch_patches.py @@ -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, diff --git a/studio/backend/core/inference/diffusion_attention.py b/studio/backend/core/inference/diffusion_attention.py index c647453862..e1e9cfc2f1 100644 --- a/studio/backend/core/inference/diffusion_attention.py +++ b/studio/backend/core/inference/diffusion_attention.py @@ -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 diff --git a/studio/backend/core/inference/diffusion_auto_policy.py b/studio/backend/core/inference/diffusion_auto_policy.py index 32e3a3ee02..be5cea23b9 100644 --- a/studio/backend/core/inference/diffusion_auto_policy.py +++ b/studio/backend/core/inference/diffusion_auto_policy.py @@ -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 diff --git a/studio/backend/core/inference/diffusion_batched.py b/studio/backend/core/inference/diffusion_batched.py index 9e91d99e75..d7ff876389 100644 --- a/studio/backend/core/inference/diffusion_batched.py +++ b/studio/backend/core/inference/diffusion_batched.py @@ -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 diff --git a/studio/backend/core/inference/diffusion_cache.py b/studio/backend/core/inference/diffusion_cache.py index f6e444755f..a27cced1b9 100644 --- a/studio/backend/core/inference/diffusion_cache.py +++ b/studio/backend/core/inference/diffusion_cache.py @@ -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 diff --git a/studio/backend/core/inference/diffusion_compile_cache.py b/studio/backend/core/inference/diffusion_compile_cache.py index ab05a3d609..392134ac30 100644 --- a/studio/backend/core/inference/diffusion_compile_cache.py +++ b/studio/backend/core/inference/diffusion_compile_cache.py @@ -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)") diff --git a/studio/backend/core/inference/diffusion_cond_cache.py b/studio/backend/core/inference/diffusion_cond_cache.py index 8d849b5593..0dfe506206 100644 --- a/studio/backend/core/inference/diffusion_cond_cache.py +++ b/studio/backend/core/inference/diffusion_cond_cache.py @@ -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), diff --git a/studio/backend/core/inference/diffusion_controlnet.py b/studio/backend/core/inference/diffusion_controlnet.py index 8947e85bc6..0ed511525b 100644 --- a/studio/backend/core/inference/diffusion_controlnet.py +++ b/studio/backend/core/inference/diffusion_controlnet.py @@ -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) diff --git a/studio/backend/core/inference/diffusion_device.py b/studio/backend/core/inference/diffusion_device.py index bf85aa4522..ccfee91667 100644 --- a/studio/backend/core/inference/diffusion_device.py +++ b/studio/backend/core/inference/diffusion_device.py @@ -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( diff --git a/studio/backend/core/inference/diffusion_eager_patches.py b/studio/backend/core/inference/diffusion_eager_patches.py index cc1d4432cd..cb3ea6dbae 100644 --- a/studio/backend/core/inference/diffusion_eager_patches.py +++ b/studio/backend/core/inference/diffusion_eager_patches.py @@ -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 diff --git a/studio/backend/core/inference/diffusion_engine_router.py b/studio/backend/core/inference/diffusion_engine_router.py index f0e6a54707..9cb460bb67 100644 --- a/studio/backend/core/inference/diffusion_engine_router.py +++ b/studio/backend/core/inference/diffusion_engine_router.py @@ -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), diff --git a/studio/backend/core/inference/diffusion_families.py b/studio/backend/core/inference/diffusion_families.py index e14c81453f..6cd621cfaf 100644 --- a/studio/backend/core/inference/diffusion_families.py +++ b/studio/backend/core/inference/diffusion_families.py @@ -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), diff --git a/studio/backend/core/inference/diffusion_hidream.py b/studio/backend/core/inference/diffusion_hidream.py index 3820fbfd23..44a3141130 100644 --- a/studio/backend/core/inference/diffusion_hidream.py +++ b/studio/backend/core/inference/diffusion_hidream.py @@ -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, diff --git a/studio/backend/core/inference/diffusion_ideogram4.py b/studio/backend/core/inference/diffusion_ideogram4.py index 2b088a1d25..445788aee2 100644 --- a/studio/backend/core/inference/diffusion_ideogram4.py +++ b/studio/backend/core/inference/diffusion_ideogram4.py @@ -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 ) diff --git a/studio/backend/core/inference/diffusion_krea2.py b/studio/backend/core/inference/diffusion_krea2.py index 1d9ac0ebff..de77de68d2 100644 --- a/studio/backend/core/inference/diffusion_krea2.py +++ b/studio/backend/core/inference/diffusion_krea2.py @@ -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 " diff --git a/studio/backend/core/inference/diffusion_lora.py b/studio/backend/core/inference/diffusion_lora.py index 546f0d254e..e58fc0d2ae 100644 --- a/studio/backend/core/inference/diffusion_lora.py +++ b/studio/backend/core/inference/diffusion_lora.py @@ -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 ``.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 ``.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") diff --git a/studio/backend/core/inference/diffusion_memory.py b/studio/backend/core/inference/diffusion_memory.py index fcee79a243..40f052d9e3 100644 --- a/studio/backend/core/inference/diffusion_memory.py +++ b/studio/backend/core/inference/diffusion_memory.py @@ -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 " diff --git a/studio/backend/core/inference/diffusion_precision.py b/studio/backend/core/inference/diffusion_precision.py index b31ac24eaa..f0f8a1e5bc 100644 --- a/studio/backend/core/inference/diffusion_precision.py +++ b/studio/backend/core/inference/diffusion_precision.py @@ -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 diff --git a/studio/backend/core/inference/diffusion_prequant.py b/studio/backend/core/inference/diffusion_prequant.py index a805458d8d..23dd8a7982 100644 --- a/studio/backend/core/inference/diffusion_prequant.py +++ b/studio/backend/core/inference/diffusion_prequant.py @@ -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 diff --git a/studio/backend/core/inference/diffusion_speed.py b/studio/backend/core/inference/diffusion_speed.py index 1146046517..bb0714a56b 100644 --- a/studio/backend/core/inference/diffusion_speed.py +++ b/studio/backend/core/inference/diffusion_speed.py @@ -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: diff --git a/studio/backend/core/inference/diffusion_te_prequant.py b/studio/backend/core/inference/diffusion_te_prequant.py index 9f83116201..d2dc76ae5f 100644 --- a/studio/backend/core/inference/diffusion_te_prequant.py +++ b/studio/backend/core/inference/diffusion_te_prequant.py @@ -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, diff --git a/studio/backend/core/inference/diffusion_transformer_quant.py b/studio/backend/core/inference/diffusion_transformer_quant.py index 0e16c6524c..f92df82340 100644 --- a/studio/backend/core/inference/diffusion_transformer_quant.py +++ b/studio/backend/core/inference/diffusion_transformer_quant.py @@ -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, diff --git a/studio/backend/core/inference/gpu_arbiter.py b/studio/backend/core/inference/gpu_arbiter.py index 3d99b07e7d..d661845816 100644 --- a/studio/backend/core/inference/gpu_arbiter.py +++ b/studio/backend/core/inference/gpu_arbiter.py @@ -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() diff --git a/studio/backend/core/inference/image_gallery.py b/studio/backend/core/inference/image_gallery.py index 131c1eafb1..a23bddc508 100644 --- a/studio/backend/core/inference/image_gallery.py +++ b/studio/backend/core/inference/image_gallery.py @@ -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: diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 04b6bfb56b..6133530247 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -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(): diff --git a/studio/backend/core/inference/sd_cpp_args.py b/studio/backend/core/inference/sd_cpp_args.py index 4380e7b67d..6285e858bb 100644 --- a/studio/backend/core/inference/sd_cpp_args.py +++ b/studio/backend/core/inference/sd_cpp_args.py @@ -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 ```` tags are unsupported server-side), so LoRAs are staged and named here. + # (prompt-embedded ```` tags are unsupported server-side), so LoRAs are staged here. if lora: req["lora"] = lora return req diff --git a/studio/backend/core/inference/sd_cpp_backend.py b/studio/backend/core/inference/sd_cpp_backend.py index 991774b940..b96bf65d4a 100644 --- a/studio/backend/core/inference/sd_cpp_backend.py +++ b/studio/backend/core/inference/sd_cpp_backend.py @@ -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" diff --git a/studio/backend/core/inference/sd_cpp_engine.py b/studio/backend/core/inference/sd_cpp_engine.py index 9b1f72188b..1ed2900d6a 100644 --- a/studio/backend/core/inference/sd_cpp_engine.py +++ b/studio/backend/core/inference/sd_cpp_engine.py @@ -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"}) diff --git a/studio/backend/core/inference/sd_cpp_server.py b/studio/backend/core/inference/sd_cpp_server.py index dff17d8b31..a7c6061de0 100644 --- a/studio/backend/core/inference/sd_cpp_server.py +++ b/studio/backend/core/inference/sd_cpp_server.py @@ -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 diff --git a/studio/backend/core/inference/video.py b/studio/backend/core/inference/video.py index f4110d0c74..b7e7de2f79 100644 --- a/studio/backend/core/inference/video.py +++ b/studio/backend/core/inference/video.py @@ -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() diff --git a/studio/backend/core/inference/video_families.py b/studio/backend/core/inference/video_families.py index b073180f68..7c6e3f02d4 100644 --- a/studio/backend/core/inference/video_families.py +++ b/studio/backend/core/inference/video_families.py @@ -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"(? 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: diff --git a/studio/backend/core/inference/video_ltx2.py b/studio/backend/core/inference/video_ltx2.py index 644e274c6f..a5544ab362 100644 --- a/studio/backend/core/inference/video_ltx2.py +++ b/studio/backend/core/inference/video_ltx2.py @@ -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 " diff --git a/studio/backend/core/training/diffusion_lora_trainer.py b/studio/backend/core/training/diffusion_lora_trainer.py index fb8d911d3a..1cf5f786c7 100644 --- a/studio/backend/core/training/diffusion_lora_trainer.py +++ b/studio/backend/core/training/diffusion_lora_trainer.py @@ -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) diff --git a/studio/backend/core/training/diffusion_train_common.py b/studio/backend/core/training/diffusion_train_common.py index 956bcb46ee..50294d98f8 100644 --- a/studio/backend/core/training/diffusion_train_common.py +++ b/studio/backend/core/training/diffusion_train_common.py @@ -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 - # ``/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): diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py index bd029ad4fc..a9d2a0470c 100644 --- a/studio/backend/core/training/training.py +++ b/studio/backend/core/training/training.py @@ -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 diff --git a/studio/backend/hub/services/download_lifecycle.py b/studio/backend/hub/services/download_lifecycle.py index 9eec24f7a7..7e3c17df6c 100644 --- a/studio/backend/hub/services/download_lifecycle.py +++ b/studio/backend/hub/services/download_lifecycle.py @@ -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)]) diff --git a/studio/backend/hub/services/models/deletion.py b/studio/backend/hub/services/models/deletion.py index 8ef6f0e174..0505af820a 100644 --- a/studio/backend/hub/services/models/deletion.py +++ b/studio/backend/hub/services/models/deletion.py @@ -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): diff --git a/studio/backend/hub/services/models/downloads.py b/studio/backend/hub/services/models/downloads.py index 19460950b9..34ffd08d21 100644 --- a/studio/backend/hub/services/models/downloads.py +++ b/studio/backend/hub/services/models/downloads.py @@ -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 = "@" diff --git a/studio/backend/hub/utils/download_registry.py b/studio/backend/hub/utils/download_registry.py index 10ec3aacbc..e2310a4c42 100644 --- a/studio/backend/hub/utils/download_registry.py +++ b/studio/backend/hub/utils/download_registry.py @@ -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 diff --git a/studio/backend/hub/workers/hf_download.py b/studio/backend/hub/workers/hf_download.py index 66ef21fd4f..c90f32dcd0 100644 --- a/studio/backend/hub/workers/hf_download.py +++ b/studio/backend/hub/workers/hf_download.py @@ -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: diff --git a/studio/backend/main.py b/studio/backend/main.py index e769fcfed4..5cdebc122e 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -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 ) diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 10f8b78c9d..40c5b9df0a 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -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), " diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 4a870e4139..f5907d2410 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -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"]}, diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index 0a50dbc6d1..1e6b25e62a 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -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; diff --git a/studio/backend/routes/settings.py b/studio/backend/routes/settings.py index 3c7f72e69a..6d5f41d110 100644 --- a/studio/backend/routes/settings.py +++ b/studio/backend/routes/settings.py @@ -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, diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index 3e0414ea61..6304619334 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -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 .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 .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 .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 .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: diff --git a/studio/backend/routes/video.py b/studio/backend/routes/video.py index 9678dcae74..d542868cc8 100644 --- a/studio/backend/routes/video.py +++ b/studio/backend/routes/video.py @@ -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", diff --git a/studio/backend/tests/conftest.py b/studio/backend/tests/conftest.py index 9cb6cfbb0a..3226ea2823 100644 --- a/studio/backend/tests/conftest.py +++ b/studio/backend/tests/conftest.py @@ -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") diff --git a/studio/backend/tests/test_cached_gguf_routes.py b/studio/backend/tests/test_cached_gguf_routes.py index c95cd98945..a1d14cd88d 100644 --- a/studio/backend/tests/test_cached_gguf_routes.py +++ b/studio/backend/tests/test_cached_gguf_routes.py @@ -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)], diff --git a/studio/backend/tests/test_diffusion_arch_patches.py b/studio/backend/tests/test_diffusion_arch_patches.py index ceeb047669..58c5dba3f0 100644 --- a/studio/backend/tests/test_diffusion_arch_patches.py +++ b/studio/backend/tests/test_diffusion_arch_patches.py @@ -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() diff --git a/studio/backend/tests/test_diffusion_attention.py b/studio/backend/tests/test_diffusion_attention.py index 1f0fff7ccc..828db2f771 100644 --- a/studio/backend/tests/test_diffusion_attention.py +++ b/studio/backend/tests/test_diffusion_attention.py @@ -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 diff --git a/studio/backend/tests/test_diffusion_auto_policy.py b/studio/backend/tests/test_diffusion_auto_policy.py index d4c0dfe85f..9db6157142 100644 --- a/studio/backend/tests/test_diffusion_auto_policy.py +++ b/studio/backend/tests/test_diffusion_auto_policy.py @@ -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 diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index 57f7783fd2..77068ff7ec 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/studio/backend/tests/test_diffusion_backend.py @@ -26,10 +26,9 @@ from core.inference.diffusion import ( _resolve_diffusion_compute_dtype, ) -# diffusion.py imports the compile/arch patch modules LAZILY (they pull torch at module level, -# and diffusion.py must stay importable on a torchless native install). Import them here at -# collection time -- under the real torch -- so they're cached in sys.modules before the -# fake-torch fixtures swap it out; else the lazy import would build them against the stub torch. +# diffusion.py imports the compile/arch patch modules LAZILY (they pull torch at module level). +# Import them here at collection time, under the real torch, so they are cached in sys.modules +# before the fake-torch fixtures swap it out. import core.inference.diffusion_eager_patches # noqa: E402,F401 import core.inference.diffusion_arch_patches # noqa: E402,F401 from core.inference.diffusion_families import ( @@ -44,12 +43,12 @@ from core.inference.diffusion_families import ( def test_clamp_max_side_bounds_oversized_init(): - # img2img / inpaint derive OUTPUT size from the uploaded image; an oversized upload (up to - # the 4096/side decode cap = 4x the txt2img 2048 ceiling) drives an OOM-scale latent. - # _clamp_max_side bounds the longest side to 2048, preserving aspect ratio. + # img2img / inpaint derive OUTPUT size from the uploaded image, and an oversized upload (up to the + # 4096/side decode cap) drives an OOM-scale latent. _clamp_max_side bounds the longest side to + # 2048, preserving aspect ratio. from PIL import Image - # A 12MP-shaped landscape photo -> longest side clamped to 2048, 4:3 aspect preserved. + # A 12MP landscape photo: longest side clamped to 2048, 4:3 aspect preserved. out = _clamp_max_side(Image.new("RGB", (4096, 3072)), 2048) assert out.size == (2048, 1536) # A portrait upload clamps on its longest (height) side. @@ -72,8 +71,8 @@ def test_detect_family_from_repo_id(): assert klein.cfg_kwarg == "guidance_scale" # Both klein sizes share the one family (base repo resolved per-variant). assert detect_family("unsloth/FLUX.2-klein-9B-GGUF").name == "flux.2-klein" - # FLUX.2-dev is the Mistral-based Flux2Pipeline, a distinct family from klein; its - # gated base repo is reachable with an HF token. It must not collide with klein. + # FLUX.2-dev is the Mistral-based Flux2Pipeline, a distinct family from klein, and must not + # collide with it. dev = detect_family("unsloth/FLUX.2-dev-GGUF") assert dev.name == "flux.2-dev" assert dev.pipeline_class == "Flux2Pipeline" @@ -82,15 +81,15 @@ def test_detect_family_from_repo_id(): # Qwen-Image guides via true_cfg_scale, not guidance_scale. assert detect_family("unsloth/Qwen-Image-2512-GGUF").cfg_kwarg == "true_cfg_scale" assert detect_family("unsloth/Z-Image-GGUF").cfg_kwarg == "guidance_scale" - # Qwen-Image-Edit is a SUPPORTED instruction-editing family (its own edit pipeline); - # the most-specific match wins so it doesn't fall back to the generic qwen-image. + # Qwen-Image-Edit is a SUPPORTED instruction-editing family; the most-specific match wins so it + # doesn't fall back to the generic qwen-image. edit = detect_family("unsloth/Qwen-Image-Edit-2511-GGUF") assert edit.name == "qwen-image-edit" assert edit.pipeline_class == "QwenImageEditPlusPipeline" assert edit.edit is True assert detect_family("unsloth/Qwen-Image-Edit-2509-GGUF").name == "qwen-image-edit" - # FLUX Kontext is a SUPPORTED editing family (FluxKontextPipeline); the "kontext" - # keyword is un-rejected for it, and it must win over the generic "flux.1" match. + # FLUX Kontext is a SUPPORTED editing family: the "kontext" keyword is un-rejected for it, and it + # must win over the generic "flux.1" match. kontext = detect_family("unsloth/FLUX.1-Kontext-dev-GGUF") assert kontext.name == "flux.1-kontext" assert kontext.pipeline_class == "FluxKontextPipeline" @@ -112,9 +111,8 @@ def test_detect_family_from_repo_id(): def test_detect_family_matches_reject_and_alias_by_segment(): - # Reject keywords and short aliases must match whole path/name segments, not raw substrings, - # so an unrelated word that CONTAINS one doesn't misroute a valid base model (regression: - # substring matching broke these). + # Reject keywords and short aliases must match whole path/name segments, not raw substrings, so + # an unrelated word containing one doesn't misroute a valid base model. assert detect_family("/models/edited/z-image-turbo-Q4_K_M.gguf").name == "z-image" assert detect_family("unsloth/Z-Image-Edition-GGUF").name == "z-image" assert detect_family("/models/kontextual/z-image-turbo-Q4_K_M.gguf").name == "z-image" @@ -129,9 +127,8 @@ def test_detect_family_matches_reject_and_alias_by_segment(): def test_detect_family_edit_keyword_scoped_to_basename(): from core.inference.diffusion_families import detect_family_for_pick - # A parent directory named `edit`/`inpaint` must NOT poison a valid pick: only the model id - # / filename basename is scanned for reject keywords. A direct local pick arrives as - # (parent_dir, filename). + # A parent directory named `edit`/`inpaint` must NOT poison a valid pick: only the model id / + # filename basename is scanned for reject keywords. assert detect_family("/models/edit") is None # the dir alone is ambiguous assert detect_family_for_pick("/models/edit", "Z-Image-Turbo-Q4.gguf").name == "z-image" assert detect_family_for_pick("/models/inpaint", "qwen-image-2512-Q4.gguf").name == "qwen-image" @@ -246,9 +243,8 @@ class _FakePipe: def enable_vae_slicing(self) -> None: self.vae_sliced = True - # Explicit signature (not just **kwargs) so generate()'s signature-gated guards for - # negative_prompt / callback_on_step_end take effect -- a **kwargs-only fake would make - # `"negative_prompt" in signature` always False. + # Explicit signature (not just **kwargs) so generate()'s signature-gated guards take effect: a + # **kwargs-only fake would make `"negative_prompt" in signature` always False. def __call__( self, *, @@ -269,8 +265,8 @@ class _FakePipe: "cfg_trunc_ratio": cfg_trunc_ratio, **kwargs, } - # Mirror diffusers batching: a prompt LIST yields one image per prompt, and - # num_images_per_prompt fans each prompt out. + # Mirror diffusers batching: a prompt LIST yields one image per prompt, and num_images_per_prompt + # fans each prompt out. n = kwargs.get("num_images_per_prompt", 1) if isinstance(prompt, list): n *= len(prompt) @@ -407,18 +403,16 @@ def fake_runtime(monkeypatch): diffusers.QwenImageInpaintPipeline = _FakeInpaintPipeline # Instruction-editing pipeline (Qwen-Image-Edit): its own pipeline IS the loaded one. diffusers.QwenImageEditPlusPipeline = _FakePipeline - # Ideogram 4, so its guidance_scale/guidance_schedule pairing is exercisable. It loads only - # as a full pipeline (two DiTs), assembled per-component by load_ideogram4_pipeline -- stub - # that to a fake pipe so the guidance path is reachable without real weights. + # Ideogram 4, so its guidance_scale/guidance_schedule pairing is exercisable. It loads only as a + # full pipeline (two DiTs), so stub the assembly to a fake pipe. diffusers.Ideogram4Pipeline = _FakePipeline diffusers.Ideogram4Transformer2DModel = _FakeTransformer - # Lumina 2, so the cfg_trunc_ratio special case is exercisable (the fake pipe's - # signature carries the kwarg, mirroring the real Lumina2Pipeline). + # Lumina 2, so the cfg_trunc_ratio special case is exercisable (the fake pipe's signature carries + # the kwarg, mirroring the real Lumina2Pipeline). diffusers.Lumina2Pipeline = _FakePipeline diffusers.Lumina2Transformer2DModel = _FakeTransformer - # SDXL: a U-Net family. Its single-file checkpoint is the whole pipeline, so the pipeline - # class carries from_single_file; UNet2DConditionModel is the denoiser class (fetched but - # unused on the pipeline/single-file-pipeline paths). + # SDXL: a U-Net family whose single-file checkpoint is the whole pipeline, so the pipeline class + # carries from_single_file and UNet2DConditionModel is the denoiser class. diffusers.StableDiffusionXLPipeline = _FakePipeline diffusers.UNet2DConditionModel = _FakeTransformer diffusers.StableDiffusionXLImg2ImgPipeline = _FakeImg2ImgPipeline @@ -431,8 +425,8 @@ def fake_runtime(monkeypatch): monkeypatch.setitem(sys.modules, "torch", torch) monkeypatch.setitem(sys.modules, "diffusers", diffusers) - # The backend imports clear_gpu_cache by reference; no-op it so unload doesn't - # run real hardware detection against the stubbed torch. + # The backend imports clear_gpu_cache by reference; no-op it so unload doesn't run real hardware + # detection against the stubbed torch. monkeypatch.setattr("core.inference.diffusion.clear_gpu_cache", lambda: None) _FakePipeline.last = {} _FakePipeline.last_single_file = {} @@ -475,8 +469,8 @@ def test_load_generate_unload_gguf(fake_runtime, tmp_path): assert gen["seed"] == 4242 # random seed reported back assert gen["repo_id"] == str(tmp_path) # echoed so the route can record the model assert len(gen["images"]) == 1 # PIL images handed to the route for persistence - # z-image guides via guidance_scale (not true_cfg_scale); the signature-gated - # negative_prompt and per-step callback both reach the pipeline call. + # z-image guides via guidance_scale (not true_cfg_scale); the signature-gated negative_prompt and + # per-step callback both reach the pipeline call. call = backend._state.pipe.last_kwargs assert call["guidance_scale"] == 3.0 and call["true_cfg_scale"] is None assert call["negative_prompt"] == "blurry" @@ -485,8 +479,8 @@ def test_load_generate_unload_gguf(fake_runtime, tmp_path): gen2 = backend.generate(prompt = "again", seed = 99) assert gen2["seed"] == 99 - # batch_size produces that many images in one call, seeded base..base+2 per image - # (matching the native engine) so each batch member replays alone from its recipe. + # batch_size produces that many images in one call, seeded base..base+2 per image (matching the + # native engine) so each batch member replays alone from its recipe. batch = backend.generate(prompt = "batch", seed = 7, batch_size = 3) assert len(batch["images"]) == 3 and batch["seed"] == 7 assert batch["seeds"] == [7, 8, 9] @@ -515,7 +509,6 @@ def test_generate_progress_active_during_setup(fake_runtime, tmp_path, monkeypat monkeypatch.setattr(DiffusionBackend, "_apply_loras", fake_apply) - # Idle before the run. assert backend.generate_progress()["active"] is False gen = backend.generate(prompt = "a sloth", steps = 4) @@ -526,13 +519,12 @@ def test_generate_progress_active_during_setup(fake_runtime, tmp_path, monkeypat assert seen["progress"]["total_steps"] == 4 assert seen["progress"]["step"] == 0 - # And it is cleared once the generation returns. assert backend.generate_progress()["active"] is False def test_generate_progress_cleared_on_setup_error(fake_runtime, tmp_path, monkeypatch): - # A setup-time failure skips the inner finally that nulls _gen, so the outer finally must - # clear the published progress; otherwise a crashed generation leaves the UI stuck "active". + # A setup-time failure skips the inner finally that nulls _gen, so the outer finally must clear + # the published progress; otherwise a crashed generation leaves the UI stuck "active". (tmp_path / "model.gguf").write_bytes(b"weights") backend = DiffusionBackend() backend.load_pipeline( @@ -556,8 +548,8 @@ def test_generate_progress_cleared_on_setup_error(fake_runtime, tmp_path, monkey def test_generate_progress_active_through_compile_cache_save(fake_runtime, tmp_path, monkeypatch): # Post-denoise work (the compile-cache save) still runs before the route persists the image, so - # progress must stay active through it; clearing early would let a reload's mount probe read idle - # and refresh the gallery before the result exists. + # progress must stay active through it, else a reload's mount probe reads idle and refreshes the + # gallery before the result exists. from core.inference import diffusion as dmod (tmp_path / "model.gguf").write_bytes(b"weights") @@ -581,17 +573,15 @@ def test_generate_progress_active_through_compile_cache_save(fake_runtime, tmp_p gen = backend.generate(prompt = "a sloth", steps = 4) assert len(gen["images"]) == 1 - # Still active while the compile-cache save ran (after the denoise, before generate() returns). + # Still active while the compile-cache save ran. assert seen["progress"]["active"] is True assert seen["progress"]["total_steps"] == 4 - # And cleared once the generation returns. assert backend.generate_progress()["active"] is False def test_dense_speed_auto_defers_compile_to_third_generation(fake_runtime, tmp_path, monkeypatch): - # Dense models with speed unset stay bit-identical eager for the first two generations; the - # 3rd engages the `default` profile mid-session (repeated use amortises the one-time compile), - # upgrading attention alongside it. + # Dense models with speed unset stay bit-identical eager for the first two generations; the 3rd + # engages the `default` profile mid-session, upgrading attention alongside it. from core.inference import diffusion as dmod monkeypatch.setattr(dmod, "compile_eligible", lambda *a, **k: True) @@ -647,10 +637,9 @@ def test_dense_speed_auto_defers_compile_to_third_generation(fake_runtime, tmp_p def test_deferred_speed_skips_when_lora_requested(fake_runtime, tmp_path, monkeypatch): - # A compiled transformer rejects LoRA (supports_lora False once compiled), and _apply_loras - # raises before its unchanged-selection no-op, so engaging the deferred compile on a LoRA - # generation would permanently break every LoRA generation on this load. The deferral must - # skip while a LoRA is requested and engage only on a later LoRA-free generation. + # A compiled transformer rejects LoRA, and _apply_loras raises before its unchanged-selection + # no-op, so engaging the deferred compile on a LoRA generation would permanently break every LoRA + # generation on this load. The deferral must skip while a LoRA is requested. from core.inference import diffusion as dmod monkeypatch.setattr(dmod, "compile_eligible", lambda *a, **k: True) @@ -674,19 +663,18 @@ def test_deferred_speed_skips_when_lora_requested(fake_runtime, tmp_path, monkey ) backend.generate(prompt = "one") backend.generate(prompt = "two") - # 3rd generation requests a LoRA: the deferral must be skipped (pipe stays eager, LoRA-capable). + # 3rd generation requests a LoRA: the deferral must be skipped (pipe stays LoRA-capable). backend.generate(prompt = "three", loras = [("adapter", 1.0)]) assert engaged == [] - # 4th generation without a LoRA: the deferral now engages (the guard is LoRA-specific, not off). + # 4th generation without a LoRA: the deferral now engages (the guard is LoRA-specific). backend.generate(prompt = "four") assert len(engaged) == 1 def test_deferred_speed_skips_while_adapter_attached(fake_runtime, tmp_path, monkeypatch): - # Even a NO-LoRA generation must defer the compile while an adapter from a PRIOR generation is - # still attached: _apply_loras runs AFTER the engage, so compiling here would bake the resident - # adapter into the graph and the later unload (swallowed on a compiled pipe) would leave it - # active forever -- silent wrong output. Defer until _apply_loras clears it. + # Even a NO-LoRA generation must defer while an adapter from a PRIOR generation is still + # attached: _apply_loras runs AFTER the engage, so compiling here would bake the resident adapter + # into the graph and the later (swallowed) unload would leave it active forever. from core.inference import diffusion as dmod monkeypatch.setattr(dmod, "compile_eligible", lambda *a, **k: True) @@ -698,7 +686,7 @@ def test_deferred_speed_skips_while_adapter_attached(fake_runtime, tmp_path, mon monkeypatch.setattr(DiffusionBackend, "_engage_deferred_speed", fake_engage) - # Track the attached set on the pipe, mirroring the real _apply_loras marker (_unsloth_loras). + # Track the attached set on the pipe, mirroring the real _apply_loras marker. def fake_apply(self, state, loras, cancel): specs = [(i, w) for (i, w) in (loras or []) if w != 0] state.pipe._unsloth_loras = tuple(specs) @@ -716,19 +704,18 @@ def test_deferred_speed_skips_while_adapter_attached(fake_runtime, tmp_path, mon # Gens 1-2 attach an adapter, so it is still resident going into gen 3. backend.generate(prompt = "one", loras = [("adapter", 1.0)]) backend.generate(prompt = "two", loras = [("adapter", 1.0)]) - # Gen 3 requests NO LoRA but the adapter is still attached -> defer (no compile-with-adapter). + # Gen 3 requests NO LoRA but the adapter is still attached, so defer. backend.generate(prompt = "three") assert engaged == [] - # Gen 3's _apply_loras([]) cleared the adapter; gen 4 is genuinely LoRA-free -> engage. + # Gen 3's _apply_loras([]) cleared the adapter; gen 4 is genuinely LoRA-free, so engage. backend.generate(prompt = "four") assert len(engaged) == 1 def test_deferred_speed_preserves_explicit_attention(fake_runtime, tmp_path, monkeypatch): - # A dense model loaded with Speed on Auto but Attention explicitly pinned (e.g. "native" to - # avoid cuDNN) must KEEP that choice when the 3rd generation engages the deferred `default` - # profile. The auto cuDNN upgrade applies only when attention was left on auto, never when the - # caller pinned a backend. + # A dense model loaded with Speed on Auto but Attention explicitly pinned must KEEP that choice + # when the 3rd generation engages the deferred profile: the auto cuDNN upgrade applies only when + # attention was left on auto. from core.inference import diffusion as dmod monkeypatch.setattr(dmod, "compile_eligible", lambda *a, **k: True) @@ -739,9 +726,8 @@ def test_deferred_speed_preserves_explicit_attention(fake_runtime, tmp_path, mon ) monkeypatch.setattr(dmod, "apply_attention_backend", lambda pipe, backend, logger = None: backend) - # A select mock that -- unlike a bare "auto -> cuDNN" stub -- HONORS an explicit request: - # "native" stays on the default (None) even under a speed profile, and only a left-unset - # ("auto"/None) request upgrades to cuDNN when speed is active. + # A select mock that HONORS an explicit request: "native" stays on the default even under a speed + # profile, and only an unset request upgrades to cuDNN. def fake_select( target, requested, @@ -776,8 +762,8 @@ def test_deferred_speed_preserves_explicit_attention(fake_runtime, tmp_path, mon assert status["resolved"]["attention_backend"]["value"] == "native" assert status["resolved"]["attention_backend"]["source"] == "explicit" - # Control: with attention left on auto, the same 3rd-generation deferral DOES upgrade - # to cuDNN -- so the assertion above is not vacuously passing. + # Control: with attention left on auto the same deferral DOES upgrade to cuDNN, so the assertion + # above is not vacuously passing. backend.unload() backend.load_pipeline( str(tmp_path), @@ -811,8 +797,8 @@ def test_generate_img2img_uses_from_pipe(fake_runtime, tmp_path): backend.load_pipeline( str(tmp_path), gguf_filename = "model.gguf", base_repo = "base/repo", family_override = "z-image" ) - # The loaded family advertises the image-conditioned workflows for UI gating - # (upscale rides the img2img pipeline, so it appears whenever img2img does). + # The loaded family advertises the image-conditioned workflows for UI gating (upscale rides the + # img2img pipeline, so it appears whenever img2img does). assert backend.status()["workflows"] == ["txt2img", "img2img", "upscale", "inpaint", "outpaint"] loaded_pipe = backend._state.pipe @@ -827,8 +813,8 @@ def test_generate_img2img_uses_from_pipe(fake_runtime, tmp_path): assert len(out["images"]) == 1 # from_pipe was handed the loaded text-to-image pipe (component reuse, no reload). assert _FakeImg2ImgPipeline.built_from is loaded_pipe - # ...and with torch_dtype=None so from_pipe SKIPS its default float32 recast, which - # both upcasts the reused bf16 modules and crashes on torchao-quantized weights. + # ...and with torch_dtype=None so from_pipe SKIPS its default float32 recast, which upcasts the + # reused bf16 modules and crashes on torchao-quantized weights. assert _FakeImg2ImgPipeline.from_pipe_kwargs.get("torch_dtype", "MISSING") is None call = _FakeImg2ImgPipe.last_kwargs assert call["image"] is not None # decoded source image passed through @@ -845,8 +831,8 @@ def test_generate_img2img_unsupported_family_raises(fake_runtime, tmp_path, monk an init_image with a clear error rather than failing deep in the pipeline.""" from core.inference.diffusion_families import DiffusionFamily - # A synthetic txt2img-only family: no img2img/inpaint pipeline, not edit, not reference. - # (Every shipped family now supports some image workflow, so build one for this case.) + # A synthetic txt2img-only family: no img2img/inpaint pipeline, not edit, not reference (every + # shipped family now supports some image workflow). plain = DiffusionFamily( name = "plain-test", pipeline_class = "ZImagePipeline", @@ -1052,8 +1038,8 @@ def test_generate_reference_uses_loaded_pipe_at_slider_size(fake_runtime, tmp_pa base_repo = "base/repo", family_override = "flux.2-klein", ) - # FLUX.2-klein: txt2img + reference (own pipe) + inpaint (dedicated pipe). No img2img class, - # so no img2img/upscale. + # FLUX.2-klein: txt2img + reference (own pipe) + inpaint (dedicated pipe). No img2img class, so + # no img2img/upscale. assert backend.status()["workflows"] == ["txt2img", "reference", "inpaint"] loaded_pipe = backend._state.pipe @@ -1076,8 +1062,8 @@ def test_generate_reference_uses_loaded_pipe_at_slider_size(fake_runtime, tmp_pa # Guidance flows via guidance_scale (FLUX.2 default behaviour). assert call["guidance_scale"] == 4.0 - # Multi-reference: extra reference_images are combined with init_image into a LIST so the - # model can blend several references (subject + style). + # Multi-reference: extra reference_images are combined with init_image into a LIST so the model + # can blend several references. backend.generate( prompt = "combine these", steps = 6, @@ -1090,8 +1076,8 @@ def test_generate_reference_uses_loaded_pipe_at_slider_size(fake_runtime, tmp_pa img_arg = loaded_pipe.last_kwargs["image"] assert isinstance(img_arg, list) and len(img_arg) == 3 # primary + 2 extras - # Branch ordering: an init image + MASK on a reference family must route to inpaint (the - # dedicated pipeline), NOT be swallowed by the reference branch (which ignores the mask). + # Branch ordering: an init image + MASK on a reference family must route to inpaint, NOT be + # swallowed by the reference branch (which ignores the mask). backend.generate( prompt = "repaint here", steps = 6, @@ -1271,7 +1257,7 @@ def test_edit_family_uses_own_pipeline_and_requires_image(fake_runtime, tmp_path base_repo = "Qwen/Qwen-Image-Edit-2511", family_override = "qwen-image-edit", ) - # Edit families advertise only the edit workflow (no txt2img / img2img / inpaint). + # Edit families advertise only the edit workflow. assert backend.status()["workflows"] == ["edit"] loaded_pipe = backend._state.pipe @@ -1283,7 +1269,7 @@ def test_edit_family_uses_own_pipeline_and_requires_image(fake_runtime, tmp_path init_image = _tiny_png_b64(), ) assert len(out["images"]) == 1 - # The loaded pipe handled it directly -- no from_pipe img2img/inpaint was built. + # The loaded pipe handled it directly: no from_pipe img2img/inpaint was built. assert backend._state.pipe is loaded_pipe assert _FakeImg2ImgPipeline.built_from is None and _FakeInpaintPipeline.built_from is None assert loaded_pipe.last_kwargs.get("image") is not None @@ -1388,8 +1374,8 @@ def test_load_sdxl_rejects_untrusted_repo(fake_runtime): def test_validate_gates_untrusted_base_repo(fake_runtime, tmp_path): # A companion base_repo also loads via from_pretrained, so a trusted GGUF model_path must not - # smuggle in an arbitrary remote base: base_repo clears the same trust bar as a non-GGUF repo - # id (mirrors the video loader), and the check runs before any GPU handoff. + # smuggle in an arbitrary remote base: base_repo clears the same trust bar, before any GPU + # handoff. backend = DiffusionBackend() with pytest.raises(ValueError, match = "base_repo"): backend.validate_load_request( @@ -1398,10 +1384,9 @@ def test_validate_gates_untrusted_base_repo(fake_runtime, tmp_path): model_kind = "gguf", base_repo = "evil/companions", ) - # A local base_repo dir that is NOT a diffusers pipeline (no model_index.json) is rejected - # HERE, before the GPU handoff: it passes the any-existing-path trust check but the base loads - # via from_pretrained (needs model_index.json), so it would else evict the resident model and - # only then fail in the background load. + # A local base_repo dir that is NOT a diffusers pipeline is rejected HERE, before the GPU handoff: + # it passes the any-existing-path trust check but from_pretrained needs model_index.json, so it + # would otherwise evict the resident model and only then fail. bad_base = tmp_path / "bare-base" bad_base.mkdir() with pytest.raises(ValueError, match = "model_index.json"): @@ -1411,7 +1396,7 @@ def test_validate_gates_untrusted_base_repo(fake_runtime, tmp_path): model_kind = "gguf", base_repo = str(bad_base), ) - # A local base_repo that IS a real pipeline dir (model_index.json) passes the gate. + # A local base_repo that IS a real pipeline dir passes the gate. (tmp_path / "model_index.json").write_text("{}") fam = backend.validate_load_request( "unsloth/Qwen-Image-2512-GGUF", @@ -1423,9 +1408,8 @@ def test_validate_gates_untrusted_base_repo(fake_runtime, tmp_path): def test_resolve_local_single_file(tmp_path): - # A bare single-file safetensors directory (no model_index.json) resolves to that checkpoint's - # basename, so the images load route can reinterpret an On-Device "pipeline" pick as a - # single_file load instead of 400ing on the missing model_index.json. + # A bare single-file safetensors directory resolves to that checkpoint's basename, so the images + # load route can reinterpret an On-Device "pipeline" pick as a single_file load. from core.inference.diffusion import resolve_local_single_file d = tmp_path / "solo" @@ -1433,11 +1417,11 @@ def test_resolve_local_single_file(tmp_path): (d / "model.safetensors").write_bytes(b"w") assert resolve_local_single_file(str(d)) == "model.safetensors" - # A real diffusers pipeline dir (has model_index.json) loads as a pipeline unchanged -> None. + # A real diffusers pipeline dir loads as a pipeline unchanged. (d / "model_index.json").write_text("{}") assert resolve_local_single_file(str(d)) is None - # Ambiguous (two checkpoints, e.g. a sharded pipeline) or empty dirs -> None (unchanged load). + # Ambiguous (two checkpoints) or empty dirs leave the load unchanged. d2 = tmp_path / "shards" d2.mkdir() (d2 / "a.safetensors").write_bytes(b"w") @@ -1447,16 +1431,15 @@ def test_resolve_local_single_file(tmp_path): # A remote repo id (not a local dir) -> None. assert resolve_local_single_file("unsloth/Qwen-Image-2512-GGUF") is None - # A PEFT LoRA adapter folder (adapter_config.json + adapter_model.safetensors), even with a - # family-token name, is NOT a base checkpoint: from_single_file would fail on the adapter - # weights AFTER the route evicted the resident model, so it must not be reinterpreted as a - # single_file pick -> None (the pipeline pick then 400s in validation, before the handoff). + # A PEFT adapter folder, even with a family-token name, is NOT a base checkpoint: from_single_file + # would fail on the adapter weights AFTER eviction, so it must not be reinterpreted as single_file + # (the pipeline pick then 400s in validation, before the handoff). adapter = tmp_path / "flux-style-lora" adapter.mkdir() (adapter / "adapter_config.json").write_text("{}") (adapter / "adapter_model.safetensors").write_bytes(b"w") assert resolve_local_single_file(str(adapter)) is None - # A bare adapter_model.safetensors (no config) is likewise not treated as the sole checkpoint. + # A bare adapter_model.safetensors is likewise not treated as the sole checkpoint. adapter2 = tmp_path / "z-image-lora" adapter2.mkdir() (adapter2 / "adapter_model.safetensors").write_bytes(b"w") @@ -1464,14 +1447,13 @@ def test_resolve_local_single_file(tmp_path): def test_resolve_base_repo_drops_untrusted_card_tag(monkeypatch): - # With no base_repo, the base is resolved from the GGUF repo's base_model card tag -- - # attacker-controlled metadata on any remote repo -- then loaded via from_pretrained. An - # untrusted tag must be dropped for the curated family default, so an attacker GGUF repo can't - # point the base at an arbitrary repo to be deserialized. + # With no base_repo the base comes from the GGUF repo's base_model card tag -- attacker-controlled + # metadata -- then loads via from_pretrained. An untrusted tag must be dropped for the curated + # family default. import core.inference.diffusion as dmod fam = detect_family("unsloth/FLUX.1-dev-GGUF") - # A malicious card tag is ignored -> the family default base is used instead. + # A malicious card tag is ignored, so the family default base is used. monkeypatch.setattr(dmod, "_hf_base_model", lambda repo_id, hf_token: "attacker/evil-pipeline") assert _resolve_base_repo("attacker/flux.1-evil-GGUF", None, fam, None) == fam.base_repo # A trusted (allowlisted) card tag is still honoured, so variant resolution is not regressed. @@ -1482,8 +1464,8 @@ def test_resolve_base_repo_drops_untrusted_card_tag(monkeypatch): _resolve_base_repo("unsloth/FLUX.1-dev-GGUF", None, fam, None) == "black-forest-labs/FLUX.1-dev" ) - # An explicit trusted base_repo wins over the card tag; an explicit untrusted one is caught - # earlier at validate_load_request (covered by test_validate_gates_untrusted_base_repo). + # An explicit trusted base_repo wins over the card tag; an untrusted one is caught earlier at + # validate_load_request. assert ( _resolve_base_repo("unsloth/FLUX.1-dev-GGUF", "unsloth/custom-base", fam, None) == "unsloth/custom-base" @@ -1491,8 +1473,8 @@ def test_resolve_base_repo_drops_untrusted_card_tag(monkeypatch): def test_detect_family_rejects_layered(): - # Qwen-Image-Layered needs a dedicated pipeline (additional_t_cond); it must be - # rejected so it fails fast at load instead of crashing at the first denoise step. + # Qwen-Image-Layered needs a dedicated pipeline (additional_t_cond), so it must be rejected at + # load instead of crashing at the first denoise step. assert detect_family("unsloth/Qwen-Image-Layered-GGUF") is None assert detect_family("unsloth/qwen_image_layered") is None @@ -1561,9 +1543,8 @@ def test_generate_without_load_raises(fake_runtime): def test_failed_load_restores_backend_flags(fake_runtime, tmp_path, monkeypatch): # A failure AFTER apply_speed_optims (here an OOM in apply_memory_plan) must go through the - # load's try/finally and restore the process-global TF32 / cudnn flags, so a later `off` load - # is still bit-identical, and must not commit partial state. Regression: a refactor dropped - # this guard, leaking the flags on a failed load. + # load's try/finally and restore the process-global TF32 / cudnn flags, so a later `off` load is + # still bit-identical, and must not commit partial state. (tmp_path / "model.gguf").write_bytes(b"x") backend = DiffusionBackend() @@ -1617,8 +1598,8 @@ def test_resolve_base_repo_prefers_caller_then_hf_tag_then_fallback(monkeypatch) def test_load_without_gguf_raises(): backend = DiffusionBackend() - # No gguf_filename -> a full-pipeline load, gated to unsloth/*; a non-unsloth repo - # is rejected before any GPU/network work. + # No gguf_filename means a full-pipeline load, gated to unsloth/*; a non-unsloth repo is rejected + # before any GPU/network work. with pytest.raises(ValueError, match = "unsloth"): backend.load_pipeline("some-org/Z-Image-bnb-4bit") @@ -1667,8 +1648,8 @@ def test_base_file_downloaded_excludes_undownloaded(): assert _base_file_downloaded("model_index.json") assert _base_file_downloaded("text_encoder/model-00001-of-00003.safetensors") assert _base_file_downloaded("vae/diffusion_pytorch_model.safetensors") - # Excluded: the GGUF supplies the transformer; docs/assets and top-level files - # are never downloaded, so counting them would peg the bar short of 100%. + # Excluded: the GGUF supplies the transformer, and docs/assets are never downloaded, so counting + # them would peg the bar short of 100%. assert not _base_file_downloaded( "transformer/diffusion_pytorch_model-00001-of-00003.safetensors" ) @@ -1678,8 +1659,8 @@ def test_base_file_downloaded_excludes_undownloaded(): def test_load_progress_fraction_clamped(monkeypatch): - # The cache scan can exceed the estimate (e.g. a second cached quant); the - # reported fraction must still clamp to 1.0 rather than overshoot. + # The cache scan can exceed the estimate (e.g. a second cached quant); the reported fraction must + # still clamp to 1.0. backend = DiffusionBackend() backend._loading = _LoadingState(repo_id = "r", base_repo = "b", expected_bytes = 1000) monkeypatch.setattr(DiffusionBackend, "_cache_bytes", staticmethod(lambda repo: 900)) @@ -1695,7 +1676,7 @@ def test_estimate_eta(): # No rate yet until a step has elapsed since the first. assert _estimate_eta(8, 1, first_step_at = 100.0, now = 100.0) is None assert _estimate_eta(8, 0, first_step_at = 0.0, now = 100.0) is None - # 3 steps in 3s since the first ⇒ 1s/step ⇒ 4 steps left ⇒ ~4s. + # 3 steps in 3s since the first: 1s/step, 4 steps left, so ~4s. assert _estimate_eta(8, 4, first_step_at = 100.0, now = 103.0) == 4.0 # Last step ⇒ 0 remaining. assert _estimate_eta(8, 8, first_step_at = 100.0, now = 107.0) == 0.0 @@ -1717,16 +1698,16 @@ def test_generate_qwen_uses_true_cfg_scale(fake_runtime, tmp_path): def _load_ideogram(backend, tmp_path): - # Ideogram 4 loads only as a full pipeline (its two DiTs are assembled per-component - # by the stubbed load_ideogram4_pipeline); a local pipeline dir is enough here. + # Ideogram 4 loads only as a full pipeline (its two DiTs are assembled per-component by the + # stubbed loader), so a local pipeline dir is enough here. (tmp_path / "model_index.json").write_text("{}") backend.load_pipeline(str(tmp_path), family_override = "ideogram-4") def test_ideogram_rejects_single_file_and_gguf_kinds(fake_runtime, tmp_path): - # Ideogram 4 needs two DiTs assembled per-component, so there's no transformer-only - # single-file or GGUF load: the explicit kinds must be rejected up front (before a load evicts - # a working model), not assembled into a pipeline missing its second DiT. + # Ideogram 4 needs two DiTs assembled per-component, so there is no transformer-only single-file + # or GGUF load: the explicit kinds must be rejected up front, not assembled into a pipeline + # missing its second DiT. backend = DiffusionBackend() (tmp_path / "model.gguf").write_bytes(b"x") with pytest.raises(ValueError, match = "full diffusers pipeline"): @@ -1744,9 +1725,9 @@ def test_ideogram_rejects_single_file_and_gguf_kinds(fake_runtime, tmp_path): def test_generate_ideogram_defaults_keep_recommended_schedule(fake_runtime, tmp_path): - # Ideogram 4's pipeline defaults to its recommended tapered guidance_schedule (45x7.0 + 3x3.0, - # valid only at 48 steps) and REJECTS guidance_scale while the schedule is set. At the family's - # advertised defaults the backend must drop the constant so the recommended taper engages. + # Ideogram 4's pipeline defaults to its tapered guidance_schedule (valid only at 48 steps) and + # REJECTS guidance_scale while the schedule is set, so at the advertised defaults the backend must + # drop the constant. backend = DiffusionBackend() _load_ideogram(backend, tmp_path) backend.generate(prompt = "a sloth", steps = 48, guidance = 7.0) @@ -1757,8 +1738,7 @@ def test_generate_ideogram_defaults_keep_recommended_schedule(fake_runtime, tmp_ def test_generate_ideogram_custom_guidance_nulls_schedule(fake_runtime, tmp_path): # Any non-default request must broadcast the constant legally: guidance_scale set AND - # guidance_schedule explicitly nulled (the pipeline raises when both are set, and its default - # schedule is non-None). + # guidance_schedule explicitly nulled (the pipeline raises when both are set). backend = DiffusionBackend() _load_ideogram(backend, tmp_path) backend.generate(prompt = "a sloth", steps = 20, guidance = 5.0) @@ -1768,16 +1748,14 @@ def test_generate_ideogram_custom_guidance_nulls_schedule(fake_runtime, tmp_path def _load_lumina(backend, tmp_path): - # Lumina 2 loads through the GENERIC pipeline path (standard diffusers layout); - # a local pipeline dir is enough here. + # Lumina 2 loads through the GENERIC pipeline path, so a local pipeline dir is enough here. (tmp_path / "model_index.json").write_text("{}") backend.load_pipeline(str(tmp_path), family_override = "lumina-2") def test_generate_lumina2_passes_cfg_trunc_ratio(fake_runtime, tmp_path): - # The card recipe truncates the CFG double-forward to the first quarter of the - # trajectory; the pipeline default (1.0) applies it everywhere. The backend passes - # the constant card value on every lumina-2 generate. + # The card recipe truncates the CFG double-forward to the first quarter of the trajectory, while + # the pipeline default (1.0) applies it everywhere, so the backend passes the card value. backend = DiffusionBackend() _load_lumina(backend, tmp_path) backend.generate(prompt = "a sloth", steps = 50, guidance = 4.0) @@ -1787,8 +1765,8 @@ def test_generate_lumina2_passes_cfg_trunc_ratio(fake_runtime, tmp_path): def test_generate_other_family_never_passes_cfg_trunc_ratio(fake_runtime, tmp_path): - # The kwarg is family-gated, not just signature-gated: another family whose pipeline - # happens to accept cfg_trunc_ratio must not inherit Lumina's recipe constant. + # The kwarg is family-gated, not just signature-gated: another family whose pipeline accepts + # cfg_trunc_ratio must not inherit Lumina's recipe constant. backend = DiffusionBackend() (tmp_path / "model.gguf").write_bytes(b"weights") backend.load_pipeline( @@ -1804,8 +1782,7 @@ def test_generate_other_family_never_passes_cfg_trunc_ratio(fake_runtime, tmp_pa def test_begin_load_rejects_concurrent(monkeypatch): backend = DiffusionBackend() - # The worker resolves the base + downloads, both over the network; stub them - # so the test is offline. + # The worker resolves the base + downloads, both over the network; stub them so this is offline. monkeypatch.setattr("core.inference.diffusion._hf_base_model", lambda *a, **k: None) monkeypatch.setattr(DiffusionBackend, "_prefetch_files", lambda self, *a, **k: None) monkeypatch.setattr( @@ -1840,10 +1817,9 @@ def test_unload_cancels_in_flight_load(fake_runtime): def test_superseded_load_does_not_cancel_live_generation(fake_runtime): - # A superseded background load (token bumped by a newer load/unload) that finally reaches - # load_pipeline must bail WITHOUT signalling the current model's in-flight generation: the - # token check must run before the cancel is set, or a stale worker aborts an unrelated, - # still-live denoise. + # A superseded background load that finally reaches load_pipeline must bail WITHOUT signalling the + # current model's in-flight generation: the token check must run before the cancel is set, or a + # stale worker aborts an unrelated, still-live denoise. import threading as _threading backend = DiffusionBackend() @@ -1882,12 +1858,12 @@ def test_unload_sets_cancel_event(fake_runtime): def test_prefetch_aborts_when_cancelled(tmp_path): - # A prefetch interrupted by unload (cancel event set) raises rather than - # downloading the whole base, so the load can be preempted mid-download. + # A prefetch interrupted by unload raises rather than downloading the whole base, so the load can + # be preempted mid-download. backend = DiffusionBackend() backend._cancel_event.set() - # Local gguf path so the transformer download is skipped; the base loop hits - # the cancel check on its first file (no network). + # Local gguf path so the transformer download is skipped; the base loop hits the cancel check on + # its first file (no network). (tmp_path / "model.gguf").write_bytes(b"x") with pytest.raises(RuntimeError, match = "Cancelled"): backend._prefetch_files( @@ -1941,7 +1917,7 @@ def test_resolve_compute_dtype_promotes_fp16_for_zimage(fake_runtime): torch = sys.modules["torch"] z = detect_family("unsloth/Z-Image-GGUF") q = detect_family("unsloth/Qwen-Image-GGUF") - # Z-Image: fp16 -> fp32; bf16 / fp32 pass through unchanged. + # Z-Image: fp16 promotes to fp32; bf16 / fp32 pass through unchanged. assert _resolve_diffusion_compute_dtype(z, torch.float16) is torch.float32 assert _resolve_diffusion_compute_dtype(z, torch.bfloat16) is torch.bfloat16 assert _resolve_diffusion_compute_dtype(z, torch.float32) is torch.float32 @@ -1952,8 +1928,8 @@ def test_resolve_compute_dtype_promotes_fp16_for_zimage(fake_runtime): def test_load_promotes_fp16_to_fp32_for_zimage_only(fake_runtime, monkeypatch, tmp_path): torch = sys.modules["torch"] - # Pre-Ampere CUDA -> the resolver picks fp16; the guard must promote Z-Image - # (and only Z-Image) to fp32 so it doesn't render a black image. + # Pre-Ampere CUDA resolves to fp16, so the guard must promote Z-Image (and only Z-Image) to fp32 + # or it renders a black image. monkeypatch.setattr(torch.cuda, "is_available", lambda: True, raising = False) monkeypatch.setattr(torch.cuda, "get_device_capability", lambda: (7, 5), raising = False) (tmp_path / "m.gguf").write_bytes(b"x") @@ -1972,8 +1948,8 @@ def test_load_promotes_fp16_to_fp32_for_zimage_only(fake_runtime, monkeypatch, t def test_bad_mode_strings_fail_before_eviction(fake_runtime): - # Every mode normalizer that can raise runs BEFORE the load evicts the previous - # pipeline, so a bad request never costs the user their working model. + # Every mode normalizer that can raise runs BEFORE the load evicts the previous pipeline, so a bad + # request never costs the user their working model. backend = DiffusionBackend() fam = detect_family("unsloth/Z-Image-GGUF") backend._state = _LoadState( @@ -2044,9 +2020,8 @@ def test_generate_lock_split_keeps_status_and_unload_responsive(fake_runtime): assert cancel_ref is not None # unload() signals THIS generation's cancel event, then waits for the denoise to exit before - # returning: callers treat its return as "VRAM is free" (the GPU arbiter hands the GPU to chat - # on it). Release the pipe once the cancel lands, standing in for a real pipeline's step - # callback. + # returning: callers treat its return as "VRAM is free". Release the pipe once the cancel lands, + # standing in for a real pipeline's step callback. releaser = threading.Thread(target = lambda: (cancel_ref.wait(5), release.set())) releaser.start() backend.unload() @@ -2055,8 +2030,8 @@ def test_generate_lock_split_keeps_status_and_unload_responsive(fake_runtime): assert backend.status()["loaded"] is False t.join(5) - # The cancelled generation raised rather than returning a now-evicted image, and - # it had already exited (deregistering its cancel) before unload() returned. + # The cancelled generation raised rather than returning a now-evicted image, and had already + # exited (deregistering its cancel) before unload() returned. assert "exc" in out and "cancelled" in str(out["exc"]).lower() assert backend._active_generate_cancel is None @@ -2119,8 +2094,8 @@ def test_callback_cancellation_interrupts_denoise(fake_runtime): backend._active_generate_cancel.set() resume.set() t.join(5) - # The next step's callback saw the cancel, flipped pipe._interrupt, and the loop - # broke early, so the generation raised instead of returning a partial image. + # The next step's callback saw the cancel, flipped pipe._interrupt and broke the loop, so the + # generation raised instead of returning a partial image. assert pipe._interrupt is True assert pipe.steps_run < 8 assert "exc" in out and "cancelled" in str(out["exc"]).lower() @@ -2146,17 +2121,16 @@ def test_validate_load_request(tmp_path): backend.validate_load_request("some-org/Z-Image", gguf_filename = "model.safetensors") with pytest.raises(ValueError, match = "family"): backend.validate_load_request("meta/Llama-3", gguf_filename = "q.gguf") - # A family-looking repo paired with a non-GGUF single-file name is rejected here, BEFORE the - # route evicts chat and hands over the GPU (else the background load would be the first to - # notice README.md is not a checkpoint). + # A family-looking repo paired with a non-GGUF single-file name is rejected here, BEFORE the route + # evicts chat (else the background load would be the first to notice README.md is no checkpoint). with pytest.raises(ValueError, match = r"\.gguf"): backend.validate_load_request("unsloth/Z-Image-Turbo-GGUF", gguf_filename = "README.md") assert ( backend.validate_load_request("unsloth/Z-Image-Turbo-GGUF", gguf_filename = "q.gguf").name == "z-image" ) - # A kind/extension mismatch fails fast here, before the route evicts chat + grabs the - # GPU only to fail in the background from_single_file path. + # A kind/extension mismatch fails fast here, before the route evicts chat and grabs the GPU only + # to fail in the background from_single_file path. with pytest.raises(ValueError, match = ".gguf"): backend.validate_load_request( "unsloth/Z-Image-Turbo-GGUF", gguf_filename = "model.safetensors", model_kind = "gguf" @@ -2165,9 +2139,8 @@ def test_validate_load_request(tmp_path): backend.validate_load_request( "unsloth/Qwen-Image-2512-FP8", gguf_filename = "q.gguf", model_kind = "single_file" ) - # A remote "*-GGUF" repo loaded as a full pipeline (no single-file name) is a single-file GGUF - # repo, so from_pretrained finds no pipeline manifest and fails after chat is evicted; reject - # it here before the GPU handoff. + # A remote "*-GGUF" repo loaded as a full pipeline is a single-file GGUF repo, so from_pretrained + # finds no manifest and fails after chat is evicted; reject it before the GPU handoff. with pytest.raises(ValueError, match = "GGUF"): backend.validate_load_request("unsloth/Z-Image-Turbo-GGUF", model_kind = "pipeline") # A local path with a missing child fails here (before any GPU/network work). @@ -2182,8 +2155,8 @@ def test_validate_load_request(tmp_path): ).name == "z-image" ) - # A path-shaped repo_id that does not exist is rejected here (it would otherwise - # be treated as remote, evict chat, and only fail in the background load). + # A path-shaped repo_id that does not exist is rejected here (it would otherwise be treated as + # remote, evict chat, and only fail in the background load). with pytest.raises(FileNotFoundError): backend.validate_load_request( "/tmp/unsloth-definitely-missing-model", @@ -2194,8 +2167,8 @@ def test_validate_load_request(tmp_path): def test_replacement_load_waits_for_inflight_generation(fake_runtime, tmp_path): # A superseding load must signal the in-flight generation's cancel AND wait for it to release - # _generate_lock before allocating, so two pipelines never sit in VRAM at once (unlike - # unload(), which returns promptly without waiting). + # _generate_lock before allocating, so two pipelines never sit in VRAM at once (unlike unload(), + # which returns promptly). import threading backend = DiffusionBackend() @@ -2241,8 +2214,8 @@ def test_replacement_load_waits_for_inflight_generation(fake_runtime, tmp_path): lt = threading.Thread(target = _load) lt.start() - # The load must NOT finish while the generation still holds _generate_lock; it - # has signalled the generation's cancel and is waiting to allocate. + # The load must NOT finish while the generation still holds _generate_lock; it has signalled the + # cancel and is waiting to allocate. assert not load_done.wait(0.5) assert backend._active_generate_cancel is not None assert backend._active_generate_cancel.is_set() @@ -2259,8 +2232,8 @@ def test_replacement_load_waits_for_inflight_generation(fake_runtime, tmp_path): def test_load_reports_memory_plan_fields_on_cpu(fake_runtime, tmp_path): - # The default stub resolves to a CPU target: no offload is possible, but VAE - # tiling is on (no separate device pool), and status carries the new fields. + # The default stub resolves to a CPU target: no offload is possible, but VAE tiling is on (no + # separate device pool), and status carries the new fields. (tmp_path / "m.gguf").write_bytes(b"weights") backend = DiffusionBackend() status = backend.load_pipeline(str(tmp_path), gguf_filename = "m.gguf", family_override = "z-image") @@ -2279,9 +2252,9 @@ def _force_cuda_target(backend, monkeypatch): def test_load_memory_mode_balanced_streams_or_falls_back(fake_runtime, tmp_path, monkeypatch): - # balanced requests streamed block-level (group) offload. Under the stub there's no real + # balanced requests streamed block-level (group) offload. Under the stub there is no real # diffusers.hooks, so group can't engage and the applier falls back to whole-module offload, - # reporting the policy actually engaged (the real "group" path is GPU-verified in the bench). + # reporting the policy actually engaged. (tmp_path / "m.gguf").write_bytes(b"x") backend = DiffusionBackend() _force_cuda_target(backend, monkeypatch) @@ -2294,8 +2267,8 @@ def test_load_memory_mode_balanced_streams_or_falls_back(fake_runtime, tmp_path, def test_load_memory_mode_low_vram_engages_model_offload(fake_runtime, tmp_path, monkeypatch): - # low_vram offloads every component (lowest VRAM); whole-module offload is the - # robust path and engages directly (no streaming, so no diffusers.hooks needed). + # low_vram offloads every component; whole-module offload is the robust path and engages directly + # (no streaming, so no diffusers.hooks needed). (tmp_path / "m.gguf").write_bytes(b"x") backend = DiffusionBackend() _force_cuda_target(backend, monkeypatch) @@ -2310,8 +2283,8 @@ def test_load_memory_mode_low_vram_engages_model_offload(fake_runtime, tmp_path, def test_load_explicit_cpu_offload_engages_model_offload_on_cuda( fake_runtime, tmp_path, monkeypatch ): - # cpu_offload=True with no mode: auto would stay resident (budget unknown under - # the stub), but the explicit flag forces whole-module offload. + # cpu_offload=True with no mode: auto would stay resident (budget unknown under the stub), but + # the explicit flag forces whole-module offload. (tmp_path / "m.gguf").write_bytes(b"x") backend = DiffusionBackend() _force_cuda_target(backend, monkeypatch) @@ -2322,9 +2295,8 @@ def test_load_explicit_cpu_offload_engages_model_offload_on_cuda( def test_load_speed_mode_gguf_auto_defaults_and_explicit(fake_runtime, tmp_path): - # No speed_mode on a GGUF model -> auto `default` (near-lossless, compile sits below the quant - # noise floor). compile only engages on CUDA, so on this CPU stub no optim engages, but the - # resolved mode is `default`. + # No speed_mode on a GGUF model resolves to auto `default`. compile only engages on CUDA, so on + # this CPU stub no optim engages, but the resolved mode is `default`. (tmp_path / "m.gguf").write_bytes(b"x") backend = DiffusionBackend() status = backend.load_pipeline(str(tmp_path), gguf_filename = "m.gguf", family_override = "z-image") @@ -2339,8 +2311,8 @@ def test_load_speed_mode_gguf_auto_defaults_and_explicit(fake_runtime, tmp_path) str(tmp_path), gguf_filename = "m.gguf", family_override = "z-image", speed_mode = "max" ) assert status2["speed_mode"] == "max" - # Text-encoder quant defaults off (None); a requested mode threads through (the - # actual engagement is GPU-verified, since it needs real torch/torchao). + # Text-encoder quant defaults off; a requested mode threads through (actual engagement is + # GPU-verified, since it needs real torch/torchao). assert status2["text_encoder_quant"] is None status3 = backend.load_pipeline( str(tmp_path), @@ -2348,7 +2320,7 @@ def test_load_speed_mode_gguf_auto_defaults_and_explicit(fake_runtime, tmp_path) family_override = "z-image", text_encoder_quant = "nvfp4", ) - # Under the CPU stub nvfp4 is unsupported, so it engages nothing -> None. + # Under the CPU stub nvfp4 is unsupported, so it engages nothing. assert status3["text_encoder_quant"] is None @@ -2382,8 +2354,8 @@ def _stub_dense_quant(monkeypatch, *, scheme = "fp8"): monkeypatch.setattr(_FakeTransformer, "from_pretrained", _from_pretrained, raising = False) monkeypatch.setattr(dmod, "dense_transformer_supported", lambda target: True) - # Resolve the scheme without the real GPU smoke probe, and configure no pre-quant - # checkpoint so the dense materialise+quantise branch is the one exercised. + # Resolve the scheme without the real GPU smoke probe, and configure no pre-quant checkpoint so + # the dense materialise+quantise branch is exercised. monkeypatch.setattr( dmod, "select_transformer_quant_scheme", lambda target, mode, family = None: scheme ) @@ -2399,8 +2371,8 @@ def _stub_dense_quant(monkeypatch, *, scheme = "fp8"): def test_default_load_autos_dense_gate_and_falls_back(fake_runtime, tmp_path, monkeypatch): - # UNSET Dtype defaults to the hardware ladder: the dense gate IS consulted, and a - # device without dense support (this fake runtime) falls back to the GGUF build. + # UNSET Dtype defaults to the hardware ladder: the dense gate IS consulted, and a device without + # dense support falls back to the GGUF build. from core.inference import diffusion as dmod consulted = {"n": 0} @@ -2419,8 +2391,8 @@ def test_default_load_autos_dense_gate_and_falls_back(fake_runtime, tmp_path, mo def test_explicit_off_load_skips_dense_quant_path(fake_runtime, tmp_path, monkeypatch): - # An EXPLICIT "none" pins running the GGUF as-is: the dense gate is never even - # consulted (short-circuit), so the pinned-off contract cannot regress. + # An EXPLICIT "none" pins running the GGUF as-is: the dense gate is never consulted, so the + # pinned-off contract cannot regress. from core.inference import diffusion as dmod monkeypatch.setattr( @@ -2441,9 +2413,8 @@ def test_explicit_off_load_skips_dense_quant_path(fake_runtime, tmp_path, monkey def test_speed_off_load_suppresses_auto_dtype_quant(fake_runtime, tmp_path, monkeypatch): - # An explicit Speed="off" (bit-exact) load with an UNSET dtype must stay GGUF-as-is: the auto - # dtype default must NOT promote it to a quantized + compiled build (silently breaking the - # bit-exact request). The dense gate must never be consulted. + # An explicit Speed="off" load with an UNSET dtype must stay GGUF-as-is: the auto dtype default + # must NOT promote it to a quantized + compiled build. The dense gate must never be consulted. from core.inference import diffusion as dmod monkeypatch.setattr( @@ -2465,8 +2436,8 @@ def test_speed_off_load_suppresses_auto_dtype_quant(fake_runtime, tmp_path, monk def test_transformer_quant_dense_path_engaged(fake_runtime, tmp_path, monkeypatch): - # transformer_quant + a CUDA resident plan -> load the DENSE transformer from the - # base repo, place it on the device, quantise it, and report the engaged scheme. + # transformer_quant + a CUDA resident plan: load the DENSE transformer from the base repo, place + # it, quantise it, and report the engaged scheme. backend = DiffusionBackend() _force_cuda_target(backend, monkeypatch) calls = _stub_dense_quant(monkeypatch, scheme = "fp8") @@ -2478,8 +2449,8 @@ def test_transformer_quant_dense_path_engaged(fake_runtime, tmp_path, monkeypatc transformer_quant = "fp8", ) assert status["transformer_quant"] == "fp8" - # No speed_mode was given, but a quantized transformer is ~30x slower eager, so the - # backend promotes it to `default` (regional compile) instead of the dense `off`. + # No speed_mode was given, but a quantized transformer is ~30x slower eager, so the backend + # promotes it to `default` (regional compile) instead of the dense `off`. assert status["speed_mode"] == "default" assert calls["from_pretrained"] == 1 and calls["quantize"] == 1 assert calls["quant_mode"] == "fp8" @@ -2492,8 +2463,8 @@ def test_transformer_quant_dense_path_engaged(fake_runtime, tmp_path, monkeypatc def test_transformer_quant_prequant_path_engaged(fake_runtime, tmp_path, monkeypatch): - # A configured pre-quant checkpoint -> load the already-quantized transformer directly; - # the dense from_pretrained and the on-device quantize_transformer are NOT used. + # A configured pre-quant checkpoint loads the already-quantized transformer directly; the dense + # from_pretrained and the on-device quantize_transformer are NOT used. from core.inference import diffusion as dmod backend = DiffusionBackend() @@ -2565,8 +2536,8 @@ def test_transformer_quant_prequant_load_fails_falls_back_to_dense( def test_transformer_quant_falls_back_to_gguf_on_failure(fake_runtime, tmp_path, monkeypatch): - # A dense/quant failure (here: quantize returns None -> unsupported) must fall back - # to the GGUF build, not error -- status reports no transformer_quant engaged. + # A dense/quant failure (here quantize returns None) must fall back to the GGUF build, not error; + # status reports no transformer_quant engaged. from core.inference import diffusion as dmod backend = DiffusionBackend() @@ -2592,9 +2563,8 @@ def test_transformer_quant_falls_back_to_gguf_on_failure(fake_runtime, tmp_path, def test_transformer_quant_skipped_when_plan_offloads(fake_runtime, tmp_path, monkeypatch): - # The dense bf16 transformer only fits resident, so when the memory plan would offload (here - # low_vram) the fast path is skipped and GGUF loads instead -- the dense transformer is never - # loaded. + # The dense bf16 transformer only fits resident, so when the memory plan would offload (low_vram) + # the fast path is skipped and GGUF loads instead. from core.inference import diffusion as dmod backend = DiffusionBackend() @@ -2622,16 +2592,16 @@ def test_transformer_quant_skipped_when_plan_offloads(fake_runtime, tmp_path, mo def test_dense_quant_skipped_when_dense_transformer_does_not_fit( fake_runtime, tmp_path, monkeypatch ): - # The GGUF fits resident (plan `none`), but the DENSE bf16 transformer the fast path - # materializes does not. The fast path must be skipped up front (preflighted against the dense - # transformer, not the GGUF), and GGUF loads RESIDENT -- not evicted, OOMed, then offloaded. + # The GGUF fits resident (plan `none`), but the DENSE bf16 transformer the fast path materializes + # does not. The fast path must be skipped up front (preflighted against the dense transformer), + # and GGUF loads RESIDENT -- not evicted, OOMed, then offloaded. from core.inference import diffusion as dmod backend = DiffusionBackend() _force_cuda_target(backend, monkeypatch) monkeypatch.setattr(dmod, "dense_transformer_supported", lambda target: True) - # A scheme resolves and there is no prequant, so the dense bf16 is materialized and the - # dense-fit re-check runs against a large (won't-fit) dense transformer. + # A scheme resolves and there is no prequant, so the dense bf16 is materialized and the dense-fit + # re-check runs against a large (won't-fit) dense transformer. monkeypatch.setattr( dmod, "select_transformer_quant_scheme", lambda target, mode, family = None: "fp8" ) @@ -2676,10 +2646,9 @@ def test_dense_quant_skipped_when_dense_transformer_does_not_fit( def test_dense_quant_prequant_proceeds_but_forbids_dense_fallback( fake_runtime, tmp_path, monkeypatch ): - # With a prequant checkpoint, the fast path loads the small quantized file, so a dense - # misfit must NOT decline the fast path -- but the dense re-check still runs to gate the - # in-loader fallback: if the prequant later fails, the loader must raise to GGUF instead of - # materialising the dense bf16 the plan never budgeted (allow_dense_fallback=False). + # With a prequant checkpoint the fast path loads the small quantized file, so a dense misfit must + # NOT decline it -- but the dense re-check still runs to gate the in-loader fallback: if the + # prequant later fails, the loader must raise to GGUF (allow_dense_fallback=False). from core.inference import diffusion as dmod backend = DiffusionBackend() @@ -2688,8 +2657,8 @@ def test_dense_quant_prequant_proceeds_but_forbids_dense_fallback( monkeypatch.setattr( dmod, "select_transformer_quant_scheme", lambda target, mode, family = None: "fp8" ) - # usable_ (not resolve_): the re-check site only honours a source the loader would - # actually accept, so the fake must present a USABLE one (e.g. a hosted repo). + # usable_ (not resolve_): the re-check site only honours a source the loader would accept, so the + # fake must present a USABLE one. monkeypatch.setattr(dmod, "usable_prequant_source", lambda fam, scheme, **kw: "prequant/path") # Large dense shards cached: if the re-check ran, it would wrongly decline the fast path. monkeypatch.setattr( @@ -2734,10 +2703,9 @@ def test_dense_quant_prequant_proceeds_but_forbids_dense_fallback( def test_dense_quant_replan_retries_once_on_transient_free_undercount( fake_runtime, tmp_path, monkeypatch ): - # A transient foreign allocation at snapshot time makes an empty card look full and the - # candidate replan declines resident -- but the candidate FITS total capacity, so the - # loader must retry the replan once with a fresh settled snapshot instead of silently - # falling back to GGUF-as-is (measured: FLUX.2-dev int8 cold load on an idle B200). + # A transient foreign allocation at snapshot time makes an empty card look full and the candidate + # replan declines resident -- but the candidate FITS total capacity, so the loader must retry the + # replan once with a fresh settled snapshot instead of falling back to GGUF-as-is. import dataclasses from core.inference import diffusion as dmod @@ -2772,8 +2740,7 @@ def test_dense_quant_replan_retries_once_on_transient_free_undercount( return dataclasses.replace(real, offload_policy = "model") replan_calls.append(True) if len(replan_calls) == 1: - # First replan: the transient undercount. Required fits total capacity - # (90,228 <= 0.85 * (183,359 - 18,335)), so a retry must follow. + # First replan: the transient undercount. Required fits total capacity, so a retry must follow. return types.SimpleNamespace( offload_policy = "model", estimates = {"resident_required_mib": 90_228, "safe_device_budget_mib": 40_000}, @@ -2906,9 +2873,8 @@ def _decline_dense_quant(backend, monkeypatch, tmp_path): def test_declined_dense_with_baked_loras_fails_instead_of_silent_drop( fake_runtime, tmp_path, monkeypatch ): - # transformer_quant + adapters, dense build declined for capacity: the GGUF fallback - # cannot bake the adapters, so completing it would silently generate WITHOUT the - # requested LoRAs behind an HTTP success. The load must fail with the recovery options. + # transformer_quant + adapters, dense build declined for capacity: the GGUF fallback cannot bake + # the adapters, so completing it would silently generate WITHOUT them behind an HTTP success. backend = DiffusionBackend() _decline_dense_quant(backend, monkeypatch, tmp_path) with pytest.raises(RuntimeError, match = "LoRA adapters could not be applied"): @@ -2922,8 +2888,8 @@ def test_declined_dense_with_baked_loras_fails_instead_of_silent_drop( def test_declined_dense_without_loras_still_falls_back_to_gguf(fake_runtime, tmp_path, monkeypatch): - # The plain decline (no adapters requested) keeps the silent GGUF fallback: weight-0 - # adapters count as "none" (an explicit disable is not a bake request). + # The plain decline (no adapters requested) keeps the silent GGUF fallback: weight-0 adapters + # count as "none". backend = DiffusionBackend() _decline_dense_quant(backend, monkeypatch, tmp_path) result = backend.load_pipeline( @@ -3027,8 +2993,8 @@ def _quant_lora_state(pipe, quant = "int8"): def test_apply_loras_quant_unbaked_requires_reload(monkeypatch): - # A quantized pipe built WITHOUT adapters cannot take one at generation time (topology is - # frozen after quantize_ + compile): clean 400 telling the client to reload. + # A quantized pipe built WITHOUT adapters cannot take one at generation time (topology frozen + # after quantize_ + compile): a clean 400 telling the client to reload. backend = DiffusionBackend() pipe = _BakePipe() with pytest.raises(ValueError, match = "Reload the model with the adapter selection"): @@ -3039,9 +3005,8 @@ def test_apply_loras_quant_unbaked_requires_reload(monkeypatch): def test_apply_loras_quant_baked_matrix(monkeypatch): - # Baked pipe: same set -> no-op; weight-only change -> set_adapters (value-level, no - # topology change); empty -> all scales 0 (reproduces the quantized base); different - # adapter set -> reload error. + # Baked pipe: same set is a no-op; weight-only change calls set_adapters; empty scales all to 0 + # (reproducing the quantized base); a different adapter set errors. backend = DiffusionBackend() monkeypatch.setattr( DiffusionBackend, @@ -3082,9 +3047,8 @@ def test_apply_loras_quant_baked_matrix(monkeypatch): def test_assemble_pipe_routes_krea2_per_component(monkeypatch): - # krea's repo ships transformers-5.x configs and no top-level tokenizer files, so - # Pipeline.from_pretrained dies in the tokenizer (vocab_file = None). The quant fast - # path must assemble per-component via load_krea2_pipeline like every other krea load. + # krea's repo ships transformers-5.x configs and no top-level tokenizer files, so from_pretrained + # dies in the tokenizer. The quant fast path must assemble per-component via load_krea2_pipeline. from core.inference import diffusion as dmod calls: dict = {} @@ -3127,10 +3091,9 @@ def test_assemble_pipe_routes_krea2_per_component(monkeypatch): def test_dense_quant_unusable_prequant_path_runs_dense_refit(fake_runtime, tmp_path, monkeypatch): - # A request-supplied transformer_prequant_path the loader refuses (missing, or outside - # UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH) resolves to NO usable prequant source, so the - # dense-transformer fit re-check MUST run: with real device budgets it declines the fast path - # up front instead of evicting the resident pipeline and OOMing in the dense bf16 fallback. + # A request-supplied transformer_prequant_path the loader refuses resolves to NO usable prequant + # source, so the dense-transformer fit re-check MUST run: it declines the fast path up front + # instead of evicting the resident pipeline and OOMing in the dense bf16 fallback. from core.inference import diffusion as dmod backend = DiffusionBackend() @@ -3139,8 +3102,8 @@ def test_dense_quant_unusable_prequant_path_runs_dense_refit(fake_runtime, tmp_p monkeypatch.setattr( dmod, "select_transformer_quant_scheme", lambda target, mode, family = None: "fp8" ) - # The REAL usable_prequant_source refuses a non-allowlisted path (unit-tested in - # test_diffusion_prequant.py); returning None here pins that outcome at this site. + # The REAL usable_prequant_source refuses a non-allowlisted path (unit-tested elsewhere); + # returning None here pins that outcome at this site. monkeypatch.setattr(dmod, "usable_prequant_source", lambda fam, scheme, **kw: None) monkeypatch.setattr( DiffusionBackend, @@ -3180,10 +3143,9 @@ def test_dense_quant_unusable_prequant_path_runs_dense_refit(fake_runtime, tmp_p def test_transformer_quant_unsupported_scheme_skips_dense_download( fake_runtime, tmp_path, monkeypatch ): - # An explicit unsupported scheme (select_transformer_quant_scheme -> None) must fail the dense - # path BEFORE materialising the multi-GB dense transformer, then fall back to GGUF -- else the - # download runs under the load lock during finalization after the old model was evicted, only - # to fail at quantize. + # An explicit unsupported scheme must fail the dense path BEFORE materialising the multi-GB dense + # transformer, then fall back to GGUF -- else the download runs under the load lock after the old + # model was evicted, only to fail at quantize. from core.inference import diffusion as dmod backend = DiffusionBackend() @@ -3212,8 +3174,8 @@ def test_transformer_quant_unsupported_scheme_skips_dense_download( def test_base_file_downloaded_include_transformer_flag(): - # Default: transformer/ shards are the GGUF's job, so they are excluded from - # the prefetch list; the dense transformer-quant path opts them back in. + # Default: transformer/ shards are the GGUF's job, so they are excluded from the prefetch list; + # the dense transformer-quant path opts them back in. from core.inference.diffusion import _base_file_downloaded assert _base_file_downloaded("transformer/diffusion_pytorch_model-00001.safetensors") is False @@ -3229,11 +3191,9 @@ def test_base_file_downloaded_include_transformer_flag(): def test_dense_quant_prefetch_capacity_gate(fake_runtime, monkeypatch): - # On a device that cannot hold even the candidate's post-quant resident set, the re-plan - # is certain to decline the dense path -- widening would fetch the multi-GB base - # transformer/ shards (measured: ~47 GB on a Qwen-Image GGUF load on a 24 GB card) only - # to run the GGUF as-is. The gate compares steady_total against TOTAL capacity (reserve + - # 0.85 margin), never the instantaneous free reading. + # On a device that cannot hold even the candidate's post-quant resident set, the re-plan is + # certain to decline the dense path, so widening would fetch the multi-GB base transformer/ shards + # only to run the GGUF as-is. The gate compares steady_total against TOTAL capacity. from core.inference import diffusion as dmod from core.inference import diffusion_memory as dmem @@ -3251,7 +3211,7 @@ def test_dense_quant_prefetch_capacity_gate(fake_runtime, monkeypatch): total_mib = 24_564, free_mib = 24_000, memory_kind = "discrete_vram" ), ) - # int8 qwen steady (~22 GB DiT + 17 GB companions = 39 GB) cannot fit a 24 GB card. + # int8 qwen steady (~22 GB DiT + 17 GB companions) cannot fit a 24 GB card. monkeypatch.setattr(dmod, "resolve_dense_quant_candidate", candidate_with(39_900)) assert backend._dense_quant_prefetch_needed(fam, {"transformer_quant": "int8"}) is False # A candidate that fits total capacity still widens. @@ -3268,9 +3228,8 @@ def test_dense_quant_prefetch_capacity_gate(fake_runtime, monkeypatch): def test_dense_quant_prefetch_needed_gates(fake_runtime, monkeypatch): # The transformer/ prefetch widens exactly when load_pipeline takes the dense-quant path: it - # defers to resolve_dense_quant_candidate (quant requested + device supported + scheme - # resolvable + no prequant checkpoint + disk for the extra bf16 shards). An explicit - # Speed="off" (bit-exact) load never widens. + # defers to resolve_dense_quant_candidate (quant requested + device supported + scheme resolvable + # + no prequant checkpoint + disk). An explicit Speed="off" load never widens. from core.inference import diffusion as dmod backend = DiffusionBackend() @@ -3290,23 +3249,21 @@ def test_dense_quant_prefetch_needed_gates(fake_runtime, monkeypatch): logger = None, ): seen.append(requested) - # A real (non-prequant) dense-quant candidate: scheme resolves AND disk fits, so the - # loader takes the dense build that needs the base repo's bf16 transformer/ shards. + # A real (non-prequant) dense-quant candidate: scheme resolves AND disk fits, so the loader takes + # the dense build that needs the base repo's bf16 transformer/ shards. return types.SimpleNamespace(prequant = False) monkeypatch.setattr(dmod, "resolve_dense_quant_candidate", fake_candidate) - # Explicit fp8 -> widens; the resolved mode is threaded through to the candidate resolver. + # Explicit fp8 widens; the resolved mode is threaded through to the candidate resolver. assert backend._dense_quant_prefetch_needed(fam, {"transformer_quant": "fp8"}) is True assert seen[-1] == "fp8" - # UNSET defaults to the hardware ladder (Dtype default-auto) -> widens, threading auto. + # UNSET defaults to the hardware ladder, so it widens, threading auto. assert backend._dense_quant_prefetch_needed(fam, {}) is True assert seen[-1] == "auto" - # A definite-offload memory policy forces load_pipeline onto offload regardless of the dense - # candidate's smaller footprint, so the dense build never runs and the widened prefetch would - # download base transformer/ shards the offloaded GGUF path never uses (disk-full there has no - # GGUF fallback). balanced / low_vram (and the legacy cpu_offload flag absent a memory_mode) - # must NOT widen, even though the candidate is dense-viable. + # A definite-offload memory policy forces load_pipeline onto offload regardless of the candidate's + # smaller footprint, so the widened prefetch would download shards the offloaded GGUF path never + # uses (and a disk-full there has no GGUF fallback). balanced / low_vram must NOT widen. before = len(seen) assert ( backend._dense_quant_prefetch_needed( @@ -3327,7 +3284,7 @@ def test_dense_quant_prefetch_needed_gates(fake_runtime, monkeypatch): # The gate short-circuits BEFORE resolving the candidate (no wasted resolve). assert len(seen) == before # An explicit memory_mode still consulting the candidate: fast/auto can flip resident, so they - # widen when the candidate is dense-viable (memory_mode="fast" does not force offload). + # widen when the candidate is dense-viable. assert ( backend._dense_quant_prefetch_needed( fam, {"transformer_quant": "fp8", "memory_mode": "fast"} @@ -3341,46 +3298,43 @@ def test_dense_quant_prefetch_needed_gates(fake_runtime, monkeypatch): ) is True ) - # An explicit off pins running the GGUF as-is -> never widen (mode resolves to None first). + # An explicit off pins running the GGUF as-is, so never widen. assert backend._dense_quant_prefetch_needed(fam, {"transformer_quant": "none"}) is False - # An explicit Speed="off" (bit-exact) load suppresses the dense path -> never widen. + # An explicit Speed="off" (bit-exact) load suppresses the dense path, so never widen. assert ( backend._dense_quant_prefetch_needed(fam, {"transformer_quant": "fp8", "speed_mode": "off"}) is False ) - # A PREQUANT candidate loads the small pre-quantized checkpoint (+ config / companions), NOT - # the base repo's dense transformer/ shards, so the widened prefetch must NOT fire -- else it - # defeats the prequant savings and can hard-fail begin_load (no GGUF fallback) on a disk-full. + # A PREQUANT candidate loads the small pre-quantized checkpoint, NOT the base repo's dense shards, + # so the widened prefetch must NOT fire -- else it defeats the prequant savings and can hard-fail + # begin_load on a disk-full. monkeypatch.setattr( dmod, "resolve_dense_quant_candidate", lambda **kw: types.SimpleNamespace(prequant = True) ) assert backend._dense_quant_prefetch_needed(fam, {"transformer_quant": "fp8"}) is False - # No viable candidate (unsupported scheme / no disk room) -> never widen. The disk guard here - # averts filling the cache volume and hard-failing the load instead of falling back to GGUF. + # No viable candidate (unsupported scheme / no disk room) never widens. The disk guard averts + # filling the cache volume and hard-failing instead of falling back to GGUF. monkeypatch.setattr(dmod, "resolve_dense_quant_candidate", lambda **kw: None) assert backend._dense_quant_prefetch_needed(fam, {"transformer_quant": "fp8"}) is False def test_diffusion_status_response_carries_resolved(): - # The backend records per-control auto-policy provenance (build_resolved_record) on - # state.resolved; the response model must DECLARE the field or Pydantic's extra='ignore' drops - # it, leaving that plumbing dead (never reaching a client). + # The backend records per-control auto-policy provenance on state.resolved; the response model + # must DECLARE the field or Pydantic's extra='ignore' drops it, leaving that plumbing dead. from models.inference import DiffusionStatusResponse rec = {"transformer_quant": {"value": "fp8", "source": "auto", "reason": "blackwell"}} resp = DiffusionStatusResponse(loaded = True, resolved = rec) - # The typed field coerces the plain record into DiffusionResolvedControl objects; the - # serialized form must round-trip back to the record, proving the field is DECLARED and not - # dropped by Pydantic's extra='ignore'. + # The typed field coerces the plain record into DiffusionResolvedControl objects; the serialized + # form must round-trip back, proving the field is DECLARED. assert resp.model_dump()["resolved"] == rec # Absent by default (nothing resolved / native engine). assert DiffusionStatusResponse(loaded = False).resolved is None def test_companion_cache_bytes_local_dir_excludes_transformer(tmp_path): - # A LOCAL diffusers base: sum the on-disk VAE / text-encoder weights so auto memory planning - # sees the resident companions, but exclude transformer/ (the GGUF supplies it) and non-weight - # files. A folded-to-zero companion could OOM a resident plan. + # A LOCAL diffusers base: sum the on-disk VAE / text-encoder weights so auto memory planning sees + # the resident companions, but exclude transformer/ (the GGUF supplies it) and non-weight files. (tmp_path / "vae").mkdir() (tmp_path / "vae" / "diffusion_pytorch_model.safetensors").write_bytes(b"x" * 100) (tmp_path / "text_encoder").mkdir() @@ -3394,26 +3348,24 @@ def test_companion_cache_bytes_local_dir_excludes_transformer(tmp_path): def test_plan_memory_dense_replan_does_not_double_count_prefetched_transformer(monkeypatch): # Re-planning the dense transformer-quant candidate: the dense path prefetches the base repo's - # transformer/ shards into the SAME blob cache _companion_cache_bytes sums. If the re-plan read - # that cache it would count the transformer TWICE (as transformer_resident_override_mib and as - # a "companion") and force offload even when the quantised artifact fits resident. The re-plan - # must use the auto-policy's companion estimate. Here the cache is stubbed to the inflated - # value; the plan must still stay resident. + # transformer/ shards into the SAME blob cache _companion_cache_bytes sums, so reading that cache + # would count the transformer TWICE and force offload even when the quantised artifact fits. Here + # the cache is stubbed to the inflated value; the plan must still stay resident. from core.inference import diffusion as dmod from core.inference.diffusion_memory import OFFLOAD_NONE, DeviceMemory backend = DiffusionBackend() target = types.SimpleNamespace(device = "cuda", backend = "cuda", supports_model_cpu_offload = True) - # 40 GiB discrete card, 40000 MiB free: comfortably fits transformer + real - # companions + headroom, but NOT a second copy of the bf16 transformer. + # 40 GiB discrete card: comfortably fits transformer + real companions + headroom, but NOT a + # second copy of the bf16 transformer. monkeypatch.setattr( dmod, "settled_snapshot_device_memory", lambda t: DeviceMemory("cuda", "cuda", "discrete_vram", 40000, 40960), ) monkeypatch.setattr(dmod, "estimate_image_runtime_mib", lambda **kw: 4000) - # The cache is inflated by the prefetched bf16 transformer (~24000) on top of the - # ~8000 real companions; if the re-plan consulted it the plan would offload. + # The cache is inflated by the prefetched bf16 transformer on top of the real companions; if the + # re-plan consulted it the plan would offload. monkeypatch.setattr( DiffusionBackend, "_companion_cache_bytes", @@ -3431,16 +3383,15 @@ def test_plan_memory_dense_replan_does_not_double_count_prefetched_transformer(m transformer_resident_override_mib = 12000, # int8 candidate transient (~half bf16) companion_override_mib = 8000, # auto-policy text-encoder + VAE estimate ) - # 12000 + 8000 + 4000 + 2048 overhead = 26048 MiB, fits the ~36 GiB budget. - # A double-count (12000 + [8000+24000] + ...) would have exceeded it and offloaded. + # 12000 + 8000 + 4000 + 2048 overhead = 26048 MiB, fits the ~36 GiB budget. A double-count would + # have exceeded it and offloaded. assert plan.offload_policy == OFFLOAD_NONE def test_reset_step_cache_helper_is_best_effort(): - # Prefers the real diffusers CacheMixin hook (_reset_stateful_cache): on a genuine Flux / - # QwenImage transformer that is the reset entry point, and reset_stateful_hooks lives only on - # the HookRegistry (getattr on the transformer returns None), so the old lookup was a silent - # no-op that left stale FBCache residuals for the next generation. + # Prefers the real diffusers CacheMixin hook (_reset_stateful_cache): reset_stateful_hooks lives + # only on the HookRegistry, so the old lookup was a silent no-op that left stale FBCache + # residuals for the next generation. calls = [] pipe = types.SimpleNamespace( transformer = types.SimpleNamespace(_reset_stateful_cache = lambda: calls.append("real")) @@ -3464,14 +3415,14 @@ def test_reset_step_cache_helper_is_best_effort(): ) DiffusionBackend._reset_step_cache(pipe) assert calls == ["fallback"] - # No transformer, or a transformer without either hook -> silent no-op (never raises). + # No transformer, or one without either hook, is a silent no-op (never raises). DiffusionBackend._reset_step_cache(types.SimpleNamespace()) DiffusionBackend._reset_step_cache(types.SimpleNamespace(transformer = object())) def test_generate_resets_step_cache_only_when_engaged(fake_runtime, tmp_path): - # FBCache residuals live on the resident transformer across generations, so each - # generate() must reset the stateful cache first -- but only when a cache is engaged. + # FBCache residuals live on the resident transformer across generations, so each generate() must + # reset the stateful cache first -- but only when a cache is engaged. (tmp_path / "model.gguf").write_bytes(b"weights") backend = DiffusionBackend() backend.load_pipeline( @@ -3481,8 +3432,8 @@ def test_generate_resets_step_cache_only_when_engaged(fake_runtime, tmp_path): family_override = "z-image", ) resets = [] - # Use the real diffusers CacheMixin entry point (_reset_stateful_cache); a genuine - # Flux/QwenImage transformer exposes this, not reset_stateful_hooks. + # Use the real diffusers CacheMixin entry point; a genuine Flux/QwenImage transformer exposes + # this, not reset_stateful_hooks. backend._state.pipe.transformer = types.SimpleNamespace( _reset_stateful_cache = lambda: resets.append(True) ) @@ -3497,8 +3448,8 @@ def test_generate_resets_step_cache_only_when_engaged(fake_runtime, tmp_path): def test_prefetch_returns_snapshot_dir_for_manifest(monkeypatch): - # The prefetched pipeline manifest's directory is the local snapshot root; a - # config-only base list (no manifest) returns None so the hub id stays in use. + # The prefetched pipeline manifest's directory is the local snapshot root; a config-only base list + # returns None so the hub id stays in use. backend = DiffusionBackend() monkeypatch.setattr( "utils.hf_xet_fallback.hf_hub_download_with_xet_fallback", @@ -3515,7 +3466,7 @@ def test_prefetch_returns_snapshot_dir_for_manifest(monkeypatch): def test_pipeline_load_uses_predownloaded_dir(fake_runtime, tmp_path): # With a prefetched snapshot, from_pretrained must receive the local dir -- its own hub sweep - # would re-download the root packaged singles the scoped prefetch skips (24 GB per FLUX.1 repo). + # would re-download the root packaged singles the scoped prefetch skips (24 GB per FLUX.1). backend = DiffusionBackend() backend.load_pipeline( "unsloth/Qwen-Image-2512-bnb-4bit", @@ -3568,7 +3519,7 @@ def test_unload_waits_for_in_flight_denoise_before_teardown(): u = threading.Thread(target = _unload) u.start() assert cancel.wait(2.0) # unload has signalled the denoise and is now waiting on _generate_lock - # unload must NOT have torn down yet -- it is blocked on the denoise's _generate_lock. + # unload must NOT have torn down yet: it is blocked on the denoise's _generate_lock. assert teardown_saw == [] assert not unloaded.wait(0.3) @@ -3576,7 +3527,7 @@ def test_unload_waits_for_in_flight_denoise_before_teardown(): d.join(2.0) u.join(2.0) assert unloaded.is_set() - # Teardown ran exactly once, and only AFTER the denoise had exited (denoise_active was False). + # Teardown ran exactly once, and only AFTER the denoise had exited. assert teardown_saw == [False] @@ -3628,8 +3579,8 @@ def test_generate_prompt_list_with_matching_seed_list(fake_runtime, tmp_path): def test_generate_single_image_keeps_scalar_generator(fake_runtime, tmp_path): - # The single-image call shape is the bit-identical reference path: scalar prompt, - # ONE scalar generator (not a 1-list), num_images_per_prompt=1. + # The single-image call shape is the bit-identical reference path: scalar prompt, ONE scalar + # generator (not a 1-list), num_images_per_prompt=1. backend = _load_zimage_backend(tmp_path) out = backend.generate(prompt = "one", seed = 5) call = backend._state.pipe.last_kwargs @@ -3640,8 +3591,8 @@ def test_generate_single_image_keeps_scalar_generator(fake_runtime, tmp_path): def test_generate_batched_seed_matches_solo_replay(fake_runtime, tmp_path): - # Per-image reproducibility contract: image i of a batched call is driven by the - # exact generator seed a solo replay of that image uses. + # Per-image reproducibility: image i of a batched call is driven by the exact generator seed a + # solo replay of that image uses. backend = _load_zimage_backend(tmp_path) backend.generate(prompt = "p", seeds = [3, 9]) batched = [g.manual for g in backend._state.pipe.last_kwargs["generator"]] @@ -3727,10 +3678,9 @@ def test_generate_non_oom_error_is_not_retried(fake_runtime, tmp_path): def test_generate_broadcasts_negative_prompt_across_a_mixed_prompt_batch(fake_runtime, tmp_path): - # A prompt LIST must carry a matching negative-prompt LIST: ZImagePipeline.encode_prompt - # asserts len(prompt) == len(negative_prompt), and the pipes that encode the negative - # separately (Qwen-Image / Krea 2 / FLUX true-CFG) would build batch-1 negative embeds - # against batch-N latents and fail in the transformer's txt/img concat. + # A prompt LIST must carry a matching negative-prompt LIST: ZImagePipeline.encode_prompt asserts + # equal lengths, and the pipes that encode the negative separately would build batch-1 negative + # embeds against batch-N latents and fail in the transformer's txt/img concat. backend = _load_zimage_backend(tmp_path) backend.generate(prompt = "fallback", prompts = ["a", "b", "c"], negative_prompt = "blurry") call = backend._state.pipe.last_kwargs @@ -3742,8 +3692,8 @@ def test_generate_broadcasts_negative_prompt_across_a_mixed_prompt_batch(fake_ru def test_generate_keeps_a_scalar_negative_prompt_off_the_list_paths(fake_runtime, tmp_path): - # Uniform-prompt and single-image forwards pass a SCALAR prompt, so the negative prompt - # must stay scalar too (a list would mismatch the batch-1 positive encode). + # Uniform-prompt and single-image forwards pass a SCALAR prompt, so the negative prompt must stay + # scalar too (a list would mismatch the batch-1 positive encode). backend = _load_zimage_backend(tmp_path) backend.generate(prompt = "a sloth", seeds = [1, 2, 3], negative_prompt = "blurry") assert backend._state.pipe.last_kwargs["prompt"] == "a sloth" @@ -3777,11 +3727,10 @@ class _TracingPipe(_CountingPipe): def test_generate_resets_the_step_cache_before_an_oom_retry(fake_runtime, tmp_path): - # A forward that RAISES skips the pipeline's end-of-__call__ maybe_free_model_hooks(), - # so its FBCache head-block residual stays on the resident transformer. Without a reset - # before each retry the halved chunk compares a batch-2 residual against the stale - # batch-4 one ("size of tensor a (2) must match ... b (4)"), turning a recoverable OOM - # into a hard failure. + # A forward that RAISES skips the pipeline's end-of-call maybe_free_model_hooks(), so its FBCache + # residual stays on the resident transformer. Without a reset before each retry the halved chunk + # compares a batch-2 residual against the stale batch-4 one, turning a recoverable OOM into a + # hard failure. backend = _load_zimage_backend(tmp_path) trace: list = [] pipe = _TracingPipe(trace, max_images = 2) @@ -3790,7 +3739,7 @@ def test_generate_resets_the_step_cache_before_an_oom_retry(fake_runtime, tmp_pa object.__setattr__(backend._state, "transformer_cache", "fbcache") out = backend.generate(prompt = "p", seeds = [1, 2, 3, 4]) assert len(out["images"]) == 4 and out["seeds"] == [1, 2, 3, 4] - # Every forward -- including both post-OOM retries -- is preceded by a reset. + # Every forward, including both post-OOM retries, is preceded by a reset. assert trace == [ ("reset",), ("call", 4), @@ -3825,8 +3774,8 @@ class _FakeInfo: GB = 1024**3 -# A FLUX-shaped base repo: the packaged root single and the transformer shards are what a -# plain snapshot_download would drag in and the loader never opens. +# A FLUX-shaped base repo: the packaged root single and the transformer shards are what a plain +# snapshot_download would drag in and the loader never opens. _FLUX_BASE_SIBLINGS = [ _FakeSibling("model_index.json", 1000), _FakeSibling("flux1-dev.safetensors", 24 * GB), @@ -3855,9 +3804,8 @@ def _fake_hf_api(monkeypatch, repos): def test_download_plan_scopes_the_base_repo_files(monkeypatch): - # The plan drives the Hub download manager, so its file list must match what the loader - # actually reads. A full snapshot would add the 24 GB root single and the transformer - # shards the GGUF replaces. + # The plan drives the Hub download manager, so its file list must match what the loader actually + # reads. A full snapshot would add the 24 GB root single and the shards the GGUF replaces. _fake_hf_api( monkeypatch, { diff --git a/studio/backend/tests/test_diffusion_base_precision.py b/studio/backend/tests/test_diffusion_base_precision.py index 09d9018364..5193093986 100644 --- a/studio/backend/tests/test_diffusion_base_precision.py +++ b/studio/backend/tests/test_diffusion_base_precision.py @@ -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 ( diff --git a/studio/backend/tests/test_diffusion_cache.py b/studio/backend/tests/test_diffusion_cache.py index eb8db319d4..a6730c04aa 100644 --- a/studio/backend/tests/test_diffusion_cache.py +++ b/studio/backend/tests/test_diffusion_cache.py @@ -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"]) diff --git a/studio/backend/tests/test_diffusion_compile_cache.py b/studio/backend/tests/test_diffusion_compile_cache.py index cbb2cab307..fe7de6bffb 100644 --- a/studio/backend/tests/test_diffusion_compile_cache.py +++ b/studio/backend/tests/test_diffusion_compile_cache.py @@ -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)) diff --git a/studio/backend/tests/test_diffusion_cond_cache.py b/studio/backend/tests/test_diffusion_cond_cache.py index 62e69041ac..456aebbe72 100644 --- a/studio/backend/tests/test_diffusion_cond_cache.py +++ b/studio/backend/tests/test_diffusion_cond_cache.py @@ -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"]) diff --git a/studio/backend/tests/test_diffusion_controlnet.py b/studio/backend/tests/test_diffusion_controlnet.py index fb11160960..7adbec3a26 100644 --- a/studio/backend/tests/test_diffusion_controlnet.py +++ b/studio/backend/tests/test_diffusion_controlnet.py @@ -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 diff --git a/studio/backend/tests/test_diffusion_dataset_api.py b/studio/backend/tests/test_diffusion_dataset_api.py index ec7a580d95..0cd4094dfa 100644 --- a/studio/backend/tests/test_diffusion_dataset_api.py +++ b/studio/backend/tests/test_diffusion_dataset_api.py @@ -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) diff --git a/studio/backend/tests/test_diffusion_device.py b/studio/backend/tests/test_diffusion_device.py index 2eb9b3a822..88f7d16c26 100644 --- a/studio/backend/tests/test_diffusion_device.py +++ b/studio/backend/tests/test_diffusion_device.py @@ -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" diff --git a/studio/backend/tests/test_diffusion_dit_trainer.py b/studio/backend/tests/test_diffusion_dit_trainer.py index 0eebe4e5a6..4c2848e10a 100644 --- a/studio/backend/tests/test_diffusion_dit_trainer.py +++ b/studio/backend/tests/test_diffusion_dit_trainer.py @@ -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) diff --git a/studio/backend/tests/test_diffusion_dit_trainer_flow_shift.py b/studio/backend/tests/test_diffusion_dit_trainer_flow_shift.py index d21a0eed62..ae4941943c 100644 --- a/studio/backend/tests/test_diffusion_dit_trainer_flow_shift.py +++ b/studio/backend/tests/test_diffusion_dit_trainer_flow_shift.py @@ -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()) diff --git a/studio/backend/tests/test_diffusion_eager_patches.py b/studio/backend/tests/test_diffusion_eager_patches.py index a32dbec2ee..ea8f614317 100644 --- a/studio/backend/tests/test_diffusion_eager_patches.py +++ b/studio/backend/tests/test_diffusion_eager_patches.py @@ -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) diff --git a/studio/backend/tests/test_diffusion_engine_router.py b/studio/backend/tests/test_diffusion_engine_router.py index dfcfa52be5..649608ba81 100644 --- a/studio/backend/tests/test_diffusion_engine_router.py +++ b/studio/backend/tests/test_diffusion_engine_router.py @@ -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") diff --git a/studio/backend/tests/test_diffusion_gguf_compile.py b/studio/backend/tests/test_diffusion_gguf_compile.py index c453713d29..cceb0fe0ef 100644 --- a/studio/backend/tests/test_diffusion_gguf_compile.py +++ b/studio/backend/tests/test_diffusion_gguf_compile.py @@ -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() diff --git a/studio/backend/tests/test_diffusion_inference_info.py b/studio/backend/tests/test_diffusion_inference_info.py index 7f2a0ae3f6..053a5346eb 100644 --- a/studio/backend/tests/test_diffusion_inference_info.py +++ b/studio/backend/tests/test_diffusion_inference_info.py @@ -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"] diff --git a/studio/backend/tests/test_diffusion_krea2.py b/studio/backend/tests/test_diffusion_krea2.py index 74556e2179..6f1728e672 100644 --- a/studio/backend/tests/test_diffusion_krea2.py +++ b/studio/backend/tests/test_diffusion_krea2.py @@ -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])) diff --git a/studio/backend/tests/test_diffusion_lora.py b/studio/backend/tests/test_diffusion_lora.py index fe54c1ce65..d342543c41 100644 --- a/studio/backend/tests/test_diffusion_lora.py +++ b/studio/backend/tests/test_diffusion_lora.py @@ -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 ", [r]) == "a cat " 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 ", [r]) assert "" not in out assert out == "a cat " @@ -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 diff --git a/studio/backend/tests/test_diffusion_lora_trainer.py b/studio/backend/tests/test_diffusion_lora_trainer.py index 4cad9388fa..6f9288eca0 100644 --- a/studio/backend/tests/test_diffusion_lora_trainer.py +++ b/studio/backend/tests/test_diffusion_lora_trainer.py @@ -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" diff --git a/studio/backend/tests/test_diffusion_memory.py b/studio/backend/tests/test_diffusion_memory.py index 7917f7a5fb..055e137701 100644 --- a/studio/backend/tests/test_diffusion_memory.py +++ b/studio/backend/tests/test_diffusion_memory.py @@ -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 diff --git a/studio/backend/tests/test_diffusion_more_families.py b/studio/backend/tests/test_diffusion_more_families.py index 2eebf34e71..0b972c2399 100644 --- a/studio/backend/tests/test_diffusion_more_families.py +++ b/studio/backend/tests/test_diffusion_more_families.py @@ -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") diff --git a/studio/backend/tests/test_diffusion_precision.py b/studio/backend/tests/test_diffusion_precision.py index b835434b37..38f1f24feb 100644 --- a/studio/backend/tests/test_diffusion_precision.py +++ b/studio/backend/tests/test_diffusion_precision.py @@ -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) diff --git a/studio/backend/tests/test_diffusion_prequant.py b/studio/backend/tests/test_diffusion_prequant.py index aa67756e94..fb88f33a7e 100644 --- a/studio/backend/tests/test_diffusion_prequant.py +++ b/studio/backend/tests/test_diffusion_prequant.py @@ -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_.pt name serves the download after - # the model-name filename 404s; both names are requested in order. + # A repo still carrying the legacy transformer_.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 = {} diff --git a/studio/backend/tests/test_diffusion_routes.py b/studio/backend/tests/test_diffusion_routes.py index 9e7306e577..360537c625 100644 --- a/studio/backend/tests/test_diffusion_routes.py +++ b/studio/backend/tests/test_diffusion_routes.py @@ -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 = {} diff --git a/studio/backend/tests/test_diffusion_sdxl.py b/studio/backend/tests/test_diffusion_sdxl.py index 6bec0a4100..37dc8d4404 100644 --- a/studio/backend/tests/test_diffusion_sdxl.py +++ b/studio/backend/tests/test_diffusion_sdxl.py @@ -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") diff --git a/studio/backend/tests/test_diffusion_speed.py b/studio/backend/tests/test_diffusion_speed.py index 746d05ac1c..953dd839a6 100644 --- a/studio/backend/tests/test_diffusion_speed.py +++ b/studio/backend/tests/test_diffusion_speed.py @@ -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 diff --git a/studio/backend/tests/test_diffusion_te_prequant.py b/studio/backend/tests/test_diffusion_te_prequant.py index e7aa0c7f1a..38f2c50e06 100644 --- a/studio/backend/tests/test_diffusion_te_prequant.py +++ b/studio/backend/tests/test_diffusion_te_prequant.py @@ -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)) diff --git a/studio/backend/tests/test_diffusion_train_perf.py b/studio/backend/tests/test_diffusion_train_perf.py index dd79fb131d..1e9a7ea0e7 100644 --- a/studio/backend/tests/test_diffusion_train_perf.py +++ b/studio/backend/tests/test_diffusion_train_perf.py @@ -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)) diff --git a/studio/backend/tests/test_diffusion_transformer_quant.py b/studio/backend/tests/test_diffusion_transformer_quant.py index 6a38dcfeb7..3a5b3f9b1a 100644 --- a/studio/backend/tests/test_diffusion_transformer_quant.py +++ b/studio/backend/tests/test_diffusion_transformer_quant.py @@ -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()) diff --git a/studio/backend/tests/test_gguf_load_cache_reuse.py b/studio/backend/tests/test_gguf_load_cache_reuse.py index 8fb49e24d7..621df67379 100644 --- a/studio/backend/tests/test_gguf_load_cache_reuse.py +++ b/studio/backend/tests/test_gguf_load_cache_reuse.py @@ -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 ( diff --git a/studio/backend/tests/test_gpu_arbiter.py b/studio/backend/tests/test_gpu_arbiter.py index 30010b11ca..4ed268b6e3 100644 --- a/studio/backend/tests/test_gpu_arbiter.py +++ b/studio/backend/tests/test_gpu_arbiter.py @@ -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( diff --git a/studio/backend/tests/test_gpu_memory_mode.py b/studio/backend/tests/test_gpu_memory_mode.py index fb731deadf..18c5205781 100644 --- a/studio/backend/tests/test_gpu_memory_mode.py +++ b/studio/backend/tests/test_gpu_memory_mode.py @@ -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 diff --git a/studio/backend/tests/test_gpu_selection.py b/studio/backend/tests/test_gpu_selection.py index b2eed1e19a..fa1dc1cf64 100644 --- a/studio/backend/tests/test_gpu_selection.py +++ b/studio/backend/tests/test_gpu_selection.py @@ -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 diff --git a/studio/backend/tests/test_hf_cache_settings.py b/studio/backend/tests/test_hf_cache_settings.py index a27cea00f5..4ebf44a78f 100644 --- a/studio/backend/tests/test_hf_cache_settings.py +++ b/studio/backend/tests/test_hf_cache_settings.py @@ -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("): diff --git a/studio/backend/tests/test_image_gallery.py b/studio/backend/tests/test_image_gallery.py index 466398cf13..2c3ad5006f 100644 --- a/studio/backend/tests/test_image_gallery.py +++ b/studio/backend/tests/test_image_gallery.py @@ -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 diff --git a/studio/backend/tests/test_inference_model_validation.py b/studio/backend/tests/test_inference_model_validation.py index aa6891d4f8..95b50c7d2e 100644 --- a/studio/backend/tests/test_inference_model_validation.py +++ b/studio/backend/tests/test_inference_model_validation.py @@ -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" diff --git a/studio/backend/tests/test_local_model_format.py b/studio/backend/tests/test_local_model_format.py index a4f79a51ce..9e1e5c7d13 100644 --- a/studio/backend/tests/test_local_model_format.py +++ b/studio/backend/tests/test_local_model_format.py @@ -118,9 +118,9 @@ def test_scan_models_dir_classifies_root_gguf_with_config(tmp_path): def test_scan_models_dir_surfaces_diffusers_pipeline_folder(tmp_path): - # A diffusers PIPELINE folder (weights in component subdirs, only model_index.json at the root) - # is loadable, so the scan must surface it or it never reaches the On Device picker. Not a GGUF, - # so model_format stays None (task tagging classifies it later). + # A diffusers PIPELINE folder (weights in component subdirs, only model_index.json at the root) is + # loadable, so the scan must surface it or it never reaches the On Device picker. Not a GGUF, so + # model_format stays None. root = tmp_path / "models" pipe = root / "my-pipeline" _touch(pipe / "model_index.json") @@ -136,8 +136,8 @@ def test_scan_models_dir_surfaces_diffusers_pipeline_folder(tmp_path): def test_scan_models_dir_surfaces_root_diffusers_pipeline(tmp_path): # A scan folder can point DIRECTLY at a diffusers pipeline, which _is_model_directory rejects; - # without admitting it the scan surfaces the component subdirs as bogus models and hides the - # real pipeline. Treat the root as one model. + # without admitting it the scan surfaces the component subdirs as bogus models and hides the real + # pipeline. root = tmp_path / "my-local-pipeline" _touch(root / "model_index.json") _touch(root / "transformer" / "config.json") @@ -152,10 +152,8 @@ def test_scan_models_dir_surfaces_root_diffusers_pipeline(tmp_path): def test_scan_models_dir_surfaces_root_single_file_checkpoint(tmp_path): # A scan folder can also point DIRECTLY at a bare single-file checkpoint dir (one loose - # .safetensors, no config.json / model_index.json). The child loop admits exactly that shape - # and the images route reinterprets it via resolve_local_single_file, so the root must be - # surfaced too -- otherwise registering the model folder itself yields no On Device row while - # registering its parent works. + # .safetensors, no config.json / model_index.json). The child loop admits exactly that shape and + # the images route reinterprets it via resolve_local_single_file, so the root must be surfaced too. root = tmp_path / "qwen-image-2509" _touch(root / "qwen-image-2509.safetensors") @@ -166,8 +164,8 @@ def test_scan_models_dir_surfaces_root_single_file_checkpoint(tmp_path): def test_scan_models_dir_root_weights_do_not_hide_child_models(tmp_path): - # A stray loose .safetensors at a models ROOT must not collapse the scan to a single row and - # hide the real model subdirs: the root fallback applies only when nothing else matched. + # A stray loose .safetensors at a models ROOT must not collapse the scan to a single row and hide + # the real model subdirs: the root fallback applies only when nothing else matched. root = tmp_path / "models" _touch(root / "stray.safetensors") _touch(root / "llama" / "config.json") @@ -199,8 +197,8 @@ def _local( def test_local_task_tags_family_named_pipeline_dir(tmp_path): - # A local diffusers pipeline (top-level model_index.json) whose id resolves to a supported - # image family loads fine, so tag it so the Images picker keeps it. + # A local diffusers pipeline whose id resolves to a supported image family loads fine, so tag it + # and the Images picker keeps it. d = tmp_path / "flux-pipeline" _touch(d / "model_index.json") _touch(d / "unet" / "diffusion_pytorch_model.safetensors") @@ -213,7 +211,7 @@ def test_local_task_tags_family_named_pipeline_dir(tmp_path): def test_local_task_none_for_familyless_pipeline_dir(tmp_path): # A generically named on-device pipeline (model_index.json, no family token) is UNLOADABLE: the # Images load path resolves no family and 400s after evicting the GPU owner, so it must stay - # untagged and never be advertised. + # untagged. d = tmp_path / "my-local-pipeline" _touch(d / "model_index.json") _touch(d / "unet" / "diffusion_pytorch_model.safetensors") @@ -222,8 +220,8 @@ def test_local_task_none_for_familyless_pipeline_dir(tmp_path): def test_local_task_tags_diffusers_by_family_id(tmp_path): - # A single-file / safetensors image checkpoint ships no model_index.json; fall back - # to the model id resolving to a known diffusion family. + # A single-file / safetensors image checkpoint ships no model_index.json, so fall back to the model + # id resolving to a known diffusion family. d = tmp_path / "flux-checkpoint" _touch(d / "flux1-dev.safetensors") assert ( @@ -241,9 +239,8 @@ def test_local_task_none_for_plain_llm(tmp_path): def test_local_task_tags_video_pipeline_dir(tmp_path): - # A local diffusers pipeline whose id resolves to a VIDEO family (LTX / Wan / Hunyuan) - # must be tagged text-to-video so it surfaces in the Video On-Device picker, mirroring the - # cached-repo path -- not text-to-image, where the image loader would reject it. + # A local diffusers pipeline whose id resolves to a VIDEO family must be tagged text-to-video so it + # surfaces in the Video On-Device picker, mirroring the cached-repo path. d = tmp_path / "wan-local" _touch(d / "model_index.json") _touch(d / "transformer" / "diffusion_pytorch_model.safetensors") @@ -254,9 +251,8 @@ def test_local_task_tags_video_pipeline_dir(tmp_path): def test_local_task_tags_video_single_file_checkpoint(tmp_path): - # A video-family dir holding a bare single-file .safetensors (no model_index.json) is loadable - # (the route loads it as a single_file), so it must be tagged text-to-video and surfaced, not - # left task=null and hidden. + # A video-family dir holding a bare single-file .safetensors is loadable (the route loads it as a + # single_file), so it must be tagged text-to-video, not left task=null and hidden. d = tmp_path / "ltx-loose" _touch(d / "ltx-2.safetensors") # loose weights, no model_index.json assert ( @@ -284,11 +280,10 @@ def test_local_task_tags_video_single_file_by_checkpoint_filename(tmp_path): def test_local_task_ignores_family_token_in_parent_path(tmp_path): - # model.id is the full on-disk path for a scanned On-Device model, and the family-token - # matcher treats any path segment as a hint. A family token in a PARENT dir (e.g. - # /models/qwen-image/misc) must NOT tag an unrelated single-file as text-to-image: that - # would surface it in the Images picker and evict the GPU owner before from_single_file - # fails on the unrelated weights. Detection is scoped to the leaf name, not the raw path. + # model.id is the full on-disk path for a scanned On-Device model, and the family-token matcher + # treats any path segment as a hint. A family token in a PARENT dir must NOT tag an unrelated + # single-file as text-to-image: that would surface it in the Images picker and evict the GPU owner + # before from_single_file fails. Detection is scoped to the leaf name, not the raw path. d = tmp_path / "misc" _touch(d / "unrelated.safetensors") # one non-family single file, no model_index.json m = _local(d, id = "/models/qwen-image/misc", display_name = "misc") diff --git a/studio/backend/tests/test_middleware.py b/studio/backend/tests/test_middleware.py index ae0784eb19..5fc8f14166 100644 --- a/studio/backend/tests/test_middleware.py +++ b/studio/backend/tests/test_middleware.py @@ -217,11 +217,10 @@ class TestMaxBodyMiddleware: assert r.json()["total"] == 512 def test_diffusion_dataset_upload_in_body_passthrough(self, main_module): - # The diffusion dataset upload route lives under the protected /api/train prefix, so it - # must be in the REAL passthrough allowlist with the DB-aware + multipart-overhead cap; - # otherwise MaxBodyMiddleware would 413 near-limit batches (and ignore a raised - # max_upload_size_mb) before the handler's own get_upload_limit_bytes() check runs. It is - # matched by EXACT path (not prefix) so its JSON sub-routes keep the normal small cap. + # The diffusion dataset upload route lives under the protected /api/train prefix, so it must be in + # the REAL passthrough allowlist with the DB-aware + multipart-overhead cap; otherwise + # MaxBodyMiddleware would 413 near-limit batches before the handler's own check runs. It is matched + # by EXACT path so its JSON sub-routes keep the normal small cap. from utils.upload_limits import ( default_request_body_limit_bytes, upload_request_limit_bytes, @@ -235,11 +234,9 @@ class TestMaxBodyMiddleware: assert cap > default_request_body_limit_bytes() # not the plain default body cap def test_diffusion_dataset_json_subroutes_keep_default_cap(self, main_module): - # The exact-path passthrough must NOT sweep in the JSON sub-routes that live under the same - # /api/train/diffusion/dataset prefix (PUT .../{name}/caption/{filename}, - # POST .../import-example). A prefix match would let a large caption/import body bypass the - # default JSON cap and be buffered + parsed up to the far larger upload limit. They must - # get the plain default body cap, and never be treated as upload passthrough. + # The exact-path passthrough must NOT sweep in the JSON sub-routes under the same + # /api/train/diffusion/dataset prefix. A prefix match would let a large caption/import body bypass + # the default JSON cap and be buffered up to the far larger upload limit. from utils.upload_limits import default_request_body_limit_bytes for path in ( "/api/train/diffusion/dataset/my-set/caption/img.png", @@ -254,10 +251,9 @@ class TestMaxBodyMiddleware: ), path def test_diffusion_dataset_trailing_slash_gets_upload_cap(self, main_module): - # The trailing-slash variant reaches the middleware BEFORE the router's - # redirect_slashes 307, so it must resolve to the same passthrough + upload cap - # as the canonical path or a large upload 413s on the default /api/train cap. - # JSON sub-routes keep extra components after normalization, so they stay capped. + # The trailing-slash variant reaches the middleware BEFORE the router's redirect_slashes 307, so it + # must resolve to the same passthrough + upload cap or a large upload 413s. JSON sub-routes keep + # extra components after normalization, so they stay capped. from utils.upload_limits import ( default_request_body_limit_bytes, upload_request_limit_bytes, @@ -267,8 +263,8 @@ class TestMaxBodyMiddleware: assert main_module._get_upload_passthrough_request_max_bytes(slashed) == ( upload_request_limit_bytes() ) - # End-to-end through the middleware: a body over the default cap but under the - # upload cap passes through on both the canonical and the slashed path. + # End-to-end through the middleware: a body over the default cap but under the upload cap passes + # on both the canonical and the slashed path. app = _make_protected_app( 128, main_module, @@ -305,10 +301,9 @@ class TestMaxBodyMiddleware: ) def test_v1_surface_is_body_protected(self, main_module): - # /images/generations is mounted at both /api/inference and /v1; the /v1 alias (and every - # other /v1 POST route) must be body-capped via the /v1 blanket prefix, or an unbounded - # ImageGenerationRequest.prompt buffers outside the Studio request limit. Also confirms - # the blanket did not drop protection for the original /v1 chat/completions route. + # /images/generations is mounted at both /api/inference and /v1; the /v1 alias (and every other /v1 + # POST route) must be body-capped via the /v1 blanket prefix, or an unbounded prompt buffers + # outside the Studio request limit. Also confirms the blanket kept /v1 chat/completions protected. for path in ( "/v1/images/generations", "/v1/audio/generate", @@ -357,9 +352,8 @@ class TestMaxBodyMiddleware: assert "Content-Length" in r.json()["detail"] def test_exact_path_passthrough_does_not_cover_subroutes(self, main_module): - # The exact-path passthrough lifts the cap for the upload path itself, but a sibling - # sub-path under the same prefix stays on the small protected cap: a large JSON body to - # the sub-route is 413'd (buffered+capped), not waved through at the upload limit. + # The exact-path passthrough lifts the cap for the upload path itself, but a sibling sub-path under + # the same prefix stays on the small protected cap. app = FastAPI() app.add_middleware( main_module.MaxBodyMiddleware, diff --git a/studio/backend/tests/test_openai_images_generations_route.py b/studio/backend/tests/test_openai_images_generations_route.py index f82586a173..541012d588 100644 --- a/studio/backend/tests/test_openai_images_generations_route.py +++ b/studio/backend/tests/test_openai_images_generations_route.py @@ -45,8 +45,8 @@ def test_default_generation_params(repo_id, expected): def test_default_generation_params_specificity_ordering(): - # The "-turbo" / "-schnell" entries must win over their broader siblings; a - # reorder that broke this would silently mis-default. + # The "-turbo" / "-schnell" entries must win over their broader siblings; a reorder that broke this + # would silently mis-default. assert default_generation_params("x/Z-Image-Turbo") != default_generation_params("x/Z-Image") assert default_generation_params("x/FLUX.1-schnell") != default_generation_params( "x/FLUX.1-dev" @@ -54,8 +54,8 @@ def test_default_generation_params_specificity_ordering(): def test_default_generation_params_falls_back_to_base_repo(): - # A local-path load: repo_id is a filesystem path that names no model, so the - # resolved base repo is what identifies it (and distinguishes dev from schnell). + # A local-path load: repo_id is a filesystem path that names no model, so the resolved base repo is + # what identifies it (and distinguishes dev from schnell). assert default_generation_params("/models/my-ckpt", "black-forest-labs/FLUX.1-dev") == (28, 3.5) assert default_generation_params("/models/my-ckpt", "black-forest-labs/FLUX.1-schnell") == ( 4, @@ -110,8 +110,8 @@ class _FakeBackend: self._base_repo = base_repo # Model the native sd.cpp engine, which returns a distinct seed per image. self._native_seeds = native_seeds - # When set, generate() raises this; unload_on_generate flips is_loaded off - # first, to model the eviction/unload race vs an in-pipeline failure (OOM). + # When set, generate() raises this; unload_on_generate flips is_loaded off first, to model the + # eviction/unload race vs an in-pipeline failure (OOM). self._generate_error = generate_error self._unload_on_generate = unload_on_generate self.calls = [] @@ -223,8 +223,8 @@ def test_b64_response_shape(client): def test_local_load_uses_base_repo_for_defaults(monkeypatch): - # repo_id is a local path that names no model; base_repo identifies FLUX.1-dev, - # so the route must pick 28 steps / 3.5 guidance, not the 9/0 fallback. + # repo_id is a local path that names no model; base_repo identifies FLUX.1-dev, so the route must + # pick 28 steps / 3.5 guidance, not the 9/0 fallback. backend = _FakeBackend(repo_id = "/models/my-flux", base_repo = "black-forest-labs/FLUX.1-dev") monkeypatch.setattr(diffusion_module, "get_diffusion_backend", lambda: backend) cli, store, _save = _make_client(backend) @@ -235,9 +235,8 @@ def test_local_load_uses_base_repo_for_defaults(monkeypatch): def test_pipeline_runtime_error_is_sanitized_500(monkeypatch): - # A RuntimeError raised inside the pipeline while the model stays loaded (e.g. - # CUDA OOM, a RuntimeError subclass) must be a sanitized 500, not a 503 that - # echoes the raw exception text. + # A RuntimeError raised inside the pipeline while the model stays loaded (e.g. CUDA OOM) must be a + # sanitized 500, not a 503 that echoes the raw exception text. oom = RuntimeError("CUDA out of memory. Tried to allocate 20.00 GiB (GPU 0; 47.5 GiB total)") backend = _FakeBackend(generate_error = oom) monkeypatch.setattr(diffusion_module, "get_diffusion_backend", lambda: backend) @@ -250,8 +249,8 @@ def test_pipeline_runtime_error_is_sanitized_500(monkeypatch): def test_unload_race_returns_503(monkeypatch): - # The model is evicted/unloaded between the readiness check and the call: the - # RuntimeError with is_loaded now False is the one case that maps to 503. + # The model is evicted/unloaded between the readiness check and the call: a RuntimeError with + # is_loaded now False is the one case that maps to 503. backend = _FakeBackend( generate_error = RuntimeError("No diffusion model is loaded."), unload_on_generate = True, @@ -268,8 +267,8 @@ def test_unload_race_returns_503(monkeypatch): def test_non_runtime_pipeline_error_is_500(monkeypatch): - # A non-RuntimeError from the pipeline (the model stays loaded) must not route - # to the 503 branch (which is gated on isinstance RuntimeError) -> sanitized 500. + # A non-RuntimeError from the pipeline (the model stays loaded) must not route to the 503 branch, + # which is gated on isinstance RuntimeError, so it is a sanitized 500. backend = _FakeBackend(generate_error = ValueError("bad tensor shape")) monkeypatch.setattr(diffusion_module, "get_diffusion_backend", lambda: backend) cli, store, _save = _make_client(backend) @@ -287,8 +286,8 @@ def test_n_maps_to_batch(client): def test_batch_persists_batch_size(monkeypatch): - # n>1 must persist batch_size in each gallery record so the Studio restore path - # can replay a batch_index>0 sibling (which shares the batch's single seed). + # n>1 must persist batch_size in each gallery record so the Studio restore path can replay a + # batch_index>0 sibling (which shares the batch's single seed). backend = _FakeBackend() monkeypatch.setattr(diffusion_module, "get_diffusion_backend", lambda: backend) cli, store, _save = _make_client(backend) @@ -301,9 +300,8 @@ def test_batch_persists_batch_size(monkeypatch): def test_uses_active_engine_not_diffusers_singleton(monkeypatch): - # On a no-GPU host the loaded model lives behind the native sd_cpp engine, not the - # diffusers singleton. The route must query get_active_diffusion_engine (like - # /images/generate) or it 503s a model that is loaded and usable. + # On a no-GPU host the loaded model lives behind the native sd_cpp engine, not the diffusers + # singleton, so the route must query get_active_diffusion_engine or it 503s a usable model. active = _FakeBackend(loaded = True) # the active (e.g. sd_cpp) engine, loaded idle_diffusers = _FakeBackend(loaded = False) # diffusers singleton, empty monkeypatch.setattr(engine_router, "get_active_diffusion_engine", lambda: active) @@ -316,9 +314,8 @@ def test_uses_active_engine_not_diffusers_singleton(monkeypatch): def test_native_batch_persists_per_image_seed(monkeypatch): - # The native sd.cpp engine returns a distinct seed per image (base+index) in - # "seeds"; each gallery record must store its own seed (like /images/generate), - # not the shared base, or a restored batch_index>0 image shows the wrong seed. + # The native sd.cpp engine returns a distinct seed per image (base+index) in "seeds", so each + # gallery record must store its own seed or a restored batch_index>0 image shows the wrong one. backend = _FakeBackend(native_seeds = True) monkeypatch.setattr(engine_router, "get_active_diffusion_engine", lambda: backend) cli, store, _save = _make_client(backend) @@ -330,7 +327,7 @@ def test_native_batch_persists_per_image_seed(monkeypatch): def test_null_fields_coalesce_to_defaults(client): - # OpenAI marks n/size/response_format nullable-with-default: null -> default. + # OpenAI marks n/size/response_format nullable-with-default: null means the default. resp = _post(client, {"prompt": "p", "n": None, "size": None, "response_format": None}) assert resp.status_code == 200 assert len(resp.json()["data"]) == 1 diff --git a/studio/backend/tests/test_personalization_settings.py b/studio/backend/tests/test_personalization_settings.py index 3a5921ee69..0d3fc7b2aa 100644 --- a/studio/backend/tests/test_personalization_settings.py +++ b/studio/backend/tests/test_personalization_settings.py @@ -441,8 +441,7 @@ def test_personalization_route_roundtrip_real_shape(monkeypatch): {"id": "chat", "visible": False}, {"id": "connections", "visible": False}, ], - # Reordered and partly unpinned, so the round-trip proves order - # survives a save. + # Reordered and partly unpinned, so the round-trip proves order survives a save. "sidebarNav": [ {"id": "images", "pinned": True}, {"id": "video", "pinned": True}, diff --git a/studio/backend/tests/test_scoped_download_job.py b/studio/backend/tests/test_scoped_download_job.py index 028f650485..7d7701d1ad 100644 --- a/studio/backend/tests/test_scoped_download_job.py +++ b/studio/backend/tests/test_scoped_download_job.py @@ -45,16 +45,16 @@ def _request(**over) -> DownloadModelRequest: def test_scope_keys_apart_from_the_full_snapshot(): - # Same repo, two jobs: the scoped one must not adopt or overwrite the full snapshot's - # manifest, or the repo would read as partial against expectations it never had. + # Same repo, two jobs: the scoped one must not adopt or overwrite the full snapshot's manifest, or + # the repo would read as partial against expectations it never had. full = dl._download_job_key("black-forest-labs/FLUX.1-dev", None) scoped = dl._download_job_key("black-forest-labs/FLUX.1-dev", dl._scope_variant("diffusion")) assert full != scoped assert scoped.endswith("@diffusion") # It rides the variant slot, so it must satisfy the same validator. assert is_valid_gguf_variant("@diffusion") - # The "@" prefix is what keeps a scope out of the quant namespace: a job scoped - # "diffusion" and a (hypothetical) quant named "diffusion" stay distinct. + # The "@" prefix keeps a scope out of the quant namespace: a job scoped "diffusion" and a + # (hypothetical) quant named "diffusion" stay distinct. assert dl._download_job_key("org/m", "diffusion") != dl._download_job_key( "org/m", dl._scope_variant("diffusion") ) @@ -105,8 +105,8 @@ def test_scoped_start_spawns_a_file_scoped_worker(monkeypatch): def test_scoped_files_survive_into_the_registry(monkeypatch): - # The XET -> HTTP retry rebuilds worker args from registry metadata alone. Without the - # file list there, a retried scoped job would silently become a full snapshot. + # The XET to HTTP retry rebuilds worker args from registry metadata alone, so without the file list + # there a retried scoped job would silently become a full snapshot. captured: dict = {} real_claim = dl._registry.claim @@ -138,9 +138,9 @@ def test_files_manifest_round_trips(): def test_a_different_file_set_is_not_adopted(monkeypatch): - # Two quants of one repo are two different downloads that share the "@diffusion" slot. - # Adopting the running one made the UI wait on the wrong file set and then load a file - # that had never been fetched, so the second request is refused while the first runs. + # Two quants of one repo are two different downloads that share the "@diffusion" slot. Adopting + # the running one made the UI wait on the wrong file set and then load a file that had never been + # fetched, so the second request is refused while the first runs. monkeypatch.setattr(dl, "_reject_if_load_in_flight", lambda repo_id: None) monkeypatch.setattr(dl, "resolve_cached_repo_id_case", lambda repo, **k: repo) monkeypatch.setattr(dl, "scoped_file_blob_hashes", lambda *a, **k: frozenset()) @@ -160,8 +160,8 @@ def test_a_different_file_set_is_not_adopted(monkeypatch): assert other_files.value.status_code == 409 assert "different" in other_files.value.detail - # The same file set is still the same download: it adopts the live job as before, - # in any order and with duplicates collapsed. + # The same file set is still the same download: it adopts the live job as before, in any order and + # with duplicates collapsed. same = asyncio.run( dl.download_model_response(_request(files = [FILES[1], FILES[0], FILES[0]])) ) @@ -171,10 +171,9 @@ def test_a_different_file_set_is_not_adopted(monkeypatch): def test_the_http_retry_keeps_the_scoped_file_list_on_the_record(monkeypatch): - # The retry reclaims the same slot with replace_active, and that claim OVERWRITES the - # stored metadata. Dropping the file list there left the record claiming an empty scope, - # so the next identical scoped start compared [] against the real list and 409'd on - # "already fetching a different set of files" instead of adopting the running download. + # The retry reclaims the same slot with replace_active, and that claim OVERWRITES the stored + # metadata. Dropping the file list there left the record claiming an empty scope, so the next + # identical scoped start compared [] against the real list and 409'd instead of adopting. monkeypatch.setattr(dl, "_reject_if_load_in_flight", lambda repo_id: None) monkeypatch.setattr(dl, "resolve_cached_repo_id_case", lambda repo, **k: repo) monkeypatch.setattr(dl, "scoped_file_blob_hashes", lambda *a, **k: frozenset()) @@ -217,8 +216,8 @@ def test_the_http_retry_keeps_the_scoped_file_list_on_the_record(monkeypatch): def test_scope_key_stays_derivable_from_the_scope_alone(): - # The download manager builds this key client-side (it polls and cancels before any - # server round-trip tells it a key), so the scope name alone must produce it. + # The download manager builds this key client-side (it polls and cancels before any server + # round-trip tells it a key), so the scope name alone must produce it. assert dl._scope_variant("diffusion") == "@diffusion" assert dl._scope_variant("video") == "@video" assert dl._scope_variant(" ") is None diff --git a/studio/backend/tests/test_sd_cpp_args.py b/studio/backend/tests/test_sd_cpp_args.py index fdc5aec863..649fd08d34 100644 --- a/studio/backend/tests/test_sd_cpp_args.py +++ b/studio/backend/tests/test_sd_cpp_args.py @@ -54,8 +54,8 @@ def test_native_speed_flags(): assert native_speed_flags(None) == [] assert native_speed_flags("off") == [] assert native_speed_flags("") == [] - # default now includes conv-direct: measured ~9% faster sampling on CPU - # (z-image Q8_0, 192 threads) with identical RSS and unchanged decode. + # default now includes conv-direct: measured ~9% faster sampling on CPU (z-image Q8_0, 192 + # threads) with identical RSS and unchanged decode. assert native_speed_flags("default") == ["--diffusion-fa", "--diffusion-conv-direct"] assert native_speed_flags("max") == ["--diffusion-fa", "--diffusion-conv-direct"] with pytest.raises(ValueError): @@ -169,8 +169,8 @@ def test_build_negative_prompt_and_batch(): params = SdCppGenParams(prompt = "x", negative_prompt = "blurry") cmd = build_sd_cpp_command("/bin/sd-cli", files, params, output_path = "/o.png") assert _pair(cmd, "--negative-prompt") == "blurry" - # A CLI batch would silently drop every image after the first (the runner only - # collects the literal --output path), so the builder rejects it outright. + # A CLI batch would silently drop every image after the first (the runner only collects the literal + # --output path), so the builder rejects it outright. with pytest.raises(ValueError, match = "single-image"): build_sd_cpp_command( "/bin/sd-cli", @@ -235,8 +235,8 @@ def test_build_inpaint_adds_mask(): def test_build_inpaint_mask_without_init_img_rejected(): - # sd-cli inpaint needs a source image; a --mask with no --init-img is invalid argv, - # so the builder must reject it up front instead of emitting a doomed command. + # sd-cli inpaint needs a source image; a --mask with no --init-img is invalid argv, so the builder + # must reject it up front instead of emitting a doomed command. files = SdCppModelFiles(diffusion_model = "/m/z.gguf") params = SdCppGenParams(prompt = "x", mask = "/in/mask.png") with pytest.raises(ValueError, match = "init_img is required"): @@ -244,8 +244,7 @@ def test_build_inpaint_mask_without_init_img_rejected(): def test_build_rejects_none_prompt(): - # A None prompt must be rejected, not coerced to the literal string "None" and - # forwarded into argv. + # A None prompt must be rejected, not coerced to the literal string "None" and forwarded into argv. files = SdCppModelFiles(diffusion_model = "/m/z.gguf") with pytest.raises(ValueError, match = "prompt is required"): build_sd_cpp_command( @@ -264,8 +263,8 @@ def test_build_edit_repeats_ref_image(): def test_img2img_unset_dims_lets_sdcpp_derive_from_source(): - # img2img/inpaint/edit with dims left unset must NOT force --width/--height, - # so sd.cpp derives the size from the input image instead of resizing it to 1024. + # img2img/inpaint/edit with dims left unset must NOT force --width/--height, so sd.cpp derives the + # size from the input image instead of resizing it to 1024. files = SdCppModelFiles(diffusion_model = "/m/z.gguf") cmd = build_sd_cpp_command( "/bin/sd-cli", diff --git a/studio/backend/tests/test_sd_cpp_backend.py b/studio/backend/tests/test_sd_cpp_backend.py index 675520db0f..a21476fd65 100644 --- a/studio/backend/tests/test_sd_cpp_backend.py +++ b/studio/backend/tests/test_sd_cpp_backend.py @@ -83,10 +83,10 @@ def _loaded_backend(fam_name = "z-image", engine = None): def test_loaded_repo_ids_includes_native_companions(): - # The one-shot native engine re-reads its companion VAE / text-encoder files from the HF - # cache on every generation, so the delete-cached guard queries loaded_repo_ids() to refuse - # deleting an in-use companion repo. It must surface the committed family's VAE + text-encoder - # repos (plus the main + base repos), not just the loaded GGUF, and be empty once unloaded. + # The one-shot native engine re-reads its companion VAE / text-encoder files from the HF cache on + # every generation, so the delete-cached guard queries loaded_repo_ids() to refuse deleting an + # in-use companion repo. It must surface the committed family's VAE + text-encoder repos (plus the + # main + base repos), not just the loaded GGUF, and be empty once unloaded. b = _loaded_backend("flux.1") ids = set(b.loaded_repo_ids()) fam = detect_family("flux.1") @@ -207,8 +207,8 @@ def test_asset_specs_cover_required_files(fam_name, expect_kinds): def test_asset_specs_flux2_klein_selects_encoder_by_variant(): - # FLUX.2-klein 4B pairs with Qwen3-4B, 9B with Qwen3-8B; the encoder must be chosen from the - # load identity, not the family's single default (a mismatched encoder fails deep in sd-cli). + # FLUX.2-klein 4B pairs with Qwen3-4B, 9B with Qwen3-8B, so the encoder must be chosen from the + # load identity, not the family default (a mismatched encoder fails deep in sd-cli). b = SdCppDiffusionBackend(engine = _FakeEngine()) fam = detect_family("flux.2-klein") @@ -323,9 +323,8 @@ def test_generate_progress_tracks_parsed_steps(): def test_generate_publishes_progress_before_lora_resolution(monkeypatch): - # LoRA resolution runs during pre-generate setup while _generate_lock is held, so a progress - # probe in that window must read ACTIVE; _gen is published before it, mirroring the diffusers - # path. + # LoRA resolution runs during pre-generate setup while _generate_lock is held, so a progress probe + # in that window must read ACTIVE; _gen is published before it, mirroring the diffusers path. from core.inference import diffusion_lora eng = _FakeEngine() @@ -371,9 +370,8 @@ def test_begin_load_requires_gguf_filename(): def test_begin_load_resolves_family_from_filename_only(monkeypatch): - # A local .gguf pick whose family keyword lives only in the basename (parent dir - # carries none) must resolve via the same filename fallback the route validated - # with -- not dead-end with "Could not infer" on a native (no-GPU) host. + # A local .gguf pick whose family keyword lives only in the basename must resolve via the same + # filename fallback the route validated with, not dead-end with "Could not infer" on a native host. b = SdCppDiffusionBackend(engine = _FakeEngine()) monkeypatch.setattr(b, "_run_load", lambda **kwargs: None) # skip the download thread b.begin_load("/models/gguf-store", gguf_filename = "Z-Image-Turbo-Q4_K_M.gguf") @@ -402,8 +400,8 @@ def test_unload_clears_state_and_signals_cancel(): def test_status_reports_offload_when_flags_active(): - # status must reflect the offload flags actually passed to sd-cli, not always "none", - # so a balanced/low_vram (or cpu_offload) load is verifiable. + # status must reflect the offload flags actually passed to sd-cli, not always "none", so a + # balanced/low_vram (or cpu_offload) load is verifiable. b = _loaded_backend() # No flags (CPU default) -> none. assert b.status()["offload_policy"] == "none" and b.status()["cpu_offload"] is False @@ -422,10 +420,9 @@ def test_status_reports_offload_when_flags_active(): def test_run_load_cancels_and_waits_for_inflight_generation(monkeypatch): - # A generation that started during the asset download is still running against the OLD - # model. _run_load must cancel it AND wait on _generate_lock before committing the new - # state, or a stale sd-cli run finishes afterward and persists an image from the previous - # model once the new load reports ready. + # A generation that started during the asset download is still running against the OLD model. + # _run_load must cancel it AND wait on _generate_lock before committing the new state, or a stale + # sd-cli run persists an image from the previous model once the new load reports ready. b = SdCppDiffusionBackend(engine = _FakeEngine()) fam = detect_family("z-image") monkeypatch.setattr(b, "_asset_specs", lambda *a, **k: []) @@ -435,8 +432,8 @@ def test_run_load_cancels_and_waits_for_inflight_generation(monkeypatch): "_fetch_assets", lambda *a, **k: {"diffusion_model": "/m/z.gguf", "vae": "/m/vae.sft", "llm": "/m/llm.sft"}, ) - # Avoid importing torch from the worker thread (its first import deadlocks off the main - # thread -- a test artifact, not a production path); the device only needs to be CPU here. + # Avoid importing torch from the worker thread (its first import deadlocks off the main thread, a + # test artifact); the device only needs to be CPU here. monkeypatch.setattr( bk, "resolve_diffusion_device_target", lambda: types.SimpleNamespace(device = "cpu") ) @@ -461,8 +458,8 @@ def test_run_load_cancels_and_waits_for_inflight_generation(monkeypatch): b._generate_lock.acquire() # simulate the live denoise holding _generate_lock try: threading.Thread(target = _load, daemon = True).start() - # The commit must block behind the live generation and not publish the new state, - # but must already have signalled the in-flight cancel. + # The commit must block behind the live generation and not publish the new state, but must already + # have signalled the in-flight cancel. assert not committed.wait(0.5) assert b._state is None assert cancel.is_set() @@ -498,8 +495,8 @@ def test_resolve_backend_falls_back_to_oneshot_without_server(monkeypatch): def test_resolve_backend_cached_fallback_engine_does_not_pin_oneshot(monkeypatch): - # A lazily cached fallback engine (NOT an explicit injection) must not force one-shot: - # once a server is available again, the next load can use it. + # A lazily cached fallback engine (NOT an explicit injection) must not force one-shot: once a + # server is available again, the next load can use it. b = SdCppDiffusionBackend() # no injected engine b._engine = _FakeEngine() # simulate a prior lazy one-shot fallback caching the engine monkeypatch.setattr(bk, "find_sd_server_binary", lambda: "/x/sd-server") @@ -571,8 +568,8 @@ def test_server_generate_uses_one_request_for_whole_batch(monkeypatch): def test_server_generate_splits_batches_above_server_limit(monkeypatch): - # A batch above the server's per-job limit is chunked (the one-shot path did these - # image-by-image); each chunk gets a timeout proportional to its image count. + # A batch above the server's per-job limit is chunked (the one-shot path did these image-by-image); + # each chunk gets a timeout proportional to its image count. b = SdCppDiffusionBackend() servers: list = [] _run_server_load(monkeypatch, b, servers) @@ -591,8 +588,8 @@ def test_server_generate_splits_batches_above_server_limit(monkeypatch): def test_server_generate_masks_large_seed(monkeypatch): - # sd.cpp's image seed is signed int64; a larger explicit seed must be masked before it - # reaches the server (the request model / diffusers accept up to 2**64 - 1). + # sd.cpp's image seed is signed int64, so a larger explicit seed must be masked before it reaches + # the server (the request model / diffusers accept up to 2**64 - 1). b = SdCppDiffusionBackend() servers: list = [] _run_server_load(monkeypatch, b, servers) @@ -684,7 +681,7 @@ def test_server_start_failure_falls_back_to_oneshot(monkeypatch): # A present-but-broken sd-server must not fail the load when sd-cli works. b = SdCppDiffusionBackend() monkeypatch.setattr(bk, "find_sd_server_binary", lambda: "/x/sd-server") - # Probe passes; the failure we exercise here is in start(), not the up-front probe. + # The probe passes; the failure exercised here is in start(), not the up-front probe. monkeypatch.setattr(bk, "_server_binary_runnable", lambda *_a, **_k: True) class _BadServer: @@ -727,8 +724,8 @@ def test_server_start_failure_falls_back_to_oneshot(monkeypatch): def test_run_load_redacts_paths_in_progress_error(monkeypatch): - # A load failure surfaced via load_progress() must run through redact_native_paths, the - # same scrub the diffusers load path applies, so a registered native path can't leak. + # A load failure surfaced via load_progress() must run through redact_native_paths, the same scrub + # the diffusers load path applies, so a registered native path can't leak. from utils import native_path_leases as npl secret_root = "/managed/native/root" @@ -811,9 +808,9 @@ def test_generate_oneshot_applies_loras_via_prompt_tags(monkeypatch): def test_generate_server_stages_loras_and_sends_structured_field(monkeypatch, tmp_path): - # Server-mode LoRA rides the structured `lora` request field (the sdcpp API ignores - # prompt tags): adapters staged into the server's --lora-model-dir, referenced - # by their path relative to it + the validated multiplier. + # Server-mode LoRA rides the structured `lora` request field (the sdcpp API ignores prompt + # tags): adapters staged into the server's --lora-model-dir, referenced by their path relative to + # it + the validated multiplier. from pathlib import Path as _P from core.inference import diffusion_lora as dl @@ -843,8 +840,8 @@ def test_generate_rejects_loras_on_unsupported_family(monkeypatch): def test_generate_zero_weight_loras_are_noop(monkeypatch): - # weight-0 rows are dropped BEFORE the support gate, so a request carrying only disabled - # adapters stays a no-op even on a family where native LoRA is unsupported. + # weight-0 rows are dropped BEFORE the support gate, so a request carrying only disabled adapters + # stays a no-op even on a family where native LoRA is unsupported. eng = _FakeEngine() b = _loaded_backend(engine = eng) _patch_lora(monkeypatch, [], supported = False) # would raise if the gate were reached @@ -854,9 +851,8 @@ def test_generate_zero_weight_loras_are_noop(monkeypatch): def test_generate_rejects_controlnet_on_native_engine(): - # ControlNet is diffusers-only. The route passes `controlnet` to whichever engine is - # active, so the native backend must reject it with a clean ValueError (-> 400) rather - # than TypeError on an unexpected kwarg (-> opaque 500). + # ControlNet is diffusers-only. The route passes `controlnet` to whichever engine is active, so the + # native backend must reject it with a clean ValueError (400) rather than TypeError (500). b = _loaded_backend(engine = _FakeEngine()) with pytest.raises(ValueError, match = "ControlNet is not yet supported on the native"): b.generate(prompt = "x", steps = 4, seed = 1, controlnet = ("id", "img", "canny", 1.0, 0.0, 1.0)) @@ -864,9 +860,9 @@ def test_generate_rejects_controlnet_on_native_engine(): @pytest.mark.parametrize("cn_strength", [0, 0.0, None]) def test_generate_treats_zero_strength_controlnet_as_disabled(cn_strength): - # strength 0 (or None) disables ControlNet -- the diffusers path treats it as plain - # txt2img and the request model documents it -- so a strength-0 spec must succeed on the - # native engine too, not 400. Only a genuinely active (strength > 0) ControlNet is rejected. + # strength 0 (or None) disables ControlNet -- the diffusers path treats it as plain txt2img -- so a + # strength-0 spec must succeed on the native engine too. Only a genuinely active ControlNet is + # rejected. eng = _FakeEngine() b = _loaded_backend(engine = eng) out = b.generate( @@ -879,8 +875,8 @@ def test_generate_treats_zero_strength_controlnet_as_disabled(cn_strength): def test_generate_rejects_image_conditioned_on_native_engine(): - # img2img / inpaint / reference / upscale are likewise diffusers-only; a direct API call - # with an init image on the native engine gets a clean ValueError, not a silent txt2img. + # img2img / inpaint / reference / upscale are likewise diffusers-only; a direct API call with an + # init image on the native engine gets a clean ValueError, not a silent txt2img. b = _loaded_backend(engine = _FakeEngine()) with pytest.raises(ValueError, match = "not yet supported on the native"): b.generate(prompt = "x", steps = 4, seed = 1, init_image = "data:image/png;base64,AAAA") diff --git a/studio/backend/tests/test_sd_cpp_engine.py b/studio/backend/tests/test_sd_cpp_engine.py index 8fee117b3a..94ace6784d 100644 --- a/studio/backend/tests/test_sd_cpp_engine.py +++ b/studio/backend/tests/test_sd_cpp_engine.py @@ -115,8 +115,8 @@ def test_find_server_path_fallback(tmp_path, monkeypatch): def test_find_server_not_confused_with_sd_cli(tmp_path, monkeypatch): - # A tree that has only sd-cli must NOT be reported as an sd-server (and vice versa), - # so the backend correctly falls back to one-shot when only the CLI is present. + # A tree that has only sd-cli must NOT be reported as an sd-server (and vice versa), so the backend + # correctly falls back to one-shot when only the CLI is present. _clear_server_env(monkeypatch) root = tmp_path / "sdcpp" (root / "build" / "bin").mkdir(parents = True) @@ -132,8 +132,8 @@ def test_find_server_not_confused_with_sd_cli(tmp_path, monkeypatch): def test_engine_unavailable_when_no_binary(monkeypatch): - # Force the "no binary anywhere" condition so the test is hermetic even on a host - # that happens to have sd-cli installed. + # Force the "no binary anywhere" condition so the test is hermetic even on a host that happens to + # have sd-cli installed. monkeypatch.setattr(eng, "find_sd_cpp_binary", lambda: None) e = SdCppEngine(binary = None) assert e.is_available() is False @@ -170,8 +170,8 @@ def test_runtime_env_prepends_binary_dir_to_lib_path(): def test_runtime_env_scrubs_native_path_lease_secret(monkeypatch): - # The sd-cli child is an external process and must never receive the native-path - # lease secret; every launch (version + generate/upscale) funnels through runtime_env. + # The sd-cli child is an external process and must never receive the native-path lease secret; + # every launch funnels through runtime_env. monkeypatch.setenv("UNSLOTH_STUDIO_NATIVE_PATH_LEASE_SECRET", "top-secret") from_os = runtime_env("/opt/sdcpp/bin/sd-cli") assert "UNSLOTH_STUDIO_NATIVE_PATH_LEASE_SECRET" not in from_os @@ -189,10 +189,9 @@ def test_runtime_env_handles_missing_lib_path(): def test_terminate_reaps_killed_child(): - # Cancellation/timeout paths call _terminate then immediately raise, so it must - # reap the killed child itself -- otherwise a burst of image cancellations leaves - # zombies until a later Popen cleanup. After _terminate the returncode is set - # (the child has been waited on), so nothing lingers. + # Cancellation/timeout paths call _terminate then immediately raise, so it must reap the killed + # child itself, otherwise a burst of image cancellations leaves zombies. After _terminate the + # returncode is set, so nothing lingers. import subprocess proc = subprocess.Popen( [sys.executable, "-c", "import time; time.sleep(30)"], @@ -324,8 +323,8 @@ def test_generate_raises_when_no_output_despite_success(tmp_path, monkeypatch): def test_generate_does_not_return_stale_preexisting_output(tmp_path, monkeypatch): - # A leftover file at the target path must not satisfy the post-run output check - # when the run itself produced nothing: the target is cleared before the run. + # A leftover file at the target path must not satisfy the post-run output check when the run itself + # produced nothing: the target is cleared before the run. e = _engine(tmp_path) out = tmp_path / "img.png" out.write_bytes(b"stale") diff --git a/studio/backend/tests/test_sd_cpp_install.py b/studio/backend/tests/test_sd_cpp_install.py index 2d251a723c..1b1f32799e 100644 --- a/studio/backend/tests/test_sd_cpp_install.py +++ b/studio/backend/tests/test_sd_cpp_install.py @@ -257,14 +257,14 @@ def test_install_downloads_verifies_extracts(tmp_path, monkeypatch): sd_cli = install(install_dir = tmp_path) assert sd_cli.name == "sd-cli" and sd_cli.is_file() assert not (tmp_path / name).exists() # archive cleaned up after extract - # Ownership marker lets the uninstaller delete a Studio-installed sd.cpp beside a custom - # root while keeping a user's own stable-diffusion.cpp checkout. + # The ownership marker lets the uninstaller delete a Studio-installed sd.cpp beside a custom root + # while keeping a user's own stable-diffusion.cpp checkout. assert (tmp_path / ".unsloth-studio-owned").is_file() def test_install_into_empty_dir_claims_ownership(tmp_path, monkeypatch): - # An empty (or freshly created) target may be adopted: the marker is written so the uninstaller - # can later remove the Studio-installed tree. + # An empty (or freshly created) target may be adopted: the marker is written so the uninstaller can + # later remove the Studio-installed tree. zb = _zip_with_sd_cli() _stub_release(monkeypatch, zip_bytes = zb, digest = "sha256:" + hashlib.sha256(zb).hexdigest()) empty = tmp_path / "sdcpp" @@ -274,8 +274,8 @@ def test_install_into_empty_dir_claims_ownership(tmp_path, monkeypatch): def test_install_into_nonempty_unowned_dir_is_refused(tmp_path, monkeypatch): - # A pre-existing, non-empty directory Studio did not create (e.g. a user's own checkout) must - # not be extracted into; install() refuses up front and leaves it untouched. + # A pre-existing, non-empty directory Studio did not create (e.g. a user's own checkout) must not + # be extracted into; install() refuses up front and leaves it untouched. zb = _zip_with_sd_cli() _stub_release(monkeypatch, zip_bytes = zb, digest = "sha256:" + hashlib.sha256(zb).hexdigest()) target = tmp_path / "stable-diffusion.cpp" @@ -293,8 +293,8 @@ def test_install_into_nonempty_unowned_dir_is_refused(tmp_path, monkeypatch): def test_reinstall_into_owned_dir_keeps_ownership(tmp_path, monkeypatch): - # A directory that already carries our marker (a prior Studio install / upgrade) stays owned - # even though it is now non-empty. + # A directory that already carries our marker (a prior Studio install / upgrade) stays owned even + # though it is now non-empty. zb = _zip_with_sd_cli() _stub_release(monkeypatch, zip_bytes = zb, digest = "sha256:" + hashlib.sha256(zb).hexdigest()) target = tmp_path / "stable-diffusion.cpp" @@ -315,10 +315,9 @@ def test_install_sha256_mismatch_raises_and_cleans_up(tmp_path, monkeypatch): def test_partial_install_failure_is_reclaimed_on_retry(tmp_path, monkeypatch): - # A crash AFTER extraction (here the post-extract cudart fetch raises) leaves the target - # non-empty. Because ownership is marked BEFORE the partial writes, the retry recognises the - # debris as ours and re-extracts instead of tripping the "not a Studio-managed directory" - # refusal, so native install self-heals without the user manually deleting the directory. + # A crash AFTER extraction (here the post-extract cudart fetch raises) leaves the target non-empty. + # Because ownership is marked BEFORE the partial writes, the retry recognises the debris as ours + # and re-extracts instead of tripping the "not a Studio-managed directory" refusal. zb = _zip_with_sd_cli() _stub_release(monkeypatch, zip_bytes = zb, digest = "sha256:" + hashlib.sha256(zb).hexdigest()) target = tmp_path / "sdcpp" @@ -334,7 +333,7 @@ def test_partial_install_failure_is_reclaimed_on_retry(tmp_path, monkeypatch): with pytest.raises(RuntimeError, match = "simulated interrupted"): install(install_dir = target) - # The partial install left extracted files AND the ownership marker (written before the writes). + # The partial install left extracted files AND the ownership marker. assert (target / ".unsloth-studio-owned").is_file() assert any(target.iterdir()) @@ -372,8 +371,8 @@ def test_safe_extractall_extracts_normal_members(tmp_path): def test_find_sd_cpp_binary_honors_studio_home(tmp_path, monkeypatch): - # A binary installed under a custom Studio root must be discovered without also - # setting UNSLOTH_SD_CPP_PATH (matches default_install_dir's env handling). + # A binary installed under a custom Studio root must be discovered without also setting + # UNSLOTH_SD_CPP_PATH (matching default_install_dir's env handling). from core.inference import sd_cpp_engine as eng monkeypatch.delenv("SD_CLI_PATH", raising = False) @@ -389,8 +388,8 @@ def test_find_sd_cpp_binary_honors_studio_home(tmp_path, monkeypatch): # ── Unsloth mirror: default source + the CPU/Apple asset set it publishes ───── _TAG = "master-741-484baa4" -# Exactly what unslothai/stable-diffusion.cpp's CI publishes (CPU + Apple only; GPU -# hosts run diffusers and never reach the native installer). +# Exactly what unslothai/stable-diffusion.cpp's CI publishes (CPU + Apple only; GPU hosts run +# diffusers and never reach the native installer). _MIRROR_ASSETS = [ f"sd-{_TAG}-bin-Darwin-macOS-arm64.zip", f"sd-{_TAG}-bin-Darwin-macOS-x86_64.zip", @@ -498,8 +497,8 @@ def test_install_falls_back_to_upstream_when_mirror_missing(tmp_path, monkeypatc sd_cli = install(install_dir = tmp_path) assert sd_cli.name == "sd-cli" and sd_cli.is_file() captured = capsys.readouterr() - # The repo-fallback diagnostic goes to stderr so --print-asset's stdout stays a single - # asset line; the "source ..." install note still prints to stdout. + # The repo-fallback diagnostic goes to stderr so --print-asset's stdout stays a single asset line; + # the "source ..." install note still prints to stdout. assert "falling back to leejet/stable-diffusion.cpp" in captured.err assert "source leejet/stable-diffusion.cpp" in captured.out @@ -516,9 +515,9 @@ def test_install_errors_when_neither_source_serves(tmp_path, monkeypatch): def test_mirror_windows_gpu_accel_is_no_match_not_cpu(): - # The mirror ships only a CPU win zip. An explicit --accelerator cuda/vulkan/rocm must - # NOT silently resolve to that CPU build; it returns None so install() falls back to - # upstream, which does build the accelerated asset. + # The mirror ships only a CPU win zip. An explicit --accelerator cuda/vulkan/rocm must NOT silently + # resolve to that CPU build; it returns None so install() falls back to upstream, which does build + # the accelerated asset. for accel in ("cuda", "vulkan", "rocm"): assert _mresolve("Windows", "AMD64", accel) is None # auto / cpu still resolve to the CPU build. @@ -532,8 +531,8 @@ def test_mirror_linux_gpu_accel_is_no_match_not_cpu(): def test_upstream_full_matrix_still_resolves_gpu_accel(): - # Regression guard: the "explicit GPU accel -> None on miss" change must not break - # the upstream matrix, which does publish the accelerated builds. + # Regression guard: the "explicit GPU accel -> None on miss" change must not break the upstream + # matrix, which does publish the accelerated builds. assert _resolve("Windows", "AMD64", "cuda") == "sd-master-8caa3f9-bin-win-cuda12-x64.zip" assert _resolve("Linux", "x86_64", "vulkan").endswith("x86_64-vulkan.zip") assert "rocm" in _resolve("Linux", "x86_64", "rocm") @@ -543,8 +542,8 @@ def test_upstream_full_matrix_still_resolves_gpu_accel(): def test_explicit_repo_override_equal_to_default_suppresses_fallback(tmp_path, monkeypatch): - # A user who pins UNSLOTH_SD_CPP_REPO (even to the default value) must get exactly that - # repo -- no surprise leejet substitution -- so a missing release errors instead. + # A user who pins UNSLOTH_SD_CPP_REPO (even to the default value) must get exactly that repo -- no + # surprise leejet substitution -- so a missing release errors instead. _stub_two_repos( monkeypatch, mirror_serves = False, upstream_serves = True, zip_bytes = b"", digest = "" ) @@ -557,8 +556,8 @@ def test_explicit_repo_override_equal_to_default_suppresses_fallback(tmp_path, m def test_pinned_tag_prefers_upstream_pin_over_mirror_latest(tmp_path, monkeypatch, capsys): - # The mirror lacks the pinned tag but has a newer latest; upstream has the pin. The - # pinned upstream build must win over the unpinned mirror-latest (reproducibility). + # The mirror lacks the pinned tag but has a newer latest; upstream has the pin. The pinned upstream + # build must win over the unpinned mirror-latest (reproducibility). monkeypatch.delenv("UNSLOTH_SD_CPP_REPO", raising = False) monkeypatch.setenv("UNSLOTH_SD_CPP_TAG", "master-999-pinned") zb = _zip_with_sd_cli() @@ -613,8 +612,8 @@ def test_pinned_tag_prefers_upstream_pin_over_mirror_latest(tmp_path, monkeypatc def test_print_asset_uses_upstream_fallback(monkeypatch, capsys): - # A host the mirror does not build (Linux Vulkan) must print the upstream asset that a - # real install would fetch, not "no matching prebuilt". + # A host the mirror does not build (Linux Vulkan) must print the upstream asset that a real install + # would fetch, not "no matching prebuilt". monkeypatch.delenv("UNSLOTH_SD_CPP_REPO", raising = False) monkeypatch.delenv("UNSLOTH_SD_CPP_TAG", raising = False) diff --git a/studio/backend/tests/test_sd_cpp_server.py b/studio/backend/tests/test_sd_cpp_server.py index 9e040814e9..f74339fbf0 100644 --- a/studio/backend/tests/test_sd_cpp_server.py +++ b/studio/backend/tests/test_sd_cpp_server.py @@ -168,10 +168,10 @@ def test_start_becomes_ready_when_capabilities_200(patched): def test_start_fails_fast_when_process_exits(patched): - # Model load failed -> process exits before listening; start must raise with the tail. + # Model load failed, so the process exits before listening; start must raise with the tail. popen = _FakePopen(lines = ["error: bad model"], exit_code = 1) patched.setattr(srv.subprocess, "Popen", lambda *a, **k: popen) - # Capabilities never answers (connection refused) -> readiness relies on exit detection. + # Capabilities never answers (connection refused), so readiness relies on exit detection. s = _server_with( popen, _FakeClient(get = lambda url: (_ for _ in ()).throw(srv.httpx.ConnectError("refused"))) ) @@ -339,8 +339,8 @@ def test_img_gen_rejected_after_stop(patched): def test_img_gen_cancelled_before_submit_reports_cancellation(patched): - # The server was stopped for a cancel/unload before submit; with the cancel event set - # this must surface as a cancellation (route -> 409), not a generic "not running" 500. + # The server was stopped for a cancel/unload before submit; with the cancel event set this must + # surface as a cancellation (route 409), not a generic "not running" 500. popen = _FakePopen() patched.setattr(srv.subprocess, "Popen", lambda *a, **k: popen) s = _server_with(popen, _FakeClient(get = lambda url: _Resp(200, {}))) @@ -353,8 +353,8 @@ def test_img_gen_cancelled_before_submit_reports_cancellation(patched): def test_img_gen_abandons_when_cancel_not_honored(patched): - # A best-effort cancel the server ignores must not pin this call (and the generate - # lock) until natural completion: after the grace window it raises cancellation. + # A best-effort cancel the server ignores must not pin this call (and the generate lock) until + # natural completion: after the grace window it raises cancellation. patched.setattr(srv, "_CANCEL_GRACE_S", 0.0) popen = _FakePopen() cancel = threading.Event() @@ -389,16 +389,16 @@ def test_img_gen_non_dict_status_json_raises(patched): def test_decode_images_tolerates_unexpected_shapes(): - # A misbehaving/older server can return non-dict result/images/items; _decode_images - # must raise a clean "no images" rather than an AttributeError on .get(). + # A misbehaving/older server can return non-dict result/images/items, so _decode_images must raise + # a clean "no images" rather than an AttributeError on .get(). for job in ({"result": ["x"]}, {"result": {"images": "nope"}}, {"result": {"images": [1, 2]}}): with pytest.raises(RuntimeError, match = "no images"): SdCppServer._decode_images(job) def test_start_aborted_by_concurrent_stop(patched): - # A stop() during the readiness wait must abort start() promptly (without waiting out - # the startup timeout) and surface as a cancellation. + # A stop() during the readiness wait must abort start() promptly (without waiting out the startup + # timeout) and surface as a cancellation. popen = _FakePopen(lines = ["loading model"]) patched.setattr(srv.subprocess, "Popen", lambda *a, **k: popen) diff --git a/studio/backend/tests/test_training_start_offload.py b/studio/backend/tests/test_training_start_offload.py index 1b53868427..9535c22852 100644 --- a/studio/backend/tests/test_training_start_offload.py +++ b/studio/backend/tests/test_training_start_offload.py @@ -36,8 +36,8 @@ class _FakeBackend: before_spawn = None, **kwargs, ): - # The real backend runs before_spawn synchronously inside this call, so the - # thread this method runs on is the thread the blocking VRAM hook runs on. + # The real backend runs before_spawn synchronously inside this call, so the thread this method runs + # on is the thread the blocking VRAM hook runs on. self.start_thread = threading.current_thread() self.hook = before_spawn self.current_job_id = job_id @@ -75,9 +75,9 @@ def test_start_route_offloads_blocking_start(monkeypatch): def test_backend_start_guard_blocks_overlapping_starts(): - # With the route offloaded to worker threads, two overlapping /train/start requests - # can reach TrainingBackend.start_training concurrently; the compare-and-set - # _start_in_progress flag must let exactly one of them spawn. + # With the route offloaded to worker threads, two overlapping /train/start requests can reach + # TrainingBackend.start_training concurrently, so the compare-and-set _start_in_progress flag must + # let exactly one of them spawn. from core.training.training import TrainingBackend backend = TrainingBackend() @@ -103,8 +103,8 @@ def test_backend_start_guard_blocks_overlapping_starts(): t = threading.Thread(target = _first, daemon = True) t.start() assert first_entered.wait(timeout = 5.0) - # Second start while the first is still inside the impl: refused by the guard, - # without ever entering the impl. + # A second start while the first is still inside the impl: refused by the guard, without ever + # entering the impl. results["second"] = backend.start_training("job-b") release_first.set() t.join(timeout = 5.0) diff --git a/studio/backend/tests/test_video_backend.py b/studio/backend/tests/test_video_backend.py index 4fb3c74d69..20c1179a3d 100644 --- a/studio/backend/tests/test_video_backend.py +++ b/studio/backend/tests/test_video_backend.py @@ -67,8 +67,8 @@ class _FakePipe: def enable_vae_tiling(self) -> None: self.vae.tiled = True - # Explicit signature so generate()'s signature-gated kwargs (negative_prompt, - # frame_rate, callback) actually engage; **kwargs would defeat the gates. + # Explicit signature so generate()'s signature-gated kwargs actually engage; **kwargs would defeat + # the gates. def __call__( self, *, @@ -125,10 +125,9 @@ class _FakeTransformer: # ── Wan2.2 fakes: a per-DiT trackable transformer so the dual-DiT optimisation tests can assert -# speed / cache / attention engaged on BOTH experts, plus two pipeline fakes -- single-DiT -# (TI2V-5B) and dual-DiT MoE (A14B). The MoE __call__ carries guidance_scale_2 so the cfg2 -# signature-gate exercises; the single-DiT __call__ omits it so the gate proves it is NOT -# threaded there. +# speed / cache / attention engaged on BOTH experts, plus single-DiT (TI2V-5B) and dual-DiT MoE +# (A14B) pipeline fakes. The MoE __call__ carries guidance_scale_2 so the cfg2 signature-gate +# exercises; the single-DiT __call__ omits it so the gate proves it is NOT threaded there. class _FakeWanDiT: @@ -294,17 +293,16 @@ class _FakeWanPipelineSingle: return _FakeWanPipeMoE() if moe else _FakeWanPipeSingle() -# ── HunyuanVideo-1.5 fakes: __call__ has NO guidance kwarg and NO callback_on_step_end -# (matching pipeline_hunyuan_video1_5.py in diffusers 0.39), a guider object carries the CFG -# scale, and the denoise loop drives scheduler.step -- so the guider write and the scheduler-wrap -# progress/cancel paths exercise. +# ── HunyuanVideo-1.5 fakes: __call__ has NO guidance kwarg and NO callback_on_step_end (matching +# diffusers 0.39), a guider object carries the CFG scale, and the denoise loop drives +# scheduler.step, so the guider write and the scheduler-wrap progress/cancel paths exercise. class _FakeHV15Scheduler: def __init__(self) -> None: self.calls = 0 - # Test hook fired from the ORIGINAL step (i.e. inside the wrapped call), - # letting a test cancel mid-denoise exactly as a user request would land. + # Test hook fired from the ORIGINAL step (inside the wrapped call), letting a test cancel + # mid-denoise exactly as a user request would land. self.on_step = None def step(self, *args, **kwargs): @@ -400,8 +398,8 @@ def fake_runtime(monkeypatch): monkeypatch.setitem(sys.modules, "torch", torch) monkeypatch.setitem(sys.modules, "diffusers", diffusers) monkeypatch.setattr("core.inference.video.clear_gpu_cache", lambda: None) - # MP4 encode needs real frames + PyAV; the backend contract under test is the - # byte handoff, so stub the encoder. + # MP4 encode needs real frames + PyAV; the contract under test is the byte handoff, so stub the + # encoder. monkeypatch.setattr( VideoBackend, "_encode_mp4", staticmethod(lambda frames, fps, audio, pipe: b"MP4") ) @@ -478,9 +476,9 @@ def test_validate_gates_base_repo_and_local_paths(tmp_path): def test_validate_rejects_kind_extension_mismatch(tmp_path): backend = VideoBackend() - # model_kind single_file with a .gguf file, or gguf with a non-.gguf file, must be rejected - # BEFORE the GPU handoff (mirrors the image loader), not fail in the wrong single-file loader - # after the route evicted the resident model. + # model_kind single_file with a .gguf file, or gguf with a non-.gguf file, must be rejected BEFORE + # the GPU handoff (mirrors the image loader), not fail in the wrong single-file loader after the + # route evicted the resident model. with pytest.raises(ValueError, match = "needs model_kind 'gguf'"): backend.validate_load_request( "unsloth/LTX-2.3-GGUF", @@ -506,11 +504,10 @@ def test_validate_rejects_kind_extension_mismatch(tmp_path): def test_validate_rejects_local_file_suffix_kind_mismatch(tmp_path): backend = VideoBackend() - # A local FILE is handed straight to the gguf/single_file loader: _resolve_checkpoint_path - # returns the file itself, IGNORING gguf_filename, so the file's OWN suffix must match the - # kind. A .gguf picked as single_file (or a .safetensors picked as gguf) slips past the - # gguf_filename suffix checks, so reject it HERE, before the route evicts the resident GPU - # owner and fails deep in from_single_file / the GGUF reader. + # A local FILE is handed straight to the gguf/single_file loader: _resolve_checkpoint_path returns + # the file itself, IGNORING gguf_filename, so the file's OWN suffix must match the kind. Such a + # mismatch slips past the gguf_filename suffix checks, so reject it HERE, before the route evicts + # the resident GPU owner. gguf_file = tmp_path / "ltx.gguf" gguf_file.write_bytes(b"weights") safetensors_file = tmp_path / "ltx.safetensors" @@ -552,9 +549,8 @@ def test_validate_rejects_local_file_suffix_kind_mismatch(tmp_path): def test_validate_rejects_windows_shaped_missing_checkpoint(tmp_path): backend = VideoBackend() - # A missing Windows-shaped local pick (backslash path, or a C:/ drive path) must fail HERE, - # not be treated as a Hub repo and fail after the route evicts the resident GPU owner. Mirrors - # the image loader's is_absolute()/backslash path-shaped check. + # A missing Windows-shaped local pick (backslash path, or a C:/ drive path) must fail HERE, not be + # treated as a Hub repo and fail after the route evicts the resident GPU owner. with pytest.raises(ValueError, match = "does not exist"): backend.validate_load_request( "C:\\models\\ltx.gguf", @@ -575,9 +571,8 @@ def test_validate_rejects_local_pipeline_without_model_index(tmp_path): d = tmp_path / "ltx-local" (d / "transformer").mkdir(parents = True) (d / "transformer" / "diffusion_pytorch_model.safetensors").write_bytes(b"x") - # A local dir resolved to a video family but missing model_index.json is not a loadable - # diffusers pipeline; it must fail preflight BEFORE the route evicts the resident model, - # mirroring the image loader's local-pipeline shape check. + # A local dir resolved to a video family but missing model_index.json is not a loadable diffusers + # pipeline; it must fail preflight BEFORE the route evicts the resident model. with pytest.raises(ValueError, match = "model_index.json"): backend.validate_load_request(str(d), family_override = "ltx-2") # With a model_index.json it is a valid local pipeline pick and passes preflight. @@ -589,9 +584,8 @@ def test_validate_rejects_local_pipeline_without_model_index(tmp_path): def test_validate_rejects_local_file_picked_as_pipeline(tmp_path): backend = VideoBackend() # A local FILE (a bare .safetensors) sent as a pipeline is not a diffusers directory, so - # from_pretrained fails deep in the background load AFTER the route evicts the resident model. - # The preflight rejects it HERE -- gating on .exists() (not .is_dir()), mirroring the image - # loader, so it catches files as well as directories. + # from_pretrained fails deep in the background load AFTER eviction. The preflight rejects it HERE, + # gating on .exists() (not .is_dir()) so it catches files as well as directories. f = tmp_path / "ltx-2.safetensors" f.write_bytes(b"x") with pytest.raises(ValueError, match = "model_index.json"): @@ -600,10 +594,9 @@ def test_validate_rejects_local_file_picked_as_pipeline(tmp_path): def test_validate_rejects_local_base_repo_without_model_index(tmp_path): backend = VideoBackend() - # A local base_repo dir that is NOT a diffusers pipeline (no model_index.json) passes the - # any-existing-path trust check, but the base loads via from_pretrained (needs model_index), so - # reject it HERE before the route hands the GPU to VIDEO -- the pipeline-kind shape check covers - # only repo_id, and an explicit base_repo is only meaningful for a gguf/single_file load. + # A local base_repo dir that is NOT a diffusers pipeline passes the any-existing-path trust check, + # but the base loads via from_pretrained, so reject it HERE before the route hands the GPU to + # VIDEO: the pipeline-kind shape check covers only repo_id. bad_base = tmp_path / "bare-base" bad_base.mkdir() with pytest.raises(ValueError, match = "model_index.json"): @@ -626,8 +619,8 @@ def test_validate_rejects_local_base_repo_without_model_index(tmp_path): def test_validate_rejects_gguf_repo_as_pipeline(): backend = VideoBackend() - # A -GGUF repo with no quant filename resolves to the pipeline kind and would - # only fail minutes later in from_pretrained, AFTER evicting the GPU owner. + # A -GGUF repo with no quant filename resolves to the pipeline kind and would only fail minutes + # later in from_pretrained, AFTER evicting the GPU owner. with pytest.raises(ValueError, match = "pick one of its .gguf files"): backend.validate_load_request("unsloth/LTX-2.3-GGUF") with pytest.raises(ValueError, match = "pick one of its .gguf files"): @@ -638,24 +631,23 @@ def test_detect_load_family_filename_fallback(): # Repo id alone carries the family. fam = _detect_load_family("Lightricks/LTX-2", None, None) assert fam is not None and fam.name == "ltx-2" - # Repo id is opaque but the picked filename carries it: fall back to the - # combined path so validate and _run_load agree on the family. + # Repo id is opaque but the picked filename carries the family: fall back to the combined path so + # validate and _run_load agree. fam = _detect_load_family("someorg/quants", "ltx-2-19b-Q4_K_M.gguf", None) assert fam is not None and fam.name == "ltx-2" # No filename and no recognisable repo id: no family. assert _detect_load_family("someorg/quants", None, None) is None - # An explicit override resolves by name/alias and skips the filename fallback: - # a bogus override stays None even when the filename would have matched. + # An explicit override resolves by name/alias and skips the filename fallback: a bogus override + # stays None even when the filename would have matched. fam = _detect_load_family("someorg/quants", "ltx-2-19b-Q4_K_M.gguf", "ltxv") assert fam is not None and fam.name == "ltx-2" assert _detect_load_family("someorg/quants", "ltx-2-19b-Q4_K_M.gguf", "bogus") is None def test_detect_load_family_cached_hub_arch_fallback(monkeypatch): - # A CACHED HUB GGUF is admitted to the picker by its general.architecture (the cached-gguf - # listing tags it text-to-video), but an opaque repo id + renamed file carry no family token, - # so name detection misses. The local-file arch read misses too (a hub repo id is not a local - # dir), so without a cache fallback the loader 400s a SUPPORTED checkpoint the picker offered. + # A CACHED HUB GGUF is admitted to the picker by its general.architecture, but an opaque repo id + + # renamed file carry no family token, so name detection misses. The local-file arch read misses + # too, so without a cache fallback the loader 400s a SUPPORTED checkpoint the picker offered. import huggingface_hub import utils.models.gguf_metadata as gguf_meta @@ -676,8 +668,8 @@ def test_detect_load_family_cached_hub_arch_fallback(monkeypatch): monkeypatch.setattr(huggingface_hub, "try_to_load_from_cache", lambda *a, **k: None) assert _detect_load_family("someorg/opaque-quants", "model.gguf", None) is None - # A recognised-but-unsupported video arch (wan has no backend family in this build) stays None, - # so an unsupported cached pick 400s just like the local-dir case. + # A recognised-but-unsupported video arch (wan has no backend family in this build) stays None, so + # an unsupported cached pick 400s just like the local-dir case. monkeypatch.setattr( huggingface_hub, "try_to_load_from_cache", lambda *a, **k: "/fake/cache/blobs/model.gguf" ) @@ -686,8 +678,8 @@ def test_detect_load_family_cached_hub_arch_fallback(monkeypatch): ) assert _detect_load_family("someorg/opaque-quants", "model.gguf", None) is None - # The blob lives in a NON-active cache root (legacy / default): the active probe (no cache_dir) - # misses, but the per-root probe finds it, so a GGUF the picker offered from any root resolves. + # The blob lives in a NON-active cache root (legacy / default): the active probe misses, but the + # per-root probe finds it, so a GGUF the picker offered from any root resolves. import hub.utils.paths as hub_paths monkeypatch.setattr(hub_paths, "legacy_hf_cache_dir", lambda: "/fake/legacy") @@ -709,15 +701,15 @@ def test_detect_load_family_cached_hub_arch_fallback(monkeypatch): def test_loading_repo_ids_guards_in_flight_delete(): # During a background load status()["loaded"] is still False, but the target repo (+ companion - # base) is downloading, so the delete-cached guard needs loading_repo_ids to refuse deletion - # and avoid yanking blobs from under the in-flight download/assembly. + # base) is downloading, so the delete-cached guard needs loading_repo_ids to avoid yanking blobs + # from under the in-flight download. from core.inference.video import _VideoLoadingState backend = VideoBackend() assert backend.loading_repo_ids() == () # idle: nothing to guard backend._loading = _VideoLoadingState(repo_id = "org/ckpt", base_repo = "Lightricks/LTX-2") assert set(backend.loading_repo_ids()) == {"org/ckpt", "Lightricks/LTX-2"} - # An errored load is no longer in flight -> the files are safe to delete. + # An errored load is no longer in flight, so the files are safe to delete. backend._loading = _VideoLoadingState( repo_id = "org/ckpt", base_repo = "Lightricks/LTX-2", error = "boom" ) @@ -760,11 +752,10 @@ def test_load_generate_unload_gguf(fake_runtime, tmp_path): def test_load_holds_generate_lock_across_placement(fake_runtime, tmp_path, monkeypatch): - # The video load must hold _generate_lock across GPU placement (apply_memory_plan) so an - # unload / arbiter eviction -- which barriers on _generate_lock before freeing -- can't hand - # the GPU to another backend while a multi-GB pipeline is still being moved onto it (mirrors - # the image backend). Verify unload() blocks until placement releases the lock, and the - # superseded load then aborts without committing. + # The video load must hold _generate_lock across GPU placement (apply_memory_plan) so an unload / + # arbiter eviction -- which barriers on _generate_lock before freeing -- can't hand the GPU to + # another backend while a multi-GB pipeline is still being moved onto it. Verify unload() blocks + # until placement releases the lock, and the superseded load then aborts without committing. import threading from core.inference import video as video_mod @@ -793,7 +784,7 @@ def test_load_holds_generate_lock_across_placement(fake_runtime, tmp_path, monke load_thread.start() assert placement_started.wait(timeout = 5), "load never reached placement" - # Placement is in flight, holding _generate_lock. unload() must block on its barrier. + # Placement is in flight, holding _generate_lock, so unload() must block on its barrier. unload_done = [] def do_unload(): @@ -805,8 +796,8 @@ def test_load_holds_generate_lock_across_placement(fake_runtime, tmp_path, monke unload_thread.join(timeout = 0.5) assert not unload_done, "unload() returned while placement still held _generate_lock (the race)" - # Release placement; unload()'s barrier then passes and its teardown runs strictly AFTER - # the load's placement+commit -- never concurrently -- so no two pipelines are ever resident. + # Release placement; unload()'s barrier then passes and its teardown runs strictly AFTER the + # load's placement+commit, so no two pipelines are ever resident. release_placement.set() unload_thread.join(timeout = 5) load_thread.join(timeout = 5) @@ -816,10 +807,10 @@ def test_load_holds_generate_lock_across_placement(fake_runtime, tmp_path, monke def test_load_records_engaged_speed_optims(fake_runtime, tmp_path, monkeypatch): - # Regression: the load tail once re-ran the already-filtered speed_optims tuple through - # ``.items()`` as if it were still the raw applied dict, so every real-GPU load (where at least - # channels_last engages) crashed with 'tuple' object has no attribute 'items'. Fake runtime - # forces every optim False, so this only reproduces when one is made to engage. + # Regression: the load tail once re-ran the already-filtered speed_optims tuple through ``.items()`` + # as if it were still the raw applied dict, so every real-GPU load crashed with 'tuple' object has + # no attribute 'items'. The fake runtime forces every optim False, so this only reproduces when + # one is made to engage. from core.inference import video as video_mod monkeypatch.setattr( @@ -847,16 +838,16 @@ def test_generate_defaults_from_variant(fake_runtime, tmp_path): call = backend._state.pipe.last_kwargs assert call["num_inference_steps"] == 8 assert call["guidance_scale"] == 1.0 - # At the distilled default step count the calibrated ltx_core curve is passed verbatim - # (the DiT was trained against it; the scheduler's own 8-step spacing lands far off). + # At the distilled default step count the calibrated ltx_core curve is passed verbatim (the DiT was + # trained against it; the scheduler's own 8-step spacing lands far off). from core.inference.video_ltx2 import LTX23_DISTILLED_SIGMAS assert call["sigmas"] == list(LTX23_DISTILLED_SIGMAS) def test_generate_distilled_custom_steps_keep_scheduler_spacing(fake_runtime, tmp_path): - # A non-default step count on the distilled DiT has no calibrated list; the scheduler's - # spacing applies and no sigmas kwarg is injected. + # A non-default step count on the distilled DiT has no calibrated list, so the scheduler's spacing + # applies and no sigmas kwarg is injected. (tmp_path / "ltx-2.3-22b-distilled-1.1-Q4_K_M.gguf").write_bytes(b"w") backend = VideoBackend() backend.load_pipeline( @@ -872,8 +863,8 @@ def test_generate_distilled_custom_steps_keep_scheduler_spacing(fake_runtime, tm def test_generate_dev_base_never_gets_distilled_sigmas(fake_runtime, tmp_path): - # The dev/base DiT uses the resolution-shifted scheduler spacing even at 8 steps: the - # calibrated list is distilled-only. + # The dev/base DiT uses the resolution-shifted scheduler spacing even at 8 steps: the calibrated + # list is distilled-only. (tmp_path / "ltx-2.3-22b-dev-Q4_K_M.gguf").write_bytes(b"w") backend = VideoBackend() backend.load_pipeline( @@ -889,8 +880,8 @@ def test_generate_dev_base_never_gets_distilled_sigmas(fake_runtime, tmp_path): def test_ltx23_verbatim_sigmas_restores_scheduler_config(): - # The context manager must neutralise exactly the transforms that distort explicit - # sigmas and put the original values back afterwards, even on error. + # The context manager must neutralise exactly the transforms that distort explicit sigmas and put + # the original values back afterwards, even on error. from core.inference.video_ltx2 import ltx23_verbatim_sigmas class _Cfg(dict): @@ -919,11 +910,10 @@ def test_ltx23_verbatim_sigmas_restores_scheduler_config(): def test_generate_resets_step_cache_only_when_engaged(fake_runtime, tmp_path): - # FBCache residuals live on the long-lived DiT(s) and survive a generation, so the next clip - # at a new resolution would crash on stale state. generate must reset them when a cache is - # engaged (diffusers 0.39 exposes _reset_stateful_cache on the transformer; reset_stateful_hooks - # only on the HookRegistry) and must not touch an uncached load. transformer_2 (the Wan dual - # expert) resets too when present. + # FBCache residuals live on the long-lived DiT(s) and survive a generation, so the next clip at a + # new resolution would crash on stale state. generate must reset them when a cache is engaged + # (diffusers 0.39 exposes _reset_stateful_cache on the transformer) and must not touch an uncached + # load. transformer_2 (the Wan dual expert) resets too when present. import dataclasses (tmp_path / "ltx-2.3-22b-distilled-1.1-Q4_K_M.gguf").write_bytes(b"w") @@ -944,16 +934,16 @@ def test_generate_resets_step_cache_only_when_engaged(fake_runtime, tmp_path): # No cache engaged -> no reset. backend.generate(prompt = "a sloth") assert resets == [] - # Cache engaged -> both resident DiTs reset before the pipe call. + # Cache engaged: both resident DiTs reset before the pipe call. backend._state = dataclasses.replace(backend._state, transformer_cache = "fbcache") backend.generate(prompt = "a sloth") assert resets == ["transformer", "transformer_2"] def test_is_ltx23_checkpoint_gguf(monkeypatch, tmp_path): - # diffusers maps every LTX-2 single file to the 2.0 config; a 2.3 checkpoint (9-row modulation - # tables in the header) must be detected so the loader routes to the full 2.3 assembly. A 2.0 - # header must not, and an unreadable header falls back to the stock path (False), never raise. + # diffusers maps every LTX-2 single file to the 2.0 config, so a 2.3 checkpoint (9-row modulation + # tables in the header) must be detected and routed to the full 2.3 assembly. A 2.0 header must + # not, and an unreadable header falls back to the stock path (False), never raises. from core.inference.video_ltx2 import is_ltx23_checkpoint def _reader_for(shapes): @@ -1054,12 +1044,11 @@ def test_ltx23_split_and_variant(tmp_path): def test_ltx23_scaled_fp8_refused(monkeypatch, tmp_path): - # The Lightricks fp8 files carry .weight_scale/.input_scale companions; a plain dtype cast - # would corrupt them, so the loader must refuse with a pointer to the supported GGUF path. + # The Lightricks fp8 files carry .weight_scale/.input_scale companions, so a plain dtype cast would + # corrupt them and the loader must refuse with a pointer to the supported GGUF path. from core.inference import video_ltx2 - # Stub the module tree so this also runs under the CI sim, which blocks the - # real diffusers import. + # Stub the module tree so this also runs under the CI sim, which blocks the real diffusers import. diffusers = types.ModuleType("diffusers") diffusers.LTX2Pipeline = object loaders = types.ModuleType("diffusers.loaders") @@ -1091,16 +1080,15 @@ def test_generate_without_load_raises(fake_runtime): def test_generate_progress_and_cancel_idle(fake_runtime): backend = VideoBackend() - # Idle shape carries the image-endpoint-compatible aliases (total_steps / fraction) - # so one poller works against both generate-progress APIs. + # The idle shape carries the image-endpoint-compatible aliases (total_steps / fraction) so one + # poller works against both generate-progress APIs. assert backend.generate_progress() == {"active": False, "total_steps": 0, "fraction": 0.0} assert backend.cancel_generate() is False def test_generate_progress_derives_total_steps_and_fraction(fake_runtime): - # A mid-denoise poll must report fraction = step / total under BOTH field names: - # a client polling the image API's shape against video used to read - # total_steps=null / fraction=0 while step advanced. + # A mid-denoise poll must report fraction = step / total under BOTH field names: a client polling + # the image API's shape against video used to read total_steps=null / fraction=0. backend = VideoBackend() backend._gen = {"active": True, "phase": "denoise", "step": 5, "total": 20} gen = backend.generate_progress() @@ -1109,10 +1097,9 @@ def test_generate_progress_derives_total_steps_and_fraction(fake_runtime): def test_failed_background_generate_retains_terminal_error(fake_runtime, tmp_path, monkeypatch): - # A page mounted AFTER a background job failed (browser reload during a minutes-long - # generation) reads the outcome from this retained terminal record -- its only diagnosis, - # since nothing else survives the reload. So a failure must stay pollable as - # active=False + phase="failed" + error, repeatedly, until the next job replaces it. + # A page mounted AFTER a background job failed (a browser reload during a minutes-long generation) + # reads the outcome from this retained terminal record -- its only diagnosis. So a failure must + # stay pollable as active=False + phase="failed" + error, repeatedly, until the next job. backend = VideoBackend() _load_gguf(backend, tmp_path) @@ -1147,9 +1134,9 @@ def test_failed_background_generate_retains_terminal_error(fake_runtime, tmp_pat def test_cache_bytes_counts_incomplete_blobs(fake_runtime, tmp_path, monkeypatch): - # scan_cache_dir skips in-flight *.incomplete blobs, so the old counter froze at the - # last completed blob for the whole multi-GB shard pull. The walk must count both, - # without double-counting snapshot symlinks. + # scan_cache_dir skips in-flight *.incomplete blobs, so the old counter froze at the last completed + # blob for the whole multi-GB shard pull. The walk must count both, without double-counting + # snapshot symlinks. import core.inference.video as video_mod repo_dir = tmp_path / "models--Wan-AI--Wan2.2-TI2V-5B-Diffusers" @@ -1160,8 +1147,8 @@ def test_cache_bytes_counts_incomplete_blobs(fake_runtime, tmp_path, monkeypatch snap = repo_dir / "snapshots" / "deadbeef" snap.mkdir(parents = True) (snap / "model_index.json").symlink_to(blobs / "aa11") # must not double-count - # The live cache root, not huggingface_hub's import-time constant: the counter follows - # a mid-session cache-folder change, which the constant does not. + # The live cache root, not huggingface_hub's import-time constant: the counter follows a + # mid-session cache-folder change, which the constant does not. monkeypatch.setattr(video_mod, "hub_cache_dir", lambda: str(tmp_path)) backend = VideoBackend() @@ -1171,8 +1158,8 @@ def test_cache_bytes_counts_incomplete_blobs(fake_runtime, tmp_path, monkeypatch def test_hv15_guider_and_scheduler_progress(fake_runtime): - # HunyuanVideo-1.5: no guidance kwarg (CFG set on the guider), no step - # callback (progress via the scheduler.step wrapper, restored afterwards). + # HunyuanVideo-1.5: no guidance kwarg (CFG set on the guider), no step callback (progress via the + # scheduler.step wrapper, restored afterwards). backend = VideoBackend() status = backend.load_pipeline( "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v", @@ -1202,24 +1189,23 @@ def test_hv15_cancel_unwinds_scheduler_loop(fake_runtime): model_kind = "pipeline", ) pipe = _FakeHV15Pipeline.instance - # Cancel lands during the FIRST real step; the next wrapped call must raise out - # of the denoise loop and generate() must surface the cancelled sentinel. + # The cancel lands during the FIRST real step; the next wrapped call must raise out of the denoise + # loop and generate() must surface the cancelled sentinel. pipe.scheduler.on_step = lambda n: backend.cancel_generate() if n == 1 else None with pytest.raises(RuntimeError, match = VIDEO_CANCELLED_MSG): backend.generate(prompt = "a fox", steps = 4) assert pipe.scheduler.calls == 1 # The wrapper must restore scheduler.step even on the exception path. assert pipe.scheduler.step.__func__ is _FakeHV15Scheduler.step - # The exception unwound pipe.__call__ before its own end-of-call cleanup, so generate() must - # have freed the offload hooks itself (VRAM would otherwise stay onloaded until the next - # request). + # The exception unwound pipe.__call__ before its own end-of-call cleanup, so generate() must have + # freed the offload hooks itself (VRAM would otherwise stay onloaded until the next request). assert pipe.hooks_freed == 1 def test_cancel_during_export_discards_clip(fake_runtime, monkeypatch): - # A cancel landing during the (blocking, uncancellable) export/mux must still discard the - # clip: cancel_generate() already reported success, so generate() must raise the cancelled - # sentinel rather than return the clip to be persisted to the gallery. + # A cancel landing during the (blocking, uncancellable) export/mux must still discard the clip: + # cancel_generate() already reported success, so generate() must raise the cancelled sentinel + # rather than return the clip to be persisted. backend = VideoBackend() backend.load_pipeline( "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v", @@ -1243,8 +1229,8 @@ def test_singleton(): def test_load_wan_ti2v_5b_pipeline(fake_runtime): - # A full-pipeline load of the single-DiT TI2V-5B repo: WanPipeline.from_pretrained, - # no audio, tiling forced on, and the 4k+1 frame lattice surfaced. + # A full-pipeline load of the single-DiT TI2V-5B repo: WanPipeline.from_pretrained, no audio, + # tiling forced on, and the 4k+1 frame lattice surfaced. backend = VideoBackend() status = backend.load_pipeline("Wan-AI/Wan2.2-TI2V-5B-Diffusers", model_kind = "pipeline") assert status["loaded"] is True @@ -1258,9 +1244,8 @@ def test_load_wan_ti2v_5b_pipeline(fake_runtime): def test_video_dense_speed_defaults_to_compile_profile(fake_runtime): - # A clip denoise amortises the one-time compile within a single run, so an UNSET speed on a - # dense (pipeline) load resolves to `default` -- never `max`, never `off`. Explicit "off" is - # still honored verbatim. + # A clip denoise amortises the one-time compile within a single run, so an UNSET speed on a dense + # (pipeline) load resolves to `default` -- never `max`, never `off`. Explicit "off" is honored. backend = VideoBackend() status = backend.load_pipeline("Wan-AI/Wan2.2-TI2V-5B-Diffusers", model_kind = "pipeline") assert status["speed_mode"] == "default" @@ -1274,10 +1259,9 @@ def test_video_dense_speed_defaults_to_compile_profile(fake_runtime): def test_video_speed_off_suppresses_auto_dtype_quant(fake_runtime, monkeypatch): - # An explicit Speed="off" (bit-exact) pipeline load with Precision left at auto must NOT - # promote the unset precision to auto-quant: that would engage torchao quantization (and force - # speed back to default), breaking the bit-exact request. Mirrors the image backend. On a - # dense-capable GPU (stubbed) quantize_transformer must not run. + # An explicit Speed="off" pipeline load with Precision left at auto must NOT promote the unset + # precision to auto-quant: that would engage torchao quantization (and force speed back to default), + # breaking the bit-exact request. Mirrors the image backend. import core.inference.video as video_mod monkeypatch.setattr(video_mod, "dense_transformer_supported", lambda target: True) @@ -1296,16 +1280,16 @@ def test_video_speed_off_suppresses_auto_dtype_quant(fake_runtime, monkeypatch): assert status["speed_mode"] == "off" backend.unload() - # Control: with speed NOT off, the auto precision promotion still engages the dense quant, so - # the suppression above is specific to speed=off (not a blanket disable). + # Control: with speed NOT off, the auto precision promotion still engages the dense quant, so the + # suppression above is specific to speed=off. backend.load_pipeline("Wan-AI/Wan2.2-TI2V-5B-Diffusers", model_kind = "pipeline") assert calls == [True] def test_video_step_cache_auto_from_default_schedule(fake_runtime, tmp_path): # Unset step cache is AUTO, decided from the model's default schedule: Wan's 50-step default - # engages FBCache at load; the LTX distilled 8-step default keeps it off. Both are re-checked - # per generation (toggle test below). + # engages FBCache at load; the LTX distilled 8-step default keeps it off. Both are re-checked per + # generation (toggle test below). backend = VideoBackend() status = backend.load_pipeline("Wan-AI/Wan2.2-TI2V-5B-Diffusers", model_kind = "pipeline") assert status["transformer_cache"] == "fbcache" @@ -1325,8 +1309,8 @@ def test_video_step_cache_auto_from_default_schedule(fake_runtime, tmp_path): def test_video_step_cache_auto_toggles_on_actual_steps(fake_runtime): - # The AUTO decision follows the ACTUAL step count of each generation: a few-step request drops - # the load-time cache, a many-step request restores it. An explicit "off" never toggles. + # The AUTO decision follows the ACTUAL step count of each generation: a few-step request drops the + # load-time cache, a many-step request restores it. An explicit "off" never toggles. backend = VideoBackend() backend.load_pipeline("Wan-AI/Wan2.2-TI2V-5B-Diffusers", model_kind = "pipeline") assert backend.status()["transformer_cache"] == "fbcache" @@ -1352,8 +1336,8 @@ def test_wan_frame_snapping_4k_plus_1(fake_runtime): backend.generate(prompt = "a sloth", width = 1000, height = 700, num_frames = 120) call = backend._state.pipe.last_kwargs assert call["num_frames"] == 117 # 4*29 + 1 - # /32 spatial snap for TI2V-5B: its VAE is 16x spatial * patch 2 = 32 (WanPipeline floors - # H/W to 32), so 1000x700 -> 992x672 (not the /16 992x688). + # /32 spatial snap for TI2V-5B: its VAE is 16x spatial * patch 2 = 32, so 1000x700 gives 992x672 + # (not the /16 992x688). assert (call["width"], call["height"]) == (992, 672) @@ -1368,9 +1352,8 @@ def test_wan_ti2v_defaults_applied(fake_runtime): def test_wan_ti2v_does_not_thread_cfg2(fake_runtime): - # The single-DiT TI2V pipeline has no guidance_scale_2 in its signature, so a request value - # must NOT be threaded (WanPipeline raises on it when boundary_ratio is None), even if the - # caller passes guidance_2. + # The single-DiT TI2V pipeline has no guidance_scale_2 in its signature, so a request value must + # NOT be threaded (WanPipeline raises on it when boundary_ratio is None). backend = VideoBackend() backend.load_pipeline("Wan-AI/Wan2.2-TI2V-5B-Diffusers", model_kind = "pipeline") backend.generate(prompt = "a sloth", guidance_2 = 3.5) @@ -1388,8 +1371,8 @@ def test_wan_a14b_dual_dit_pipeline_loads(fake_runtime): def test_wan_a14b_cfg2_threaded_when_signature_has_it(fake_runtime): - # The MoE pipeline's __call__ carries guidance_scale_2, so an explicit guidance_2 - # is threaded through as that kwarg (the cfg2_kwarg the family declares). + # The MoE pipeline's __call__ carries guidance_scale_2, so an explicit guidance_2 is threaded + # through as that kwarg (the cfg2_kwarg the family declares). backend = VideoBackend() backend.load_pipeline("Wan-AI/Wan2.2-T2V-A14B-Diffusers", model_kind = "pipeline") backend.generate(prompt = "a sloth", guidance = 5.0, guidance_2 = 3.0) @@ -1404,8 +1387,8 @@ def test_wan_a14b_cfg2_threaded_when_signature_has_it(fake_runtime): def test_wan_a14b_step_cache_applies_to_both_dits(fake_runtime): - # A dual-DiT MoE load must engage the step cache on BOTH experts, not just the - # first: transformer_2 handles the low-noise steps and would otherwise run uncached. + # A dual-DiT MoE load must engage the step cache on BOTH experts, not just the first: transformer_2 + # handles the low-noise steps and would otherwise run uncached. backend = VideoBackend() status = backend.load_pipeline( "Wan-AI/Wan2.2-T2V-A14B-Diffusers", @@ -1420,8 +1403,7 @@ def test_wan_a14b_step_cache_applies_to_both_dits(fake_runtime): def test_wan_a14b_attention_applies_to_both_dits(fake_runtime, monkeypatch): # An explicit attention backend must be set on both experts. The fake runtime is a CPU target, - # where the NVIDIA gate drops explicit kernels; pin the gate open so the explicit-set path is - # what this test exercises. + # where the NVIDIA gate drops explicit kernels, so pin the gate open. from core.inference import diffusion_attention as attn_mod monkeypatch.setattr(attn_mod, "_is_cuda_nvidia", lambda target: True) @@ -1452,9 +1434,9 @@ def test_wan_ti2v_single_dit_only_touches_one(fake_runtime): def test_wan_a14b_dense_quant_applies_to_both_dits(fake_runtime, monkeypatch): - # transformer_quant on a pipeline load quantises the dense DiT(s). On CPU the real dense path - # is unsupported, so stub the two quant seams to record which pipe view each helper saw: BOTH - # experts must be quantised (via the _SecondDiTView proxy), and status must report the scheme. + # transformer_quant on a pipeline load quantises the dense DiT(s). On CPU the real dense path is + # unsupported, so stub the two quant seams to record which pipe view each helper saw: BOTH experts + # must be quantised (via the _SecondDiTView proxy), and status must report the scheme. import core.inference.video as video_mod monkeypatch.setattr(video_mod, "dense_transformer_supported", lambda target: True) @@ -1468,8 +1450,8 @@ def test_wan_a14b_dense_quant_applies_to_both_dits(fake_runtime, monkeypatch): family, logger = None, ): - # The helper reads view.transformer; record the object it would quantise so the - # test proves the second expert was reached through the proxy. + # The helper reads view.transformer; record the object it would quantise so the test proves the + # second expert was reached through the proxy. quantised.append(view.transformer) return "int8" @@ -1488,10 +1470,9 @@ def test_wan_a14b_dense_quant_applies_to_both_dits(fake_runtime, monkeypatch): def test_dense_quant_skipped_under_offload(fake_runtime, monkeypatch): - # Offload hooks move modules with Module.to(), which torchao quantized tensors reject - # (observed as a hard crash on the A14B gate run). When the memory plan resolves to any offload - # policy, quant must be SKIPPED, not attempted: the load succeeds dense and the record explains - # why. + # Offload hooks move modules with Module.to(), which torchao quantized tensors reject (a hard crash + # on the A14B gate run). When the memory plan resolves to any offload policy, quant must be + # SKIPPED, not attempted: the load succeeds dense and the record explains why. import core.inference.video as video_mod monkeypatch.setattr(video_mod, "dense_transformer_supported", lambda target: True) @@ -1509,9 +1490,8 @@ def test_dense_quant_skipped_under_offload(fake_runtime, monkeypatch): return "int8" monkeypatch.setattr(video_mod, "quantize_transformer", _fake_quant) - # The CPU fake target never plans an offload, so force one at the plan seam (frozen dataclass - # -> dataclasses.replace) and stub the apply step, which would else call offload hooks the fake - # pipe lacks. + # The CPU fake target never plans an offload, so force one at the plan seam and stub the apply + # step, which would else call offload hooks the fake pipe lacks. import dataclasses real_plan = video_mod.plan_diffusion_memory @@ -1539,9 +1519,9 @@ def test_dense_quant_skipped_under_offload(fake_runtime, monkeypatch): def test_wan_a14b_partial_quant_fails_the_load(fake_runtime, monkeypatch): - # If the first expert quantises but the second doesn't, the pipe is left at mismatched - # precision with no way back (in-place mutation), so the load must fail cleanly rather than run - # mixed with quant reported off. + # If the first expert quantises but the second doesn't, the pipe is left at mismatched precision + # with no way back (in-place mutation), so the load must fail cleanly rather than run mixed with + # quant reported off. import core.inference.video as video_mod monkeypatch.setattr(video_mod, "dense_transformer_supported", lambda target: True) @@ -1593,8 +1573,8 @@ def test_wan_ti2v_dense_quant_applies_to_single_dit(fake_runtime, monkeypatch): def test_wan_validate_trusted_repos(fake_runtime): - # The two Wan base repos are trusted for non-GGUF (pipeline) loads; an unrelated - # repo carrying the family name is not. + # The two Wan base repos are trusted for non-GGUF (pipeline) loads; an unrelated repo carrying the + # family name is not. backend = VideoBackend() fam = backend.validate_load_request("Wan-AI/Wan2.2-TI2V-5B-Diffusers", model_kind = "pipeline") assert fam.name == "wan2.2-ti2v-5b" @@ -1612,9 +1592,9 @@ def test_wan_validate_trusted_repos(fake_runtime): def test_wan_a14b_refuses_single_file_loads(fake_runtime): - # A single gguf/safetensors checkpoint carries only one of the A14B's two experts; the - # pipeline would pull the other dense bf16 from the base repo outside the memory plan, so - # validate refuses it up front (before any download). + # A single gguf/safetensors checkpoint carries only one of the A14B's two experts, and the pipeline + # would pull the other dense bf16 from the base repo outside the memory plan, so validate refuses + # it up front. backend = VideoBackend() with pytest.raises(ValueError, match = "dual-expert"): backend.validate_load_request( @@ -1631,8 +1611,7 @@ def test_wan_a14b_refuses_single_file_loads(fake_runtime): def test_second_dit_view_write_through(): # Attribute writes on the proxy must land on the real pipe (a helper's side effect would else - # vanish with the temporary view); a ``transformer`` write mirrors the read property onto the - # second expert. + # vanish with the temporary view); a ``transformer`` write mirrors onto the second expert. from core.inference.video import _SecondDiTView pipe = types.SimpleNamespace(transformer = "t1", transformer_2 = "t2", flag = None) @@ -1669,8 +1648,8 @@ _LTX2_SIBLINGS = [ def test_base_download_files_scopes_pipeline_pull(): - # A pipeline load skips the packaged root checkpoint, the duplicate - # text-encoder shard naming, and non-weight assets -- and keeps everything else. + # A pipeline load skips the packaged root checkpoint, the duplicate text-encoder shard naming and + # non-weight assets, and keeps everything else. info = types.SimpleNamespace(siblings = _LTX2_SIBLINGS) files = dict(VideoBackend._base_download_files(info, "pipeline")) assert "ltx-2-19b-packaged-fp8.safetensors" not in files @@ -1678,8 +1657,8 @@ def test_base_download_files_scopes_pipeline_pull(): assert "assets/example.mp4" not in files assert files["text_encoder/model-00001-of-00002.safetensors"] == 25 assert files["transformer/diffusion_pytorch_model-00001-of-00002.safetensors"] == 20 - # The standalone chat template must survive the whitelist: apply_chat_template - # reads it at generation time and it is not embedded in tokenizer_config.json. + # The standalone chat template must survive the whitelist: apply_chat_template reads it at + # generation time and it is not embedded in tokenizer_config.json. assert "tokenizer/chat_template.jinja" in files assert sum(files.values()) == 10 + 1 + 20 + 18 + 25 + 25 + 3 + 1 + 1 @@ -1693,8 +1672,8 @@ def test_base_download_files_gguf_drops_transformer(): def test_load_progress_clamps_overshoot(fake_runtime, monkeypatch): - # The cache scan counts blobs a broader previous pull left behind; the reported - # counter must never exceed the scoped estimate (no "282 GB of 263 GB"). + # The cache scan counts blobs a broader previous pull left behind, so the reported counter must + # never exceed the scoped estimate (no "282 GB of 263 GB"). backend = VideoBackend() backend._loading = types.SimpleNamespace( repo_id = "Lightricks/LTX-2", base_repo = None, expected_bytes = 100, error = None @@ -1707,8 +1686,8 @@ def test_load_progress_clamps_overshoot(fake_runtime, monkeypatch): def test_pipeline_load_uses_predownloaded_dir(fake_runtime, tmp_path): - # When the scoped pre-download produced a local snapshot, from_pretrained must - # receive that dir (keeping diffusers' own broader snapshot sweep off the hub). + # When the scoped pre-download produced a local snapshot, from_pretrained must receive that dir, + # keeping diffusers' own broader snapshot sweep off the hub. backend = VideoBackend() backend.load_pipeline( "Lightricks/LTX-2", @@ -1720,8 +1699,8 @@ def test_pipeline_load_uses_predownloaded_dir(fake_runtime, tmp_path): def test_base_download_files_ltx23_keeps_only_shared_components(): - # A 2.3 checkpoint supplies the DiT, connectors, both VAEs and the vocoder, so - # the base pull shrinks to scheduler + text encoder + tokenizer (+ root manifest). + # A 2.3 checkpoint supplies the DiT, connectors, both VAEs and the vocoder, so the base pull + # shrinks to scheduler + text encoder + tokenizer (+ root manifest). siblings = _LTX2_SIBLINGS + [ _sibling("scheduler/scheduler_config.json", 1), _sibling("connectors/diffusion_pytorch_model.safetensors", 3), @@ -1739,8 +1718,8 @@ def test_base_download_files_ltx23_keeps_only_shared_components(): def test_hv15_720p_repo_gets_720p_family_defaults(): - # The 720p repack is trusted, but it must resolve its OWN family entry: the - # generic hunyuanvideo-1.5 entry would default generation to 832x480. + # The 720p repack is trusted, but it must resolve its OWN family entry: the generic + # hunyuanvideo-1.5 entry would default generation to 832x480. from core.inference.video_families import detect_video_family fam = detect_video_family("hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v") @@ -1754,8 +1733,8 @@ def test_hv15_720p_repo_gets_720p_family_defaults(): def test_predownload_base_honors_cancel_between_files(monkeypatch): - # A warm-cache sweep returns each file instantly without consulting the event, - # so the loop must check it explicitly or an unload mid-predownload is ignored. + # A warm-cache sweep returns each file instantly without consulting the event, so the loop must + # check it explicitly or an unload mid-predownload is ignored. backend = VideoBackend() backend._cancel_event.set() calls: list = [] @@ -1789,9 +1768,9 @@ def test_predownload_base_honors_cancel_between_files(monkeypatch): def test_detect_load_family_arch_fallback_for_local_gguf(tmp_path, monkeypatch): - # A local GGUF is admitted to the Video picker by its general.architecture, but its path name - # may carry no whole-segment family token (a renamed "model.gguf"). The loader must resolve the - # same family the picker offered by reading the arch, not only the name. + # A local GGUF is admitted to the Video picker by its general.architecture, but its path name may + # carry no whole-segment family token (a renamed "model.gguf"). The loader must resolve the same + # family the picker offered by reading the arch, not only the name. from core.inference import video as vid from core.inference.video_families import detect_video_family @@ -1811,7 +1790,7 @@ def test_detect_load_family_arch_fallback_for_local_gguf(tmp_path, monkeypatch): fam = vid._detect_load_family(str(d), "model.gguf", None) assert fam is not None and fam.name == "ltx-2" - # A video arch with no backend family (wan) stays None -> the loader 400s as before. + # A video arch with no backend family (wan) stays None, so the loader 400s as before. monkeypatch.setattr( "utils.models.gguf_metadata.read_gguf_general_metadata", lambda p: {"general.architecture": "wan"}, @@ -1873,10 +1852,10 @@ def _plan_api(monkeypatch, repos): def test_download_plan_narrows_an_ltx23_pick_and_stages_its_extras(monkeypatch): - # A 2.3 checkpoint brings its own VAEs, vocoder and connectors, so staging the 2.0 base's - # copies downloads gigabytes the pipeline never opens -- and the companion files it DOES - # read were left out of the plan entirely, so they were pulled inline at load time, - # outside the panel's progress, cancel and disk preflight. + # A 2.3 checkpoint brings its own VAEs, vocoder and connectors, so staging the 2.0 base's copies + # downloads gigabytes the pipeline never opens -- and the companion files it DOES read were left + # out of the plan entirely, so they were pulled inline at load time, outside the panel's progress, + # cancel and disk preflight. _plan_api( monkeypatch, { @@ -1892,8 +1871,8 @@ def test_download_plan_narrows_an_ltx23_pick_and_stages_its_extras(monkeypatch): ) by_repo = {e["repo_id"]: e for e in plan["entries"]} - # Checkpoint and extras share one repo, so they must be ONE entry: two jobs for the same - # repo would collide on the scoped job key. + # Checkpoint and extras share one repo, so they must be ONE entry: two jobs for the same repo would + # collide on the scoped job key. assert set(by_repo) == {"unsloth/LTX-2.3-GGUF", "Lightricks/LTX-2"} ckpt = by_repo["unsloth/LTX-2.3-GGUF"] assert ckpt["gguf_filename"] == "ltx-2.3-22b-distilled.gguf" diff --git a/studio/backend/tests/test_video_families.py b/studio/backend/tests/test_video_families.py index 3455abe7bc..5271964845 100644 --- a/studio/backend/tests/test_video_families.py +++ b/studio/backend/tests/test_video_families.py @@ -74,8 +74,8 @@ def test_detect_wan_ti2v_5b(repo_id): ], ) def test_detect_wan_t2v_a14b(repo_id): - # The A14B repo ids route to the dual-expert MoE family: a second DiT + a second - # guidance kwarg (guidance_scale_2, verified present in diffusers 0.39). + # The A14B repo ids route to the dual-expert MoE family: a second DiT + a second guidance kwarg + # (guidance_scale_2, verified present in diffusers 0.39). fam = detect_video_family(repo_id) assert fam is not None and fam.name == "wan2.2-t2v-a14b" assert fam.pipeline_class == "WanPipeline" @@ -102,8 +102,8 @@ def test_wan_and_ltx_do_not_cross_route(): def test_sentinels_are_video_specific(): - # The routes match these EXACTLY for 409s; they must not collide with the - # image sentinels or a video 409 would be mis-attributed. + # The routes match these EXACTLY for 409s; they must not collide with the image sentinels or a + # video 409 would be mis-attributed. assert "video" in VIDEO_NOT_LOADED_MSG.lower() assert "video" in VIDEO_CANCELLED_MSG.lower() @@ -117,8 +117,8 @@ def test_resolve_base_repo(): def test_snap_num_frames_lattice(): fam = detect_video_family("unsloth/LTX-2.3-GGUF") - # Valid counts are k * 8 + 1: on-lattice values pass through, everything - # else floors to the previous lattice point, never below 1. + # Valid counts are k * 8 + 1: on-lattice values pass through, everything else floors to the + # previous lattice point, never below 1. assert snap_num_frames(fam, 121) == 121 assert snap_num_frames(fam, 120) == 113 assert snap_num_frames(fam, 122) == 121 @@ -136,8 +136,8 @@ def test_snap_video_size_multiple(): def test_generation_defaults_distilled_vs_dev(): - # The distilled checkpoints run few-step with CFG off; the dev-config base - # repo wants the full schedule. The picked filename wins over the base repo. + # The distilled checkpoints run few-step with CFG off; the dev-config base repo wants the full + # schedule. The picked filename wins over the base repo. assert default_video_generation_params( "distilled-1.1/ltx-2.3-22b-distilled-1.1-Q4_K_M.gguf", "Lightricks/LTX-2" ) == (8, 1.0) @@ -170,8 +170,8 @@ def test_wan_snap_video_size_16(): # Wan patchifies at spatial factor 8 * patch 2 = 16; sizes floor to /16. fam = detect_video_family("Wan-AI/Wan2.2-T2V-A14B-Diffusers") assert fam.resolution_multiple == 16 - # A14B's native 720p is the true 16:9 1280x720 (720 = 45*16 renders exactly on the /16 grid), - # NOT the 1280x704 that TI2V-5B's /32 VAE floors to. The default preset is that native 720p. + # A14B's native 720p is the true 16:9 1280x720 (720 = 45*16 renders exactly on the /16 grid), NOT + # the 1280x704 that TI2V-5B's /32 VAE floors to. The default preset is that native 720p. assert fam.resolution_presets[0] == (1280, 720) assert snap_video_size(fam, 1280, 720) == (1280, 720) # native 720p, on-grid assert snap_video_size(fam, 1000, 700) == (992, 688) @@ -188,9 +188,9 @@ def test_wan_generation_defaults(): def test_generation_defaults_fallback_honors_family(): - # When no identifier names a known variant (a Wan model loaded from an opaque local path under - # an explicit family_override), the fallback -- the resolved family's own default -- is used, - # not the hardcoded LTX 40/4.0. Without a family fallback a Wan model would wrongly run 40/4.0. + # When no identifier names a known variant (a Wan model loaded from an opaque local path under an + # explicit family_override), the resolved family's own default is used, not the hardcoded LTX + # 40/4.0. assert default_video_generation_params("/models/my-clip", "/models/my-clip") == (40, 4.0) assert default_video_generation_params( "/models/my-clip", "/models/my-clip", fallback = (50, 5.0) @@ -200,9 +200,9 @@ def test_generation_defaults_fallback_honors_family(): def test_generation_defaults_wan_is_segment_not_substring(): - # "wan" must match as a name segment, not a raw substring, so an opaque non-Wan - # repo/path whose name merely contains the letters "wan" ("swan", "taiwan") does - # NOT silently pick up Wan's 50-step/CFG-5 schedule ahead of its canonical base repo. + # "wan" must match as a name segment, not a raw substring, so an opaque non-Wan repo/path whose + # name merely contains the letters ("swan", "taiwan") does NOT pick up Wan's 50-step/CFG-5 + # schedule ahead of its canonical base repo. assert default_video_generation_params( "user/swan-video", "Lightricks/LTX-2", fallback = (40, 4.0) ) == (40, 4.0) @@ -221,11 +221,10 @@ def test_wan_size_tables_present(): ti2v = detect_video_family("Wan-AI/Wan2.2-TI2V-5B-Diffusers") a14b = detect_video_family("Wan-AI/Wan2.2-T2V-A14B-Diffusers") assert ti2v.bf16_components_gb is not None and a14b.bf16_components_gb is not None - # bf16-RESIDENT transformer sizes. The Wan transformers ship FP32 on disk (safetensors - # headers are F32; TI2V index 20.0 GB = 5B x 4, A14B 57.15 GB per expert = 14.3B x 4), so the - # table must hold the HALVED bf16-resident sizes -- ti2v ~10.0, a14b two experts ~57.2 -- NOT - # the fp32 on-disk sums (20.0 / 114.3). A regression to those over-budgets the plan ~2x and - # forces needless offload on an 80 GB GPU. + # bf16-RESIDENT transformer sizes. The Wan transformers ship FP32 on disk (TI2V index 20.0 GB = + # 5B x 4, A14B 57.15 GB per expert), so the table must hold the HALVED bf16-resident sizes -- + # ti2v ~10.0, a14b two experts ~57.2 -- NOT the fp32 on-disk sums. A regression to those + # over-budgets the plan ~2x and forces needless offload on an 80 GB GPU. assert ti2v.bf16_components_gb[0] == 10.0 assert a14b.bf16_components_gb[0] == 57.2 # The A14B DiT total (two experts) still dwarfs the single TI2V-5B DiT. @@ -235,9 +234,8 @@ def test_wan_size_tables_present(): def test_wan_ti2v_5b_snaps_to_32_not_16(): - # TI2V-5B's VAE is 16x spatial (vae/config.json scale_factor_spatial=16) and the transformer - # patch is 2, so WanPipeline floors H/W to 16*2 = 32 (pipeline_wan.py:505). The backend must - # snap to /32 too, or a /16-but-not-/32 request (e.g. 720) is recorded but rendered at 704, + # TI2V-5B's VAE is 16x spatial and the transformer patch is 2, so WanPipeline floors H/W to 32. The + # backend must snap to /32 too, or a /16-but-not-/32 request is recorded but rendered at 704, # desyncing gallery metadata from the actual clip. fam = detect_video_family("Wan-AI/Wan2.2-TI2V-5B-Diffusers") assert fam.resolution_multiple == 32 @@ -249,9 +247,9 @@ def test_wan_ti2v_5b_snaps_to_32_not_16(): def test_wan_families_force_vae_fp32(): - # Wan's VAE decodes in float32 (diffusers loads AutoencoderKLWan at torch.float32 while the - # pipe runs bf16); the loader pins the VAE back to fp32 for these families to avoid banding / - # black frames. LTX-2 keeps the default (its VAE is bf16-native). + # Wan's VAE decodes in float32 (diffusers loads AutoencoderKLWan at torch.float32 while the pipe + # runs bf16), so the loader pins the VAE back to fp32 for these families to avoid banding / black + # frames. LTX-2 keeps the default (its VAE is bf16-native). assert detect_video_family("Wan-AI/Wan2.2-TI2V-5B-Diffusers").vae_force_fp32 is True assert detect_video_family("Wan-AI/Wan2.2-T2V-A14B-Diffusers").vae_force_fp32 is True assert detect_video_family("unsloth/LTX-2.3-GGUF").vae_force_fp32 is False @@ -261,10 +259,9 @@ def test_family_size_table_present(): fam = detect_video_family("unsloth/LTX-2.3-GGUF") assert fam.bf16_components_gb is not None transformer_gb, text_encoder_gb, companions_gb = fam.bf16_components_gb - # RESIDENT bf16 figures: the 37.8 GB DiT and the Gemma3-12B TE at ~24.4 GB once cast - # to bf16 (the fp32 hub store is ~49 GB but never sits on device). A table that - # regressed to the fp32 download size would push auto planning to offload on cards - # that fit the real footprint. + # RESIDENT bf16 figures: the 37.8 GB DiT and the Gemma3-12B TE at ~24.4 GB once cast to bf16 (the + # fp32 hub store is ~49 GB but never sits on device). A table that regressed to the fp32 download + # size would push auto planning to offload on cards that fit the real footprint. assert transformer_gb > text_encoder_gb > 20.0 assert text_encoder_gb < 30.0 assert companions_gb > 0.0 @@ -273,20 +270,20 @@ def test_family_size_table_present(): def test_hv15_detection_and_flags(): fam = detect_video_family("hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v") assert fam is not None and fam.name == "hunyuanvideo-1.5" - # CFG lives on the guider component (no guidance kwarg in __call__), and the - # HV15 VAE compresses 16x spatial / 4x temporal. + # CFG lives on the guider component (no guidance kwarg in __call__), and the HV15 VAE compresses + # 16x spatial / 4x temporal. assert fam.guidance_via_guider is True assert fam.frame_step == 4 and fam.resolution_multiple == 16 assert fam.has_audio is False assert detect_video_family("x/y", override = "hv15") is fam - # The incompatible HunyuanVideo 1.0 repos must NOT be claimed: their - # model_index pins HunyuanVideoPipeline, which this family cannot load. + # The incompatible HunyuanVideo 1.0 repos must NOT be claimed: their model_index pins + # HunyuanVideoPipeline, which this family cannot load. assert detect_video_family("hunyuanvideo-community/HunyuanVideo") is None def test_hv15_generation_defaults(): - # The community repacks ship a guider with guidance_scale 6.0 and the - # pipeline's own 50-step schedule. + # The community repacks ship a guider with guidance_scale 6.0 and the pipeline's own 50-step + # schedule. assert default_video_generation_params( None, "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v" ) == (50, 6.0)