mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-09-21 11:11:26 -03:00
Will Miao
41302e75ba
The public REST API rewrites files[].name to "{model}_{version}" for
non-LoRA model types, so every precision variant of a multi-file version
shared one name and landed on disk with a random short-hash suffix.
Fetch the raw stored filename from the model-versions/mini endpoint
(always pinned with modelFileId) and use it for the on-disk name and
metadata when available; fall back silently to the REST name otherwise.
CivArchive already serves raw names and is skipped.
942 lines
37 KiB
Python
942 lines
37 KiB
Python
from abc import ABC, abstractmethod
|
|
import asyncio
|
|
import json
|
|
import logging
|
|
import random
|
|
from typing import Optional, Dict, Tuple, Any, List, Sequence
|
|
from .downloader import get_downloader
|
|
from .errors import RateLimitError, ResourceNotFoundError
|
|
|
|
try:
|
|
from bs4 import BeautifulSoup
|
|
except ImportError as exc:
|
|
BeautifulSoup = None # pyright: ignore[reportAssignmentType]
|
|
_BS4_IMPORT_ERROR = exc
|
|
else:
|
|
_BS4_IMPORT_ERROR = None
|
|
|
|
try:
|
|
import aiosqlite
|
|
except ImportError as exc:
|
|
aiosqlite = None # pyright: ignore[reportAssignmentType]
|
|
_AIOSQLITE_IMPORT_ERROR = exc
|
|
else:
|
|
_AIOSQLITE_IMPORT_ERROR = None
|
|
|
|
def _require_beautifulsoup() -> Any:
|
|
if BeautifulSoup is None:
|
|
raise RuntimeError(
|
|
"BeautifulSoup (bs4) is required for CivArchiveModelMetadataProvider. "
|
|
"Install it with 'pip install beautifulsoup4'."
|
|
) from _BS4_IMPORT_ERROR
|
|
return BeautifulSoup
|
|
|
|
def _require_aiosqlite() -> Any:
|
|
if aiosqlite is None:
|
|
raise RuntimeError(
|
|
"aiosqlite is required for SQLiteModelMetadataProvider. "
|
|
"Install it with 'pip install aiosqlite'."
|
|
) from _AIOSQLITE_IMPORT_ERROR
|
|
return aiosqlite
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class _RateLimitRetryHelper:
|
|
"""Coordinate exponential backoff retries after rate limiting."""
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
retry_limit: int = 3,
|
|
base_delay: float = 1.5,
|
|
max_delay: float = 30.0,
|
|
jitter_ratio: float = 0.2,
|
|
) -> None:
|
|
self._retry_limit = max(1, retry_limit)
|
|
self._base_delay = base_delay
|
|
self._max_delay = max_delay
|
|
self._jitter_ratio = max(0.0, jitter_ratio)
|
|
|
|
async def run(self, label: str, func, *args, **kwargs):
|
|
attempt = 0
|
|
while True:
|
|
try:
|
|
return await func(*args, **kwargs)
|
|
except RateLimitError as exc:
|
|
attempt += 1
|
|
|
|
# The downloader's rate-limit gate already applied the wait
|
|
# policy for this request (waited out the vendor window or
|
|
# deliberately refused because it exceeds the cap). Sleeping
|
|
# again here would double the wait — just propagate.
|
|
if getattr(exc, "gate_handled", False):
|
|
exc.provider = exc.provider or label
|
|
raise
|
|
|
|
# Determine effective retry limit based on rate-limit magnitude
|
|
effective_retry_limit = self._retry_limit # default: 3
|
|
if exc.retry_after is not None and exc.retry_after >= 120.0:
|
|
# Long rate-limit window (>=2 min) — retries are futile
|
|
effective_retry_limit = 1 # total 1 attempt = 0 retries
|
|
|
|
if attempt >= effective_retry_limit:
|
|
exc.provider = exc.provider or label
|
|
raise
|
|
|
|
delay = self._calculate_delay(exc.retry_after, attempt)
|
|
logger.warning(
|
|
"Provider %s rate limited request; retrying in %.2fs (attempt %s/%s)",
|
|
label,
|
|
delay,
|
|
attempt,
|
|
self._retry_limit,
|
|
)
|
|
await asyncio.sleep(delay)
|
|
|
|
def _calculate_delay(self, retry_after: Optional[float], attempt: int) -> float:
|
|
if retry_after is not None:
|
|
# Cap at 1800s (30 min) as a safety ceiling. The old 30s cap was
|
|
# too low — CivArchive can return retry_after ~1500s, causing all
|
|
# retries to fail. A generous ceiling protects against pathological
|
|
# server values while still respecting the server's guidance.
|
|
return min(1800.0, max(0.0, retry_after))
|
|
|
|
base_delay = self._base_delay * (2 ** max(0, attempt - 1))
|
|
jitter_span = base_delay * self._jitter_ratio
|
|
if jitter_span > 0:
|
|
base_delay += random.uniform(-jitter_span, jitter_span)
|
|
|
|
return min(self._max_delay, max(0.0, base_delay))
|
|
|
|
|
|
# Labels of providers that are free to consult even while a network provider
|
|
# is rate-limited (local lookups, no vendor cost).
|
|
_LOCAL_PROVIDER_LABELS = frozenset({"sqlite"})
|
|
|
|
|
|
class ModelMetadataProvider(ABC):
|
|
"""Base abstract class for all model metadata providers"""
|
|
|
|
@abstractmethod
|
|
async def get_model_by_hash(self, model_hash: str) -> Tuple[Optional[Dict[str, Any]], Optional[str]]:
|
|
"""Find model by hash value"""
|
|
pass
|
|
|
|
@abstractmethod
|
|
async def get_model_versions(self, model_id: str) -> Optional[Dict[str, Any]]:
|
|
"""Get all versions of a model with their details"""
|
|
pass
|
|
|
|
async def get_model_versions_bulk(
|
|
self, model_ids: Sequence[int]
|
|
) -> Optional[Dict[int, Dict[str, Any]]]:
|
|
"""Fetch model versions for multiple model ids when supported."""
|
|
raise NotImplementedError
|
|
|
|
async def get_model_versions_by_hashes(
|
|
self, hashes: List[str]
|
|
) -> Optional[List[Dict[str, Any]]]:
|
|
"""Fetch full version details for multiple SHA256 hashes.
|
|
|
|
Used specifically to retrieve ``usageControl`` which is only
|
|
available from the per-version / by-hash API, not from model-level
|
|
responses. Providers that cannot resolve hashes should let the
|
|
default ``NotImplementedError`` propagate.
|
|
"""
|
|
raise NotImplementedError
|
|
|
|
@abstractmethod
|
|
async def get_model_version(self, model_id: Optional[int] = None, version_id: Optional[int] = None) -> Optional[Dict[str, Any]]:
|
|
"""Get specific model version with additional metadata"""
|
|
pass
|
|
|
|
@abstractmethod
|
|
async def get_model_version_info(self, version_id: str) -> Tuple[Optional[Dict[str, Any]], Optional[str]]:
|
|
"""Fetch model version metadata"""
|
|
pass
|
|
|
|
@abstractmethod
|
|
async def get_user_models(self, username: str, cursor: Optional[str] = None) -> Optional[Dict[str, Any]]:
|
|
"""Fetch one page of models owned by the specified user.
|
|
|
|
Returns ``{"items": [...], "nextCursor": <str|None>}`` on success,
|
|
or None when unsupported/failed. ``cursor`` continues a previous page.
|
|
"""
|
|
pass
|
|
|
|
async def get_creator_model_count(self, username: str) -> Optional[int]:
|
|
"""Published model count for the user; None when unsupported."""
|
|
return None
|
|
|
|
async def get_version_file_mini(
|
|
self, version_id: int, file_id: int
|
|
) -> Optional[Dict[str, Any]]:
|
|
"""Fetch raw stored file info via CivitAI's model-versions/mini endpoint.
|
|
|
|
Only the CivitAI provider implements this (#1100); other providers
|
|
already serve raw file names (CivArchive) or cannot resolve this
|
|
lookup (SQLite), so the default is None.
|
|
"""
|
|
return None
|
|
|
|
class CivitaiModelMetadataProvider(ModelMetadataProvider):
|
|
"""Provider that uses Civitai API for metadata"""
|
|
|
|
def __init__(self, civitai_client):
|
|
self.client = civitai_client
|
|
|
|
async def get_model_by_hash(self, model_hash: str) -> Tuple[Optional[Dict[str, Any]], Optional[str]]:
|
|
return await self.client.get_model_by_hash(model_hash)
|
|
|
|
async def get_model_versions(self, model_id: str) -> Optional[Dict[str, Any]]:
|
|
return await self.client.get_model_versions(model_id)
|
|
|
|
async def get_model_versions_bulk(
|
|
self, model_ids: Sequence[int]
|
|
) -> Optional[Dict[int, Dict[str, Any]]]:
|
|
return await self.client.get_model_versions_bulk(model_ids)
|
|
|
|
async def get_model_versions_by_hashes(
|
|
self, hashes: List[str]
|
|
) -> Optional[List[Dict[str, Any]]]:
|
|
return await self.client.get_model_versions_by_hashes(hashes)
|
|
|
|
async def get_model_version(self, model_id: Optional[int] = None, version_id: Optional[int] = None) -> Optional[Dict[str, Any]]:
|
|
return await self.client.get_model_version(model_id, version_id)
|
|
|
|
async def get_model_version_info(self, version_id: str) -> Tuple[Optional[Dict[str, Any]], Optional[str]]:
|
|
return await self.client.get_model_version_info(version_id)
|
|
|
|
async def get_user_models(self, username: str, cursor: Optional[str] = None) -> Optional[Dict[str, Any]]:
|
|
return await self.client.get_user_models(username, cursor)
|
|
|
|
async def get_creator_model_count(self, username: str) -> Optional[int]:
|
|
return await self.client.get_creator_model_count(username)
|
|
|
|
async def get_version_file_mini(
|
|
self, version_id: int, file_id: int
|
|
) -> Optional[Dict[str, Any]]:
|
|
return await self.client.get_version_file_mini(version_id, file_id)
|
|
|
|
class CivArchiveModelMetadataProvider(ModelMetadataProvider):
|
|
"""Provider that uses CivArchive API for metadata"""
|
|
|
|
def __init__(self, civarchive_client):
|
|
self.client = civarchive_client
|
|
|
|
async def get_model_by_hash(self, model_hash: str) -> Tuple[Optional[Dict[str, Any]], Optional[str]]:
|
|
return await self.client.get_model_by_hash(model_hash)
|
|
|
|
async def get_model_versions(self, model_id: str) -> Optional[Dict[str, Any]]:
|
|
return await self.client.get_model_versions(model_id)
|
|
|
|
async def get_model_version(self, model_id: Optional[int] = None, version_id: Optional[int] = None) -> Optional[Dict[str, Any]]:
|
|
return await self.client.get_model_version(model_id, version_id)
|
|
|
|
async def get_model_version_info(self, version_id: str) -> Tuple[Optional[Dict[str, Any]], Optional[str]]:
|
|
return await self.client.get_model_version_info(version_id)
|
|
|
|
async def get_user_models(self, username: str, cursor: Optional[str] = None) -> Optional[Dict[str, Any]]:
|
|
"""Not supported by CivArchive provider"""
|
|
return None
|
|
|
|
class SQLiteModelMetadataProvider(ModelMetadataProvider):
|
|
"""Provider that uses SQLite database for metadata"""
|
|
|
|
def __init__(self, db_path: str):
|
|
self.db_path = db_path
|
|
self._aiosqlite = _require_aiosqlite()
|
|
|
|
async def get_model_by_hash(self, model_hash: str) -> Tuple[Optional[Dict[str, Any]], Optional[str]]:
|
|
"""Find model by hash value from SQLite database"""
|
|
async with self._aiosqlite.connect(self.db_path) as db:
|
|
# Look up in model_files table to get model_id and version_id
|
|
query = """
|
|
SELECT model_id, version_id
|
|
FROM model_files
|
|
WHERE sha256 = ?
|
|
LIMIT 1
|
|
"""
|
|
db.row_factory = self._aiosqlite.Row
|
|
cursor = await db.execute(query, (model_hash.upper(),))
|
|
file_row = await cursor.fetchone()
|
|
|
|
if not file_row:
|
|
return None, "Model not found"
|
|
|
|
# Get version details
|
|
model_id = file_row['model_id']
|
|
version_id = file_row['version_id']
|
|
|
|
# Build response in the same format as Civitai API
|
|
result = await self._get_version_with_model_data(db, model_id, version_id)
|
|
return result, None if result else "Error retrieving model data"
|
|
|
|
async def get_model_versions(self, model_id: str) -> Optional[Dict[str, Any]]:
|
|
"""Get all versions of a model from SQLite database"""
|
|
async with self._aiosqlite.connect(self.db_path) as db:
|
|
db.row_factory = self._aiosqlite.Row
|
|
|
|
# First check if model exists
|
|
model_query = "SELECT * FROM models WHERE id = ?"
|
|
cursor = await db.execute(model_query, (model_id,))
|
|
model_row = await cursor.fetchone()
|
|
|
|
if not model_row:
|
|
return None
|
|
|
|
model_data = json.loads(model_row['data'])
|
|
model_type = model_row['type']
|
|
model_name = model_row['name']
|
|
|
|
# Get all versions for this model
|
|
versions_query = """
|
|
SELECT id, name, base_model, data, position, published_at
|
|
FROM model_versions
|
|
WHERE model_id = ?
|
|
ORDER BY position ASC
|
|
"""
|
|
cursor = await db.execute(versions_query, (model_id,))
|
|
version_rows = await cursor.fetchall()
|
|
|
|
if not version_rows:
|
|
return {'modelVersions': [], 'type': model_type}
|
|
|
|
# Format versions similar to Civitai API
|
|
model_versions = []
|
|
for row in version_rows:
|
|
version_data = json.loads(row['data'])
|
|
# Add fields from the row to ensure we have the basic fields
|
|
version_entry = {
|
|
'id': row['id'],
|
|
'modelId': int(model_id),
|
|
'name': row['name'],
|
|
'baseModel': row['base_model'],
|
|
'model': {
|
|
'name': model_row['name'],
|
|
'type': model_type,
|
|
},
|
|
'source': 'archive_db'
|
|
}
|
|
# Update with any additional data
|
|
version_entry.update(version_data)
|
|
model_versions.append(version_entry)
|
|
|
|
return {
|
|
'modelVersions': model_versions,
|
|
'type': model_type,
|
|
'name': model_name
|
|
}
|
|
|
|
async def get_model_version(self, model_id: Optional[int] = None, version_id: Optional[int] = None) -> Optional[Dict[str, Any]]:
|
|
"""Get specific model version with additional metadata from SQLite database"""
|
|
if not model_id and not version_id:
|
|
return None
|
|
|
|
async with self._aiosqlite.connect(self.db_path) as db:
|
|
db.row_factory = self._aiosqlite.Row
|
|
|
|
# Case 1: Only version_id is provided
|
|
if model_id is None and version_id is not None:
|
|
# First get the version info to extract model_id
|
|
version_query = "SELECT model_id FROM model_versions WHERE id = ?"
|
|
cursor = await db.execute(version_query, (version_id,))
|
|
version_row = await cursor.fetchone()
|
|
|
|
if not version_row:
|
|
return None
|
|
|
|
model_id = version_row['model_id']
|
|
|
|
# Case 2: model_id is provided but version_id is not
|
|
elif model_id is not None and version_id is None:
|
|
# Find the latest version
|
|
version_query = """
|
|
SELECT id FROM model_versions
|
|
WHERE model_id = ?
|
|
ORDER BY position ASC
|
|
LIMIT 1
|
|
"""
|
|
cursor = await db.execute(version_query, (model_id,))
|
|
version_row = await cursor.fetchone()
|
|
|
|
if not version_row:
|
|
return None
|
|
|
|
version_id = version_row['id']
|
|
|
|
# Now we have both model_id and version_id, get the full data
|
|
return await self._get_version_with_model_data(db, model_id, version_id)
|
|
|
|
async def get_model_version_info(self, version_id: str) -> Tuple[Optional[Dict[str, Any]], Optional[str]]:
|
|
"""Fetch model version metadata from SQLite database"""
|
|
async with self._aiosqlite.connect(self.db_path) as db:
|
|
db.row_factory = self._aiosqlite.Row
|
|
|
|
# Get version details
|
|
version_query = "SELECT model_id FROM model_versions WHERE id = ?"
|
|
cursor = await db.execute(version_query, (version_id,))
|
|
version_row = await cursor.fetchone()
|
|
|
|
if not version_row:
|
|
return None, "Model version not found"
|
|
|
|
model_id = version_row['model_id']
|
|
|
|
# Build complete version data with model info
|
|
version_data = await self._get_version_with_model_data(db, model_id, version_id)
|
|
return version_data, None
|
|
|
|
async def get_user_models(self, username: str, cursor: Optional[str] = None) -> Optional[Dict[str, Any]]:
|
|
"""Listing models by username is not supported for archive database"""
|
|
return None
|
|
|
|
async def _get_version_with_model_data(self, db, model_id, version_id) -> Optional[Dict[str, Any]]:
|
|
"""Helper to build version data with model information"""
|
|
# Get version details
|
|
version_query = "SELECT name, base_model, data FROM model_versions WHERE id = ? AND model_id = ?"
|
|
cursor = await db.execute(version_query, (version_id, model_id))
|
|
version_row = await cursor.fetchone()
|
|
|
|
if not version_row:
|
|
return None
|
|
|
|
# Get model details
|
|
model_query = "SELECT name, type, data, username FROM models WHERE id = ?"
|
|
cursor = await db.execute(model_query, (model_id,))
|
|
model_row = await cursor.fetchone()
|
|
|
|
if not model_row:
|
|
return None
|
|
|
|
# Parse JSON data
|
|
try:
|
|
version_data = json.loads(version_row['data'])
|
|
model_data = json.loads(model_row['data'])
|
|
|
|
# Build response
|
|
result = {
|
|
"id": int(version_id),
|
|
"modelId": int(model_id),
|
|
"name": version_row['name'],
|
|
"baseModel": version_row['base_model'],
|
|
"model": {
|
|
"name": model_row['name'],
|
|
"description": model_data.get("description"),
|
|
"type": model_row['type'],
|
|
"tags": model_data.get("tags", [])
|
|
},
|
|
"creator": {
|
|
"username": model_row['username'] or model_data.get("creator", {}).get("username"),
|
|
"image": model_data.get("creator", {}).get("image")
|
|
},
|
|
"source": "archive_db"
|
|
}
|
|
|
|
# Add any additional fields from version data
|
|
result.update(version_data)
|
|
|
|
# Attach files associated with this version from model_files table
|
|
files_query = """
|
|
SELECT data
|
|
FROM model_files
|
|
WHERE version_id = ? AND type = 'Model'
|
|
ORDER BY id ASC
|
|
"""
|
|
cursor = await db.execute(files_query, (version_id,))
|
|
file_rows = await cursor.fetchall()
|
|
|
|
files = []
|
|
for file_row in file_rows:
|
|
try:
|
|
file_data = json.loads(file_row['data'])
|
|
except json.JSONDecodeError:
|
|
logger.warning(
|
|
"Skipping model_files entry with invalid JSON for version_id %s", version_id
|
|
)
|
|
continue
|
|
# Remove 'modelId' and 'modelVersionId' fields if present
|
|
file_data.pop('modelId', None)
|
|
file_data.pop('modelVersionId', None)
|
|
files.append(file_data)
|
|
|
|
if 'files' in result:
|
|
existing_files = result['files']
|
|
if isinstance(existing_files, list):
|
|
existing_files.extend(files)
|
|
result['files'] = existing_files
|
|
else:
|
|
merged_files = files.copy()
|
|
if existing_files:
|
|
merged_files.insert(0, existing_files)
|
|
result['files'] = merged_files
|
|
elif files:
|
|
result['files'] = files
|
|
else:
|
|
result['files'] = []
|
|
|
|
return result
|
|
except json.JSONDecodeError:
|
|
return None
|
|
|
|
class FallbackMetadataProvider(ModelMetadataProvider):
|
|
"""Try providers in order, return first successful result.
|
|
|
|
Rate-limit policy (#1085): once a *network* provider raises
|
|
``RateLimitError``, the chain stops consulting further network providers —
|
|
failing over would just spread the flood to the next vendor. Local-only
|
|
providers (see ``_LOCAL_PROVIDER_LABELS``) are still allowed as a last
|
|
resort because they cost the vendor nothing.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
providers: Sequence[ModelMetadataProvider | Tuple[str, ModelMetadataProvider]],
|
|
*,
|
|
rate_limit_retry_limit: int = 3,
|
|
rate_limit_base_delay: float = 1.5,
|
|
rate_limit_max_delay: float = 30.0,
|
|
rate_limit_jitter_ratio: float = 0.2,
|
|
) -> None:
|
|
self.providers: List[ModelMetadataProvider] = []
|
|
self._provider_labels: List[str] = []
|
|
|
|
for entry in providers:
|
|
if isinstance(entry, tuple) and len(entry) == 2:
|
|
name, provider = entry
|
|
else:
|
|
provider = entry
|
|
name = provider.__class__.__name__
|
|
self.providers.append(provider)
|
|
self._provider_labels.append(str(name))
|
|
|
|
self._rate_limit_retry_limit = max(1, rate_limit_retry_limit)
|
|
self._rate_limit_base_delay = rate_limit_base_delay
|
|
self._rate_limit_max_delay = rate_limit_max_delay
|
|
self._rate_limit_jitter_ratio = max(0.0, rate_limit_jitter_ratio)
|
|
self._rate_limit_helper = _RateLimitRetryHelper(
|
|
retry_limit=self._rate_limit_retry_limit,
|
|
base_delay=self._rate_limit_base_delay,
|
|
max_delay=self._rate_limit_max_delay,
|
|
jitter_ratio=self._rate_limit_jitter_ratio,
|
|
)
|
|
|
|
async def get_model_by_hash(self, model_hash: str) -> Tuple[Optional[Dict[str, Any]], Optional[str]]:
|
|
rate_limited = False
|
|
for provider, label in self._iter_providers():
|
|
if rate_limited and label not in _LOCAL_PROVIDER_LABELS:
|
|
continue
|
|
try:
|
|
result, error = await self._call_with_rate_limit(
|
|
label,
|
|
provider.get_model_by_hash,
|
|
model_hash,
|
|
)
|
|
if result:
|
|
return result, error
|
|
except RateLimitError as exc:
|
|
rate_limited = True
|
|
logger.warning(
|
|
"Provider %s is rate-limited (retry_after=%.0fs); not failing over to other network providers",
|
|
label,
|
|
exc.retry_after or 0,
|
|
)
|
|
continue
|
|
except Exception as e:
|
|
logger.debug("Provider %s failed for get_model_by_hash: %s", label, e)
|
|
continue
|
|
if rate_limited:
|
|
# Distinct from "Model not found": callers must not mistake a
|
|
# rate-limited lookup for a confirmed deletion.
|
|
return None, "Rate limited"
|
|
return None, "Model not found"
|
|
|
|
async def get_model_versions(self, model_id: str) -> Optional[Dict[str, Any]]:
|
|
not_found_confirmed = False
|
|
rate_limited = False
|
|
for provider, label in self._iter_providers():
|
|
if rate_limited and label not in _LOCAL_PROVIDER_LABELS:
|
|
continue
|
|
try:
|
|
result = await self._call_with_rate_limit(
|
|
label,
|
|
provider.get_model_versions,
|
|
model_id,
|
|
)
|
|
if result:
|
|
return result
|
|
except RateLimitError as exc:
|
|
rate_limited = True
|
|
logger.warning(
|
|
"Provider %s is rate-limited (retry_after=%.0fs); not failing over to other network providers",
|
|
label,
|
|
exc.retry_after or 0,
|
|
)
|
|
continue
|
|
except ResourceNotFoundError:
|
|
not_found_confirmed = True
|
|
logger.debug(
|
|
"Provider %s reports model %s as not found",
|
|
label,
|
|
model_id,
|
|
)
|
|
continue
|
|
except Exception as e:
|
|
logger.debug("Provider %s failed for get_model_versions: %s", label, e)
|
|
continue
|
|
return None
|
|
|
|
async def get_model_version(self, model_id: Optional[int] = None, version_id: Optional[int] = None) -> Optional[Dict[str, Any]]:
|
|
rate_limited = False
|
|
for provider, label in self._iter_providers():
|
|
if rate_limited and label not in _LOCAL_PROVIDER_LABELS:
|
|
continue
|
|
try:
|
|
result = await self._call_with_rate_limit(
|
|
label,
|
|
provider.get_model_version,
|
|
model_id,
|
|
version_id,
|
|
)
|
|
if result:
|
|
return result
|
|
except RateLimitError as exc:
|
|
rate_limited = True
|
|
logger.warning(
|
|
"Provider %s is rate-limited (retry_after=%.0fs); not failing over to other network providers",
|
|
label,
|
|
exc.retry_after or 0,
|
|
)
|
|
continue
|
|
except Exception as e:
|
|
logger.debug("Provider %s failed for get_model_version: %s", label, e)
|
|
continue
|
|
return None
|
|
|
|
async def get_model_version_info(self, version_id: str) -> Tuple[Optional[Dict[str, Any]], Optional[str]]:
|
|
rate_limited = False
|
|
for provider, label in self._iter_providers():
|
|
if rate_limited and label not in _LOCAL_PROVIDER_LABELS:
|
|
continue
|
|
try:
|
|
result, error = await self._call_with_rate_limit(
|
|
label,
|
|
provider.get_model_version_info,
|
|
version_id,
|
|
)
|
|
if result:
|
|
return result, error
|
|
except RateLimitError as exc:
|
|
rate_limited = True
|
|
logger.warning(
|
|
"Provider %s is rate-limited (retry_after=%.0fs); not failing over to other network providers",
|
|
label,
|
|
exc.retry_after or 0,
|
|
)
|
|
continue
|
|
except Exception as e:
|
|
logger.debug("Provider %s failed for get_model_version_info: %s", label, e)
|
|
continue
|
|
if rate_limited:
|
|
return None, "Rate limited"
|
|
return None, "No provider could retrieve the data"
|
|
|
|
async def get_model_versions_by_hashes(
|
|
self, hashes: List[str]
|
|
) -> Optional[List[Dict[str, Any]]]:
|
|
rate_limited = False
|
|
for provider, label in self._iter_providers():
|
|
if rate_limited and label not in _LOCAL_PROVIDER_LABELS:
|
|
continue
|
|
try:
|
|
result = await self._call_with_rate_limit(
|
|
label,
|
|
provider.get_model_versions_by_hashes,
|
|
hashes,
|
|
)
|
|
if result is not None:
|
|
return result
|
|
except NotImplementedError:
|
|
continue
|
|
except RateLimitError as exc:
|
|
rate_limited = True
|
|
logger.warning(
|
|
"Provider %s is rate-limited (retry_after=%.0fs); not failing over to other network providers",
|
|
label,
|
|
exc.retry_after or 0,
|
|
)
|
|
continue
|
|
except Exception as e:
|
|
logger.debug(
|
|
"Provider %s failed for get_model_versions_by_hashes: %s",
|
|
label,
|
|
e,
|
|
)
|
|
continue
|
|
return None
|
|
|
|
async def get_user_models(self, username: str, cursor: Optional[str] = None) -> Optional[Dict[str, Any]]:
|
|
rate_limited = False
|
|
for provider, label in self._iter_providers():
|
|
if rate_limited and label not in _LOCAL_PROVIDER_LABELS:
|
|
continue
|
|
try:
|
|
result = await self._call_with_rate_limit(
|
|
label,
|
|
provider.get_user_models,
|
|
username,
|
|
cursor=cursor,
|
|
)
|
|
if result is not None:
|
|
return result
|
|
except RateLimitError as exc:
|
|
rate_limited = True
|
|
logger.warning(
|
|
"Provider %s is rate-limited (retry_after=%.0fs); not failing over to other network providers",
|
|
label,
|
|
exc.retry_after or 0,
|
|
)
|
|
continue
|
|
except Exception as e:
|
|
logger.debug("Provider %s failed for get_user_models: %s", label, e)
|
|
continue
|
|
return None
|
|
|
|
async def get_creator_model_count(self, username: str) -> Optional[int]:
|
|
for provider, label in self._iter_providers():
|
|
try:
|
|
result = await provider.get_creator_model_count(username)
|
|
if result is not None:
|
|
return result
|
|
except Exception as e:
|
|
logger.debug(
|
|
"Provider %s failed for get_creator_model_count: %s", label, e
|
|
)
|
|
continue
|
|
return None
|
|
|
|
async def get_version_file_mini(
|
|
self, version_id: int, file_id: int
|
|
) -> Optional[Dict[str, Any]]:
|
|
rate_limited = False
|
|
for provider, label in self._iter_providers():
|
|
if rate_limited and label not in _LOCAL_PROVIDER_LABELS:
|
|
continue
|
|
try:
|
|
result = await self._call_with_rate_limit(
|
|
label,
|
|
provider.get_version_file_mini,
|
|
version_id,
|
|
file_id,
|
|
)
|
|
if result:
|
|
return result
|
|
except RateLimitError as exc:
|
|
rate_limited = True
|
|
logger.warning(
|
|
"Provider %s is rate-limited (retry_after=%.0fs); not failing over to other network providers",
|
|
label,
|
|
exc.retry_after or 0,
|
|
)
|
|
continue
|
|
except Exception as e:
|
|
logger.debug(
|
|
"Provider %s failed for get_version_file_mini: %s", label, e
|
|
)
|
|
continue
|
|
return None
|
|
|
|
def _iter_providers(self):
|
|
return zip(self.providers, self._provider_labels)
|
|
|
|
async def _call_with_rate_limit(self, label: str, func, *args, **kwargs):
|
|
return await self._rate_limit_helper.run(label, func, *args, **kwargs)
|
|
|
|
|
|
class RateLimitRetryingProvider(ModelMetadataProvider):
|
|
"""Adapter that retries individual provider calls after rate limiting."""
|
|
|
|
def __init__(
|
|
self,
|
|
provider: ModelMetadataProvider,
|
|
label: Optional[str] = None,
|
|
*,
|
|
rate_limit_retry_limit: int = 3,
|
|
rate_limit_base_delay: float = 1.5,
|
|
rate_limit_max_delay: float = 30.0,
|
|
rate_limit_jitter_ratio: float = 0.2,
|
|
) -> None:
|
|
self._provider = provider
|
|
self._label = label or provider.__class__.__name__
|
|
self._rate_limit_helper = _RateLimitRetryHelper(
|
|
retry_limit=rate_limit_retry_limit,
|
|
base_delay=rate_limit_base_delay,
|
|
max_delay=rate_limit_max_delay,
|
|
jitter_ratio=rate_limit_jitter_ratio,
|
|
)
|
|
|
|
def __getattr__(self, item):
|
|
return getattr(self._provider, item)
|
|
|
|
async def get_model_by_hash(self, model_hash: str) -> Tuple[Optional[Dict[str, Any]], Optional[str]]:
|
|
return await self._rate_limit_helper.run(
|
|
self._label,
|
|
self._provider.get_model_by_hash,
|
|
model_hash,
|
|
)
|
|
|
|
async def get_model_versions(self, model_id: str) -> Optional[Dict[str, Any]]:
|
|
return await self._rate_limit_helper.run(
|
|
self._label,
|
|
self._provider.get_model_versions,
|
|
model_id,
|
|
)
|
|
|
|
async def get_model_versions_bulk(
|
|
self,
|
|
model_ids: Sequence[int],
|
|
) -> Optional[Dict[int, Dict[str, Any]]]:
|
|
return await self._rate_limit_helper.run(
|
|
self._label,
|
|
self._provider.get_model_versions_bulk,
|
|
model_ids,
|
|
)
|
|
|
|
async def get_model_versions_by_hashes(
|
|
self, hashes: List[str]
|
|
) -> Optional[List[Dict[str, Any]]]:
|
|
return await self._rate_limit_helper.run(
|
|
self._label,
|
|
self._provider.get_model_versions_by_hashes,
|
|
hashes,
|
|
)
|
|
|
|
async def get_model_version(self, model_id: Optional[int] = None, version_id: Optional[int] = None) -> Optional[Dict[str, Any]]:
|
|
return await self._rate_limit_helper.run(
|
|
self._label,
|
|
self._provider.get_model_version,
|
|
model_id,
|
|
version_id,
|
|
)
|
|
|
|
async def get_model_version_info(self, version_id: str) -> Tuple[Optional[Dict[str, Any]], Optional[str]]:
|
|
return await self._rate_limit_helper.run(
|
|
self._label,
|
|
self._provider.get_model_version_info,
|
|
version_id,
|
|
)
|
|
|
|
async def get_user_models(self, username: str, cursor: Optional[str] = None) -> Optional[Dict[str, Any]]:
|
|
return await self._rate_limit_helper.run(
|
|
self._label,
|
|
self._provider.get_user_models,
|
|
username,
|
|
cursor=cursor,
|
|
)
|
|
|
|
async def get_creator_model_count(self, username: str) -> Optional[int]:
|
|
return await self._provider.get_creator_model_count(username)
|
|
|
|
async def get_version_file_mini(
|
|
self, version_id: int, file_id: int
|
|
) -> Optional[Dict[str, Any]]:
|
|
return await self._rate_limit_helper.run(
|
|
self._label,
|
|
self._provider.get_version_file_mini,
|
|
version_id,
|
|
file_id,
|
|
)
|
|
|
|
class ModelMetadataProviderManager:
|
|
"""Manager for selecting and using model metadata providers"""
|
|
|
|
_instance = None
|
|
|
|
@classmethod
|
|
async def get_instance(cls):
|
|
"""Get singleton instance of ModelMetadataProviderManager"""
|
|
if cls._instance is None:
|
|
cls._instance = cls()
|
|
return cls._instance
|
|
|
|
def __init__(self):
|
|
self.providers = {}
|
|
self.default_provider = None
|
|
|
|
def register_provider(self, name: str, provider: ModelMetadataProvider, is_default: bool = False):
|
|
"""Register a metadata provider"""
|
|
self.providers[name] = provider
|
|
if is_default or self.default_provider is None:
|
|
self.default_provider = name
|
|
|
|
async def get_model_by_hash(self, model_hash: str, provider_name: Optional[str] = None) -> Tuple[Optional[Dict[str, Any]], Optional[str]]:
|
|
"""Find model by hash using specified or default provider"""
|
|
provider = self._get_provider(provider_name)
|
|
return await provider.get_model_by_hash(model_hash)
|
|
|
|
async def get_model_versions(self, model_id: str, provider_name: Optional[str] = None) -> Optional[Dict[str, Any]]:
|
|
"""Get model versions using specified or default provider"""
|
|
provider = self._get_provider(provider_name)
|
|
return await provider.get_model_versions(model_id)
|
|
|
|
async def get_model_versions_bulk(
|
|
self,
|
|
model_ids: Sequence[int],
|
|
provider_name: Optional[str] = None,
|
|
) -> Optional[Dict[int, Dict[str, Any]]]:
|
|
"""Fetch model versions for multiple model ids when supported by provider."""
|
|
provider = self._get_provider(provider_name)
|
|
try:
|
|
return await provider.get_model_versions_bulk(model_ids)
|
|
except NotImplementedError:
|
|
return None
|
|
|
|
async def get_model_version(self, model_id: Optional[int] = None, version_id: Optional[int] = None, provider_name: Optional[str] = None) -> Optional[Dict[str, Any]]:
|
|
"""Get specific model version using specified or default provider"""
|
|
provider = self._get_provider(provider_name)
|
|
return await provider.get_model_version(model_id, version_id)
|
|
|
|
async def get_model_version_info(self, version_id: str, provider_name: Optional[str] = None) -> Tuple[Optional[Dict[str, Any]], Optional[str]]:
|
|
"""Fetch model version info using specified or default provider"""
|
|
provider = self._get_provider(provider_name)
|
|
return await provider.get_model_version_info(version_id)
|
|
|
|
async def get_model_versions_by_hashes(
|
|
self,
|
|
hashes: List[str],
|
|
provider_name: Optional[str] = None,
|
|
) -> Optional[List[Dict[str, Any]]]:
|
|
provider = self._get_provider(provider_name)
|
|
try:
|
|
return await provider.get_model_versions_by_hashes(hashes)
|
|
except NotImplementedError:
|
|
return None
|
|
|
|
async def get_user_models(
|
|
self,
|
|
username: str,
|
|
provider_name: Optional[str] = None,
|
|
cursor: Optional[str] = None,
|
|
) -> Optional[Dict[str, Any]]:
|
|
"""Fetch one page of models owned by the specified user"""
|
|
provider = self._get_provider(provider_name)
|
|
return await provider.get_user_models(username, cursor)
|
|
|
|
async def get_creator_model_count(self, username: str, provider_name: Optional[str] = None) -> Optional[int]:
|
|
"""Best-effort published model count for the specified user"""
|
|
provider = self._get_provider(provider_name)
|
|
return await provider.get_creator_model_count(username)
|
|
|
|
def _get_provider(self, provider_name: Optional[str] = None) -> ModelMetadataProvider:
|
|
"""Get provider by name or default provider"""
|
|
if provider_name:
|
|
if provider_name not in self.providers:
|
|
raise ValueError(f"Provider '{provider_name}' is not registered")
|
|
return self.providers[provider_name]
|
|
|
|
if self.default_provider is None:
|
|
raise ValueError("No default provider set and no valid provider specified")
|
|
|
|
return self.providers[self.default_provider]
|