From bfabbbe7b64a4a93039368a748bf0d1dec173720 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 3 Jul 2026 11:42:56 -0700 Subject: [PATCH] Studio diffusion: LoRA adapters for the Images workflow (#6771) * Studio diffusion: cross-platform device policy, fp16 guard, lock split, validate-before-evict Phase 1 of porting the richer diffusion stack onto the image-generation backend. - Add a compartmentalized device/dtype policy module (diffusion_device.py) resolving CUDA/ROCm/XPU/MPS/CPU with capability flags. Keeps the NVIDIA capability-based bf16 choice; ROCm and XPU are isolated; MPS uses bf16 or fp32, never a silent fp16 that renders a black image. - Add a per-family fp16_incompatible flag (Z-Image) and promote a resolved float16 to float32 for those families so they do not produce black images. - Split the backend locks: a generation holds only _generate_lock, so status, unload, and a new load are never blocked by a long denoise. Add per-generation cancellation via callback_on_step_end so an eviction or a superseding load preempts a running generation; a replacement load waits for it to stop before allocating, so two pipelines never sit in VRAM at once. - Validate a load request before the GPU handoff so an unloadable pick never evicts a working chat model, and reject missing local paths up front. - Add CPU-only tests for the device policy, dtype guard, lock split and cancellation, and validate-before-evict, plus a GPU benchmark/regression script (scripts/diffusion_bench.py) measuring latency, peak VRAM, and PSNR against a saved reference. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 2A): measured-budget memory planner + offload/VAE policy Add a lean, backend-agnostic memory policy that picks a CPU-offload policy and VAE tiling/slicing from measured free device memory vs the model's estimated resident footprint, then applies it to the built pipeline. auto stays resident when the model fits (byte-identical to the prior resident path), and falls to whole-module offload when tight; fast/balanced/low_vram are explicit overrides. Sequential submodule offload is unreliable for GGUF transformers on diffusers 0.38, so it falls back to whole-module offload and status reports the policy actually engaged. Verified on Z-Image-Turbo Q4_K_M (B200): auto reproduces the resident image with no VRAM/latency regression (PSNR inf); balanced/low_vram cut generation peak VRAM 47.9% (15951 -> 8318 MB) with byte-identical output, at the expected latency cost. 73 prior + 35 new CPU tests pass. * Studio diffusion (Phase 2D): streamed block-level offload + functional VAE tiling Add a streamed 'group' offload tier (diffusers apply_group_offloading, block_level, use_stream) that keeps the transformer flowing through the GPU a few blocks at a time while the text encoder / VAE stay resident, and fix VAE tiling to drive the VAE submodule (pipelines like Z-Image expose enable_tiling on pipe.vae, not the pipeline). apply_memory_plan now returns the (policy, tiling) actually engaged so status never overstates either, and group falls back to whole-module offload when the transformer can't be streamed. Measured on Z-Image (B200), all lossless (PSNR inf vs resident): balanced/group cuts generation peak VRAM 32% (15951 -> 10840 MB) at near-resident speed (2.07 -> 2.99s); low_vram/model cuts it 48% (-> 8318 MB) but is slower (7.99s). Mode names now match that tradeoff: balanced = stream the transformer, low_vram = offload every component. auto picks group when the companions fit resident, else model. 112 CPU tests pass. * Studio diffusion (Phase 5): image quality-vs-quant accuracy harness Add scripts/diffusion_quality.py, the accuracy analogue of the KLD workflow: hold prompt + seed fixed, render a grid with a reference quant (default BF16), then render each candidate quant and measure drift from the reference. Records mean PSNR + SSIM (pure-numpy, no skimage/scipy) and optional CLIP text-alignment + image-similarity (transformers, --clip), plus file size, latency, and peak VRAM, then prints a quality-vs-cost table and recommends the smallest quant within a quality budget. --selftest validates the metrics on synthetic images with no GPU or model. Verified on Z-Image (B200): the table degrades monotonically with quant size (Q8 -> Q4 -> Q2: PSNR 21.7 -> 15.5, SSIM 0.82 -> 0.61), while CLIP-text stays flat (~0.34) -- quantization erodes fine detail far more than prompt adherence. * Studio diffusion (Phase 3): opt-in speed layer (channels_last / compile / TF32) Add a speed_mode knob (off by default, so the render path stays bit-identical): default applies channels_last VAE + regional torch.compile of the denoiser's repeated block where eligible; max also enables TF32 matmul and fused QKV. Regional compile is gated off for the GGUF transformer (dequantises per-op) and for families flagged not compile-friendly (a new supports_torch_compile flag, False for Z-Image), so it activates automatically only once a non-GGUF bf16 transformer is loaded. Speed optims run before placement/offload, per the diffusers composition order. status now reports speed_mode + the optims actually engaged. Verified on Z-Image (B200): default -> ['channels_last'], max -> ['channels_last', 'tf32'], compile correctly skipped for GGUF; generation works in every mode. 121 CPU tests pass. * Studio diffusion (Phase 2B): opt-in fp8 text-encoder layerwise casting Add a text_encoder_fp8 knob that casts the companion text encoder(s) to fp8 (e4m3) storage via diffusers apply_layerwise_casting, upcasting per layer to the bf16 compute dtype while normalisations and embeddings stay full precision. Applied before placement, gated to CUDA + bf16, best-effort (a failure leaves the encoder dense). status reports which encoders were cast. Verified on Z-Image (B200, balanced/group mode where the encoder stays resident): generation peak VRAM dropped 37% (10840 -> 6791 MB, below the lowest-VRAM offload) at near-resident speed. It is a memory-vs-quality tradeoff, not free -- ~20 dB PSNR vs the bf16 encoder, a larger shift than one transformer quant step -- so it is off by default and documented as such, with the Phase 5 harness to size the cost. 127 CPU tests pass. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 2C): NVFP4 text-encoder quant (+ generalise fp8 knob) Generalise the text-encoder precision knob from a fp8 bool to text_encoder_quant (fp8 | nvfp4). nvfp4 quantises the companion text encoder to 4-bit via torchao NVFP4 weight-only (two-level microscaling) on Blackwell's FP4 tensor cores; fp8 stays the broader-hardware path (cc>=8.9). Both are gated, best-effort, and run before placement; status reports the mode actually engaged. This is the lean realisation of GGUF-native text-encoder quant: 4-bit on the encoder without the 3045-line port. Verified on Z-Image (B200, balanced/group where the encoder stays resident), vs the bf16 encoder: nvfp4 cut generation peak VRAM 48% (10840 -> 5593 MB, the lowest TE option, below whole-model offload) at near-fp8 quality (16.4 vs 17.1 dB PSNR), and both quants ran faster than bf16. A memory-vs-quality tradeoff (off by default); size it per model with the Phase 5 quality harness. diffusion_bench gains --text-encoder-quant. 129 CPU tests pass. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 4): native stable-diffusion.cpp engine for CPU/Mac Adds the CPU / Apple-Silicon tier of the two-engine strategy, mirroring the chat backend's llama.cpp shell-out. Diffusers stays the default on CUDA / ROCm / XPU; this covers the hardware diffusers serves poorly, consuming the same split GGUF assets Studio already curates. - sd_cpp_args.py: pure sd-cli command builder. Maps the family to its text-encoder flag (Z-Image Qwen3 to --llm, Qwen-Image to --qwen2vl, FLUX.1 CLIP-L + T5), and the diffusers memory policy (none/group/model/sequential) to sd.cpp's offload flags (--offload-to-cpu / --clip-on-cpu / --vae-on-cpu / --vae-tiling / --diffusion-fa), so one user knob drives both engines. - sd_cpp_engine.py: SdCppEngine over a located sd-cli. find_sd_cpp_binary() with the same precedence as the llama finder (env override, then the Studio install root, then in-tree, then PATH), an is_available/version probe, and a one-shot subprocess generate that streams progress and returns the PNG. runtime_env() prepends the binary's directory to the platform library path so a prebuilt's bundled libstable-diffusion.so resolves. select_diffusion_engine() is the pure routing decision (GPU backends to diffusers, CPU/MPS to native when present). - install_sd_cpp_prebuilt.py: resolve + download the per-host prebuilt (macOS-arm64/Metal, Linux x86_64 CPU, Vulkan/ROCm/Windows variants) into the Studio install root. resolve_release_asset() is a pure, unit-tested host-to-asset matrix. - scripts/sd_cpp_smoke.py: end-to-end native generation harness. Tests (CPU-only, subprocess/filesystem stubbed): 49 new across args, engine, routing, runtime env, and the installer resolver. Full diffusion suite 166 passing. Verified on a B200 box: built sd-cli (CUDA) and the prebuilt (CPU) both generate Z-Image-Turbo Q4_K end to end through SdCppEngine: balanced (group offload, 5.0s gen), low_vram (full CPU offload + VAE tiling, 13.4s), and the dynamically-linked CPU prebuilt (50.4s on CPU), all producing coherent images. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 6): img2img / inpaint / edit / LoRA / upscale on the native engine Builds on Phase 4's native stable-diffusion.cpp engine, extending it from text-to-image to the wider feature surface, since sd.cpp supports all of these through the binary already. Pure command-builder additions plus one engine method, so the txt2img path is unchanged. - sd_cpp_args.py: SdCppGenParams gains image-conditioning fields. init_img + strength make a run img2img, adding mask makes it inpaint, ref_images drives FLUX-Kontext / Qwen-Image-Edit style editing (repeated --ref-image), and lora_dir + the prompt syntax select LoRAs. New SdCppUpscaleParams + build_sd_cpp_upscale_command for the ESRGAN upscale run mode (input image + esrgan model, no prompt / text encoders). - sd_cpp_engine.py: the subprocess runner is factored into a shared _run() so generate() (now carrying the conditioning flags) and a new upscale() reuse the same streaming / error / output-check path. - scripts/sd_cpp_smoke.py: --task {txt2img,img2img,upscale} with --init-img / --strength / --upscale-model / --upscale-repeats. Tests: 10 new across the img2img / inpaint / edit / LoRA flag construction, the upscale builder and its validation, and the engine's img2img + upscale paths. Full diffusion suite 176 passing. Verified on a B200 box through SdCppEngine: img2img (Z-Image-Turbo Q4_K, the init image conditioned at strength 0.6, 4.8s) and ESRGAN upscale (512x512 -> 2048x2048 via RealESRGAN_x4plus_anime_6B, 2.7s), both producing coherent images. Video and the diffusers-path feature wiring are deferred. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 7): accuracy-preserving speed pass Re-review of the diffusion stack (#6675/#6679/#6680) surfaced one real accuracy bug and a dead-on-arrival speed path; this fixes both and adds the lossless / near-lossless wins, all measured on a B200. Correctness: - TF32 global-state leak (fix). speed_mode=max flipped torch.backends.*.allow_tf32 process-wide and never restored them, so a later `off` load silently inherited TF32 and was no longer bit-identical. Added snapshot_backend_flags / restore_backend_flags (TF32 + cudnn.benchmark), captured before the speed layer runs and restored on unload. Verified: load max -> unload -> load off is now byte-identical (PSNR inf) to a fresh off. - sd-cli timeout could hang forever. _run() blocked in `for line in stdout` and only checked the timeout after EOF, so a child stuck in model load / GPU init with no output ignored the timeout. Drained stdout on a reader thread with a wall-clock deadline. Added a silent-hang regression test. Speed (diffusers path), near-lossless, opt-in tiers: - Regional torch.compile now runs on the GGUF transformer. The is_gguf gate (and Z-Image's supports_torch_compile=False) were stale: compile_repeated_blocks compiles and runs ~2.2x faster on the GGUF Z-Image transformer on torch 2.9.1 / diffusers 0.38 (the per-op dequant stays eager, the rest of the block compiles). Measured: off 1.80s -> default 0.82s/gen (+54.7%), PSNR 37.7 dB vs eager -- far above the Q4 quant noise floor (~21 dB), so it does not move output quality. Gate relaxed; default tier delivers it. - cudnn.benchmark added to the default tier (autotunes the fixed-shape VAE convs). - torch.inference_mode() around the pipeline call (lossless, strictly faster than the no_grad diffusers uses internally). Memory path: - VAE tiling (not bit-identical >1MP) restricted to the model/sequential/CPU tiers; the balanced (group) tier keeps exact slicing only, so it is now bit-identical to the resident image (verified PSNR inf) and slightly faster. - Group offload adds non_blocking + record_stream on the CUDA stream path to overlap each block's H2D copy with compute (lossless; gated on the installed diffusers signature so older versions still work). Native (sd.cpp) path: - native_speed_flags: a first-class speed knob (default -> --diffusion-fa, a near-lossless CUDA win that was previously only added on offload tiers; max also -> --diffusion-conv-direct). conv-direct stays opt-in: measured +45% on CUDA, so it is never auto-on. Engine generate() merges it, de-duped against offload flags. Default profile: a GGUF model with no explicit speed_mode now resolves to the `default` profile (resolve_speed_mode), since compile's perturbation sits below the quantisation noise floor and so does not reduce quality versus the dense reference; out of the box a GGUF Z-Image generation drops from 1.80s to 0.81s. Dense models stay `off` / bit-identical, and an explicit speed_mode -- including "off" -- is always honored, so the byte-identical path remains one flag away and is the regression reference. Tooling: scripts/compile_probe.py (eager vs compiled GGUF probe), scripts/ perf_verify.py (the B200 verification above), and diffusion_bench.py gains --speed-mode so the speed tiers are benchmarkable. Tests: 183 passing (was 166); new coverage for the backend-flag snapshot/restore, GGUF compile eligibility, the balanced tiling/slicing split, native_speed_flags + the engine de-dup, and the sd-cli silent-hang timeout. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 7): max tier uses max-autotune-no-cudagraphs + engine/lever benchmarks The opt-in `max` speed tier now compiles the repeated block with mode=max-autotune-no-cudagraphs (dynamic=False) instead of the default mode: Triton autotuning for GEMM/conv-heavier models, gated to the tier where a longer cold compile is acceptable. CUDA-graph modes (reduce-overhead / max-autotune) are deliberately avoided -- both crash on the regionally-compiled block (its static output buffer is overwritten across denoise steps), measured. Adds two reproducible benchmarks used to validate the optimization research: - scripts/compare_engines.py: PyTorch (diffusers GGUF) vs native sd.cpp head-to-head. - scripts/leverage_probe.py: coordinate_descent_tuning + FirstBlockCache probes. Measured on B200 (Z-Image Q4_K_M, 1024px, 8 steps): default compile 0.80s/gen; coordinate_descent_tuning 0.79s (within noise, already covered by max-autotune); FirstBlockCache does not run on Z-Image (diffusers 0.38 block-detection / Dynamo). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 8): opt-in fast transformer (torchao int8/fp8/fp4 on a dense source) Add an opt-in transformer_quant mode that loads the dense bf16 transformer and torchao-quantises it onto the low-precision tensor cores, instead of the GGUF transformer (which dequantises to bf16 per matmul and so runs at bf16 rate). On a B200 (Z-Image-Turbo, 1024px/8 steps): auto picks fp8 at 0.614s vs GGUF+compile's 0.823s (1.34x), int8 0.626s (1.32x), both at lower LPIPS than GGUF's own 4-bit floor. GGUF+compile stays the low-memory default and the fallback. The mode is gated on CUDA + bf16 + resident VRAM headroom (the dense load peaks ~21GB vs GGUF's 13GB); any unsupported arch/scheme, OOM, or quant failure falls back to GGUF with a logged reason. auto picks the best scheme per GPU via a real quantise+matmul smoke probe (Blackwell nvfp4/fp8/mxfp8, Ada/Hopper fp8, Ampere int8); a min-features filter skips the tiny projections that crash int8's torch._int_mm. New module mirrors diffusion_precision.py; quant runs before compile before placement. 184 -> tests pass; new test_diffusion_transformer_quant.py plus backend/route coverage. scripts/diffusion_bench.py gains --transformer-quant; scripts/quant_probe.py is the standalone torchao lever probe. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 8): consumer-GPU tuning - lock fp8 fast accumulate, prefer fp8 over mxfp8, reject 2:4 sparsity Consumer Blackwell halves tensor-core throughput on FP32 accumulate (fp8 419 vs 838 TFLOPS with FP16 accumulate; bf16 209), so: - fp8 config locks use_fast_accum=True (Float8MMConfig). torchao already defaults it on; pinning it guards consumer cards against a default change. On B200 it is identical speed and slightly better quality (LPIPS 0.050 vs 0.091). - the Blackwell auto ladder prefers fp8 over mxfp8 (measured faster + more accurate). 2:4 semi-structured sparsity evaluated and rejected (scripts/sparse_accum_probe.py): 2:4 magnitude-prune + fp8 gives LPIPS 0.858 (broken image) with no fine-tune, the cuSPARSELt kernel errors on torch 2.9, and it does not compose with torch.compile (our main ~2x). Documented as a dead end, not shipped. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 8): add fp8 fast-accum overflow verification probe scripts/fp8_overflow_check.py hooks every quantised linear during a real Z-Image generation and reports max-abs + non-finite counts for use_fast_accum True vs False. Confirms fast accumulation is an accumulation-precision knob, not an overflow one: across 276 linears, including Z-Image's ~1.0e6 activation peaks (which overflow FP16), 0 non-finite elements and identical max-abs for both modes. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 8): detect consumer vs data-center GPU for fp8 accumulate, with user override Consumer/workstation GPUs (GDDR) halve fp8 FP32-accumulate throughput, so they want fast (FP16) accumulate; data-center HBM parts (B200/H100/A100/L40) are not nerfed and prefer the higher-precision FP32 accumulate. Add _is_consumer_gpu() (token-exact match on the device name per NVIDIA's GPU list, so workstation A4000 != data-center A40; GeForce/TITAN and unknown default to consumer) and gate the fp8 use_fast_accum on it. Measured: fast accumulate is ~2x on consumer Blackwell and ~8% on B200 (0.608 vs 0.665s), no overflow, quality below the quant noise floor. So the default leans to accuracy on data-center; a new request field transformer_quant_fast_accum (null=auto, true/false=force) lets the operator override per load (scripts/diffusion_bench.py --fp8-fast-accum auto|on|off). 187 diffusion tests pass (+ consumer detection, _resolve_fast_accum, and the override threading). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 8): add NVFP4 probe documenting it is not yet a win on torch 2.9 scripts/nvfp4_probe.py measures NVFP4 via torchao on the real Z-Image transformer. Finding (B200, 1024px/8 steps): NVFP4 is a torchao feature and DOES run with use_triton_kernel=False (the default triton path needs the missing MSLK library), but only at bf16-compile rate (0.667s vs fp8 0.592s) -- it dequantises FP4->bf16 rather than using the FP4 tensor cores. The real FP4 speedup needs MSLK or torch>=2.11 + torchao's CUTLASS FP4 GEMM. The smoke probe (default triton=True) already keeps NVFP4 out of auto on this env, so auto correctly stays on fp8; NVFP4 activates automatically once fast. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 8): prefer fp8 over nvfp4 in Blackwell auto ladder Validated NVFP4 on torch 2.11 + torchao CUTLASS FP4 in an isolated env. The FP4 tensor-core GEMM is genuinely active there (a 16384^3 GEMM hits ~3826 TFLOPS, 2.52x bf16 and 1.37x fp8), but it only beats fp8 on very large GEMMs. At the diffusion transformer's shapes (hidden ~3072, MLP ~12288, M~4096) NVFP4 is both slower (0.81x fp8 end to end on Z-Image 1024px) and less accurate (LPIPS 0.166 vs fp8's 0.044). Reorder the Blackwell auto ladder to fp8 before nvfp4 so auto is correct even on a future MSLK-equipped box; nvfp4 stays an explicit opt-in. Add scripts/nvfp4_t211_probe.py (extension diagnostics + GEMM micro + end-to-end). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 9): pre-quantized transformer loading The Phase 8 fast transformer_quant path materialises the dense bf16 transformer on the GPU and torchao-quantises it in place, so its load peak is ~2x GGUF's (~21 vs 13.4 GB) plus a ~12 GB download. Add a pre-quantized branch: quantise once offline (scripts/build_prequant_checkpoint.py) and at runtime build the transformer skeleton on the meta device (accelerate.init_empty_weights) and load_state_dict(assign=True) the quantized weights, so the dense bf16 never touches the GPU. Measured (B200, Z-Image fp8): full-pipeline GPU load peak 21.2 -> 14.6 GB (matching GGUF's 13.4), on-disk 12 -> 6.28 GB, output bit-identical (LPIPS 0.0). It is the same torchao config + min_features filter the runtime path uses, applied ahead of time. New core/inference/diffusion_prequant.py (resolve_prequant_source + load_prequantized_transformer, best-effort, lazy imports). diffusion.py _load_dense_quant_pipeline tries the pre-quant source first and falls back to the dense materialise+quantise path, then to GGUF, so the default is unchanged. DiffusionLoadRequest gains transformer_prequant_path; DiffusionFamily gains an empty prequant_repos map for hosted checkpoints (hosting deferred). Hermetic CPU tests for the resolver, the meta-init+assign loader, and the backend branch selection + fallbacks; GPU verification via scripts/verify_prequant_backend.py. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 10): attention-backend selection Add a selectable attention kernel via the diffusers set_attention_backend dispatcher. Attention is memory-bandwidth bound, so a better kernel is an end-to-end win orthogonal to the linear-weight quantisation (it speeds the QK/PV matmuls torchao never touches) and composes with torch.compile. auto picks the best exact backend for the device: cuDNN fused attention (_native_cudnn) on NVIDIA when a speed profile is active, measured ~1.18x end-to-end on a B200 (Z-Image 1024px/8 steps) with LPIPS ~0.004 vs the default (below the compile/quant noise floor); native SDPA elsewhere and when speed=off (so off stays bit-identical). Explicit native/cudnn/flash/flash3/flash4/sage/ xformers/aiter are honored, and an unavailable kernel falls back to the default rather than failing the load. New core/inference/diffusion_attention.py (normalize + per-device select + apply, best-effort, lazy imports). Set on pipe.transformer BEFORE compile in load_pipeline; attention_backend threads through begin_load / load_pipeline / status like the other load knobs. New request field attention_backend + status field. Hermetic CPU tests for normalize / select policy / apply fallback, plus route threading + 422. Measured via scripts/perf_levers_probe.py. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 11): prefer int8 on consumer GPUs in the auto ladder Consumer / workstation GPUs halve fp8 (and fp16/bf16) FP32-accumulate tensor-core throughput, while int8 runs at full rate (int32 accumulate is not nerfed). Public benchmarks (SDNQ across RTX 3090/4090/5090, AMD, Intel) confirm int8 via torch._int_mm is as fast or faster than fp8 on every consumer part, and the only path on pre-Ada consumer cards without fp8 tensor cores. So when transformer_quant=auto, reorder the arch tier to put int8 first on a consumer/workstation GPU (detected by the existing _is_consumer_gpu name heuristic), while data-center HBM parts keep fp8 first. Pure ladder reorder via _prefer_consumer_scheme; no new flags. Verified non-regression on a B200 (still picks fp8). Hermetic tests for consumer Blackwell/Ada/workstation (-> int8) and data-center Ada/Hopper/Blackwell (-> fp8). * Studio diffusion (Phase 12): First-Block-Cache step caching for many-step DiT Add opt-in step caching (First-Block-Cache) for the diffusion transformer. Across denoise steps a DiT's output settles, so once the first block's residual barely changes the remaining blocks are skipped and their cached output reused. diffusers ships it natively (FirstBlockCacheConfig + transformer.enable_cache, with the standalone apply_first_block_cache hook as a fallback). Measured on Flux.1-dev (28 steps, 1024px): ~1.4x on top of torch.compile (2.83 -> 2.03s) at LPIPS ~0.08 vs the no-cache output, well inside the quality bar. OFF by default and a per-load opt-in: the win scales with step count, so it is for many-step models (Flux / Qwen-Image) and pointless for few-step distilled models (e.g. Z-Image-Turbo at ~8 steps), where a single skipped step is a large fraction of the trajectory. It composes with regional compile only with fullgraph=False (the cache's per-step decision is a torch.compiler.disable graph break), which the speed layer now switches to automatically when a cache is engaged. Best-effort: a model whose block signature the hook does not recognise is caught and the load proceeds uncached. - new core/inference/diffusion_cache.py: normalize_transformer_cache + apply_step_cache (enable_cache / apply_first_block_cache fallback; threshold auto-raised for a quantised transformer per ParaAttention's fp8 guidance; lazy diffusers import). - diffusion_speed.py: apply_speed_optims takes cache_active; compile drops fullgraph when a cache is engaged. - diffusion.py: apply_step_cache before compile; thread transformer_cache / transformer_cache_threshold through begin_load -> load_pipeline and report the engaged mode in status(). - models/inference.py + routes/inference.py: transformer_cache (off | fbcache) and transformer_cache_threshold request fields, engaged mode in the status response. - hermetic tests for normalisation, the enable_cache / hook-fallback paths, threshold selection, and best-effort failure handling, plus route threading + validation. - scripts/fbcache_flux_probe.py: the Flux validation probe (latency / speedup / VRAM / LPIPS vs the compiled no-cache baseline). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 14): fix int8 dense quant on Flux / Qwen (skip M=1 modulation linears) The opt-in dense int8 transformer path crashed on Flux.1 and Qwen-Image with 'torch._int_mm: self.size(0) needs to be greater than 16, but got 1'. int8 dynamic quant goes through torch._int_mm, which requires the activation row count M > 16. A DiT's AdaLN modulation projections (Flux norm1.linear 3072->18432, Qwen img_mod.1 / txt_mod.1, Flux.2 *_modulation.linear) and its timestep / guidance / pooled-text conditioning embedders are computed once from the [batch, dim] conditioning vector (M = batch = 1), not per token, so they hit _int_mm at M=1 and crash. Their feature dims are large, so the existing min_features filter did not exclude them. Fix: the int8 filter now also skips any Linear whose fully-qualified name matches a modulation / conditioning-embedder token (norm, _mod, modulation, timestep_embed, guidance_embed, time_text_embed, pooled). These layers run at M=1 once per block and are a negligible share of the FLOPs, so int8 keeps the full speedup on the attention / FFN layers (M = sequence length). fp8 / nvfp4 / mxfp8 use scaled_mm, which has no M>16 limit and quantises these layers fine, so the exclusion is int8-only. Sequence embedders (context_embedder / x_embedder / txt_in, M = seq) are deliberately not excluded -- note 'context_embedder' contains the substring 'text_embed', which is why the token is the specific 'time_text_embed', not 'text_embed'. Measured on a B200 (1024px, transformer_quant=int8 + speed=default), int8 now runs on every supported model and is the fastest dense path on Flux/Qwen (int8 runs full-rate vs fp8's FP32-accumulate): FLUX.1-dev 9.62s eager -> 1.98s (4.86x, vs fp8 2.15s), Qwen-Image -> 1.87s (5.57x, vs fp8 2.09s), FLUX.1-schnell -> 0.41s (3.59x). Z-Image and Flux.2-klein (already working) are unchanged. - diffusion_transformer_quant.py: add _INT8_EXCLUDE_NAME_TOKENS; make_filter_fn takes exclude_name_tokens; quantize_transformer passes it for int8 only. - hermetic test that the int8 filter excludes the modulation / embedder linears (and keeps attention / FFN / sequence-embedder linears), while fp8 keeps them. - scripts/int8_linear_probe.py: the meta-device probe used to enumerate each transformer's Linear layers and derive the exclusion list. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 15): build int8 pre-quantized checkpoints (skip M=1 modulation linears) The prequant-checkpoint builder applied the dense quant filter without the int8-only M=1 modulation / conditioning-embedder exclusion the runtime path uses, so a built int8 checkpoint baked those projections as int8 and crashed (torch._int_mm needs M>16) at the first denoise step on Flux / Qwen. Factor the scheme->exclusion decision into a shared exclude_tokens_for_scheme() used by both the runtime quantise path and the offline builder so they can never drift, and apply it in build_prequant_checkpoint.py. int8 prequant now produces a working checkpoint on every supported model, giving int8 (the consumer-preferred scheme) the same ~2x load-VRAM and download reduction fp8 already had. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 16): route no-GPU loads to the native sd.cpp engine When no CUDA/ROCm/XPU GPU is available, route diffusion load/generate to the native stable-diffusion.cpp engine instead of diffusers, with diffusers as the guaranteed fallback. On CPU sd.cpp is 1.4-2.8x faster and uses 1.5-2.2x less RAM. - diffusion_engine_router: centralised engine selection (built on the existing select_diffusion_engine), env opt-outs, MPS gating, recorded fallback reason. - sd_cpp_backend (SdCppDiffusionBackend): the diffusers backend method surface backed by sd-cli, with lazy binary install, registry-driven asset fetch, step-progress parsing, and cancellation. - diffusion_families: per-family single-file VAE + text-encoder asset mapping. - sd_cpp_engine: cancellation support (process-group kill + SdCppCancelled). - routes/inference + gpu_arbiter: drive the active engine via the router; the API now reports the active engine and any fallback reason. - tests for the backend, router, route selection, and cancellation. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Phase 16 review fixes: engine-switch unload, sd.cpp error mapping, per-image seeds, Qwen sampler Address review feedback on #6724: - engine router: unload the engine being deactivated on a switch, so the old model is not left resident-but-unreachable (the evictor only targets the active engine). - generate route: sd.cpp execution errors (nonzero exit / timeout / missing output) now map to 500, not 409 (which only means not-loaded / cancelled). - native batch: return per-image seeds and persist the actual seed for each image so every batch image is reproducible. - Qwen-Image native path: apply --sampling-method euler --flow-shift 3 per the stable-diffusion.cpp docs; other families keep sd-cli defaults. - honor speed_mode (native --diffusion-fa) and, off-CPU, memory_mode/cpu_offload offload flags on the native load instead of hardcoding them off. - fail the load when the sd-cli binary is present but not runnable (version() now returns None on exec error / nonzero exit). - size estimate: only treat the transformer asset as a possible local path. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 9): gate request-supplied local prequant paths behind operator opt-in load_prequantized_transformer ends in torch.load(weights_only=False), which executes arbitrary code from the pickle. The transformer_prequant_path load-request field reached that unpickle for any local file an authenticated caller named, so a request could trigger remote code execution. Refuse the source.kind=='path' branch unless the operator sets UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH=1; the first-party hosted-repo checkpoint stays trusted and unaffected. Document the requirement on the API field and add gate tests. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 10): reset the global attention backend on native, gate arch-specific kernels, accept sdpa - apply_attention_backend now restores the native default when no backend is requested or a kernel fails. diffusers keeps a process-wide active attention backend that set_attention_backend updates, and a fresh transformer's processors follow it, so a load that wanted native could silently inherit a backend (e.g. cuDNN) an earlier speed-profile load pinned, breaking the bit-identical/off guarantee. - select_attention_backend drops flash3/flash4 up front when the CUDA capability is below Hopper/Blackwell. diffusers only checks the kernels package at set time, so an explicit request on the wrong card set fine then crashed mid-generation; it now falls back to native. - Add the sdpa alias to the attention_backend Literal so an API request with sdpa (already a valid alias of native) is accepted instead of 422-rejected by Pydantic. - Drop the dead replace('-','_') normalization (no alias uses dashes/underscores). - perf_levers_probe.py output dir is now relative to the script, not a hardcoded path. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 12): only engage FBCache on context-aware transformers; quantized threshold for GGUF - apply_step_cache now engages only via the transformer's native enable_cache (the diffusers CacheMixin path), which exists exactly when the pipeline wraps the transformer call in a cache_context. The standalone apply_first_block_cache fallback installed on non-CacheMixin transformers too (e.g. Z-Image), whose pipeline opens no cache_context, so the load reported transformer_cache=fbcache and then the first generation crashed inside the hook. Such a model now runs uncached per the best-effort contract. - GGUF transformers are quantized (the default Studio load path), so they now use the higher quantized FBCache threshold when the caller leaves it unset, instead of the dense default that could keep the cache from triggering. - fbcache_flux_probe.py: compile cached runs with fullgraph=False (FBCache is a graph break, so fullgraph=True failed warmup and silently measured an eager cached run); output dir is now relative to the script, not a hardcoded path. * Studio diffusion (Phase 11): keep professional RTX cards on the fp8 ladder _is_consumer_gpu treated professional parts (RTX PRO 6000 Blackwell, RTX 6000 Ada) as consumer because their names carry no datacenter token, so the auto ladder moved int8 ahead of fp8 and the fp8 path chose fast accumulate for them. The rest of the backend already classifies these as datacenter/professional (llama_cpp.py _DATACENTER_GPU_RE), so detect the same RTX PRO 6000 / RTX 6000 Ada markers here and keep fp8 first with precise accumulate. Also fix the consumer-Blackwell test to use compute capability (10, 0) instead of (12, 0). * Studio diffusion (Phase 8): tolerate missing torch.float8_e4m3fn in the mxfp8 config Accessing torch.float8_e4m3fn raises AttributeError on a torch build without it (not just TypeError on older torchao), which would break the mxfp8 config helper instead of falling back to the default. Catch both so the fallback is robust. quant_probe.py: same AttributeError fallback; run LPIPS on CPU so the scorer never holds CUDA memory during the per-row VRAM probe; output dir relative to the script. * Studio diffusion (Phase 7): robust backend-flag snapshot/restore and restore on failed speeded load - snapshot_backend_flags reads each flag defensively (getattr + hasattr), so a build/platform missing one (no cuda.matmul on CPU/MPS) still captures the rest instead of skipping the whole snapshot. restore_backend_flags restores each flag independently so one failure can't leave the others leaked process-wide. - load_pipeline restores the flags (and clears the GPU cache) when the build fails after apply_speed_optims mutated the process-wide flags but before _state captured them for unload to restore -- otherwise a failed default/max load left cudnn.benchmark/TF32 on and contaminated later off generations. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 4): enforce the sd-cli timeout while reading output Iterating proc.stdout directly blocks until the stream closes, so a sd-cli that hangs without producing output (or without closing stdout) would never reach proc.wait and the wall-clock timeout was silently bypassed. Drain stdout on a daemon thread and wait on the PROCESS, so the main thread always enforces the timeout and kills a hung process (which closes the pipe and ends the reader). Add a test that times out even when stdout blocks, and make the no-binary test hermetic so a host-installed sd-cli can't leak in. * Studio diffusion (Phase 14): guard the int8 exclusion filter against a None fqn The filter callback can be invoked without a module name, so fqn.lower() would raise AttributeError on None. Fall back to an empty name (nothing matches the exclusion tokens, so the linear is kept) instead of crashing the quantise pass. * Studio diffusion (Phase 16) review fixes: native engine robustness - sd_cpp_backend: stop truncating explicit seeds to 53 bits (mask to int64); a large requested seed was silently collapsed (2**53 -> 0) and distinct seeds aliased to the same image. Random seeds stay 53-bit (JS-safe). - sd_cpp_backend: sanitize empty/whitespace hf_token to None so HfApi/hf_hub fall back to anonymous instead of failing auth on a blank token. - sd_cpp_backend: a superseding load now cancels the in-flight generation, so the old sd-cli can no longer return/persist an image from the previous model. - diffusion_engine_router: run the previous engine's unload() OUTSIDE the lock so a slow 10+ GB free / CUDA sync does not block engine selection. - diffusion_engine_router: probe sd-cli runnability (version()) before committing to native, so a present-but-unrunnable binary falls back to diffusers at selection. - diffusion_device: resolve a torch-free CPU target when torch is unavailable, so a CPU-only install can still reach the native sd.cpp engine instead of failing load. - tests updated for the runnability probe + a not-runnable fallback case. * Studio diffusion (Phase 9) review fixes: prequant safety + validation - SECURITY: a request-supplied local pre-quant path is now unpickled only when it resolves inside an operator-configured ALLOWLIST of directories (UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH = dir[:dir...]). The previous boolean opt-in, once enabled for one trusted checkpoint, allowed torch.load(weights_only=False) on any path a load request named (arbitrary code execution). realpath() blocks symlink escapes; a bare on/off toggle is no longer a wildcard. - Validate the checkpoint's min_features against the runtime Linear filter, so a checkpoint that quantised a different layer set is rejected instead of silently loading a model that mismatches the dense path while reporting the same scheme. - Tolerant base_model_id compare (exact or same final path/repo segment), so a local path or fork of the canonical base is accepted instead of falling back to dense. - _has_meta_tensors uses any(chain(...)) (no intermediate lists). - prequant verify/probe scripts use repo-relative paths (+ env overrides), not the author's absolute /mnt paths. - tests: allowlist-dir opt-in, outside-allowlist refusal, min_features mismatch, fork tail. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 7) review fixes: offload fallback + bench scripts - diffusion_memory: when group offload is unavailable and the plan falls back to whole-module offload, enable VAE tiling (the group plan left it off, but the fallback is the low-VRAM path where the decode spike can OOM). Covers both the group and sequential fallback branches. - perf_verify: include the balanced-vs-off PSNR in the pass/fail condition, so a balanced bit-identity regression actually fails the check instead of exiting 0. - compare_engines: --vae/--llm default to None (were author-absolute /mnt paths), and the load-progress poll has a 30 min deadline instead of looping forever on a hang. - test for the group->model fallback enabling VAE tiling. * Studio diffusion (Phase 8) review fixes: quant compile + nvfp4 path - diffusion: a torchao-quantized transformer is committed only compiled. A dense model resolves to speed_mode=off, which would run the quant eager (~30x slower than the GGUF it replaced), so when transformer_quant engaged and speed resolved to off, promote to default (regional compile); warn loudly if compile still does not engage. - diffusion_transformer_quant: build the nvfp4 config with use_triton_kernel=False so the CUTLASS FP4 path is used (torchao defaults to the Triton kernel, which needs MSLK); otherwise the smoke probe fails on CUTLASS-only Blackwell and silently drops to GGUF. - nvfp4_probe: repo-relative output dir + --out-dir (was an author-absolute /mnt path). - test asserts the eager-quant -> default-compile promotion. * Studio diffusion (Phase 10) review fixes: attention gating + probe isolation - diffusion_attention: gate the auto cuDNN-attention upgrade on SM80+; on pre-Ampere NVIDIA (T4/V100) cuDNN fused SDPA is accepted at set time but fails at first generation, so auto now stays on native SDPA there. - diffusion_attention: _active_attention_backend handles get_active_backend() returning an enum/None (not a tuple); the old unpack always raised and was swallowed, so the native-restore short-circuit never fired. - perf_levers_probe: free the resident pipe on a skipped (attn/fbcache) variant; run LPIPS on CPU so it isn't charged to every variant's peak VRAM; reset force_fuse_int_mm_with_mul so the inductor_flags variant doesn't leak into later compiled rows. - tests for the SM80 cuDNN gate. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 4) review fixes: sd.cpp installer + engine hardening - install_sd_cpp_prebuilt: download the release archive with urlopen + an explicit timeout + copyfileobj (urlretrieve has no timeout and hangs on a stalled socket); extract through a per-member containment check (Zip-Slip guard); expanduser the --install-dir so a tilde path is not taken literally; and on Windows CUDA also fetch the separately-published cudart runtime DLL archive so sd-cli.exe can start. - sd_cpp_engine: find_sd_cpp_binary honors UNSLOTH_STUDIO_HOME / STUDIO_HOME like the installer, so a custom-root install is discovered without UNSLOTH_SD_CPP_PATH; start sd-cli with the parent-death child_popen_kwargs so it is not orphaned on a backend crash; reap the SIGKILLed child (proc.wait) so a cancel/timeout does not leave a zombie. - tests: Zip-Slip rejection, normal extraction, studio-home discovery. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 4) review round 2: collect sd-cli batch outputs Codex review: when batch_count > 1, stable-diffusion.cpp's save_results() writes the numbered files _ (base_0.png, base_1.png, ...) instead of the literal --output path. SdCppEngine.generate checked only the literal path, so a batch generation would exit 0 and then raise 'no image' (or return a stale file). generate now returns the literal path when present and otherwise falls back to the numbered siblings; single-image behavior is unchanged. Test: a fake sd-cli that writes img_0.png/img_1.png (not img.png) is collected without error. * Studio diffusion (Phase 6) review round 2: img2img source dims + upscale repeats Codex review on the native engine arg builder: - build_sd_cpp_command emitted --width/--height unconditionally, so an img2img/inpaint/edit run that left dims unset forced a 1024x1024 resize/crop of the input. width/height are now Optional (None = unset): an image-conditioned run (init_img or ref_images) with unset dims omits the flags so sd.cpp derives the size from the input image (set_width_and_height_if_unset); a plain txt2img run with unset dims keeps the prior 1024x1024 default; explicit dims are always honored. width/height are read only by the builder, so the type change is local. - build_sd_cpp_upscale_command used a truthiness guard (params.repeats and ...) that silently swallowed repeats=0 into sd-cli's default of one pass, turning an explicit no-op into a real upscale. It now rejects repeats < 1 with ValueError and emits the flag for any explicit value != 1. Tests: img2img unset dims omit width/height (init_img and ref_images), explicit dims emitted, txt2img keeps 1024; upscale rejects repeats=0 and omits the flag at the default. (Two pre-existing binary-discovery tests fail only because a real sd-cli is installed in this dev environment; unrelated to this change.) * Studio diffusion (Phase 9) review round 2: correct prequant allowlist doc Codex review: the transformer_prequant_path field description still told operators to enable local checkpoints with UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH=1, but the prior security fix made that variable a directory allowlist -- _allowed_prequant_roots deliberately drops bare on/off toggle tokens (1/true/yes/...). An operator following the documented =1 would have every transformer_prequant_path request silently refused. The description now states it must name one or more allowlisted directories and that a bare on/off value is not accepted. Test: asserts the field help references UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH, does not say =1, and describes an allowlist/directory (guards against doc drift). * Studio diffusion (Phase 10) review round 2: cudnn/flash3 gating + registry reset Codex review on attention-backend selection: - Explicit attention_backend=cudnn skipped the SM80 gate that auto applies, so on pre-Ampere NVIDIA (T4 SM75 / V100 SM70) it set fine then crashed at the first generation with no fallback. select_attention_backend now applies _cudnn_attention_supported() to an explicit cuDNN request too. - flash3 used a minimum-only capability gate (>= SM90), so an explicit flash3 on a Blackwell B200 (SM100) passed and then failed at generation -- FlashAttention 3 is a Hopper-SM90 rewrite with no Blackwell kernel. The arch gate is now a (min, max-exclusive) range: flash3 is SM9x-only, flash4 stays SM100+. - apply_attention_backend's success path left diffusers' process-wide active backend pinned to the kernel it set; a later component whose processors are unconfigured (backend None) would inherit it. It now resets the global registry to native after a successful per-transformer set (the transformer keeps its own backend), best-effort. Also fixed _active_attention_backend: get_active_backend() returns a (name, fn) tuple, so the prior code stringified the tuple and never matched a name, defeating the native-restore short-circuit. Tests: explicit cudnn dropped below SM80; flash3 dropped on SM100 and allowed on SM90; global registry reset after a successful set; _active_attention_backend reads the tuple return. * Studio diffusion (Phase 11) review round 2: keep GH200/B300 on the fp8 ladder Codex review: _DATACENTER_GPU_TOKENS omitted GH200 (Grace-Hopper) and B300 (Blackwell Ultra), though it has the distinct GB200/GB300 superchip tokens. So _is_consumer_gpu returned True for 'NVIDIA GH200 480GB' / 'NVIDIA B300', and the auto ladder moved int8 ahead of fp8 on those data-center parts -- contradicting llama_cpp.py's datacenter regex, which lists both. Added GH200 and B300 so they are treated as data-center class and keep the intended fp8-first behavior. Test: extends the datacenter parametrize with 'NVIDIA B300' and 'NVIDIA GH200 480GB' (now _is_consumer_gpu False). * Studio diffusion (Phase 14) review round 2: apply int8 M=1 exclusion in the builder Codex review: the M=1 modulation/embedder exclusion was wired only into the dense runtime quantiser; the offline builder scripts/build_prequant_checkpoint.py called make_filter_fn(min_features) with no exclusion. So an int8 prequant checkpoint quantised the AdaLN modulation and conditioning-embedder linears, and loading it via transformer_prequant_path (the load path only loads already-quantised tensors, it can't re-skip them) reintroduced the torch._int_mm M=1 crash this phase fixes for the runtime path. Extracted int8_exclude_name_tokens(scheme) as the single source of truth (int8 -> the M=1 exclusion, every other scheme -> none) and use it in both the runtime quantiser and the builder, so a prequant artifact's quantised-layer set always matches the runtime. fp8/fp4/mx artifacts are byte-identical (empty exclusion). Test: int8_exclude_name_tokens returns the exclusion for int8 and () for fp8/nvfp4/mxfp8. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 16) review round 2: native CPU arbiter, status offload, load race Codex review on the native-engine routing: - The /images/load route took the GPU arbiter (acquire_for(DIFFUSION) -> evict chat) unconditionally after engine selection. A native sd.cpp load on a pure-CPU host never touches the GPU, so that needlessly tore down the resident chat model. The handoff is now gated: diffusers always takes it, a force-native sd.cpp load on a CUDA/XPU/MPS box still takes it, but a native sd.cpp load on a CPU host skips it. - sd_cpp status() hardcoded offload_policy 'none' / cpu_offload False even when _run_load computed real offload flags (balanced/low_vram/cpu_offload off-CPU), so the setting was unverifiable. status now derives them from state.offload_flags (still 'none' on CPU, where the flags are empty). - _run_load committed the new state without cancelling/waiting on a generation that started during the (slow) asset download, so a stale sd-cli run against the OLD model could finish afterward and persist an image from the previous model once the new load reported ready. The commit now signals the in-flight cancel and waits on _generate_lock before swapping _state (taken only at commit, so the download never serialises against generation), mirroring the diffusers load path. Tests: CPU native load skips the arbiter while a GPU native load takes it; status reports offload active when flags are set; _run_load cancels and waits for an in-flight generation before committing. * Studio diffusion (Phase 14) review round 2: align helper name with the stack Rename the int8 exclusion helper to exclude_tokens_for_scheme, matching the identical helper already present higher in the diffusion stack (Phase 16). The helper definition, the runtime quantiser call, and the offline builder are now byte-identical to that version, so the two branches no longer introduce a divergent name for the same single-source-of-truth and the stack merges without a conflict on this fix. No behavior change. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion: eager patches + torch.compile cache speed phase Adds the opt-in speed path for the GGUF diffusion transformer behind a selectable speed mode (default off, so output is unchanged until a profile is chosen): - diffusion_eager_patches.py: shared eager fast-paths (channels_last, attention/backend selection, fused norms and QKV) installed at load and rolled back on unload or failed load. - diffusion_compile_cache.py / diffusion_gguf_compile.py: a persistent torch.compile cache and the GGUF-transformer compile wiring. - diffusion_arch_patches.py: architecture-specific patches. - diffusion_patch_backend.py: shared install/restore plumbing. - diffusion_speed.py: speed-profile planning. Tests for each module plus the benchmarking and probe scripts used to measure speed, memory, and accuracy of the path. * Studio diffusion: image workflows (safetensors, image-conditioned, editing) + Images UI Backend: - Load non-GGUF safetensors models: full bnb-4bit pipelines and single-file fp8 transformers, gated to the unsloth org plus a curated allowlist. - Image-conditioned workflows built with Pipeline.from_pipe so they reuse the loaded transformer/VAE/text-encoder with no extra VRAM: img2img, inpaint, outpaint, and a hires-fix upscale pass. - Instruction editing as its own family kind (Qwen-Image-Edit-2511, FLUX.1-Kontext-dev) and FLUX.2-klein reference conditioning (single and multi-reference) plus klein inpaint. - Auto-resize odd-sized inputs to a multiple of 16 (and resize the matched mask) so img2img/inpaint/edit no longer reject non-/16 uploads. Bound the decoded image size and cap upscale output to avoid OOM on large inputs. - Fixes: from_pipe defaulting to a float32 recast that crashed torchao quantized transformers; image-conditioned calls forcing the slider size onto the input image. Native sd.cpp engine rejects image-conditioned and reference requests it cannot serve. Frontend: - Redesigned Images page with capability-gated workflow tabs (Create, Transform, Inpaint, Extend, Upscale, Reference, Edit), a brush mask editor, client-side outpaint, and a multi-reference picker. - Advanced options moved to a right-docked panel mirroring Chat: closed by default, toggled by a single fixed top-bar button that stays in place. sd.cpp installer: pin the release, verify each download's sha256, add a download timeout, and make the source repo configurable for a future mirror. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio Images: correct the Advanced panel comment (closed by default, fixed toggle) * Studio diffusion: LoRA adapters for the Images workflow Add community LoRA support across both diffusion backends, the single biggest step toward broad image-workflow coverage. Backend - New shared module core/inference/diffusion_lora.py: adapter discovery (local scan + curated catalog + owner/name[:file] Hub refs), download via hf_hub_download_with_xet_fallback, alias sanitization, native managed-dir materialization with collision-broken aliases, prompt-tag injection (deduped against user-typed tags), and a supports_lora gate. - Native sd-cli: resolve + materialize selected LoRAs into a per-run managed dir, inject tags, pass --lora-model-dir with --lora-apply-mode auto. The arg builder already emitted these flags. - Diffusers: non-fused load_lora_weights + set_adapters manager, tracked on the pipe so an unchanged selection is a no-op and a model swap resets; cleared on unload. Never fuses (breaks quantized transformers and blocks live weight tweaks). - Gated off where unsupported: torchao fp8/int8 dense, GGUF-via-diffusers, and native Qwen-Image (no LoRA name-conversion branch upstream). - Request contract: optional loras on DiffusionGenerateRequest; empty or omitted is identical to today. supports_lora surfaced in status; chosen LoRAs persisted in gallery recipe metadata. - New GET /api/models/diffusion-loras for the picker (family-filtered). Frontend - Repeatable multi-LoRA picker (adapter select + weight slider 0..2 + remove), gated by the loaded model's supports_lora and family, max 8. Tests - New test_diffusion_lora.py (14): helpers, request validation, native tag/dir wiring, diffusers set_adapters manager, supports_lora matrix. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: do not force diffusers pipelines cross-tagged gguf into the GGUF variant expander Some diffusers image repos (e.g. unsloth/Qwen-Image-2512-unsloth-bnb-4bit) carry a stray "gguf" tag on the Hub but ship no .gguf files. The model search classified them as GGUF from the bare tag, so the picker rendered the GGUF variant expander, which then dead-ended at "No GGUF variants found." Trust the bare gguf tag only when the repo is not a diffusers pipeline; the -GGUF name suffix and real gguf metadata (populated via expand=gguf) remain authoritative, so genuine GGUF repos are unaffected. * Studio Images: load non-curated unsloth/on-device diffusers repos instead of no-op handleModelSelect only loaded curated safetensors ids and GGUF variant picks; any other non-GGUF pick (an on-device diffusers folder, or a future unsloth diffusers image repo surfaced by search) silently did nothing. Treat such a pick as a full diffusers pipeline load when the id is unsloth-hosted or on-device (the backend infers the family + base repo and gates loads to unsloth/* or local paths), and show a clear message otherwise instead of silently ignoring the click. Curated and GGUF paths are unchanged. * Studio Images: keep curated safetensors models in Recommended after download The curated bnb-4bit / fp8 diffusion rows were filtered out of the Images picker's Recommended list once cached (curatedSafetensorsRows dropped anything in downloadedSet), so they vanished from the picker after the first load and could only be found by typing an exact search. The row already renders a downloaded badge, matching how GGUF Recommended rows stay visible when cached. Drop the exclusion so the curated safetensors always list. * Studio diffusion LoRA: sanitize dots out of adapter aliases The LoRA alias is used as the diffusers PEFT adapter name, and PEFT rejects names containing "." (module name can't contain "."). sanitize_alias kept dots, so a LoRA whose filename carries a version tag (e.g. Qwen-Image-2512-Lightning-8steps-V1.0-bf16) failed to apply with a 400. Replace dots too; the alias stays a valid native filename stem. Adds regression coverage for internal dots. * Studio Images: clarify the GGUF transformer-quant Advanced control Renamed the confusing "Transformer quant / GGUF default" control to "GGUF speed mode" with an "Off (run the GGUF)" default, and reworded the hint to state plainly that FP8/INT8/ FP4 load the FULL base model (larger download + more VRAM) rather than re-packing the GGUF, falling back to the GGUF if it can't fit. Behavior unchanged; labels/hint only. * Studio Images: list on-device unsloth diffusion models in the picker The Images picker's On Device tab hid every non-GGUF cached repo whenever a task filter was active, so downloaded unsloth diffusion pipelines (bnb-4bit and FP8 safetensors) never showed up there. List cached repos that pass the task gate, limited under a filter to unsloth-hosted ones so base repos (which fail the diffusion load trust gate) don't appear only to dead-end on click. Chat behavior is unchanged: the task gate still drops image repos there. * Studio: hide single-file image checkpoints from the chat model picker The chat picker treats a cached repo as an image model, and hides it, only when it ships a diffusers model_index.json. Single-file, ComfyUI, and ControlNet image checkpoints (an FP8 Qwen-Image, a z-image safetensors, a Qwen-Image ControlNet) carry none, so they surfaced as loadable chat models. Fall back to resolving the repo id against the known diffusion families, the same resolver the Images backend loads from, so these checkpoints are tagged text-to-image and stay in the Images picker only. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio Images: add the FLUX.2-dev model family Loading unsloth/FLUX.2-dev-GGUF failed because detect_family knew only the Qwen3-based FLUX.2-klein, so FLUX.2-dev (the full, Mistral-based Flux2Pipeline) resolved to nothing and the load errored. Add a flux.2-dev family: Flux2Pipeline + Flux2Transformer2DModel over the black-forest-labs/FLUX.2-dev base repo (gated, reachable with an HF token), with its FLUX.2 32-channel VAE and Mistral text encoder wired for the sd-cli path from the open Comfy-Org/flux2-dev mirror. text-to-image only: diffusers 0.38 ships no Flux2 img2img / inpaint pipeline for dev. Frontend gets sensible dev defaults (28 steps, guidance 4), distinct from klein's turbo defaults. Verified live: GGUF load resolves the family + gated base repo and generates a real 1024x1024 image on GPU. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio Images: clearer error for an unsupported diffusion model When a repo id resolves to no diffusion family the load raised 'Could not infer a diffusion family... Pass family_override (z-image)', which points at an unrelated family and doesn't say what is supported. Replace it with a message that lists the supported families (from a new supported_family_names helper) and notes that video models and image models whose diffusers transformer has no single-file loader are not supported. Applies to both the diffusers and native sd.cpp load paths. Also refreshes two stale family-registry comments that still called FLUX.2-dev omitted. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Remove stray async task scratch outputs committed by mistake * Diffusion LoRA: harden resolution, native tag precedence, and diffusers teardown Address review findings on the LoRA path: - resolve_one: normalise a blank/whitespace hf_token to None (anonymous access) and reject a client-supplied weight file with traversal / absolute path. - resolve_specs: convert FileNotFoundError from an unknown/stale id to ValueError so the route returns 400 instead of a generic 500. - _scan_local: disambiguate local adapters that share a stem (foo.safetensors vs foo.gguf) so each is uniquely addressable. - inject_prompt_tags: the backend-validated weight now wins over a user-typed for a selected adapter; unselected user tags are left alone. - diffusers _apply_loras: reject a .gguf adapter with a clear error before touching the pipe (diffusers loads safetensors only). - _unload_locked: drop the explicit unload_lora_weights() on teardown; the pipe is dropped wholesale (freeing adapters), so the previous call could race an in-flight denoise on the same pipe. - Images page: use a stable LoRA key and clear the selection (not just the options) when the catalog refresh fails. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Diffusion: guard trust check against OSError and validate conditioning inputs - _is_trusted_diffusion_repo: wrap Path.exists() so a repo id with invalid characters (or a bare owner/name id) can't raise OSError; treat any failure as not-a-local-path and fall through to the unsloth/ allowlist. validate_load_request still raises the clear FileNotFoundError for a genuinely missing local pick. - generate(): reject mask_image / upscale / reference_images supplied without an input image, and reject reference_images on a family that does not support reference conditioning, instead of silently degrading to txt2img / img2img. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Address Codex review findings on the image-workflows PR Keep diffusion.py importable without torch: the compile/arch patch modules import torch at module level, so import them lazily at their load/unload call sites instead of at module load. This restores the torchless contract so get_diffusion_backend() works on a CPU/native sd.cpp install. Match family reject keywords and aliases as whole path/name segments, not raw substrings, so an unrelated word like edited, edition, or kontextual no longer misroutes or hides a valid base image model, while supported edit families (Qwen-Image-Edit, FLUX Kontext) still resolve. Mirror the same segment matching in the picker task filter. Route FLUX.2-dev native guidance through --guidance like the other FLUX families rather than --cfg-scale. Reject native upscale requests that have no input image. Read image header dimensions and reject over-limit inputs before decoding pixels, so a crafted small-payload image cannot spike memory. Reject an upscale that would shrink the source below its input size. Validate the model_kind against the filename extension before the GPU handoff. Estimate a local diffusers pipeline's size from its on-disk weights so auto memory planning does not skip offload and OOM. Report workflows: [txt2img] from the native backend status so the Create tab stays enabled for a loaded native model. Clamp the outpaint canvas to the backend's 4096px decode limit. Adds regression tests for segment matching and kind/extension validation. * Harden diffusion LoRA handling on the diffusers and native paths Reject LoRA on a torch.compile'd diffusers transformer (Speed=default/max): diffusers requires the adapter loaded before compilation, so applying one to the already-compiled module fails with adapter-key mismatches. The status gate now hides the picker and generate raises a clear message instead. Convert a cancelled Hub LoRA download (RuntimeError Cancelled) to the diffusion cancellation sentinel in resolve_specs, so an unload/superseding load during resolution maps to a 409 instead of a generic server error. Drop weight-0 LoRA rows before the native support gate so a request carrying only disabled adapters stays a no-op on families where native LoRA is unsupported, matching the diffusers path. Reject duplicate LoRA ids in the request model: both apply paths suffix colliding names, so a repeated id would stack the same adapter past its per-adapter weight bound. Strip all user-typed prompt tags on the native path (only the selected adapters are materialized in the managed lora-model-dir, so an unselected tag can never resolve), and restore saved LoRA selections from a gallery recipe so restore reproduces a LoRA image. * Address further Codex findings on the image-workflows PR - Persist the actual output image size in the gallery recipe instead of the request sliders: Transform/Inpaint/Edit derive the size from the uploaded image, Extend grows the canvas, and Upscale resizes it, so the sliders recorded (and later restored) the wrong dimensions for those workflows. - Reject a remote '*-GGUF' repo loaded as a full pipeline (no single-file name) in validate_load_request, so the unloadable pick fails before chat is evicted rather than deep in from_pretrained. - Only publish an image-conditioned from_pipe wrapper to the shared aux cache when the load is still current: from_pipe runs under the generate lock but not the state lock, so an unload racing its construction could otherwise cache a wrapper over torn-down modules that a later load would reuse. - Verify the Windows CUDA runtime archive checksum before extracting it, like the main sd-cli archive, so a corrupt or tampered runtime is rejected rather than extracted next to the binary. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: oobabooga <112222186+oobabooga@users.noreply.github.com> --- installer_files | 1 + start.sh | 51 +++ studio/backend/core/inference/diffusion.py | 108 +++++ .../backend/core/inference/diffusion_lora.py | 378 ++++++++++++++++++ studio/backend/core/inference/sd_cpp_args.py | 6 + .../backend/core/inference/sd_cpp_backend.py | 164 ++++++-- .../backend/core/inference/sd_cpp_server.py | 8 + studio/backend/models/inference.py | 49 +++ studio/backend/routes/inference.py | 4 + studio/backend/routes/models.py | 34 ++ studio/backend/tests/test_diffusion_lora.py | 375 +++++++++++++++++ studio/backend/tests/test_sd_cpp_backend.py | 84 ++++ studio/frontend/src/features/images/api.ts | 33 ++ .../src/features/images/images-page.tsx | 137 ++++++- 14 files changed, 1395 insertions(+), 37 deletions(-) create mode 120000 installer_files create mode 100755 start.sh create mode 100644 studio/backend/core/inference/diffusion_lora.py create mode 100644 studio/backend/tests/test_diffusion_lora.py diff --git a/installer_files b/installer_files new file mode 120000 index 0000000000..7c77e509a9 --- /dev/null +++ b/installer_files @@ -0,0 +1 @@ +/home/rafael/unsloth/installer_files \ No newline at end of file diff --git a/start.sh b/start.sh new file mode 100755 index 0000000000..803fe1b749 --- /dev/null +++ b/start.sh @@ -0,0 +1,51 @@ +#!/usr/bin/env bash +# Build the studio frontend from local source, then serve it via the REPO's +# studio backend (studio/backend/run.py) run under the installed studio venv, +# so local BACKEND edits take effect (plain `unsloth studio` runs the installed +# package's backend instead). Uses the self-contained Node in +# installer_files/node so the system Node is never touched. +set -euo pipefail + +PORT="${PORT:-8888}" +HOST="${HOST:-0.0.0.0}" + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +FRONTEND_DIR="$SCRIPT_DIR/studio/frontend" +NODE_BIN="$SCRIPT_DIR/installer_files/node/bin" + +if [ ! -x "$NODE_BIN/node" ]; then + echo "ERROR: local Node not found at $NODE_BIN" >&2 + echo "Install it once, e.g.:" >&2 + echo " curl -fL https://nodejs.org/dist/v22.12.0/node-v22.12.0-linux-x64.tar.xz | tar -xJ -C /tmp" >&2 + echo " mv /tmp/node-v22.12.0-linux-x64 $SCRIPT_DIR/installer_files/node" >&2 + exit 1 +fi + +export PATH="$NODE_BIN:$PATH" +echo "==> Using $(node -v) / npm $(npm -v)" + +cd "$FRONTEND_DIR" + +if [ ! -d node_modules ] || ! cmp -s package-lock.json node_modules/.lockfile-stamp; then + echo "==> Installing frontend dependencies" + npm ci + cp package-lock.json node_modules/.lockfile-stamp +fi + +echo "==> Building frontend ($FRONTEND_DIR)" +npm run build + +STUDIO_VENV_PY="${UNSLOTH_STUDIO_HOME:-$HOME/.unsloth/studio}/unsloth_studio/bin/python" +RUN_PY="$SCRIPT_DIR/studio/backend/run.py" + +if [ ! -x "$STUDIO_VENV_PY" ]; then + echo "ERROR: studio venv python not found at $STUDIO_VENV_PY" >&2 + echo "Run 'unsloth studio' once to create the venv, or set UNSLOTH_STUDIO_HOME." >&2 + exit 1 +fi + +echo "==> Starting repo backend on $HOST:$PORT (frontend: $FRONTEND_DIR/dist)" +echo " backend: $RUN_PY" +echo " python: $STUDIO_VENV_PY" +cd "$SCRIPT_DIR/studio/backend" +exec "$STUDIO_VENV_PY" "$RUN_PY" --host "$HOST" --port "$PORT" --frontend "$FRONTEND_DIR/dist" diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index eb4bb33034..c06e3cc359 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -246,6 +246,8 @@ class _LoadState: # Pre-warmed torch.compile cache context (diffusion_compile_cache.CacheContext) when a # compiled tier ran, else None. Carries the per-key inductor dir + bundle for save/restore. compile_cache_ctx: Any = None + # Token kept so LoRA adapters selected at generate time can be fetched from the Hub. + hf_token: Optional[str] = None @dataclass @@ -1145,6 +1147,7 @@ class DiffusionBackend: transformer_cache = cache_engaged, eager_patched = eager_patched, compile_cache_ctx = compile_ctx, + hf_token = hf_token, ) state_committed = True finally: @@ -1398,6 +1401,87 @@ class DiffusionBackend: except (StopIteration, AttributeError, RuntimeError): pass + def _apply_loras( + self, state: Any, loras: Optional[list[tuple[str, float]]], cancel: threading.Event + ) -> None: + """Load + activate requested LoRA adapters on ``state.pipe`` (non-fused), or clear + them when none are requested. + + The applied set is recorded on the pipe object, so an unchanged selection is a no-op + and a model swap (a fresh pipe with no marker) resets naturally. Never fuses: fusing + breaks on quantized (bnb-4bit / torchao) transformers and blocks live weight tweaks. + """ + from core.inference import diffusion_lora + + pipe = state.pipe + current = getattr(pipe, "_unsloth_loras", ()) + specs = [(i, w) for (i, w) in (loras or []) if w != 0] + + if not specs: + if current: + try: + pipe.unload_lora_weights() + except Exception: # noqa: BLE001 -- best-effort clear + pass + pipe._unsloth_loras = () + return + + if not diffusion_lora.supports_lora( + engine = "diffusers", + family = getattr(state.family, "name", None), + model_kind = state.kind, + transformer_quant = state.transformer_quant, + compiled = "compiled" in (getattr(state, "speed_optims", ()) or ()), + ): + raise ValueError( + "LoRA is not supported for this model/quantisation on the diffusers engine " + "(GGUF-via-diffusers, torchao fp8/int8, or a torch.compile'd Speed=default/max " + "load). Use a bf16 or bnb-4bit load at Speed=off/eager, or the native engine " + "for GGUF models." + ) + + resolved = diffusion_lora.resolve_specs(specs, hf_token = state.hf_token, cancel_event = cancel) + # The shared catalog scans both .safetensors and .gguf, but diffusers' + # load_lora_weights only takes safetensors; a .gguf adapter would otherwise fail + # deep in generation. Reject it here as a clean 400 before touching the pipe. + bad = [r.id for r in resolved if r.fmt != "safetensors"] + if bad: + raise ValueError( + "GGUF LoRA adapters are not supported on the diffusers engine " + f"({', '.join(bad)}); use a .safetensors adapter, or the native engine." + ) + # Unique adapter names (diffusers requires distinct names; sanitized stems can collide). + uniq: list[tuple[str, str, float]] = [] + seen: set[str] = set() + for r in resolved: + name = r.alias + n = 1 + while name in seen: + n += 1 + name = f"{r.alias}_{n}" + seen.add(name) + uniq.append((name, r.path, r.weight)) + + desired = tuple(uniq) + if desired == current: + return + try: + if current: + pipe.unload_lora_weights() + for name, path, _weight in uniq: + pipe.load_lora_weights(path, adapter_name = name) + pipe.set_adapters( + [name for name, _p, _w in uniq], adapter_weights = [w for _n, _p, w in uniq] + ) + except Exception as exc: # noqa: BLE001 -- surface as a clean 400 + try: + pipe.unload_lora_weights() + except Exception: # noqa: BLE001 + pass + pipe._unsloth_loras = () + raise ValueError(f"Failed to apply LoRA: {exc}") from exc + pipe._unsloth_loras = desired + @staticmethod def _reset_step_cache(pipe: Any) -> None: """Clear the transformer's stateful step cache (FBCache) before a generation. @@ -1445,6 +1529,9 @@ class DiffusionBackend: # pipeline accepts a list, so multiple references can be combined (subject + style, # character + scene). Ignored by non-reference workflows. reference_images: Optional[list[str]] = None, + # LoRA adapters as (id, weight) pairs; loaded onto the pipe (non-fused) and activated + # with set_adapters for this generation. None/empty = no LoRA (adapters cleared). + loras: Optional[list[tuple[str, float]]] = None, ) -> dict[str, Any]: import torch from PIL import Image @@ -1477,6 +1564,10 @@ class DiffusionBackend: seed = int(seed) generator.manual_seed(seed) + # Apply/adjust LoRA adapters on the resident pipe (non-fused) before picking + # the workflow pipe; from_pipe pipes share the transformer, so it propagates. + self._apply_loras(state, loras, cancel) + # Select the pipeline for this workflow. txt2img uses the loaded pipe; # img2img/inpaint reuse its resident modules via from_pipe (no reload); # an edit model's OWN loaded pipe is already the edit pipeline. @@ -1745,6 +1836,13 @@ class DiffusionBackend: uninstall_patches() uninstall_arch_patches() + # NOTE: we deliberately do NOT call state.pipe.unload_lora_weights() here. unload() + # sets the cancel event but does not take _generate_lock, so a LoRA-backed denoise + # can still be running on this same pipe for up to one more callback; mutating its + # adapter layers now would race that in-flight generation. The whole pipe is dropped + # just below (self._state = None; del state; clear_gpu_cache()), so the adapter + # tensors are freed with it -- no explicit unload is needed for memory or for a + # later load (which builds a fresh pipe). # Drop the workflow pipes built around this load's modules so they don't pin the # freed pipeline (they only re-wire its components, but holding the wrappers # would keep the modules alive past unload). @@ -1775,7 +1873,10 @@ class DiffusionBackend: "attention_backend": None, "transformer_cache": None, "workflows": [], + "supports_lora": False, } + from core.inference import diffusion_lora + return { "loaded": True, "repo_id": state.repo_id, @@ -1797,6 +1898,13 @@ class DiffusionBackend: # Image-conditioned workflows the loaded family supports, so the UI can gate # its tabs. txt2img is always available on the diffusers engine. "workflows": _family_workflows(state.family), + "supports_lora": diffusion_lora.supports_lora( + engine = "diffusers", + family = state.family.name, + model_kind = state.kind, + transformer_quant = state.transformer_quant, + compiled = "compiled" in (getattr(state, "speed_optims", ()) or ()), + ), } diff --git a/studio/backend/core/inference/diffusion_lora.py b/studio/backend/core/inference/diffusion_lora.py new file mode 100644 index 0000000000..7db07dcd60 --- /dev/null +++ b/studio/backend/core/inference/diffusion_lora.py @@ -0,0 +1,378 @@ +"""Shared LoRA support for the Studio diffusion backends. + +Both engines apply LoRA differently -- the native stable-diffusion.cpp CLI selects +adapters by `` prompt tags resolved against a `--lora-model-dir`, +while diffusers loads them with `load_lora_weights()` + `set_adapters()`. This module +holds the parts they share: a curated + local catalog, id->file resolution (with HF +download), a managed directory materialiser for the native tier, deterministic native +alias naming, and the single `supports_lora()` gate the UI/status and both backends use. + +The request layer only ever passes a LoRA *id* (a discovery id, local stem, or HF repo +id) plus a weight -- never a raw filesystem path -- so a client cannot make the backend +read an arbitrary file. Resolution validates the id against the catalog / the local LoRA +directory / the HF hub before anything is loaded. +""" + +from __future__ import annotations + +import os +import re +import threading +from dataclasses import dataclass, field +from pathlib import Path +from typing import Optional + +from utils.hf_xet_fallback import hf_hub_download_with_xet_fallback +from utils.paths.storage_roots import studio_root + +from .diffusion_families import DIFFUSION_CANCELLED_MSG + +# LoRA file formats we accept. sd-cli probes .safetensors/.gguf/.pt; diffusers loads +# .safetensors. We expose safetensors + gguf (pt is legacy/pickled -> excluded for safety). +_NATIVE_EXTS = (".safetensors", ".gguf") +_DIFFUSERS_EXTS = (".safetensors",) +_ALL_EXTS = (".safetensors", ".gguf") + + +@dataclass(frozen = True) +class LoraCatalogEntry: + """One discoverable LoRA adapter.""" + + id: str + display_name: str + source: str # "local" | "hub" + fmt: str # "safetensors" | "gguf" + # Family names this LoRA is compatible with (see diffusion_families). Empty = unknown + # (shown but not family-gated). Used by the UI to grey out incompatible adapters. + families: tuple[str, ...] = () + repo_id: Optional[str] = None # for source == "hub" + weight_name: Optional[str] = None # file within the repo (source == "hub") + local_path: Optional[str] = None # for source == "local" + size_bytes: int = 0 + weight_default: float = 1.0 + + +@dataclass(frozen = True) +class ResolvedLora: + """A LoRA resolved to a concrete local file, ready to apply.""" + + id: str + alias: str # sanitized stem used for the native tag / diffusers adapter name + path: str + fmt: str + weight: float + + +# Curated, family-tagged catalog of known-good diffusion LoRAs. Kept intentionally small +# and data-driven; extend as unsloth hosts/curates more. Entries are HF repos with a +# single-file weight. (Left minimal on purpose -- local discovery is the primary source, +# and users can also reference any public HF LoRA repo id directly.) +_CURATED: tuple[LoraCatalogEntry, ...] = () + + +def loras_dir() -> Path: + """Local directory Studio scans for user-provided diffusion LoRA files.""" + d = studio_root() / "loras" / "diffusion" + d.mkdir(parents = True, exist_ok = True) + return d + + +def sanitize_alias(raw: str) -> str: + """Deterministic, filesystem- and prompt-tag-safe alias from an id/stem. + + The native `` tag resolves NAME as a filename stem, so the alias must + contain no path separators, spaces, colons, or angle brackets. It is also used as the + diffusers PEFT adapter name, which additionally forbids "." (PEFT treats it as a module + path separator), so dots are replaced too -- many real LoRA filenames carry a version + like "V1.0". Collisions across sources are broken by the caller (materialize_native_dir + / the diffusers manager) with a numeric suffix. + """ + stem = raw.rsplit("/", 1)[-1] + for ext in _ALL_EXTS: + if stem.lower().endswith(ext): + stem = stem[: -len(ext)] + break + stem = re.sub(r"[^A-Za-z0-9_-]+", "_", stem).strip("_-") + return stem or "lora" + + +def _scan_local() -> list[LoraCatalogEntry]: + root = loras_dir() + try: + children = sorted(root.iterdir()) + except OSError: + return [] + files = [p for p in children if p.is_file() and p.suffix.lower() in _ALL_EXTS] + # Two files that share a stem but differ in extension (foo.safetensors + foo.gguf) + # would collide on id (== stem), so the frontend select value and resolve_one's + # id->entry lookup could only ever address one of them. Disambiguate a colliding + # stem by keeping the full filename as the id; a unique stem stays the clean stem. + stem_counts: dict[str, int] = {} + for p in files: + stem_counts[p.stem] = stem_counts.get(p.stem, 0) + 1 + entries: list[LoraCatalogEntry] = [] + for p in files: + ext = p.suffix.lower() + try: + size = p.stat().st_size + except OSError: + size = 0 + entry_id = p.name if stem_counts.get(p.stem, 0) > 1 else p.stem + entries.append( + LoraCatalogEntry( + id = entry_id, + display_name = entry_id, + source = "local", + fmt = "gguf" if ext == ".gguf" else "safetensors", + local_path = str(p), + size_bytes = size, + ) + ) + return entries + + +def list_loras(*, family: Optional[str] = None) -> list[LoraCatalogEntry]: + """Return the merged catalog (curated + local), optionally family-filtered. + + Cheap: a single directory scan plus the in-memory curated list. Network is only + touched later, on resolve(), when a hub adapter is actually selected. + """ + merged = list(_CURATED) + _scan_local() + if family: + fam = family.strip().lower() + merged = [e for e in merged if not e.families or fam in {f.lower() for f in e.families}] + # Stable order: local first (user intent), then by display name. + merged.sort(key = lambda e: (e.source != "local", e.display_name.lower())) + return merged + + +def _catalog_by_id() -> dict[str, LoraCatalogEntry]: + return {e.id: e for e in (list(_CURATED) + _scan_local())} + + +def resolve_one( + spec_id: str, + weight: float, + *, + hf_token: Optional[str] = None, + cancel_event: Optional[threading.Event] = None, +) -> ResolvedLora: + """Resolve a request LoRA id + weight to a concrete local file. + + Accepts, in order: a catalog/local id, or a bare HF repo id (``owner/name`` with an + optional ``owner/name:weight_file.safetensors`` suffix). Downloads hub weights via the + shared xet-fallback helper. Raises FileNotFoundError/ValueError on an unresolvable or + unsupported id -- the caller maps that to a clear 400. + """ + # An empty / whitespace token sent verbatim to HfApi triggers an auth error instead + # of falling back to anonymous access; normalise it 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: + if entry.source == "local": + path = entry.local_path or "" + if not path or not os.path.exists(path): + raise FileNotFoundError(f"LoRA '{spec_id}' is no longer present on disk") + return ResolvedLora(spec_id, sanitize_alias(spec_id), path, entry.fmt, weight) + # hub catalog entry + if not entry.repo_id or not entry.weight_name: + raise ValueError(f"LoRA '{spec_id}' has no downloadable weight") + path = hf_hub_download_with_xet_fallback( + entry.repo_id, entry.weight_name, hf_token, cancel_event = cancel_event + ) + return ResolvedLora(spec_id, sanitize_alias(spec_id), path, entry.fmt, weight) + + # Not in the catalog: allow a bare public HF repo id (owner/name[:weight_file]). + if "/" in spec_id: + 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 never resolve outside the HF + # cache dir once handed to the downloader. + if ( + ".." in weight_name + or weight_name.startswith(("/", "\\", "~")) + or "\\" in weight_name + or os.path.isabs(weight_name) + ): + raise ValueError(f"invalid LoRA weight file path '{weight_name}'") + if weight_name is None: + weight_name = _pick_repo_weight_file(repo_id, hf_token) + ext = os.path.splitext(weight_name)[1].lower() + if ext not in _ALL_EXTS: + raise ValueError(f"unsupported LoRA file '{weight_name}' (need .safetensors/.gguf)") + path = hf_hub_download_with_xet_fallback( + repo_id, weight_name, hf_token, cancel_event = cancel_event + ) + fmt = "gguf" if ext == ".gguf" else "safetensors" + return ResolvedLora(spec_id, sanitize_alias(repo_id), path, fmt, weight) + + raise FileNotFoundError( + f"unknown LoRA '{spec_id}': not a local adapter, catalog entry, or HF repo id" + ) + + +def _pick_repo_weight_file(repo_id: str, hf_token: Optional[str]) -> str: + """Pick the single LoRA weight file in an HF repo (prefer safetensors).""" + from huggingface_hub import HfApi + + files = HfApi(token = hf_token).list_repo_files(repo_id) + safes = [f for f in files if f.lower().endswith(".safetensors") and "/" not in f] + if len(safes) == 1: + return safes[0] + # Prefer a filename hinting at a lora, else the first safetensors, else a gguf. + for f in safes: + if "lora" in f.lower(): + return f + if safes: + return safes[0] + ggufs = [f for f in files if f.lower().endswith(".gguf") and "/" not in f] + if ggufs: + return ggufs[0] + raise FileNotFoundError(f"no .safetensors/.gguf LoRA file found in '{repo_id}'") + + +def resolve_specs( + specs: list[tuple[str, float]], + *, + hf_token: Optional[str] = None, + cancel_event: Optional[threading.Event] = None, +) -> list[ResolvedLora]: + """Resolve request (id, weight) pairs, dropping zero-weight entries. + + A stale / unknown id raises FileNotFoundError inside resolve_one; convert it to + ValueError so the route (which maps only ValueError to a 400) reports bad client + input instead of a generic 500. A Hub download can also raise + ``RuntimeError("Cancelled")`` when the user unloads / starts a superseding load + mid-download; convert that to the diffusion cancellation sentinel so the route + maps it to a 409 instead of a generic server error toast.""" + out: list[ResolvedLora] = [] + try: + for spec_id, weight in specs: + if weight == 0: + continue + out.append(resolve_one(spec_id, weight, hf_token = hf_token, cancel_event = cancel_event)) + except FileNotFoundError as exc: + raise ValueError(str(exc)) from exc + except RuntimeError as exc: + if str(exc) == "Cancelled": + raise RuntimeError(DIFFUSION_CANCELLED_MSG) from exc + raise + return out + + +def materialize_native_dir(resolved: list[ResolvedLora], dest: Path) -> list[ResolvedLora]: + """Populate ``dest`` with symlinks (copy fallback) to the resolved LoRA files. + + sd-cli scans ``--lora-model-dir`` and resolves ```` against filenames in + it, so each selected adapter needs a uniquely-named file there. We use a dedicated + managed directory (never a broad cache dir), which keeps the directory scan small and + safe. Returns the resolved list with aliases updated to the (collision-broken) stems + actually written, so the caller injects matching prompt tags. + """ + dest.mkdir(parents = True, exist_ok = True) + used: set[str] = set() + out: list[ResolvedLora] = [] + for r in resolved: + alias = r.alias + n = 1 + while alias in used: + n += 1 + alias = f"{r.alias}_{n}" + used.add(alias) + ext = os.path.splitext(r.path)[1].lower() or ( + ".gguf" if r.fmt == "gguf" else ".safetensors" + ) + link = dest / f"{alias}{ext}" + try: + if link.exists() or link.is_symlink(): + link.unlink() + os.symlink(os.path.realpath(r.path), link) + except OSError: + import shutil + shutil.copy2(r.path, link) + out.append(ResolvedLora(r.id, alias, str(link), r.fmt, r.weight)) + return out + + +_TAG_RE = re.compile(r"]+):([^>]+)>") + + +def inject_prompt_tags(prompt: str, resolved: list[ResolvedLora]) -> str: + """Append `` tags for the selected adapters, using the backend- + validated weights. + + sd-cli strips these tags before they reach the model, so appending them is safe and + deterministic. A selected adapter's weight is validated (0-2) and recorded in the + request/gallery, so the injected tag must WIN over any `` the user + typed. Strip ALL user-typed tags first: only the selected adapters are materialized in + the managed `--lora-model-dir`, so a tag for an unselected alias can never resolve + anyway (sd-cli's extract_and_remove_lora silently removes unresolved tags), and a tag + for a selected alias must not override the validated weight. Then append the validated + tags for the selected adapters. + """ + # Drop every user-typed tag: unselected ones are dead (not in the managed dir) and + # selected ones must not override the validated weight / 0-2 bounds. + cleaned = _TAG_RE.sub("", prompt) + # Collapse whitespace left by stripped tags without disturbing the user's text. + cleaned = re.sub(r"[ \t]{2,}", " ", cleaned).strip() + tags = [f"" for r in resolved] + if not tags: + return cleaned + sep = "" if not cleaned or cleaned.endswith(" ") else " " + return f"{cleaned}{sep}{' '.join(tags)}" + + +def _fmt_weight(w: float) -> str: + # Stable, compact float formatting (no trailing zeros): 1.0 -> "1", 0.75 -> "0.75". + s = f"{w:.4f}".rstrip("0").rstrip(".") + return s or "0" + + +# Families the native sd-cli LoRA name-conversion supports (SD1.5/SD2/SDXL/SD3/FLUX/ +# z-image). Qwen-Image has no LoRA branch in stable-diffusion.cpp -> excluded until +# validated. Matched by substring against the resolved family name. +_NATIVE_LORA_FAMILY_TOKENS = ( + "flux.1", + "flux.2", + "z-image", + "sd1", + "sd2", + "sdxl", + "sd3", + "stable-diffusion", +) +# Diffusers quant schemes that cannot take LoRA cleanly (torchao tensor-subclass weights). +_DIFFUSERS_LORA_BLOCKED_QUANT = ("int8", "fp8", "nvfp4", "mxfp8") + + +def supports_lora( + *, + engine: Optional[str], + family: Optional[str], + model_kind: Optional[str], + transformer_quant: Optional[str], + compiled: bool = False, +) -> bool: + """Single gate for whether the current load can apply LoRA (used by status + backends). + + Native (sd_cpp): GGUF via sd-cli, for the LoRA-capable families only (Qwen excluded). + Diffusers: bf16 or bnb-4bit transformers, but NOT the dense torchao fp8/int8 fast path + (tensor-subclass weights) and NOT GGUF-via-diffusers. A diffusers transformer that was + torch.compile'd at load (Speed=default/max) also can't take a non-hotswap adapter: + diffusers requires the adapter to be loaded BEFORE compilation, so applying one to the + already-compiled module fails with adapter-key mismatches. ``compiled`` is diffusers-only + (the native sd-cli path has no torch compile). + """ + fam = (family or "").lower() + if engine == "sd_cpp": + return any(tok in fam for tok in _NATIVE_LORA_FAMILY_TOKENS) + # diffusers + if model_kind == "gguf": + return False # GGUF diffusers transformer: use the native engine for LoRA + if transformer_quant and transformer_quant.lower() in _DIFFUSERS_LORA_BLOCKED_QUANT: + return False + if compiled: + return False # can't load an adapter onto an already-compiled transformer + return True diff --git a/studio/backend/core/inference/sd_cpp_args.py b/studio/backend/core/inference/sd_cpp_args.py index ae729a881a..2edd230f26 100644 --- a/studio/backend/core/inference/sd_cpp_args.py +++ b/studio/backend/core/inference/sd_cpp_args.py @@ -412,6 +412,7 @@ def build_img_gen_request( cfg_scale: Optional[float] = None, distilled_guidance: Optional[float] = None, output_format: str = "png", + lora: Optional[list[dict]] = None, ) -> dict: """Build the ``POST /sdcpp/v1/img_gen`` JSON body for one text-to-image request. @@ -453,6 +454,11 @@ def build_img_gen_request( req["seed"] = int(seed) if sample_params: req["sample_params"] = sample_params + # Structured LoRA list: the sdcpp API resolves each ``path`` against the server's + # scanned ``--lora-model-dir`` (prompt-embedded ```` tags are intentionally + # unsupported server-side), so LoRAs must be staged there and named 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 757a28a3a0..059b85974c 100644 --- a/studio/backend/core/inference/sd_cpp_backend.py +++ b/studio/backend/core/inference/sd_cpp_backend.py @@ -213,6 +213,8 @@ class _SdState: flow_shift: Optional[float] = None server: Optional[SdCppServer] = None mode: str = "server" + # Token kept so LoRA adapters selected at generate time can be fetched from the Hub. + hf_token: Optional[str] = None def _memory_policy(memory_mode: Optional[str], cpu_offload: bool) -> str: @@ -571,6 +573,7 @@ class SdCppDiffusionBackend: flow_shift = fam.sd_cpp_flow_shift, server = server, mode = mode, + hf_token = hf_token, ) with self._lock: if self._load_token != _load_token: @@ -706,11 +709,17 @@ class SdCppDiffusionBackend: upscale: Optional[float] = None, # Reference workflow is GPU/diffusers-only (FLUX.2); accepted for interface parity. reference_images: Optional[list[str]] = None, + # LoRA adapters as (id, weight) pairs; resolved up front, then applied per engine + # path: prompt tags for one-shot sd-cli, structured `lora` entries + # for the resident sd-server. None/empty = no LoRA. + loras: Optional[list[tuple[str, float]]] = None, ) -> dict[str, Any]: import tempfile from PIL import Image + from core.inference import diffusion_lora + if ( init_image is not None or mask_image is not None @@ -748,6 +757,27 @@ class SdCppDiffusionBackend: else: seed = int(seed) cfg_scale, flux_guidance = _map_guidance(state.family, guidance) + # Resolve any selected LoRA adapters up front (downloads land in the HF + # cache; a bad id fails here as a clear 400 before we generate). Drop + # weight-0 rows BEFORE the support gate: weight 0 disables an adapter, so a + # request carrying only disabled rows stays a no-op even on a family where + # native LoRA is unsupported, rather than 400 on a dead selection. + lora_resolved: list = [] + active_loras = [(i, w) for (i, w) in (loras or []) if w != 0] + if active_loras: + if not diffusion_lora.supports_lora( + engine = "sd_cpp", + family = state.family.name, + model_kind = "gguf", + transformer_quant = None, + ): + raise ValueError( + f"LoRA is not supported for {state.family.name} on the native " + "sd.cpp engine." + ) + lora_resolved = diffusion_lora.resolve_specs( + active_loras, hf_token = state.hf_token, cancel_event = cancel + ) self._gen = _SdGen(total_steps = int(steps)) if state.mode == "server" and state.server is not None: images, seeds = self._generate_server( @@ -761,6 +791,7 @@ class SdCppDiffusionBackend: batch_size = batch_size, cfg_scale = cfg_scale, flux_guidance = flux_guidance, + lora_resolved = lora_resolved, cancel = cancel, ) else: @@ -775,6 +806,7 @@ class SdCppDiffusionBackend: batch_size = batch_size, cfg_scale = cfg_scale, flux_guidance = flux_guidance, + lora_resolved = lora_resolved, cancel = cancel, ) if cancel.is_set(): @@ -808,6 +840,7 @@ class SdCppDiffusionBackend: batch_size: int, cfg_scale: Optional[float], flux_guidance: Optional[float], + lora_resolved: list, cancel: threading.Event, ) -> tuple[list, list[int]]: """Generate via the resident sd-server (no model reload). @@ -818,11 +851,20 @@ class SdCppDiffusionBackend: signed-int64 range (the request model / diffusers accept larger seeds), and each chunk is submitted at base+offset so the per-image seeds stay reproducible. Each chunk gets a timeout proportional to its image count so a slow CPU batch is not - cancelled partway through on one fixed deadline.""" + cancelled partway through on one fixed deadline. + + LoRA on the server goes through the structured ``lora`` request field, NOT prompt + tags (the sdcpp API intentionally ignores ```` in the prompt). Selected + adapters are staged into the server's ``--lora-model-dir`` scratch dir, which the + server rescans per request, and referenced by their staged filename.""" import io + import os + import shutil from PIL import Image + from core.inference import diffusion_lora + assert state.server is not None total = max(1, int(batch_size)) # sd.cpp's image seed is signed int64; mask the base (and every derived seed) so a @@ -830,40 +872,60 @@ class SdCppDiffusionBackend: base_seed = int(seed) & ((1 << 63) - 1) images: list = [] seeds: list[int] = [] - for offset in range(0, total, _MAX_SERVER_BATCH): - if cancel.is_set(): - raise SdCppCancelled("sd-server generation was cancelled.") - count = min(_MAX_SERVER_BATCH, total - offset) - chunk_seed = (base_seed + offset) & ((1 << 63) - 1) - payload = build_img_gen_request( - prompt = prompt, - negative_prompt = negative_prompt or None, - width = int(width), - height = int(height), - steps = int(steps), - seed = chunk_seed, - batch_count = count, - sample_method = state.sampling_method, - flow_shift = state.flow_shift, - cfg_scale = cfg_scale, - distilled_guidance = flux_guidance, - ) - blobs = state.server.img_gen( - payload, - on_step = self._on_log, - cancel_event = cancel, - total_timeout = _SERVER_PER_IMAGE_TIMEOUT_S * count, - ) - # All-or-nothing per chunk, like the one-shot path: if the server returns fewer - # blobs than requested (e.g. one image in the batch failed to encode), fail - # rather than silently dropping images from the user's requested batch. - if not cancel.is_set() and len(blobs) != count: - raise RuntimeError( - f"sd-server returned {len(blobs)} of {count} requested images in the batch." + # Stage selected LoRAs into a per-request subdir of the server's lora-model-dir so a + # previous request's adapters can't leak into this one; reference them by the path + # relative to that dir (what the server's recursive scan resolves against). The + # subdir is removed after the batch. supports_lora already gated the family upstream. + lora_payload: Optional[list[dict]] = None + lora_stage: Optional[Path] = None + if lora_resolved: + server_lora_dir = state.server.lora_dir + if server_lora_dir: + lora_stage = Path(server_lora_dir) / f"gen_{os.urandom(6).hex()}" + materialized = diffusion_lora.materialize_native_dir(lora_resolved, lora_stage) + lora_payload = [ + {"path": f"{lora_stage.name}/{Path(m.path).name}", "multiplier": float(m.weight)} + for m in materialized + ] + try: + for offset in range(0, total, _MAX_SERVER_BATCH): + if cancel.is_set(): + raise SdCppCancelled("sd-server generation was cancelled.") + count = min(_MAX_SERVER_BATCH, total - offset) + chunk_seed = (base_seed + offset) & ((1 << 63) - 1) + payload = build_img_gen_request( + prompt = prompt, + negative_prompt = negative_prompt or None, + width = int(width), + height = int(height), + steps = int(steps), + seed = chunk_seed, + batch_count = count, + sample_method = state.sampling_method, + flow_shift = state.flow_shift, + cfg_scale = cfg_scale, + distilled_guidance = flux_guidance, + lora = lora_payload, ) - images.extend(Image.open(io.BytesIO(b)).convert("RGB") for b in blobs) - # sd.cpp advances the seed per image within a job, so report chunk_seed+i. - seeds.extend((chunk_seed + i) & ((1 << 63) - 1) for i in range(len(blobs))) + blobs = state.server.img_gen( + payload, + on_step = self._on_log, + cancel_event = cancel, + total_timeout = _SERVER_PER_IMAGE_TIMEOUT_S * count, + ) + # All-or-nothing per chunk, like the one-shot path: if the server returns fewer + # blobs than requested (e.g. one image in the batch failed to encode), fail + # rather than silently dropping images from the user's requested batch. + if not cancel.is_set() and len(blobs) != count: + raise RuntimeError( + f"sd-server returned {len(blobs)} of {count} requested images in the batch." + ) + images.extend(Image.open(io.BytesIO(b)).convert("RGB") for b in blobs) + # sd.cpp advances the seed per image within a job, so report chunk_seed+i. + seeds.extend((chunk_seed + i) & ((1 << 63) - 1) for i in range(len(blobs))) + finally: + if lora_stage is not None: + shutil.rmtree(lora_stage, ignore_errors = True) return images, seeds def _generate_oneshot( @@ -879,13 +941,21 @@ class SdCppDiffusionBackend: batch_size: int, cfg_scale: Optional[float], flux_guidance: Optional[float], + lora_resolved: list, cancel: threading.Event, ) -> tuple[list, list[int]]: - """Fallback path: re-run one-shot sd-cli per image (reloads the model each time).""" + """Fallback path: re-run one-shot sd-cli per image (reloads the model each time). + + LoRA on the one-shot path uses sd-cli's own mechanism: materialize the selected + adapters into a ``--lora-model-dir`` and inject matching ```` tags + into the prompt (sd-cli parses and strips them). supports_lora already gated the + family upstream, so a non-empty ``lora_resolved`` is safe to apply here.""" import tempfile from PIL import Image + from core.inference import diffusion_lora + engine = self._resolve_engine() extra_args: list[str] = [] if state.vae_format: @@ -896,6 +966,17 @@ class SdCppDiffusionBackend: images = [] seeds: list[int] = [] with tempfile.TemporaryDirectory(prefix = "sdcpp_gen_") as tmpdir: + # Materialize selected LoRAs into a managed dir sd-cli can scan, and inject + # matching tags into the prompt (deduped against any the user + # typed). Empty -> prompt/dir unchanged. + eff_prompt = prompt + lora_dir: Optional[str] = None + if lora_resolved: + materialized = diffusion_lora.materialize_native_dir( + lora_resolved, Path(tmpdir) / "loras" + ) + eff_prompt = diffusion_lora.inject_prompt_tags(prompt, materialized) + lora_dir = str(Path(tmpdir) / "loras") for index in range(max(1, int(batch_size))): if cancel.is_set(): raise RuntimeError(DIFFUSION_CANCELLED_MSG) @@ -906,7 +987,7 @@ class SdCppDiffusionBackend: seed_i = (seed + index) & ((1 << 63) - 1) out_path = str(Path(tmpdir) / f"img_{index}.png") params = SdCppGenParams( - prompt = prompt, + prompt = eff_prompt, negative_prompt = negative_prompt or None, width = int(width), height = int(height), @@ -916,6 +997,8 @@ class SdCppDiffusionBackend: seed = seed_i, sampling_method = state.sampling_method, batch_count = 1, + lora_dir = lora_dir, + lora_apply_mode = "auto" if lora_dir else None, ) engine.generate( state.files, @@ -1024,8 +1107,11 @@ class SdCppDiffusionBackend: "transformer_cache": None, "engine": "sd_cpp", "native_mode": None, + "supports_lora": False, "workflows": [], } + from core.inference import diffusion_lora + return { "loaded": True, "repo_id": state.repo_id, @@ -1048,6 +1134,12 @@ class SdCppDiffusionBackend: "attention_backend": None, "transformer_cache": None, "engine": "sd_cpp", + "supports_lora": diffusion_lora.supports_lora( + engine = "sd_cpp", + family = state.family.name, + model_kind = "gguf", + transformer_quant = None, + ), # "server" = resident sd-server (load once); "oneshot" = legacy per-image sd-cli. "native_mode": state.mode, # The native engine supports plain text-to-image only (generate() rejects diff --git a/studio/backend/core/inference/sd_cpp_server.py b/studio/backend/core/inference/sd_cpp_server.py index 17b269e71d..d66f3cdded 100644 --- a/studio/backend/core/inference/sd_cpp_server.py +++ b/studio/backend/core/inference/sd_cpp_server.py @@ -121,6 +121,14 @@ class SdCppServer: def base_url(self) -> str: return f"http://{self.host}:{self.port}" + @property + def lora_dir(self) -> Optional[str]: + """The server's ``--lora-model-dir`` scratch dir. Adapters staged here are picked + up per request (the sdcpp ``img_gen`` route refreshes its LoRA scan each call), so + server-mode LoRA works by materializing here and referencing the files via the + structured ``lora`` request field (the server ignores ```` prompt tags).""" + return self._scratch_dir + def is_alive(self) -> bool: return self._process is not None and self._process.poll() is None diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 41faf29303..d3feff32f7 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -1813,6 +1813,22 @@ class DiffusionLoadRequest(BaseModel): return value.strip().lower() if isinstance(value, str) else value +class LoraSpec(BaseModel): + """One LoRA adapter to apply for a generation, referenced by its discovery id. + + The id is resolved against the backend's own LoRA catalog + local scan (see + core/inference/diffusion_lora.py); the client never supplies a raw filesystem + path, so an arbitrary file can't be loaded. Weight 0 disables the adapter. + """ + + id: str = Field( + ..., min_length = 1, max_length = 512, description = "LoRA discovery id (repo id or local stem)" + ) + weight: float = Field( + 1.0, ge = 0.0, le = 2.0, description = "Adapter strength; 0 disables, 1.0 is full strength" + ) + + class DiffusionGenerateRequest(BaseModel): """Request to generate one image from the loaded diffusion model.""" @@ -1873,6 +1889,30 @@ class DiffusionGenerateRequest(BaseModel): description = "Additional reference images (base64/data-URL) for the FLUX.2 reference " "workflow, combined with init_image. Up to 3; ignored by other workflows.", ) + loras: Optional[list[LoraSpec]] = Field( + None, + max_length = 8, + description = "LoRA adapters to apply for this generation (by discovery id + weight). " + "Omitted/empty applies none and behaves exactly as before. Rejected with a clear " + "message when the loaded model or its quantisation can't apply LoRA.", + ) + + @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 as several distinct adapters and stack + # its effect past the per-adapter weight bound. The UI already prevents duplicates; + # reject them for API clients too so each adapter takes effect at most once. + if value: + seen: set[str] = set() + for spec in value: + if spec.id in seen: + raise ValueError( + f"duplicate LoRA id '{spec.id}'; list each adapter at most once" + ) + seen.add(spec.id) + return value @field_validator("reference_images") @classmethod @@ -1913,6 +1953,9 @@ class GalleryImage(BaseModel): 1, description = "Batch size used; with batch_index it lets restore replay this image" ) model: Optional[str] = Field(None, description = "Model repo id that produced it") + loras: list[str] = Field( + default_factory = list, description = "LoRA adapters applied, formatted as 'id:weight'" + ) created_at: float = Field(..., description = "Creation time (epoch seconds)") @@ -2002,6 +2045,12 @@ class DiffusionStatusResponse(BaseModel): None, description = "Why diffusers was chosen over the native sd.cpp engine (null when none)", ) + supports_lora: bool = Field( + False, + description = "Whether the loaded model + quantisation can apply LoRA adapters (drives the " + "LoRA picker's enabled state). False on unsupported families/quant (e.g. torchao fp8/int8 " + "dense, GGUF-via-diffusers, or Qwen-Image on the native engine).", + ) # ── OpenAI-compatible images API (POST /v1/images/generations) ── diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 2c0c2bf6b0..929aa19afa 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -11161,6 +11161,7 @@ async def generate_diffusion_image( strength = request.strength, upscale = request.upscale, reference_images = request.reference_images, + loras = [(l.id, l.weight) for l in request.loras] if request.loras else None, ) except ValueError as exc: # Bad client input (undecodable image/mask, or a workflow the loaded family @@ -11220,6 +11221,9 @@ async def generate_diffusion_image( # needs the original batch_size: persist it so restore can replay. "batch_size": request.batch_size, "model": result.get("repo_id"), + "loras": ( + [f"{l.id}:{l.weight:g}" for l in request.loras] if request.loras else [] + ), "created_at": created_at, }, ) diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index cfba8a1181..b22d47c545 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -2031,6 +2031,40 @@ async def scan_loras( ) +@router.get("/diffusion-loras") +async def scan_diffusion_loras( + family: Optional[str] = Query( + default = None, description = "Filter to LoRAs compatible with this diffusion family" + ), + current_subject: str = Depends(get_current_subject), +): + """List diffusion image LoRA adapters for the Images workflow. + + Merges the curated catalog with local files in ``/loras/diffusion``, + optionally filtered to the loaded model's family. Cheap: one directory scan, no network + (a hub adapter is only downloaded when actually selected for a generation). Distinct from + ``/loras`` above, which lists trained/exported TEXT adapters. + """ + from core.inference import diffusion_lora + + entries = diffusion_lora.list_loras(family = family) + return { + "loras": [ + { + "id": e.id, + "display_name": e.display_name, + "source": e.source, + "format": e.fmt, + "families": list(e.families), + "size_bytes": e.size_bytes, + "weight_default": e.weight_default, + } + for e in entries + ], + "loras_dir": str(diffusion_lora.loras_dir()), + } + + def _is_path_under(path: Path, root: Path) -> bool: try: path.resolve().relative_to(root.resolve()) diff --git a/studio/backend/tests/test_diffusion_lora.py b/studio/backend/tests/test_diffusion_lora.py new file mode 100644 index 0000000000..f99739eff9 --- /dev/null +++ b/studio/backend/tests/test_diffusion_lora.py @@ -0,0 +1,375 @@ +"""Tests for diffusion LoRA support: the shared helpers, request-model validation, the +native prompt-tag/dir wiring, and the diffusers set_adapters manager.""" + +from __future__ import annotations + +import os +import types +from pathlib import Path + +import pytest + +from core.inference import diffusion_lora as dl + + +# ── Pure helpers ──────────────────────────────────────────────────────────── + + +def test_sanitize_alias_strips_path_ext_and_unsafe_chars(): + assert dl.sanitize_alias("My Cool/LoRA v2.safetensors") == "LoRA_v2" + 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. + assert ( + dl.sanitize_alias("Qwen-Image-2512-Lightning-8steps-V1.0-bf16") + == "Qwen-Image-2512-Lightning-8steps-V1_0-bf16" + ) + assert "." not in dl.sanitize_alias("model.v1.0.safetensors") + + +def test_inject_prompt_tags_appends_with_spacing(): + r = dl.ResolvedLora("id", "style", "/p.safetensors", "safetensors", 0.8) + assert dl.inject_prompt_tags("a cat", [r]) == "a cat " + # weight formatting: 1.0 -> "1", trailing zeros trimmed + r1 = dl.ResolvedLora("id", "s", "/p", "safetensors", 1.0) + assert dl.inject_prompt_tags("x", [r1]) == "x " + + +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. + 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. + out = dl.inject_prompt_tags("a cat ", [r]) + assert "" not in out + assert out == "a cat " + + +def test_inject_prompt_tags_empty_returns_prompt(): + assert dl.inject_prompt_tags("hello", []) == "hello" + + +def test_supports_lora_matrix(): + # native: flux/z-image yes, qwen no + assert dl.supports_lora( + engine = "sd_cpp", family = "flux.1", model_kind = "gguf", transformer_quant = None + ) + assert dl.supports_lora( + engine = "sd_cpp", family = "z-image", model_kind = "gguf", transformer_quant = None + ) + assert not dl.supports_lora( + engine = "sd_cpp", family = "qwen-image", model_kind = "gguf", transformer_quant = None + ) + # diffusers: bf16 yes, fp8/int8 dense no, gguf-diffusers no + assert dl.supports_lora( + engine = "diffusers", family = "flux.1", model_kind = "pipeline", transformer_quant = None + ) + assert dl.supports_lora( + engine = "diffusers", family = "flux.1", model_kind = "single_file", transformer_quant = None + ) + assert not dl.supports_lora( + engine = "diffusers", family = "flux.1", model_kind = "single_file", transformer_quant = "fp8" + ) + assert not dl.supports_lora( + engine = "diffusers", family = "flux.1", model_kind = "single_file", transformer_quant = "int8" + ) + 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. + assert not dl.supports_lora( + engine = "diffusers", + family = "flux.1", + model_kind = "pipeline", + transformer_quant = None, + compiled = True, + ) + # compiled is diffusers-only; the native path ignores it. + assert dl.supports_lora( + engine = "sd_cpp", + family = "flux.1", + model_kind = "gguf", + transformer_quant = None, + compiled = True, + ) + + +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. + def _boom(spec_id, weight, **kw): + raise RuntimeError("Cancelled") + + monkeypatch.setattr(dl, "resolve_one", _boom) + with pytest.raises(RuntimeError) as ei: + dl.resolve_specs([("a", 1.0)]) + assert str(ei.value) == dl.DIFFUSION_CANCELLED_MSG + + # A non-cancellation RuntimeError is left untouched. + def _other(spec_id, weight, **kw): + raise RuntimeError("disk full") + + monkeypatch.setattr(dl, "resolve_one", _other) + with pytest.raises(RuntimeError) as ei2: + dl.resolve_specs([("a", 1.0)]) + assert str(ei2.value) == "disk full" + + +def test_materialize_native_dir_symlinks_and_breaks_collisions(tmp_path): + a = tmp_path / "a.safetensors" + a.write_bytes(b"x") + b = tmp_path / "sub" + b.mkdir() + b2 = b / "a.safetensors" # same stem as `a` -> alias collision + b2.write_bytes(b"y") + resolved = [ + dl.ResolvedLora("a", "a", str(a), "safetensors", 1.0), + dl.ResolvedLora("a2", "a", str(b2), "safetensors", 0.5), + ] + dest = tmp_path / "managed" + out = dl.materialize_native_dir(resolved, dest) + aliases = [r.alias for r in out] + assert aliases == ["a", "a_2"] # collision broken + for r in out: + assert os.path.exists(r.path) + assert Path(r.path).parent == dest + + +def test_list_loras_scans_local(tmp_path, monkeypatch): + d = tmp_path / "loras" + d.mkdir() + (d / "mystyle.safetensors").write_bytes(b"x") + (d / "other.gguf").write_bytes(b"y") + (d / "ignore.txt").write_bytes(b"z") + monkeypatch.setattr(dl, "loras_dir", lambda: d) + ids = {e.id for e in dl.list_loras()} + assert ids == {"mystyle", "other"} + fmts = {e.id: e.fmt for e in dl.list_loras()} + assert fmts["other"] == "gguf" and fmts["mystyle"] == "safetensors" + + +def test_resolve_one_local_and_unknown(tmp_path, monkeypatch): + d = tmp_path / "loras" + d.mkdir() + (d / "mystyle.safetensors").write_bytes(b"x") + monkeypatch.setattr(dl, "loras_dir", lambda: d) + r = dl.resolve_one("mystyle", 0.7) + assert r.path.endswith("mystyle.safetensors") and r.weight == 0.7 + with pytest.raises(FileNotFoundError): + dl.resolve_one("does-not-exist", 1.0) + + +def test_resolve_specs_drops_zero_weight(tmp_path, monkeypatch): + d = tmp_path / "loras" + d.mkdir() + (d / "a.safetensors").write_bytes(b"x") + monkeypatch.setattr(dl, "loras_dir", lambda: d) + out = dl.resolve_specs([("a", 0.0), ("a", 1.0)]) + assert len(out) == 1 and out[0].weight == 1.0 + + +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. + d = tmp_path / "loras" + d.mkdir() + monkeypatch.setattr(dl, "loras_dir", lambda: d) + with pytest.raises(ValueError): + dl.resolve_specs([("nope", 1.0)]) + + +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. + d = tmp_path / "loras" + d.mkdir() + (d / "foo.safetensors").write_bytes(b"x") + (d / "foo.gguf").write_bytes(b"y") + (d / "solo.safetensors").write_bytes(b"z") + monkeypatch.setattr(dl, "loras_dir", lambda: d) + by_id = {e.id: e for e in dl.list_loras()} + assert "foo.safetensors" in by_id and "foo.gguf" in by_id + assert by_id["foo.safetensors"].fmt == "safetensors" + assert by_id["foo.gguf"].fmt == "gguf" + assert "solo" in by_id # unique stem is untouched + + +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). + monkeypatch.setattr(dl, "loras_dir", lambda: tmp_path) + for bad in ("owner/name:../secret.safetensors", "owner/name:/etc/x.safetensors"): + with pytest.raises(ValueError): + dl.resolve_one(bad, 1.0) + + +# ── Request-model validation ──────────────────────────────────────────────── + + +def test_lora_spec_and_request_validation(): + from models.inference import DiffusionGenerateRequest, LoraSpec + + # empty / missing loras -> unchanged behaviour + assert DiffusionGenerateRequest(prompt = "x").loras is None + req = DiffusionGenerateRequest( + prompt = "x", loras = [{"id": "a", "weight": 0.5}, {"id": "b", "weight": 1.0}] + ) + assert [l.id for l in req.loras] == ["a", "b"] + # weight bounds enforced + with pytest.raises(Exception): + LoraSpec(id = "a", weight = 3.0) + with pytest.raises(Exception): + 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. + with pytest.raises(Exception): + DiffusionGenerateRequest( + prompt = "x", loras = [{"id": "a", "weight": 0.5}, {"id": "a", "weight": 1.0}] + ) + + +# ── Diffusers apply manager ───────────────────────────────────────────────── + + +class _FakePipe: + def __init__(self): + self.loaded: list[tuple[str, str]] = [] + self.active = None + self.unloaded = 0 + + def load_lora_weights( + self, + path, + adapter_name = None, + ): + self.loaded.append((path, adapter_name)) + + def set_adapters( + self, + names, + adapter_weights = None, + ): + self.active = (list(names), list(adapter_weights) if adapter_weights else None) + + def unload_lora_weights(self): + self.unloaded += 1 + self.loaded = [] + self.active = None + + +def _fake_state( + pipe, + *, + kind = "pipeline", + quant = None, +): + fam = types.SimpleNamespace(name = "flux.1") + return types.SimpleNamespace( + pipe = pipe, family = fam, kind = kind, transformer_quant = quant, hf_token = None + ) + + +def _backend(): + from core.inference.diffusion import DiffusionBackend + return DiffusionBackend() + + +def test_diffusers_apply_loads_and_sets_adapters(monkeypatch): + import threading + + monkeypatch.setattr( + dl, + "resolve_specs", + lambda specs, **_: [ + dl.ResolvedLora(i, dl.sanitize_alias(i), f"/{i}.safetensors", "safetensors", w) + for i, w in specs + ], + ) + pipe = _FakePipe() + _backend()._apply_loras( + _fake_state(pipe), [("styleA", 0.8), ("styleB", 1.0)], threading.Event() + ) + assert [n for _p, n in pipe.loaded] == ["styleA", "styleB"] + assert pipe.active == (["styleA", "styleB"], [0.8, 1.0]) + assert getattr(pipe, "_unsloth_loras") # marker recorded + + +def test_diffusers_apply_noop_when_unchanged(monkeypatch): + import threading + + monkeypatch.setattr( + dl, + "resolve_specs", + lambda specs, **_: [ + dl.ResolvedLora(i, dl.sanitize_alias(i), f"/{i}.safetensors", "safetensors", w) + for i, w in specs + ], + ) + pipe = _FakePipe() + b = _backend() + b._apply_loras(_fake_state(pipe), [("styleA", 0.8)], threading.Event()) + first_loaded = list(pipe.loaded) + b._apply_loras(_fake_state(pipe), [("styleA", 0.8)], threading.Event()) + assert pipe.loaded == first_loaded # not reloaded + assert pipe.unloaded == 0 + + +def test_diffusers_apply_clears_when_empty(monkeypatch): + import threading + + monkeypatch.setattr( + dl, + "resolve_specs", + lambda specs, **_: [ + dl.ResolvedLora(i, dl.sanitize_alias(i), f"/{i}.safetensors", "safetensors", w) + for i, w in specs + ], + ) + pipe = _FakePipe() + b = _backend() + b._apply_loras(_fake_state(pipe), [("styleA", 0.8)], threading.Event()) + b._apply_loras(_fake_state(pipe), [], threading.Event()) + assert pipe.unloaded == 1 + assert pipe._unsloth_loras == () + + +def test_diffusers_apply_rejects_unsupported_quant(): + import threading + pipe = _FakePipe() + with pytest.raises(ValueError, match = "not supported"): + _backend()._apply_loras( + _fake_state(pipe, kind = "single_file", quant = "fp8"), + [("styleA", 1.0)], + threading.Event(), + ) + + +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. + import threading + + monkeypatch.setattr( + dl, + "resolve_specs", + lambda specs, **_: [ + dl.ResolvedLora(i, dl.sanitize_alias(i), f"/{i}.gguf", "gguf", w) for i, w in specs + ], + ) + pipe = _FakePipe() + with pytest.raises(ValueError, match = "GGUF LoRA"): + _backend()._apply_loras(_fake_state(pipe), [("styleA", 1.0)], threading.Event()) + assert pipe.loaded == [] # never touched the pipe diff --git a/studio/backend/tests/test_sd_cpp_backend.py b/studio/backend/tests/test_sd_cpp_backend.py index fb64f1324b..109724ca0b 100644 --- a/studio/backend/tests/test_sd_cpp_backend.py +++ b/studio/backend/tests/test_sd_cpp_backend.py @@ -92,6 +92,7 @@ class _FakeServer: self.payloads = [] self.timeouts = [] self.alive = True + self.lora_dir = None # set by a test to the server's --lora-model-dir scratch dir def is_alive(self): return self.alive and not self.stopped @@ -668,3 +669,86 @@ def test_run_load_redacts_paths_in_progress_error(monkeypatch): with npl._REDACTION_LOCK: if secret_root in npl._NATIVE_PATH_REDACTIONS: npl._NATIVE_PATH_REDACTIONS.remove(secret_root) + + +# ── LoRA (native engine) ──────────────────────────────────────────────────────── + + +def _fake_materialize(resolved, dest): + """Stand-in for diffusion_lora.materialize_native_dir: write a stub file per adapter + into ``dest`` and return the resolved list pointing at the written paths (mirroring the + real helper's contract without touching the Hub / real weights).""" + from pathlib import Path as _P + + from core.inference import diffusion_lora as dl + + dest.mkdir(parents = True, exist_ok = True) + out = [] + for r in resolved: + p = _P(dest) / f"{r.alias}.safetensors" + p.write_bytes(b"stub") + out.append(dl.ResolvedLora(r.id, r.alias, str(p), r.fmt, r.weight)) + return out + + +def _patch_lora(monkeypatch, resolved, supported = True): + from core.inference import diffusion_lora as dl + + monkeypatch.setattr(dl, "supports_lora", lambda **k: supported) + monkeypatch.setattr(dl, "resolve_specs", lambda specs, **k: list(resolved)) + monkeypatch.setattr(dl, "materialize_native_dir", _fake_materialize) + + +def test_generate_oneshot_applies_loras_via_prompt_tags(monkeypatch): + # One-shot sd-cli LoRA: adapters materialized into a --lora-model-dir and selected with + # tags injected into the prompt (the real inject_prompt_tags runs here). + from core.inference import diffusion_lora as dl + + eng = _FakeEngine() + b = _loaded_backend(engine = eng) # mode = "oneshot" + _patch_lora(monkeypatch, [dl.ResolvedLora("id1", "myalias", "/x/a.safetensors", "safetensors", 0.8)]) + b.generate(prompt = "a fox", steps = 4, seed = 1, loras = [("id1", 0.8)]) + _, params, _, _ = eng.calls[0] + assert params.lora_dir is not None and params.lora_apply_mode == "auto" + assert "" in params.prompt + + +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. + from pathlib import Path as _P + + from core.inference import diffusion_lora as dl + + b = SdCppDiffusionBackend() + servers: list = [] + _run_server_load(monkeypatch, b, servers) + servers[0].lora_dir = str(tmp_path) + _patch_lora(monkeypatch, [dl.ResolvedLora("id1", "myalias", "/x/a.safetensors", "safetensors", 0.7)]) + b.generate(prompt = "x", steps = 4, seed = 1, batch_size = 1, loras = [("id1", 0.7)]) + payload = servers[0].payloads[0] + assert "lora" in payload and len(payload["lora"]) == 1 + assert payload["lora"][0]["multiplier"] == 0.7 + assert payload["lora"][0]["path"].endswith("myalias.safetensors") + assert " { return parseJson(await authFetch("/api/inference/images/unload", { method: "POST" })); } +/** List diffusion LoRA adapters, optionally filtered to a model family. */ +export async function listDiffusionLoras(family?: string): Promise { + const qs = family ? `?family=${encodeURIComponent(family)}` : ""; + const data = await parseJson<{ loras: DiffusionLoraInfo[] }>( + await authFetch(`/api/models/diffusion-loras${qs}`), + ); + return data.loras ?? []; +} + export interface GalleryPage { images: GalleryImage[]; has_more: boolean; diff --git a/studio/frontend/src/features/images/images-page.tsx b/studio/frontend/src/features/images/images-page.tsx index a563bbd909..58340003c9 100644 --- a/studio/frontend/src/features/images/images-page.tsx +++ b/studio/frontend/src/features/images/images-page.tsx @@ -48,8 +48,10 @@ import { toast } from "@/lib/toast"; import { type DiffusionGenerateProgress, type DiffusionLoadProgress, + type DiffusionLoraInfo, type DiffusionStatus, type GalleryImage, + type LoraSpecInput, deleteGalleryImage, fetchGalleryObjectUrl, generateDiffusionImage, @@ -57,6 +59,7 @@ import { getDiffusionStatus, getGallery, getGenerateProgress, + listDiffusionLoras, loadDiffusionModel, unloadDiffusionModel, } from "./api"; @@ -910,6 +913,10 @@ export function ImagesPage({ active = true }: { active?: boolean }) { // Reference (FLUX.2): up to 3 ADDITIONAL reference images beyond the primary one, combined // by the model (subject + style, character + scene). const [referenceImages, setReferenceImages] = useState([]); + // LoRA adapters selected for the next generation (id + weight), plus the list the picker + // offers. Applied at generate time; available adapters are refreshed per loaded family. + const [loras, setLoras] = useState([]); + const [availableLoras, setAvailableLoras] = useState([]); // Advanced options live in a right-docked panel (like Chat's settings panel). Closed by // default; a single fixed toggle in the top bar opens/closes it (the icon never moves). const [advancedOpen, setAdvancedOpen] = useState(false); @@ -985,6 +992,36 @@ export function ImagesPage({ active = true }: { active?: boolean }) { galleryCache.quant = quant; }, [images, hasMore, selectedId, quant]); + // Refresh the LoRA picker's options when the loaded model (family) changes, and drop any + // selected adapters the new model can't use so a stale/incompatible LoRA is never sent. + const loraCapable = Boolean(status?.loaded && status?.supports_lora); + useEffect(() => { + if (!loraCapable) { + setAvailableLoras([]); + setLoras([]); + return; + } + let cancelled = false; + listDiffusionLoras(status?.family ?? undefined) + .then((list) => { + if (cancelled) return; + setAvailableLoras(list); + const ids = new Set(list.map((l) => l.id)); + setLoras((prev) => prev.filter((s) => ids.has(s.id))); + }) + .catch(() => { + if (cancelled) return; + // Clear the SELECTED adapters too, not just the options: leaving a stale `loras` + // selection in state (with the picker now hidden/empty) would still be posted by + // handleGenerate and could apply adapters from the previous model, or fail. + setAvailableLoras([]); + setLoras([]); + }); + return () => { + cancelled = true; + }; + }, [loraCapable, status?.family]); + const selected = useMemo( () => images.find((i) => i.id === selectedId) ?? images[0] ?? null, [images, selectedId], @@ -1081,6 +1118,19 @@ export function ImagesPage({ active = true }: { active?: boolean }) { const m = matchAspect(image.width, image.height); setAspect(m.key); setPortrait(m.portrait); + // Restore selected LoRA adapters from the recipe ("id:weight" strings); split on the + // LAST colon so an id that itself contains ':' is preserved. Unparseable entries are + // skipped, and a recipe with no LoRAs clears the current selection so the restore + // reproduces the image faithfully rather than leaking a stale form selection. + const restoredLoras: LoraSpecInput[] = []; + for (const entry of image.loras ?? []) { + const idx = entry.lastIndexOf(":"); + if (idx <= 0) continue; + const id = entry.slice(0, idx); + const weight = Number(entry.slice(idx + 1)); + if (id && Number.isFinite(weight)) restoredLoras.push({ id, weight }); + } + setLoras(restoredLoras); toast.success("Settings restored to inputs"); }, []); @@ -1570,6 +1620,8 @@ export function ImagesPage({ active = true }: { active?: boolean }) { strength: condStrength, upscale: condUpscale, reference_images: condRefImages, + // Drop zero-weight rows so the recipe records only adapters that actually applied. + loras: loras.length ? loras.filter((l) => l.weight > 0) : undefined, }); if (!isMounted.current) break; // Prepend this run's records (newest first) and load their blobs. @@ -1587,7 +1639,7 @@ export function ImagesPage({ active = true }: { active?: boolean }) { setGenDone(null); setGenStep(null); } - }, [prompt, negativePrompt, width, height, steps, guidance, seed, batchSize, count, workflow, initImage, maskImage, strength, extendPct, extendSides, upscaleFactor, upscaleStrength, referenceImages, ensureSrc]); + }, [prompt, negativePrompt, width, height, steps, guidance, seed, batchSize, count, workflow, initImage, maskImage, strength, extendPct, extendSides, upscaleFactor, upscaleStrength, referenceImages, loras, ensureSrc]); // Keep the active workflow valid for the loaded model: an edit-only model (Qwen-Image- // Edit) has no Create/Transform tabs, a base model has no Edit tab. Snap to the first @@ -2025,6 +2077,89 @@ export function ImagesPage({ active = true }: { active?: boolean }) { onChange={(e) => setPrompt(e.target.value)} /> + {/* LoRA adapters: shown only when the loaded model + quant can apply them and at + least one adapter is discoverable. Stack multiple, each with a 0-2 weight. The + backend owns how they apply (native prompt tags / diffusers set_adapters); the + UI only sends {id, weight}. */} + {loraCapable && availableLoras.length > 0 && ( + +
+ {loras.map((sel, i) => ( +
+
+ + +
+ + setLoras((prev) => prev.map((p, j) => (j === i ? { ...p, weight: v } : p))) + } + /> +
+ ))} + {loras.length < Math.min(availableLoras.length, 8) && ( + + )} +
+ + )} {/* A negative prompt only does anything with guidance on, so hide it at guidance 0 (Z-Image-Turbo's default) instead of showing a dead field. */} {guidance > 0 && (