Merge branch 'diffusion-train-perf2' into diffusion-auto-policy

studio/backend/core/inference/diffusion.py

@@ -457,11 +457,17 @@ class DiffusionBackend:
         base: str,
         base_files: list[str],
         hf_token: Optional[str],
-    ) -> None:
+    ) -> Optional[str]:
         """Pre-download the GGUF + the given ``base_files`` into the HF cache,
         WITHOUT the lock and honoring ``_cancel_event``, so load_pipeline's
         from_single_file / from_pretrained hit the cache and the heavy download can
-        be preempted by an unload/eviction. Raises ``RuntimeError("Cancelled")``."""
+        be preempted by an unload/eviction. Raises ``RuntimeError("Cancelled")``.
+
+        Returns the base repo's local snapshot dir when the prefetched set includes
+        the pipeline manifest, so from_pretrained can load from disk instead of
+        re-sweeping the hub (its own sweep also pulls files the scoped list skips,
+        e.g. the 24 GB packaged root singles in each FLUX.1 repo); None otherwise
+        (estimate failure, config-only base, local repo) -> hub id as before."""
         from utils.hf_xet_fallback import hf_hub_download_with_xet_fallback
 
         # GGUF transformer (hub repos only; a local path is already on disk).
@@ -470,12 +476,16 @@ class DiffusionBackend:
                 repo_id, gguf_filename, hf_token, cancel_event = self._cancel_event
             )
         # Base repo (VAE / text-encoder / scheduler); list comes from the estimate.
+        snapshot_root: Optional[str] = None
         for rfilename in base_files:
             if self._cancel_event.is_set():
                 raise RuntimeError("Cancelled")
-            hf_hub_download_with_xet_fallback(
+            local = hf_hub_download_with_xet_fallback(
                 base, rfilename, hf_token, cancel_event = self._cancel_event
             )
+            if rfilename == "model_index.json":
+                snapshot_root = str(Path(local).parent)
+        return snapshot_root
 
     def validate_load_request(
         self,
@@ -692,7 +702,7 @@ class DiffusionBackend:
                     self._loading.expected_bytes = expected
             # Download outside the lock so unload()/an eviction can preempt the
             # multi-GB pull; load_pipeline below then assembles from the cache.
-            self._prefetch_files(
+            kwargs["_base_local_dir"] = self._prefetch_files(
                 kwargs["repo_id"],
                 kwargs.get("gguf_filename"),
                 base,
@@ -905,6 +915,7 @@ class DiffusionBackend:
         transformer_cache_threshold: Optional[float] = None,
         model_kind: Optional[str] = None,
         _load_token: Optional[int] = None,
+        _base_local_dir: Optional[str] = None,
     ) -> dict[str, Any]:
         # A blank / whitespace-only token must degrade to anonymous access, not be passed
         # as an explicit credential (from_single_file / from_pretrained / the Hub client
@@ -1069,6 +1080,7 @@ class DiffusionBackend:
                             transformer_quant,
                             transformer_quant_fast_accum,
                             fam = fam,
+                            base_local_dir = _base_local_dir,
                             prequant_path = transformer_prequant_path,
                         )
                     except Exception as exc:  # noqa: BLE001 — fall back to the GGUF build
@@ -1109,7 +1121,12 @@ class DiffusionBackend:
                             pipe_kwargs: dict[str, Any] = {"torch_dtype": dtype}
                             if hf_token:
                                 pipe_kwargs["token"] = hf_token
-                            pipe = pipeline_cls.from_pretrained(repo_id, **pipe_kwargs)
+                            # The prefetched snapshot dir keeps from_pretrained off the
+                            # hub: its own snapshot sweep re-downloads files the scoped
+                            # prefetch skipped (packaged root singles, 24 GB per FLUX.1).
+                            pipe = pipeline_cls.from_pretrained(
+                                _base_local_dir or repo_id, **pipe_kwargs
+                            )
                     elif kind == "single_file" and fam.single_file_is_pipeline:
                         # A single-file SDXL-style checkpoint is the WHOLE pipeline
                         # (U-Net + VAE + both text encoders), not a transformer-only file,
@@ -1150,7 +1167,9 @@ class DiffusionBackend:
                             pipe_kwargs = {"torch_dtype": dtype, "transformer": transformer}
                             if hf_token:
                                 pipe_kwargs["token"] = hf_token
-                            pipe = pipeline_cls.from_pretrained(base, **pipe_kwargs)
+                            pipe = pipeline_cls.from_pretrained(
+                                _base_local_dir or base, **pipe_kwargs
+                            )
 
                 # Resolve the effective speed mode: GGUF models default to the
                 # near-lossless `default` profile (compile is ~2.2x and sits below
@@ -1434,6 +1453,7 @@ class DiffusionBackend:
         *,
         fam: Optional[DiffusionFamily] = None,
         prequant_path: Optional[str] = None,
+        base_local_dir: Optional[str] = None,
     ) -> tuple[Any, str]:
         """Build the opt-in fast pipeline and return ``(pipe, engaged_scheme)``.
 
@@ -1481,7 +1501,7 @@ class DiffusionBackend:
                 )
                 if transformer is not None:
                     pipe = self._assemble_pipe(
-                        pipeline_cls, base, transformer, dtype, hf_token, device
+                        pipeline_cls, base, transformer, dtype, hf_token, device, base_local_dir
                     )
                     return pipe, scheme
 
@@ -1489,7 +1509,9 @@ class DiffusionBackend:
         transformer = transformer_cls.from_pretrained(
             base, subfolder = "transformer", torch_dtype = dtype, token = hf_token
         )
-        pipe = self._assemble_pipe(pipeline_cls, base, transformer, dtype, hf_token, device)
+        pipe = self._assemble_pipe(
+            pipeline_cls, base, transformer, dtype, hf_token, device, base_local_dir
+        )
         scheme = quantize_transformer(
             pipe,
             target,
@@ -1510,13 +1532,14 @@ class DiffusionBackend:
         dtype: Any,
         hf_token: Optional[str],
         device: str,
+        base_local_dir: Optional[str] = None,
     ) -> Any:
         """Assemble the diffusers pipeline around ``transformer`` and place it on ``device``
         (a no-op for an already-placed pre-quantized transformer; it moves the companions)."""
         pipe_kwargs: dict[str, Any] = {"torch_dtype": dtype, "transformer": transformer}
         if hf_token:
             pipe_kwargs["token"] = hf_token
-        pipe = pipeline_cls.from_pretrained(base, **pipe_kwargs)
+        pipe = pipeline_cls.from_pretrained(base_local_dir or base, **pipe_kwargs)
         pipe.to(device)
         return pipe
 

studio/backend/tests/test_diffusion_backend.py

@@ -2156,3 +2156,35 @@ def test_generate_resets_step_cache_only_when_engaged(fake_runtime, tmp_path):
     backend.generate(prompt = "a sloth")
     backend.generate(prompt = "another sloth")
     assert resets == [True, True]
+
+
+def test_prefetch_returns_snapshot_dir_for_manifest(monkeypatch):
+    # The prefetched pipeline manifest's directory is the local snapshot root; a
+    # config-only base list (no manifest) returns None so the hub id stays in use.
+    backend = DiffusionBackend()
+    monkeypatch.setattr(
+        "utils.hf_xet_fallback.hf_hub_download_with_xet_fallback",
+        lambda repo, fn, tok, **k: f"/cache/snap/{fn}",
+    )
+    root = backend._prefetch_files(
+        "base/repo", None, "base/repo", ["model_index.json", "vae/x.safetensors"], None
+    )
+    assert root == "/cache/snap"
+    assert (
+        backend._prefetch_files("base/repo", None, "base/repo", ["vae/x.safetensors"], None)
+        is None
+    )
+
+
+def test_pipeline_load_uses_predownloaded_dir(fake_runtime, tmp_path):
+    # With a prefetched snapshot, from_pretrained must receive the local dir --
+    # its own hub sweep would re-download the root packaged singles the scoped
+    # prefetch skips (24 GB per FLUX.1 repo).
+    backend = DiffusionBackend()
+    backend.load_pipeline(
+        "unsloth/Qwen-Image-2512-bnb-4bit",
+        model_kind = "pipeline",
+        _base_local_dir = str(tmp_path),
+    )
+    assert _FakePipeline.last["base"] == str(tmp_path)
+    backend.unload()