mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-03-29 08:58:53 -03:00
feat: Dynamic base model fetching from Civitai API (#854)
Implement automatic fetching of base models from Civitai API to keep data up-to-date without manual updates. Backend: - Add CivitaiBaseModelService with 7-day TTL caching - Add /api/lm/base-models endpoints for fetching and refreshing - Merge hardcoded and remote models for backward compatibility - Smart abbreviation generation for unknown models Frontend: - Add civitaiBaseModelApi client for API communication - Dynamic base model loading on app initialization - Update SettingsManager to use merged model lists - Add support for 8 new models: Anima, CogVideoX, LTXV 2.3, Mochi, Pony V7, Wan Video 2.5 T2V/I2V API Endpoints: - GET /api/lm/base-models - Get merged models - POST /api/lm/base-models/refresh - Force refresh - GET /api/lm/base-models/categories - Get categories - GET /api/lm/base-models/cache-status - Check cache status Closes #854
This commit is contained in:
Will Miao
430
py/services/civitai_base_model_service.py
Normal file
430
py/services/civitai_base_model_service.py
Normal file
@@ -0,0 +1,430 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import re
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any, Dict, List, Optional, Set, Tuple
|
||||
|
||||
from ..utils.constants import SUPPORTED_DOWNLOAD_SKIP_BASE_MODELS
|
||||
from .downloader import get_downloader
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class CivitaiBaseModelService:
|
||||
"""Service for fetching and managing Civitai base models.
|
||||
|
||||
This service provides:
|
||||
- Fetching base models from Civitai API
|
||||
- Caching with TTL (7 days default)
|
||||
- Merging hardcoded and remote base models
|
||||
- Generating abbreviations for new/unknown models
|
||||
"""
|
||||
|
||||
_instance: Optional[CivitaiBaseModelService] = None
|
||||
_lock = asyncio.Lock()
|
||||
|
||||
# Default TTL for cache in seconds (7 days)
|
||||
DEFAULT_CACHE_TTL = 7 * 24 * 60 * 60
|
||||
|
||||
# Civitai API endpoint for enums
|
||||
CIVITAI_ENUMS_URL = "https://civitai.com/api/v1/enums"
|
||||
|
||||
@classmethod
|
||||
async def get_instance(cls) -> CivitaiBaseModelService:
|
||||
"""Get singleton instance of the service."""
|
||||
async with cls._lock:
|
||||
if cls._instance is None:
|
||||
cls._instance = cls()
|
||||
return cls._instance
|
||||
|
||||
def __init__(self):
|
||||
"""Initialize the service."""
|
||||
if hasattr(self, "_initialized"):
|
||||
return
|
||||
self._initialized = True
|
||||
|
||||
# Cache storage
|
||||
self._cache: Optional[Dict[str, Any]] = None
|
||||
self._cache_timestamp: Optional[datetime] = None
|
||||
self._cache_ttl = self.DEFAULT_CACHE_TTL
|
||||
|
||||
# Hardcoded models for fallback
|
||||
self._hardcoded_models = set(SUPPORTED_DOWNLOAD_SKIP_BASE_MODELS)
|
||||
|
||||
logger.info("CivitaiBaseModelService initialized")
|
||||
|
||||
async def get_base_models(self, force_refresh: bool = False) -> Dict[str, Any]:
|
||||
"""Get merged base models (hardcoded + remote).
|
||||
|
||||
Args:
|
||||
force_refresh: If True, fetch from API regardless of cache state.
|
||||
|
||||
Returns:
|
||||
Dictionary containing:
|
||||
- models: List of merged base model names
|
||||
- source: 'cache', 'api', or 'fallback'
|
||||
- last_updated: ISO timestamp of last successful API fetch
|
||||
- hardcoded_count: Number of hardcoded models
|
||||
- remote_count: Number of remote models
|
||||
- merged_count: Total unique models
|
||||
"""
|
||||
# Check if cache is valid
|
||||
if not force_refresh and self._is_cache_valid():
|
||||
logger.debug("Returning cached base models")
|
||||
return self._build_response("cache")
|
||||
|
||||
# Try to fetch from API
|
||||
try:
|
||||
remote_models = await self._fetch_from_civitai()
|
||||
if remote_models:
|
||||
self._update_cache(remote_models)
|
||||
return self._build_response("api")
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to fetch base models from Civitai: {e}")
|
||||
|
||||
# Fallback to hardcoded models
|
||||
return self._build_response("fallback")
|
||||
|
||||
async def refresh_cache(self) -> Dict[str, Any]:
|
||||
"""Force refresh the cache from Civitai API.
|
||||
|
||||
Returns:
|
||||
Response dict same as get_base_models()
|
||||
"""
|
||||
return await self.get_base_models(force_refresh=True)
|
||||
|
||||
def get_cache_status(self) -> Dict[str, Any]:
|
||||
"""Get current cache status.
|
||||
|
||||
Returns:
|
||||
Dictionary containing:
|
||||
- has_cache: Whether cache exists
|
||||
- last_updated: ISO timestamp or None
|
||||
- is_expired: Whether cache is expired
|
||||
- ttl_seconds: TTL in seconds
|
||||
- age_seconds: Age of cache in seconds (if exists)
|
||||
"""
|
||||
if self._cache is None or self._cache_timestamp is None:
|
||||
return {
|
||||
"has_cache": False,
|
||||
"last_updated": None,
|
||||
"is_expired": True,
|
||||
"ttl_seconds": self._cache_ttl,
|
||||
"age_seconds": None,
|
||||
}
|
||||
|
||||
age = (datetime.now(timezone.utc) - self._cache_timestamp).total_seconds()
|
||||
return {
|
||||
"has_cache": True,
|
||||
"last_updated": self._cache_timestamp.isoformat(),
|
||||
"is_expired": age > self._cache_ttl,
|
||||
"ttl_seconds": self._cache_ttl,
|
||||
"age_seconds": int(age),
|
||||
}
|
||||
|
||||
def generate_abbreviation(self, model_name: str) -> str:
|
||||
"""Generate abbreviation for a base model name.
|
||||
|
||||
Algorithm:
|
||||
1. Extract version patterns (e.g., "2.5" from "Wan Video 2.5")
|
||||
2. Extract main acronym (e.g., "SD" from "SD 1.5")
|
||||
3. Handle special cases (Flux, Wan, etc.)
|
||||
4. Fallback to first letters of words (max 4 chars)
|
||||
|
||||
Args:
|
||||
model_name: Full base model name
|
||||
|
||||
Returns:
|
||||
Generated abbreviation (max 4 characters)
|
||||
"""
|
||||
if not model_name or not isinstance(model_name, str):
|
||||
return "OTH"
|
||||
|
||||
name = model_name.strip()
|
||||
if not name:
|
||||
return "OTH"
|
||||
|
||||
# Check if it's already in hardcoded abbreviations
|
||||
# This is a simplified check - in practice you'd have a mapping
|
||||
lower_name = name.lower()
|
||||
|
||||
# Special cases
|
||||
special_cases = {
|
||||
"sd 1.4": "SD1",
|
||||
"sd 1.5": "SD1",
|
||||
"sd 1.5 lcm": "SD1",
|
||||
"sd 1.5 hyper": "SD1",
|
||||
"sd 2.0": "SD2",
|
||||
"sd 2.1": "SD2",
|
||||
"sd 3": "SD3",
|
||||
"sd 3.5": "SD3",
|
||||
"sd 3.5 medium": "SD3",
|
||||
"sd 3.5 large": "SD3",
|
||||
"sd 3.5 large turbo": "SD3",
|
||||
"sdxl 1.0": "XL",
|
||||
"sdxl lightning": "XL",
|
||||
"sdxl hyper": "XL",
|
||||
"flux.1 d": "F1D",
|
||||
"flux.1 s": "F1S",
|
||||
"flux.1 krea": "F1KR",
|
||||
"flux.1 kontext": "F1KX",
|
||||
"flux.2 d": "F2D",
|
||||
"flux.2 klein 9b": "FK9",
|
||||
"flux.2 klein 9b-base": "FK9B",
|
||||
"flux.2 klein 4b": "FK4",
|
||||
"flux.2 klein 4b-base": "FK4B",
|
||||
"auraflow": "AF",
|
||||
"chroma": "CHR",
|
||||
"pixart a": "PXA",
|
||||
"pixart e": "PXE",
|
||||
"hunyuan 1": "HY",
|
||||
"hunyuan video": "HYV",
|
||||
"lumina": "L",
|
||||
"kolors": "KLR",
|
||||
"noobai": "NAI",
|
||||
"illustrious": "IL",
|
||||
"pony": "PONY",
|
||||
"pony v7": "PNY7",
|
||||
"hidream": "HID",
|
||||
"qwen": "QWEN",
|
||||
"zimageturbo": "ZIT",
|
||||
"zimagebase": "ZIB",
|
||||
"anima": "ANI",
|
||||
"svd": "SVD",
|
||||
"ltxv": "LTXV",
|
||||
"ltxv2": "LTV2",
|
||||
"ltxv 2.3": "LTX",
|
||||
"cogvideox": "CVX",
|
||||
"mochi": "MCHI",
|
||||
"wan video": "WAN",
|
||||
"wan video 1.3b t2v": "WAN",
|
||||
"wan video 14b t2v": "WAN",
|
||||
"wan video 14b i2v 480p": "WAN",
|
||||
"wan video 14b i2v 720p": "WAN",
|
||||
"wan video 2.2 ti2v-5b": "WAN",
|
||||
"wan video 2.2 t2v-a14b": "WAN",
|
||||
"wan video 2.2 i2v-a14b": "WAN",
|
||||
"wan video 2.5 t2v": "WAN",
|
||||
"wan video 2.5 i2v": "WAN",
|
||||
}
|
||||
|
||||
if lower_name in special_cases:
|
||||
return special_cases[lower_name]
|
||||
|
||||
# Try to extract acronym from version pattern
|
||||
# e.g., "Model Name 2.5" -> "MN25"
|
||||
version_match = re.search(r"(\d+(?:\.\d+)?)", name)
|
||||
version = version_match.group(1) if version_match else ""
|
||||
|
||||
# Remove version and common words
|
||||
words = re.sub(r"\d+(?:\.\d+)?", "", name)
|
||||
words = re.sub(
|
||||
r"\b(model|video|diffusion|checkpoint|textualinversion)\b",
|
||||
"",
|
||||
words,
|
||||
flags=re.I,
|
||||
)
|
||||
words = words.strip()
|
||||
|
||||
# Get first letters of remaining words
|
||||
tokens = re.findall(r"[A-Za-z]+", words)
|
||||
if tokens:
|
||||
# Build abbreviation from first letters
|
||||
abbrev = "".join(token[0].upper() for token in tokens)
|
||||
# Add version if present
|
||||
if version:
|
||||
# Clean version (remove dots for abbreviation)
|
||||
version_clean = version.replace(".", "")
|
||||
abbrev = abbrev[: 4 - len(version_clean)] + version_clean
|
||||
return abbrev[:4]
|
||||
|
||||
# Final fallback: just take first 4 alphanumeric chars
|
||||
alphanumeric = re.sub(r"[^A-Za-z0-9]", "", name)
|
||||
if alphanumeric:
|
||||
return alphanumeric[:4].upper()
|
||||
|
||||
return "OTH"
|
||||
|
||||
async def _fetch_from_civitai(self) -> Optional[Set[str]]:
|
||||
"""Fetch base models from Civitai API.
|
||||
|
||||
Returns:
|
||||
Set of base model names, or None if failed
|
||||
"""
|
||||
try:
|
||||
downloader = await get_downloader()
|
||||
success, result = await downloader.make_request(
|
||||
"GET",
|
||||
self.CIVITAI_ENUMS_URL,
|
||||
use_auth=False, # enums endpoint doesn't require auth
|
||||
)
|
||||
|
||||
if not success:
|
||||
logger.warning(f"Failed to fetch enums from Civitai: {result}")
|
||||
return None
|
||||
|
||||
if isinstance(result, str):
|
||||
data = json.loads(result)
|
||||
else:
|
||||
data = result
|
||||
|
||||
# Extract base models from response
|
||||
base_models = set()
|
||||
|
||||
# Use ActiveBaseModel if available (recommended active models)
|
||||
if "ActiveBaseModel" in data:
|
||||
base_models.update(data["ActiveBaseModel"])
|
||||
logger.info(f"Fetched {len(base_models)} models from ActiveBaseModel")
|
||||
# Fallback to full BaseModel list
|
||||
elif "BaseModel" in data:
|
||||
base_models.update(data["BaseModel"])
|
||||
logger.info(f"Fetched {len(base_models)} models from BaseModel")
|
||||
else:
|
||||
logger.warning("No base model data found in Civitai response")
|
||||
return None
|
||||
|
||||
return base_models
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error fetching from Civitai: {e}")
|
||||
return None
|
||||
|
||||
def _update_cache(self, remote_models: Set[str]) -> None:
|
||||
"""Update internal cache with fetched models.
|
||||
|
||||
Args:
|
||||
remote_models: Set of base model names from API
|
||||
"""
|
||||
self._cache = {
|
||||
"remote_models": sorted(remote_models),
|
||||
"hardcoded_models": sorted(self._hardcoded_models),
|
||||
}
|
||||
self._cache_timestamp = datetime.now(timezone.utc)
|
||||
logger.info(f"Cache updated with {len(remote_models)} remote models")
|
||||
|
||||
def _is_cache_valid(self) -> bool:
|
||||
"""Check if current cache is valid (not expired).
|
||||
|
||||
Returns:
|
||||
True if cache exists and is not expired
|
||||
"""
|
||||
if self._cache is None or self._cache_timestamp is None:
|
||||
return False
|
||||
|
||||
age = (datetime.now(timezone.utc) - self._cache_timestamp).total_seconds()
|
||||
return age <= self._cache_ttl
|
||||
|
||||
def _build_response(self, source: str) -> Dict[str, Any]:
|
||||
"""Build response dictionary.
|
||||
|
||||
Args:
|
||||
source: 'cache', 'api', or 'fallback'
|
||||
|
||||
Returns:
|
||||
Response dictionary
|
||||
"""
|
||||
if source == "fallback" or self._cache is None:
|
||||
# Use only hardcoded models
|
||||
merged = sorted(self._hardcoded_models)
|
||||
return {
|
||||
"models": merged,
|
||||
"source": source,
|
||||
"last_updated": None,
|
||||
"hardcoded_count": len(self._hardcoded_models),
|
||||
"remote_count": 0,
|
||||
"merged_count": len(merged),
|
||||
}
|
||||
|
||||
# Merge hardcoded and remote models
|
||||
remote_set = set(self._cache.get("remote_models", []))
|
||||
merged = sorted(self._hardcoded_models | remote_set)
|
||||
|
||||
return {
|
||||
"models": merged,
|
||||
"source": source,
|
||||
"last_updated": self._cache_timestamp.isoformat()
|
||||
if self._cache_timestamp
|
||||
else None,
|
||||
"hardcoded_count": len(self._hardcoded_models),
|
||||
"remote_count": len(remote_set),
|
||||
"merged_count": len(merged),
|
||||
}
|
||||
|
||||
def get_model_categories(self) -> Dict[str, List[str]]:
|
||||
"""Get categorized base models.
|
||||
|
||||
Returns:
|
||||
Dictionary mapping category names to lists of model names
|
||||
"""
|
||||
# Define category patterns
|
||||
categories = {
|
||||
"Stable Diffusion 1.x": ["SD 1.4", "SD 1.5", "SD 1.5 LCM", "SD 1.5 Hyper"],
|
||||
"Stable Diffusion 2.x": ["SD 2.0", "SD 2.1"],
|
||||
"Stable Diffusion 3.x": [
|
||||
"SD 3",
|
||||
"SD 3.5",
|
||||
"SD 3.5 Medium",
|
||||
"SD 3.5 Large",
|
||||
"SD 3.5 Large Turbo",
|
||||
],
|
||||
"SDXL": ["SDXL 1.0", "SDXL Lightning", "SDXL Hyper"],
|
||||
"Flux Models": [
|
||||
"Flux.1 D",
|
||||
"Flux.1 S",
|
||||
"Flux.1 Krea",
|
||||
"Flux.1 Kontext",
|
||||
"Flux.2 D",
|
||||
"Flux.2 Klein 9B",
|
||||
"Flux.2 Klein 9B-base",
|
||||
"Flux.2 Klein 4B",
|
||||
"Flux.2 Klein 4B-base",
|
||||
],
|
||||
"Video Models": [
|
||||
"SVD",
|
||||
"LTXV",
|
||||
"LTXV2",
|
||||
"LTXV 2.3",
|
||||
"CogVideoX",
|
||||
"Mochi",
|
||||
"Hunyuan Video",
|
||||
"Wan Video",
|
||||
"Wan Video 1.3B t2v",
|
||||
"Wan Video 14B t2v",
|
||||
"Wan Video 14B i2v 480p",
|
||||
"Wan Video 14B i2v 720p",
|
||||
"Wan Video 2.2 TI2V-5B",
|
||||
"Wan Video 2.2 T2V-A14B",
|
||||
"Wan Video 2.2 I2V-A14B",
|
||||
"Wan Video 2.5 T2V",
|
||||
"Wan Video 2.5 I2V",
|
||||
],
|
||||
"Other Models": [
|
||||
"Illustrious",
|
||||
"Pony",
|
||||
"Pony V7",
|
||||
"HiDream",
|
||||
"Qwen",
|
||||
"AuraFlow",
|
||||
"Chroma",
|
||||
"ZImageTurbo",
|
||||
"ZImageBase",
|
||||
"PixArt a",
|
||||
"PixArt E",
|
||||
"Hunyuan 1",
|
||||
"Lumina",
|
||||
"Kolors",
|
||||
"NoobAI",
|
||||
"Anima",
|
||||
],
|
||||
}
|
||||
|
||||
return categories
|
||||
|
||||
|
||||
# Convenience function for getting the singleton instance
|
||||
async def get_civitai_base_model_service() -> CivitaiBaseModelService:
|
||||
"""Get the singleton instance of CivitaiBaseModelService."""
|
||||
return await CivitaiBaseModelService.get_instance()
|
||||
Reference in New Issue
Block a user