fixes
This commit is contained in:
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user