diffusion: don't double-count the prefetched transformer in the dense re-plan

The dense-quant re-plan passes transformer_resident_override_mib (the bf16 build
peak) AND computes companions via _companion_cache_bytes(base), which sums every
flat blob in the HF cache. Because the dense path prefetches the base transformer/
shards into that same cache before load_pipeline runs, the transformer is counted
twice, inflating the footprint (~44 GB instead of ~20 GB in the reproduction) and
wrongly forcing offload for models that fit resident -- the case this path exists
to enable. Add companion_override_mib and pass the auto-policy's own text-encoder
plus VAE estimate on the re-plan so the cache (with its prefetched transformer) is
not read for this artifact.

studio/backend/core/inference/diffusion.py

@@ -1059,6 +1059,10 @@ class DiffusionBackend:
                             kind = kind,
                             repo_id = repo_id,
                             transformer_resident_override_mib = (candidate.transient_transformer_mib),
+                            # The dense path prefetches the base transformer/ shards into the
+                            # cache _companion_cache_bytes reads; pass the auto-policy's own
+                            # companion estimate so the re-plan does not double-count them.
+                            companion_override_mib = candidate.companions_mib,
                         )
                         if replanned.offload_policy == OFFLOAD_NONE:
                             quant_plan = replanned
@@ -1555,6 +1559,7 @@ class DiffusionBackend:
         kind: str = "gguf",
         repo_id: Optional[str] = None,
         transformer_resident_override_mib: Optional[int] = None,
+        companion_override_mib: Optional[int] = None,
     ):
         """Build the memory plan for this load: snapshot free device memory and
         estimate the model's resident footprint, then let the planner pick an
@@ -1567,7 +1572,11 @@ class DiffusionBackend:
         pipeline is one cached download (transformer + companions), already compressed.
         ``transformer_resident_override_mib`` replaces the file-size transformer estimate
         when the loader is planning for a DIFFERENT artifact than the file on disk (the
-        dense transformer-quant candidate, whose footprint the auto-policy estimates)."""
+        dense transformer-quant candidate, whose footprint the auto-policy estimates);
+        ``companion_override_mib`` likewise replaces the cached companion total on that
+        re-plan, so the base repo's PREFETCHED transformer/ shards -- which land in the
+        same blob cache _companion_cache_bytes sums -- are not counted as companions on
+        top of transformer_resident_override_mib (a double-count of the transformer)."""
         device_memory = snapshot_device_memory(target)
         if kind == "pipeline":
             # The whole repo (transformer + companions) is one cached download; the
@@ -1600,8 +1609,17 @@ class DiffusionBackend:
             # LOCAL diffusers base -- the on-disk component weights (the blob cache is
             # empty for a local path, which would otherwise fold multi-GB companions to 0
             # and let auto planning pick a resident placement that OOMs).
-            companion = self._companion_cache_bytes(base)
-            companion_mib = int(companion // (1024 * 1024)) if companion else None
+            if companion_override_mib is not None:
+                # Re-planning the dense transformer-quant candidate: the dense path
+                # prefetches the base repo's transformer/ shards into the SAME blob cache
+                # _companion_cache_bytes sums, so reading it here would count the
+                # transformer AGAIN on top of transformer_resident_override_mib and make
+                # the resident quant plan look far too large. Use the auto-policy's own
+                # companion (text-encoder + VAE) estimate for this artifact instead.
+                companion_mib = companion_override_mib
+            else:
+                companion = self._companion_cache_bytes(base)
+                companion_mib = int(companion // (1024 * 1024)) if companion else None
             model_dense_mib = None
             if transformer_resident is not None:
                 model_dense_mib = transformer_resident + (companion_mib or 0)

studio/backend/tests/test_diffusion_backend.py

@@ -2120,6 +2120,53 @@ def test_companion_cache_bytes_local_dir_excludes_transformer(tmp_path):
     assert total == 150  # vae + text_encoder only; transformer/ and json excluded
 
 
+def test_plan_memory_dense_replan_does_not_double_count_prefetched_transformer(monkeypatch):
+    # Re-planning the dense transformer-quant candidate: the dense path prefetches the
+    # base repo's transformer/ shards into the SAME blob cache _companion_cache_bytes
+    # sums. If the re-plan read that cache it would count the transformer TWICE (once as
+    # transformer_resident_override_mib, once as a "companion") and force offload even
+    # when the quantised artifact fits resident. The re-plan must use the auto-policy's
+    # companion estimate instead. Here the cache is stubbed to the transformer-inflated
+    # value; the plan must still stay resident.
+    from core.inference import diffusion as dmod
+    from core.inference.diffusion_memory import OFFLOAD_NONE, DeviceMemory
+
+    backend = DiffusionBackend()
+    target = types.SimpleNamespace(
+        device = "cuda", backend = "cuda", supports_model_cpu_offload = True
+    )
+    # 40 GiB discrete card, 40000 MiB free: comfortably fits transformer + real
+    # companions + headroom, but NOT a second copy of the bf16 transformer.
+    monkeypatch.setattr(
+        dmod,
+        "snapshot_device_memory",
+        lambda t: DeviceMemory("cuda", "cuda", "discrete_vram", 40000, 40960),
+    )
+    monkeypatch.setattr(dmod, "estimate_image_runtime_mib", lambda **kw: 4000)
+    # The cache is inflated by the prefetched bf16 transformer (~24000) on top of the
+    # ~8000 real companions; if the re-plan consulted it the plan would offload.
+    monkeypatch.setattr(
+        DiffusionBackend,
+        "_companion_cache_bytes",
+        staticmethod(lambda base: (8000 + 24000) * 1024 * 1024),
+    )
+    fam = types.SimpleNamespace(name = "z-image")
+    plan = backend._plan_memory(
+        target,
+        None,
+        "org/base",
+        fam,
+        None,
+        False,
+        kind = "gguf",
+        transformer_resident_override_mib = 12000,  # int8 candidate transient (~half bf16)
+        companion_override_mib = 8000,  # auto-policy text-encoder + VAE estimate
+    )
+    # 12000 + 8000 + 4000 + 2048 overhead = 26048 MiB, fits the ~36 GiB budget.
+    # A double-count (12000 + [8000+24000] + ...) would have exceeded it and offloaded.
+    assert plan.offload_policy == OFFLOAD_NONE
+
+
 def test_reset_step_cache_helper_is_best_effort():
     # Calls the transformer's reset hook when present.
     calls = []