From 7f0ccdbf0120e3bc2b0dc5bcc19328d1a7e67344 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 22 Jul 2026 07:03:38 +0000 Subject: [PATCH] Batch diffusion inference with per-image seeds, an inference conditioning cache, and GGUF loader fixes Batched generation: /images/generate takes a prompts list (one image per prompt, txt2img only) or a seeds list (one prompt, one image per seed); the legacy batch_size path derives per-image seeds base..base+n-1 like the native engine. Every image gets its own torch.Generator so any batch member replays alone from its gallery recipe; the whole list runs as one forward by default with OOM backoff that halves a failed chunk, and an explicit batch_size caps images per forward. Validated 10-22x over serial engines on 32-image suites with LPIPS deltas within 0.002. Conditioning cache on the inference path: UNSLOTH_DIFFUSION_COND_CACHE_DIR (the inference sibling of the trainers' cond_cache_dir, same persistent store) wraps encode_prompt so repeated prompts skip the text-encoder forward entirely; verified bit-identical outputs. Bypassed while LoRA adapters are attached; tensor-argument calls pass through uncached. Compile cache: GGUF loads fingerprint their own bundles (quant=gguf, a different compiled graph than the dense family) and batched calls register every distinct (w, h, batch) chunk shape they ran, so the heavy GGUF batched warmups (~159 s at batch 32 on 12B-class, ~655 s on 20B CFG-batched) are paid once ever. GGUF loader: strip the sd.cpp model.diffusion_model. container prefix in the single-file converter; diffusers' FLUX.2 converter KeyErrors on it and the Qwen-Image identity mapping strands the model on meta. --- studio/backend/core/inference/diffusion.py | 227 +++++++++++++++--- .../core/inference/diffusion_batched.py | 138 +++++++++++ .../core/inference/diffusion_compile_cache.py | 9 + .../core/inference/diffusion_cond_cache.py | 223 +++++++++++++++++ .../backend/core/inference/sd_cpp_backend.py | 11 + studio/backend/models/inference.py | 46 ++++ studio/backend/routes/inference.py | 16 +- .../backend/tests/test_diffusion_backend.py | 146 ++++++++++- .../backend/tests/test_diffusion_batched.py | 153 ++++++++++++ .../tests/test_diffusion_compile_cache.py | 39 +++ .../tests/test_diffusion_cond_cache.py | 165 +++++++++++++ 11 files changed, 1138 insertions(+), 35 deletions(-) create mode 100644 studio/backend/core/inference/diffusion_batched.py create mode 100644 studio/backend/core/inference/diffusion_cond_cache.py create mode 100644 studio/backend/tests/test_diffusion_batched.py create mode 100644 studio/backend/tests/test_diffusion_cond_cache.py diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index f9bc5ea135..e4aab31754 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -83,7 +83,15 @@ from .diffusion_attention import ( _ensure_attention_backend_installed, ) from . import diffusion_compile_cache as compile_cache +from . import diffusion_cond_cache as cond_cache from . import diffusion_gguf_compile as gguf_compile +from .diffusion_batched import ( + chunk_jobs, + is_oom_error, + resolve_batch_jobs, + split_chunk, + uniform_prompt, +) from .diffusion_cache import ( FBCACHE_MIN_STEPS, TC_AUTO, @@ -469,6 +477,46 @@ def _resolve_diffusion_compute_dtype(fam: Optional[DiffusionFamily], dtype: Any) return torch.float32 if dtype == torch.float16 else dtype +def _install_gguf_prefix_strip(transformer_cls: Any, logger: Any) -> None: + """Wrap the class's diffusers single-file converter to strip the + ``model.diffusion_model.`` container prefix that sd.cpp-converted GGUFs + carry on every tensor. + + diffusers (<= 0.39) handles the prefix inconsistently: the FLUX.1 converter + strips it natively, but the FLUX.2 converter never does and KeyErrors on the + prefixed keys (``'double_blocks.0.img_attn.norm.key_norm'`` misparse in + ``convert_flux2_transformer_checkpoint_to_diffusers``), and the Qwen-Image + mapping fn is an identity, so every prefixed tensor is reported "not used", + the model stays on meta, and ``.to(cuda)`` dies with "Cannot copy out of + meta tensor". Stripping the prefix when present is a no-op for already-clean + checkpoints, so the shim applies to every GGUF transformer class uniformly. + Idempotent (the wrapper is marked) and best-effort.""" + try: + from diffusers.loaders import single_file_model as sfm + + entry = sfm.SINGLE_FILE_LOADABLE_CLASSES.get( + getattr(transformer_cls, "__name__", str(transformer_cls)) + ) + if not isinstance(entry, dict): + return + original = entry.get("checkpoint_mapping_fn") + if not callable(original) or getattr(original, "_unsloth_prefix_strip", False): + return + prefix = "model.diffusion_model." + + def _stripped_mapping_fn(checkpoint = None, **kwargs): + checkpoint = { + (key[len(prefix):] if key.startswith(prefix) else key): value + for key, value in (checkpoint or {}).items() + } + return original(checkpoint = checkpoint, **kwargs) + + _stripped_mapping_fn._unsloth_prefix_strip = True + entry["checkpoint_mapping_fn"] = _stripped_mapping_fn + except Exception as exc: # noqa: BLE001 — loader-compat shim only, never fail the load + logger.warning("diffusion.gguf: prefix-strip shim not installed: %s", exc) + + class DiffusionBackend: """Holds at most one loaded diffusers pipeline. All mutations are serialised.""" @@ -1502,6 +1550,9 @@ class DiffusionBackend: sf_kwargs["quantization_config"] = diffusers.GGUFQuantizationConfig( compute_dtype = dtype ) + # sd.cpp-converted GGUFs prefix tensors with model.diffusion_model.; + # the FLUX.2 / Qwen-Image converters crash or meta-strand on it. + _install_gguf_prefix_strip(transformer_cls, logger) # A safetensors single-file (fp8) carries its own dtype: no GGUF dequant config. transformer = transformer_cls.from_single_file( single_file_path, **sf_kwargs @@ -1673,7 +1724,11 @@ class DiffusionBackend: transformer = getattr(pipe, "transformer", None) or getattr(pipe, "unet", None), dtype = getattr(target, "dtype", None), - quant = transformer_quant_engaged, + # A GGUF transformer compiles a DIFFERENT graph (the dequant op + # chain under `default`, dequant+block fusion under `max`) than a + # dense load of the same family, so it must key its own bundles + # instead of sharing (and cross-hitting) the dense fingerprint. + quant = "gguf" if gguf_transformer else transformer_quant_engaged, attention_backend = attention_engaged, compile_kwargs = { # Mirrors apply_speed_optims' fullgraph decision: an active or @@ -1720,6 +1775,22 @@ class DiffusionBackend: logger = logger, ) + # Inference-side persistent conditioning cache (opt-in via + # UNSLOTH_DIFFUSION_COND_CACHE_DIR, the inference sibling of the + # trainers' cond_cache_dir): repeated prompts skip the text-encoder + # forward entirely, so warm prompts never onload the multi-GB encoders + # under an offload policy. Installed AFTER the TE quant above so the + # cache key reflects the encoders that actually run. Instance-level; + # dies with the pipe on unload. + cond_cache.install( + pipe, + family = fam.name, + repo_id = repo_id, + dtype = dtype, + te_quant = te_quant, + logger = logger, + ) + # Apply the planned placement; apply_memory_plan returns the (policy, tiling) # ACTUALLY engaged so status stays honest. Idempotent for the `none` policy. effective_policy, effective_tiling = apply_memory_plan( @@ -2528,7 +2599,8 @@ class DiffusionBackend: transformer = getattr(state.pipe, "transformer", None) or getattr(state.pipe, "unet", None), dtype = getattr(target, "dtype", None), - quant = state.transformer_quant, + # Same GGUF-vs-dense graph distinction as the load-time begin(). + quant = "gguf" if gguf_transformer else state.transformer_quant, attention_backend = attention_engaged, compile_kwargs = { # Mirrors the load-time fullgraph decision: an engaged or @@ -2586,6 +2658,13 @@ class DiffusionBackend: guidance: float = 0.0, seed: Optional[int] = None, batch_size: int = 1, + # Batched multi-image generation (see diffusion_batched.py): ``prompts`` renders one + # image per prompt (txt2img only); ``seeds`` renders one image per seed. Each image + # gets its OWN torch.Generator, so any image regenerated alone with its recorded seed + # is bit-identical to its batched rendition. With neither, ``batch_size`` derives + # per-image seeds as base..base+batch_size-1 (matching the native engine). + prompts: Optional[list[str]] = None, + seeds: Optional[list[int]] = None, # Image-conditioned (base64/data-URL): init alone = img2img; init + mask = inpaint. # ``strength`` is the denoise strength (0 = keep source, 1 = full redraw). None = txt2img. init_image: Optional[str] = None, @@ -2620,14 +2699,19 @@ class DiffusionBackend: self._gen = _GenState(total_steps = steps) try: # The local `state` ref keeps the pipe alive even if unload() nulls _state. - generator = torch.Generator(device = state.device) - if seed is None: - # Keep the seed in JS's safe-integer range (< 2**53) so it round-trips - # through JSON and reproduces the image (a raw 64-bit seed loses precision). - seed = generator.seed() & ((1 << 53) - 1) - else: - seed = int(seed) - generator.manual_seed(seed) + # Resolve the per-image (prompt, seed) jobs up front: N prompts, one prompt x + # N seeds, or the legacy single prompt whose batch derives per-image seeds as + # base..base+batch_size-1 (matching the native sd.cpp engine, so a gallery + # recipe replays any batch member alone). A fresh base seed is drawn in JS's + # safe-integer range (< 2**53) so it round-trips through JSON. + jobs, seed = resolve_batch_jobs( + prompt = prompt, + prompts = prompts, + seed = seed, + seeds = seeds, + batch_size = batch_size, + draw_seed = torch.Generator(device = state.device).seed, + ) # Deferred speed auto: engage the compile profile on the 3rd image, before the LoRA/ # workflow wiring (load-time ordering). Best-effort; a failure stays eager and never retries. @@ -2784,6 +2868,15 @@ class DiffusionBackend: cn_scale, cn_gstart, cn_gend = cn_strength, cn_gs, cn_ge # Flux Union CN selects its head by an integer control_mode; map the type. cn_mode = diffusion_controlnet.union_control_mode(cn_id, cn_type) + # A prompt LIST batches plain text-to-image only: the image-conditioned and + # ControlNet workflows are tuned around one conditioning image per call, and a + # silent broadcast would pair every prompt with the same image. Multiple SEEDS + # of one prompt remain valid everywhere (per-image generators only). + if uniform_prompt(jobs) is None and workflow != "txt2img": + raise ValueError( + "A prompts list is supported for plain text-to-image only; the " + f"{workflow} workflow takes one prompt per call (seed lists still work)." + ) # Snap odd-sized inputs to a multiple of 16 for workflows whose OUTPUT size comes # from the input image (img2img/inpaint/edit); txt2img/reference use the slider, # upscale already produced a /16 target. Mask is matched to the snapped image. @@ -2804,13 +2897,12 @@ class DiffusionBackend: call_params = inspect.signature(pipe.__call__).parameters kwargs: dict[str, Any] = { - "prompt": prompt, + # prompt / generator / num_images_per_prompt are set per chunk below: + # each per-forward chunk carries its own prompt slice and one + # torch.Generator PER IMAGE (individual seed reproducibility). "num_inference_steps": steps, # Most pipelines use "guidance_scale"; Qwen-Image uses "true_cfg_scale". state.family.cfg_kwarg: guidance, - "generator": generator, - # Whole batch in one forward pass; all share this call's seed. - "num_images_per_prompt": batch_size, } if state.family.name == IDEOGRAM4_FAMILY_NAME: # Ideogram 4 drives CFG via EITHER a constant guidance_scale OR a per-step @@ -2864,12 +2956,20 @@ class DiffusionBackend: if "control_mode" in call_params and cn_mode is not None: kwargs["control_mode"] = cn_mode - gen = _GenState(total_steps = steps) + # Per-forward chunks: the whole job list in ONE forward by default (the + # measured sweet spot: batch 32 on 4-step models, ~10-22x over serial + # engines), bounded by an explicit batch_size cap; OOM backoff below + # halves a failed chunk instead of failing the call. + chunks = chunk_jobs(jobs, batch_size) + gen = _GenState(total_steps = steps * len(chunks)) + # Steps completed by FINISHED chunks, so the progress bar spans the whole + # multi-chunk call (mutable cell: _on_step closes over it). + steps_done = [0] def _on_step(pipe, step_index, timestep, callback_kwargs): # Monotonic: a wall-clock adjustment (NTP) mid-denoise would skew the ETA. now = time.monotonic() - gen.step = step_index + 1 + gen.step = steps_done[0] + step_index + 1 if gen.first_step_at == 0.0: gen.first_step_at = now gen.eta_seconds = _estimate_eta( @@ -2922,37 +3022,104 @@ class DiffusionBackend: self._reset_step_cache(state.pipe) self._gen = gen - # inference_mode is faster than no_grad and numerically identical here. - with torch.inference_mode(): - images = pipe(**kwargs).images + images: list[Any] = [] + per_image_seeds: list[int] = [] + chunk_shapes: list[int] = [] + pending = list(chunks) + while pending: + chunk = pending.pop(0) + chunk_kwargs = dict(kwargs) + shared = uniform_prompt(chunk) + generators = [ + torch.Generator(device = state.device).manual_seed(s) + for _, s in chunk + ] + if len(jobs) == 1: + # Single image: scalar prompt + generator, exactly the pre-batching + # call shape (the regression harness checks this path bit-identical). + chunk_kwargs["prompt"] = shared + chunk_kwargs["generator"] = generators[0] + chunk_kwargs["num_images_per_prompt"] = 1 + elif shared is not None: + # Uniform prompt: encode it ONCE and fan out per-image generators. + chunk_kwargs["prompt"] = shared + chunk_kwargs["generator"] = generators + chunk_kwargs["num_images_per_prompt"] = len(chunk) + else: + # Distinct prompts: one image per prompt in a single forward. + chunk_kwargs["prompt"] = [p for p, _ in chunk] + chunk_kwargs["generator"] = generators + chunk_kwargs["num_images_per_prompt"] = 1 + try: + # inference_mode is faster than no_grad and numerically identical here. + with torch.inference_mode(): + out = pipe(**chunk_kwargs).images + except Exception as exc: # noqa: BLE001 — reraised unless a splittable OOM + if len(chunk) < 2 or not is_oom_error(exc): + raise + # OOM backoff: halve the failed chunk and retry; finished chunks keep + # their images and per-image seeds keep every retry reproducible. + empty_cache = getattr( + getattr(torch, "cuda", None), "empty_cache", None + ) + if callable(empty_cache): + empty_cache() + first_half, second_half = split_chunk(chunk) + pending[:0] = [first_half, second_half] + gen.total_steps += steps # one extra chunk to run + logger.warning( + "diffusion.generate: batch of %d hit OOM; retrying as %d + %d", + len(chunk), + len(first_half), + len(second_half), + ) + continue + # A cancelled denoise returns a partial/garbage image; don't persist it + # (nor run any remaining chunks). + if cancel.is_set(): + raise RuntimeError(DIFFUSION_CANCELLED_MSG) + images.extend(out) + per_image_seeds.extend(s for _, s in chunk) + chunk_shapes.append(len(chunk)) + steps_done[0] += steps # Keep progress ACTIVE through the post-denoise work below (the compile-cache save can # take a moment) -- don't null _gen here. The route persists the image AFTER this # returns, so a reload's mount probe that read idle now would refresh the gallery # before the result exists. The outer finally clears _gen on every exit (return/raise). - # A cancelled denoise returns a partial/garbage image; don't persist it. - if cancel.is_set(): - raise RuntimeError(DIFFUSION_CANCELLED_MSG) # Persist the warm torch.compile bundle after the first compiled generation. A # STATIC compile makes new artifacts per (w,h,batch), so register this shape first # (an uncovered shape re-dirties the context and the save rewrites the bundle). # Idempotent + best-effort. try: # Register the dims the forward ACTUALLY compiled with (image-conditioned - # workflows run at the input image's size, not the slider; see _compile_shape_dims). + # workflows run at the input image's size, not the slider; see + # _compile_shape_dims). A static compile makes one artifact per BATCH size + # too, so every distinct chunk size this call ran (OOM backoff can produce + # several) registers its own (w, h, batch) shape. reg_width, reg_height = _compile_shape_dims(workflow, init_pil, width, height) - compile_cache.register_shape( - state.compile_cache_ctx, - (reg_width, reg_height, int(batch_size)), - static = "compiled" in (state.speed_optims or ()) - and compiled_shapes_are_static(state.pipe, state.speed_mode), - ) + static_shapes = "compiled" in ( + state.speed_optims or () + ) and compiled_shapes_are_static(state.pipe, state.speed_mode) + for chunk_batch in sorted(set(chunk_shapes)): + compile_cache.register_shape( + state.compile_cache_ctx, + (reg_width, reg_height, int(chunk_batch)), + static = static_shapes, + ) compile_cache.save(state.compile_cache_ctx, logger = logger) except Exception: # noqa: BLE001 — cache persistence is best-effort pass # Count the finished generation (drives deferred speed); a batch is one generation. object.__setattr__(state, "generation_count", state.generation_count + 1) # Return the PIL images (unencoded); the route embeds recipes and persists them. - return {"images": list(images), "seed": int(seed), "repo_id": state.repo_id} + # ``seeds`` records each image's OWN seed (the native engine already reports + # these), so every batch member replays individually from its recipe. + return { + "images": list(images), + "seed": int(seed), + "seeds": [int(s) for s in per_image_seeds], + "repo_id": state.repo_id, + } finally: # Deregister so a later unload/load can't poke a finished generation (if still ours). with self._lock: diff --git a/studio/backend/core/inference/diffusion_batched.py b/studio/backend/core/inference/diffusion_batched.py new file mode 100644 index 0000000000..f4651d2322 --- /dev/null +++ b/studio/backend/core/inference/diffusion_batched.py @@ -0,0 +1,138 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Batched multi-image planning for the local diffusion backend. + +One generation call can produce N images three ways: a prompt LIST (one image per +prompt), one prompt x a seed LIST, or the legacy single prompt + ``batch_size`` +(whose per-image seeds derive as base..base+batch_size-1, matching the native +sd.cpp engine and the gallery recipe replay). These helpers turn the request into +an explicit per-image ``(prompt, seed)`` job list, chunk it into per-forward +batches, and support OOM backoff by splitting a failed chunk in half. + +Measured on the 32-image eval suites (diffusers 0.39 / torch 2.10): one batched +forward with per-image ``torch.Generator``s is numerics-safe (LPIPS deltas within +0.002 of serial) and 10-22x faster end-to-end than serial per-image engines -- +batch 32 fits 4-step 12B-class models on one GPU, batch 8 fits a 20B model at +1024px with CFG batching. Per-image generators keep every image individually +reproducible: image ``i`` regenerated alone with its own seed is bit-identical +to its batched rendition. + +Pure and torch-free so the CPU unit tests stay light; the engine owns the +torch.Generator construction and the actual pipeline calls. +""" + +from __future__ import annotations + +from typing import Any, Callable, Optional + +# Upper bound on images per generation call (mirrors the route's batch_size cap): +# a prompt/seed list beyond this is a client error, not an OOM to back off from. +MAX_BATCH_IMAGES = 32 + +# Seeds stay in JS's safe-integer range so they round-trip through the JSON +# gallery recipes and reproduce the image (a raw 64-bit seed loses precision). +SEED_MASK = (1 << 53) - 1 + + +def resolve_batch_jobs( + *, + prompt: str, + prompts: Optional[list[str]], + seed: Optional[int], + seeds: Optional[list[int]], + batch_size: int, + draw_seed: Callable[[], int], +) -> tuple[list[tuple[str, int]], int]: + """The per-image ``(prompt, seed)`` jobs plus the base seed for this call. + + - ``prompts`` (list): one image per prompt. With ``seeds`` too, lengths must + match (seed i drives prompt i); without, seeds derive from the base. + - ``seeds`` (list) alone: one image per seed, all with ``prompt``. + - neither: ``batch_size`` images of ``prompt`` with derived seeds + base..base+batch_size-1 (each masked JSON-safe). + + ``draw_seed`` supplies a fresh random base when the caller sent none (the + engine passes a ``torch.Generator`` draw). Raises ``ValueError`` on empty / + oversized lists, a length mismatch, or an out-of-range seed.""" + if prompts is not None: + if not prompts or not all(isinstance(p, str) and p.strip() for p in prompts): + raise ValueError("prompts must be a non-empty list of non-empty strings") + if len(prompts) > MAX_BATCH_IMAGES: + raise ValueError(f"prompts supports at most {MAX_BATCH_IMAGES} entries per call") + if seeds is not None: + if not seeds: + raise ValueError("seeds must be a non-empty list of integers") + if len(seeds) > MAX_BATCH_IMAGES: + raise ValueError(f"seeds supports at most {MAX_BATCH_IMAGES} entries per call") + seeds = [int(s) for s in seeds] + if any(s < 0 or s > SEED_MASK for s in seeds): + raise ValueError("every seed must be between 0 and 2**53 - 1 (JSON-safe)") + if prompts is not None and len(seeds) != len(prompts): + raise ValueError( + f"prompts and seeds must have the same length " + f"(got {len(prompts)} prompts, {len(seeds)} seeds)" + ) + + if prompts is not None: + count = len(prompts) + elif seeds is not None: + count = len(seeds) + else: + count = max(1, int(batch_size)) + + if seeds is not None: + job_seeds = seeds + base_seed = seeds[0] + else: + base_seed = int(seed) if seed is not None else int(draw_seed()) & SEED_MASK + job_seeds = [(base_seed + i) & SEED_MASK for i in range(count)] + + job_prompts = prompts if prompts is not None else [prompt] * count + return list(zip(job_prompts, job_seeds)), base_seed + + +def chunk_jobs( + jobs: list[tuple[str, int]], batch_size: int +) -> list[list[tuple[str, int]]]: + """Split the jobs into per-forward chunks. + + ``batch_size`` doubles as the per-forward cap when a prompt/seed list drives + the image count: an explicit ``batch_size > 1`` bounds each forward, while + the untouched default (1) lets the whole list run as ONE forward -- the + measured sweet spot (batch 32 on 4-step models) -- with OOM backoff as the + safety net rather than a serial default.""" + if not jobs: + return [] + per_forward = len(jobs) if batch_size <= 1 else min(int(batch_size), len(jobs)) + return [jobs[i : i + per_forward] for i in range(0, len(jobs), per_forward)] + + +def split_chunk( + chunk: list[tuple[str, int]], +) -> tuple[list[tuple[str, int]], list[tuple[str, int]]]: + """Halve a chunk for OOM backoff (first half never smaller than the second, + so repeated splits terminate at singletons). Raises on an unsplittable chunk.""" + if len(chunk) < 2: + raise ValueError("cannot split a chunk of fewer than 2 jobs") + mid = (len(chunk) + 1) // 2 + return chunk[:mid], chunk[mid:] + + +def uniform_prompt(chunk: list[tuple[str, int]]) -> Optional[str]: + """The chunk's single shared prompt, or None when prompts differ. + + A uniform chunk encodes its prompt ONCE (``num_images_per_prompt`` fans it + out); a mixed chunk passes the prompt list with one image per prompt.""" + first = chunk[0][0] + return first if all(p == first for p, _ in chunk) else None + + +def is_oom_error(exc: BaseException) -> bool: + """Whether an exception is a CUDA/accelerator out-of-memory, worth a smaller + retry. Matched structurally (class name across torch versions / devices) and + by message, so the caller needn't import torch to classify.""" + for klass in type(exc).__mro__: + if klass.__name__ == "OutOfMemoryError": + return True + return "out of memory" in str(exc).lower() diff --git a/studio/backend/core/inference/diffusion_compile_cache.py b/studio/backend/core/inference/diffusion_compile_cache.py index f1e540e9a5..ab05a3d609 100644 --- a/studio/backend/core/inference/diffusion_compile_cache.py +++ b/studio/backend/core/inference/diffusion_compile_cache.py @@ -15,6 +15,15 @@ error. Hence an EXACT-MATCH fingerprint with SILENT FALLBACK to local compile (a normal), and per-arch bundles keyed by the full fingerprint. See ``outputs/compile_cache/DISTRIBUTION.md``. +GGUF loads participate too (fingerprinted ``quant="gguf"``, a different compiled graph +than the dense family): ``torch.compile`` mode="default" runs clean over diffusers' +GGUF dequant path (measured on torch 2.10 / diffusers 0.39), but the cold warmup is +HEAVY at batched shapes -- ~159 s first-pass wall on a 12B-class 4-step model at +batch 32, up to ~655 s on a 20B CFG-batched model at 1024px -- which is exactly the +cost this bundle amortises to once-ever. Batched generation registers every distinct +(width, height, batch) chunk shape it runs via ``register_shape`` so a static-compile +bundle grows to cover the batch sizes actually used, OOM-backoff halves included. + Lifecycle (around ``_compile_repeated_blocks``): ``begin`` builds the fingerprint, points ``TORCHINDUCTOR_CACHE_DIR`` at a per-key dir, and loads a matching bundle (before the first compiled forward); ``save`` writes the bundle + manifest after the warmup forward diff --git a/studio/backend/core/inference/diffusion_cond_cache.py b/studio/backend/core/inference/diffusion_cond_cache.py new file mode 100644 index 0000000000..311dda4037 --- /dev/null +++ b/studio/backend/core/inference/diffusion_cond_cache.py @@ -0,0 +1,223 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Persistent prompt-conditioning cache for the diffusion INFERENCE path. + +The inference sibling of the trainers' ``cond_cache_dir`` (see +``diffusion_train_common.DiffusionLoraConfig``), reusing the SAME on-disk store +(``diffusion_train_extras.PersistentConditioningCache``): one safetensors file +per encoded prompt, keyed by content hash + family. Enabled by pointing +``UNSLOTH_DIFFUSION_COND_CACHE_DIR`` at a directory -- unset/blank means off, +exactly like the trainer knob, so the default render path stays untouched. + +When enabled, the loaded pipeline's ``encode_prompt`` is wrapped per-instance: +a repeated prompt returns the stored embeddings and never runs the text-encoder +forward, so under an offload policy the multi-GB encoders stay OFF the GPU for +warm prompts entirely. Verified bit-identical on the 32-image eval suites +(FLUX.1-schnell / FLUX.2-klein / Qwen-Image-2512): renders from cached +embeddings match inline-encoded renders exactly, and warm runs load the whole +pipeline with no text encoder resident (8-18s load, ~7B-VL VRAM saved on +Qwen-Image). + +Safety gates: the wrapper only caches calls whose arguments are all plain +JSON-safe values (a tensor argument such as pre-supplied ``prompt_embeds`` +passes straight through), keys on everything that changes the embedding +numerics (family, repo, dtype, text-encoder quant, diffusers version, and the +full argument set minus device/generator), and bypasses entirely while LoRA +adapters are attached (an adapter may target the text encoders). Best-effort +throughout: any cache failure falls back to the real encode. torch is imported +lazily so this stays importable in a no-torch runtime. +""" + +from __future__ import annotations + +import hashlib +import inspect +import json +import os +from typing import Any, Optional + +_ENV_DIR = "UNSLOTH_DIFFUSION_COND_CACHE_DIR" + +# Bound arguments that never change the returned embeddings: the target device +# is a placement detail (the hit is moved there) and no encode path draws RNG. +_KEY_EXCLUDED_ARGS = frozenset({"device", "generator"}) + + +def cache_dir() -> Optional[str]: + """The configured cache directory, or None when the cache is off. A blank + value means off (matching the trainers' ``cond_cache_dir`` semantics).""" + value = (os.environ.get(_ENV_DIR) or "").strip() + return value or None + + +def _json_safe(value: Any) -> bool: + if value is None or isinstance(value, (bool, int, float, str)): + return True + if isinstance(value, (list, tuple)): + return all(_json_safe(v) for v in value) + return False + + +# Per-slot layout codes for _flatten/_unflatten (stored as the leading int64 tensor): +# -1 = None slot, -2 = a bare tensor, n >= 0 = a LIST of n tensors (Z-Image returns +# its per-prompt embeddings as a list, so one nesting level must round-trip). +_SLOT_NONE = -1 +_SLOT_TENSOR = -2 + + +def _flatten(result: Any) -> Optional[list]: + """Flatten an ``encode_prompt`` result tuple of tensors / None / tensor-lists + into ``[layout_tensor, *tensors]`` for the conditioning cache, or None when the + result holds anything else (those calls simply aren't cached).""" + if not isinstance(result, tuple) or not result: + return None + layout: list[int] = [] + flat: list[Any] = [] + for item in result: + if item is None: + layout.append(_SLOT_NONE) + elif isinstance(item, (list, tuple)): + if not all(hasattr(t, "detach") for t in item): + return None + layout.append(len(item)) + flat.extend(item) + elif hasattr(item, "detach"): + layout.append(_SLOT_TENSOR) + flat.append(item) + else: + return None + import torch # noqa: PLC0415 + + return [torch.tensor(layout, dtype = torch.int64), *flat] + + +def _unflatten(stored: tuple, device: Any) -> tuple: + """Rebuild the ``encode_prompt`` result from a cached ``_flatten`` record, + moving every tensor to ``device`` (dtype is preserved as stored).""" + layout = stored[0].tolist() + tensors = list(stored[1:]) + out: list[Any] = [] + index = 0 + for code in layout: + if code == _SLOT_NONE: + out.append(None) + elif code == _SLOT_TENSOR: + out.append(tensors[index].to(device)) + index += 1 + else: + out.append([t.to(device) for t in tensors[index : index + code]]) + index += code + return tuple(out) + + +def install( + pipe: Any, + *, + family: str, + repo_id: str, + dtype: Any, + te_quant: Any = None, + logger: Any = None, +) -> bool: + """Wrap ``pipe.encode_prompt`` with the persistent cache. No-op (False) when + the env knob is unset, the pipe has no ``encode_prompt``, or setup fails. + + Instance-level assignment only: the wrapper dies with the pipe on unload and + never mutates the pipeline class.""" + root = cache_dir() + if root is None: + return False + encode = getattr(pipe, "encode_prompt", None) + if not callable(encode): + return False + try: + # Lazy: core.training imports parts of core.inference, so the module-level + # import would be circular; the extras module itself is stdlib-only. + from core.training.diffusion_train_extras import PersistentConditioningCache + + signature = inspect.signature(encode) + cache = PersistentConditioningCache(root, family, 0) + except Exception as exc: # noqa: BLE001 — cache is best-effort + if logger is not None: + logger.warning("diffusion.cond_cache: install failed: %s", exc) + return False + + # Everything beyond the call arguments that changes the embedding numerics. + load_fp = { + "repo": str(repo_id), + "dtype": str(dtype), + "te_quant": str(te_quant) if te_quant is not None else "none", + "diffusers": _diffusers_version(), + } + stats = {"hits": 0, "misses": 0} + + def cached_encode_prompt(*args: Any, **kwargs: Any) -> Any: + try: + # LoRA may target the text encoders; adapters attached -> always encode. + if getattr(pipe, "_unsloth_loras", ()): + return encode(*args, **kwargs) + bound = signature.bind(*args, **kwargs) + bound.apply_defaults() + keyed = { + name: value + for name, value in bound.arguments.items() + if name not in _KEY_EXCLUDED_ARGS + } + if not all(_json_safe(v) for v in keyed.values()): + # Tensor/object arguments (pre-supplied embeds, images) are not keyable. + return encode(*args, **kwargs) + payload = json.dumps( + {"load": load_fp, "args": keyed}, sort_keys = True, default = str + ) + key = cache.text_key( + f"inference::{hashlib.sha256(payload.encode('utf-8')).hexdigest()}" + ) + hit = cache.get(key) + if hit is not None: + stats["hits"] += 1 + return _unflatten(hit, _target_device(pipe, bound)) + except Exception as exc: # noqa: BLE001 — never fail a generation over the cache + if logger is not None: + logger.warning("diffusion.cond_cache: lookup failed: %s", exc) + return encode(*args, **kwargs) + result = encode(*args, **kwargs) + try: + flat = _flatten(result) + if flat is not None: + cache.put(key, flat) + stats["misses"] += 1 + except Exception as exc: # noqa: BLE001 — a failed write only skips reuse + if logger is not None: + logger.warning("diffusion.cond_cache: store failed: %s", exc) + return result + + pipe.encode_prompt = cached_encode_prompt + pipe._unsloth_cond_cache_stats = stats + if logger is not None: + logger.info( + "diffusion.cond_cache: enabled at %s (family=%s); repeated prompts skip " + "the text-encoder forward", + root, + family, + ) + return True + + +def _target_device(pipe: Any, bound: inspect.BoundArguments) -> Any: + """The device a hit's tensors must land on: the caller's explicit ``device`` + argument when given, else the pipeline's execution device, else CPU.""" + device = bound.arguments.get("device") + if device is not None: + return device + device = getattr(pipe, "_execution_device", None) + return device if device is not None else "cpu" + + +def _diffusers_version() -> Optional[str]: + try: + import diffusers # noqa: PLC0415 + + return str(getattr(diffusers, "__version__", None)) + except Exception: # noqa: BLE001 + return None diff --git a/studio/backend/core/inference/sd_cpp_backend.py b/studio/backend/core/inference/sd_cpp_backend.py index 020e63cb5a..991774b940 100644 --- a/studio/backend/core/inference/sd_cpp_backend.py +++ b/studio/backend/core/inference/sd_cpp_backend.py @@ -728,6 +728,11 @@ class SdCppDiffusionBackend: guidance: float = 0.0, seed: Optional[int] = None, batch_size: int = 1, + # Batched prompt/seed lists are diffusers-engine features (one batched DiT forward); + # accepted for interface parity and rejected clearly below (sd-cli renders serially, + # so a "batched" list here would silently be a slow loop). + prompts: Optional[list[str]] = None, + seeds: Optional[list[int]] = None, # Accepted for interface parity; native is text-to-image only, so image-conditioned # requests are rejected clearly below rather than silently dropped. init_image: Optional[str] = None, @@ -757,6 +762,12 @@ class SdCppDiffusionBackend: "img2img / inpaint / reference / upscale are not yet supported on the native " "sd.cpp engine; run on a GPU (diffusers) for image-conditioned workflows." ) + if prompts is not None or seeds is not None: + raise ValueError( + "Batched prompt/seed lists are not supported on the native sd.cpp engine " + "(it renders serially); run on a GPU (diffusers) for batched generation, " + "or use batch_size for a serial native batch." + ) # strength 0/None disables ControlNet (matches diffusers), so no-op it rather than 400. if controlnet is not None and controlnet[3] in (None, 0, 0.0): controlnet = None diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 9a9fe04789..ea8a21ea6a 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -1994,6 +1994,52 @@ class DiffusionGenerateRequest(BaseModel): batch_size: int = Field( 1, ge = 1, le = 32, description = "Images generated in one forward pass (VRAM-heavy)" ) + # Batched multi-image generation (diffusers engine): a prompt list renders one image per + # prompt in a single batched forward (txt2img only); a seed list renders one image per + # seed. Each image carries its OWN generator seed, so any batch member replays alone. + prompts: Optional[list[str]] = Field( + None, + min_length = 1, + max_length = 32, + description = "Prompt list for batched generation: one image per prompt in a single " + "forward pass (plain text-to-image only). Overrides `prompt` for the images; " + "`prompt` is still required as the fallback/display value.", + ) + seeds: Optional[list[int]] = Field( + None, + min_length = 1, + max_length = 32, + description = "Per-image seeds for batched generation: one image per seed (with " + "`prompts`, lengths must match; alone, every image uses `prompt`). Each image is " + "individually reproducible from its own seed.", + ) + + @field_validator("prompts") + @classmethod + def _non_empty_prompts(cls, value: Optional[list[str]]) -> Optional[list[str]]: + if value is not None and any(not p.strip() for p in value): + raise ValueError("every prompt in prompts must be non-empty") + return value + + @field_validator("seeds") + @classmethod + def _seeds_json_safe(cls, value: Optional[list[int]]) -> Optional[list[int]]: + # Same JSON safe-integer bound as `seed`, so every per-image seed survives the + # round-trip through the gallery recipe. + if value is not None and any(s < 0 or s > 2**53 - 1 for s in value): + raise ValueError("every seed must be between 0 and 2**53 - 1") + return value + + @model_validator(mode = "after") + def _prompts_seeds_lengths_match(self) -> "DiffusionGenerateRequest": + if self.prompts is not None and self.seeds is not None and len(self.prompts) != len( + self.seeds + ): + raise ValueError( + f"prompts and seeds must have the same length (got {len(self.prompts)} " + f"prompts, {len(self.seeds)} seeds)" + ) + return self # Image-conditioned workflows (base64 or data-URL): init_image alone runs img2img, # init_image + mask_image runs inpaint. Both require a family with the matching pipeline or # the load is rejected. Cap each base64 string so one request can't buffer a multi-GB payload diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index b929a07519..e3b57485e3 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -14413,6 +14413,8 @@ async def generate_diffusion_image( guidance = request.guidance, seed = request.seed, batch_size = request.batch_size, + prompts = request.prompts, + seeds = request.seeds, init_image = request.init_image, mask_image = request.mask_image, strength = request.strength, @@ -14451,9 +14453,9 @@ async def generate_diffusion_image( logger.error("diffusion.generate_failed: %s", exc, exc_info = True) raise HTTPException(status_code = 500, detail = "Image generation failed.") - # Persist each image with its full recipe embedded. The diffusers batch shares one seed - # (drawn sequentially from one generator); the native sd.cpp batch uses a distinct seed - # per image, returned in ``seeds`` so each is reproducible. + # Persist each image with its full recipe embedded. BOTH engines batch with a distinct + # seed per image (diffusers via one torch.Generator per image, native sd.cpp via + # base + index), returned in ``seeds`` so each image is individually reproducible. created_at = time.time() per_image_seeds = result.get("seeds") @@ -14469,7 +14471,13 @@ async def generate_diffusion_image( image_gallery.save( image, { - "prompt": request.prompt, + # A prompts-list batch records each image's OWN prompt so its recipe + # replays exactly; otherwise every image shares the single prompt. + "prompt": ( + request.prompts[index] + if request.prompts and index < len(request.prompts) + else request.prompt + ), "negative_prompt": request.negative_prompt, # Persist the ACTUAL output size, not the request sliders: # Transform/Inpaint/Edit derive it from the uploaded image, Extend grows diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index c7ddaea868..b8ae22b831 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/studio/backend/tests/test_diffusion_backend.py @@ -269,7 +269,11 @@ class _FakePipe: "cfg_trunc_ratio": cfg_trunc_ratio, **kwargs, } + # Mirror diffusers batching: a prompt LIST yields one image per prompt, and + # num_images_per_prompt fans each prompt out. n = kwargs.get("num_images_per_prompt", 1) + if isinstance(prompt, list): + n *= len(prompt) return types.SimpleNamespace(images = [_FakeImage() for _ in range(n)]) @@ -481,9 +485,11 @@ def test_load_generate_unload_gguf(fake_runtime, tmp_path): gen2 = backend.generate(prompt = "again", seed = 99) assert gen2["seed"] == 99 - # batch_size produces that many images in one call, all sharing the seed. + # batch_size produces that many images in one call, seeded base..base+2 per image + # (matching the native engine) so each batch member replays alone from its recipe. batch = backend.generate(prompt = "batch", seed = 7, batch_size = 3) assert len(batch["images"]) == 3 and batch["seed"] == 7 + assert batch["seeds"] == [7, 8, 9] assert backend.unload()["loaded"] is False assert backend.is_loaded is False @@ -3551,3 +3557,141 @@ def test_unload_waits_for_in_flight_denoise_before_teardown(): assert unloaded.is_set() # Teardown ran exactly once, and only AFTER the denoise had exited (denoise_active was False). assert teardown_saw == [False] + + +# Batched generation (prompt/seed lists, per-image generators, OOM backoff) + + +def _load_zimage_backend(tmp_path): + (tmp_path / "model.gguf").write_bytes(b"weights") + backend = DiffusionBackend() + backend.load_pipeline( + str(tmp_path), + gguf_filename = "model.gguf", + base_repo = "base/repo", + family_override = "z-image", + ) + return backend + + +def test_generate_seed_list_uses_one_generator_per_image(fake_runtime, tmp_path): + backend = _load_zimage_backend(tmp_path) + out = backend.generate(prompt = "a sloth", seeds = [11, 22, 33, 44]) + assert len(out["images"]) == 4 + assert out["seeds"] == [11, 22, 33, 44] + assert out["seed"] == 11 # base seed = first per-image seed + call = backend._state.pipe.last_kwargs + # A uniform prompt is encoded ONCE and fanned out; each image gets its own generator. + assert call["prompt"] == "a sloth" + assert call["num_images_per_prompt"] == 4 + assert [g.manual for g in call["generator"]] == [11, 22, 33, 44] + + +def test_generate_prompt_list_one_image_per_prompt(fake_runtime, tmp_path): + backend = _load_zimage_backend(tmp_path) + out = backend.generate(prompt = "fallback", prompts = ["a", "b", "c"], seed = 100) + assert len(out["images"]) == 3 + assert out["seeds"] == [100, 101, 102] # derived from the base seed + call = backend._state.pipe.last_kwargs + assert call["prompt"] == ["a", "b", "c"] + assert call["num_images_per_prompt"] == 1 + assert [g.manual for g in call["generator"]] == [100, 101, 102] + + +def test_generate_prompt_list_with_matching_seed_list(fake_runtime, tmp_path): + backend = _load_zimage_backend(tmp_path) + out = backend.generate(prompt = "fallback", prompts = ["a", "b"], seeds = [5, 6]) + assert out["seeds"] == [5, 6] + with pytest.raises(ValueError, match = "same length"): + backend.generate(prompt = "fallback", prompts = ["a", "b"], seeds = [5]) + + +def test_generate_single_image_keeps_scalar_generator(fake_runtime, tmp_path): + # The single-image call shape is the bit-identical reference path: scalar prompt, + # ONE scalar generator (not a 1-list), num_images_per_prompt=1. + backend = _load_zimage_backend(tmp_path) + out = backend.generate(prompt = "one", seed = 5) + call = backend._state.pipe.last_kwargs + assert not isinstance(call["generator"], list) + assert call["generator"].manual == 5 + assert call["num_images_per_prompt"] == 1 + assert out["seeds"] == [5] + + +def test_generate_batched_seed_matches_solo_replay(fake_runtime, tmp_path): + # Per-image reproducibility contract: image i of a batched call is driven by the + # exact generator seed a solo replay of that image uses. + backend = _load_zimage_backend(tmp_path) + backend.generate(prompt = "p", seeds = [3, 9]) + batched = [g.manual for g in backend._state.pipe.last_kwargs["generator"]] + backend.generate(prompt = "p", seed = 9) + solo = backend._state.pipe.last_kwargs["generator"].manual + assert batched[1] == solo == 9 + + +def test_generate_prompt_list_rejected_off_txt2img(fake_runtime, tmp_path): + backend = _load_zimage_backend(tmp_path) + with pytest.raises(ValueError, match = "text-to-image only"): + backend.generate( + prompt = "x", prompts = ["a", "b"], init_image = _tiny_png_b64() + ) + + +class _CountingPipe(_FakePipe): + """Records each forward's image count; optionally OOMs above ``max_images``.""" + + def __init__(self, max_images = None): + super().__init__() + self.batch_attempts = [] + self.max_images = max_images + + def __call__(self, *, prompt = None, **kwargs): + n = kwargs.get("num_images_per_prompt", 1) + if isinstance(prompt, list): + n *= len(prompt) + self.batch_attempts.append(n) + if self.max_images is not None and n > self.max_images: + raise _FakeOutOfMemoryError("CUDA out of memory. Tried to allocate everything") + return super().__call__(prompt = prompt, **kwargs) + + +# Structural stand-in for torch.cuda.OutOfMemoryError (matched by class name). +_FakeOutOfMemoryError = type("OutOfMemoryError", (RuntimeError,), {}) + + +def test_generate_explicit_batch_size_caps_per_forward(fake_runtime, tmp_path): + backend = _load_zimage_backend(tmp_path) + pipe = _CountingPipe() + object.__setattr__(backend._state, "pipe", pipe) + out = backend.generate(prompt = "p", seeds = [1, 2, 3, 4], batch_size = 2) + assert pipe.batch_attempts == [2, 2] + assert len(out["images"]) == 4 + assert out["seeds"] == [1, 2, 3, 4] + + +def test_generate_oom_backoff_halves_the_batch(fake_runtime, tmp_path): + backend = _load_zimage_backend(tmp_path) + pipe = _CountingPipe(max_images = 2) + object.__setattr__(backend._state, "pipe", pipe) + out = backend.generate(prompt = "p", seeds = [1, 2, 3, 4]) + # The full batch OOMs once, then both halves run; images + seeds stay complete. + assert pipe.batch_attempts == [4, 2, 2] + assert len(out["images"]) == 4 + assert out["seeds"] == [1, 2, 3, 4] + + +class _BoomPipe(_CountingPipe): + """Fails every forward with a NON-OOM error (must not trigger backoff).""" + + def __call__(self, *, prompt = None, **kwargs): + self.batch_attempts.append(kwargs.get("num_images_per_prompt", 1)) + raise RuntimeError("shape mismatch") + + +def test_generate_non_oom_error_is_not_retried(fake_runtime, tmp_path): + backend = _load_zimage_backend(tmp_path) + pipe = _BoomPipe() + object.__setattr__(backend._state, "pipe", pipe) + with pytest.raises(RuntimeError, match = "shape mismatch"): + backend.generate(prompt = "p", seeds = [1, 2, 3, 4]) + assert pipe.batch_attempts == [4] # no backoff retries on a non-OOM error diff --git a/studio/backend/tests/test_diffusion_batched.py b/studio/backend/tests/test_diffusion_batched.py new file mode 100644 index 0000000000..26df099b08 --- /dev/null +++ b/studio/backend/tests/test_diffusion_batched.py @@ -0,0 +1,153 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Unit tests for the batched multi-image planning helpers (``diffusion_batched.py``). + +Pure and torch-free: job resolution (prompt lists / seed lists / legacy batch_size), +chunking, OOM split, and the OOM classifier.""" + +from __future__ import annotations + +import pytest + +from core.inference.diffusion_batched import ( + MAX_BATCH_IMAGES, + SEED_MASK, + chunk_jobs, + is_oom_error, + resolve_batch_jobs, + split_chunk, + uniform_prompt, +) + + +def _draw(): + raise AssertionError("draw_seed must not be called when seed material was supplied") + + +# --------------------------------------------------------------------------- job resolution +def test_legacy_batch_derives_sequential_seeds(): + jobs, base = resolve_batch_jobs( + prompt = "p", prompts = None, seed = 7, seeds = None, batch_size = 3, draw_seed = _draw + ) + assert jobs == [("p", 7), ("p", 8), ("p", 9)] + assert base == 7 + + +def test_single_image_draws_a_masked_random_seed(): + jobs, base = resolve_batch_jobs( + prompt = "p", + prompts = None, + seed = None, + seeds = None, + batch_size = 1, + draw_seed = lambda: (1 << 60) + 5, # over JS's safe range: must be masked + ) + assert base == ((1 << 60) + 5) & SEED_MASK + assert jobs == [("p", base)] + + +def test_prompt_list_one_job_per_prompt(): + jobs, base = resolve_batch_jobs( + prompt = "unused", + prompts = ["a", "b"], + seed = 100, + seeds = None, + batch_size = 1, + draw_seed = _draw, + ) + assert jobs == [("a", 100), ("b", 101)] + assert base == 100 + + +def test_seed_list_one_job_per_seed(): + jobs, base = resolve_batch_jobs( + prompt = "p", prompts = None, seed = None, seeds = [5, 6, 7], batch_size = 1, draw_seed = _draw + ) + assert jobs == [("p", 5), ("p", 6), ("p", 7)] + assert base == 5 + + +def test_prompt_and_seed_lists_pair_elementwise(): + jobs, base = resolve_batch_jobs( + prompt = "unused", + prompts = ["a", "b"], + seed = None, + seeds = [9, 3], + batch_size = 1, + draw_seed = _draw, + ) + assert jobs == [("a", 9), ("b", 3)] + assert base == 9 # base seed = first per-image seed + + +def test_derived_seeds_stay_json_safe_at_the_cap(): + jobs, _ = resolve_batch_jobs( + prompt = "p", + prompts = None, + seed = SEED_MASK, + seeds = None, + batch_size = 2, + draw_seed = _draw, + ) + assert all(0 <= s <= SEED_MASK for _, s in jobs) + + +@pytest.mark.parametrize( + "kwargs,match", + [ + (dict(prompts = []), "non-empty"), + (dict(prompts = ["ok", " "]), "non-empty"), + (dict(prompts = ["p"] * (MAX_BATCH_IMAGES + 1)), "at most"), + (dict(seeds = []), "non-empty"), + (dict(seeds = [1] * (MAX_BATCH_IMAGES + 1)), "at most"), + (dict(seeds = [-1]), "between 0"), + (dict(seeds = [SEED_MASK + 1]), "between 0"), + (dict(prompts = ["a", "b"], seeds = [1]), "same length"), + ], +) +def test_invalid_lists_rejected(kwargs, match): + base = dict(prompt = "p", prompts = None, seed = None, seeds = None, batch_size = 1) + base.update(kwargs) + with pytest.raises(ValueError, match = match): + resolve_batch_jobs(draw_seed = lambda: 0, **base) + + +# --------------------------------------------------------------------------------- chunking +def test_default_batch_size_runs_everything_in_one_forward(): + jobs = [("p", i) for i in range(8)] + assert chunk_jobs(jobs, 1) == [jobs] + + +def test_explicit_batch_size_caps_each_chunk(): + jobs = [("p", i) for i in range(5)] + chunks = chunk_jobs(jobs, 2) + assert [len(c) for c in chunks] == [2, 2, 1] + assert [s for c in chunks for _, s in c] == list(range(5)) # order preserved + + +def test_chunk_jobs_empty(): + assert chunk_jobs([], 4) == [] + + +def test_split_chunk_halves_and_terminates(): + chunk = [("p", i) for i in range(5)] + first, second = split_chunk(chunk) + assert first + second == chunk + assert len(first) == 3 and len(second) == 2 # first never smaller: splits terminate + with pytest.raises(ValueError): + split_chunk([("p", 0)]) + + +def test_uniform_prompt(): + assert uniform_prompt([("a", 1), ("a", 2)]) == "a" + assert uniform_prompt([("a", 1), ("b", 2)]) is None + + +# ---------------------------------------------------------------------------- OOM classifier +def test_is_oom_error_matches_class_name_and_message(): + oom_cls = type("OutOfMemoryError", (RuntimeError,), {}) + assert is_oom_error(oom_cls("boom")) + assert is_oom_error(RuntimeError("CUDA out of memory. Tried to allocate 2 GiB")) + assert not is_oom_error(RuntimeError("shape mismatch")) + assert not is_oom_error(ValueError("bad prompt")) diff --git a/studio/backend/tests/test_diffusion_compile_cache.py b/studio/backend/tests/test_diffusion_compile_cache.py index 1c4ac25795..cbb2cab307 100644 --- a/studio/backend/tests/test_diffusion_compile_cache.py +++ b/studio/backend/tests/test_diffusion_compile_cache.py @@ -254,6 +254,45 @@ def test_new_static_shape_redirties_a_hit(monkeypatch, tmp_path, fake_megacache) assert manifest["shapes"] == [[768, 768, 1], [1024, 1024, 1]] +def test_new_batch_size_is_its_own_static_shape(monkeypatch, tmp_path, fake_megacache): + # A static compile produces one artifact PER (w, h, batch): a batched generation at a + # batch size the bundle has not seen (incl. an OOM-backoff half) must re-dirty it, and + # the same (w, h) at the covered batch must not. + monkeypatch.setenv(cc._ENV_MODE, "auto") + monkeypatch.delenv(cc._ENV_SAVE, raising = False) + monkeypatch.setenv(cc._ENV_DIR, str(tmp_path)) + ctx = cc.begin(transformer = _transformer(), **_BEGIN_KW) + cc.register_shape(ctx, (1024, 1024, 8), static = True) + assert cc.save(ctx) is True + + ctx2 = cc.begin(transformer = _transformer(), **_BEGIN_KW) + assert ctx2.hit is True + cc.register_shape(ctx2, (1024, 1024, 8), static = True) + assert cc.save(ctx2) is False # covered batch: nothing new + cc.register_shape(ctx2, (1024, 1024, 32), static = True) + assert ctx2.saved is False # new batch size: new artifacts to persist + assert cc.save(ctx2) is True + manifest = json.loads(ctx2.manifest_path.read_text()) + assert manifest["shapes"] == [[1024, 1024, 8], [1024, 1024, 32]] + + +def test_gguf_quant_keys_apart_from_dense(): + # A GGUF transformer compiles a different graph (the dequant chain) than the dense + # family; the load path fingerprints it quant="gguf" so bundles never cross-hit. + efp = cc.environment_fingerprint() + base = dict( + family = "flux.1", + transformer = _transformer(), + dtype = "torch.bfloat16", + quant = None, + attention_backend = "x", + compile_kwargs = {"fullgraph": True, "dynamic": True}, + ) + dense = cc.model_fingerprint(**base) + gguf = cc.model_fingerprint(**{**base, "quant": "gguf"}) + assert cc.cache_key(efp, dense) != cc.cache_key(efp, gguf) + + def test_dynamic_compile_never_dirties(monkeypatch, tmp_path, fake_megacache): monkeypatch.setenv(cc._ENV_MODE, "auto") monkeypatch.setenv(cc._ENV_DIR, str(tmp_path)) diff --git a/studio/backend/tests/test_diffusion_cond_cache.py b/studio/backend/tests/test_diffusion_cond_cache.py new file mode 100644 index 0000000000..5ee7023490 --- /dev/null +++ b/studio/backend/tests/test_diffusion_cond_cache.py @@ -0,0 +1,165 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Unit tests for the inference-side conditioning cache (``diffusion_cond_cache.py``). + +Runs against the real torch/safetensors on CPU with a stub ``encode_prompt`` pipe, so +the wrapper's hit/miss/bypass behaviour and the on-disk reuse (the reason warm repeats +never run the text encoder) are exercised without any model weights.""" + +from __future__ import annotations + +import pytest +import torch + +from core.inference import diffusion_cond_cache as cond_cache + + +class _EncodePipe: + """A pipe exposing a deterministic ``encode_prompt`` that counts its calls.""" + + def __init__(self): + self.calls = 0 + self._execution_device = "cpu" + + def encode_prompt( + self, + prompt, + device = None, + num_images_per_prompt = 1, + max_sequence_length = 256, + prompt_embeds = None, + ): + if prompt_embeds is not None: + return (prompt_embeds, None) + self.calls += 1 + value = float(sum(map(ord, str(prompt)))) + return ( + torch.full((num_images_per_prompt, 4), value), + None, # mask-less returns must round-trip (None slots) + ) + + +@pytest.fixture +def cache_env(tmp_path, monkeypatch): + monkeypatch.setenv("UNSLOTH_DIFFUSION_COND_CACHE_DIR", str(tmp_path)) + return tmp_path + + +def _install(pipe, **overrides): + kwargs = dict(family = "flux.1", repo_id = "unsloth/repo", dtype = "torch.bfloat16") + kwargs.update(overrides) + return cond_cache.install(pipe, **kwargs) + + +def test_off_by_default(monkeypatch): + monkeypatch.delenv("UNSLOTH_DIFFUSION_COND_CACHE_DIR", raising = False) + pipe = _EncodePipe() + assert _install(pipe) is False + assert pipe.encode_prompt.__func__ is _EncodePipe.encode_prompt # untouched + + +def test_blank_dir_means_off(monkeypatch): + # Same semantics as the trainers' cond_cache_dir: blank is "off", not cwd. + monkeypatch.setenv("UNSLOTH_DIFFUSION_COND_CACHE_DIR", " ") + assert cond_cache.cache_dir() is None + assert _install(_EncodePipe()) is False + + +def test_repeated_prompt_skips_the_encode_forward(cache_env): + pipe = _EncodePipe() + assert _install(pipe) is True + first = pipe.encode_prompt("a sloth", device = "cpu") + second = pipe.encode_prompt("a sloth", device = "cpu") + assert pipe.calls == 1 # warm repeat never ran the text encoder + assert torch.equal(first[0], second[0]) + assert first[1] is None and second[1] is None # None slot round-trips + assert pipe._unsloth_cond_cache_stats == {"hits": 1, "misses": 1} + + +def test_distinct_prompts_and_arguments_key_separately(cache_env): + pipe = _EncodePipe() + _install(pipe) + pipe.encode_prompt("a sloth") + pipe.encode_prompt("a fox") + pipe.encode_prompt("a sloth", num_images_per_prompt = 4) # shape-changing arg + assert pipe.calls == 3 + + +def test_device_argument_excluded_from_the_key(cache_env): + pipe = _EncodePipe() + _install(pipe) + pipe.encode_prompt("a sloth", device = "cpu") + out = pipe.encode_prompt("a sloth", device = torch.device("cpu")) + assert pipe.calls == 1 # placement detail: still a hit, moved to the target + assert out[0].device.type == "cpu" + + +def test_warm_reuse_across_installs(cache_env): + # A NEW pipe (fresh load) over the same directory hits the persisted entry without + # ever encoding -- the property that lets warm loads keep the text encoder off GPU. + first = _EncodePipe() + _install(first) + reference = first.encode_prompt("a sloth") + second = _EncodePipe() + _install(second) + warm = second.encode_prompt("a sloth") + assert second.calls == 0 + assert torch.equal(reference[0], warm[0]) + + +def test_load_fingerprint_keys_apart(cache_env): + # A different repo / TE quant produces different embeddings: never cross-hit. + a = _EncodePipe() + _install(a) + a.encode_prompt("a sloth") + b = _EncodePipe() + _install(b, repo_id = "unsloth/other-repo") + b.encode_prompt("a sloth") + c = _EncodePipe() + _install(c, te_quant = "fp8") + c.encode_prompt("a sloth") + assert (a.calls, b.calls, c.calls) == (1, 1, 1) + + +def test_lora_attached_bypasses_the_cache(cache_env): + pipe = _EncodePipe() + _install(pipe) + pipe._unsloth_loras = ("style",) # adapters may target the text encoders + pipe.encode_prompt("a sloth") + pipe.encode_prompt("a sloth") + assert pipe.calls == 2 + assert pipe._unsloth_cond_cache_stats == {"hits": 0, "misses": 0} + + +class _ListEncodePipe(_EncodePipe): + """Returns per-prompt embedding LISTS like Z-Image's ``encode_prompt``.""" + + def encode_prompt(self, prompt, device = None, do_classifier_free_guidance = True): + self.calls += 1 + prompts = prompt if isinstance(prompt, list) else [prompt] + embeds = [torch.full((1, 4), float(sum(map(ord, p)))) for p in prompts] + return (embeds, None) + + +def test_tensor_list_slots_round_trip(cache_env): + # Z-Image returns list-of-tensors slots; the flatten/unflatten layout must + # reproduce them exactly on a warm hit. + pipe = _ListEncodePipe() + _install(pipe) + cold = pipe.encode_prompt(["a", "bb"]) + warm = pipe.encode_prompt(["a", "bb"]) + assert pipe.calls == 1 + assert isinstance(warm[0], list) and len(warm[0]) == 2 + assert all(torch.equal(c, w) for c, w in zip(cold[0], warm[0])) + assert warm[1] is None + + +def test_tensor_arguments_pass_through_uncached(cache_env): + pipe = _EncodePipe() + _install(pipe) + supplied = torch.ones(1, 4) + out = pipe.encode_prompt("a sloth", prompt_embeds = supplied) + assert out[0] is supplied + assert pipe.calls == 0 + assert pipe._unsloth_cond_cache_stats == {"hits": 0, "misses": 0}