From 8655f925153315de9682ecdb3f647470ecd74683 Mon Sep 17 00:00:00 2001 From: Dariusz L Date: Thu, 25 Jun 2026 16:45:55 +0200 Subject: [PATCH] refactor: centralize category mapping and atomic JSON writing --- __init__.py | 49 +--------------------- core/downloader.py | 84 ++----------------------------------- core/path_utils.py | 36 +++++++++++++++- core/settings.py | 10 +++-- core/sources/huggingface.py | 11 ++--- core/sources/model_list.py | 22 ++-------- 6 files changed, 56 insertions(+), 156 deletions(-) diff --git a/__init__.py b/__init__.py index 9e65d72..084afdc 100644 --- a/__init__.py +++ b/__init__.py @@ -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 ) diff --git a/core/downloader.py b/core/downloader.py index 3831313..8919818 100644 --- a/core/downloader.py +++ b/core/downloader.py @@ -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]: diff --git a/core/path_utils.py b/core/path_utils.py index 8add723..cfa6ed6 100644 --- a/core/path_utils.py +++ b/core/path_utils.py @@ -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 diff --git a/core/settings.py b/core/settings.py index 3127a54..e327b7e 100644 --- a/core/settings.py +++ b/core/settings.py @@ -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 diff --git a/core/sources/huggingface.py b/core/sources/huggingface.py index 716e2f1..4207cba 100644 --- a/core/sources/huggingface.py +++ b/core/sources/huggingface.py @@ -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}") diff --git a/core/sources/model_list.py b/core/sources/model_list.py index 45c098f..12708c2 100644 --- a/core/sources/model_list.py +++ b/core/sources/model_list.py @@ -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]: