From c4191569dfa8adecf908acb9174931028acb22a1 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 1 Jul 2026 23:55:12 +0000 Subject: [PATCH] Pipeline prefetch: fetch only the default torch weights A full-pipeline prefetch kept every repo file outside assets/, so an official repo that ships multiple formats (SDXL Base: fp16 variants, ONNX, OpenVINO, Flax, a top-level single-file twin) downloaded tens of GB from_pretrained never loads. Skip non-torch exports and dtype-variant twins in _pipeline_file_downloaded, and drop a component .bin when the same directory carries a picked safetensors weight (diffusers' own preference). --- studio/backend/core/inference/diffusion.py | 36 +++++++++++++++++---- studio/backend/tests/test_diffusion_sdxl.py | 20 ++++++++++++ 2 files changed, 49 insertions(+), 7 deletions(-) diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index d6642d949b..96b7616d21 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -650,10 +650,19 @@ class DiffusionBackend: try: if kind == "pipeline": info = api.model_info(repo_id, files_metadata = True, token = hf_token) - for s in info.siblings: - if _pipeline_file_downloaded(s.rfilename): - base_files.append(s.rfilename) - total += s.size or 0 + picked = [s for s in info.siblings if _pipeline_file_downloaded(s.rfilename)] + # diffusers prefers safetensors per component: drop a .bin whose + # directory also carries a picked .safetensors weight. + st_dirs = { + s.rfilename.rsplit("/", 1)[0] + for s in picked + if s.rfilename.endswith(".safetensors") + } + for s in picked: + if s.rfilename.endswith(".bin") and s.rfilename.rsplit("/", 1)[0] in st_dirs: + continue + base_files.append(s.rfilename) + total += s.size or 0 return total, base_files # Skip the Hub size lookup for a LOCAL gguf path: model_info(repo_id) would # raise on a filesystem path and (caught below) skip the base-repo lookup too, @@ -1867,12 +1876,25 @@ def _pipeline_file_downloaded(rfilename: str) -> bool: Like ``_base_file_downloaded`` but for the ``pipeline`` kind, where the repo supplies its OWN transformer weights, so the ``transformer/`` subfolder is kept. - Top-level docs (README/PDF/images) and ``assets/`` are still skipped so the - progress estimate matches what actually lands on disk. + Top-level docs (README/PDF/images) and ``assets/`` are skipped, and so are + artifacts the torch loader never touches -- ONNX / OpenVINO / Flax exports and + dtype-variant twins (``*.fp16.safetensors``: the loader requests the default + variant) -- so an official repo that ships many formats (e.g. SDXL Base) does + not prefetch tens of GB it will not load. """ if "/" not in rfilename: # top-level: only the pipeline manifest is fetched return rfilename == "model_index.json" - return not rfilename.startswith("assets/") + lower = rfilename.lower() + if lower.startswith(("assets/", "onnx/", "openvino/")): + return False + name = lower.rsplit("/", 1)[1] + if name.startswith(("openvino_", "flax_")): + return False + if name.endswith((".onnx", ".onnx_data", ".pb", ".msgpack", ".h5", ".ckpt")): + return False + if ".fp16." in name or ".bf16." in name or ".non_ema." in name: + return False + return True def _progress( diff --git a/studio/backend/tests/test_diffusion_sdxl.py b/studio/backend/tests/test_diffusion_sdxl.py index 41ee904b94..2a78c0c7de 100644 --- a/studio/backend/tests/test_diffusion_sdxl.py +++ b/studio/backend/tests/test_diffusion_sdxl.py @@ -127,3 +127,23 @@ def test_sdxl_lora_supported_on_diffusers(): assert diffusion_lora.supports_lora( engine = "diffusers", family = "sdxl", model_kind = "single_file", transformer_quant = None ) + + +def test_pipeline_prefetch_skips_non_torch_artifacts(): + # The SDXL Base repo ships fp16 variants, ONNX, OpenVINO and Flax exports next to + # the default safetensors; from_pretrained (no variant kwarg) loads only the + # default torch weights, so the prefetch filter must skip everything else or a + # catalog load pulls tens of GB of unused artifacts. + from core.inference.diffusion import _pipeline_file_downloaded as keep + + assert keep("model_index.json") + assert keep("unet/diffusion_pytorch_model.safetensors") + assert keep("text_encoder/model.safetensors") + assert keep("scheduler/scheduler_config.json") + assert not keep("sd_xl_base_1.0.safetensors") # top-level single-file twin + assert not keep("unet/diffusion_pytorch_model.fp16.safetensors") + assert not keep("text_encoder/model.onnx") + assert not keep("text_encoder/openvino_model.bin") + assert not keep("unet/flax_model.msgpack") + assert not keep("vae_decoder/model.onnx_data") + assert not keep("assets/preview.png")