diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index d6642d949b..addf528819 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, @@ -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( diff --git a/studio/backend/core/inference/diffusion_controlnet.py b/studio/backend/core/inference/diffusion_controlnet.py index 22a936daac..bf103a9206 100644 --- a/studio/backend/core/inference/diffusion_controlnet.py +++ b/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( diff --git a/studio/backend/core/training/diffusion_lora_trainer.py b/studio/backend/core/training/diffusion_lora_trainer.py index f798496f43..ee095e63d2 100644 --- a/studio/backend/core/training/diffusion_lora_trainer.py +++ b/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 diff --git a/studio/backend/core/training/diffusion_training_service.py b/studio/backend/core/training/diffusion_training_service.py index bdefcfe25b..091f7b881c 100644 --- a/studio/backend/core/training/diffusion_training_service.py +++ b/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": diff --git a/studio/backend/tests/test_diffusion_controlnet.py b/studio/backend/tests/test_diffusion_controlnet.py index 9c978a483c..f9f3d5bd36 100644 --- a/studio/backend/tests/test_diffusion_controlnet.py +++ b/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() 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") diff --git a/studio/backend/tests/test_diffusion_training.py b/studio/backend/tests/test_diffusion_training.py index ffc231673c..32d239ecde 100644 --- a/studio/backend/tests/test_diffusion_training.py +++ b/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