""" Popular Models Database Curated list of common models with known download URLs. """ import os import re from typing import Any, Dict, List, Optional from ..log_system import create_module_logger log = create_module_logger(__name__) from ..catalog_manager import CatalogManager from ..network_utils import request_source_json from ..path_utils import METADATA_DIR, read_json_safe POPULAR_MODELS_FILE = os.path.join(METADATA_DIR, "popular-models.json") MODEL_ALIASES_FILE = os.path.join(METADATA_DIR, "model-aliases.json") BASE_MODELS_FILE = os.path.join(METADATA_DIR, "base-models.json") BASE_MODELS_META_FILE = os.path.join(METADATA_DIR, "base-models.meta.json") # Cache for loaded data _popular_models_cache: Optional[Dict] = None _model_aliases_cache: Optional[Dict] = None _base_models_aliases_cache: Optional[Dict[str, List[str]]] = None def _load_popular_models() -> Dict[str, Any]: """Load popular models database.""" global _popular_models_cache if _popular_models_cache is not None: return _popular_models_cache data = read_json_safe(POPULAR_MODELS_FILE, {}) _popular_models_cache = data.get("models", {}) return _popular_models_cache def _load_model_aliases() -> Dict[str, List[str]]: """Load model aliases database.""" global _model_aliases_cache if _model_aliases_cache is not None: return _model_aliases_cache data = read_json_safe(MODEL_ALIASES_FILE, {}) _model_aliases_cache = data.get("aliases", {}) return _model_aliases_cache from ..matcher import normalize_base_model as _normalize_base_model base_models_mgr = CatalogManager(BASE_MODELS_FILE, BASE_MODELS_META_FILE, "base_models") _raw_base_models_cache: Optional[List[Dict[str, Any]]] = None def load_raw_base_models() -> List[Dict[str, Any]]: """Loads and caches the raw base models database.""" global _raw_base_models_cache if _raw_base_models_cache is None: try: data = base_models_mgr.read_data() _raw_base_models_cache = data.get("base_models", []) except Exception as e: log.error(f"Error loading raw base models: {e}") _raw_base_models_cache = [] return _raw_base_models_cache def load_base_model_aliases() -> Dict[str, List[str]]: """Load and normalize base model aliases from base-models.json.""" global _base_models_aliases_cache if _base_models_aliases_cache is not None: return _base_models_aliases_cache aliases_dict = {} try: base_models = load_raw_base_models() for model in base_models: name = model.get("name", "") normalized_name = _normalize_base_model(name) if not normalized_name: continue aliases = model.get("aliases", []) # Normalize each alias token and deduplicate normalized_tokens = [] seen_tokens = set() for alias in aliases: if not alias: continue normalized_alias = _normalize_base_model(alias) if normalized_alias and normalized_alias not in seen_tokens: seen_tokens.add(normalized_alias) normalized_tokens.append(normalized_alias) # Ensure normalized name itself is in preferred tokens if not already if normalized_name not in seen_tokens: normalized_tokens.append(normalized_name) aliases_dict[normalized_name] = normalized_tokens except Exception as e: log.error(f"Error loading base model aliases: {e}") # Fallback to hardcoded defaults if file is missing/corrupted if not aliases_dict: aliases_dict = { "zimage": ["zimage", "zimageturbo"], "pony": ["pony", "ponyxl"], "illustrious": ["illustrious", "illustriousxl"], "sdxl10": ["sdxl", "sdxl10", "sdxl100", "sdxl1"], "sd15": ["sd15", "sd1", "stablediffusion15"], "flux1d": ["flux", "flux1", "flux1d", "fluxdev"], "flux1s": ["flux1s", "fluxschnell"], "qwenimage": ["qwenimage"], "hunyuan1": ["hunyuan", "hunyuan1"], "wanvideo": ["wan", "wanvideo"], "noobai": ["noobai"], "hidream": ["hidream"], } _base_models_aliases_cache = aliases_dict return _base_models_aliases_cache def get_popular_model_url(filename: str) -> Optional[Dict[str, Any]]: """ Look up a model filename in the popular models database. Args: filename: Model filename to look up Returns: Dictionary with url, type, directory if found, None otherwise """ models = _load_popular_models() # Direct lookup if filename in models: return models[filename].copy() # Try lowercase filename_lower = filename.lower() for name, info in models.items(): if name.lower() == filename_lower: return info.copy() # Try aliases aliases = _load_model_aliases() for canonical, alias_list in aliases.items(): if ( filename in alias_list or filename_lower in [a.lower() for a in alias_list] ) and canonical in models: result = models[canonical].copy() result["canonical_name"] = canonical return result return None def search_popular_models(query: str, limit: int = 10) -> List[Dict[str, Any]]: """ Search popular models database by filename pattern. Args: query: Search query (partial filename) limit: Maximum results to return Returns: List of matching models with url info """ models = _load_popular_models() query_lower = query.lower() results = [] for name, info in models.items(): if query_lower in name.lower(): result = info.copy() result["filename"] = name results.append(result) if len(results) >= limit: break return results def reload_databases(): """Force reload of all databases.""" global _popular_models_cache, _model_aliases_cache, _base_models_aliases_cache, _raw_base_models_cache _popular_models_cache = None _model_aliases_cache = None _base_models_aliases_cache = None _raw_base_models_cache = None from ..path_utils import clear_base_model_alias_cache clear_base_model_alias_cache() _load_popular_models() _load_model_aliases() load_base_model_aliases() def get_base_models_config() -> Dict[str, Any]: """Return the raw base-models config (base_models array) for the frontend dropdown.""" return _read_base_models_file() def _read_base_models_file() -> Dict[str, Any]: return base_models_mgr.read_data() def _read_base_models_meta() -> Dict[str, Any]: return base_models_mgr.read_meta() def generate_aliases(name: str) -> List[str]: aliases = {name.lower()} # Replace dots, dashes, underscores with spaces, then normalize spaces normalized = re.sub(r"[\._\-]+", " ", name.lower()) normalized = re.sub(r"\s+", " ", normalized).strip() aliases.add(normalized) version_trimmed = re.sub(r"(\d+)(?:[\s\._\-]+0)+(?!\d)", r"\1", normalized) aliases.add(version_trimmed) # Remove all spaces/dashes/dots/underscores collapsed = re.sub(r"[\s\._\-]+", "", name.lower()) aliases.add(collapsed) aliases.add(re.sub(r"[\s\._\-]+", "", version_trimmed)) # If there are dots/dashes/underscores, add specific clean replacements if "." in name: aliases.add(name.lower().replace(".", "")) aliases.add(name.lower().replace(".", " ")) if "-" in name: aliases.add(name.lower().replace("-", "")) aliases.add(name.lower().replace("-", " ")) if "_": aliases.add(name.lower().replace("_", "")) aliases.add(name.lower().replace("_", " ")) # Extract first word if it has multiple words words = name.lower().split() if len(words) > 1: first_word = words[0] if len(first_word) > 2 and first_word not in {"stable", "flux", "pony", "sdxl"}: aliases.add(first_word) return sorted(list({a for a in aliases if len(a) > 1})) def get_base_models_status(check_remote: bool = False) -> Dict[str, Any]: """Return local base-models metadata and optionally compare with CivitAI.""" local_data = _read_base_models_file() base_models = local_data.get("base_models", []) meta = _read_base_models_meta() local_count = len(base_models) local_updated_at = meta.get("updated_at") or "" status = { "local_count": local_count, "local_updated_at": local_updated_at, "remote_checked_at": meta.get("last_checked_at") or "", "update_available": False, } if check_remote: try: # Fetch from CivitAI API using unified helper enums = request_source_json("https://civitai.com/api/v1/enums", timeout=15, log_name="CivitAI Enums") if enums: remote_models = enums.get("BaseModel", []) # Check if there are new models by comparing normalized names and aliases all_existing_normalized = set() for m in base_models: all_existing_normalized.add(_normalize_base_model(m.get("name", ""))) for alias in m.get("aliases", []): all_existing_normalized.add(_normalize_base_model(alias)) new_models_found = False for remote_name in remote_models: norm_remote = _normalize_base_model(remote_name) if norm_remote and norm_remote not in all_existing_normalized: new_models_found = True break from ..type_utils import utc_now_iso status.update({ "remote_checked_at": utc_now_iso(), "update_available": new_models_found }) except Exception as e: log.warning(f"Error checking remote base models: {e}") return status def update_base_models_from_remote() -> Dict[str, Any]: """Fetch live BaseModel enums from CivitAI and merge into base-models.json.""" enums = request_source_json("https://civitai.com/api/v1/enums", timeout=15, log_name="CivitAI Enums", raise_on_error=True) if not enums: raise ValueError("CivitAI enums did not return a valid response") remote_names = enums.get("BaseModel", []) if not remote_names or not isinstance(remote_names, list): raise ValueError("CivitAI enums did not return a valid BaseModel list") local_data = _read_base_models_file() base_models = local_data.get("base_models", []) # 1. Build a set of all normalized aliases currently in base-models.json all_known_normalized = set() for m in base_models: for alias in m.get("aliases", []): all_known_normalized.add(_normalize_base_model(alias)) all_known_normalized.add(_normalize_base_model(m.get("name", ""))) updated_models = list(base_models) new_added_count = 0 new_added_names = [] for name in remote_names: norm_name = _normalize_base_model(name) if not norm_name: continue # If the normalized name is already a known alias or name, skip it if norm_name in all_known_normalized: continue # Generate candidate aliases candidates = generate_aliases(name) # Filter out aliases that are already registered/known filtered_aliases = [] for alias in candidates: norm_alias = _normalize_base_model(alias) if norm_alias not in all_known_normalized: filtered_aliases.append(alias) all_known_normalized.add(norm_alias) # If we have valid aliases, add the new base model entry if filtered_aliases: if name.lower() not in [a.lower() for a in filtered_aliases]: filtered_aliases.insert(0, name) all_known_normalized.add(norm_name) updated_models.append({ "name": name, "aliases": filtered_aliases }) new_added_count += 1 new_added_names.append(name) from ..type_utils import utc_now_iso now_str = utc_now_iso() meta = { "updated_at": now_str, "last_checked_at": now_str, "local_count": len(updated_models), "new_models_added": new_added_count, "new_models_added_list": new_added_names, } base_models_mgr.sync_catalog({"base_models": updated_models}, meta, indent=2) reload_databases() return { "local_count": len(updated_models), "local_updated_at": now_str, "new_models_added": new_added_count, "new_models_added_list": new_added_names, "updated": True }