From 7085d421c249b1a1a3a42186426730241b6bdc2d Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 26 Jul 2026 08:11:20 +0000 Subject: [PATCH] Fix batched generation crashes, cache keying and unreplayable recipes Four bugs in the batched inference path, all found by review: - A mixed-prompt batch sent a scalar negative prompt against a prompt list. Z-Image asserts on the length, and Qwen-Image, Krea 2 and FLUX true-CFG encode a batch-1 negative against batch-N latents and fail in the transformer's text/image concat. Broadcast it to match the batch. - The FBCache step-cache reset sat above the chunk loop. diffusers only resets that state at the end of a successful call, so a forward that raised (the OOM the backoff is meant to recover) left its own residual behind and the halved retry died on a shape mismatch. Reset before every forward instead. - The conditioning cache keyed on the checkpoint alone, but a GGUF or single-file load takes its text encoders from the companion base, so the same checkpoint against a different base reused the previous base's embeddings. Key the base too. - Gallery records stored the base seed and the requested batch size even when a prompts/seeds list drove the run, so restoring the second image of seeds=[5, 99] replayed seed 5. List-driven outputs now record as single-image recipes on their own seed. Also bound strength above 0: every img2img pipeline derives its step count from it, so 0 leaves zero denoising steps and either raises or, on SDXL, crashes on empty latents. --- studio/backend/core/inference/diffusion.py | 52 ++++++++---- .../core/inference/diffusion_cond_cache.py | 13 ++- studio/backend/models/inference.py | 20 +++-- studio/backend/routes/inference.py | 14 +++- .../backend/tests/test_diffusion_backend.py | 82 +++++++++++++++++++ .../tests/test_diffusion_cond_cache.py | 18 ++++ studio/backend/tests/test_diffusion_routes.py | 80 ++++++++++++++++++ 7 files changed, 250 insertions(+), 29 deletions(-) diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 1b86527622..2fc5178ada 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -1781,12 +1781,16 @@ class DiffusionBackend: # 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; + # cache key reflects the encoders that actually run. ``base`` keys the + # companion repo the TEXT ENCODERS came from (a GGUF/single-file load + # takes them from the base, so the same checkpoint against a different + # base must not reuse the previous base's embeddings). Instance-level; # dies with the pipe on unload. cond_cache.install( pipe, family = fam.name, repo_id = repo_id, + base_repo = base, dtype = dtype, te_quant = te_quant, logger = logger, @@ -2556,18 +2560,19 @@ class DiffusionBackend: @staticmethod def _reset_step_cache(pipe: Any) -> None: - """Clear the transformer's stateful step cache (FBCache) before a generation. + """Clear the transformer's stateful step cache (FBCache) before a forward. diffusers keys FBCache residuals by cache context ("cond"/"uncond") on the - long-lived transformer, and neither the pipeline nor the context exit resets - them. The transformer-level reset entry point is ``_reset_stateful_cache`` in - diffusers 0.39 (``reset_stateful_hooks`` lives only on the HookRegistry, so a - getattr for it on the transformer is a silent no-op), and no pipeline calls it. - This backend reuses one resident pipe across generations, so without a reset the - next generation's first step compares its first-block residual against the - PREVIOUS request's -- a tensor-shape mismatch when the resolution/batch changed, - or a stale-cache reuse otherwise. Best-effort: a transformer without the hook - (uncached load) is a silent no-op.""" + long-lived transformer. The context exit does NOT reset them; the end of a + pipeline ``__call__`` does, via ``maybe_free_model_hooks()`` -- but only when + the call RETURNS. A call that raised (an OOM this generate() backs off from, a + cancelled denoise, a failed prior request) leaves its own batch's residual on + the resident transformer, and the next forward's first step then compares + against it: a tensor-shape mismatch when the resolution/batch changed, or a + stale-cache reuse otherwise. The transformer-level reset entry point is + ``_reset_stateful_cache`` in diffusers 0.39 (``reset_stateful_hooks`` lives only + on the HookRegistry, so a getattr for it on the transformer is a silent no-op). + Best-effort: a transformer without the hook (uncached load) is a silent no-op.""" transformer = getattr(pipe, "transformer", None) reset = getattr(transformer, "_reset_stateful_cache", None) or getattr( transformer, "reset_stateful_hooks", None @@ -3034,11 +3039,6 @@ class DiffusionBackend: + ("reaches" if toggled else "is below") + f" {FBCACHE_MIN_STEPS}" ) - # Start each generation from a clean step cache: prior FBCache residuals would - # otherwise be compared against this first step (shape mismatch / stale reuse). - if state.transformer_cache: - self._reset_step_cache(state.pipe) - self._gen = gen images: list[Any] = [] per_image_seeds: list[int] = [] @@ -3063,10 +3063,28 @@ class DiffusionBackend: chunk_kwargs["generator"] = generators chunk_kwargs["num_images_per_prompt"] = len(chunk) else: - # Distinct prompts: one image per prompt in a single forward. + # Distinct prompts: one image per prompt in a single forward. The + # negative prompt must be broadcast to match: a pipeline either + # asserts on the length (Z-Image) or encodes a batch-1 negative + # against batch-N latents and fails in the transformer's txt/img + # concat (Qwen-Image, Krea 2, FLUX true-CFG). Only the pipes that + # already expand a scalar themselves (SDXL, Lumina 2, HiDream) are + # unaffected, so broadcast here for all of them. chunk_kwargs["prompt"] = [p for p, _ in chunk] chunk_kwargs["generator"] = generators chunk_kwargs["num_images_per_prompt"] = 1 + if isinstance(chunk_kwargs.get("negative_prompt"), str): + chunk_kwargs["negative_prompt"] = [ + chunk_kwargs["negative_prompt"] + ] * len(chunk) + # Start every forward from a clean step cache. diffusers resets the + # transformer's FBCache state at the END of a successful __call__ + # (maybe_free_model_hooks), but a call that RAISED -- the OOM below, a + # cancelled denoise, a failed prior generation -- leaves the residual of + # its own batch behind, and the next (differently sized) forward then + # compares against it: "size of tensor a (16) must match ... b (32)". + if state.transformer_cache: + self._reset_step_cache(state.pipe) try: # inference_mode is faster than no_grad and numerically identical here. with torch.inference_mode(): diff --git a/studio/backend/core/inference/diffusion_cond_cache.py b/studio/backend/core/inference/diffusion_cond_cache.py index e65ab29015..608169f4a5 100644 --- a/studio/backend/core/inference/diffusion_cond_cache.py +++ b/studio/backend/core/inference/diffusion_cond_cache.py @@ -22,8 +22,9 @@ 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 +numerics (family, checkpoint repo, the companion base repo that supplies the +text encoders, 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. @@ -117,6 +118,7 @@ def install( family: str, repo_id: str, dtype: Any, + base_repo: Any = None, te_quant: Any = None, logger: Any = None, ) -> bool: @@ -142,9 +144,14 @@ def install( logger.warning("diffusion.cond_cache: install failed: %s", exc) return False - # Everything beyond the call arguments that changes the embedding numerics. + # Everything beyond the call arguments that changes the embedding numerics. A GGUF / + # single-file checkpoint takes its TEXT ENCODERS from the companion base repo, so the + # base identity must key the cache too: the same checkpoint reloaded against a different + # base would otherwise hit entries encoded by the previous base's encoder. Defaults to + # the checkpoint itself (a full pipeline is its own base). load_fp = { "repo": str(repo_id), + "base": str(base_repo) if base_repo else str(repo_id), "dtype": str(dtype), "te_quant": str(te_quant) if te_quant is not None else "none", "diffusers": _diffusers_version(), diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 8701be6e4c..6c0bfbd1d1 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -2455,10 +2455,16 @@ class DiffusionGenerateRequest(BaseModel): ) strength: Optional[float] = Field( None, - ge = 0.0, + # EXCLUSIVE lower bound: strength 0 does not "keep the source". Every diffusers + # img2img/inpaint pipeline derives its step count from it (t_start = + # num_inference_steps - int(num_inference_steps * strength)), so 0 leaves zero + # denoising steps: FLUX/Qwen/Z-Image raise "the number of pipeline steps is 0 which + # is < 1", and SDXL img2img has no such guard and crashes on empty latents (a 500). + # The UI slider already starts at 0.1; reject 0 as a 422 instead of a pipeline error. + gt = 0.0, le = 1.0, - description = "img2img/inpaint denoise strength: 0 keeps the source, 1 fully " - "redraws it. Ignored for txt2img.", + description = "img2img/inpaint denoise strength: low values stay close to the " + "source, 1 fully redraws it. Must be greater than 0. Ignored for txt2img.", ) upscale: Optional[float] = Field( None, @@ -2554,9 +2560,11 @@ class GalleryImage(BaseModel): batch_seed: Optional[int] = Field( None, description = ( - "Base seed the batch was launched with. The native engine derives per-image seeds as " - "base + index, so restore must replay from this base, not from the derived per-image " - "seed; older records without it fall back to seed." + "Seed restore must replay this image from. For a batch_size batch that is the base " + "seed the batch launched with (the native engine derives per-image seeds as " + "base + index, so the derived seed alone would not reproduce it); for a " + "prompts/seeds list, where each image carries its own seed, it equals seed. " + "Older records without it fall back to seed." ), ) batch_index: int = Field(0, description = "Position within its batch (0-based)") diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index ed4605573c..0df099c300 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -16136,6 +16136,12 @@ async def generate_diffusion_image( # base + index), returned in ``seeds`` so each image is individually reproducible. created_at = time.time() per_image_seeds = result.get("seeds") + # A prompts/seeds LIST drives the image count and each image's own seed, so + # ``batch_size`` is only a per-forward cap there and the base seed no longer replays + # image i (seeds=[5, 99] would restore 5 for the 99 image). Persist those outputs as + # single-image recipes keyed on their OWN seed instead, so the gallery's Restore + # reproduces every image and not just the first. + list_driven = bool(request.prompts or request.seeds) def _persist() -> list[dict]: records = [] @@ -16169,13 +16175,15 @@ async def generate_diffusion_image( # Base seed the batch launched with. The native engine derives per-image # seeds as base + index, so ``seed`` above is already advanced for index>0; # restore replays from this base (diffusers shares one seed, so base == seed). - "batch_seed": result["seed"], + # A list-driven image carries its OWN seed instead (see ``list_driven``). + "batch_seed": seed if list_driven else result["seed"], # Position within the batch (shared timestamp), so the export filename # stays unique. "batch_index": index, # The batch shares one seed, so reproducing a batch_index>0 image needs - # the original batch_size: persist it so restore can replay. - "batch_size": request.batch_size, + # the original batch_size: persist it so restore can replay. A list-driven + # image needs no replay -- it restores as a single image on its own seed. + "batch_size": 1 if list_driven else request.batch_size, "model": result.get("repo_id"), "loras": ( [f"{l.id}:{l.weight:g}" for l in request.loras] if request.loras else [] diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index 4f0085daf2..56511ffcd0 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/studio/backend/tests/test_diffusion_backend.py @@ -3724,3 +3724,85 @@ def test_generate_non_oom_error_is_not_retried(fake_runtime, tmp_path): 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 + + +def test_generate_broadcasts_negative_prompt_across_a_mixed_prompt_batch(fake_runtime, tmp_path): + # A prompt LIST must carry a matching negative-prompt LIST: ZImagePipeline.encode_prompt + # asserts len(prompt) == len(negative_prompt), and the pipes that encode the negative + # separately (Qwen-Image / Krea 2 / FLUX true-CFG) would build batch-1 negative embeds + # against batch-N latents and fail in the transformer's txt/img concat. + backend = _load_zimage_backend(tmp_path) + backend.generate(prompt = "fallback", prompts = ["a", "b", "c"], negative_prompt = "blurry") + call = backend._state.pipe.last_kwargs + assert call["prompt"] == ["a", "b", "c"] + assert call["negative_prompt"] == ["blurry", "blurry", "blurry"] + # An empty negative prompt is still omitted entirely (never sent as [""] * n). + backend.generate(prompt = "fallback", prompts = ["a", "b"]) + assert backend._state.pipe.last_kwargs["negative_prompt"] is None + + +def test_generate_keeps_a_scalar_negative_prompt_off_the_list_paths(fake_runtime, tmp_path): + # Uniform-prompt and single-image forwards pass a SCALAR prompt, so the negative prompt + # must stay scalar too (a list would mismatch the batch-1 positive encode). + backend = _load_zimage_backend(tmp_path) + backend.generate(prompt = "a sloth", seeds = [1, 2, 3], negative_prompt = "blurry") + assert backend._state.pipe.last_kwargs["prompt"] == "a sloth" + assert backend._state.pipe.last_kwargs["negative_prompt"] == "blurry" + backend.generate(prompt = "a sloth", seed = 1, negative_prompt = "blurry") + assert backend._state.pipe.last_kwargs["negative_prompt"] == "blurry" + + +class _TracingPipe(_CountingPipe): + """Appends ``("call", n)`` to a shared trace so resets can be interleaved with forwards.""" + + def __init__(self, trace, max_images = None): + super().__init__(max_images = max_images) + self.trace = trace + + def __call__(self, *, prompt = None, **kwargs): + n = kwargs.get("num_images_per_prompt", 1) + if isinstance(prompt, list): + n *= len(prompt) + self.trace.append(("call", n)) + return super().__call__(prompt = prompt, **kwargs) + + +def test_generate_resets_the_step_cache_before_an_oom_retry(fake_runtime, tmp_path): + # A forward that RAISES skips the pipeline's end-of-__call__ maybe_free_model_hooks(), + # so its FBCache head-block residual stays on the resident transformer. Without a reset + # before each retry the halved chunk compares a batch-2 residual against the stale + # batch-4 one ("size of tensor a (2) must match ... b (4)"), turning a recoverable OOM + # into a hard failure. + backend = _load_zimage_backend(tmp_path) + trace: list = [] + pipe = _TracingPipe(trace, max_images = 2) + pipe.transformer = types.SimpleNamespace( + _reset_stateful_cache = lambda: trace.append(("reset",)) + ) + object.__setattr__(backend._state, "pipe", pipe) + object.__setattr__(backend._state, "transformer_cache", "fbcache") + out = backend.generate(prompt = "p", seeds = [1, 2, 3, 4]) + assert len(out["images"]) == 4 and out["seeds"] == [1, 2, 3, 4] + # Every forward -- including both post-OOM retries -- is preceded by a reset. + assert trace == [ + ("reset",), + ("call", 4), + ("reset",), + ("call", 2), + ("reset",), + ("call", 2), + ] + + +def test_generate_resets_the_step_cache_before_every_chunk(fake_runtime, tmp_path): + # Same guarantee for an explicit per-forward cap (no OOM involved). + backend = _load_zimage_backend(tmp_path) + trace: list = [] + pipe = _TracingPipe(trace) + pipe.transformer = types.SimpleNamespace( + _reset_stateful_cache = lambda: trace.append(("reset",)) + ) + object.__setattr__(backend._state, "pipe", pipe) + object.__setattr__(backend._state, "transformer_cache", "fbcache") + backend.generate(prompt = "p", seeds = [1, 2, 3], batch_size = 2) + assert trace == [("reset",), ("call", 2), ("reset",), ("call", 1)] diff --git a/studio/backend/tests/test_diffusion_cond_cache.py b/studio/backend/tests/test_diffusion_cond_cache.py index ea7f9b410f..a786939196 100644 --- a/studio/backend/tests/test_diffusion_cond_cache.py +++ b/studio/backend/tests/test_diffusion_cond_cache.py @@ -122,6 +122,24 @@ def test_load_fingerprint_keys_apart(cache_env): assert (a.calls, b.calls, c.calls) == (1, 1, 1) +def test_companion_base_keys_apart(cache_env): + # A GGUF / single-file checkpoint takes its TEXT ENCODERS from the companion base, so + # the SAME checkpoint reloaded against a different base must re-encode rather than reuse + # the previous base's embeddings (silently different conditioning otherwise). + first = _EncodePipe() + _install(first, repo_id = "org/model-GGUF", base_repo = "base/one") + first.encode_prompt("a sloth") + second = _EncodePipe() + _install(second, repo_id = "org/model-GGUF", base_repo = "base/two") + second.encode_prompt("a sloth") + assert (first.calls, second.calls) == (1, 1) + # The same base is still a warm hit (the whole point of the cache). + third = _EncodePipe() + _install(third, repo_id = "org/model-GGUF", base_repo = "base/one") + third.encode_prompt("a sloth") + assert third.calls == 0 + + def test_lora_attached_bypasses_the_cache(cache_env): pipe = _EncodePipe() _install(pipe) diff --git a/studio/backend/tests/test_diffusion_routes.py b/studio/backend/tests/test_diffusion_routes.py index 0b225f0376..a7ac02d1f2 100644 --- a/studio/backend/tests/test_diffusion_routes.py +++ b/studio/backend/tests/test_diffusion_routes.py @@ -99,10 +99,24 @@ class _FakeBackend: *, seed = None, batch_size = 1, + prompts = None, + seeds = None, **kwargs, ): if not self.loaded: raise RuntimeError("No diffusion model is loaded.") + if prompts is not None or seeds is not None: + # List-driven batch: the LIST sets the image count and each image's own seed + # (batch_size is only a per-forward cap), exactly as the real engine reports it. + base = seeds[0] if seeds else (seed if seed is not None else 4242) + count = len(prompts) if prompts is not None else len(seeds) + per_image = seeds if seeds is not None else [base + i for i in range(count)] + return { + "images": [object() for _ in range(count)], + "seed": base, + "seeds": list(per_image), + "repo_id": "x/z-image", + } # The real backend returns the PIL images; the route persists them. The # fake returns sentinels since image_gallery is stubbed in the fixture. return { @@ -369,6 +383,72 @@ def test_generate_batch_size_persists_each_image(client): assert len(client.get("/api/inference/images/gallery").json()["images"]) == 3 +def test_generate_seed_list_records_replay_from_each_own_seed(client): + # A seeds LIST sets each image's own seed, so the recipe must NOT claim the base seed + + # the request's batch_size: restore prefers batch_seed (frontend restoreSettings does + # `batch_seed ?? seed`), which would regenerate seed 5 for the seed-99 image. + client.post( + "/api/inference/images/load", json = {"model_path": "x/z-image", "gguf_filename": "q.gguf"} + ) + resp = client.post( + "/api/inference/images/generate", + json = {"prompt": "p", "seeds": [5, 99]}, + ) + assert resp.status_code == 200 + images = resp.json()["images"] + assert [i["seed"] for i in images] == [5, 99] + assert [i["batch_seed"] for i in images] == [5, 99] # replays THIS image, not the base + assert [i["batch_size"] for i in images] == [1, 1] # as a single image, not a batch + + +def test_generate_prompt_list_records_each_prompt_and_seed(client): + client.post( + "/api/inference/images/load", json = {"model_path": "x/z-image", "gguf_filename": "q.gguf"} + ) + resp = client.post( + "/api/inference/images/generate", + json = {"prompt": "unused", "prompts": ["a cat", "a dog"], "seed": 10}, + ) + assert resp.status_code == 200 + images = resp.json()["images"] + assert [i["prompt"] for i in images] == ["a cat", "a dog"] + assert [i["seed"] for i in images] == [10, 11] + assert [i["batch_seed"] for i in images] == [10, 11] + assert [i["batch_size"] for i in images] == [1, 1] + + +def test_generate_legacy_batch_still_records_the_base_seed_and_size(client): + # The batch_size path is unchanged: those images DO share one base seed, so restore + # must replay the whole batch (base seed + batch_size), not the derived per-image seed. + client.post( + "/api/inference/images/load", json = {"model_path": "x/z-image", "gguf_filename": "q.gguf"} + ) + resp = client.post( + "/api/inference/images/generate", + json = {"prompt": "p", "batch_size": 3, "seed": 5}, + ) + images = resp.json()["images"] + assert all(i["batch_seed"] == 5 for i in images) + assert all(i["batch_size"] == 3 for i in images) + assert [i["batch_index"] for i in images] == [0, 1, 2] + + +def test_generate_request_rejects_zero_denoise_strength(): + # strength 0 does NOT keep the source: every diffusers img2img/inpaint pipeline derives + # t_start = steps - int(steps * strength), so 0 leaves zero denoising steps (FLUX/Qwen/ + # Z-Image raise "the number of pipeline steps is 0 which is < 1"; SDXL img2img has no + # such guard and crashes on empty latents). Reject it as a 422 up front. + import pydantic + + from models.inference import DiffusionGenerateRequest + + with pytest.raises(pydantic.ValidationError): + DiffusionGenerateRequest(prompt = "x", strength = 0.0) + assert DiffusionGenerateRequest(prompt = "x", strength = 0.1).strength == 0.1 + assert DiffusionGenerateRequest(prompt = "x", strength = 1.0).strength == 1.0 + assert DiffusionGenerateRequest(prompt = "x").strength is None # unset stays the pipe default + + def test_gallery_pagination(client): client.post( "/api/inference/images/load", json = {"model_path": "x/z-image", "gguf_filename": "q.gguf"}