""" Popular Models Database Curated list of common models with known download URLs. """ import os import json import re import shutil import tempfile from datetime import datetime, timezone from typing import Dict, Any, Optional, List from ..log_system.log_funcs import ( log_debug, log_info, log_warn, log_error, log_exception, ) # Path to metadata directory METADATA_DIR = os.path.join( os.path.dirname(os.path.dirname(os.path.dirname(__file__))), "metadata" ) 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") # 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 try: if os.path.exists(POPULAR_MODELS_FILE): with open(POPULAR_MODELS_FILE, "r", encoding="utf-8") as f: data = json.load(f) _popular_models_cache = data.get("models", {}) return _popular_models_cache except Exception as e: log_error(f"Error loading popular models: {e}") _popular_models_cache = {} 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 try: if os.path.exists(MODEL_ALIASES_FILE): with open(MODEL_ALIASES_FILE, "r", encoding="utf-8") as f: data = json.load(f) _model_aliases_cache = data.get("aliases", {}) return _model_aliases_cache except Exception as e: log_error(f"Error loading model aliases: {e}") _model_aliases_cache = {} return _model_aliases_cache def _normalize_base_model(value: str) -> str: import re return re.sub(r"[^a-z0-9]+", "", str(value or "").lower()) 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: if os.path.exists(BASE_MODELS_FILE): with open(BASE_MODELS_FILE, "r", encoding="utf-8") as f: data = json.load(f) base_models = data.get("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]: if 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 get_all_popular_models() -> Dict[str, Any]: """Get all popular models.""" return _load_popular_models().copy() def reload_databases(): """Force reload of all databases.""" global _popular_models_cache, _model_aliases_cache, _base_models_aliases_cache _popular_models_cache = None _model_aliases_cache = None _base_models_aliases_cache = None _load_popular_models() _load_model_aliases() load_base_model_aliases() BASE_MODELS_META_FILE = os.path.join(METADATA_DIR, "base-models.meta.json") 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]: try: if os.path.exists(BASE_MODELS_FILE): with open(BASE_MODELS_FILE, "r", encoding="utf-8") as f: return json.load(f) except Exception as e: log_warn(f"Error reading base-models.json: {e}") return {"base_models": []} def _read_base_models_meta() -> Dict[str, Any]: try: if os.path.exists(BASE_MODELS_META_FILE): with open(BASE_MODELS_META_FILE, "r", encoding="utf-8") as f: return json.load(f) except Exception as e: log_warn(f"Error reading base-models.meta.json: {e}") return {} 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) # Remove all spaces/dashes/dots/underscores collapsed = re.sub(r"[\s\._\-]+", "", name.lower()) aliases.add(collapsed) # 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: import requests # Fetch from CivitAI API response = requests.get("https://civitai.com/api/v1/enums", timeout=15) if response.status_code == 200: enums = response.json() 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 status.update({ "remote_checked_at": datetime.now(timezone.utc).replace(microsecond=0).isoformat(), "update_available": new_models_found }) except Exception as e: log_warn(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.""" import requests response = requests.get("https://civitai.com/api/v1/enums", timeout=15) if response.status_code != 200: raise ValueError(f"CivitAI enums endpoint returned HTTP {response.status_code}") enums = response.json() 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 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 # Write updated models to base-models.json atomically os.makedirs(METADATA_DIR, exist_ok=True) # Backup existing if os.path.exists(BASE_MODELS_FILE): shutil.copy2(BASE_MODELS_FILE, f"{BASE_MODELS_FILE}.bak") # Write atomically fd, tmp_path = tempfile.mkstemp( prefix=f".{os.path.basename(BASE_MODELS_FILE)}.", suffix=".tmp", dir=os.path.dirname(BASE_MODELS_FILE), text=True, ) try: with os.fdopen(fd, "w", encoding="utf-8") as f: json.dump({"base_models": updated_models}, f, indent=2, ensure_ascii=False) f.write("\n") os.replace(tmp_path, BASE_MODELS_FILE) except Exception: try: os.remove(tmp_path) except OSError: pass raise # Write meta file now_str = datetime.now(timezone.utc).replace(microsecond=0).isoformat() meta = { "updated_at": now_str, "last_checked_at": now_str, "local_count": len(updated_models), "new_models_added": new_added_count } fd_meta, tmp_meta_path = tempfile.mkstemp( prefix=f".{os.path.basename(BASE_MODELS_META_FILE)}.", suffix=".tmp", dir=os.path.dirname(BASE_MODELS_META_FILE), text=True, ) try: with os.fdopen(fd_meta, "w", encoding="utf-8") as f: json.dump(meta, f, indent=2, ensure_ascii=False) f.write("\n") os.replace(tmp_meta_path, BASE_MODELS_META_FILE) except Exception: try: os.remove(tmp_meta_path) except OSError: pass raise # Force reload database in-memory cache reload_databases() return { "local_count": len(updated_models), "local_updated_at": now_str, "new_models_added": new_added_count, "updated": True }