refactor: centralize type maps and normalize_download_category in type_utils
Move CATEGORY_MAP, normalize_download_category, URN_TYPE_MAP, CIVITAI_API_TYPE_MAP, and CIVARCHIVE_API_TYPE_MAP into type_utils.py so they are defined once and imported by settings, civitai, civarchive, and workflow_analyzer instead of being duplicated across modules.
This commit is contained in:
Dariusz L
+2
-75
@@ -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:
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
+2
-15
@@ -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
|
||||
|
||||
|
||||
@@ -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")
|
||||
|
||||
|
||||
|
||||
|
||||
@@ -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+)$")
|
||||
|
||||
|
||||
Reference in New Issue
Block a user