Merge remote-tracking branch 'origin/diffusion-lora-ux' into diffusion-lora-training

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,
@@ -1243,12 +1252,17 @@ class DiffusionBackend:
         cn_model = self._cn_models.get(resolved_cn.id)
         if cn_model is None:
             if cancel.is_set():
-                raise RuntimeError("Diffusion generation was cancelled.")
+                raise RuntimeError(DIFFUSION_CANCELLED_MSG)
             cn_model = (
                 getattr(diffusers, model_cls_name)
                 .from_pretrained(resolved_cn.path, torch_dtype = state.dtype, token = state.hf_token)
                 .to(state.device)
             )
+            if cancel.is_set():
+                # An unload raced the blocking download above and already cleared the
+                # ControlNet caches; caching now would pin the module past the unload.
+                del cn_model
+                raise RuntimeError(DIFFUSION_CANCELLED_MSG)
             self._cn_models[resolved_cn.id] = cn_model
         key = (pipe_cls_name, resolved_cn.id)
         pipe = self._cn_pipes.get(key)
@@ -1867,12 +1881,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(

studio/backend/core/inference/diffusion_controlnet.py

@@ -162,8 +162,11 @@ def resolve_controlnet(
             raise ValueError(f"ControlNet '{spec_id}' has no repo")
         return ResolvedControlNet(spec_id, entry.repo_id, is_local = False)
 
-    # A bare public HF repo id (owner/name).
-    if "/" in spec_id and " " not in spec_id:
+    # A bare public HF repo id (owner/name). STRICT shape -- exactly one slash and
+    # alphanumeric-leading segments -- so a filesystem-looking id (/tmp/x, ../x, ~/x,
+    # C:\x) can never reach from_pretrained, which would happily treat it as a local
+    # directory and bypass the controlnets_dir() no-raw-path contract.
+    if re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9_.-]*/[A-Za-z0-9][A-Za-z0-9_.-]*", spec_id):
         return ResolvedControlNet(spec_id, spec_id, is_local = False)
 
     raise FileNotFoundError(

studio/backend/tests/test_diffusion_controlnet.py

@@ -37,6 +37,14 @@ def test_resolve_controlnet_catalog_bare_repo_and_unknown():
         dc.resolve_controlnet("not-a-known-id")
 
 
+def test_resolve_controlnet_rejects_filesystem_like_ids():
+    # The bare-repo fallback must never accept a path-shaped id: from_pretrained
+    # would treat it as a local directory, bypassing the controlnets_dir() contract.
+    for bad in ("/tmp/model", "../some/model", "./x/y", "~/x/y", "a/b/c", "C:\\x/y", ".hidden/x"):
+        with pytest.raises(FileNotFoundError):
+            dc.resolve_controlnet(bad)
+
+
 def test_resolve_controlnet_local(tmp_path, monkeypatch):
     d = tmp_path / "controlnets"
     d.mkdir()

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")