Do not widen the dense prefetch under a forced-offload memory policy

_dense_quant_prefetch_needed widened the prefetch to pull the base repo's transformer/
shards whenever a dense-quant candidate resolved, but balanced/low_vram (and the legacy
cpu_offload flag) force load_pipeline onto offload unconditionally in plan_diffusion_memory,
so its re-plan never flips to OFFLOAD_NONE and the dense build never runs. The offloaded GGUF
path then never touches those shards, so the widened prefetch only wastes a multi-GB download,
and a disk-full on that begin_load pull has no GGUF fallback (unlike the in-load_pipeline dense
failure). Mirror plan_diffusion_memory's definite-offload gates so the prefetch stays scoped.

studio/backend/core/inference/diffusion.py

@@ -44,12 +44,15 @@ from .diffusion_device import (
 )
 from .diffusion_krea2 import KREA2_FAMILY_NAME, load_krea2_pipeline
 from .diffusion_memory import (
+    MEMORY_MODE_BALANCED,
+    MEMORY_MODE_LOW_VRAM,
     OFFLOAD_NONE,
     apply_memory_plan,
     estimate_gguf_resident_mib,
     estimate_image_runtime_mib,
     estimate_safetensors_dense_mib,
     file_size_mib,
+    normalize_memory_mode,
     plan_diffusion_memory,
     snapshot_device_memory,
 )
@@ -441,6 +444,19 @@ class DiffusionBackend:
         if speed is not None and str(speed).strip().lower() == SPEED_OFF:
             return False
         try:
+            # A definite-offload memory policy forces load_pipeline onto offload regardless of the
+            # dense candidate's smaller footprint, so its re-plan never flips to OFFLOAD_NONE and
+            # the dense build never runs. balanced -> OFFLOAD_GROUP and low_vram -> OFFLOAD_MODEL are
+            # set unconditionally in plan_diffusion_memory; the legacy cpu_offload flag forces
+            # OFFLOAD_MODEL when no memory_mode overrides it. In those cases the GGUF path runs
+            # offloaded and never touches the base transformer/ shards, so widening the prefetch only
+            # wastes a multi-GB download -- and a disk-full on that begin_load pull has NO GGUF
+            # fallback (unlike the in-load_pipeline dense failure). Mirror those offload gates here.
+            mm = normalize_memory_mode(kwargs.get("memory_mode"))
+            if mm in (MEMORY_MODE_BALANCED, MEMORY_MODE_LOW_VRAM):
+                return False
+            if mm is None and kwargs.get("cpu_offload"):
+                return False
             target = self._resolve_device_target(fam)
             # Only widen the prefetch when the loader would actually take the dense path: resolve
             # the SAME dense-quant candidate load_pipeline re-plans against, which also checks the

studio/backend/tests/test_diffusion_backend.py

@@ -2127,6 +2127,39 @@ def test_dense_quant_prefetch_needed_gates(fake_runtime, monkeypatch):
     # UNSET defaults to the hardware ladder (Dtype default-auto) -> widens, threading auto.
     assert backend._dense_quant_prefetch_needed(fam, {}) is True
     assert seen[-1] == "auto"
+    # A definite-offload memory policy forces load_pipeline onto offload regardless of the dense
+    # candidate's smaller footprint, so the dense build never runs and the widened prefetch would
+    # download base transformer/ shards the offloaded GGUF path never uses (and a disk-full there
+    # has no GGUF fallback). balanced / low_vram (and the legacy cpu_offload flag when no
+    # memory_mode overrides it) must NOT widen, even though the candidate itself is dense-viable.
+    before = len(seen)
+    assert (
+        backend._dense_quant_prefetch_needed(fam, {"transformer_quant": "fp8", "memory_mode": "balanced"})
+        is False
+    )
+    assert (
+        backend._dense_quant_prefetch_needed(fam, {"transformer_quant": "fp8", "memory_mode": "low_vram"})
+        is False
+    )
+    assert (
+        backend._dense_quant_prefetch_needed(fam, {"transformer_quant": "fp8", "cpu_offload": True})
+        is False
+    )
+    # The gate short-circuits BEFORE resolving the candidate (no wasted resolve).
+    assert len(seen) == before
+    # An explicit memory_mode still consulting the candidate: fast/auto can flip resident, so they
+    # widen when the candidate is dense-viable (memory_mode="fast" does not force offload).
+    assert (
+        backend._dense_quant_prefetch_needed(fam, {"transformer_quant": "fp8", "memory_mode": "fast"})
+        is True
+    )
+    # A cpu_offload flag is overridden by an explicit resident memory_mode, so it still widens.
+    assert (
+        backend._dense_quant_prefetch_needed(
+            fam, {"transformer_quant": "fp8", "memory_mode": "fast", "cpu_offload": True}
+        )
+        is True
+    )
     # An explicit off pins running the GGUF as-is -> never widen (mode resolves to None first).
     assert backend._dense_quant_prefetch_needed(fam, {"transformer_quant": "none"}) is False
     # An explicit Speed="off" (bit-exact) load suppresses the dense path -> never widen.