assets.py•4.06 kB
"""Asset discovery and validation utilities."""
from __future__ import annotations
import difflib
from dataclasses import dataclass, field
from pathlib import Path
from typing import Dict, List, Mapping, MutableMapping, Optional, Sequence
from .config import DirectorySettings
SAFE_EXTENSIONS: Mapping[str, Sequence[str]] = {
"checkpoints": (".safetensors", ".ckpt", ".pt"),
"loras": (".safetensors", ".ckpt", ".pt"),
"vaes": (".safetensors", ".ckpt", ".pt"),
"text_encoders": (".safetensors", ".pt"),
"embeddings": (".safetensors", ".pt", ".bin"),
}
@dataclass(slots=True)
class AssetCatalog:
"""Represents discovered assets for ComfyUI components."""
directories: DirectorySettings
base_path: Path = field(default_factory=Path.cwd)
assets: MutableMapping[str, List[str]] = field(default_factory=dict)
def __post_init__(self) -> None: # pragma: no cover - dataclass hook
self.refresh()
def refresh(self) -> None:
"""Rescan configured directories and rebuild the asset cache."""
catalog: MutableMapping[str, List[str]] = {}
directory_map = self._directory_mapping()
for category, path in directory_map.items():
catalog[category] = self._scan_directory(path, SAFE_EXTENSIONS.get(category, ()))
self.assets = catalog
def list(self) -> Dict[str, List[str]]:
"""Return a copy of the cached asset lists."""
return {key: list(values) for key, values in self.assets.items()}
def validate_directories(self) -> None:
"""Ensure configured directories exist and are readable."""
errors: List[str] = []
for category, path in self._directory_mapping().items():
if path is None:
continue
resolved = self._resolve(path)
if not resolved.exists():
errors.append(f"{category} directory '{resolved}' does not exist")
elif not resolved.is_dir():
errors.append(f"{category} directory '{resolved}' is not a directory")
if errors:
raise FileNotFoundError("; ".join(errors))
def ensure_exists(self, category: str, name: str) -> None:
"""Ensure ``name`` exists in ``category`` or raise a descriptive error."""
if not name:
raise ValueError("Asset name must be provided")
if not _is_safe_name(name):
raise ValueError("Asset names must not include path separators")
available = self.assets.get(category)
if not available:
# If nothing discovered, skip strict validation.
return
if name in available:
return
suggestions = difflib.get_close_matches(name, available, n=3, cutoff=0.0)
suggestion_text = f". Closest matches: {', '.join(suggestions)}" if suggestions else ""
raise KeyError(f"Unknown {category.rstrip('s')} asset '{name}'{suggestion_text}")
def _directory_mapping(self) -> Mapping[str, Optional[Path]]:
return {
"checkpoints": self.directories.checkpoints,
"loras": self.directories.loras,
"vaes": self.directories.vaes,
"text_encoders": self.directories.text_encoders,
"embeddings": self.directories.embeddings,
}
def _scan_directory(self, path: Optional[Path], extensions: Sequence[str]) -> List[str]:
if path is None:
return []
resolved = self._resolve(path)
if not resolved.exists() or not resolved.is_dir():
return []
results: List[str] = []
for candidate in resolved.iterdir():
if candidate.is_file() and candidate.suffix.lower() in extensions:
results.append(candidate.name)
results.sort()
return results
def _resolve(self, path: Path) -> Path:
if path.is_absolute():
return path
return (self.base_path / path).resolve()
def _is_safe_name(name: str) -> bool:
return Path(name).name == name
__all__ = ["AssetCatalog"]