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

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

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

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

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 []

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

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)

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