mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-07-03 21:20:49 +08:00
351119eb05
Lets ComfyUI fetch the models a workflow needs directly on the server, so users no longer have to locate each file and drop it into the correct folder by hand. Crucially it supports gated HuggingFace repositories: the user logs in once via HuggingFace, after which the server can download models that require license acceptance or authentication — previously a manual, error-prone step. The frontend can surface per-model availability and download progress through the accompanying API.
112 lines
4.6 KiB
Python
112 lines
4.6 KiB
Python
"""Path resolution for model downloads.
|
|
|
|
Model identifiers used across the download API are *relative destination
|
|
paths* of the form ``<directory>/<filename>`` (e.g. ``loras/my_lora.safetensors``).
|
|
This module turns one of those identifiers into an absolute on-disk path
|
|
under one of ComfyUI's registered model folders, while rejecting unknown
|
|
folders, path traversal, and other ill-formed inputs.
|
|
"""
|
|
|
|
import os
|
|
import re
|
|
from typing import Optional, Tuple
|
|
|
|
import folder_paths
|
|
|
|
|
|
# Constrain components so a model_id can never escape its target directory.
|
|
# - directory: a single path segment of safe chars
|
|
# - filename: a single path segment of safe chars, must end with a model extension
|
|
_SEGMENT_RE = re.compile(r"^[A-Za-z0-9._-]+$")
|
|
|
|
# Destination filename must name a model file (same set as the URL allowlist),
|
|
# so a download can't land as e.g. ``foo.txt`` that ComfyUI won't recognise.
|
|
_MODEL_EXTENSIONS = (".safetensors", ".sft", ".ckpt", ".pth", ".pt")
|
|
|
|
# Distinctive temp suffix so the startup orphan-sweep only removes files THIS
|
|
# subsystem created — never unrelated ``*.tmp`` files in the model dirs.
|
|
_TMP_SUFFIX = ".comfy-download.tmp"
|
|
|
|
|
|
class InvalidModelId(ValueError):
|
|
"""Raised when a model_id is syntactically invalid or refers to an
|
|
unknown model folder."""
|
|
|
|
|
|
def parse_model_id(model_id: str) -> Tuple[str, str]:
|
|
"""Split ``<directory>/<filename>`` and validate both components.
|
|
|
|
Returns ``(directory, filename)``. Raises ``InvalidModelId`` on
|
|
malformed input. Does NOT touch the filesystem.
|
|
"""
|
|
if not isinstance(model_id, str) or "/" not in model_id:
|
|
raise InvalidModelId(f"model_id must be '<directory>/<filename>', got {model_id!r}")
|
|
directory, _, filename = model_id.partition("/")
|
|
if "/" in filename or not directory or not filename:
|
|
raise InvalidModelId(f"model_id must be exactly one '/' separator, got {model_id!r}")
|
|
if not _SEGMENT_RE.match(directory):
|
|
raise InvalidModelId(f"invalid directory segment {directory!r}")
|
|
if not _SEGMENT_RE.match(filename):
|
|
raise InvalidModelId(f"invalid filename segment {filename!r}")
|
|
if not filename.endswith(_MODEL_EXTENSIONS):
|
|
raise InvalidModelId(
|
|
f"filename must end with a model extension {_MODEL_EXTENSIONS}, got {filename!r}"
|
|
)
|
|
if directory not in folder_paths.folder_names_and_paths:
|
|
raise InvalidModelId(f"unknown model folder {directory!r}")
|
|
return directory, filename
|
|
|
|
|
|
def resolve_existing(model_id: str) -> Optional[str]:
|
|
"""Return the absolute path of an installed model, or None if missing.
|
|
|
|
Honours ``extra_model_paths.yaml`` transparently via
|
|
``folder_paths.get_full_path``.
|
|
"""
|
|
directory, filename = parse_model_id(model_id)
|
|
return folder_paths.get_full_path(directory, filename)
|
|
|
|
|
|
def resolve_destination(model_id: str, epoch: int = 0) -> Tuple[str, str]:
|
|
"""Return ``(final_path, tmp_path)`` for a download.
|
|
|
|
Downloads land at the first registered path for the model's directory
|
|
(the "primary" location). The temp sibling is the write target, atomically
|
|
renamed onto ``final_path`` on success.
|
|
|
|
``tmp_path`` embeds the session ``epoch`` so a cancel+retry of the same
|
|
model never shares a temp path between the old (cancelling) worker and the
|
|
new attempt — otherwise the old worker's rollback could delete the new
|
|
worker's in-progress file. The distinctive suffix scopes the orphan sweep.
|
|
"""
|
|
directory, filename = parse_model_id(model_id)
|
|
roots = folder_paths.get_folder_paths(directory)
|
|
if not roots:
|
|
raise InvalidModelId(f"no on-disk path registered for folder {directory!r}")
|
|
root = roots[0]
|
|
final_path = os.path.join(root, filename)
|
|
tmp_path = f"{final_path}.{epoch}{_TMP_SUFFIX}"
|
|
return final_path, tmp_path
|
|
|
|
|
|
def iter_all_tmp_paths():
|
|
"""Yield this subsystem's temp files under every registered model folder.
|
|
|
|
Matches only our distinctive ``_TMP_SUFFIX`` (not every ``*.tmp``) so the
|
|
startup orphan-sweep can't delete temp files created by other tools.
|
|
"""
|
|
seen_roots: set[str] = set()
|
|
for directory in folder_paths.folder_names_and_paths.keys():
|
|
for root in folder_paths.get_folder_paths(directory):
|
|
if root in seen_roots or not os.path.isdir(root):
|
|
continue
|
|
seen_roots.add(root)
|
|
try:
|
|
for entry in os.scandir(root):
|
|
if entry.is_file() and entry.name.endswith(_TMP_SUFFIX):
|
|
yield entry.path
|
|
except OSError:
|
|
# Folder might be unreadable / missing on certain mounts —
|
|
# not fatal, just skip it.
|
|
continue
|