DemonGatanjieu-Anomalous_Model_Browser/api/metadata.py

import os
import sys
import json
import urllib.parse
import subprocess
import threading
import asyncio
import copy
from functools import lru_cache
from aiohttp import web
import folder_paths
import struct
try:
    from ..model_identity import USER_INFO_SUFFIX, is_unmatched, read_json, sidecar_file_hash, sidecar_info
except ImportError:
    from model_identity import USER_INFO_SUFFIX, is_unmatched, read_json, sidecar_file_hash, sidecar_info

# Fields the model editor sets: <model>.anomalous.json key -> metadata key.
USER_FIELDS = {"custom_name": "custom_name", "custom_notes": "custom_notes", "source_url": "source_url"}

def _select_info_file(data, file_path):
    """Select the Civitai file entry that actually describes ``file_path``.

    A model version may contain a checkpoint, text encoder, and VAE in the same
    ``files`` array.  Taking the first SHA256 silently assigns the dependency's
    hash to every local sidecar copied from that version.  Physical byte size is
    the strongest cheap discriminator available here and remains valid after a
    local rename.  Filenames are intentionally excluded: local renames are the
    reason provenance recovery exists and cannot be identity evidence.
    """
    entries = []
    for file_info in data.get("files", []):
        if not isinstance(file_info, dict):
            continue
        hashes = file_info.get("hashes", {})
        sha256 = hashes.get("SHA256") if isinstance(hashes, dict) else None
        if not sha256:
            continue

        size_bytes = None
        try:
            if file_info.get("sizeKB") is not None:
                size_bytes = int(round(float(file_info["sizeKB"]) * 1024))
        except (TypeError, ValueError, OverflowError):
            pass
        entries.append((file_info, str(sha256), size_bytes))

    if not entries:
        return None

    try:
        physical_size = os.path.getsize(file_path)
    except OSError:
        physical_size = None

    if physical_size is not None:
        size_matches = [entry for entry in entries if entry[2] == physical_size]
        if len(size_matches) == 1:
            return size_matches[0][0]
        if len(size_matches) > 1:
            return None

    # Offline-generated metadata commonly contains one injected hash without
    # Civitai size/name data.  It is safe only when there is a single candidate.
    if len(entries) == 1:
        return entries[0][0]
    return None


def _select_info_hash(data, file_path):
    selected = _select_info_file(data, file_path)
    if not selected:
        return ""
    hashes = selected.get("hashes", {})
    return str(hashes.get("SHA256", "")) if isinstance(hashes, dict) else ""


def _positive_civitai_id(value):
    """Return a real Civitai identifier, excluding offline sentinel values."""
    if isinstance(value, bool):
        return None
    try:
        parsed = int(value)
    except (TypeError, ValueError, OverflowError):
        return None
    return parsed if parsed > 0 else None

def _read_metadata(file_path):
    base_path = os.path.splitext(file_path)[0]
    metadata = {
        "name": os.path.basename(base_path),
        "description": "",
        "notes": "",
        "trainedWords": [],
        "baseModel": "",
        "civitai_url": "",
        "source_url": "",
        "hash": "",
        "hash_algorithm": "sha256",
        "hash_scope": "file",
        "hash_source": "",
        "custom_name": "",
        "custom_notes": "",
        # "civitai" (matched), "local" (only what the file tells: no match, or offline) or "" (not scanned)
        "info_source": "",
        "user_fields": [],
    }
    
    info_files = [f"{base_path}.info", f"{base_path}.civitai.info"]
    for info_file in info_files:
        if os.path.exists(info_file):
            try:
                with open(info_file, 'r', encoding='utf-8') as f:
                    data = json.load(f)
                    
                    name = data.get("name", "")
                    if "model" in data and isinstance(data["model"], dict):
                        model_name = data["model"].get("name", "")
                        name = f"{model_name} - {name}".strip(' -')
                        
                    description = data.get("description", "") or ""
                    if not description and "model" in data and isinstance(data["model"], dict):
                        description = data["model"].get("description", "")
                    notes = data.get("notes", "") 
                    trained_words = data.get("trainedWords", [])
                    base_model = data.get("baseModel", "")
                    
                    model_id = _positive_civitai_id(data.get("modelId"))
                    version_id = _positive_civitai_id(data.get("id"))
                    civitai_url = ""
                    if model_id is not None:
                        # Handle Civitai's new mature content policy
                        nsfw_level = data.get("nsfwLevel", 1)
                        is_nsfw = False
                        if "model" in data and isinstance(data["model"], dict):
                            is_nsfw = data["model"].get("nsfw", False)
                        
                        domain = "civitai.red" if (nsfw_level > 1 or is_nsfw) else "civitai.com"
                        civitai_url = f"https://{domain}/models/{model_id}"
                        if version_id is not None:
                            civitai_url += f"?modelVersionId={version_id}"
                    
                    
                    selected_file = _select_info_file(data, file_path)
                    hash_val, hash_source = sidecar_file_hash(data, file_path, selected_file)
                    if hash_val:
                        metadata["hash_source"] = hash_source

                    if name: metadata["name"] = name
                    if description: metadata["description"] = description
                    if notes: metadata["notes"] = notes
                    if trained_words: metadata["trainedWords"] = trained_words
                    if base_model: metadata["baseModel"] = base_model
                    if civitai_url: metadata["civitai_url"] = civitai_url
                    if model_id is not None: metadata["model_id"] = model_id
                    if version_id is not None: metadata["version_id"] = version_id
                    if hash_val: metadata["hash"] = hash_val
                    if "anomalous_custom_name" in data and data["anomalous_custom_name"]: metadata["custom_name"] = data["anomalous_custom_name"]
                    if "anomalous_custom_notes" in data and data["anomalous_custom_notes"]: metadata["custom_notes"] = data["anomalous_custom_notes"]
                    if "anomalous_source_url" in data and data["anomalous_source_url"]: metadata["source_url"] = data["anomalous_source_url"]
            except Exception:
                pass

    scan = sidecar_info(base_path)
    if scan is not None:
        metadata["info_source"] = "local" if is_unmatched(scan) else "civitai"
    # The user's own layer comes last and wins, also when a field was cleared on purpose ("").
    user = read_json(base_path + USER_INFO_SUFFIX) or {}
    for key, field in USER_FIELDS.items():
        if isinstance(user.get(key), str):
            metadata[field] = user[key]
            metadata["user_fields"].append(field)
    return metadata


def _file_signature(file_path):
    try:
        stat = os.stat(file_path)
        return (stat.st_size, stat.st_mtime_ns, stat.st_ctime_ns)
    except OSError:
        return None


def _metadata_signature(file_path):
    base_path = os.path.splitext(file_path)[0]
    return (
        _file_signature(file_path),
        _file_signature(f"{base_path}.info"),
        _file_signature(f"{base_path}.civitai.info"),
        _file_signature(base_path + USER_INFO_SUFFIX),
    )


@lru_cache(maxsize=4096)
def _get_metadata_cached(file_path, signature):
    # ``signature`` is deliberately unused by the parser itself; it is part of
    # the cache key and changes whenever the model or either metadata sidecar
    # changes. The cache is bounded so removed/renamed models cannot grow it
    # without limit.
    return _read_metadata(file_path)


def get_metadata(file_path):
    normalized_path = os.path.realpath(file_path)
    signature = _metadata_signature(normalized_path)
    # Metadata contains mutable lists. Never expose the cached object directly.
    return copy.deepcopy(_get_metadata_cached(normalized_path, signature))


def clear_metadata_cache():
    _get_metadata_cached.cache_clear()


def get_metadata_cache_info():
    return _get_metadata_cached.cache_info()