Add the diffusion download plan endpoint

Reports the repos and exact files a pick needs so the download manager can stage
them with the loader's own file scope. A plain snapshot would add the packaged
root single, transformer shards and fp16 twins the loader never opens.

studio/backend/core/inference/diffusion.py

@@ -996,10 +996,14 @@ class DiffusionBackend:
         kind: str = "gguf",
         single_file_is_pipeline: bool = False,
         include_transformer: bool = False,
+        sizes_out: Optional[dict[str, int]] = None,
     ) -> tuple[int, list[str]]:
         """Total download size for the progress bar, plus the base-repo files to
         fetch (the prefetch reuses this list, so the base is listed only once).
 
+        ``sizes_out``, when given, is filled with per-repo byte totals so the download
+        plan can size one job per repo off this same single pair of Hub lookups.
+
         For a ``pipeline`` load the whole repo IS the pipeline (``base_repo`` is the
         repo itself), so the transformer/ subfolder is INCLUDED -- unlike the GGUF /
         single-file paths, where the transformer is the single file and the base repo
@@ -1026,12 +1030,17 @@ class DiffusionBackend:
                         continue
                     base_files.append(s.rfilename)
                     total += s.size or 0
+                if sizes_out is not None:
+                    sizes_out[repo_id] = total
                 return total, base_files
             # Skip the Hub size lookup for a LOCAL gguf path: model_info would raise on a
             # filesystem path and (caught below) skip the base lookup, forcing a synchronous companion pull.
             if gguf_filename and not Path(repo_id).expanduser().exists():
                 info = api.model_info(repo_id, files_metadata = True, token = hf_token)
-                total += sum(s.size or 0 for s in info.siblings if s.rfilename == gguf_filename)
+                gguf_bytes = sum(s.size or 0 for s in info.siblings if s.rfilename == gguf_filename)
+                total += gguf_bytes
+                if sizes_out is not None:
+                    sizes_out[repo_id] = gguf_bytes
             # A whole-pipeline single file (SDXL) needs only the base's config/tokenizer, not its weights.
             if kind == "single_file" and single_file_is_pipeline:
                 base_filter = _base_config_file_downloaded
@@ -1041,14 +1050,75 @@ class DiffusionBackend:
                     return _base_file_downloaded(rfilename, include_transformer = include_transformer)
 
             base_info = api.model_info(base_repo, files_metadata = True, token = hf_token)
+            base_bytes = 0
             for s in base_info.siblings:
                 if base_filter(s.rfilename):
                     base_files.append(s.rfilename)
-                    total += s.size or 0
+                    base_bytes += s.size or 0
+            total += base_bytes
+            if sizes_out is not None:
+                sizes_out[base_repo] = base_bytes
         except Exception as exc:  # noqa: BLE001 — estimate is best-effort
             logger.warning("diffusion.size_estimate_failed: %s", exc)
         return total, base_files
 
+    def download_plan(
+        self,
+        repo_id: str,
+        *,
+        gguf_filename: Optional[str] = None,
+        base_repo: Optional[str] = None,
+        family_override: Optional[str] = None,
+        model_kind: Optional[str] = None,
+        hf_token: Optional[str] = None,
+        **load_kwargs: Any,
+    ) -> dict[str, Any]:
+        """The repos + exact files this pick needs, so the Hub download manager can fetch
+        them with the same file scope the loader would.
+
+        A plain snapshot_download would also pull what the loader deliberately skips (the
+        packaged root single, transformer/ shards, fp16 twins) -- tens of GB per FLUX repo.
+        Resolves family/kind/base exactly as ``_run_load`` does, so the plan and the load
+        agree. Local paths are already on disk and yield no entries."""
+        fam = detect_family_for_pick(repo_id, gguf_filename, family_override)
+        kind = resolve_model_kind(gguf_filename, model_kind)
+        if kind == "pipeline":
+            base = repo_id  # the full pipeline IS the repo
+        else:
+            base = _resolve_base_repo(repo_id, base_repo, fam, hf_token)
+        sizes: dict[str, int] = {}
+        total, base_files = self._estimate_download_bytes(
+            repo_id,
+            gguf_filename,
+            base,
+            hf_token,
+            kind = kind,
+            single_file_is_pipeline = bool(fam and fam.single_file_is_pipeline),
+            include_transformer = kind == "gguf"
+            and self._dense_quant_prefetch_needed(fam, load_kwargs),
+            sizes_out = sizes,
+        )
+        entries: list[dict[str, Any]] = []
+        if gguf_filename and not Path(repo_id).expanduser().exists():
+            entries.append(
+                {
+                    "repo_id": repo_id,
+                    "files": [gguf_filename],
+                    "bytes": int(sizes.get(repo_id, 0)),
+                    "gguf_filename": gguf_filename,
+                }
+            )
+        if base_files and not Path(base).expanduser().exists():
+            entries.append(
+                {
+                    "repo_id": base,
+                    "files": base_files,
+                    "bytes": int(sizes.get(base, 0)),
+                    "gguf_filename": None,
+                }
+            )
+        return {"entries": entries, "total_bytes": int(total)}
+
     @staticmethod
     def _hub_cache_repo_dir(repo_id: str) -> Path:
         """Local HF hub cache dir for ``repo_id``.

studio/backend/models/inference.py

@@ -2632,6 +2632,30 @@ class DiffusionResolvedControl(BaseModel):
     reason: str = Field("", description = "Short human-readable reason for the resolved value.")
 
 
+class DiffusionDownloadPlanEntry(BaseModel):
+    """One repo the pick needs, with the exact files to fetch from it."""
+
+    repo_id: str = Field(..., description = "Repo to download from")
+    files: List[str] = Field(
+        default_factory = list,
+        description = "Exact files the loader reads. Scoped on purpose: a full snapshot would "
+        "also pull the packaged root single, transformer/ shards and fp16 twins the loader "
+        "never opens (tens of GB on a FLUX repo).",
+    )
+    bytes: int = Field(0, description = "Declared size of those files, 0 when unknown")
+    gguf_filename: Optional[str] = Field(
+        None, description = "Set when this entry is the single-file GGUF checkpoint"
+    )
+
+
+class DiffusionDownloadPlanResponse(BaseModel):
+    """What to download before a load, so the Hub download manager can fetch it with the
+    same file scope the loader would. Empty entries mean nothing to download (local path)."""
+
+    entries: List[DiffusionDownloadPlanEntry] = Field(default_factory = list)
+    total_bytes: int = Field(0, description = "Sum across entries, 0 when the estimate failed")
+
+
 class DiffusionStatusResponse(BaseModel):
     """Current diffusion backend state."""
 

studio/backend/routes/inference.py

@@ -1702,6 +1702,7 @@ from models.inference import (
     DiffusionGenerateResponse,
     DiffusionGenerateProgressResponse,
     DiffusionStatusResponse,
+    DiffusionDownloadPlanResponse,
     DiffusionInferenceInfoResponse,
     DiffusionLoadProgressResponse,
     GalleryImage,
@@ -15956,6 +15957,56 @@ def _guard_diffusion_load_against_training() -> None:
     )
 
 
+@studio_router.post("/images/download-plan", response_model = DiffusionDownloadPlanResponse)
+async def diffusion_download_plan(
+    request: DiffusionLoadRequest, current_subject: str = Depends(get_current_subject)
+):
+    """The repos + files this pick needs, so the frontend can stage them through the Hub
+    download manager (one mechanism, one panel) instead of the loader downloading inline.
+
+    Validates the same way /images/load does, so an unloadable pick fails here rather than
+    after a multi-GB download."""
+    from core.inference.diffusion import (
+        get_diffusion_backend,
+        resolve_local_single_file,
+        resolve_model_kind,
+    )
+    from utils.native_path_leases import redact_native_paths
+
+    backend = get_diffusion_backend()
+    try:
+        kind = resolve_model_kind(request.gguf_filename, request.model_kind)
+        # Same bare-single-file-directory reinterpretation as the load route, so the plan
+        # describes the load that will actually run.
+        if kind == "pipeline" and not request.gguf_filename:
+            sole = await asyncio.to_thread(resolve_local_single_file, request.model_path)
+            if sole is not None:
+                request.gguf_filename = sole
+                kind = resolve_model_kind(sole)
+        await asyncio.to_thread(
+            backend.validate_load_request,
+            request.model_path,
+            gguf_filename = request.gguf_filename,
+            family_override = request.family_override,
+            model_kind = kind,
+            base_repo = request.base_repo,
+        )
+        plan = await asyncio.to_thread(
+            backend.download_plan,
+            request.model_path,
+            gguf_filename = request.gguf_filename,
+            base_repo = request.base_repo,
+            family_override = request.family_override,
+            model_kind = kind,
+            hf_token = request.hf_token,
+            transformer_quant = request.transformer_quant,
+            speed_mode = request.speed_mode,
+        )
+        return DiffusionDownloadPlanResponse(**plan)
+    except (ValueError, FileNotFoundError) as exc:
+        raise HTTPException(status_code = 400, detail = redact_native_paths(str(exc)))
+
+
 @studio_router.post("/images/load", response_model = DiffusionStatusResponse)
 async def load_diffusion_model(
     request: DiffusionLoadRequest, current_subject: str = Depends(get_current_subject)

studio/backend/tests/test_diffusion_backend.py

@@ -3806,3 +3806,104 @@ def test_generate_resets_the_step_cache_before_every_chunk(fake_runtime, tmp_pat
     object.__setattr__(backend._state, "transformer_cache", "fbcache")
     backend.generate(prompt = "p", seeds = [1, 2, 3], batch_size = 2)
     assert trace == [("reset",), ("call", 2), ("reset",), ("call", 1)]
+class _FakeSibling:
+    def __init__(self, rfilename, size):
+        self.rfilename = rfilename
+        self.size = size
+
+
+class _FakeInfo:
+    def __init__(self, siblings):
+        self.siblings = siblings
+
+
+GB = 1024 ** 3
+# A FLUX-shaped base repo: the packaged root single and the transformer shards are what a
+# plain snapshot_download would drag in and the loader never opens.
+_FLUX_BASE_SIBLINGS = [
+    _FakeSibling("model_index.json", 1000),
+    _FakeSibling("flux1-dev.safetensors", 24 * GB),
+    _FakeSibling("transformer/diffusion_pytorch_model-00001-of-00003.safetensors", 8 * GB),
+    _FakeSibling("text_encoder/model.safetensors", 2 * GB),
+    _FakeSibling("text_encoder/model.fp16.safetensors", 1 * GB),
+    _FakeSibling("vae/diffusion_pytorch_model.safetensors", 300),
+    _FakeSibling("assets/gallery.pdf", 5000),
+    _FakeSibling("README.md", 200),
+]
+
+
+def _fake_hf_api(monkeypatch, repos):
+    """Point HfApi.model_info at a canned sibling list per repo id."""
+    class _Api:
+        def model_info(self, repo_id, files_metadata = False, token = None):
+            return _FakeInfo(repos[repo_id])
+
+    monkeypatch.setattr("huggingface_hub.HfApi", lambda *a, **k: _Api())
+
+
+def test_download_plan_scopes_the_base_repo_files(monkeypatch):
+    # The plan drives the Hub download manager, so its file list must match what the loader
+    # actually reads. A full snapshot would add the 24 GB root single and the transformer
+    # shards the GGUF replaces.
+    _fake_hf_api(
+        monkeypatch,
+        {
+            "unsloth/FLUX.1-dev-GGUF": [_FakeSibling("flux1-dev-Q4_K_M.gguf", 7 * GB)],
+            "black-forest-labs/FLUX.1-dev": _FLUX_BASE_SIBLINGS,
+        },
+    )
+    monkeypatch.setattr(
+        "core.inference.diffusion._resolve_base_repo",
+        lambda *a, **k: "black-forest-labs/FLUX.1-dev",
+    )
+    monkeypatch.setattr(
+        DiffusionBackend, "_dense_quant_prefetch_needed", lambda self, fam, kwargs: False
+    )
+
+    plan = DiffusionBackend().download_plan(
+        "unsloth/FLUX.1-dev-GGUF", gguf_filename = "flux1-dev-Q4_K_M.gguf"
+    )
+
+    assert [e["repo_id"] for e in plan["entries"]] == [
+        "unsloth/FLUX.1-dev-GGUF",
+        "black-forest-labs/FLUX.1-dev",
+    ]
+    checkpoint, base = plan["entries"]
+    assert checkpoint["files"] == ["flux1-dev-Q4_K_M.gguf"]
+    assert checkpoint["bytes"] == 7 * GB
+    assert "flux1-dev.safetensors" not in base["files"]
+    assert not any(f.startswith("transformer/") for f in base["files"])
+    assert not any(f.startswith("assets/") for f in base["files"])
+    assert "model_index.json" in base["files"]
+    assert "text_encoder/model.safetensors" in base["files"]
+    # Sized per repo, so each download job gets its own expected bytes.
+    assert base["bytes"] < 24 * GB
+    assert plan["total_bytes"] == checkpoint["bytes"] + base["bytes"]
+
+
+def test_download_plan_pipeline_kind_is_one_entry(monkeypatch):
+    # A pipeline load has no separate checkpoint repo: the repo IS the pipeline.
+    _fake_hf_api(monkeypatch, {"unsloth/some-pipeline": _FLUX_BASE_SIBLINGS})
+
+    plan = DiffusionBackend().download_plan("unsloth/some-pipeline", model_kind = "pipeline")
+
+    assert len(plan["entries"]) == 1
+    files = plan["entries"][0]["files"]
+    # The pipeline keeps its own transformer, but still drops fp16 twins and the root single.
+    assert any(f.startswith("transformer/") for f in files)
+    assert "flux1-dev.safetensors" not in files
+    assert "text_encoder/model.fp16.safetensors" not in files
+
+
+def test_download_plan_is_empty_for_a_local_path(tmp_path, monkeypatch):
+    # Nothing to stage: the files are already on disk.
+    local = tmp_path / "my-model"
+    (local / "transformer").mkdir(parents = True)
+    (local / "model_index.json").write_text("{}", encoding = "utf-8")
+    monkeypatch.setattr("core.inference.diffusion._resolve_base_repo", lambda *a, **k: str(local))
+    monkeypatch.setattr(
+        DiffusionBackend, "_estimate_download_bytes", staticmethod(lambda *a, **k: (0, []))
+    )
+
+    plan = DiffusionBackend().download_plan(str(local), gguf_filename = "weights.gguf")
+    assert plan["entries"] == []