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

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/core/training/diffusion_lora_trainer.py

@@ -261,9 +261,12 @@ def run_diffusion_lora_training(
     torch.manual_seed(cfg.seed)
 
     device = "cuda" if torch.cuda.is_available() else "cpu"
-    weight_dtype = {"bf16": torch.bfloat16, "fp16": torch.float16, "no": torch.float32}[
-        cfg.mixed_precision if device == "cuda" else "no"
-    ]
+    precision = cfg.mixed_precision if device == "cuda" else "no"
+    if precision == "bf16" and device == "cuda" and not torch.cuda.is_bf16_supported():
+        # The default is bf16, but pre-Ampere GPUs (T4 / V100 / RTX 20xx) have no
+        # bf16 compute; fall back to fp16 there instead of failing at load/forward.
+        precision = "fp16"
+    weight_dtype = {"bf16": torch.bfloat16, "fp16": torch.float16, "no": torch.float32}[precision]
 
     pairs = discover_image_caption_pairs(
         cfg.data_dir, instance_prompt = cfg.instance_prompt, caption_column = cfg.caption_column

studio/backend/core/training/diffusion_training_service.py

@@ -89,11 +89,21 @@ class DiffusionTrainingService:
 
         _config_from_dict(config).normalized()
 
+        # Join a finished job's pump OUTSIDE the lock: its final state writes take this
+        # lock (via _apply_event / the exit handler), so joining under it would stall
+        # the start for the whole timeout and then let the stale pump overwrite the new
+        # job's state once the lock was released.
         with self._lock:
             if self._proc is not None and self._proc.is_alive():
                 raise RuntimeError("A diffusion training job is already running.")
-            if self._pump is not None and self._pump.is_alive():
-                self._pump.join(timeout = 5.0)
+            pump = self._pump
+        if pump is not None and pump.is_alive():
+            pump.join(timeout = 5.0)
+
+        with self._lock:
+            # Re-check: another start() may have won the race while we joined.
+            if self._proc is not None and self._proc.is_alive():
+                raise RuntimeError("A diffusion training job is already running.")
 
             job_id = uuid.uuid4().hex
             event_queue = self._ctx.Queue()
@@ -162,11 +172,13 @@ class DiffusionTrainingService:
                     drained = False
                     while True:
                         try:
-                            self._apply_event(event_queue.get_nowait())
+                            self._apply_event(event_queue.get_nowait(), proc = proc)
                             drained = True
                         except Exception:  # noqa: BLE001
                             break
                     with self._lock:
+                        if self._proc is not proc:
+                            return  # superseded by a newer job; don't touch its state
                         if self._state.get("status") not in ("completed", "stopped", "error"):
                             self._state.update(
                                 active = False,
@@ -177,15 +189,19 @@ class DiffusionTrainingService:
                     _ = drained
                     return
                 continue
-            self._apply_event(ev)
+            self._apply_event(ev, proc = proc)
             if ev.get("type") in _TERMINAL:
                 return
 
-    def _apply_event(self, ev: dict[str, Any]) -> None:
+    def _apply_event(self, ev: dict[str, Any], proc: Any = None) -> None:
         """Fold one trainer event into the status snapshot. Pure state update -- unit
-        tested by feeding events directly."""
+        tested by feeding events directly. ``proc`` (when given) fences a stale pump:
+        an event from a superseded job's process must not touch the current job's
+        state."""
         etype = ev.get("type")
         with self._lock:
+            if proc is not None and self._proc is not proc:
+                return
             s = self._state
             s["updated_at"] = time.time()
             if etype == "model_load_started":

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

studio/backend/tests/test_diffusion_training.py

@@ -289,3 +289,31 @@ def test_route_status_and_stop(client):
     # After stopping, a stop with nothing running reports idle.
     st2 = client.post("/api/train/diffusion/stop")
     assert st2.json()["status"] == "idle"
+
+
+def test_service_restart_after_completion():
+    # A finished job's pump is joined OUTSIDE the lock (it needs the lock for its
+    # final state writes), so a second start neither stalls nor deadlocks.
+    svc = DiffusionTrainingService(ctx = _FakeCtx(), target = _happy_target)
+    svc.start(dict(_CFG))
+    _wait_status(svc, "completed")
+    t0 = time.time()
+    job2 = svc.start(dict(_CFG))
+    assert job2
+    assert time.time() - t0 < 4.0  # no 5s join-under-lock stall
+    st = _wait_status(svc, "completed")
+    assert st["status"] == "completed"
+
+
+def test_stale_pump_events_cannot_corrupt_new_job():
+    # An event carrying a superseded job's proc identity must be dropped, so a
+    # straggler pump can never overwrite the state of a newly started job.
+    svc = DiffusionTrainingService(ctx = _FakeCtx(), target = _happy_target)
+    svc.start(dict(_CFG))
+    _wait_status(svc, "completed")
+    current = svc._proc
+    svc._apply_event({"type": "error", "message": "stale boom"}, proc = object())
+    assert svc.status()["message"] != "stale boom"
+    # The current job's events still apply.
+    svc._apply_event({"type": "progress", "step": 9}, proc = current)
+    assert svc.status()["step"] == 9