This commit is contained in:
2026-06-28 17:12:14 -07:00
parent 66d085ad96
commit c718a80651

View File

@@ -172,7 +172,56 @@ def _fetch_json(url, timeout=10):
return json.loads(resp.read().decode())
def lookup_source(filepath, filename, cache):
# URL pattern for extracting embedded links from free-text fields
_URL_RE = re.compile(r'https?://\S+', re.IGNORECASE)
# Metadata fields that may contain a source URL, checked in order
_META_URL_FIELDS = [
"modelspec.source",
"source",
"url",
"model_url",
"ss_training_comment",
"note",
"description",
"comment",
]
# Map hostname → display label
_SOURCE_LABELS = {
"civitai.com": "CivitAI",
"huggingface.co": "HuggingFace",
"modelscope.cn": "ModelScope",
"tensor.art": "Tensor.Art",
}
def _label_for_url(url):
from urllib.parse import urlparse
host = urlparse(url).netloc.lstrip("www.")
for domain, label in _SOURCE_LABELS.items():
if host == domain or host.endswith("." + domain):
return label
return host or "Link"
def _extract_meta_url(meta):
"""Return (url, label) from embedded metadata fields, or (None, None)."""
if not meta:
return None, None
for field in _META_URL_FIELDS:
val = meta.get(field)
if not val or not isinstance(val, str):
continue
# Field may be a bare URL or free text containing one
m = _URL_RE.search(val.strip())
if m:
url = m.group(0).rstrip(".,;)")
return url, _label_for_url(url)
return None, None
def lookup_source(filepath, filename, cache, meta=None):
now = int(time.time())
cached = cache.get(filename)
if cached:
@@ -182,6 +231,12 @@ def lookup_source(filepath, filename, cache):
if (now - last) < 86400 * 30:
return (url, label) if url else (None, None)
# 0. Embedded URL in safetensors __metadata__ (free, no network call)
url, label = _extract_meta_url(meta)
if url:
cache[filename] = {"url": url, "label": label, "last_checked": now}
return url, label
# CivitAI model-version lookup by BLAKE2b file hash
try:
h = file_blake2b_hash(filepath)
@@ -227,6 +282,30 @@ def lookup_source(filepath, filename, cache):
except (HTTPError, URLError, json.JSONDecodeError, OSError):
pass
# Fallback: ModelScope model search by filename
try:
q = (
filename.replace(".safetensors", "")
.replace("_", " ")
.replace("-", " ")
)
data = _fetch_json(
f"https://www.modelscope.cn/api/v1/models?Name={quote(q)}&PageSize=3&PageNumber=1"
)
models = (data or {}).get("Data", {}).get("Models") or []
if models:
m = models[0]
model_id = m.get("Path") or m.get("Name", "")
if model_id:
url = f"https://www.modelscope.cn/models/{model_id}"
cache[filename] = {
"url": url, "label": "ModelScope",
"model_name": model_id, "last_checked": now,
}
return url, "ModelScope"
except (HTTPError, URLError, json.JSONDecodeError, OSError):
pass
cache[filename] = {"url": None, "label": None, "last_checked": now}
return None, None
@@ -472,9 +551,9 @@ def main():
cache = load_source_cache()
sources = {}
total = len(loras)
for i, (fname, _) in enumerate(loras, 1):
for i, (fname, meta) in enumerate(loras, 1):
fpath = os.path.join(target, fname)
url, label = lookup_source(fpath, fname, cache)
url, label = lookup_source(fpath, fname, cache, meta)
if url:
sources[fname] = (url, label)
if i % 10 == 0 or i == total: