ComfyUI-Lora-Manager/tests/enrich_hf_validation/metadata_constructor.py

"""Construct initial ``.metadata.json`` sidecars for HF model repos.

Each HF repo ID gets a minimal metadata file — no real model file is needed.
The enrichment pipeline reads only the sidecar.
"""

from __future__ import annotations

import json
import logging
import os
from typing import Dict, List

from .config import CIVITAI_MODEL_TAGS

logger = logging.getLogger(__name__)


def load_repo_ids(path: str, max_models: int | None = None) -> List[str]:
    """Read HF repo IDs from *path* (one per line, ignoring blanks/comments)."""
    path = os.path.expanduser(path)
    if not os.path.exists(path):
        raise FileNotFoundError(f"Models file not found: {path}")

    repos: List[str] = []
    with open(path, "r", encoding="utf-8") as fh:
        for line in fh:
            line = line.strip()
            if not line or line.startswith("#"):
                continue
            repos.append(line)

    if max_models is not None and max_models > 0:
        repos = repos[:max_models]

    logger.info("Loaded %d HF repo IDs from %s", len(repos), path)
    return repos


def sanitize_repo_id(repo_id: str) -> str:
    """Turn ``user/repo-name`` into a safe directory name."""
    return repo_id.replace("/", "__").replace(".", "_")


def build_model_dir(output_dir: str, repo_id: str) -> str:
    """Return the per-model working directory."""
    return os.path.join(output_dir, "models", sanitize_repo_id(repo_id))


def build_model_path(model_dir: str) -> str:
    """Return a synthetic model file path (no real file will exist)."""
    return os.path.join(model_dir, "model.safetensors")


def build_metadata_path(model_path: str) -> str:
    """Return the sidecar path for a model file.

    This MUST match the convention used by ``MetadataManager`` /
    ``apply_metadata_updates``, which derives the sidecar path via
    ``os.path.splitext(model_path)[0] + '.metadata.json'``.
    For a model file ``model.safetensors`` the sidecar is
    ``model.metadata.json`` — *not* ``model.safetensors.metadata.json``.
    """
    return f"{os.path.splitext(model_path)[0]}.metadata.json"


def create_initial_metadata(
    output_dir: str,
    repo_id: str,
) -> str:
    """Write a minimal ``.metadata.json`` for *repo_id*.

    Returns the **model path** (the ``.safetensors`` path whose sidecar was
    written).  The caller passes this path to ``AgentService.execute_skill``.
    """
    model_dir = build_model_dir(output_dir, repo_id)
    os.makedirs(model_dir, exist_ok=True)
    model_path = build_model_path(model_dir)
    metadata_path = build_metadata_path(model_path)

    hf_url = f"https://huggingface.co/{repo_id}"
    file_name = repo_id.split("/")[-1]

    metadata: Dict = {
        "file_name": file_name,
        "model_name": file_name,
        "file_path": model_path.replace(os.sep, "/"),
        "size": 0,
        "modified": 0,
        "sha256": "",
        "base_model": "Unknown",
        "preview_url": "",
        "preview_nsfw_level": 0,
        "notes": "",
        "from_civitai": False,
        "civitai": {},
        "tags": [],
        "modelDescription": "",
        "civitai_deleted": False,
        "favorite": False,
        "exclude": False,
        "db_checked": False,
        "skip_metadata_refresh": False,
        "metadata_source": "",
        "last_checked_at": 0,
        "hash_status": "completed",
        "trainedWords": [],
        "hf_url": hf_url,
        "usage_tips": "{}",
    }

    with open(metadata_path, "w", encoding="utf-8") as fh:
        json.dump(metadata, fh, indent=2, ensure_ascii=False)

    logger.debug("Created initial metadata for %s -> %s", repo_id, metadata_path)
    return model_path


def create_all_initial_metadata(
    repos: List[str],
    output_dir: str,
    *,
    skip_existing: bool = True,
) -> List[str]:
    """Create initial metadata for every repo in *repos*.

    Returns a list of model paths in the same order as *repos*.
    ``skip_existing=True`` skips repos whose metadata already exists,
    allowing safe re-run.
    """
    model_paths: List[str] = []
    for repo_id in repos:
        model_dir = build_model_dir(output_dir, repo_id)
        model_path = build_model_path(model_dir)
        metadata_path = build_metadata_path(model_path)

        if skip_existing and os.path.exists(metadata_path):
            model_paths.append(model_path)
            continue

        model_paths.append(create_initial_metadata(output_dir, repo_id))

    logger.info(
        "Constructed initial metadata for %d/%d repos",
        len(model_paths),
        len(repos),
    )
    return model_paths