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

Batch diffusion inference with per-image seeds, an inference conditioning cache, and GGUF loader fixes

Browse source 
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.
This commit is contained in:
Daniel Han 2026-07-22 07:03:38 +00:00
commit 7f0ccdbf01
11 changed files with 1138 additions and 35 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -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:

138
studio/backend/core/inference/diffusion_batched.py Normal file
View file

@ -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()

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

@ -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

223
studio/backend/core/inference/diffusion_cond_cache.py Normal file
View file

@ -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

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

@ -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

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

@ -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

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

@ -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

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

@ -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

153
studio/backend/tests/test_diffusion_batched.py Normal file
View file

@ -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"))

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

@ -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))

165
studio/backend/tests/test_diffusion_cond_cache.py Normal file
View file

@ -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}