diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 8956828684..fa88706e25 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -1215,12 +1215,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) 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/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()