From ca3592062cb353d919671fe9ceb2c3bdd2fe988d Mon Sep 17 00:00:00 2001 From: Unsloth Date: Sun, 26 Jul 2026 03:58:56 -0700 Subject: [PATCH] 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 | 74 ++++++++++++- studio/backend/models/inference.py | 24 +++++ studio/backend/routes/inference.py | 51 +++++++++ .../backend/tests/test_diffusion_backend.py | 101 ++++++++++++++++++ 4 files changed, 248 insertions(+), 2 deletions(-) diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 188d14016e..44b2f842aa 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/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``. diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 6c0bfbd1d1..10f8b78c9d 100644 --- a/studio/backend/models/inference.py +++ b/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.""" diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 0df099c300..43d8257839 100644 --- a/studio/backend/routes/inference.py +++ b/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) diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index 56511ffcd0..5866e4472e 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/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"] == []