ControlNet: reject filesystem-like ids and do not cache a model past an unload race
Two review findings on the ControlNet path: - resolve_controlnet's bare-repo fallback accepted any id with a slash, so a path-shaped id (/tmp/x, ../x) reached from_pretrained as a local directory. Restrict the fallback to a strict owner/name HF repo id shape. - _controlnet_pipe now re-checks the cancel event after the blocking from_pretrained: an unload that raced the download had already cleared the caches, so caching the late module would pin it past the unload.
This commit is contained in:
parent
52ab033c00
commit
060fac0a9d
3 changed files with 19 additions and 3 deletions
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue