nodes/Azornes-Comfyui-Model-Resolver
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

refactor: centralize category mapping and atomic JSON writing

Browse Source
This commit is contained in:
Dariusz L
2026-06-25 16:45:55 +02:00
parent 96036f0d6e
commit 8655f92515
6 changed files with 56 additions and 156 deletions
Download Patch File Download Diff File Expand all files Collapse all files
__init__.py
+1 -48
View File
@@ -766,54 +766,7 @@ class ModelResolverExtension:
try:
import folder_paths
# Map category to folder_paths type
category_map = {
"loras": "loras",
"checkpoints": "checkpoints",
"vae": "vae",
"controlnet": "controlnet",
"clip": "text_encoders",
"clips": "text_encoders",
"text_encoder": "text_encoders",
"text_encoders": "text_encoders",
"diffusion_model": "diffusion_models",
"diffusion_models": "diffusion_models",
"unet": "diffusion_models",
"upscale_models": "upscale_models",
"upscale_model": "upscale_models",
"latent_upscale_model": "latent_upscale_models",
"latent_upscale_models": "latent_upscale_models",
"style_model": "style_models",
"style_models": "style_models",
"gligen": "gligen",
"diffusers": "diffusers",
"vae_approx": "vae_approx",
"sam": "sams",
"sam_model": "sams",
"sam_models": "sams",
"sams": "sams",
"ultralytics": "ultralytics",
"ultralytics_bbox": "ultralytics",
"ultralytics_segm": "ultralytics",
"yolo": "ultralytics",
"audio_encoder": "audio_encoders",
"audio_encoders": "audio_encoders",
"background_removal": "background_removal",
"background_removal_model": "background_removal",
"frame_interpolation": "frame_interpolation",
"frame_interpolation_model": "frame_interpolation",
"geometry_estimation": "geometry_estimation",
"geometry_estimation_model": "geometry_estimation",
"detection": "detection",
"model_patch": "model_patches",
"model_patches": "model_patches",
"photomaker": "photomaker",
"optical_flow": "optical_flow",
"optical_flow_model": "optical_flow",
}
folder_type = category_map.get(
category.lower(), category.lower()
)
folder_type = normalize_download_category(category)
file_path = folder_paths.get_full_path(
folder_type, filename
)
core/downloader.py
+4 -80
View File
@@ -23,7 +23,7 @@ from .log_system.log_funcs import (
log_exception,
)
from .resolver import normalize_sha256
from .path_utils import is_path_within, get_path_identity
from .path_utils import is_path_within, get_path_identity, write_json_atomic
from .type_utils import as_dict, as_list
try:
@@ -41,61 +41,7 @@ SPEED_HISTORY_SIZE = 5 # Number of samples for smoothing
CHUNK_SIZE = 1024 * 1024 # 1MB chunks for faster downloads
CLI_LOG_INTERVAL = 5 # Log progress to CLI every N seconds
# Map common category names to folder_paths keys. "unet" and "clip" are legacy
# in ComfyUI; new files should download into diffusion_models/text_encoders.
CATEGORY_MAP = {
"checkpoint": "checkpoints",
"checkpoints": "checkpoints",
"lora": "loras",
"loras": "loras",
"vae": "vae",
"controlnet": "controlnet",
"clip": "text_encoders",
"clips": "text_encoders",
"clip_vision": "clip_vision",
"upscaler": "upscale_models",
"upscale_models": "upscale_models",
"embeddings": "embeddings",
"embedding": "embeddings",
"diffusion_model": "diffusion_models",
"diffusion_models": "diffusion_models",
"unet": "diffusion_models",
"latent_upscale_model": "latent_upscale_models",
"latent_upscale_models": "latent_upscale_models",
"text_encoders": "text_encoders",
"text_encoder": "text_encoders",
"style_model": "style_models",
"style_models": "style_models",
"gligen": "gligen",
"diffusers": "diffusers",
"vae_approx": "vae_approx",
"sam": "sams",
"sam_model": "sams",
"sam_models": "sams",
"sams": "sams",
"ultralytics": "ultralytics",
"ultralytics_bbox": "ultralytics",
"ultralytics_segm": "ultralytics",
"yolo": "ultralytics",
"audio_encoder": "audio_encoders",
"audio_encoders": "audio_encoders",
"background_removal": "background_removal",
"background_removal_model": "background_removal",
"frame_interpolation": "frame_interpolation",
"frame_interpolation_model": "frame_interpolation",
"geometry_estimation": "geometry_estimation",
"geometry_estimation_model": "geometry_estimation",
"detection": "detection",
"model_patch": "model_patches",
"model_patches": "model_patches",
"photomaker": "photomaker",
"optical_flow": "optical_flow",
"optical_flow_model": "optical_flow",
"ipadapter": "ipadapter",
"ip_adapter": "ipadapter",
# Legacy/unsupported categories map to checkpoints as fallback
"default": "upscale_models", # Model list uses "default" for upscalers
}
from .settings import CATEGORY_MAP, normalize_download_category
SENSITIVE_METADATA_KEYS = {
"authorization",
@@ -573,22 +519,14 @@ def write_lora_manager_metadata(
) -> Optional[str]:
"""Write the LoRA Manager-compatible sidecar metadata next to a model file."""
metadata_path = get_metadata_sidecar_path(dest_path)
temp_path = f"{metadata_path}.tmp"
try:
payload = build_lora_manager_metadata(dest_path, metadata, category, source_url)
with open(temp_path, "w", encoding="utf-8") as handle:
json.dump(payload, handle, indent=2, ensure_ascii=False)
os.replace(temp_path, metadata_path)
write_json_atomic(metadata_path, payload, indent=2)
log_info(f"Metadata saved: {metadata_path}")
return metadata_path
except Exception as e:
log_warn(f"Could not save metadata sidecar for {dest_path}: {e}")
try:
if os.path.exists(temp_path):
os.remove(temp_path)
except Exception:
pass
return None
@@ -605,21 +543,7 @@ def format_bytes(bytes_value: int) -> str:
return f"{bytes_value:.1f} {sizes[i]}"
def normalize_download_category(category: str) -> str:
"""Return the canonical ComfyUI folder_paths key for a download category."""
token = (
str(category or "")
.strip()
.lower()
.replace("\\", "_")
.replace("/", "_")
.replace("-", "_")
.replace(" ", "_")
)
while "__" in token:
token = token.replace("__", "_")
token = token.strip("_")
return CATEGORY_MAP.get(token, token or "checkpoints")
# Imported from .settings
def get_download_directory(category: str, preferred_base_directory: str = "") -> Optional[str]:
core/path_utils.py
+35 -1
View File
@@ -6,7 +6,7 @@ and directory matching to be shared across all modules.
"""
import os
from typing import Any
from typing import Any, Tuple, Optional
def get_path_identity(path: str) -> str:
@@ -114,3 +114,37 @@ def dedupe_local_base_directories(
by_identity[path_identity] = path_abs
return [by_identity[key] for key in ordered_identities]
def write_json_atomic(
file_path: str,
data: Any,
indent: Optional[int] = None,
separators: Optional[Tuple[str, str]] = None,
) -> None:
"""Safely write data to a JSON file using a temporary file and atomic replace."""
import tempfile
import json
dir_name = os.path.dirname(file_path)
if dir_name:
os.makedirs(dir_name, exist_ok=True)
fd, tmp_path = tempfile.mkstemp(
prefix=f".{os.path.basename(file_path)}.",
suffix=".tmp",
dir=dir_name or None,
text=True,
)
try:
with os.fdopen(fd, "w", encoding="utf-8") as f:
json.dump(data, f, indent=indent, separators=separators, ensure_ascii=False)
if indent is not None:
f.write("\n")
os.replace(tmp_path, file_path)
except Exception:
try:
os.remove(tmp_path)
except OSError:
pass
raise
core/settings.py
+6 -4
View File
@@ -6,6 +6,7 @@ import json
import re
from pathlib import Path
from typing import Any, Dict, Iterable, Mapping, Optional
from .path_utils import write_json_atomic
SETTINGS_FILE = Path(__file__).resolve().parents[1] / "model_resolver_settings.json"
@@ -61,6 +62,10 @@ CATEGORY_MAP = {
"photomaker": "photomaker",
"optical_flow": "optical_flow",
"optical_flow_model": "optical_flow",
"clip_vision": "clip_vision",
"ipadapter": "ipadapter",
"ip_adapter": "ipadapter",
"default": "upscale_models",
}
DEFAULT_DOWNLOAD_PATH_TEMPLATES: Dict[str, str] = {
@@ -263,10 +268,7 @@ def save_settings(payload: Mapping[str, Any]) -> Dict[str, Any]:
current = _read_settings_file()
current.update({str(key): value for key, value in payload.items() if key})
normalized = normalize_settings(current)
SETTINGS_FILE.write_text(
json.dumps(normalized, indent=2, ensure_ascii=False),
encoding="utf-8",
)
write_json_atomic(str(SETTINGS_FILE), normalized, indent=2)
return normalized
core/sources/huggingface.py
+6 -5
View File
@@ -21,6 +21,7 @@ from ..log_system.log_funcs import (
log_error,
log_exception,
)
from ..path_utils import write_json_atomic
HF_API_URL = "https://huggingface.co/api"
HF_AUTHOR_FALLBACKS = ["Comfy-Org"]
@@ -235,13 +236,13 @@ def _read_persistent_author_indexes() -> Dict[str, Any]:
def _write_persistent_author_index(author: str, index: Dict[str, Any]) -> None:
try:
os.makedirs(METADATA_DIR, exist_ok=True)
data = _read_persistent_author_indexes()
data["authors"][author] = index
tmp_path = f"{HF_AUTHOR_INDEX_CACHE_PATH}.tmp"
with open(tmp_path, "w", encoding="utf-8") as f:
json.dump(data, f, ensure_ascii=False, separators=(",", ":"))
os.replace(tmp_path, HF_AUTHOR_INDEX_CACHE_PATH)
write_json_atomic(
HF_AUTHOR_INDEX_CACHE_PATH,
data,
separators=(",", ":"),
)
except Exception as e:
log_debug(f"Error writing HuggingFace author index cache: {e}")
core/sources/model_list.py
+4 -18
View File
@@ -67,25 +67,11 @@ def _read_json_file(path: str, default: Any) -> Any:
return default
from ..path_utils import write_json_atomic
def _write_json_file_atomic(path: str, data: Any):
os.makedirs(os.path.dirname(path), exist_ok=True)
fd, tmp_path = tempfile.mkstemp(
prefix=f".{os.path.basename(path)}.",
suffix=".tmp",
dir=os.path.dirname(path),
text=True,
)
try:
with os.fdopen(fd, "w", encoding="utf-8") as f:
json.dump(data, f, indent=2, ensure_ascii=False)
f.write("\n")
os.replace(tmp_path, path)
except Exception:
try:
os.remove(tmp_path)
except OSError:
pass
raise
write_json_atomic(path, data, indent=2)
def _read_model_list_file() -> Dict[str, Any]: