mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-02-06 19:42:34 +08:00
Refactor helpers.py: move functions to their respective modules
- Move scanner-only functions to scanner.py - Move query-only functions (is_scalar, project_kv) to asset_info.py - Move get_query_dict to routes.py - Create path_utils.py service for path-related functions - Reduce helpers.py to shared utilities only Amp-Thread-ID: https://ampcode.com/threads/T-019c2510-33fa-7199-ae4b-bc31102277a7 Co-authored-by: Amp <amp@ampcode.com>
This commit is contained in:
parent
a02f160e20
commit
2eb100adf9
@ -10,8 +10,8 @@ from pydantic import ValidationError
|
|||||||
import app.assets.manager as manager
|
import app.assets.manager as manager
|
||||||
from app import user_manager
|
from app import user_manager
|
||||||
from app.assets.api import schemas_in
|
from app.assets.api import schemas_in
|
||||||
from app.assets.helpers import get_query_dict
|
|
||||||
from app.assets.services.scanner import seed_assets
|
from app.assets.services.scanner import seed_assets
|
||||||
|
from typing import Any
|
||||||
|
|
||||||
import folder_paths
|
import folder_paths
|
||||||
|
|
||||||
@ -21,6 +21,18 @@ USER_MANAGER: user_manager.UserManager | None = None
|
|||||||
# UUID regex (canonical hyphenated form, case-insensitive)
|
# UUID regex (canonical hyphenated form, case-insensitive)
|
||||||
UUID_RE = r"[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}"
|
UUID_RE = r"[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}"
|
||||||
|
|
||||||
|
def get_query_dict(request: web.Request) -> dict[str, Any]:
|
||||||
|
"""
|
||||||
|
Gets a dictionary of query parameters from the request.
|
||||||
|
|
||||||
|
'request.query' is a MultiMapping[str], needs to be converted to a dictionary to be validated by Pydantic.
|
||||||
|
"""
|
||||||
|
query_dict = {
|
||||||
|
key: request.query.getall(key) if len(request.query.getall(key)) > 1 else request.query.get(key)
|
||||||
|
for key in request.query.keys()
|
||||||
|
}
|
||||||
|
return query_dict
|
||||||
|
|
||||||
# Note to any custom node developers reading this code:
|
# Note to any custom node developers reading this code:
|
||||||
# The assets system is not yet fully implemented, do not rely on the code in /app/assets remaining the same.
|
# The assets system is not yet fully implemented, do not rely on the code in /app/assets remaining the same.
|
||||||
|
|
||||||
|
|||||||
@ -18,7 +18,70 @@ from sqlalchemy.orm import Session, contains_eager, noload
|
|||||||
from app.assets.database.models import (
|
from app.assets.database.models import (
|
||||||
Asset, AssetInfo, AssetInfoMeta, AssetInfoTag, Tag
|
Asset, AssetInfo, AssetInfoMeta, AssetInfoTag, Tag
|
||||||
)
|
)
|
||||||
from app.assets.helpers import escape_like_prefix, normalize_tags, project_kv, utcnow
|
from app.assets.helpers import escape_like_prefix, normalize_tags, utcnow
|
||||||
|
|
||||||
|
|
||||||
|
def is_scalar(v):
|
||||||
|
if v is None:
|
||||||
|
return True
|
||||||
|
if isinstance(v, bool):
|
||||||
|
return True
|
||||||
|
if isinstance(v, (int, float, Decimal, str)):
|
||||||
|
return True
|
||||||
|
return False
|
||||||
|
|
||||||
|
|
||||||
|
def project_kv(key: str, value):
|
||||||
|
"""
|
||||||
|
Turn a metadata key/value into typed projection rows.
|
||||||
|
Returns list[dict] with keys:
|
||||||
|
key, ordinal, and one of val_str / val_num / val_bool / val_json (others None)
|
||||||
|
"""
|
||||||
|
rows: list[dict] = []
|
||||||
|
|
||||||
|
def _null_row(ordinal: int) -> dict:
|
||||||
|
return {
|
||||||
|
"key": key, "ordinal": ordinal,
|
||||||
|
"val_str": None, "val_num": None, "val_bool": None, "val_json": None
|
||||||
|
}
|
||||||
|
|
||||||
|
if value is None:
|
||||||
|
rows.append(_null_row(0))
|
||||||
|
return rows
|
||||||
|
|
||||||
|
if is_scalar(value):
|
||||||
|
if isinstance(value, bool):
|
||||||
|
rows.append({"key": key, "ordinal": 0, "val_bool": bool(value)})
|
||||||
|
elif isinstance(value, (int, float, Decimal)):
|
||||||
|
num = value if isinstance(value, Decimal) else Decimal(str(value))
|
||||||
|
rows.append({"key": key, "ordinal": 0, "val_num": num})
|
||||||
|
elif isinstance(value, str):
|
||||||
|
rows.append({"key": key, "ordinal": 0, "val_str": value})
|
||||||
|
else:
|
||||||
|
rows.append({"key": key, "ordinal": 0, "val_json": value})
|
||||||
|
return rows
|
||||||
|
|
||||||
|
if isinstance(value, list):
|
||||||
|
if all(is_scalar(x) for x in value):
|
||||||
|
for i, x in enumerate(value):
|
||||||
|
if x is None:
|
||||||
|
rows.append(_null_row(i))
|
||||||
|
elif isinstance(x, bool):
|
||||||
|
rows.append({"key": key, "ordinal": i, "val_bool": bool(x)})
|
||||||
|
elif isinstance(x, (int, float, Decimal)):
|
||||||
|
num = x if isinstance(x, Decimal) else Decimal(str(x))
|
||||||
|
rows.append({"key": key, "ordinal": i, "val_num": num})
|
||||||
|
elif isinstance(x, str):
|
||||||
|
rows.append({"key": key, "ordinal": i, "val_str": x})
|
||||||
|
else:
|
||||||
|
rows.append({"key": key, "ordinal": i, "val_json": x})
|
||||||
|
return rows
|
||||||
|
for i, x in enumerate(value):
|
||||||
|
rows.append({"key": key, "ordinal": i, "val_json": x})
|
||||||
|
return rows
|
||||||
|
|
||||||
|
rows.append({"key": key, "ordinal": 0, "val_json": value})
|
||||||
|
return rows
|
||||||
|
|
||||||
MAX_BIND_PARAMS = 800
|
MAX_BIND_PARAMS = 800
|
||||||
|
|
||||||
|
|||||||
@ -1,12 +1,6 @@
|
|||||||
import contextlib
|
|
||||||
import os
|
import os
|
||||||
from decimal import Decimal
|
|
||||||
from aiohttp import web
|
|
||||||
from datetime import datetime, timezone
|
from datetime import datetime, timezone
|
||||||
from pathlib import Path
|
from typing import Literal, Sequence
|
||||||
from typing import Literal, Any, Sequence
|
|
||||||
|
|
||||||
import folder_paths
|
|
||||||
|
|
||||||
|
|
||||||
def pick_best_live_path(states: Sequence) -> str:
|
def pick_best_live_path(states: Sequence) -> str:
|
||||||
@ -25,42 +19,8 @@ def pick_best_live_path(states: Sequence) -> str:
|
|||||||
return alive[0].file_path
|
return alive[0].file_path
|
||||||
|
|
||||||
|
|
||||||
RootType = Literal["models", "input", "output"]
|
ALLOWED_ROOTS: tuple[Literal["models", "input", "output"], ...] = ("models", "input", "output")
|
||||||
ALLOWED_ROOTS: tuple[RootType, ...] = ("models", "input", "output")
|
|
||||||
|
|
||||||
def get_query_dict(request: web.Request) -> dict[str, Any]:
|
|
||||||
"""
|
|
||||||
Gets a dictionary of query parameters from the request.
|
|
||||||
|
|
||||||
'request.query' is a MultiMapping[str], needs to be converted to a dictionary to be validated by Pydantic.
|
|
||||||
"""
|
|
||||||
query_dict = {
|
|
||||||
key: request.query.getall(key) if len(request.query.getall(key)) > 1 else request.query.get(key)
|
|
||||||
for key in request.query.keys()
|
|
||||||
}
|
|
||||||
return query_dict
|
|
||||||
|
|
||||||
def list_tree(base_dir: str) -> list[str]:
|
|
||||||
out: list[str] = []
|
|
||||||
base_abs = os.path.abspath(base_dir)
|
|
||||||
if not os.path.isdir(base_abs):
|
|
||||||
return out
|
|
||||||
for dirpath, _subdirs, filenames in os.walk(base_abs, topdown=True, followlinks=False):
|
|
||||||
for name in filenames:
|
|
||||||
out.append(os.path.abspath(os.path.join(dirpath, name)))
|
|
||||||
return out
|
|
||||||
|
|
||||||
def prefixes_for_root(root: RootType) -> list[str]:
|
|
||||||
if root == "models":
|
|
||||||
bases: list[str] = []
|
|
||||||
for _bucket, paths in get_comfy_models_folders():
|
|
||||||
bases.extend(paths)
|
|
||||||
return [os.path.abspath(p) for p in bases]
|
|
||||||
if root == "input":
|
|
||||||
return [os.path.abspath(folder_paths.get_input_directory())]
|
|
||||||
if root == "output":
|
|
||||||
return [os.path.abspath(folder_paths.get_output_directory())]
|
|
||||||
return []
|
|
||||||
|
|
||||||
def escape_like_prefix(s: str, escape: str = "!") -> tuple[str, str]:
|
def escape_like_prefix(s: str, escape: str = "!") -> tuple[str, str]:
|
||||||
"""Escapes %, _ and the escape char itself in a LIKE prefix.
|
"""Escapes %, _ and the escape char itself in a LIKE prefix.
|
||||||
@ -70,172 +30,11 @@ def escape_like_prefix(s: str, escape: str = "!") -> tuple[str, str]:
|
|||||||
s = s.replace("%", escape + "%").replace("_", escape + "_") # escape LIKE wildcards
|
s = s.replace("%", escape + "%").replace("_", escape + "_") # escape LIKE wildcards
|
||||||
return s, escape
|
return s, escape
|
||||||
|
|
||||||
def fast_asset_file_check(
|
|
||||||
mtime_db: int | None,
|
|
||||||
size_db: int | None,
|
|
||||||
stat_result: os.stat_result,
|
|
||||||
) -> bool:
|
|
||||||
if mtime_db is None:
|
|
||||||
return False
|
|
||||||
actual_mtime_ns = getattr(stat_result, "st_mtime_ns", int(stat_result.st_mtime * 1_000_000_000))
|
|
||||||
if int(mtime_db) != int(actual_mtime_ns):
|
|
||||||
return False
|
|
||||||
sz = int(size_db or 0)
|
|
||||||
if sz > 0:
|
|
||||||
return int(stat_result.st_size) == sz
|
|
||||||
return True
|
|
||||||
|
|
||||||
def utcnow() -> datetime:
|
def utcnow() -> datetime:
|
||||||
"""Naive UTC timestamp (no tzinfo). We always treat DB datetimes as UTC."""
|
"""Naive UTC timestamp (no tzinfo). We always treat DB datetimes as UTC."""
|
||||||
return datetime.now(timezone.utc).replace(tzinfo=None)
|
return datetime.now(timezone.utc).replace(tzinfo=None)
|
||||||
|
|
||||||
def get_comfy_models_folders() -> list[tuple[str, list[str]]]:
|
|
||||||
"""Build a list of (folder_name, base_paths[]) categories that are configured for model locations.
|
|
||||||
|
|
||||||
We trust `folder_paths.folder_names_and_paths` and include a category if
|
|
||||||
*any* of its base paths lies under the Comfy `models_dir`.
|
|
||||||
"""
|
|
||||||
targets: list[tuple[str, list[str]]] = []
|
|
||||||
models_root = os.path.abspath(folder_paths.models_dir)
|
|
||||||
for name, values in folder_paths.folder_names_and_paths.items():
|
|
||||||
paths, _exts = values[0], values[1] # NOTE: this prevents nodepacks that hackily edit folder_... from breaking ComfyUI
|
|
||||||
if any(os.path.abspath(p).startswith(models_root + os.sep) for p in paths):
|
|
||||||
targets.append((name, paths))
|
|
||||||
return targets
|
|
||||||
|
|
||||||
def resolve_destination_from_tags(tags: list[str]) -> tuple[str, list[str]]:
|
|
||||||
"""Validates and maps tags -> (base_dir, subdirs_for_fs)"""
|
|
||||||
root = tags[0]
|
|
||||||
if root == "models":
|
|
||||||
if len(tags) < 2:
|
|
||||||
raise ValueError("at least two tags required for model asset")
|
|
||||||
try:
|
|
||||||
bases = folder_paths.folder_names_and_paths[tags[1]][0]
|
|
||||||
except KeyError:
|
|
||||||
raise ValueError(f"unknown model category '{tags[1]}'")
|
|
||||||
if not bases:
|
|
||||||
raise ValueError(f"no base path configured for category '{tags[1]}'")
|
|
||||||
base_dir = os.path.abspath(bases[0])
|
|
||||||
raw_subdirs = tags[2:]
|
|
||||||
else:
|
|
||||||
base_dir = os.path.abspath(
|
|
||||||
folder_paths.get_input_directory() if root == "input" else folder_paths.get_output_directory()
|
|
||||||
)
|
|
||||||
raw_subdirs = tags[1:]
|
|
||||||
for i in raw_subdirs:
|
|
||||||
if i in (".", ".."):
|
|
||||||
raise ValueError("invalid path component in tags")
|
|
||||||
|
|
||||||
return base_dir, raw_subdirs if raw_subdirs else []
|
|
||||||
|
|
||||||
def ensure_within_base(candidate: str, base: str) -> None:
|
|
||||||
cand_abs = os.path.abspath(candidate)
|
|
||||||
base_abs = os.path.abspath(base)
|
|
||||||
try:
|
|
||||||
if os.path.commonpath([cand_abs, base_abs]) != base_abs:
|
|
||||||
raise ValueError("destination escapes base directory")
|
|
||||||
except Exception:
|
|
||||||
raise ValueError("invalid destination path")
|
|
||||||
|
|
||||||
def compute_relative_filename(file_path: str) -> str | None:
|
|
||||||
"""
|
|
||||||
Return the model's path relative to the last well-known folder (the model category),
|
|
||||||
using forward slashes, eg:
|
|
||||||
/.../models/checkpoints/flux/123/flux.safetensors -> "flux/123/flux.safetensors"
|
|
||||||
/.../models/text_encoders/clip_g.safetensors -> "clip_g.safetensors"
|
|
||||||
|
|
||||||
For non-model paths, returns None.
|
|
||||||
NOTE: this is a temporary helper, used only for initializing metadata["filename"] field.
|
|
||||||
"""
|
|
||||||
try:
|
|
||||||
root_category, rel_path = get_relative_to_root_category_path_of_asset(file_path)
|
|
||||||
except ValueError:
|
|
||||||
return None
|
|
||||||
|
|
||||||
p = Path(rel_path)
|
|
||||||
parts = [seg for seg in p.parts if seg not in (".", "..", p.anchor)]
|
|
||||||
if not parts:
|
|
||||||
return None
|
|
||||||
|
|
||||||
if root_category == "models":
|
|
||||||
# parts[0] is the category ("checkpoints", "vae", etc) – drop it
|
|
||||||
inside = parts[1:] if len(parts) > 1 else [parts[0]]
|
|
||||||
return "/".join(inside)
|
|
||||||
return "/".join(parts) # input/output: keep all parts
|
|
||||||
|
|
||||||
def get_relative_to_root_category_path_of_asset(file_path: str) -> tuple[Literal["input", "output", "models"], str]:
|
|
||||||
"""Given an absolute or relative file path, determine which root category the path belongs to:
|
|
||||||
- 'input' if the file resides under `folder_paths.get_input_directory()`
|
|
||||||
- 'output' if the file resides under `folder_paths.get_output_directory()`
|
|
||||||
- 'models' if the file resides under any base path of categories returned by `get_comfy_models_folders()`
|
|
||||||
|
|
||||||
Returns:
|
|
||||||
(root_category, relative_path_inside_that_root)
|
|
||||||
For 'models', the relative path is prefixed with the category name:
|
|
||||||
e.g. ('models', 'vae/test/sub/ae.safetensors')
|
|
||||||
|
|
||||||
Raises:
|
|
||||||
ValueError: if the path does not belong to input, output, or configured model bases.
|
|
||||||
"""
|
|
||||||
fp_abs = os.path.abspath(file_path)
|
|
||||||
|
|
||||||
def _is_within(child: str, parent: str) -> bool:
|
|
||||||
try:
|
|
||||||
return os.path.commonpath([child, parent]) == parent
|
|
||||||
except Exception:
|
|
||||||
return False
|
|
||||||
|
|
||||||
def _rel(child: str, parent: str) -> str:
|
|
||||||
return os.path.relpath(os.path.join(os.sep, os.path.relpath(child, parent)), os.sep)
|
|
||||||
|
|
||||||
# 1) input
|
|
||||||
input_base = os.path.abspath(folder_paths.get_input_directory())
|
|
||||||
if _is_within(fp_abs, input_base):
|
|
||||||
return "input", _rel(fp_abs, input_base)
|
|
||||||
|
|
||||||
# 2) output
|
|
||||||
output_base = os.path.abspath(folder_paths.get_output_directory())
|
|
||||||
if _is_within(fp_abs, output_base):
|
|
||||||
return "output", _rel(fp_abs, output_base)
|
|
||||||
|
|
||||||
# 3) models (check deepest matching base to avoid ambiguity)
|
|
||||||
best: tuple[int, str, str] | None = None # (base_len, bucket, rel_inside_bucket)
|
|
||||||
for bucket, bases in get_comfy_models_folders():
|
|
||||||
for b in bases:
|
|
||||||
base_abs = os.path.abspath(b)
|
|
||||||
if not _is_within(fp_abs, base_abs):
|
|
||||||
continue
|
|
||||||
cand = (len(base_abs), bucket, _rel(fp_abs, base_abs))
|
|
||||||
if best is None or cand[0] > best[0]:
|
|
||||||
best = cand
|
|
||||||
|
|
||||||
if best is not None:
|
|
||||||
_, bucket, rel_inside = best
|
|
||||||
combined = os.path.join(bucket, rel_inside)
|
|
||||||
return "models", os.path.relpath(os.path.join(os.sep, combined), os.sep)
|
|
||||||
|
|
||||||
raise ValueError(f"Path is not within input, output, or configured model bases: {file_path}")
|
|
||||||
|
|
||||||
def get_name_and_tags_from_asset_path(file_path: str) -> tuple[str, list[str]]:
|
|
||||||
"""Return a tuple (name, tags) derived from a filesystem path.
|
|
||||||
|
|
||||||
Semantics:
|
|
||||||
- Root category is determined by `get_relative_to_root_category_path_of_asset`.
|
|
||||||
- The returned `name` is the base filename with extension from the relative path.
|
|
||||||
- The returned `tags` are:
|
|
||||||
[root_category] + parent folders of the relative path (in order)
|
|
||||||
For 'models', this means:
|
|
||||||
file '/.../ModelsDir/vae/test_tag/ae.safetensors'
|
|
||||||
-> root_category='models', some_path='vae/test_tag/ae.safetensors'
|
|
||||||
-> name='ae.safetensors', tags=['models', 'vae', 'test_tag']
|
|
||||||
|
|
||||||
Raises:
|
|
||||||
ValueError: if the path does not belong to input, output, or configured model bases.
|
|
||||||
"""
|
|
||||||
root_category, some_path = get_relative_to_root_category_path_of_asset(file_path)
|
|
||||||
p = Path(some_path)
|
|
||||||
parent_parts = [part for part in p.parent.parts if part not in (".", "..", p.anchor)]
|
|
||||||
return p.name, list(dict.fromkeys(normalize_tags([root_category, *parent_parts])))
|
|
||||||
|
|
||||||
def normalize_tags(tags: list[str] | None) -> list[str]:
|
def normalize_tags(tags: list[str] | None) -> list[str]:
|
||||||
"""
|
"""
|
||||||
@ -245,83 +44,3 @@ def normalize_tags(tags: list[str] | None) -> list[str]:
|
|||||||
"""
|
"""
|
||||||
return [t.strip().lower() for t in (tags or []) if (t or "").strip()]
|
return [t.strip().lower() for t in (tags or []) if (t or "").strip()]
|
||||||
|
|
||||||
def collect_models_files() -> list[str]:
|
|
||||||
out: list[str] = []
|
|
||||||
for folder_name, bases in get_comfy_models_folders():
|
|
||||||
rel_files = folder_paths.get_filename_list(folder_name) or []
|
|
||||||
for rel_path in rel_files:
|
|
||||||
abs_path = folder_paths.get_full_path(folder_name, rel_path)
|
|
||||||
if not abs_path:
|
|
||||||
continue
|
|
||||||
abs_path = os.path.abspath(abs_path)
|
|
||||||
allowed = False
|
|
||||||
for b in bases:
|
|
||||||
base_abs = os.path.abspath(b)
|
|
||||||
with contextlib.suppress(Exception):
|
|
||||||
if os.path.commonpath([abs_path, base_abs]) == base_abs:
|
|
||||||
allowed = True
|
|
||||||
break
|
|
||||||
if allowed:
|
|
||||||
out.append(abs_path)
|
|
||||||
return out
|
|
||||||
|
|
||||||
def is_scalar(v):
|
|
||||||
if v is None:
|
|
||||||
return True
|
|
||||||
if isinstance(v, bool):
|
|
||||||
return True
|
|
||||||
if isinstance(v, (int, float, Decimal, str)):
|
|
||||||
return True
|
|
||||||
return False
|
|
||||||
|
|
||||||
def project_kv(key: str, value):
|
|
||||||
"""
|
|
||||||
Turn a metadata key/value into typed projection rows.
|
|
||||||
Returns list[dict] with keys:
|
|
||||||
key, ordinal, and one of val_str / val_num / val_bool / val_json (others None)
|
|
||||||
"""
|
|
||||||
rows: list[dict] = []
|
|
||||||
|
|
||||||
def _null_row(ordinal: int) -> dict:
|
|
||||||
return {
|
|
||||||
"key": key, "ordinal": ordinal,
|
|
||||||
"val_str": None, "val_num": None, "val_bool": None, "val_json": None
|
|
||||||
}
|
|
||||||
|
|
||||||
if value is None:
|
|
||||||
rows.append(_null_row(0))
|
|
||||||
return rows
|
|
||||||
|
|
||||||
if is_scalar(value):
|
|
||||||
if isinstance(value, bool):
|
|
||||||
rows.append({"key": key, "ordinal": 0, "val_bool": bool(value)})
|
|
||||||
elif isinstance(value, (int, float, Decimal)):
|
|
||||||
num = value if isinstance(value, Decimal) else Decimal(str(value))
|
|
||||||
rows.append({"key": key, "ordinal": 0, "val_num": num})
|
|
||||||
elif isinstance(value, str):
|
|
||||||
rows.append({"key": key, "ordinal": 0, "val_str": value})
|
|
||||||
else:
|
|
||||||
rows.append({"key": key, "ordinal": 0, "val_json": value})
|
|
||||||
return rows
|
|
||||||
|
|
||||||
if isinstance(value, list):
|
|
||||||
if all(is_scalar(x) for x in value):
|
|
||||||
for i, x in enumerate(value):
|
|
||||||
if x is None:
|
|
||||||
rows.append(_null_row(i))
|
|
||||||
elif isinstance(x, bool):
|
|
||||||
rows.append({"key": key, "ordinal": i, "val_bool": bool(x)})
|
|
||||||
elif isinstance(x, (int, float, Decimal)):
|
|
||||||
num = x if isinstance(x, Decimal) else Decimal(str(x))
|
|
||||||
rows.append({"key": key, "ordinal": i, "val_num": num})
|
|
||||||
elif isinstance(x, str):
|
|
||||||
rows.append({"key": key, "ordinal": i, "val_str": x})
|
|
||||||
else:
|
|
||||||
rows.append({"key": key, "ordinal": i, "val_json": x})
|
|
||||||
return rows
|
|
||||||
for i, x in enumerate(value):
|
|
||||||
rows.append({"key": key, "ordinal": i, "val_json": x})
|
|
||||||
return rows
|
|
||||||
|
|
||||||
rows.append({"key": key, "ordinal": 0, "val_json": value})
|
|
||||||
return rows
|
|
||||||
|
|||||||
@ -26,9 +26,9 @@ from app.assets.database.queries import (
|
|||||||
list_cache_states_by_asset_id,
|
list_cache_states_by_asset_id,
|
||||||
touch_asset_info_by_id,
|
touch_asset_info_by_id,
|
||||||
)
|
)
|
||||||
from app.assets.helpers import (
|
from app.assets.helpers import pick_best_live_path
|
||||||
|
from app.assets.services.path_utils import (
|
||||||
ensure_within_base,
|
ensure_within_base,
|
||||||
pick_best_live_path,
|
|
||||||
resolve_destination_from_tags,
|
resolve_destination_from_tags,
|
||||||
)
|
)
|
||||||
from app.assets.services import (
|
from app.assets.services import (
|
||||||
|
|||||||
@ -13,11 +13,8 @@ from typing import Sequence
|
|||||||
|
|
||||||
from app.assets.database.models import Asset
|
from app.assets.database.models import Asset
|
||||||
from app.database.db import create_session
|
from app.database.db import create_session
|
||||||
from app.assets.helpers import (
|
from app.assets.helpers import pick_best_live_path, utcnow
|
||||||
compute_relative_filename,
|
from app.assets.services.path_utils import compute_relative_filename
|
||||||
pick_best_live_path,
|
|
||||||
utcnow,
|
|
||||||
)
|
|
||||||
from app.assets.database.queries import (
|
from app.assets.database.queries import (
|
||||||
asset_info_exists_for_asset_id,
|
asset_info_exists_for_asset_id,
|
||||||
delete_asset_info_by_id,
|
delete_asset_info_by_id,
|
||||||
|
|||||||
@ -13,12 +13,8 @@ from sqlalchemy import select
|
|||||||
|
|
||||||
from app.assets.database.models import Asset, Tag
|
from app.assets.database.models import Asset, Tag
|
||||||
from app.database.db import create_session
|
from app.database.db import create_session
|
||||||
from app.assets.helpers import (
|
from app.assets.helpers import normalize_tags, pick_best_live_path, utcnow
|
||||||
compute_relative_filename,
|
from app.assets.services.path_utils import compute_relative_filename
|
||||||
normalize_tags,
|
|
||||||
pick_best_live_path,
|
|
||||||
utcnow,
|
|
||||||
)
|
|
||||||
from app.assets.database.queries import (
|
from app.assets.database.queries import (
|
||||||
get_asset_by_hash,
|
get_asset_by_hash,
|
||||||
get_or_create_asset_info,
|
get_or_create_asset_info,
|
||||||
|
|||||||
148
app/assets/services/path_utils.py
Normal file
148
app/assets/services/path_utils.py
Normal file
@ -0,0 +1,148 @@
|
|||||||
|
import os
|
||||||
|
from pathlib import Path
|
||||||
|
from typing import Literal
|
||||||
|
|
||||||
|
import folder_paths
|
||||||
|
|
||||||
|
from app.assets.helpers import normalize_tags
|
||||||
|
|
||||||
|
|
||||||
|
def resolve_destination_from_tags(tags: list[str]) -> tuple[str, list[str]]:
|
||||||
|
"""Validates and maps tags -> (base_dir, subdirs_for_fs)"""
|
||||||
|
root = tags[0]
|
||||||
|
if root == "models":
|
||||||
|
if len(tags) < 2:
|
||||||
|
raise ValueError("at least two tags required for model asset")
|
||||||
|
try:
|
||||||
|
bases = folder_paths.folder_names_and_paths[tags[1]][0]
|
||||||
|
except KeyError:
|
||||||
|
raise ValueError(f"unknown model category '{tags[1]}'")
|
||||||
|
if not bases:
|
||||||
|
raise ValueError(f"no base path configured for category '{tags[1]}'")
|
||||||
|
base_dir = os.path.abspath(bases[0])
|
||||||
|
raw_subdirs = tags[2:]
|
||||||
|
else:
|
||||||
|
base_dir = os.path.abspath(
|
||||||
|
folder_paths.get_input_directory() if root == "input" else folder_paths.get_output_directory()
|
||||||
|
)
|
||||||
|
raw_subdirs = tags[1:]
|
||||||
|
for i in raw_subdirs:
|
||||||
|
if i in (".", ".."):
|
||||||
|
raise ValueError("invalid path component in tags")
|
||||||
|
|
||||||
|
return base_dir, raw_subdirs if raw_subdirs else []
|
||||||
|
|
||||||
|
|
||||||
|
def ensure_within_base(candidate: str, base: str) -> None:
|
||||||
|
cand_abs = os.path.abspath(candidate)
|
||||||
|
base_abs = os.path.abspath(base)
|
||||||
|
try:
|
||||||
|
if os.path.commonpath([cand_abs, base_abs]) != base_abs:
|
||||||
|
raise ValueError("destination escapes base directory")
|
||||||
|
except Exception:
|
||||||
|
raise ValueError("invalid destination path")
|
||||||
|
|
||||||
|
|
||||||
|
def compute_relative_filename(file_path: str) -> str | None:
|
||||||
|
"""
|
||||||
|
Return the model's path relative to the last well-known folder (the model category),
|
||||||
|
using forward slashes, eg:
|
||||||
|
/.../models/checkpoints/flux/123/flux.safetensors -> "flux/123/flux.safetensors"
|
||||||
|
/.../models/text_encoders/clip_g.safetensors -> "clip_g.safetensors"
|
||||||
|
|
||||||
|
For non-model paths, returns None.
|
||||||
|
NOTE: this is a temporary helper, used only for initializing metadata["filename"] field.
|
||||||
|
"""
|
||||||
|
try:
|
||||||
|
root_category, rel_path = get_relative_to_root_category_path_of_asset(file_path)
|
||||||
|
except ValueError:
|
||||||
|
return None
|
||||||
|
|
||||||
|
p = Path(rel_path)
|
||||||
|
parts = [seg for seg in p.parts if seg not in (".", "..", p.anchor)]
|
||||||
|
if not parts:
|
||||||
|
return None
|
||||||
|
|
||||||
|
if root_category == "models":
|
||||||
|
# parts[0] is the category ("checkpoints", "vae", etc) – drop it
|
||||||
|
inside = parts[1:] if len(parts) > 1 else [parts[0]]
|
||||||
|
return "/".join(inside)
|
||||||
|
return "/".join(parts) # input/output: keep all parts
|
||||||
|
|
||||||
|
|
||||||
|
def get_relative_to_root_category_path_of_asset(file_path: str) -> tuple[Literal["input", "output", "models"], str]:
|
||||||
|
"""Given an absolute or relative file path, determine which root category the path belongs to:
|
||||||
|
- 'input' if the file resides under `folder_paths.get_input_directory()`
|
||||||
|
- 'output' if the file resides under `folder_paths.get_output_directory()`
|
||||||
|
- 'models' if the file resides under any base path of categories returned by `get_comfy_models_folders()`
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
(root_category, relative_path_inside_that_root)
|
||||||
|
For 'models', the relative path is prefixed with the category name:
|
||||||
|
e.g. ('models', 'vae/test/sub/ae.safetensors')
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
ValueError: if the path does not belong to input, output, or configured model bases.
|
||||||
|
"""
|
||||||
|
from app.assets.services.scanner import get_comfy_models_folders
|
||||||
|
|
||||||
|
fp_abs = os.path.abspath(file_path)
|
||||||
|
|
||||||
|
def _is_within(child: str, parent: str) -> bool:
|
||||||
|
try:
|
||||||
|
return os.path.commonpath([child, parent]) == parent
|
||||||
|
except Exception:
|
||||||
|
return False
|
||||||
|
|
||||||
|
def _rel(child: str, parent: str) -> str:
|
||||||
|
return os.path.relpath(os.path.join(os.sep, os.path.relpath(child, parent)), os.sep)
|
||||||
|
|
||||||
|
# 1) input
|
||||||
|
input_base = os.path.abspath(folder_paths.get_input_directory())
|
||||||
|
if _is_within(fp_abs, input_base):
|
||||||
|
return "input", _rel(fp_abs, input_base)
|
||||||
|
|
||||||
|
# 2) output
|
||||||
|
output_base = os.path.abspath(folder_paths.get_output_directory())
|
||||||
|
if _is_within(fp_abs, output_base):
|
||||||
|
return "output", _rel(fp_abs, output_base)
|
||||||
|
|
||||||
|
# 3) models (check deepest matching base to avoid ambiguity)
|
||||||
|
best: tuple[int, str, str] | None = None # (base_len, bucket, rel_inside_bucket)
|
||||||
|
for bucket, bases in get_comfy_models_folders():
|
||||||
|
for b in bases:
|
||||||
|
base_abs = os.path.abspath(b)
|
||||||
|
if not _is_within(fp_abs, base_abs):
|
||||||
|
continue
|
||||||
|
cand = (len(base_abs), bucket, _rel(fp_abs, base_abs))
|
||||||
|
if best is None or cand[0] > best[0]:
|
||||||
|
best = cand
|
||||||
|
|
||||||
|
if best is not None:
|
||||||
|
_, bucket, rel_inside = best
|
||||||
|
combined = os.path.join(bucket, rel_inside)
|
||||||
|
return "models", os.path.relpath(os.path.join(os.sep, combined), os.sep)
|
||||||
|
|
||||||
|
raise ValueError(f"Path is not within input, output, or configured model bases: {file_path}")
|
||||||
|
|
||||||
|
|
||||||
|
def get_name_and_tags_from_asset_path(file_path: str) -> tuple[str, list[str]]:
|
||||||
|
"""Return a tuple (name, tags) derived from a filesystem path.
|
||||||
|
|
||||||
|
Semantics:
|
||||||
|
- Root category is determined by `get_relative_to_root_category_path_of_asset`.
|
||||||
|
- The returned `name` is the base filename with extension from the relative path.
|
||||||
|
- The returned `tags` are:
|
||||||
|
[root_category] + parent folders of the relative path (in order)
|
||||||
|
For 'models', this means:
|
||||||
|
file '/.../ModelsDir/vae/test_tag/ae.safetensors'
|
||||||
|
-> root_category='models', some_path='vae/test_tag/ae.safetensors'
|
||||||
|
-> name='ae.safetensors', tags=['models', 'vae', 'test_tag']
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
ValueError: if the path does not belong to input, output, or configured model bases.
|
||||||
|
"""
|
||||||
|
root_category, some_path = get_relative_to_root_category_path_of_asset(file_path)
|
||||||
|
p = Path(some_path)
|
||||||
|
parent_parts = [part for part in p.parent.parts if part not in (".", "..", p.anchor)]
|
||||||
|
return p.name, list(dict.fromkeys(normalize_tags([root_category, *parent_parts])))
|
||||||
@ -3,6 +3,7 @@ import logging
|
|||||||
import os
|
import os
|
||||||
import time
|
import time
|
||||||
import uuid
|
import uuid
|
||||||
|
from typing import Literal
|
||||||
|
|
||||||
import folder_paths
|
import folder_paths
|
||||||
from sqlalchemy.orm import Session
|
from sqlalchemy.orm import Session
|
||||||
@ -25,19 +26,90 @@ from app.assets.database.queries import (
|
|||||||
get_asset_info_ids_by_ids,
|
get_asset_info_ids_by_ids,
|
||||||
bulk_insert_tags_and_meta,
|
bulk_insert_tags_and_meta,
|
||||||
)
|
)
|
||||||
from app.assets.helpers import (
|
from app.assets.helpers import utcnow
|
||||||
collect_models_files,
|
from app.assets.services.path_utils import compute_relative_filename, get_name_and_tags_from_asset_path
|
||||||
compute_relative_filename,
|
|
||||||
fast_asset_file_check,
|
|
||||||
get_name_and_tags_from_asset_path,
|
|
||||||
list_tree,
|
|
||||||
prefixes_for_root,
|
|
||||||
RootType,
|
|
||||||
utcnow,
|
|
||||||
)
|
|
||||||
from app.database.db import create_session, dependencies_available
|
from app.database.db import create_session, dependencies_available
|
||||||
|
|
||||||
|
|
||||||
|
RootType = Literal["models", "input", "output"]
|
||||||
|
|
||||||
|
|
||||||
|
def fast_asset_file_check(
|
||||||
|
mtime_db: int | None,
|
||||||
|
size_db: int | None,
|
||||||
|
stat_result: os.stat_result,
|
||||||
|
) -> bool:
|
||||||
|
if mtime_db is None:
|
||||||
|
return False
|
||||||
|
actual_mtime_ns = getattr(stat_result, "st_mtime_ns", int(stat_result.st_mtime * 1_000_000_000))
|
||||||
|
if int(mtime_db) != int(actual_mtime_ns):
|
||||||
|
return False
|
||||||
|
sz = int(size_db or 0)
|
||||||
|
if sz > 0:
|
||||||
|
return int(stat_result.st_size) == sz
|
||||||
|
return True
|
||||||
|
|
||||||
|
|
||||||
|
def list_tree(base_dir: str) -> list[str]:
|
||||||
|
out: list[str] = []
|
||||||
|
base_abs = os.path.abspath(base_dir)
|
||||||
|
if not os.path.isdir(base_abs):
|
||||||
|
return out
|
||||||
|
for dirpath, _subdirs, filenames in os.walk(base_abs, topdown=True, followlinks=False):
|
||||||
|
for name in filenames:
|
||||||
|
out.append(os.path.abspath(os.path.join(dirpath, name)))
|
||||||
|
return out
|
||||||
|
|
||||||
|
|
||||||
|
def get_comfy_models_folders() -> list[tuple[str, list[str]]]:
|
||||||
|
"""Build a list of (folder_name, base_paths[]) categories that are configured for model locations.
|
||||||
|
|
||||||
|
We trust `folder_paths.folder_names_and_paths` and include a category if
|
||||||
|
*any* of its base paths lies under the Comfy `models_dir`.
|
||||||
|
"""
|
||||||
|
targets: list[tuple[str, list[str]]] = []
|
||||||
|
models_root = os.path.abspath(folder_paths.models_dir)
|
||||||
|
for name, values in folder_paths.folder_names_and_paths.items():
|
||||||
|
paths, _exts = values[0], values[1] # NOTE: this prevents nodepacks that hackily edit folder_... from breaking ComfyUI
|
||||||
|
if any(os.path.abspath(p).startswith(models_root + os.sep) for p in paths):
|
||||||
|
targets.append((name, paths))
|
||||||
|
return targets
|
||||||
|
|
||||||
|
|
||||||
|
def prefixes_for_root(root: RootType) -> list[str]:
|
||||||
|
if root == "models":
|
||||||
|
bases: list[str] = []
|
||||||
|
for _bucket, paths in get_comfy_models_folders():
|
||||||
|
bases.extend(paths)
|
||||||
|
return [os.path.abspath(p) for p in bases]
|
||||||
|
if root == "input":
|
||||||
|
return [os.path.abspath(folder_paths.get_input_directory())]
|
||||||
|
if root == "output":
|
||||||
|
return [os.path.abspath(folder_paths.get_output_directory())]
|
||||||
|
return []
|
||||||
|
|
||||||
|
|
||||||
|
def collect_models_files() -> list[str]:
|
||||||
|
out: list[str] = []
|
||||||
|
for folder_name, bases in get_comfy_models_folders():
|
||||||
|
rel_files = folder_paths.get_filename_list(folder_name) or []
|
||||||
|
for rel_path in rel_files:
|
||||||
|
abs_path = folder_paths.get_full_path(folder_name, rel_path)
|
||||||
|
if not abs_path:
|
||||||
|
continue
|
||||||
|
abs_path = os.path.abspath(abs_path)
|
||||||
|
allowed = False
|
||||||
|
for b in bases:
|
||||||
|
base_abs = os.path.abspath(b)
|
||||||
|
with contextlib.suppress(Exception):
|
||||||
|
if os.path.commonpath([abs_path, base_abs]) == base_abs:
|
||||||
|
allowed = True
|
||||||
|
break
|
||||||
|
if allowed:
|
||||||
|
out.append(abs_path)
|
||||||
|
return out
|
||||||
|
|
||||||
|
|
||||||
def _seed_from_paths_batch(
|
def _seed_from_paths_batch(
|
||||||
session: Session,
|
session: Session,
|
||||||
specs: list[dict],
|
specs: list[dict],
|
||||||
|
|||||||
Loading…
Reference in New Issue
Block a user