diff --git a/core/settings.py b/core/settings.py index ae009be..1579de1 100644 --- a/core/settings.py +++ b/core/settings.py @@ -11,62 +11,9 @@ from .path_utils import write_json_atomic, read_json_safe SETTINGS_FILE = Path(__file__).resolve().parents[1] / "model_resolver_settings.json" -DOWNLOAD_PATH_MODES = {"suggested", "template", "manual"} +from .type_utils import to_bool, CATEGORY_MAP, normalize_download_category -CATEGORY_MAP = { - "checkpoint": "checkpoints", - "checkpoints": "checkpoints", - "lora": "loras", - "loras": "loras", - "embedding": "embeddings", - "embeddings": "embeddings", - "diffusion_model": "diffusion_models", - "diffusion_models": "diffusion_models", - "unet": "diffusion_models", - "clip": "text_encoders", - "clips": "text_encoders", - "text_encoder": "text_encoders", - "text_encoders": "text_encoders", - "controlnet": "controlnet", - "control_net": "controlnet", - "vae": "vae", - "upscaler": "upscale_models", - "upscale_model": "upscale_models", - "upscale_models": "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", - "clip_vision": "clip_vision", - "ipadapter": "ipadapter", - "ip_adapter": "ipadapter", - "default": "upscale_models", -} +DOWNLOAD_PATH_MODES = {"suggested", "template", "manual"} DEFAULT_DOWNLOAD_PATH_TEMPLATES: Dict[str, str] = { "loras": "{base_model}/{first_tag}", @@ -120,30 +67,10 @@ _INVALID_SEGMENT_CHARS_RE = re.compile(r'[<>:"|?*\x00-\x1f]+') _UNKNOWN_PLACEHOLDER_RE = re.compile(r"\{[^{}]+\}") -from .type_utils import to_bool - - def bool_setting(value: Any, default: bool = True) -> bool: return to_bool(value, default) - -def normalize_download_category(category: str) -> str: - 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") - - def sanitize_folder_name(value: Any, fallback: str = "") -> str: text = str(value or "").strip() if not text: diff --git a/core/sources/civarchive.py b/core/sources/civarchive.py index 16a229b..7ab1bc9 100644 --- a/core/sources/civarchive.py +++ b/core/sources/civarchive.py @@ -23,7 +23,7 @@ from ..matcher import ( base_model_score as _base_model_score, calculate_candidate_rank, ) -from ..type_utils import to_int +from ..type_utils import to_int, CIVARCHIVE_API_TYPE_MAP from ..progress import report_progress from ..log_system.log_funcs import ( log_debug, @@ -56,23 +56,7 @@ REQUEST_HEADERS = { ), } -CIVARCHIVE_TYPE_MAP = { - "checkpoint": "Checkpoint", - "checkpoints": "Checkpoint", - "lora": "LORA", - "loras": "LORA", - "locon": "LoCon", - "lycoris": "LoCon", - "vae": "VAE", - "controlnet": "Controlnet", - "embedding": "TextualInversion", - "embeddings": "TextualInversion", - "textualinversion": "TextualInversion", - "upscaler": "Upscaler", - "upscale_models": "Upscaler", - "workflow": "Workflows", - "workflows": "Workflows", -} +# Imported CIVARCHIVE_API_TYPE_MAP class CivArchiveSearchError(Exception): @@ -303,7 +287,7 @@ def _search_page( "rating": "all", "platform": "all", } - civarchive_type = CIVARCHIVE_TYPE_MAP.get(str(model_type or "").lower()) + civarchive_type = CIVARCHIVE_API_TYPE_MAP.get(str(model_type or "").lower()) if civarchive_type: params["type"] = civarchive_type diff --git a/core/sources/civitai.py b/core/sources/civitai.py index d5f9b90..30fddc9 100644 --- a/core/sources/civitai.py +++ b/core/sources/civitai.py @@ -20,7 +20,7 @@ from ..matcher import ( base_model_matches as _base_model_matches, base_model_score as _base_model_score, ) -from ..type_utils import first_non_empty +from ..type_utils import first_non_empty, CIVITAI_API_TYPE_MAP from ..progress import report_progress from ..path_utils import calculate_file_sha256 from ..log_system.log_funcs import ( @@ -1082,24 +1082,11 @@ def search_civitai( """ results = [] - type_map = { - "checkpoint": "Checkpoint", - "checkpoints": "Checkpoint", - "lora": "LORA", - "loras": "LORA", - "vae": "VAE", - "controlnet": "Controlnet", - "embedding": "TextualInversion", - "embeddings": "TextualInversion", - "upscaler": "Upscaler", - "upscale_models": "Upscaler", - } - try: params = {"query": query, "limit": limit, "nsfw": "false"} if model_type: - civitai_type = type_map.get(model_type.lower()) + civitai_type = CIVITAI_API_TYPE_MAP.get(model_type.lower()) if civitai_type: params["types"] = civitai_type diff --git a/core/type_utils.py b/core/type_utils.py index d53d572..2048f37 100644 --- a/core/type_utils.py +++ b/core/type_utils.py @@ -103,3 +103,121 @@ def to_bool(value: Any, default: bool = False) -> bool: return bool(value) +# Map of AIR URN type strings to internal ComfyUI folder_paths categories +URN_TYPE_MAP = { + "checkpoint": "checkpoints", + "lora": "loras", + "vae": "vae", + "upscaler": "upscale_models", + "upscale_model": "upscale_models", + "latent_upscale_model": "latent_upscale_models", + "embedding": "embeddings", + "hypernetwork": "hypernetworks", + "controlnet": "controlnet", + "clip": "text_encoders", + "clip_vision": "clip_vision", + "diffusers": "diffusers", +} + +# General mapping of raw/unnormalized keys to canonical categories +CATEGORY_MAP = { + "checkpoints": "checkpoints", + "checkpoint": "checkpoints", + "loras": "loras", + "lora": "loras", + "embeddings": "embeddings", + "embedding": "embeddings", + "hypernetworks": "hypernetworks", + "hypernetwork": "hypernetworks", + "controlnet": "controlnet", + "control_net": "controlnet", + "vae": "vae", + "upscaler": "upscale_models", + "upscale_model": "upscale_models", + "upscale_models": "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", + "clip_vision": "clip_vision", + "ipadapter": "ipadapter", + "ip_adapter": "ipadapter", + "default": "upscale_models", +} + +# Strict case-sensitive types for CivitAI search API (HTTP 400 on lowercase/mismatch) +CIVITAI_API_TYPE_MAP = { + "checkpoint": "Checkpoint", + "checkpoints": "Checkpoint", + "lora": "LORA", + "loras": "LORA", + "vae": "VAE", + "controlnet": "Controlnet", + "embedding": "TextualInversion", + "embeddings": "TextualInversion", + "upscaler": "Upscaler", + "upscale_models": "Upscaler", +} + +# Case-sensitive types for CivArchive search API (mapping filter support) +CIVARCHIVE_API_TYPE_MAP = { + "checkpoint": "Checkpoint", + "checkpoints": "Checkpoint", + "lora": "LORA", + "loras": "LORA", + "locon": "LoCon", + "lycoris": "LoCon", + "vae": "VAE", + "controlnet": "Controlnet", + "embedding": "TextualInversion", + "embeddings": "TextualInversion", + "textualinversion": "TextualInversion", + "upscaler": "Upscaler", + "upscale_models": "Upscaler", + "workflow": "Workflows", + "workflows": "Workflows", +} + +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") + + + diff --git a/core/workflow_analyzer.py b/core/workflow_analyzer.py index 06c12c8..dfe98f5 100644 --- a/core/workflow_analyzer.py +++ b/core/workflow_analyzer.py @@ -1,4 +1,4 @@ -""" +""" Workflow Analyzer Module Extracts model references from workflow JSON and identifies missing models. @@ -39,21 +39,7 @@ MODEL_EXTENSIONS = { ".gguf", } -# URN Type to ComfyUI category mapping -URN_TYPE_MAP = { - "checkpoint": "checkpoints", - "lora": "loras", - "vae": "vae", - "upscaler": "upscale_models", - "upscale_model": "upscale_models", - "latent_upscale_model": "latent_upscale_models", - "embedding": "embeddings", - "hypernetwork": "hypernetworks", - "controlnet": "controlnet", - "clip": "text_encoders", - "clip_vision": "clip_vision", - "diffusers": "diffusers", -} +from .type_utils import URN_TYPE_MAP URN_REGEX = re.compile(r"^urn:air:([^:]+):([^:]+):([^:]+):(\d+)@(\d+)$")