From d0c5cf6e073848caa646b8a633e04c42c0793fae Mon Sep 17 00:00:00 2001
From: oobabooga <112222186+oobabooga@users.noreply.github.com>
Date: Wed, 1 Jul 2026 16:58:39 -0300
Subject: [PATCH 1/4] Fix diffusion flag leak, sd-cli orphan, and native family
 fallback

Six correctness fixes to the diffusion stack, found reviewing the merged
phase PRs on this branch:

- load_pipeline: restore the try/finally guard around the speed/quant/
  placement span. A failure after apply_speed_optims (e.g. OOM in quant or
  the memory plan) left TF32/cudnn flags flipped process-wide and the
  half-built pipe resident in VRAM. Now restores the flags and frees VRAM
  on a failed load.
- sd-cli Popen binds to the parent (PR_SET_PDEATHSIG via child_popen_kwargs,
  matching the llama.cpp sites), so a parent crash mid-generation can't
  orphan it holding VRAM/RAM.
- Native begin_load uses the filename-fallback family detector the route
  validated with, so a local .gguf whose family keyword lives only in the
  basename no longer dead-ends 400 on a no-GPU host.
- Generate error handler matches exact sentinel messages instead of the
  "cancelled" substring, fixing a 409 misroute and a raw sd-cli output leak.
- find_sd_cpp_binary honors UNSLOTH_STUDIO_HOME/STUDIO_HOME like the
  installer, so a custom Studio home resolves.
- Drop the redundant _tf32_prev bookkeeping; snapshot/restore_backend_flags
  is now the single owner of the TF32/cudnn restore.

The two client-state messages are now shared constants so the 409-vs-500
contract can't drift. Adds a regression test for each behavioral fix.
---
 studio/backend/core/inference/diffusion.py    | 151 +++++++++---------
 .../core/inference/diffusion_families.py      |  24 +++
 .../backend/core/inference/diffusion_speed.py |  44 +----
 .../backend/core/inference/sd_cpp_backend.py  |  17 +-
 .../backend/core/inference/sd_cpp_engine.py   |  16 +-
 studio/backend/routes/inference.py            |  16 +-
 .../backend/tests/test_diffusion_backend.py   |  32 ++++
 studio/backend/tests/test_diffusion_routes.py |  31 ++++
 studio/backend/tests/test_sd_cpp_backend.py   |  11 ++
 9 files changed, 215 insertions(+), 127 deletions(-)

diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py
index cba3ba2af5..6d728df035 100644
--- a/studio/backend/core/inference/diffusion.py
+++ b/studio/backend/core/inference/diffusion.py
@@ -24,8 +24,10 @@ from loggers import get_logger
 from utils.hardware import clear_gpu_cache
 
 from .diffusion_families import (
+    DIFFUSION_CANCELLED_MSG,
+    DIFFUSION_NOT_LOADED_MSG,
     DiffusionFamily,
-    detect_family,
+    detect_family_for_pick,
     resolve_base_repo,
     resolve_local_gguf_child,
 )
@@ -245,22 +247,6 @@ class DiffusionBackend:
                 base, rfilename, hf_token, cancel_event = self._cancel_event
             )
 
-    @staticmethod
-    def _detect_family_for_pick(
-        repo_id: str, gguf_filename: Optional[str], family_override: Optional[str]
-    ) -> Optional[DiffusionFamily]:
-        """Detect the family from the repo id, falling back to the combined
-        path/filename for a direct local .gguf pick. The frontend splits such a
-        pick into (parent dir, basename), so the family keyword can live only in
-        the filename (e.g. /models/z-image-turbo-Q4_K_M.gguf) while the parent
-        directory carries none; scan it too when the directory alone is
-        undetectable. Only used as a fallback, so remote 'org/name' picks and
-        explicit overrides behave exactly as before."""
-        fam = detect_family(repo_id, family_override)
-        if fam is None and gguf_filename and not family_override:
-            fam = detect_family(f"{repo_id}/{gguf_filename}", family_override)
-        return fam
-
     def validate_load_request(
         self,
         repo_id: str,
@@ -277,7 +263,7 @@ class DiffusionBackend:
             raise ValueError(
                 "gguf_filename is required: this backend loads single-file GGUF checkpoints only."
             )
-        fam = self._detect_family_for_pick(repo_id, gguf_filename, family_override)
+        fam = detect_family_for_pick(repo_id, gguf_filename, family_override)
         if fam is None:
             raise ValueError(
                 f"Could not infer a diffusion family for '{repo_id}'. Pass family_override (z-image)."
@@ -368,7 +354,7 @@ class DiffusionBackend:
             # Resolve the base repo and estimate sizes on this thread (both network
             # calls) so begin_load returns instantly; the bar shows raw bytes until
             # the total lands. This is the only writer of _loading's fields here.
-            fam = self._detect_family_for_pick(
+            fam = detect_family_for_pick(
                 kwargs["repo_id"], kwargs.get("gguf_filename"), kwargs.get("family_override")
             )
             base = _resolve_base_repo(
@@ -641,63 +627,78 @@ class DiffusionBackend:
                     quant_active = transformer_quant_engaged is not None or bool(gguf_filename),
                     logger = logger,
                 )
-                speed_applied = apply_speed_optims(
-                    pipe,
-                    target,
-                    is_gguf = bool(gguf_filename),
-                    family = fam,
-                    speed_mode = effective_speed,
-                    cache_active = cache_engaged is not None,
-                    logger = logger,
-                )
-                if transformer_quant_engaged is not None and not speed_applied.get("compiled"):
-                    # Promotion above could not engage compile (e.g. the family is not
-                    # compile-friendly, or compile_repeated_blocks failed): the quantized
-                    # transformer is now running eager, which is far slower than the GGUF
-                    # path it replaced. Surface it loudly rather than hiding the regression.
-                    logger.warning(
-                        "diffusion.transformer_quant: %s engaged but the transformer is NOT "
-                        "compiled; eager torchao quant is ~30x slower than GGUF here",
-                        transformer_quant_engaged,
+                # apply_speed_optims flips the process-global TF32 / cudnn.benchmark
+                # flags. If a later step here (text-encoder quant, memory plan) then
+                # raises -- e.g. OOM -- those flags would leak flipped and a subsequent
+                # `off` load would no longer be bit-identical. Restore the snapshot
+                # unless we reach the commit (unload restores on the happy path).
+                committed = False
+                try:
+                    speed_applied = apply_speed_optims(
+                        pipe,
+                        target,
+                        is_gguf = bool(gguf_filename),
+                        family = fam,
+                        speed_mode = effective_speed,
+                        cache_active = cache_engaged is not None,
+                        logger = logger,
+                    )
+                    if transformer_quant_engaged is not None and not speed_applied.get("compiled"):
+                        # Promotion above could not engage compile (e.g. the family is not
+                        # compile-friendly, or compile_repeated_blocks failed): the quantized
+                        # transformer is now running eager, which is far slower than the GGUF
+                        # path it replaced. Surface it loudly rather than hiding the regression.
+                        logger.warning(
+                            "diffusion.transformer_quant: %s engaged but the transformer is NOT "
+                            "compiled; eager torchao quant is ~30x slower than GGUF here",
+                            transformer_quant_engaged,
+                        )
+                    # Quantise the dense companion text encoder(s) (opt-in fp8 / nvfp4),
+                    # also before placement so the offload hooks move the smaller weights.
+                    te_quant = quantize_text_encoders(
+                        pipe,
+                        target,
+                        mode = text_encoder_quant,
+                        logger = logger,
                     )
-                # Quantise the dense companion text encoder(s) (opt-in fp8 / nvfp4),
-                # also before placement so the offload hooks move the smaller weights.
-                te_quant = quantize_text_encoders(
-                    pipe,
-                    target,
-                    mode = text_encoder_quant,
-                    logger = logger,
-                )
 
-                # Apply the placement planned above (from MEASURED free device memory vs
-                # the model's estimated resident size). apply_memory_plan returns the
-                # (policy, tiling) ACTUALLY engaged (it may fall back to whole-module
-                # offload, and tiling is a no-op on a pipeline with no tiling control), so
-                # status stays honest. The dense fast path already placed the pipe resident;
-                # for the `none` policy this is an idempotent re-placement.
-                effective_policy, effective_tiling = apply_memory_plan(
-                    pipe, plan, device = device, logger = logger
-                )
+                    # Apply the placement planned above (from MEASURED free device memory vs
+                    # the model's estimated resident size). apply_memory_plan returns the
+                    # (policy, tiling) ACTUALLY engaged (it may fall back to whole-module
+                    # offload, and tiling is a no-op on a pipeline with no tiling control), so
+                    # status stays honest. The dense fast path already placed the pipe resident;
+                    # for the `none` policy this is an idempotent re-placement.
+                    effective_policy, effective_tiling = apply_memory_plan(
+                        pipe, plan, device = device, logger = logger
+                    )
 
-                self._state = _LoadState(
-                    pipe = pipe,
-                    family = fam,
-                    repo_id = repo_id,
-                    base_repo = base,
-                    device = device,
-                    dtype = str(dtype).replace("torch.", ""),
-                    cpu_offload = effective_policy != OFFLOAD_NONE,
-                    offload_policy = effective_policy,
-                    vae_tiling = effective_tiling,
-                    memory_mode = plan.requested_mode,
-                    speed_mode = effective_speed,
-                    speed_optims = tuple(k for k, v in speed_applied.items() if v),
-                    backend_flags_before = backend_flags_before,
-                    text_encoder_quant = te_quant,
-                    transformer_quant = transformer_quant_engaged,
-                    attention_backend = attention_engaged,
-                    transformer_cache = cache_engaged,
-                )
+                    self._state = _LoadState(
+                        pipe = pipe,
+                        family = fam,
+                        repo_id = repo_id,
+                        base_repo = base,
+                        device = device,
+                        dtype = str(dtype).replace("torch.", ""),
+                        cpu_offload = effective_policy != OFFLOAD_NONE,
+                        offload_policy = effective_policy,
+                        vae_tiling = effective_tiling,
+                        memory_mode = plan.requested_mode,
+                        speed_mode = effective_speed,
+                        speed_optims = tuple(k for k, v in speed_applied.items() if v),
+                        backend_flags_before = backend_flags_before,
+                        text_encoder_quant = te_quant,
+                        transformer_quant = transformer_quant_engaged,
+                        attention_backend = attention_engaged,
+                        transformer_cache = cache_engaged,
+                    )
+                    committed = True
+                finally:
+                    if not committed:
+                        # Restore the flags AND free the half-built pipe's VRAM: the
+                        # failed load never commits _state, so nothing else reclaims it
+                        # until the next unload.
+                        restore_backend_flags(backend_flags_before)
+                        clear_gpu_cache()
 
         logger.info(
             "diffusion.loaded: repo=%s base=%s device=%s offload=%s tiling=%s reasons=%s",
@@ -854,7 +855,7 @@ class DiffusionBackend:
             with self._lock:
                 state = self._state
                 if state is None:
-                    raise RuntimeError("No diffusion model is loaded.")
+                    raise RuntimeError(DIFFUSION_NOT_LOADED_MSG)
                 # Register under _lock so unload()/a load can signal THIS generation.
                 # A cancel that arrived before now either nulled _state (we raised
                 # above) or targets an older generation, so nothing is lost.
@@ -923,7 +924,7 @@ class DiffusionBackend:
                 # A cancelled denoise returns early with a partial/garbage image;
                 # don't hand it back to be persisted.
                 if cancel.is_set():
-                    raise RuntimeError("Diffusion generation was cancelled.")
+                    raise RuntimeError(DIFFUSION_CANCELLED_MSG)
                 # Return the PIL images (not yet encoded): the route embeds each
                 # image's recipe and persists it via the gallery.
                 return {"images": list(images), "seed": int(seed), "repo_id": state.repo_id}
diff --git a/studio/backend/core/inference/diffusion_families.py b/studio/backend/core/inference/diffusion_families.py
index 3afbe9b495..3ac238e399 100644
--- a/studio/backend/core/inference/diffusion_families.py
+++ b/studio/backend/core/inference/diffusion_families.py
@@ -20,6 +20,14 @@ from pathlib import Path, PurePosixPath
 from typing import Optional
 
 
+# Runtime->route contract: the RuntimeError messages a backend raises for
+# client-recoverable generate states. The /images/generate route matches these
+# EXACTLY to return 409 (vs a sanitized 500 for real failures), so both engines
+# must raise them verbatim -- keep them named here, not as scattered literals.
+DIFFUSION_NOT_LOADED_MSG = "No diffusion model is loaded."
+DIFFUSION_CANCELLED_MSG = "Diffusion generation was cancelled."
+
+
 @dataclass(frozen = True)
 class DiffusionFamily:
     name: str
@@ -173,6 +181,22 @@ def detect_family(repo_id: str, override: Optional[str] = None) -> Optional[Diff
     return None
 
 
+def detect_family_for_pick(
+    repo_id: str, gguf_filename: Optional[str] = None, override: Optional[str] = None
+) -> Optional[DiffusionFamily]:
+    """``detect_family``, falling back to the combined path/filename for a direct
+    local ``.gguf`` pick. The frontend splits such a pick into (parent dir, basename),
+    so the family keyword can live only in the filename (e.g.
+    ``/models/z-image-turbo-Q4_K_M.gguf``) while the parent directory carries none;
+    scan the combined string too when the directory alone is undetectable. Only a
+    fallback, so remote ``org/name`` picks and explicit overrides behave exactly as
+    ``detect_family``. Shared by both engines so validation and load can't diverge."""
+    fam = detect_family(repo_id, override)
+    if fam is None and gguf_filename and not override:
+        fam = detect_family(f"{repo_id}/{gguf_filename}", override)
+    return fam
+
+
 def resolve_base_repo(fam: DiffusionFamily, base_repo: Optional[str]) -> str:
     """The companion diffusers repo: caller-supplied if given, else the family fallback."""
     base = (base_repo or "").strip()
diff --git a/studio/backend/core/inference/diffusion_speed.py b/studio/backend/core/inference/diffusion_speed.py
index 208cb8b513..5f1ef64445 100644
--- a/studio/backend/core/inference/diffusion_speed.py
+++ b/studio/backend/core/inference/diffusion_speed.py
@@ -161,12 +161,11 @@ def apply_speed_optims(
         "compiled": False,
     }
     mode = normalize_speed_mode(speed_mode)
-    # TF32 is the one PROCESS-GLOBAL flag we flip (on max). Restore it whenever this
-    # load isn't max, so a later default/off diffusion load -- or chat inference in the
-    # same long-lived process -- doesn't silently inherit a prior max load's TF32 and
-    # lose the bit-identical default the regression harness checks.
-    if mode != SPEED_MAX:
-        _restore_tf32(logger)
+    # TF32 and cudnn.benchmark are the process-global flags this may flip (TF32 on max,
+    # cudnn.benchmark on any non-off CUDA load). The caller snapshots them before this
+    # call and restores on unload / failed load via snapshot_backend_flags /
+    # restore_backend_flags, so a later `off` load -- or chat inference in the same
+    # process -- never inherits them. We keep no separate bookkeeping here.
     if mode == SPEED_OFF:
         return applied
 
@@ -251,22 +250,9 @@ def _enable_cudnn_benchmark(logger: Any) -> bool:
         return False
 
 
-# The TF32 flag values from before the first max load flipped them, so a later
-# non-max load / unload can put the process back exactly as it found it (rather than
-# forcing a hardcoded default that might clobber another component's choice).
-_tf32_prev: Optional[tuple[bool, bool]] = None
-
-
 def _enable_tf32(logger: Any) -> bool:
-    global _tf32_prev
     try:
         import torch
-
-        if _tf32_prev is None:
-            _tf32_prev = (
-                torch.backends.cuda.matmul.allow_tf32,
-                torch.backends.cudnn.allow_tf32,
-            )
         torch.backends.cuda.matmul.allow_tf32 = True
         torch.backends.cudnn.allow_tf32 = True
         return True
@@ -275,26 +261,6 @@ def _enable_tf32(logger: Any) -> bool:
         return False
 
 
-def restore_tf32(logger: Any = None) -> None:
-    """Put the process-global TF32 flags back to their pre-max-load values. No-op if
-    a max load never set them. Called on a non-max load and on unload."""
-    _restore_tf32(logger)
-
-
-def _restore_tf32(logger: Any) -> None:
-    global _tf32_prev
-    if _tf32_prev is None:
-        return
-    try:
-        import torch
-        torch.backends.cuda.matmul.allow_tf32 = _tf32_prev[0]
-        torch.backends.cudnn.allow_tf32 = _tf32_prev[1]
-    except Exception as exc:  # noqa: BLE001 — best-effort restore
-        _warn(logger, "tf32_restore", exc)
-    finally:
-        _tf32_prev = None
-
-
 def _fuse_qkv(pipe: Any, logger: Any) -> bool:
     for owner in (pipe, getattr(pipe, "transformer", None)):
         fn = getattr(owner, "fuse_qkv_projections", None)
diff --git a/studio/backend/core/inference/sd_cpp_backend.py b/studio/backend/core/inference/sd_cpp_backend.py
index 4d0bbcbfa7..88661f2be1 100644
--- a/studio/backend/core/inference/sd_cpp_backend.py
+++ b/studio/backend/core/inference/sd_cpp_backend.py
@@ -35,8 +35,10 @@ from typing import Any, Optional
 
 from core.inference.diffusion_device import resolve_diffusion_device_target
 from core.inference.diffusion_families import (
+    DIFFUSION_CANCELLED_MSG,
+    DIFFUSION_NOT_LOADED_MSG,
     DiffusionFamily,
-    detect_family,
+    detect_family_for_pick,
     family_sd_cpp_supported,
     resolve_base_repo,
     resolve_local_gguf_child,
@@ -242,7 +244,10 @@ class SdCppDiffusionBackend:
             raise ValueError(
                 "gguf_filename is required: the native engine loads single-file GGUF checkpoints only."
             )
-        fam = detect_family(repo_id, family_override)
+        # Use the filename-fallback detector the route validated with, so a local
+        # .gguf pick whose family keyword lives only in the basename doesn't pass
+        # validation and then dead-end here on a no-GPU (native-routed) host.
+        fam = detect_family_for_pick(repo_id, gguf_filename, family_override)
         if fam is None:
             raise ValueError(f"Could not infer a diffusion family for '{repo_id}'.")
         if not family_sd_cpp_supported(fam):
@@ -464,7 +469,7 @@ class SdCppDiffusionBackend:
             with self._lock:
                 state = self._state
                 if state is None:
-                    raise RuntimeError("No diffusion model is loaded.")
+                    raise RuntimeError(DIFFUSION_NOT_LOADED_MSG)
                 self._active_generate_cancel = cancel
             engine = self._resolve_engine()
             try:
@@ -485,7 +490,7 @@ class SdCppDiffusionBackend:
                 with tempfile.TemporaryDirectory(prefix = "sdcpp_gen_") as tmpdir:
                     for index in range(max(1, int(batch_size))):
                         if cancel.is_set():
-                            raise RuntimeError("Diffusion generation was cancelled.")
+                            raise RuntimeError(DIFFUSION_CANCELLED_MSG)
                         # Distinct seed per batch image (sd-cli is one image/run here),
                         # so a batch is reproducible image-by-image from the base seed.
                         # Mask to sd-cli's int64 range, NOT 53 bits: the request model and
@@ -522,7 +527,7 @@ class SdCppDiffusionBackend:
                             images.append(im.copy())
                         seeds.append(seed_i)
                 if cancel.is_set():
-                    raise RuntimeError("Diffusion generation was cancelled.")
+                    raise RuntimeError(DIFFUSION_CANCELLED_MSG)
                 # ``seeds`` is the per-image seed (each sd-cli run used seed+index), so
                 # the route can persist the real seed for every image in the batch.
                 return {
@@ -532,7 +537,7 @@ class SdCppDiffusionBackend:
                     "repo_id": state.repo_id,
                 }
             except SdCppCancelled as exc:
-                raise RuntimeError("Diffusion generation was cancelled.") from exc
+                raise RuntimeError(DIFFUSION_CANCELLED_MSG) from exc
             finally:
                 self._gen = None
                 with self._lock:
diff --git a/studio/backend/core/inference/sd_cpp_engine.py b/studio/backend/core/inference/sd_cpp_engine.py
index 6714151045..557f6fcb3e 100644
--- a/studio/backend/core/inference/sd_cpp_engine.py
+++ b/studio/backend/core/inference/sd_cpp_engine.py
@@ -34,6 +34,7 @@ import time
 from pathlib import Path
 from typing import Callable, Optional
 
+from utils.process_lifetime import child_popen_kwargs
 from core.inference.sd_cpp_args import (
     SdCppGenParams,
     SdCppModelFiles,
@@ -125,7 +126,8 @@ def find_sd_cpp_binary() -> Optional[str]:
     both engines look):
       1. ``SD_CLI_PATH`` env -- a direct path to the binary.
       2. ``UNSLOTH_SD_CPP_PATH`` env -- a stable-diffusion.cpp install dir.
-      3. ``~/.unsloth/stable-diffusion.cpp`` build layouts (the installer target).
+      3. the installer target: ``<UNSLOTH_STUDIO_HOME>/../stable-diffusion.cpp`` when
+         that env (or ``STUDIO_HOME``) is set, else ``~/.unsloth/stable-diffusion.cpp``.
       4. ``./stable-diffusion.cpp`` in-tree build (developer checkout).
       5. ``sd-cli`` (then legacy ``sd``) on PATH.
     """
@@ -151,8 +153,12 @@ def find_sd_cpp_binary() -> Optional[str]:
         if hit:
             return hit
 
-    # 3. Default install root (sibling of ~/.unsloth/llama.cpp).
-    hit = _first_file(_layout_candidates(Path.home() / ".unsloth" / "stable-diffusion.cpp"))
+    # 3. Default install root: the installer's default_install_dir() -- a sibling of
+    #    the llama.cpp install under UNSLOTH_STUDIO_HOME / STUDIO_HOME when set, else
+    #    ~/.unsloth. Mirror that env resolution or a custom Studio home never resolves.
+    studio_home = os.environ.get("UNSLOTH_STUDIO_HOME") or os.environ.get("STUDIO_HOME")
+    default_base = Path(studio_home).parent if studio_home else Path.home() / ".unsloth"
+    hit = _first_file(_layout_candidates(default_base / "stable-diffusion.cpp"))
     if hit:
         return hit
 
@@ -342,6 +348,10 @@ class SdCppEngine:
             # Own session/process group so cancellation/timeout can kill the whole
             # tree, not just the parent (POSIX only; harmless flag elsewhere).
             start_new_session = (os.name == "posix"),
+            # Bind the child to the parent's lifetime (Linux PR_SET_PDEATHSIG), so a
+            # hard parent crash mid-generation can't orphan sd-cli holding VRAM/RAM --
+            # matching every llama.cpp Popen site. Composes with start_new_session.
+            **child_popen_kwargs(),
         )
         # Drain stdout on a reader thread so the timeout is enforced even when the
         # child hangs WITHOUT printing (e.g. stuck in model load / GPU init): a plain
diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py
index 400dc239fc..e20a2f1a49 100644
--- a/studio/backend/routes/inference.py
+++ b/studio/backend/routes/inference.py
@@ -10370,6 +10370,10 @@ async def generate_diffusion_image(
 ):
     from core.inference import image_gallery
     from core.inference.diffusion_engine_router import get_active_diffusion_engine
+    from core.inference.diffusion_families import (
+        DIFFUSION_CANCELLED_MSG,
+        DIFFUSION_NOT_LOADED_MSG,
+    )
 
     backend = get_active_diffusion_engine()
     try:
@@ -10385,11 +10389,15 @@ async def generate_diffusion_image(
             batch_size = request.batch_size,
         )
     except RuntimeError as exc:
-        # Only "no model loaded" / cancelled are client-state (409). The native
-        # sd.cpp engine also raises RuntimeError for execution failures (nonzero
-        # exit, timeout, missing output), which are server errors (500).
+        # Only "no model loaded" / user-cancelled are client-state (409); both engines
+        # raise these two EXACT messages. The native sd.cpp engine also raises
+        # RuntimeError for execution failures (nonzero exit, timeout, missing output)
+        # whose text can embed the raw sd-cli tail (local paths / argv) -- those are
+        # server errors (500) returned as a fixed literal, never echoed. Match the
+        # sentinels exactly, not as a substring, so an sd-cli failure that merely
+        # contains "cancelled" can't misroute to 409 and leak that output.
         msg = str(exc)
-        if "No diffusion model is loaded" in msg or "cancelled" in msg.lower():
+        if msg in (DIFFUSION_NOT_LOADED_MSG, DIFFUSION_CANCELLED_MSG):
             raise HTTPException(status_code = 409, detail = msg)
         logger.error("diffusion.generate_failed: %s", exc)
         raise HTTPException(status_code = 500, detail = "Image generation failed.")
diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py
index f8d0172173..5ec8bb3954 100644
--- a/studio/backend/tests/test_diffusion_backend.py
+++ b/studio/backend/tests/test_diffusion_backend.py
@@ -307,6 +307,38 @@ def test_generate_without_load_raises(fake_runtime):
         backend.generate(prompt = "x")
 
 
+def test_failed_load_restores_backend_flags(fake_runtime, tmp_path, monkeypatch):
+    # A failure AFTER apply_speed_optims (here an OOM in apply_memory_plan) must go
+    # through the load's try/finally and restore the process-global TF32 / cudnn flags,
+    # so a later `off` load is still bit-identical, and must not commit a partial state.
+    # Regression: a refactor dropped this guard, leaking the flags on a failed load.
+    (tmp_path / "model.gguf").write_bytes(b"x")
+    backend = DiffusionBackend()
+
+    restored: list = []
+    cleared: list = []
+    monkeypatch.setattr(
+        "core.inference.diffusion.restore_backend_flags", lambda snap: restored.append(snap)
+    )
+    monkeypatch.setattr("core.inference.diffusion.clear_gpu_cache", lambda: cleared.append(True))
+    monkeypatch.setattr(
+        "core.inference.diffusion.apply_memory_plan",
+        lambda *a, **k: (_ for _ in ()).throw(RuntimeError("CUDA out of memory")),
+    )
+
+    with pytest.raises(RuntimeError, match = "out of memory"):
+        backend.load_pipeline(
+            str(tmp_path),
+            gguf_filename = "model.gguf",
+            family_override = "z-image",
+            base_repo = "base/repo",
+            speed_mode = "max",
+        )
+    assert restored, "restore_backend_flags was not called on the failed-load path"
+    assert cleared, "clear_gpu_cache was not called on the failed-load path (VRAM leak)"
+    assert backend._state is None and backend.is_loaded is False
+
+
 def test_resolve_base_repo_prefers_caller_then_hf_tag_then_fallback(monkeypatch):
     from core.inference import diffusion
     from core.inference.diffusion_families import detect_family
diff --git a/studio/backend/tests/test_diffusion_routes.py b/studio/backend/tests/test_diffusion_routes.py
index d0b6997d07..a3c78a21aa 100644
--- a/studio/backend/tests/test_diffusion_routes.py
+++ b/studio/backend/tests/test_diffusion_routes.py
@@ -281,6 +281,37 @@ def test_generate_pipeline_error_returns_sanitized_500(client, monkeypatch):
     assert "CUDA" not in resp.json()["detail"]
 
 
+def test_generate_execution_error_with_cancelled_substring_is_sanitized_500(client, monkeypatch):
+    # A native sd-cli execution failure whose raw tail merely CONTAINS "cancelled"
+    # must stay a sanitized 500, not misroute to 409 and echo that output (path/arg
+    # leak). Regression: the handler matched "cancelled" as a substring.
+    backend = diffusion_module.get_diffusion_backend()
+    backend.loaded = True
+
+    def _fail(**kwargs):
+        raise RuntimeError("sd-cli exited 1. Last output:\nop cancelled at /home/u/models/x.gguf")
+
+    monkeypatch.setattr(backend, "generate", _fail)
+    resp = client.post("/api/inference/images/generate", json = {"prompt": "p"})
+    assert resp.status_code == 500
+    assert resp.json()["detail"] == "Image generation failed."
+    assert "cancelled" not in resp.json()["detail"] and "models" not in resp.json()["detail"]
+
+
+def test_generate_user_cancellation_returns_409(client, monkeypatch):
+    # The exact cancellation sentinel both engines raise is client-state (409).
+    backend = diffusion_module.get_diffusion_backend()
+    backend.loaded = True
+
+    def _cancel(**kwargs):
+        raise RuntimeError("Diffusion generation was cancelled.")
+
+    monkeypatch.setattr(backend, "generate", _cancel)
+    resp = client.post("/api/inference/images/generate", json = {"prompt": "p"})
+    assert resp.status_code == 409
+    assert resp.json()["detail"] == "Diffusion generation was cancelled."
+
+
 def test_load_unknown_family_returns_400(client, monkeypatch):
     def _raise(*a, **k):
         raise ValueError("'x/y' isn't a supported image-generation model. Supported: Z-Image.")
diff --git a/studio/backend/tests/test_sd_cpp_backend.py b/studio/backend/tests/test_sd_cpp_backend.py
index 79464dd6be..84e1b37dd2 100644
--- a/studio/backend/tests/test_sd_cpp_backend.py
+++ b/studio/backend/tests/test_sd_cpp_backend.py
@@ -219,6 +219,17 @@ def test_begin_load_requires_gguf_filename():
         b.begin_load("unsloth/Z-Image-Turbo-GGUF")
 
 
+def test_begin_load_resolves_family_from_filename_only(monkeypatch):
+    # A local .gguf pick whose family keyword lives only in the basename (parent dir
+    # carries none) must resolve via the same filename fallback the route validated
+    # with -- not dead-end with "Could not infer" on a native (no-GPU) host.
+    b = SdCppDiffusionBackend(engine = _FakeEngine())
+    monkeypatch.setattr(b, "_run_load", lambda **kwargs: None)  # skip the download thread
+    b.begin_load("/models/gguf-store", gguf_filename = "Z-Image-Turbo-Q4_K_M.gguf")
+    # Validation passed (no ValueError) and the family was inferred from the filename.
+    assert b._loading is not None and b._loading.repo_id == "/models/gguf-store"
+
+
 def test_ensure_binary_returns_found(monkeypatch):
     monkeypatch.setattr(bk, "find_sd_cpp_binary", lambda: "/usr/bin/sd-cli")
     assert ensure_sd_cpp_binary() == "/usr/bin/sd-cli"

From c34020bfec21557fdba87e77ca039f697a24376c Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Wed, 1 Jul 2026 20:02:34 +0000
Subject: [PATCH 2/4] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 studio/backend/core/inference/diffusion_families.py | 4 +++-
 studio/backend/core/inference/diffusion_speed.py    | 1 +
 2 files changed, 4 insertions(+), 1 deletion(-)

diff --git a/studio/backend/core/inference/diffusion_families.py b/studio/backend/core/inference/diffusion_families.py
index 3ac238e399..7257c69e9e 100644
--- a/studio/backend/core/inference/diffusion_families.py
+++ b/studio/backend/core/inference/diffusion_families.py
@@ -182,7 +182,9 @@ def detect_family(repo_id: str, override: Optional[str] = None) -> Optional[Diff
 
 
 def detect_family_for_pick(
-    repo_id: str, gguf_filename: Optional[str] = None, override: Optional[str] = None
+    repo_id: str,
+    gguf_filename: Optional[str] = None,
+    override: Optional[str] = None,
 ) -> Optional[DiffusionFamily]:
     """``detect_family``, falling back to the combined path/filename for a direct
     local ``.gguf`` pick. The frontend splits such a pick into (parent dir, basename),
diff --git a/studio/backend/core/inference/diffusion_speed.py b/studio/backend/core/inference/diffusion_speed.py
index 5f1ef64445..fba962f3ba 100644
--- a/studio/backend/core/inference/diffusion_speed.py
+++ b/studio/backend/core/inference/diffusion_speed.py
@@ -253,6 +253,7 @@ def _enable_cudnn_benchmark(logger: Any) -> bool:
 def _enable_tf32(logger: Any) -> bool:
     try:
         import torch
+
         torch.backends.cuda.matmul.allow_tf32 = True
         torch.backends.cudnn.allow_tf32 = True
         return True

From 8d16ef977b77d3a5e16654ab07eeb7db58368dd3 Mon Sep 17 00:00:00 2001
From: oobabooga <112222186+oobabooga@users.noreply.github.com>
Date: Wed, 1 Jul 2026 18:01:12 -0300
Subject: [PATCH 3/4] Fix diffusion GGUF memory over-estimate and torch.compile
 crashes

Three chained bugs that made Z-Image (and other GGUF DiTs) crash at generation
on anything but a huge, fully-idle GPU. Verified end to end on an RTX 6000 Ada:
Q2_K now plans resident and generates a real 1024x1024 PNG on both the resident
and forced-group-offload paths.

- Memory planner over-estimated the GGUF transformer's resident size. diffusers
  keeps GGUF weights PACKED (uint8 GGUFParameter) and dequantises per-matmul
  transiently, so resident VRAM is ~= the on-disk size, not the unpacked bf16
  size (measured: Q2_K 3.64->3.68 GiB, Q8_0 7.22->7.25 GiB). The old per-quant
  expansion (x8 for Q2) over-estimated ~7.6x, so a 3.6 GB model on a 48 GB-free
  card was judged a "tight fit" and forced into group offload. Replace the
  multiplier table with estimate_gguf_resident_mib = storage * 1.05 (matches
  diffusers' own get_memory_footprint of a loaded GGUF model).

- torch.compile with fullgraph=True crashed under CPU offload: group/model/
  sequential offload installs a @torch.compiler.disable'd ModuleGroup.onload_
  hook, which graph-breaks. Drop fullgraph when offloading is planned, same as
  the existing step-cache case (fullgraph = not (cache_active or offload_active)).
  This mirrors diffusers' documented compile+offload guidance.

- compile_repeated_blocks compiles one graph per distinct block shape, but
  Z-Image's "repeated" blocks are heterogeneous (~11 variants), above dynamo's
  default recompile_limit of 8, so a resident load hard-errored under fullgraph.
  Raise the limit (diffusers' documented fix for regional-compile recompilation).
  Confirmed force_parameter_static_shapes=False is the wrong lever: same variant
  count, ~6x slower compile.

Also drops the now-dead infer_gguf_quant_label / gguf_filename plumbing and adds
regression tests for the estimate and the offload fullgraph drop.
---
 studio/backend/core/inference/diffusion.py    | 17 +++---
 .../core/inference/diffusion_memory.py        | 53 +++++--------------
 .../backend/core/inference/diffusion_speed.py | 44 ++++++++++++---
 studio/backend/tests/test_diffusion_memory.py | 35 ++++--------
 studio/backend/tests/test_diffusion_speed.py  | 18 +++++++
 5 files changed, 85 insertions(+), 82 deletions(-)

diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py
index 6d728df035..ecf1a1dcf9 100644
--- a/studio/backend/core/inference/diffusion.py
+++ b/studio/backend/core/inference/diffusion.py
@@ -39,10 +39,9 @@ from .diffusion_device import (
 from .diffusion_memory import (
     OFFLOAD_NONE,
     apply_memory_plan,
-    estimate_gguf_dense_mib,
+    estimate_gguf_resident_mib,
     estimate_image_runtime_mib,
     file_size_mib,
-    infer_gguf_quant_label,
     plan_diffusion_memory,
     snapshot_device_memory,
 )
@@ -522,7 +521,7 @@ class DiffusionBackend:
                 # dense bf16 transformer must fit resident, so the fast path is offered only
                 # when the plan is `none`.
                 plan = self._plan_memory(
-                    target, gguf_path, gguf_filename, base, fam, memory_mode, cpu_offload
+                    target, gguf_path, base, fam, memory_mode, cpu_offload
                 )
 
                 # Opt-in fast path: load the DENSE bf16 transformer and torchao-quantise it
@@ -641,6 +640,9 @@ class DiffusionBackend:
                         family = fam,
                         speed_mode = effective_speed,
                         cache_active = cache_engaged is not None,
+                        # The planned offload policy: group/model/sequential offload installs
+                        # compiler-disabled onload hooks, so compile must drop fullgraph.
+                        offload_active = plan.offload_policy != OFFLOAD_NONE,
                         logger = logger,
                     )
                     if transformer_quant_engaged is not None and not speed_applied.get("compiled"):
@@ -797,7 +799,6 @@ class DiffusionBackend:
         self,
         target: DiffusionDeviceTarget,
         gguf_path: str,
-        gguf_filename: Optional[str],
         base: str,
         fam: DiffusionFamily,
         memory_mode: Optional[str],
@@ -808,16 +809,14 @@ class DiffusionBackend:
         offload policy + VAE memory savers. Kept on the backend so the cached base
         repo (companion text-encoder / VAE) feeds the size estimate."""
         device_memory = snapshot_device_memory(target)
-        transformer_dense = estimate_gguf_dense_mib(
-            file_size_mib(gguf_path), infer_gguf_quant_label(gguf_filename)
-        )
+        transformer_resident = estimate_gguf_resident_mib(file_size_mib(gguf_path))
         # The companion components (VAE + text encoders) load near their on-disk
         # size; sum whatever the prefetch already placed in the base-repo cache.
         companion = self._cache_bytes(base)
         companion_mib = int(companion // (1024 * 1024)) if companion else None
         model_dense_mib = None
-        if transformer_dense is not None:
-            model_dense_mib = transformer_dense + (companion_mib or 0)
+        if transformer_resident is not None:
+            model_dense_mib = transformer_resident + (companion_mib or 0)
         runtime_headroom = estimate_image_runtime_mib(width = None, height = None, family = fam.name)
         return plan_diffusion_memory(
             target = target,
diff --git a/studio/backend/core/inference/diffusion_memory.py b/studio/backend/core/inference/diffusion_memory.py
index 9172ef7cf8..b2c1e25f11 100644
--- a/studio/backend/core/inference/diffusion_memory.py
+++ b/studio/backend/core/inference/diffusion_memory.py
@@ -218,51 +218,22 @@ def file_size_mib(path: Any) -> Optional[int]:
         return None
 
 
-def infer_gguf_quant_label(filename: Optional[str]) -> Optional[str]:
-    """Pull a quant tag (Q4_K_M, Q8_0, BF16, ...) out of a GGUF filename."""
-    if not filename:
-        return None
-    from pathlib import Path
+def estimate_gguf_resident_mib(storage_mib: Optional[int]) -> Optional[int]:
+    """Approximate the RESIDENT device size of a GGUF transformer loaded through
+    diffusers' ``GGUFQuantizationConfig``.
 
-    stem = Path(filename).name
-    if stem.lower().endswith(".gguf"):
-        stem = stem[:-5]
-    parts = [p.upper() for p in stem.replace("-", "_").split("_") if p]
-    for index, part in enumerate(parts):
-        if part in ("BF16", "F16", "FP16", "FP8", "Q8", "Q6", "Q5", "Q4", "Q3", "Q2"):
-            suffix = parts[index + 1 :]
-            # Quant names carry either a K-family suffix (Q4_K_M) or a legacy
-            # numeric one (Q8_0, Q5_1); keep up to two suffix tokens.
-            if suffix and suffix[0] in ("K", "M", "S", "L", "XS", "XXS", "0", "1"):
-                return "_".join([part] + suffix[:2])
-            return part
-        if part.startswith("IQ") or part.startswith("UD"):
-            return "_".join(parts[index : index + 3])
-    return None
+    The weights stay PACKED on the device as quantised bytes (``GGUFParameter`` /
+    uint8); ``GGUFLinear.forward`` dequantises each weight to the bf16 compute dtype
+    transiently for its matmul and frees it immediately, so the persistent footprint
+    is ~= the on-disk tensor size, NOT the unpacked bf16 size. Measured on
+    Z-Image-Turbo: Q2_K 3.64 GiB -> 3.68 GiB, Q8_0 7.22 GiB -> 7.25 GiB resident.
+    The transient per-op dequant is covered by the separate runtime headroom.
 
-
-def estimate_gguf_dense_mib(storage_mib: Optional[int], quant: Optional[str]) -> Optional[int]:
-    """Approximate the dequantised (device) size of a GGUF from its on-disk size
-    and quant label. The compute dtype is bf16/fp16, so a 4-bit file roughly
-    quadruples once unpacked; higher-bit quants expand less."""
+    (The prior per-quant expansion assumed a full unpack that never happens on this
+    path; it over-estimated e.g. Q2 ~7.6x, forcing needless offload.)"""
     if storage_mib is None:
         return None
-    q = (quant or "").upper()
-    if any(t in q for t in ("BF16", "F16", "FP16")):
-        return storage_mib
-    if "FP8" in q or "Q8" in q:
-        return int(storage_mib * 2.0)
-    if "Q6" in q:
-        return int(storage_mib * 2.8)
-    if "Q5" in q:
-        return int(storage_mib * 3.3)
-    if "Q4" in q or "IQ4" in q or "UD" in q:
-        return int(storage_mib * 4.0)
-    if "Q3" in q or "IQ3" in q:
-        return int(storage_mib * 5.3)
-    if "Q2" in q or "Q1" in q or "IQ2" in q or "IQ1" in q:
-        return int(storage_mib * 8.0)
-    return int(storage_mib * 4.0)  # unknown: assume 4-bit-ish
+    return int(storage_mib * 1.05)  # small margin for allocator + bf16 norms/biases
 
 
 def estimate_image_runtime_mib(
diff --git a/studio/backend/core/inference/diffusion_speed.py b/studio/backend/core/inference/diffusion_speed.py
index fba962f3ba..7c4a9018d8 100644
--- a/studio/backend/core/inference/diffusion_speed.py
+++ b/studio/backend/core/inference/diffusion_speed.py
@@ -148,11 +148,17 @@ def apply_speed_optims(
     family: Any,
     speed_mode: str = SPEED_OFF,
     cache_active: bool = False,
+    offload_active: bool = False,
     logger: Any = None,
 ) -> dict[str, bool]:
     """Apply the opt-in speed optimisations for ``speed_mode`` to a built pipeline,
     BEFORE placement / offload. Returns which optimisations actually engaged. Every
-    step is best-effort: a pipeline that doesn't support one is simply skipped."""
+    step is best-effort: a pipeline that doesn't support one is simply skipped.
+
+    ``offload_active`` is the planned offload policy != none: group/model/sequential
+    offloading installs ``@torch.compiler.disable``d onload hooks, so the compile must
+    drop ``fullgraph`` (same reason as an active step cache) or it crashes at the first
+    denoise step."""
     applied = {
         "channels_last": False,
         "cudnn_benchmark": False,
@@ -182,7 +188,11 @@ def apply_speed_optims(
     # max-autotune (longer compile, autotuned kernels).
     if compile_eligible(target, is_gguf = is_gguf, family = family):
         applied["compiled"] = _compile_repeated_blocks(
-            pipe, logger, max_autotune = mode == SPEED_MAX, cache_active = cache_active
+            pipe,
+            logger,
+            max_autotune = mode == SPEED_MAX,
+            cache_active = cache_active,
+            offload_active = offload_active,
         )
 
     if mode == SPEED_MAX:
@@ -213,6 +223,7 @@ def _compile_repeated_blocks(
     *,
     max_autotune: bool = False,
     cache_active: bool = False,
+    offload_active: bool = False,
 ) -> bool:
     transformer = getattr(pipe, "transformer", None)
     fn = getattr(transformer, "compile_repeated_blocks", None)
@@ -225,14 +236,33 @@ def _compile_repeated_blocks(
     # / max-autotune) are deliberately NOT used: they crash on the regionally-compiled
     # block because its static output buffer is overwritten across denoise steps.
     #
-    # fullgraph drops to False when a step cache is engaged: FBCache's per-step decision is
-    # ``@torch.compiler.disable``d, i.e. a graph break, which fullgraph=True rejects ("Skip
-    # inlining torch.compiler.disable()d function"). The break is cheap and the rest of the
-    # block still compiles.
-    kwargs: dict[str, Any] = {"fullgraph": not cache_active, "dynamic": not max_autotune}
+    # fullgraph drops to False when a step cache OR CPU offloading is engaged: both insert
+    # an ``@torch.compiler.disable``d function into the forward -- FBCache's per-step
+    # decision, and group/model/sequential offload's ``ModuleGroup.onload_`` streaming hook
+    # -- i.e. a graph break, which fullgraph=True rejects ("Skip inlining
+    # torch.compiler.disable()d function"). The break is cheap and the rest of the block
+    # still compiles.
+    kwargs: dict[str, Any] = {
+        "fullgraph": not (cache_active or offload_active),
+        "dynamic": not max_autotune,
+    }
     if max_autotune:
         kwargs["mode"] = "max-autotune-no-cudagraphs"
     try:
+        import torch
+        # Heterogeneous-block DiTs (e.g. Z-Image) compile ~one graph per distinct block
+        # shape through compile_repeated_blocks; Z-Image needs ~11, above dynamo's default
+        # recompile_limit of 8. Once the limit is hit a resident load hard-errors under
+        # fullgraph (and an offload/cache load silently drops the overflow blocks to eager),
+        # so raise it well past that (64) for headroom on larger heterogeneous DiTs. This is
+        # diffusers' own documented fix for regional-compile recompilation (their guide bumps
+        # cache_size_limit). Deliberately NOT force_parameter_static_shapes=False: it doesn't
+        # cut the variant count here and makes each compile ~6x slower (24s -> 143s cold).
+        dynamo_cfg = getattr(getattr(torch, "_dynamo", None), "config", None)
+        if dynamo_cfg is not None:
+            for _limit_attr in ("recompile_limit", "cache_size_limit"):  # name varies by torch ver
+                if hasattr(dynamo_cfg, _limit_attr):
+                    setattr(dynamo_cfg, _limit_attr, max(getattr(dynamo_cfg, _limit_attr) or 0, 64))
         fn(**kwargs)
         return True
     except Exception as exc:  # noqa: BLE001 — optimisation only
diff --git a/studio/backend/tests/test_diffusion_memory.py b/studio/backend/tests/test_diffusion_memory.py
index bc09ffbd0f..cfdddb051e 100644
--- a/studio/backend/tests/test_diffusion_memory.py
+++ b/studio/backend/tests/test_diffusion_memory.py
@@ -27,9 +27,8 @@ from core.inference.diffusion_memory import (
     DeviceMemory,
     MemoryPlan,
     apply_memory_plan,
-    estimate_gguf_dense_mib,
+    estimate_gguf_resident_mib,
     estimate_image_runtime_mib,
-    infer_gguf_quant_label,
     normalize_memory_mode,
     plan_diffusion_memory,
     snapshot_device_memory,
@@ -70,29 +69,15 @@ def test_normalize_memory_mode_accepts_and_rejects():
 # ── filename / size estimates ─────────────────────────────────────────────────
 
 
-@pytest.mark.parametrize(
-    "filename,expected",
-    [
-        ("z-image-turbo-Q4_K_M.gguf", "Q4_K_M"),
-        ("flux1-dev-Q8_0.gguf", "Q8_0"),
-        ("model-BF16.gguf", "BF16"),
-        ("qwen-image-IQ4_XS.gguf", "IQ4_XS"),
-        ("no-quant-here.gguf", None),
-        (None, None),
-    ],
-)
-def test_infer_gguf_quant_label(filename, expected):
-    assert infer_gguf_quant_label(filename) == expected
-
-
-def test_estimate_gguf_dense_mib_expansion():
-    # 4-bit roughly quadruples once dequantised to bf16; F16 is already dense.
-    assert estimate_gguf_dense_mib(1000, "Q4_K_M") == 4000
-    assert estimate_gguf_dense_mib(1000, "Q8_0") == 2000
-    assert estimate_gguf_dense_mib(1000, "BF16") == 1000
-    assert estimate_gguf_dense_mib(None, "Q4_K_M") is None
-    # Unknown quant falls back to the conservative 4-bit-ish factor.
-    assert estimate_gguf_dense_mib(1000, None) == 4000
+def test_estimate_gguf_resident_mib_matches_packed_size():
+    # GGUF weights stay packed (uint8) on-device; diffusers dequantises per-matmul
+    # transiently, so the resident footprint ~= the on-disk size regardless of quant
+    # level (measured on Z-Image-Turbo: Q2_K 3.64->3.68 GiB, Q8_0 7.22->7.25 GiB). A
+    # small margin covers allocator overhead. The prior per-quant expansion over-
+    # estimated (Q2 ~7.6x) and forced needless offload on a roomy card.
+    assert estimate_gguf_resident_mib(1000) == 1050
+    assert estimate_gguf_resident_mib(7220) == 7581
+    assert estimate_gguf_resident_mib(None) is None
 
 
 def test_estimate_image_runtime_scales_with_pixels_and_family():
diff --git a/studio/backend/tests/test_diffusion_speed.py b/studio/backend/tests/test_diffusion_speed.py
index b0ca10af4c..1f64b31b7a 100644
--- a/studio/backend/tests/test_diffusion_speed.py
+++ b/studio/backend/tests/test_diffusion_speed.py
@@ -216,6 +216,24 @@ def test_speed_default_channels_last_compile_and_cudnn_benchmark(monkeypatch):
     assert applied["tf32"] is False and applied["fused_qkv"] is False
 
 
+def test_offload_active_drops_fullgraph(monkeypatch):
+    # Group/model/sequential offload installs a torch.compiler.disable'd onload hook;
+    # compiling with fullgraph=True then crashes at the first denoise step. Same reason
+    # as an active step cache -> fullgraph must drop to False when offload is planned.
+    _stub_torch(monkeypatch)
+    pipe = _Pipe(with_compile = True)
+    applied = apply_speed_optims(
+        pipe,
+        _target(),
+        is_gguf = True,
+        family = _family(),
+        speed_mode = SPEED_DEFAULT,
+        offload_active = True,
+    )
+    assert applied["compiled"] is True
+    assert pipe.compile_kwargs["fullgraph"] is False
+
+
 def test_speed_default_compiles_gguf(monkeypatch):
     _stub_torch(monkeypatch)
     pipe = _Pipe(with_compile = True)

From a34d48d1fd498d73180797f399f0c3c678f9e903 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Wed, 1 Jul 2026 21:05:52 +0000
Subject: [PATCH 4/4] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 studio/backend/core/inference/diffusion.py       | 4 +---
 studio/backend/core/inference/diffusion_speed.py | 1 +
 2 files changed, 2 insertions(+), 3 deletions(-)

diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py
index ecf1a1dcf9..894e3ce692 100644
--- a/studio/backend/core/inference/diffusion.py
+++ b/studio/backend/core/inference/diffusion.py
@@ -520,9 +520,7 @@ class DiffusionBackend:
                 # the real budget) -- this also doubles as the dense-quant preflight: the
                 # dense bf16 transformer must fit resident, so the fast path is offered only
                 # when the plan is `none`.
-                plan = self._plan_memory(
-                    target, gguf_path, base, fam, memory_mode, cpu_offload
-                )
+                plan = self._plan_memory(target, gguf_path, base, fam, memory_mode, cpu_offload)
 
                 # Opt-in fast path: load the DENSE bf16 transformer and torchao-quantise it
                 # (int8 / fp8 / fp4 tensor cores), which beats GGUF's bf16-rate per-matmul
diff --git a/studio/backend/core/inference/diffusion_speed.py b/studio/backend/core/inference/diffusion_speed.py
index 7c4a9018d8..99f506450a 100644
--- a/studio/backend/core/inference/diffusion_speed.py
+++ b/studio/backend/core/inference/diffusion_speed.py
@@ -250,6 +250,7 @@ def _compile_repeated_blocks(
         kwargs["mode"] = "max-autotune-no-cudagraphs"
     try:
         import torch
+
         # Heterogeneous-block DiTs (e.g. Z-Image) compile ~one graph per distinct block
         # shape through compile_repeated_blocks; Z-Image needs ~11, above dynamo's default
         # recompile_limit of 8. Once the limit is hit a resident load hard-errors under