Merge remote-tracking branch 'origin/diffusion-controlnet' into diffusion-sdxl

studio/backend/core/inference/diffusion.py

@@ -520,6 +520,9 @@ class DiffusionBackend:
         model_kind: Optional[str] = None,
     ) -> dict[str, Any]:
         """Validate, then run the (slow) load on a daemon thread. Returns at once."""
+        # A blank token (the Studio default when none is configured) must mean
+        # "anonymous", not an explicit empty credential the Hub rejects with 401.
+        hf_token = (hf_token.strip() if isinstance(hf_token, str) else hf_token) or None
         fam = self.validate_load_request(
             repo_id,
             gguf_filename = gguf_filename,
@@ -649,6 +652,18 @@ class DiffusionBackend:
         fraction = min(downloaded / expected, 1.0) if expected > 0 else 0.0
         return _progress("downloading", downloaded, expected, fraction)
 
+    def loading_repo_ids(self) -> tuple[str, ...]:
+        """Repo ids an in-flight background load is downloading (empty when idle).
+
+        The delete-cached guard needs this: during a load ``status()["loaded"]`` is
+        still False, but deleting the target repo (or its companion base) would yank
+        blobs and snapshot files from under the download/assembly."""
+        with self._lock:
+            loading = self._loading
+            if loading is None or loading.error is not None:
+                return ()
+            return tuple(r for r in (loading.repo_id, loading.base_repo) if r)
+
     @staticmethod
     def _estimate_download_bytes(
         repo_id: str,
@@ -760,7 +775,10 @@ class DiffusionBackend:
         hf_token = hf_token or None
 
         # Validate first (cheap, no torch/diffusers) so a direct call with a bad
-        # family fails with ValueError even in a no-diffusers runtime.
+        # family fails with ValueError even in a no-diffusers runtime. Sanitize the
+        # token here too (direct callers bypass begin_load): a blank string must
+        # load anonymously, not 401 as an explicit empty credential.
+        hf_token = (hf_token.strip() if isinstance(hf_token, str) else hf_token) or None
         fam = self.validate_load_request(
             repo_id,
             gguf_filename = gguf_filename,
@@ -784,12 +802,18 @@ class DiffusionBackend:
         # The cancel makes that wait ~one step (or the rest of the denoise for a
         # pipeline that ignores the step callback).
         with self._lock:
+            # Bail BEFORE signalling any cancel if this load was already superseded (an
+            # unload/eviction or a newer load bumped the token while we were resolving /
+            # downloading). Otherwise a stale worker would abort an unrelated, still-live
+            # generation from the CURRENT model and only then discover it has nothing to do.
+            if _load_token is not None and _load_token != self._load_token:
+                raise RuntimeError("Diffusion load was cancelled.")
             if self._active_generate_cancel is not None:
                 self._active_generate_cancel.set()
         with self._generate_lock:
             with self._lock:
-                # Bail before the (slow, VRAM-heavy) build if an unload/eviction or a
-                # newer load superseded this one while we were resolving/downloading.
+                # Re-check under the generate lock: a newer load/unload may have superseded
+                # this one while we waited for the in-flight denoise to exit.
                 if _load_token is not None and _load_token != self._load_token:
                     raise RuntimeError("Diffusion load was cancelled.")
 
@@ -858,6 +882,12 @@ class DiffusionBackend:
                         )
                         pipe = None
                         transformer_quant_engaged = None
+                        # Drop the exception (and its traceback) BEFORE clearing the cache:
+                        # exc.__traceback__ keeps _load_dense_quant_pipeline's frame -- and
+                        # thus its partially-built dense bf16 transformer/pipe -- alive, so
+                        # clear_gpu_cache() could not otherwise reclaim that VRAM before the
+                        # GGUF build (the OOM-fallback path this cleanup exists for).
+                        del exc
                         clear_gpu_cache()
 
                 if pipe is None:
@@ -1435,6 +1465,26 @@ class DiffusionBackend:
             raise ValueError(f"Failed to apply LoRA: {exc}") from exc
         pipe._unsloth_loras = desired
 
+    @staticmethod
+    def _reset_step_cache(pipe: Any) -> None:
+        """Clear the transformer's stateful step cache (FBCache) before a generation.
+
+        diffusers keys FBCache residuals by cache context ("cond"/"uncond") on the
+        long-lived transformer, and neither the pipeline nor the context exit resets
+        them (``StateManager`` only clears via ``reset_stateful_hooks``, which no
+        pipeline calls). 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."""
+        transformer = getattr(pipe, "transformer", None)
+        reset = getattr(transformer, "reset_stateful_hooks", None)
+        if callable(reset):
+            try:
+                reset()
+            except Exception:  # noqa: BLE001 — reset is best-effort, never fail a generation
+                pass
+
     def generate(
         self,
         *,
@@ -1744,6 +1794,13 @@ class DiffusionBackend:
                 if "callback_on_step_end" in call_params:
                     kwargs["callback_on_step_end"] = _on_step
 
+                # Start each generation from a clean step cache: FBCache residuals from
+                # a prior request on this resident pipe would otherwise be compared
+                # against this generation's first step (shape mismatch on a resolution/
+                # batch change, or stale reuse). No-op when no cache is engaged.
+                if state.transformer_cache:
+                    self._reset_step_cache(state.pipe)
+
                 self._gen = gen
                 try:
                     # inference_mode is strictly faster than the no_grad diffusers

studio/backend/core/inference/diffusion_memory.py

@@ -434,20 +434,20 @@ def apply_memory_plan(
         # is in the low-VRAM situation where the decode-time spike can OOM, so turn VAE
         # tiling on now (if not already engaged) to cap it.
         nonlocal tiling_engaged
-        pipe.enable_model_cpu_offload()
+        pipe.enable_model_cpu_offload(device = device)
         if not tiling_engaged:
             tiling_engaged = _enable_vae_saver(pipe, "enable_vae_tiling", "enable_tiling", logger)
 
     policy = plan.offload_policy
     if policy == OFFLOAD_MODEL:
-        pipe.enable_model_cpu_offload()
+        pipe.enable_model_cpu_offload(device = device)
     elif policy == OFFLOAD_GROUP:
         if not _apply_group_offload(pipe, device, logger):
             _fallback_to_model_offload()
             policy = OFFLOAD_MODEL
     elif policy == OFFLOAD_SEQUENTIAL:
         try:
-            pipe.enable_sequential_cpu_offload()
+            pipe.enable_sequential_cpu_offload(device = device)
         except Exception as exc:  # noqa: BLE001 — keep the model loadable
             if logger is not None:
                 logger.warning(

studio/backend/core/inference/diffusion_prequant.py

@@ -193,6 +193,15 @@ def load_prequantized_transformer(
             transformer.load_state_dict(state_dict, strict = True, assign = True)
 
         transformer = transformer.to(device)
+        # Built via from_config (not from_pretrained), so it starts in TRAIN mode; the
+        # dense and GGUF paths load through from_pretrained, which diffusers documents as
+        # returning an eval()'d module. Match that here so any train/eval-sensitive layer
+        # (e.g. dropout) can't make prequant inference nondeterministic or diverge from
+        # the other load paths.
+        try:
+            transformer.eval()
+        except Exception:  # noqa: BLE001 — eval() is best-effort
+            pass
         try:  # diagnostic marker, mirrors the runtime-quant path
             transformer._unsloth_runtime_quant = scheme
         except Exception:  # noqa: BLE001 — marker is best-effort

studio/backend/core/inference/sd_cpp_args.py

@@ -207,8 +207,15 @@ def build_sd_cpp_command(
     """
     if not files.diffusion_model:
         raise ValueError("diffusion_model path is required")
-    if not str(params.prompt).strip():
+    # ``(prompt or "")`` so a None prompt is rejected here rather than slipping past
+    # ``str(None)`` == "None" (truthy) and landing in argv as a literal "None".
+    if not (params.prompt or "").strip():
         raise ValueError("prompt is required")
+    # sd-cli inpaint needs the source image too: a --mask with no --init-img is an
+    # invalid invocation (sd-cli has nothing to inpaint into), so reject it here with a
+    # clear error instead of emitting a command that fails deep in sd-cli.
+    if params.mask and not params.init_img:
+        raise ValueError("init_img is required when mask is set (inpaint needs a source image)")
 
     cmd: list[str] = [binary, "--mode", DEFAULT_MODE, "--diffusion-model", files.diffusion_model]
     for flag, value in (

studio/backend/core/inference/sd_cpp_backend.py

@@ -457,6 +457,16 @@ class SdCppDiffusionBackend:
         fraction = min(downloaded / expected, 1.0) if expected > 0 else 0.0
         return _progress("downloading", downloaded, expected, fraction)
 
+    def loading_repo_ids(self) -> tuple[str, ...]:
+        """Repo ids an in-flight background load is downloading (empty when idle).
+        Mirrors the diffusers backend so the delete-cached guard can query whichever
+        engine is active without caring which one it got."""
+        with self._lock:
+            loading = self._loading
+            if loading is None or loading.error is not None:
+                return ()
+            return tuple(r for r in (loading.repo_id, loading.base_repo) if r)
+
     # ── Generate ───────────────────────────────────────────────────────────
 
     def generate(