mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-07-06 09:21:16 -03:00
feat(agent): enrich_hf_metadata — filename-aware section matching, preview extraction for markdown/HTML/widget, JSON salvage, instance_prompt fallback, and validation suite
- extract_relevant_section(): trim README to model-filename-matching section for collection repos (download link, anchor ID, heading strategies) - _strip_standalone_images(): preserve markdown image URLs so LLM can extract preview_url; strip only HTML <img> tags - extract_simple_markdown_images(): extract civitai.images from ![]() body - extract_html_img_tags(): extract from <img src="..."> (deadman44-style) - extract_gallery_images(): fix widget parser for YAML - output: dash prefix - _is_heading: exclude </hN> closing tags from boundary detection - _extract_section: start at matching heading when match IS a heading line - _try_salvage_json(): recover truncated JSON (close braces/brackets in LIFO order, close unterminated strings, strip trailing commas) - PostProcessor: store _llm_confidence, add instance_prompt YAML fallback - agent_service: pass model_basename to prompt, trim README via extract_relevant_section before clean_readme_for_llm - Add tests/enrich_hf_validation/ suite: 100-model pipeline with progress checkpoint/resume, per-field scoring, markdown+JSON reporting - Fix evaluation_engine: read _llm_confidence (not _llm_response)
This commit is contained in:
148
tests/enrich_hf_validation/metadata_constructor.py
Normal file
148
tests/enrich_hf_validation/metadata_constructor.py
Normal file
@@ -0,0 +1,148 @@
|
||||
"""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
|
||||
Reference in New Issue
Block a user