mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-10-04 00:55:32 -03:00
feat(metadata): add OpenModelDB metadata provider and model source for upscalers
Add OpenModelDB (openmodeldb.info) as a metadata and download source for
the existing upscaler model type.
Metadata:
- New OpenModelDBClient: fetches the site's bulk JSON dumps, caches them
on disk (24h TTL + ETag revalidation), and builds a local sha256 index
- New OpenModelDBModelMetadataProvider adapts catalogue entries to the
CivitAI-shaped version dict contract; registered in the fallback chain
behind the enable_openmodeldb_api setting (default on), gated to the
upscaler sub-type so other model types never trigger the dump download
- Persisted provenance uses metadata_source "openmodeldb" plus a nested
openmodeldb block (page URL, architecture, scale, license)
Images: paired-image LR/SR URLs are ephemeral imgdiff.net sessions, so
displayable images come from the site-hosted auto-generated thumbnails
(model-level cover leads images[], per-image thumbs for the rest); the
original comparison URL is kept in meta.comparisonUrl.
Downloads:
- New OpenModelDBSource (flat model ids, omdb: group prefix) with
resource filename derivation that recovers names hidden mid-path
(mediafire) or synthesizes {id}.{type} for folder links
- HTML-gateway mirrors (mediafire/mega/drive) are rejected with a clear
manual-download hint instead of silently saving an HTML page as .pth
- ModelSource base gains is_valid_source_id / default_subdir_parts /
resolve_download_url hooks so flat-id sources need no platform branches
UI: "View on OpenModelDB" link in the model modal (downloaded and
hash-enriched models), settings toggle next to the CivArchive one.
This commit is contained in:
Will Miao
@@ -0,0 +1,762 @@
|
||||
"""Client for the OpenModelDB bulk JSON API.
|
||||
|
||||
OpenModelDB (https://openmodeldb.info) is a static catalogue of upscaler
|
||||
models. It exposes no per-model or by-hash endpoint — only bulk JSON dumps
|
||||
(``/api/v1/models.json`` and friends), so this client downloads the dumps
|
||||
once, caches them on disk with a TTL, honors ETag/Last-Modified on refresh,
|
||||
and builds an in-memory SHA256 -> model index for read-only metadata lookups.
|
||||
|
||||
Every catalogue resource carries a ``sha256`` and a byte ``size``, which is
|
||||
what makes hash-based matching against local files possible. Lookups degrade
|
||||
gracefully: when the catalogue cannot be fetched (offline, upstream failure)
|
||||
the stale disk cache is used, and if there is no cache at all the lookup
|
||||
reports "not found" so the metadata fallback chain simply moves on.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
from typing import Any, Dict, List, Optional, Tuple
|
||||
from urllib.parse import urlparse
|
||||
|
||||
from .downloader import get_downloader
|
||||
from .errors import RateLimitError
|
||||
from ..utils.cache_paths import get_cache_base_dir
|
||||
from ..utils.constants import MODEL_FILE_EXTENSIONS
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
OPENMODELDB_API_BASE = "https://openmodeldb.info/api/v1"
|
||||
OPENMODELDB_SITE_BASE = "https://openmodeldb.info"
|
||||
|
||||
#: Value emitted as ``source`` in the synthesized version dict; the metadata
|
||||
#: sync service persists it as the model's ``metadata_source``.
|
||||
METADATA_SOURCE_VALUE = "openmodeldb"
|
||||
|
||||
#: Hosts whose URLs serve an HTML interstitial page instead of the raw file
|
||||
#: bytes. Downloading from them would silently save a web page as ``.pth``
|
||||
#: (the download flow does not verify the sha256 afterwards), so the model
|
||||
#: source layer rejects them with a manual-download hint. Kept deliberately
|
||||
#: small and explicit.
|
||||
HTML_GATEWAY_HOSTS = frozenset({"mediafire.com", "mega.nz", "drive.google.com"})
|
||||
|
||||
|
||||
def is_html_gateway_url(url: str) -> bool:
|
||||
"""Return ``True`` when *url* points at a known HTML-gateway host."""
|
||||
|
||||
if not isinstance(url, str) or not url:
|
||||
return False
|
||||
try:
|
||||
host = urlparse(url).netloc.lower()
|
||||
except ValueError:
|
||||
return False
|
||||
return any(host == g or host.endswith(f".{g}") for g in HTML_GATEWAY_HOSTS)
|
||||
|
||||
|
||||
def _is_ephemeral_viewer_url(url: str) -> bool:
|
||||
"""Return ``True`` for imgdiff.net session URLs.
|
||||
|
||||
Paired comparisons are hosted as ephemeral imgdiff viewer sessions
|
||||
(``/api/image.php?id=...``) that expire shortly after the site build;
|
||||
they 404 when used as an ``<img>`` source and must never be emitted as a
|
||||
displayable image URL.
|
||||
"""
|
||||
|
||||
return isinstance(url, str) and "imgdiff.net/api/" in url
|
||||
|
||||
#: Bulk dumps consumed by the client. Only ``models`` is strictly required;
|
||||
#: the rest resolve ids to human-readable names and degrade to raw ids.
|
||||
_DUMP_NAMES = ("models", "users", "tags", "architectures")
|
||||
|
||||
#: How long a fetched catalogue is considered fresh before a revalidation
|
||||
#: request is made. The site only changes when it is rebuilt (hours to days),
|
||||
#: so a daily TTL avoids re-downloading the ~1.4MB models dump on every
|
||||
#: lookup while still picking up new models reasonably fast.
|
||||
CACHE_TTL_SECONDS = 24 * 60 * 60
|
||||
|
||||
_META_FILENAME = "_meta.json"
|
||||
|
||||
|
||||
class OpenModelDBClient:
|
||||
"""Hash-lookup client over a locally cached OpenModelDB catalogue dump."""
|
||||
|
||||
_instance: Optional["OpenModelDBClient"] = None
|
||||
_instance_lock = asyncio.Lock()
|
||||
|
||||
@classmethod
|
||||
async def get_instance(cls) -> "OpenModelDBClient":
|
||||
"""Get the singleton instance of OpenModelDBClient."""
|
||||
async with cls._instance_lock:
|
||||
if cls._instance is None:
|
||||
cls._instance = cls()
|
||||
|
||||
# Register this client as a metadata provider (mirrors the
|
||||
# CivitAI/CivArchive client bootstrap).
|
||||
from .model_metadata_provider import (
|
||||
ModelMetadataProviderManager,
|
||||
OpenModelDBModelMetadataProvider,
|
||||
)
|
||||
|
||||
provider_manager = await ModelMetadataProviderManager.get_instance()
|
||||
provider_manager.register_provider(
|
||||
"openmodeldb",
|
||||
OpenModelDBModelMetadataProvider(cls._instance),
|
||||
False,
|
||||
)
|
||||
|
||||
return cls._instance
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
cache_dir: Optional[str] = None,
|
||||
ttl_seconds: float = CACHE_TTL_SECONDS,
|
||||
) -> None:
|
||||
# Guard re-initialization for the singleton pattern.
|
||||
if hasattr(self, "_initialized"):
|
||||
return
|
||||
self._initialized = True
|
||||
|
||||
self._cache_dir_override = cache_dir
|
||||
self._ttl_seconds = ttl_seconds
|
||||
|
||||
self._models: Dict[str, Dict[str, Any]] = {}
|
||||
self._users: Dict[str, Dict[str, Any]] = {}
|
||||
self._tags: Dict[str, Dict[str, Any]] = {}
|
||||
self._architectures: Dict[str, Dict[str, Any]] = {}
|
||||
# sha256 (lowercase) -> (model_id, model entry, matching resource)
|
||||
self._index: Dict[str, Tuple[str, Dict[str, Any], Dict[str, Any]]] = {}
|
||||
self._loaded_at: float = 0.0
|
||||
self._load_lock = asyncio.Lock()
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Public API
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def get_model_by_hash(
|
||||
self, model_hash: str
|
||||
) -> Tuple[Optional[Dict[str, Any]], Optional[str]]:
|
||||
"""Find an upscaler model by SHA256 hash.
|
||||
|
||||
Returns a CivitAI-shaped version dict (same contract as the other
|
||||
metadata providers) or ``(None, reason)``.
|
||||
"""
|
||||
if not model_hash or not isinstance(model_hash, str):
|
||||
return None, "Model not found"
|
||||
|
||||
try:
|
||||
loaded = await self._ensure_loaded()
|
||||
except RateLimitError:
|
||||
raise
|
||||
except Exception as exc:
|
||||
logger.error("OpenModelDB lookup failed for %s: %s", model_hash[:10], exc)
|
||||
return None, str(exc)
|
||||
|
||||
if not loaded:
|
||||
return None, "OpenModelDB catalogue unavailable"
|
||||
|
||||
hit = self._index.get(model_hash.lower())
|
||||
if hit is None:
|
||||
return None, "Model not found"
|
||||
|
||||
model_id, model_entry, resource = hit
|
||||
return self._to_civitai_version(model_id, model_entry, resource), None
|
||||
|
||||
async def catalogue_ready(self) -> bool:
|
||||
"""Return ``True`` when the catalogue is loaded (or loadable)."""
|
||||
|
||||
try:
|
||||
return await self._ensure_loaded()
|
||||
except RateLimitError:
|
||||
raise
|
||||
except Exception as exc:
|
||||
logger.error("OpenModelDB catalogue load failed: %s", exc)
|
||||
return False
|
||||
|
||||
async def get_model_entry(
|
||||
self, model_id: str
|
||||
) -> Optional[Tuple[str, Dict[str, Any]]]:
|
||||
"""Return ``(model_id, catalogue entry)`` for *model_id*, or ``None``.
|
||||
|
||||
Loads the catalogue on first use; an unavailable catalogue and an
|
||||
unknown id both yield ``None`` (callers that need to distinguish the
|
||||
two can check :meth:`catalogue_ready` first).
|
||||
"""
|
||||
if not model_id or not isinstance(model_id, str):
|
||||
return None
|
||||
if not await self.catalogue_ready():
|
||||
return None
|
||||
entry = self._models.get(model_id)
|
||||
if not isinstance(entry, dict):
|
||||
return None
|
||||
return model_id, entry
|
||||
|
||||
def find_resource_by_filename(
|
||||
self, model_id: str, entry: Dict[str, Any], filename: str
|
||||
) -> Optional[Dict[str, Any]]:
|
||||
"""Match a resource by its derived filename (see :meth:`resource_filename`)."""
|
||||
target = (filename or "").strip().lower()
|
||||
if not target:
|
||||
return None
|
||||
for resource in entry.get("resources") or []:
|
||||
if not isinstance(resource, dict):
|
||||
continue
|
||||
if self.resource_filename(model_id, resource).lower() == target:
|
||||
return resource
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def resource_filename(model_id: str, resource: Dict[str, Any]) -> str:
|
||||
"""Derive the local filename for a catalogue resource.
|
||||
|
||||
The download URL's basename is not authoritative — mirrors like
|
||||
mediafire put the real filename mid-path
|
||||
(``/file/<key>/90s_Sonic_2x.pth/file``) and folder links (mega.nz)
|
||||
have no filename at all. Strategy: the first URL path segment whose
|
||||
extension is a known model format, else ``{model_id}.{type}`` (the
|
||||
catalogue's ``type`` field is authoritative).
|
||||
"""
|
||||
urls = resource.get("urls")
|
||||
for url in urls if isinstance(urls, list) else []:
|
||||
if not isinstance(url, str):
|
||||
continue
|
||||
path = url.split("?", 1)[0].split("#", 1)[0]
|
||||
for segment in path.split("/"):
|
||||
if os.path.splitext(segment)[1].lower() in MODEL_FILE_EXTENSIONS:
|
||||
return segment
|
||||
resource_type = str(resource.get("type") or "").lower()
|
||||
extension = (
|
||||
resource_type if resource_type in {"pth", "safetensors", "onnx"} else "bin"
|
||||
)
|
||||
return f"{model_id}.{extension}"
|
||||
|
||||
@staticmethod
|
||||
def primary_url(resource: Dict[str, Any]) -> str:
|
||||
"""Return the resource's primary download URL.
|
||||
|
||||
Only the first URL is used: additional entries are mirrors that may
|
||||
need site-specific handling (e.g. mega.nz) and are never tried
|
||||
automatically.
|
||||
"""
|
||||
urls = resource.get("urls")
|
||||
if isinstance(urls, list):
|
||||
for url in urls:
|
||||
if isinstance(url, str) and url.startswith("http"):
|
||||
return url
|
||||
return ""
|
||||
|
||||
@staticmethod
|
||||
def direct_url(resource: Dict[str, Any]) -> str:
|
||||
"""Return the first URL that serves raw bytes, or ``""``.
|
||||
|
||||
HTML-gateway hosts (mediafire, mega.nz, drive.google.com — see
|
||||
:data:`HTML_GATEWAY_HOSTS`) serve an interstitial page instead of the
|
||||
file, so they are skipped here and reported to the user instead.
|
||||
"""
|
||||
urls = resource.get("urls")
|
||||
if isinstance(urls, list):
|
||||
for url in urls:
|
||||
if (
|
||||
isinstance(url, str)
|
||||
and url.startswith("http")
|
||||
and not is_html_gateway_url(url)
|
||||
):
|
||||
return url
|
||||
return ""
|
||||
|
||||
@staticmethod
|
||||
def _absolutize(url: str) -> str:
|
||||
"""Turn a site-relative path (``/thumbs/...``) into an absolute URL."""
|
||||
|
||||
if isinstance(url, str) and url.startswith("/"):
|
||||
return f"{OPENMODELDB_SITE_BASE}{url}"
|
||||
return url
|
||||
|
||||
def _paired_display_url(self, image: Dict[str, Any]) -> str:
|
||||
"""Return the displayable URL for a paired comparison image.
|
||||
|
||||
Prefers the site-hosted thumbnail: the ``LR``/``SR`` originals are
|
||||
frequently ephemeral imgdiff session URLs that 404 outside the
|
||||
viewer. Falls back to the SR (then LR) original only when it is not
|
||||
one of those session URLs.
|
||||
"""
|
||||
thumbnail = image.get("thumbnail")
|
||||
if isinstance(thumbnail, str) and thumbnail:
|
||||
return self._absolutize(thumbnail)
|
||||
for key in ("SR", "LR"):
|
||||
original = image.get(key)
|
||||
if (
|
||||
isinstance(original, str)
|
||||
and original
|
||||
and not _is_ephemeral_viewer_url(original)
|
||||
):
|
||||
return original
|
||||
return ""
|
||||
|
||||
def _model_preview_url(self, entry: Dict[str, Any]) -> str:
|
||||
"""Return the model-level thumbnail URL, mirroring the site's own
|
||||
``getPreviewImage`` precedence (paired → SR, standalone → url)."""
|
||||
thumbnail = entry.get("thumbnail")
|
||||
if not isinstance(thumbnail, dict):
|
||||
return ""
|
||||
if thumbnail.get("type") == "paired":
|
||||
url = thumbnail.get("SR") or thumbnail.get("LR")
|
||||
else:
|
||||
url = thumbnail.get("url")
|
||||
if isinstance(url, str) and url:
|
||||
return self._absolutize(url)
|
||||
return ""
|
||||
|
||||
def example_image_urls(self, entry: Dict[str, Any]) -> List[str]:
|
||||
"""Return displayable example-image URLs, model thumbnail first.
|
||||
|
||||
Never contains ephemeral imgdiff session URLs; standalone images keep
|
||||
their direct URLs (regular image hosts are hotlinkable).
|
||||
"""
|
||||
urls: List[str] = []
|
||||
lead = self._model_preview_url(entry)
|
||||
if lead:
|
||||
urls.append(lead)
|
||||
for image in entry.get("images") or []:
|
||||
if not isinstance(image, dict):
|
||||
continue
|
||||
if image.get("type") == "paired":
|
||||
url = self._paired_display_url(image)
|
||||
else:
|
||||
url = image.get("url")
|
||||
if isinstance(url, str) and url and url not in urls:
|
||||
urls.append(url)
|
||||
return urls
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Catalogue loading
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def _ensure_loaded(self) -> bool:
|
||||
"""Ensure the in-memory index is built, refreshing stale caches."""
|
||||
async with self._load_lock:
|
||||
if self._index and (time.monotonic() - self._loaded_at) < self._ttl_seconds:
|
||||
return True
|
||||
|
||||
meta = self._read_meta()
|
||||
fetched_at = float(meta.get("fetched_at") or 0.0)
|
||||
disk_fresh = (
|
||||
fetched_at > 0
|
||||
and (time.time() - fetched_at) < self._ttl_seconds
|
||||
and all(os.path.exists(self._dump_path(name)) for name in _DUMP_NAMES)
|
||||
)
|
||||
|
||||
if disk_fresh:
|
||||
if self._load_from_disk():
|
||||
return True
|
||||
# Corrupt disk cache: fall through to a network refresh.
|
||||
|
||||
if await self._refresh_from_network(meta):
|
||||
return True
|
||||
|
||||
# Network failed or was blocked: fall back to whatever is on disk,
|
||||
# however stale — old metadata beats none.
|
||||
if fetched_at > 0 and self._load_from_disk():
|
||||
logger.info("Using stale OpenModelDB cache (network refresh failed)")
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
def _load_from_disk(self) -> bool:
|
||||
"""Load all dumps from the disk cache and rebuild the index."""
|
||||
payloads: Dict[str, Dict[str, Any]] = {}
|
||||
for name in _DUMP_NAMES:
|
||||
path = self._dump_path(name)
|
||||
try:
|
||||
with open(path, "r", encoding="utf-8") as handle:
|
||||
data = json.load(handle)
|
||||
except FileNotFoundError:
|
||||
if name == "models":
|
||||
return False
|
||||
data = {}
|
||||
except (OSError, json.JSONDecodeError) as exc:
|
||||
logger.warning("Failed to read OpenModelDB cache %s: %s", path, exc)
|
||||
if name == "models":
|
||||
return False
|
||||
data = {}
|
||||
payloads[name] = data if isinstance(data, dict) else {}
|
||||
|
||||
if not payloads["models"]:
|
||||
return False
|
||||
|
||||
self._install_payloads(payloads)
|
||||
return True
|
||||
|
||||
async def _refresh_from_network(self, meta: Dict[str, Any]) -> bool:
|
||||
"""Revalidate cached dumps against the site and rebuild the index.
|
||||
|
||||
Honors ETag/Last-Modified via a HEAD probe: an unchanged dump keeps
|
||||
its cached body, so a TTL expiry without upstream changes costs one
|
||||
tiny request per dump instead of a full download.
|
||||
"""
|
||||
payloads: Dict[str, Dict[str, Any]] = {}
|
||||
etags: Dict[str, str] = dict(meta.get("etags") or {})
|
||||
last_modified: Dict[str, str] = dict(meta.get("last_modified") or {})
|
||||
|
||||
for name in _DUMP_NAMES:
|
||||
payload, etag, modified = await self._fetch_dump(
|
||||
name,
|
||||
known_etag=etags.get(name) or "",
|
||||
known_last_modified=last_modified.get(name) or "",
|
||||
)
|
||||
if payload is None:
|
||||
if name == "models":
|
||||
return False
|
||||
payload = {}
|
||||
payloads[name] = payload
|
||||
if etag:
|
||||
etags[name] = etag
|
||||
if modified:
|
||||
last_modified[name] = modified
|
||||
|
||||
self._install_payloads(payloads)
|
||||
self._write_cache(payloads, etags, last_modified)
|
||||
return True
|
||||
|
||||
async def _fetch_dump(
|
||||
self,
|
||||
name: str,
|
||||
*,
|
||||
known_etag: str,
|
||||
known_last_modified: str,
|
||||
) -> Tuple[Optional[Dict[str, Any]], str, str]:
|
||||
"""Fetch one dump, returning ``(body, etag, last_modified)``.
|
||||
|
||||
``body`` is ``None`` when the fetch failed and there is no usable
|
||||
cached copy. When the HEAD probe shows the resource unchanged, the
|
||||
cached body is returned without a full download.
|
||||
"""
|
||||
url = f"{OPENMODELDB_API_BASE}/{name}.json"
|
||||
disk_path = self._dump_path(name)
|
||||
have_cached = os.path.exists(disk_path)
|
||||
|
||||
downloader = await get_downloader()
|
||||
|
||||
head_etag = ""
|
||||
head_modified = ""
|
||||
try:
|
||||
head_ok, head_headers = await downloader.get_response_headers(url)
|
||||
except Exception as exc: # pragma: no cover - defensive guard
|
||||
logger.debug("OpenModelDB HEAD probe failed for %s: %s", url, exc)
|
||||
head_ok, head_headers = False, {}
|
||||
|
||||
if head_ok and isinstance(head_headers, dict):
|
||||
# aiohttp headers are case-insensitive; plain dicts in tests are not.
|
||||
head_etag = str(head_headers.get("ETag") or head_headers.get("etag") or "")
|
||||
head_modified = str(
|
||||
head_headers.get("Last-Modified") or head_headers.get("last-modified") or ""
|
||||
)
|
||||
|
||||
if (
|
||||
have_cached
|
||||
and known_etag
|
||||
and head_etag
|
||||
and head_etag == known_etag
|
||||
):
|
||||
cached = self._read_dump_file(disk_path)
|
||||
if cached is not None:
|
||||
logger.debug("OpenModelDB %s unchanged (etag match); using cache", name)
|
||||
return cached, known_etag, known_last_modified or head_modified
|
||||
|
||||
success, payload = await downloader.make_request("GET", url, use_auth=False)
|
||||
if isinstance(payload, RateLimitError):
|
||||
raise payload
|
||||
if not success or not isinstance(payload, dict):
|
||||
logger.warning(
|
||||
"OpenModelDB %s fetch failed: %s",
|
||||
name,
|
||||
payload if isinstance(payload, str) else "unexpected payload",
|
||||
)
|
||||
if have_cached:
|
||||
cached = self._read_dump_file(disk_path)
|
||||
if cached is not None:
|
||||
return cached, known_etag, known_last_modified
|
||||
return None, known_etag, known_last_modified
|
||||
|
||||
return payload, head_etag or known_etag, head_modified or known_last_modified
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Index and transformation
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _install_payloads(self, payloads: Dict[str, Dict[str, Any]]) -> None:
|
||||
"""Install dump payloads and rebuild the sha256 index."""
|
||||
self._models = payloads.get("models") or {}
|
||||
self._users = payloads.get("users") or {}
|
||||
self._tags = payloads.get("tags") or {}
|
||||
self._architectures = payloads.get("architectures") or {}
|
||||
self._index = self._build_index(self._models)
|
||||
self._loaded_at = time.monotonic()
|
||||
logger.debug(
|
||||
"OpenModelDB catalogue loaded: %d models, %d indexed hashes",
|
||||
len(self._models),
|
||||
len(self._index),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _build_index(
|
||||
models: Dict[str, Dict[str, Any]]
|
||||
) -> Dict[str, Tuple[str, Dict[str, Any], Dict[str, Any]]]:
|
||||
"""Build the sha256 -> (model_id, model entry, resource) index."""
|
||||
index: Dict[str, Tuple[str, Dict[str, Any], Dict[str, Any]]] = {}
|
||||
for model_id, entry in models.items():
|
||||
if not isinstance(entry, dict):
|
||||
continue
|
||||
resources = entry.get("resources")
|
||||
if not isinstance(resources, list):
|
||||
continue
|
||||
for resource in resources:
|
||||
if not isinstance(resource, dict):
|
||||
continue
|
||||
sha256 = resource.get("sha256")
|
||||
if not isinstance(sha256, str) or not sha256:
|
||||
continue
|
||||
# First writer wins: duplicate hashes across catalogue entries
|
||||
# are ambiguous and cannot be disambiguated locally.
|
||||
index.setdefault(sha256.lower(), (model_id, entry, resource))
|
||||
return index
|
||||
|
||||
def _resolve_authors(self, entry: Dict[str, Any]) -> Tuple[str, List[str]]:
|
||||
"""Resolve the author field to a display name plus the raw user ids."""
|
||||
raw = entry.get("author")
|
||||
author_ids = raw if isinstance(raw, list) else [raw]
|
||||
ids = [str(a) for a in author_ids if isinstance(a, str) and a]
|
||||
names: List[str] = []
|
||||
for author_id in ids:
|
||||
user = self._users.get(author_id)
|
||||
name = user.get("name") if isinstance(user, dict) else None
|
||||
names.append(name if isinstance(name, str) and name else author_id)
|
||||
return ", ".join(names), ids
|
||||
|
||||
def _resolve_tags(self, entry: Dict[str, Any]) -> List[str]:
|
||||
"""Resolve tag ids to their display names."""
|
||||
raw_tags = entry.get("tags")
|
||||
if not isinstance(raw_tags, list):
|
||||
return []
|
||||
resolved: List[str] = []
|
||||
for tag_id in raw_tags:
|
||||
if not isinstance(tag_id, str) or not tag_id:
|
||||
continue
|
||||
tag = self._tags.get(tag_id)
|
||||
name = tag.get("name") if isinstance(tag, dict) else None
|
||||
resolved.append(name if isinstance(name, str) and name else tag_id)
|
||||
return resolved
|
||||
|
||||
def _resolve_architecture_name(self, entry: Dict[str, Any]) -> str:
|
||||
"""Resolve the architecture id to its display name."""
|
||||
arch_id = entry.get("architecture")
|
||||
if not isinstance(arch_id, str) or not arch_id:
|
||||
return ""
|
||||
arch = self._architectures.get(arch_id)
|
||||
if isinstance(arch, dict):
|
||||
name = arch.get("name")
|
||||
if isinstance(name, str) and name:
|
||||
return name
|
||||
return arch_id
|
||||
|
||||
@staticmethod
|
||||
def _resource_format(resource: Dict[str, Any]) -> str:
|
||||
"""Map an OpenModelDB resource type to a CivitAI file metadata format."""
|
||||
resource_type = str(resource.get("type") or "").lower()
|
||||
if resource_type == "safetensors":
|
||||
return "SafeTensor"
|
||||
if resource_type in ("pth", "pt", "ckpt"):
|
||||
return "PickleTensor"
|
||||
return "Other"
|
||||
|
||||
def _to_civitai_version(
|
||||
self,
|
||||
model_id: str,
|
||||
entry: Dict[str, Any],
|
||||
matched_resource: Dict[str, Any],
|
||||
) -> Dict[str, Any]:
|
||||
"""Map an OpenModelDB catalogue entry to a CivitAI-shaped version dict.
|
||||
|
||||
Follows the same contract as the CivArchive/SQLite providers so the
|
||||
metadata sync service can merge it unchanged. Numeric ``id``/``modelId``
|
||||
are deliberately omitted: OpenModelDB ids are strings, and consumers
|
||||
treat a missing ``modelId`` as "not a CivitAI model" (no CivitAI page
|
||||
link, no update checks).
|
||||
"""
|
||||
author_display, author_ids = self._resolve_authors(entry)
|
||||
tags = self._resolve_tags(entry)
|
||||
architecture_id = entry.get("architecture")
|
||||
architecture_name = self._resolve_architecture_name(entry)
|
||||
description = entry.get("description")
|
||||
license_name = entry.get("license")
|
||||
page_url = f"{OPENMODELDB_SITE_BASE}/models/{model_id}"
|
||||
|
||||
files: List[Dict[str, Any]] = []
|
||||
resources = entry.get("resources")
|
||||
for resource in resources if isinstance(resources, list) else []:
|
||||
if not isinstance(resource, dict):
|
||||
continue
|
||||
# The displayable download URL prefers a direct-bytes mirror when
|
||||
# one exists; the filename is derived (never the raw URL basename,
|
||||
# which mediafire-style mirrors leave as "file").
|
||||
download_url = self.direct_url(resource) or self.primary_url(resource)
|
||||
sha256 = resource.get("sha256")
|
||||
size_bytes = resource.get("size")
|
||||
files.append(
|
||||
{
|
||||
"name": self.resource_filename(model_id, resource),
|
||||
"type": "Model",
|
||||
"sizeKB": (size_bytes / 1024.0)
|
||||
if isinstance(size_bytes, (int, float))
|
||||
else 0,
|
||||
"downloadUrl": download_url,
|
||||
"primary": resource is matched_resource,
|
||||
"hashes": {"SHA256": str(sha256).upper()} if sha256 else {},
|
||||
"metadata": {"format": self._resource_format(resource)},
|
||||
}
|
||||
)
|
||||
|
||||
images: List[Dict[str, Any]] = []
|
||||
# The model-level thumbnail is the site's own preview pick and larger
|
||||
# than the per-image small thumbs; the card preview derives from
|
||||
# images[0], so it leads the list.
|
||||
lead = self._model_preview_url(entry)
|
||||
if lead:
|
||||
images.append({"url": lead, "nsfwLevel": 1, "type": "image"})
|
||||
raw_images = entry.get("images")
|
||||
for image in raw_images if isinstance(raw_images, list) else []:
|
||||
if not isinstance(image, dict):
|
||||
continue
|
||||
paired = image.get("type") == "paired"
|
||||
# Paired entries show the upscaled (SR) result as the preview.
|
||||
url = self._paired_display_url(image) if paired else image.get("url")
|
||||
if not isinstance(url, str) or not url:
|
||||
continue
|
||||
if any(existing["url"] == url for existing in images):
|
||||
continue
|
||||
mapped: Dict[str, Any] = {"url": url, "nsfwLevel": 1, "type": "image"}
|
||||
thumbnail = image.get("thumbnail")
|
||||
if isinstance(thumbnail, str) and thumbnail:
|
||||
thumbnail_url = self._absolutize(thumbnail)
|
||||
if thumbnail_url != url:
|
||||
mapped["thumbnailUrl"] = thumbnail_url
|
||||
meta: Dict[str, Any] = {}
|
||||
if paired:
|
||||
comparison = image.get("SR") or image.get("LR")
|
||||
if (
|
||||
isinstance(comparison, str)
|
||||
and comparison
|
||||
and comparison != url
|
||||
and _is_ephemeral_viewer_url(comparison)
|
||||
):
|
||||
# Ephemeral imgdiff viewer session, kept for reference
|
||||
# only — it 404s outside the session and is never
|
||||
# displayable.
|
||||
meta["comparisonUrl"] = comparison
|
||||
caption = image.get("caption")
|
||||
if isinstance(caption, str) and caption:
|
||||
meta["caption"] = caption
|
||||
if meta:
|
||||
mapped["meta"] = meta
|
||||
images.append(mapped)
|
||||
|
||||
return {
|
||||
"name": entry.get("name") or model_id,
|
||||
# Upscalers are not tied to a diffusion base model.
|
||||
"baseModel": "Other",
|
||||
"description": description or "",
|
||||
"publishedAt": entry.get("date"),
|
||||
"trainedWords": [],
|
||||
"model": {
|
||||
"name": entry.get("name") or model_id,
|
||||
"type": "Upscaler",
|
||||
"nsfw": False,
|
||||
"description": description,
|
||||
"tags": tags,
|
||||
"license": license_name or "",
|
||||
},
|
||||
"creator": {"username": author_display, "image": None},
|
||||
"files": files,
|
||||
"images": images,
|
||||
"source": METADATA_SOURCE_VALUE,
|
||||
# OpenModelDB-native provenance, kept inside the persisted payload
|
||||
# so the UI can link to the model page in a later phase.
|
||||
"openmodeldb": {
|
||||
"id": model_id,
|
||||
"url": page_url,
|
||||
"authors": author_ids,
|
||||
"architecture": architecture_id or "",
|
||||
"architectureName": architecture_name,
|
||||
"scale": entry.get("scale"),
|
||||
"inputChannels": entry.get("inputChannels"),
|
||||
"outputChannels": entry.get("outputChannels"),
|
||||
"size": entry.get("size") or [],
|
||||
"license": license_name or "",
|
||||
"date": entry.get("date"),
|
||||
},
|
||||
}
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Disk cache
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _cache_dir(self) -> str:
|
||||
base = self._cache_dir_override or os.path.join(
|
||||
get_cache_base_dir(), "openmodeldb"
|
||||
)
|
||||
os.makedirs(base, exist_ok=True)
|
||||
return base
|
||||
|
||||
def _dump_path(self, name: str) -> str:
|
||||
return os.path.join(self._cache_dir(), f"{name}.json")
|
||||
|
||||
def _meta_path(self) -> str:
|
||||
return os.path.join(self._cache_dir(), _META_FILENAME)
|
||||
|
||||
def _read_meta(self) -> Dict[str, Any]:
|
||||
try:
|
||||
with open(self._meta_path(), "r", encoding="utf-8") as handle:
|
||||
meta = json.load(handle)
|
||||
return meta if isinstance(meta, dict) else {}
|
||||
except FileNotFoundError:
|
||||
return {}
|
||||
except (OSError, json.JSONDecodeError) as exc:
|
||||
logger.warning("Failed to read OpenModelDB cache meta: %s", exc)
|
||||
return {}
|
||||
|
||||
def _read_dump_file(self, path: str) -> Optional[Dict[str, Any]]:
|
||||
try:
|
||||
with open(path, "r", encoding="utf-8") as handle:
|
||||
data = json.load(handle)
|
||||
return data if isinstance(data, dict) else None
|
||||
except (OSError, json.JSONDecodeError) as exc:
|
||||
logger.warning("Failed to read OpenModelDB cache %s: %s", path, exc)
|
||||
return None
|
||||
|
||||
def _write_cache(
|
||||
self,
|
||||
payloads: Dict[str, Dict[str, Any]],
|
||||
etags: Dict[str, str],
|
||||
last_modified: Dict[str, str],
|
||||
) -> None:
|
||||
for name, payload in payloads.items():
|
||||
path = self._dump_path(name)
|
||||
try:
|
||||
with open(path, "w", encoding="utf-8") as handle:
|
||||
json.dump(payload, handle)
|
||||
except OSError as exc:
|
||||
logger.warning("Failed to write OpenModelDB cache %s: %s", path, exc)
|
||||
|
||||
meta = {
|
||||
"fetched_at": time.time(),
|
||||
"etags": etags,
|
||||
"last_modified": last_modified,
|
||||
}
|
||||
try:
|
||||
with open(self._meta_path(), "w", encoding="utf-8") as handle:
|
||||
json.dump(meta, handle, indent=2)
|
||||
except OSError as exc:
|
||||
logger.warning("Failed to write OpenModelDB cache meta: %s", exc)
|
||||
Reference in New Issue
Block a user