mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-07-06 09:21:16 -03:00
feat(agent): optimize enrich_hf_metadata with README cleaning, Ollama native API, and expanded fields
- Add clean_readme_for_llm() to strip noise from README before LLM injection - Keep widget section text (valuable tag signal) and unmarked code blocks (trigger words) - Preserve standalone image alt text instead of removing entirely - Switch Ollama to native /api/chat with think:false to fix empty content on thinking models - Extract Sample Gallery table images and deduplicate with widget images - Only strip code blocks with explicit language tags (bash) - Add notes and usage_tips fields to SKILL.md output format and post-processor - Clean up dead code, fix regex edge cases, remove double type annotation
This commit is contained in:
@@ -28,6 +28,7 @@ from ..llm_service import LLMService
|
||||
from ..websocket_manager import ws_manager
|
||||
from .post_processor import PostProcessor
|
||||
from .skill_registry import SkillRegistry
|
||||
from .skills.enrich_hf_metadata.md_to_html import clean_readme_for_llm
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -368,7 +369,8 @@ class AgentService:
|
||||
context["repo"] = repo or ""
|
||||
if repo:
|
||||
readme = await self._fetch_readme(repo)
|
||||
context["readme_content"] = readme[:8000] if readme else "(README not available)"
|
||||
cleaned = clean_readme_for_llm(readme) if readme else ""
|
||||
context["readme_content"] = cleaned if cleaned else "(README not available)"
|
||||
context["readme_content_full"] = readme or ""
|
||||
|
||||
try:
|
||||
|
||||
@@ -10,6 +10,7 @@ refresh cache). All actual I/O is delegated to :mod:`~py.agent_cli`.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
from datetime import datetime, timezone
|
||||
@@ -79,6 +80,7 @@ class PostProcessor:
|
||||
from .skills.enrich_hf_metadata.md_to_html import (
|
||||
convert_readme_to_html,
|
||||
extract_gallery_images,
|
||||
extract_gallery_table_images,
|
||||
extract_repo_from_hf_url,
|
||||
)
|
||||
|
||||
@@ -127,23 +129,38 @@ class PostProcessor:
|
||||
desc_civitai["description"] = short_desc
|
||||
updates["civitai"] = desc_civitai
|
||||
|
||||
# gallery images → civitai.images (from YAML frontmatter widget entries)
|
||||
# gallery images → civitai.images (from YAML frontmatter widget entries
|
||||
# and Sample Gallery markdown tables in the README body)
|
||||
gallery_images: List[Dict[str, Any]] = []
|
||||
if readme_content and is_hf_model:
|
||||
hf_url = metadata.get("hf_url", "") or ""
|
||||
repo = extract_repo_from_hf_url(hf_url)
|
||||
if repo:
|
||||
rec_w = llm_output.get("recommended_width") or 0
|
||||
rec_h = llm_output.get("recommended_height") or 0
|
||||
|
||||
# 1. Widget images (YAML frontmatter)
|
||||
gallery = extract_gallery_images(
|
||||
readme_content, repo,
|
||||
default_width=rec_w, default_height=rec_h,
|
||||
)
|
||||
if gallery:
|
||||
|
||||
# 2. Sample Gallery table images (markdown body), deduplicated
|
||||
existing_urls = {img["url"] for img in gallery if img.get("url")}
|
||||
table_images = extract_gallery_table_images(
|
||||
readme_content, repo,
|
||||
existing_urls=existing_urls,
|
||||
default_width=rec_w, default_height=rec_h,
|
||||
)
|
||||
|
||||
all_images = gallery + table_images
|
||||
if all_images:
|
||||
gallery_images = all_images
|
||||
current_civitai = metadata.get("civitai") or {}
|
||||
gallery_civitai = dict(current_civitai)
|
||||
if "civitai" in updates and isinstance(updates["civitai"], dict):
|
||||
gallery_civitai.update(updates["civitai"])
|
||||
gallery_civitai["images"] = gallery
|
||||
gallery_civitai["images"] = all_images
|
||||
updates["civitai"] = gallery_civitai
|
||||
|
||||
# tags
|
||||
@@ -159,6 +176,11 @@ class PostProcessor:
|
||||
updates["llm_enriched_at"] = datetime.now(timezone.utc).isoformat()
|
||||
|
||||
preview_remote_url = (llm_output.get("preview_url") or "").strip()
|
||||
# Fallback: if the LLM couldn't find a preview image in the cleaned
|
||||
# README, use the first gallery image extracted from the YAML widget
|
||||
# section.
|
||||
if not preview_remote_url and gallery_images:
|
||||
preview_remote_url = gallery_images[0].get("url", "")
|
||||
current_preview = metadata.get("preview_url") or ""
|
||||
if preview_remote_url and not (current_preview and os.path.exists(current_preview)):
|
||||
local_path = await download_preview(model_path, preview_remote_url)
|
||||
@@ -166,6 +188,22 @@ class PostProcessor:
|
||||
preview_downloaded = True
|
||||
updates["preview_url"] = local_path
|
||||
|
||||
# notes — plain-text summary of usage info from the LLM
|
||||
new_notes = (llm_output.get("notes") or "").strip()
|
||||
if new_notes:
|
||||
updates["notes"] = new_notes
|
||||
|
||||
# usage_tips — JSON string (e.g. {"strength_min":0.85,"strength_max":1.4})
|
||||
raw_tips = (llm_output.get("usage_tips") or "").strip()
|
||||
if raw_tips and raw_tips != "{}":
|
||||
try:
|
||||
json.loads(raw_tips)
|
||||
updates["usage_tips"] = raw_tips
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
logger.warning(
|
||||
"LLM returned invalid usage_tips JSON: %s", raw_tips[:200]
|
||||
)
|
||||
|
||||
if updates:
|
||||
updated_fields = await apply_metadata_updates(model_path, updates)
|
||||
|
||||
|
||||
@@ -84,6 +84,25 @@ The recommended image generation resolution for this model, in pixels. Look for
|
||||
### preview_url
|
||||
The URL of the most suitable preview image from the README. Look for image tags (e.g. ``) and the YAML frontmatter `widget:` section (which often has `output.url` fields). Choose the first image that appears to be a generation example (not a logo or diagram). Construct the absolute URL as `https://huggingface.co/{{repo}}/resolve/main/{filename}`. If no suitable image is found, return an empty string.
|
||||
|
||||
### notes
|
||||
A plain-text summary of the model card's key practical usage information. Combine trigger words, style modifiers, recommended parameters (steps, CFG, resolution, sampler), and any setup tips into a readable paragraph. Return empty string if the README has no useful usage info.
|
||||
|
||||
### usage_tips
|
||||
A JSON string with structured usage recommendations. Extract from the README any explicit ranges or recommended values (e.g. "Set LoRA strength: **0.85 - 1.4**", "CLIP strength: 0.5"). Possible fields (include only those you can determine):
|
||||
|
||||
```json
|
||||
{
|
||||
"strength_min": 0.85,
|
||||
"strength_max": 1.4,
|
||||
"strength_range": "0.85-1.4",
|
||||
"strength": 0.6,
|
||||
"clip_strength": 0.5,
|
||||
"clip_skip": 2
|
||||
}
|
||||
```
|
||||
|
||||
Return the JSON string (e.g. `'{"strength_min":0.85,"strength_max":1.4}'`). Return `"{}"` if nothing useful is found.
|
||||
|
||||
### confidence
|
||||
Your confidence level in the extracted data:
|
||||
- "high" — most fields were explicitly stated in the README
|
||||
@@ -104,6 +123,8 @@ Return ONLY a JSON object with exactly these fields (no markdown fences, no extr
|
||||
"recommended_width": 768,
|
||||
"recommended_height": 1024,
|
||||
"preview_url": "<image URL or empty string>",
|
||||
"notes": "<plain-text usage summary or empty string>",
|
||||
"usage_tips": "<JSON string like '{\"strength_min\":0.85,\"strength_max\":1.4}' or '{}'>",
|
||||
"confidence": "<high|medium|low>"
|
||||
}
|
||||
```
|
||||
|
||||
@@ -1,8 +1,13 @@
|
||||
"""Inline markdown-to-HTML converter for HF README content.
|
||||
"""Inline markdown-to-HTML converter and LLM-prompt cleaner for HF README content.
|
||||
|
||||
No external dependencies. Strips YAML frontmatter, ``<Gallery />`` sections,
|
||||
badge images, and HTML comments before rendering. Only used by the
|
||||
``enrich_hf_metadata`` skill.
|
||||
|
||||
Also provides :func:`clean_readme_for_llm` which pre-processes the raw README
|
||||
before it is injected into the LLM prompt, removing content that has zero value
|
||||
for metadata extraction (widget sections, code blocks, training tables,
|
||||
boilerplate, massive lists, etc.).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -118,6 +123,88 @@ def extract_gallery_images(
|
||||
return images
|
||||
|
||||
|
||||
def extract_gallery_table_images(
|
||||
markdown_text: str,
|
||||
repo: str,
|
||||
existing_urls: set | None = None,
|
||||
default_width: int = 512,
|
||||
default_height: int = 512,
|
||||
) -> list[dict]:
|
||||
"""Extract images from ``| Preview | Prompt |`` markdown gallery tables.
|
||||
|
||||
Many HF READMEs include a sample-gallery table in the body (outside
|
||||
the YAML frontmatter) that shows generation examples with their
|
||||
prompts. This function parses those tables and merges results with
|
||||
the widget-sourced images from :func:`extract_gallery_images`.
|
||||
|
||||
Returns a list of dicts in the same ``civitai.images`` format as
|
||||
:func:`extract_gallery_images`. Already-seen URLs (from *existing_urls*)
|
||||
are skipped.
|
||||
"""
|
||||
if not markdown_text or not repo:
|
||||
return []
|
||||
|
||||
base_url = f"https://huggingface.co/{repo}/resolve/main"
|
||||
images: list[dict] = []
|
||||
seen_urls: set = set(existing_urls) if existing_urls else set()
|
||||
lines = markdown_text.split("\n")
|
||||
n = len(lines)
|
||||
i = 0
|
||||
|
||||
while i < n:
|
||||
line = lines[i]
|
||||
if "|" not in line or i + 1 >= n:
|
||||
i += 1
|
||||
continue
|
||||
|
||||
# Check for table separator row
|
||||
if not re.match(r"^\|[\s:-]+\|", lines[i + 1]):
|
||||
i += 1
|
||||
continue
|
||||
|
||||
header_lower = line.strip().lower()
|
||||
first_cell = header_lower.strip("|").split("|")[0].strip() if "|" in header_lower else ""
|
||||
is_gallery = any(kw in first_cell for kw in ("preview", "sample", "gallery", "image", "thumbnail"))
|
||||
if not is_gallery:
|
||||
i += 1
|
||||
continue
|
||||
|
||||
# Skip header + separator
|
||||
i += 2
|
||||
while i < n and "|" in lines[i]:
|
||||
cells = [c.strip() for c in lines[i].strip().strip("|").split("|")]
|
||||
if len(cells) >= 2:
|
||||
first = cells[0]
|
||||
prompt = cells[1]
|
||||
|
||||
url_match = re.search(r"!\[([^\]]*)\]\(([^)]+)\)", first)
|
||||
if url_match:
|
||||
raw_path = url_match.group(2)
|
||||
if raw_path.startswith("http"):
|
||||
url = raw_path
|
||||
else:
|
||||
# Normalise: remove leading / and ./ prefixes
|
||||
clean = raw_path.lstrip("./").lstrip("/")
|
||||
url = f"{base_url}/{clean}"
|
||||
|
||||
if url not in seen_urls:
|
||||
seen_urls.add(url)
|
||||
images.append({
|
||||
"url": url,
|
||||
"type": "image",
|
||||
"nsfwLevel": 0,
|
||||
"width": default_width,
|
||||
"height": default_height,
|
||||
"meta": {"prompt": prompt, "negativePrompt": ""},
|
||||
"hasMeta": bool(prompt),
|
||||
"hasPositivePrompt": bool(prompt),
|
||||
})
|
||||
i += 1
|
||||
continue
|
||||
|
||||
return images
|
||||
|
||||
|
||||
def _extract_frontmatter(text: str) -> str:
|
||||
"""Return the YAML frontmatter content (without the ``---`` delimiters).
|
||||
|
||||
@@ -145,7 +232,260 @@ def convert_readme_to_html(markdown_text: str | None) -> str:
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Pre-processing: strip unwanted sections
|
||||
# README cleaning for LLM prompt injection
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
#: Section headers that signal boilerplate content with zero metadata value.
|
||||
_BOILERPLATE_HEADERS: tuple[str, ...] = (
|
||||
"download model",
|
||||
"license",
|
||||
"citation",
|
||||
"links",
|
||||
"disclaimer",
|
||||
"architecture notes",
|
||||
"training details",
|
||||
"dataset",
|
||||
"provenance",
|
||||
)
|
||||
|
||||
#: Table header keywords that identify training-parameter tables.
|
||||
_TRAINING_PARAM_KEYWORDS: tuple[str, ...] = (
|
||||
"lr scheduler",
|
||||
"optimizer",
|
||||
"network dim",
|
||||
"network alpha",
|
||||
"noise offset",
|
||||
"multires noise",
|
||||
"repeat",
|
||||
"epoch",
|
||||
"batch size",
|
||||
"gradient accumulation",
|
||||
"learning rate",
|
||||
"rslora",
|
||||
"dtype",
|
||||
)
|
||||
|
||||
#: Maximum chars before a single-line comma list is considered massive.
|
||||
_MASSIVE_LIST_LINE_MIN_LEN = 150
|
||||
#: Minimum consecutive enumeration lines to trigger massive-list stripping.
|
||||
_MASSIVE_LIST_THRESHOLD = 8
|
||||
|
||||
|
||||
def clean_readme_for_llm(markdown_text: str | None, max_length: int = 6000) -> str:
|
||||
"""Clean a HF README for injection into an LLM metadata-extraction prompt.
|
||||
|
||||
Removes content that carries no signal for inferring base model,
|
||||
trigger words, short description, tags, or a preview image URL:
|
||||
|
||||
* ``widget:`` YAML block (example prompts + output URLs)
|
||||
* ``<Gallery />`` tags and wrappers
|
||||
* Fenced code blocks (Python / bash / bibtex / yaml)
|
||||
* Standalone ```` image lines and ``<img>`` tags
|
||||
* Training-parameter tables
|
||||
* Boilerplate sections (Download / License / Citation / …)
|
||||
* Massive enumeration lists (e.g. 3000+ celebrity names)
|
||||
|
||||
The post-processor still receives the **full** raw README via
|
||||
``readme_content_full``, so nothing is lost for HTML conversion or
|
||||
gallery-image extraction.
|
||||
|
||||
Args:
|
||||
markdown_text: Raw README.md content from HuggingFace.
|
||||
max_length: Hard ceiling on output length (default 6 000 chars).
|
||||
|
||||
Returns:
|
||||
Cleaned markdown, truncated to *max_length*.
|
||||
"""
|
||||
if not markdown_text:
|
||||
return ""
|
||||
|
||||
text = markdown_text
|
||||
|
||||
# Order matters — broader strips first, then finer ones.
|
||||
text = _strip_gallery(text)
|
||||
text = _strip_fenced_code_blocks(text)
|
||||
text = _strip_standalone_images(text)
|
||||
text = _strip_training_tables(text)
|
||||
text = _strip_boilerplate_sections(text)
|
||||
text = _strip_massive_lists(text)
|
||||
text = _strip_badge_images(text)
|
||||
text = _strip_html_comments(text)
|
||||
text = _compress_blank_lines(text)
|
||||
|
||||
if len(text) > max_length:
|
||||
text = text[:max_length]
|
||||
|
||||
return text.strip()
|
||||
|
||||
|
||||
def _strip_fenced_code_blocks(text: str) -> str:
|
||||
"""Strip fenced code blocks that have an explicit programming-language tag.
|
||||
|
||||
Blocks without a language tag (just `` ``` ``) are preserved — they
|
||||
often contain trigger words, example prompts, or config snippets
|
||||
rather than actual runnable code.
|
||||
"""
|
||||
# Match opening ``` immediately followed by a word character (the language
|
||||
# tag), then any content, then closing ```. Plain ``` at the start of a
|
||||
# line is left intact. A leading \n is optional (handles blocks at the
|
||||
# start of the text).
|
||||
return re.sub(
|
||||
r"(?:\n|^)```[a-zA-Z_][a-zA-Z0-9_]*\s*\n.*?\n```",
|
||||
"",
|
||||
text,
|
||||
flags=re.DOTALL,
|
||||
)
|
||||
|
||||
|
||||
def _strip_standalone_images(text: str) -> str:
|
||||
"""Strip image embeds that occupy their own line.
|
||||
|
||||
Preserves the alt text from markdown images (```` → ``alt``)
|
||||
since it often describes what the model generates, which is useful signal
|
||||
for tag/description extraction.
|
||||
"""
|
||||
# Markdown: ```` on its own line → keep alt text
|
||||
text = re.sub(
|
||||
r"^\s*!\[([^\]]*)\]\([^)]+\)\s*$",
|
||||
r"\1",
|
||||
text,
|
||||
flags=re.MULTILINE,
|
||||
)
|
||||
# HTML: ``<img src="..." ...>`` on its own line → remove entirely
|
||||
text = re.sub(
|
||||
r'^\s*<img\s[^>]+/?>(?:</img>)?\s*$',
|
||||
"",
|
||||
text,
|
||||
flags=re.MULTILINE | re.IGNORECASE,
|
||||
)
|
||||
return text
|
||||
|
||||
|
||||
def _strip_training_tables(text: str) -> str:
|
||||
"""Strip markdown tables whose header row mentions training parameters.
|
||||
|
||||
Checks the header row (first line of a detected table) against
|
||||
``_TRAINING_PARAM_KEYWORDS``. Non-training tables (e.g. "Best
|
||||
Dimensions") are preserved.
|
||||
"""
|
||||
lines = text.split("\n")
|
||||
out: list[str] = []
|
||||
i = 0
|
||||
n = len(lines)
|
||||
|
||||
while i < n:
|
||||
line = lines[i]
|
||||
if "|" in line and i + 1 < n and re.match(r"^\|[\s:-]+\|", lines[i + 1]):
|
||||
table_lines = [line]
|
||||
i += 1
|
||||
while i < n and "|" in lines[i]:
|
||||
table_lines.append(lines[i])
|
||||
i += 1
|
||||
|
||||
# Check header + first data row for training keywords
|
||||
header_and_first = (line + "\n" + (table_lines[2] if len(table_lines) > 2 else "")).lower()
|
||||
if any(kw in header_and_first for kw in _TRAINING_PARAM_KEYWORDS):
|
||||
continue
|
||||
out.extend(table_lines)
|
||||
else:
|
||||
out.append(line)
|
||||
i += 1
|
||||
|
||||
return "\n".join(out)
|
||||
|
||||
|
||||
def _strip_boilerplate_sections(text: str) -> str:
|
||||
"""Strip sections whose headings match known boilerplate patterns.
|
||||
|
||||
When a heading (``## Download model``, ``## License``, etc.) is
|
||||
detected, the heading and all content until the next heading of
|
||||
equal-or-higher level is removed.
|
||||
"""
|
||||
lines = text.split("\n")
|
||||
out: list[str] = []
|
||||
i = 0
|
||||
n = len(lines)
|
||||
skip_until_level: int | None = None
|
||||
|
||||
while i < n:
|
||||
line = lines[i]
|
||||
h_match = re.match(r"^(#{1,4})\s+(.+?)\s*#*$", line)
|
||||
if h_match:
|
||||
level = len(h_match.group(1))
|
||||
title = h_match.group(2).strip().lower()
|
||||
|
||||
is_boilerplate = any(
|
||||
title == kw or title.startswith(kw + " ") or title.startswith(kw + ":")
|
||||
for kw in _BOILERPLATE_HEADERS
|
||||
)
|
||||
|
||||
if is_boilerplate:
|
||||
skip_until_level = level
|
||||
i += 1
|
||||
continue
|
||||
|
||||
if skip_until_level is not None and level <= skip_until_level:
|
||||
skip_until_level = None
|
||||
|
||||
if skip_until_level is None:
|
||||
out.append(line)
|
||||
i += 1
|
||||
|
||||
return "\n".join(out)
|
||||
|
||||
|
||||
def _strip_massive_lists(text: str) -> str:
|
||||
"""Strip blocks of 8+ consecutive enumeration-style lines.
|
||||
|
||||
Targets long comma-separated name lists (e.g. the 3000+ celebrity
|
||||
names in some Z-Image READMEs) and dense bullet enumerations.
|
||||
"""
|
||||
lines = text.split("\n")
|
||||
out: list[str] = []
|
||||
i = 0
|
||||
n = len(lines)
|
||||
|
||||
while i < n:
|
||||
stripped = lines[i].strip()
|
||||
|
||||
# A "list-like" line ends with comma or is a bullet with commas
|
||||
is_list_like = bool(stripped) and (
|
||||
stripped.endswith(",")
|
||||
or len(stripped) >= _MASSIVE_LIST_LINE_MIN_LEN
|
||||
or (bool(re.match(r"^[-*+]\s", stripped)) and "," in stripped)
|
||||
)
|
||||
|
||||
if is_list_like:
|
||||
count = 1
|
||||
j = i + 1
|
||||
while j < n:
|
||||
s = lines[j].strip()
|
||||
if not s:
|
||||
j += 1
|
||||
continue
|
||||
if s.endswith(",") or (bool(re.match(r"^[-*+]\s", s)) and "," in s):
|
||||
count += 1
|
||||
j += 1
|
||||
else:
|
||||
break
|
||||
|
||||
if count >= _MASSIVE_LIST_THRESHOLD:
|
||||
i = j
|
||||
continue
|
||||
|
||||
out.append(lines[i])
|
||||
i += 1
|
||||
|
||||
return "\n".join(out)
|
||||
|
||||
|
||||
def _compress_blank_lines(text: str) -> str:
|
||||
"""Collapse runs of 3+ blank lines down to 2."""
|
||||
return re.sub(r"\n{3,}", "\n\n", text)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Pre-processing: strip unwanted sections (HTML conversion helpers)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user