1277 lines
50 KiB
Python
1277 lines
50 KiB
Python
"""
|
|
Model search API for CivitAI, Hugging Face, and GitHub.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import os
|
|
import re
|
|
from typing import Any, Dict, List, Optional, Tuple
|
|
from urllib.parse import quote, urlencode, urlparse
|
|
from urllib.request import HTTPRedirectHandler, Request, build_opener
|
|
import json
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
import time
|
|
import threading
|
|
|
|
# --- Hugging Face File Search ---
|
|
|
|
# Whitelist of HF repos that are known to host useful model files directly
|
|
# (as opposed to being 'model cards' that point to other repos).
|
|
# This is a fallback search mechanism.
|
|
HF_FILE_SEARCH_REPOS = [
|
|
"Kijai/",
|
|
"GitMylo/",
|
|
"Comfy-Org/",
|
|
"mcmonkey/",
|
|
"City96/",
|
|
"stabilityai/sd-vae-ft-mse-original",
|
|
"stabilityai/sdxl-vae",
|
|
# Add other repos known to host raw .safetensors/.ckpt files here
|
|
]
|
|
|
|
# Cache for Hugging Face tree listings to avoid redundant API calls.
|
|
# Format: { "repo_id": (timestamp, file_list) }
|
|
_hf_tree_cache: Dict[str, Tuple[float, List[Dict[str, Any]]]] = {}
|
|
_hf_tree_cache_lock = threading.Lock()
|
|
HF_TREE_CACHE_TTL = 600 # 10 minutes
|
|
|
|
SEARCH_TIMEOUT = 30 # seconds - increased for more thorough search
|
|
|
|
|
|
class _NoAuthCrossHostRedirects(HTTPRedirectHandler):
|
|
def redirect_request(self, req, fp, code, msg, headers, newurl):
|
|
newreq = super().redirect_request(req, fp, code, msg, headers, newurl)
|
|
if newreq is None:
|
|
return None
|
|
try:
|
|
old_host = urlparse(req.full_url).hostname or ""
|
|
new_host = urlparse(newreq.full_url).hostname or ""
|
|
except Exception:
|
|
old_host, new_host = "", ""
|
|
if old_host and new_host and old_host.lower() != new_host.lower():
|
|
if "Authorization" in newreq.headers:
|
|
del newreq.headers["Authorization"]
|
|
return newreq
|
|
|
|
|
|
def calculate_match_score(query: str, candidate: str, filename: str = "") -> Tuple[float, str]:
|
|
"""
|
|
Calculate how well a candidate matches the query.
|
|
Handles special characters, underscores, and long names.
|
|
|
|
Returns:
|
|
Tuple of (score from 0-100, match_level description)
|
|
"""
|
|
if not query or not candidate:
|
|
return 0.0, "No match"
|
|
|
|
query_lower = query.lower().strip()
|
|
candidate_lower = candidate.lower().strip()
|
|
filename_lower = filename.lower().strip()
|
|
|
|
# Normalize: remove common separators and extra spaces
|
|
def normalize(text):
|
|
# Remove extensions
|
|
text = re.sub(r'\.(safetensors|ckpt|pt|pth|bin)$', '', text, flags=re.IGNORECASE)
|
|
# Replace ALL separators (hyphens, underscores, dots) with spaces
|
|
text = re.sub(r'[-_\.]+', ' ', text)
|
|
# Remove special characters except alphanumeric and spaces
|
|
text = re.sub(r'[^a-z0-9\s]', ' ', text)
|
|
# Remove extra spaces
|
|
text = re.sub(r'\s+', ' ', text)
|
|
return text.strip()
|
|
|
|
query_norm = normalize(query_lower)
|
|
candidate_norm = normalize(candidate_lower)
|
|
filename_norm = normalize(filename_lower)
|
|
|
|
# Exact match (100%)
|
|
if query_norm == candidate_norm or query_norm == filename_norm:
|
|
return 100.0, "Exact match"
|
|
|
|
# Check if query is contained in candidate or filename
|
|
# Only give high score if query is significant portion of candidate
|
|
if query_norm in candidate_norm or query_norm in filename_norm:
|
|
# Calculate ratio to avoid false positives
|
|
target = candidate_norm if query_norm in candidate_norm else filename_norm
|
|
ratio = len(query_norm) / max(len(target), 1)
|
|
if ratio > 0.7: # Query is most of the candidate
|
|
return 95.0, "Near exact match"
|
|
elif ratio > 0.4: # Query is significant part
|
|
return 90.0, "Contains query"
|
|
else: # Query is small part of candidate
|
|
return 70.0, "Contains query (partial)"
|
|
|
|
# Check if candidate contains query - usually means query is too broad or mismatched
|
|
if candidate_norm in query_norm or filename_norm in query_norm:
|
|
# Calculate how much of query is the candidate
|
|
source = query_norm if candidate_norm in query_norm else query_norm
|
|
target = candidate_norm if candidate_norm in query_norm else filename_norm
|
|
ratio = len(target) / max(len(source), 1)
|
|
if ratio > 0.5: # Candidate is significant part of query
|
|
return 60.0, "Subset match"
|
|
else:
|
|
# Likely a false positive (e.g., "vision" in "completely_wrong_name")
|
|
return 20.0, "Weak match"
|
|
|
|
# Word-by-word matching: Give high scores for near-complete matches
|
|
query_words = set(query_norm.split())
|
|
candidate_words = set(candidate_norm.split())
|
|
filename_words = set(filename_norm.split())
|
|
|
|
all_words = candidate_words | filename_words
|
|
|
|
# Case 1: All query words are present in the candidate/filename.
|
|
# This handles when the query is missing words (e.g., 'visual').
|
|
if query_words and query_words.issubset(all_words):
|
|
extra_words = len(all_words - query_words)
|
|
# Few extra words means a very strong match.
|
|
if extra_words <= 2:
|
|
return 96.0 - (extra_words * 3), "Near-exact word match"
|
|
# More extra words, still a good match.
|
|
else:
|
|
return 88.0, "All query words match"
|
|
|
|
# Case 2: All candidate/filename words are present in the query.
|
|
# This handles when the query has extra, irrelevant words.
|
|
if all_words and all_words.issubset(query_words):
|
|
extra_words = len(query_words - all_words)
|
|
if extra_words <= 2:
|
|
return 94.0 - (extra_words * 3), "Candidate is subset of query"
|
|
else:
|
|
return 85.0, "Subset word match"
|
|
|
|
# Case 3: Partial overlap, use Jaccard similarity.
|
|
if query_words or all_words:
|
|
common_words = query_words.intersection(all_words)
|
|
jaccard_sim = len(common_words) / len(query_words.union(all_words)) if (query_words or all_words) else 0
|
|
|
|
if jaccard_sim >= 0.7:
|
|
return 80.0, "High word similarity"
|
|
elif jaccard_sim >= 0.5:
|
|
return 70.0, "Good word similarity"
|
|
elif jaccard_sim > 0.3:
|
|
return 60.0, "Some words match"
|
|
|
|
# Fuzzy substring match (30-50%)
|
|
# Check for partial word matches
|
|
partial_score = 0
|
|
for q_word in query_words:
|
|
for c_word in all_words:
|
|
if len(q_word) >= 4 and q_word in c_word:
|
|
partial_score += 10
|
|
elif len(c_word) >= 4 and c_word in q_word:
|
|
partial_score += 8
|
|
|
|
if partial_score > 0:
|
|
score = min(50.0, 30.0 + partial_score)
|
|
return score, "Partial match"
|
|
|
|
# Very low or no match
|
|
return 10.0, "Poor match"
|
|
|
|
|
|
def _make_request(url: str, headers: Optional[Dict[str, str]] = None) -> Dict[str, Any]:
|
|
"""Make HTTP request and return JSON response."""
|
|
if headers is None:
|
|
headers = {}
|
|
headers.setdefault("User-Agent", "ComfyUI-Majoor-Downloader")
|
|
|
|
try:
|
|
request = Request(url, headers=headers)
|
|
opener = build_opener(_NoAuthCrossHostRedirects())
|
|
with opener.open(request, timeout=SEARCH_TIMEOUT) as resp:
|
|
data = resp.read()
|
|
return json.loads(data.decode("utf-8"))
|
|
except Exception as e:
|
|
logger.warning("Request failed for %s: %s", url, str(e))
|
|
return {}
|
|
|
|
|
|
def search_civitai(query: str, limit: int = 5) -> List[Dict[str, Any]]:
|
|
"""
|
|
Search for models on CivitAI.
|
|
|
|
Args:
|
|
query: Model name to search for
|
|
limit: Maximum number of results
|
|
|
|
Returns:
|
|
List of results with format:
|
|
{
|
|
"platform": "civitai",
|
|
"name": str,
|
|
"url": str, # download URL
|
|
"page_url": str, # web page URL
|
|
"type": str, # model type (checkpoint, lora, etc)
|
|
"version": str,
|
|
"size_mb": int,
|
|
"sha256": str (optional)
|
|
}
|
|
"""
|
|
results = []
|
|
|
|
try:
|
|
# CivitAI API search endpoint
|
|
params = {
|
|
"query": query,
|
|
"limit": limit,
|
|
"nsfw": "false"
|
|
}
|
|
url = f"https://civitai.com/api/v1/models?{urlencode(params)}"
|
|
|
|
data = _make_request(url)
|
|
items = data.get("items", [])
|
|
|
|
for item in items[:limit]:
|
|
model_id = item.get("id")
|
|
model_name = item.get("name", "")
|
|
model_type = item.get("type", "").lower()
|
|
|
|
# Get latest version
|
|
versions = item.get("modelVersions", [])
|
|
if not versions:
|
|
continue
|
|
|
|
latest = versions[0]
|
|
version_name = latest.get("name", "")
|
|
files = latest.get("files", [])
|
|
|
|
# Find primary file
|
|
primary_file = None
|
|
for f in files:
|
|
if f.get("primary", False):
|
|
primary_file = f
|
|
break
|
|
|
|
if not primary_file and files:
|
|
primary_file = files[0]
|
|
|
|
if not primary_file:
|
|
continue
|
|
|
|
download_url = primary_file.get("downloadUrl", "")
|
|
if not download_url:
|
|
continue
|
|
|
|
# Map CivitAI types to ComfyUI kinds
|
|
type_mapping = {
|
|
"checkpoint": "checkpoints",
|
|
"lora": "loras",
|
|
"textualinversion": "embeddings",
|
|
"hypernetwork": "loras",
|
|
"aestheticgradient": "loras",
|
|
"controlnet": "controlnet",
|
|
"vae": "vae",
|
|
"upscaler": "upscale_models",
|
|
}
|
|
kind = type_mapping.get(model_type, "checkpoints")
|
|
|
|
size_kb = primary_file.get("sizeKB", 0)
|
|
size_mb = int(size_kb / 1024) if size_kb else 0
|
|
|
|
# Get hashes
|
|
hashes = primary_file.get("hashes", {})
|
|
sha256 = hashes.get("SHA256", "")
|
|
|
|
# Calculate match score
|
|
full_name = f"{model_name} - {version_name}"
|
|
score, match_level = calculate_match_score(query, full_name, primary_file.get("name", ""))
|
|
|
|
results.append({
|
|
"platform": "civitai",
|
|
"name": full_name,
|
|
"filename": primary_file.get("name", ""),
|
|
"url": download_url,
|
|
"page_url": f"https://civitai.com/models/{model_id}",
|
|
"type": kind,
|
|
"version": version_name,
|
|
"size_mb": size_mb,
|
|
"sha256": sha256.lower() if sha256 else None,
|
|
"match_score": score,
|
|
"match_level": match_level,
|
|
})
|
|
|
|
except Exception as e:
|
|
logger.error("CivitAI search failed: %s", str(e))
|
|
|
|
# Sort by match score (highest first) and filter out poor matches
|
|
results.sort(key=lambda x: x.get("match_score", 0), reverse=True)
|
|
# Only return results with score >= 30 (configurable threshold)
|
|
results = [r for r in results if r.get("match_score", 0) >= 30]
|
|
|
|
return results[:limit]
|
|
|
|
|
|
def _get_official_sd_models() -> Dict[str, Dict[str, Any]]:
|
|
"""
|
|
Return known official Stable Diffusion models on Hugging Face.
|
|
This helps find models that might not appear in search results.
|
|
"""
|
|
return {
|
|
"v1-5-pruned-emaonly": {
|
|
"model_id": "runwayml/stable-diffusion-v1-5",
|
|
"filename": "v1-5-pruned-emaonly.safetensors",
|
|
"alt_filename": "v1-5-pruned-emaonly.ckpt",
|
|
"name": "Stable Diffusion v1.5",
|
|
},
|
|
"v1-5-pruned": {
|
|
"model_id": "runwayml/stable-diffusion-v1-5",
|
|
"filename": "v1-5-pruned.safetensors",
|
|
"name": "Stable Diffusion v1.5",
|
|
},
|
|
"v2-1_768-ema-pruned": {
|
|
"model_id": "stabilityai/stable-diffusion-2-1",
|
|
"filename": "v2-1_768-ema-pruned.safetensors",
|
|
"name": "Stable Diffusion v2.1",
|
|
},
|
|
"v2-1_512-ema-pruned": {
|
|
"model_id": "stabilityai/stable-diffusion-2-1-base",
|
|
"filename": "v2-1_512-ema-pruned.safetensors",
|
|
"name": "Stable Diffusion v2.1 Base",
|
|
},
|
|
"sd_xl_base_1.0": {
|
|
"model_id": "stabilityai/stable-diffusion-xl-base-1.0",
|
|
"filename": "sd_xl_base_1.0.safetensors",
|
|
"name": "Stable Diffusion XL Base 1.0",
|
|
},
|
|
"sd_xl_refiner_1.0": {
|
|
"model_id": "stabilityai/stable-diffusion-xl-refiner-1.0",
|
|
"filename": "sd_xl_refiner_1.0.safetensors",
|
|
"name": "Stable Diffusion XL Refiner 1.0",
|
|
},
|
|
}
|
|
|
|
|
|
def canonicalize_hf_url(url: str) -> str:
|
|
"""Converts a Hugging Face blob URL to a resolve URL for direct downloads."""
|
|
if not isinstance(url, str):
|
|
return url
|
|
|
|
# Example: https://huggingface.co/Kijai/WanVideo_comfy/blob/main/open-clip-xlm-roberta-large-vit-huge-14_visual_fp16.safetensors
|
|
# -> https://huggingface.co/Kijai/WanVideo_comfy/resolve/main/open-clip-xlm-roberta-large-vit-huge-14_visual_fp16.safetensors
|
|
return re.sub(r'huggingface\.co/([^/]+)/([^/]+)/blob/', r'huggingface.co/\1/\2/resolve/', url)
|
|
|
|
def _expand_hf_repo_wildcards(
|
|
wildcards: List[str], query: str, limit: int
|
|
) -> List[str]:
|
|
"""
|
|
Expand repo wildcards like "Kijai/" or "Kijai/*" into concrete repo ids.
|
|
|
|
Uses Hugging Face models search API and caps results to keep the fallback
|
|
search bounded.
|
|
"""
|
|
if not wildcards:
|
|
return []
|
|
q = str(query or "").strip()
|
|
if len(q) < 4:
|
|
return []
|
|
|
|
expanded: List[str] = []
|
|
seen = set()
|
|
|
|
# Small cap to prevent excessive API calls / slowdowns.
|
|
per_author_cap = max(1, min(5, int(limit) * 2))
|
|
|
|
for wc in wildcards:
|
|
wc = str(wc or "").strip()
|
|
if not wc:
|
|
continue
|
|
author = wc.rstrip("/").rstrip("*").strip()
|
|
if not author:
|
|
continue
|
|
|
|
# Prefer author-filtered search if supported; otherwise filter client-side.
|
|
urls = [
|
|
f"https://huggingface.co/api/models?{urlencode({'search': q, 'author': author, 'limit': per_author_cap})}",
|
|
f"https://huggingface.co/api/models?{urlencode({'search': q, 'limit': per_author_cap * 2})}",
|
|
]
|
|
|
|
data: Any = []
|
|
for url in urls:
|
|
data = _make_request(url)
|
|
if isinstance(data, list) and data:
|
|
break
|
|
|
|
if not isinstance(data, list):
|
|
continue
|
|
|
|
for item in data:
|
|
model_id = (item or {}).get("id")
|
|
if not isinstance(model_id, str) or not model_id:
|
|
continue
|
|
if not model_id.startswith(f"{author}/"):
|
|
continue
|
|
if model_id in seen:
|
|
continue
|
|
seen.add(model_id)
|
|
expanded.append(model_id)
|
|
if len(expanded) >= per_author_cap:
|
|
break
|
|
|
|
return expanded[:per_author_cap]
|
|
|
|
|
|
def _search_duckduckgo_repos(query: str, limit: int = 5) -> List[str]:
|
|
"""
|
|
Search DuckDuckGo for Hugging Face repositories matching the query.
|
|
Acts as a 'Google fallback' to find repos not indexed by HF API or whitelists.
|
|
"""
|
|
candidates = []
|
|
try:
|
|
# Search global web with DuckDuckGo
|
|
# We search specifically for the term to find download pages
|
|
search_q = quote(query)
|
|
url = f"https://html.duckduckgo.com/html/?q={search_q}"
|
|
|
|
# Headers to strictly mimic a browser to avoid 403s
|
|
headers = {
|
|
"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36",
|
|
"Accept": "text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,*/*;q=0.8",
|
|
"Referer": "https://html.duckduckgo.com/",
|
|
"Accept-Language": "en-US,en;q=0.5"
|
|
}
|
|
|
|
req = Request(url, headers=headers)
|
|
with build_opener().open(req, timeout=15) as resp:
|
|
html = resp.read().decode("utf-8", errors="ignore")
|
|
|
|
# Regex to find huggingface.co links
|
|
# Matches: huggingface.co/User/Repo
|
|
# Excludes: /blob/, /resolve/, /tree/, /api/, etc. by pattern matching
|
|
|
|
# We look for ANY huggingface.co link in the search results
|
|
import re
|
|
|
|
# Find all HF links
|
|
# Pattern extracts User and Repo from URLs
|
|
matches = re.findall(r'huggingface\.co/([^/"\'\s<>?#&]+)/([^/"\'\s<>?#&]+)', html)
|
|
|
|
ignored_users = {
|
|
"search", "datasets", "spaces", "api", "login", "join",
|
|
"docs", "posts", "models", "users", "settings", "blog", "chat", "pricing"
|
|
}
|
|
|
|
for user, repo in matches:
|
|
# Basic validation
|
|
if user in ignored_users: continue
|
|
if repo in {"blob", "resolve", "tree", "main", "full", "discussions"}: continue
|
|
|
|
full_id = f"{user}/{repo}"
|
|
if full_id not in candidates:
|
|
candidates.append(full_id)
|
|
if len(candidates) >= limit:
|
|
break
|
|
|
|
if candidates:
|
|
logger.info(f"DuckDuckGo found potential HF repos: {candidates}")
|
|
|
|
except Exception as e:
|
|
logger.warning(f"DuckDuckGo fallback search failed: {e}")
|
|
|
|
return candidates
|
|
|
|
|
|
def _search_huggingface_repo_files(query: str, limit: int = 5, extra_repos: List[str] = None) -> List[Dict[str, Any]]:
|
|
"""
|
|
Fallback search: Scans whitelisted AND dynamically found Hugging Face repos for individual files.
|
|
This is effective for repos that are collections of files rather than a single model.
|
|
"""
|
|
if len(query) < 4:
|
|
return []
|
|
|
|
logger.info(f"Running fallback Hugging Face file search for query: {query}")
|
|
results = []
|
|
|
|
# Start with base whitelist
|
|
base_repos = list(HF_FILE_SEARCH_REPOS)
|
|
|
|
# Add any extra repos found dynamically (e.g. from broad search)
|
|
if extra_repos:
|
|
base_repos.extend(extra_repos)
|
|
|
|
wildcard_specs = [
|
|
r for r in base_repos if isinstance(r, str) and (r.endswith("/") or r.endswith("/*"))
|
|
]
|
|
concrete_repos = [r for r in base_repos if r not in wildcard_specs]
|
|
|
|
# Expand wildcards (e.g., "Kijai/") into a small set of repos relevant to the query.
|
|
expanded = _expand_hf_repo_wildcards(wildcard_specs, query, limit)
|
|
|
|
# Prioritize expanded/dynamic repos, then static whitelist.
|
|
repos_to_scan = expanded + concrete_repos
|
|
# De-dup while preserving order.
|
|
seen_repo = set()
|
|
repos_to_scan = [r for r in repos_to_scan if not (r in seen_repo or seen_repo.add(r))]
|
|
|
|
if not repos_to_scan:
|
|
return []
|
|
|
|
logger.info(f"Scanning {len(repos_to_scan)} repos for files matching '{query}': {repos_to_scan}")
|
|
|
|
for repo_id in repos_to_scan:
|
|
if len(results) >= limit:
|
|
break
|
|
|
|
files = []
|
|
with _hf_tree_cache_lock:
|
|
cached = _hf_tree_cache.get(repo_id)
|
|
if cached and (time.time() - cached[0]) < HF_TREE_CACHE_TTL:
|
|
files = cached[1]
|
|
logger.debug(f"Using cached file list for HF repo: {repo_id}")
|
|
else:
|
|
try:
|
|
# Fetch file list from HF API
|
|
api_url = f"https://huggingface.co/api/models/{repo_id}/tree/main?recursive=True"
|
|
|
|
# Get HF token if available
|
|
token = (
|
|
os.environ.get("HUGGINGFACE_HUB_TOKEN")
|
|
or os.environ.get("HF_TOKEN")
|
|
or os.environ.get("HUGGINGFACE_TOKEN")
|
|
)
|
|
headers = {}
|
|
if token:
|
|
headers["Authorization"] = f"Bearer {token}"
|
|
|
|
data = _make_request(api_url, headers)
|
|
|
|
if isinstance(data, list):
|
|
files = data
|
|
_hf_tree_cache[repo_id] = (time.time(), files)
|
|
logger.info(f"Fetched and cached file list for HF repo: {repo_id} ({len(files)} files)")
|
|
except Exception as e:
|
|
logger.error(f"Failed to fetch file tree for HF repo {repo_id}: {e}")
|
|
# Cache the failure for a shorter time to avoid hammering the API
|
|
_hf_tree_cache[repo_id] = (time.time(), [])
|
|
|
|
# Filter and score files
|
|
for file_info in files:
|
|
filepath = file_info.get("path", "")
|
|
|
|
# Limit scan depth to avoid performance issues on huge repos
|
|
if len(results) > limit * 2: # Stop scanning files if we have a decent number of candidates
|
|
break
|
|
|
|
if not filepath.endswith((".safetensors", ".ckpt", ".pt", ".pth", ".bin")):
|
|
continue
|
|
|
|
filename = filepath.split('/')[-1]
|
|
score, match_level = calculate_match_score(query, filename, filename)
|
|
|
|
if score >= 75: # Use a slightly lower threshold for this targeted search
|
|
# Construct download URL and ensure it's canonical
|
|
download_url = canonicalize_hf_url(
|
|
f"https://huggingface.co/{repo_id}/resolve/main/{filepath}"
|
|
)
|
|
|
|
# Guess type from filename
|
|
tags = repo_id.lower().split('/') + filename.lower().split('.')
|
|
kind = "checkpoints"
|
|
if "lora" in tags: kind = "loras"
|
|
elif "vae" in tags: kind = "vae"
|
|
elif "controlnet" in tags: kind = "controlnet"
|
|
elif "clip" in tags: kind = "clip"
|
|
|
|
results.append({
|
|
"platform": "huggingface",
|
|
"name": f"{repo_id} (file)",
|
|
"filename": filename,
|
|
"url": download_url,
|
|
"page_url": f"https://huggingface.co/{repo_id}/tree/main",
|
|
"type": kind,
|
|
"version": file_info.get("lastCommit", {}).get("oid", "main")[:7],
|
|
"size_mb": int(file_info.get("size", 0) / (1024 * 1024)),
|
|
"sha256": file_info.get("lfs", {}).get("oid"),
|
|
"match_score": score,
|
|
"match_level": f"{match_level} (in {repo_id})",
|
|
})
|
|
|
|
results.sort(key=lambda x: x.get("match_score", 0), reverse=True)
|
|
return results[:limit]
|
|
|
|
|
|
def search_huggingface(query: str, limit: int = 5) -> List[Dict[str, Any]]:
|
|
"""
|
|
Search for models on Hugging Face.
|
|
|
|
Args:
|
|
query: Model name to search for
|
|
limit: Maximum number of results
|
|
|
|
Returns:
|
|
List of results with same format as search_civitai
|
|
"""
|
|
results = []
|
|
|
|
# First, check if this matches a known official model
|
|
official_models = _get_official_sd_models()
|
|
query_normalized = query.lower().strip().replace("_", "-").replace(" ", "-")
|
|
|
|
for model_key, model_info in official_models.items():
|
|
model_key_normalized = model_key.lower().replace("_", "-")
|
|
if model_key_normalized in query_normalized or query_normalized in model_key_normalized:
|
|
# Direct match with official model
|
|
model_id = model_info["model_id"]
|
|
filename = model_info["filename"]
|
|
alt_filename = model_info.get("alt_filename")
|
|
|
|
# Try both filenames if available
|
|
for fname in [filename, alt_filename] if alt_filename else [filename]:
|
|
if fname:
|
|
download_url = f"https://huggingface.co/{model_id}/resolve/main/{fname}"
|
|
|
|
# Calculate match score
|
|
score, match_level = calculate_match_score(query, model_info["name"], fname)
|
|
|
|
# Boost score for exact filename matches
|
|
fname_base = fname.lower().replace(".safetensors", "").replace(".ckpt", "")
|
|
if query_normalized == fname_base:
|
|
score = 100.0
|
|
match_level = "Exact match (official)"
|
|
elif query_normalized.replace("-fp16", "").replace("-fp32", "") == fname_base:
|
|
# Match but with precision suffix
|
|
score = 98.0
|
|
match_level = "Exact match (fp16/fp32 variant)"
|
|
|
|
results.append({
|
|
"platform": "huggingface",
|
|
"name": f"{model_info['name']} (Official)",
|
|
"filename": fname,
|
|
"url": download_url,
|
|
"page_url": f"https://huggingface.co/{model_id}",
|
|
"type": "checkpoints",
|
|
"version": "official",
|
|
"size_mb": 0, # Size not pre-computed
|
|
"sha256": None,
|
|
"match_score": score,
|
|
"match_level": match_level,
|
|
})
|
|
|
|
# If we found exact official matches, return them
|
|
if any(r.get("match_score", 0) >= 95 for r in results):
|
|
# Ensure URLs are downloadable
|
|
for r in results:
|
|
r["url"] = canonicalize_hf_url(r["url"])
|
|
return results[:limit]
|
|
|
|
# Otherwise, continue with API search
|
|
potential_repos_from_search = []
|
|
|
|
# 1. Search via Hugging Face Internal API
|
|
try:
|
|
# Hugging Face API search endpoint
|
|
params = {"search": query, "limit": limit * 2}
|
|
url = f"https://huggingface.co/api/models?{urlencode(params)}"
|
|
|
|
token = os.environ.get("HUGGINGFACE_HUB_TOKEN") or os.environ.get("HF_TOKEN")
|
|
headers = {}
|
|
if token:
|
|
headers["Authorization"] = f"Bearer {token}"
|
|
|
|
data = _make_request(url, headers)
|
|
|
|
if isinstance(data, list):
|
|
# Collect top repos for fallback scan
|
|
for i, item in enumerate(data):
|
|
if i < 3: # Keep top 3 for deep scan
|
|
mid = item.get("id")
|
|
if mid:
|
|
potential_repos_from_search.append(mid)
|
|
|
|
for item in data:
|
|
model_id = item.get("id", "")
|
|
if not model_id:
|
|
continue
|
|
|
|
# For model search, we often don't get a specific file, so we score based on the model ID/name
|
|
score, match_level = calculate_match_score(query, model_id)
|
|
|
|
# Find a suitable file to download (prefer safetensors)
|
|
siblings = item.get("siblings", [])
|
|
best_file = None
|
|
for s in siblings:
|
|
rfilename = s.get("rfilename", "")
|
|
if rfilename.endswith(".safetensors"):
|
|
best_file = rfilename
|
|
break
|
|
elif rfilename.endswith((".ckpt", ".pt", ".bin")):
|
|
best_file = rfilename
|
|
|
|
if not best_file:
|
|
continue
|
|
|
|
download_url = canonicalize_hf_url(f"https://huggingface.co/{model_id}/resolve/main/{best_file}")
|
|
|
|
results.append({
|
|
"platform": "huggingface",
|
|
"name": model_id,
|
|
"filename": best_file,
|
|
"url": download_url,
|
|
"page_url": f"https://huggingface.co/{model_id}",
|
|
"type": "checkpoints", # Guess, can be refined
|
|
"version": "main",
|
|
"size_mb": 0, # Size is not in this API response
|
|
"sha256": None,
|
|
"match_score": score,
|
|
"match_level": f"{match_level} (model)",
|
|
})
|
|
|
|
except Exception as e:
|
|
logger.error("Hugging Face model search failed: %s", str(e))
|
|
|
|
# 2. Search via DuckDuckGo (Google Fallback) if needed
|
|
# If we don't have good API results, try to find repos via search engine
|
|
if not any(r.get("match_score", 0) >= 80 for r in results):
|
|
logger.info(f"No high-confidence matches, engaging DuckDuckGo/Google fallback for: {query}")
|
|
ddg_repos = _search_duckduckgo_repos(query, limit=3)
|
|
for repo in ddg_repos:
|
|
if repo not in potential_repos_from_search:
|
|
potential_repos_from_search.append(repo)
|
|
|
|
# --- Fallback to file search in whitelisted repos ---
|
|
# If no high-confidence results from model search, OR if we want to be thorough
|
|
has_good_results = any(r.get("match_score", 0) >= 80 for r in results)
|
|
|
|
# Always try fallback search if we found potential repos, to find specific files inside them
|
|
# irrespective of whether we found "models" (which might just correspond to the repo root)
|
|
if not has_good_results or potential_repos_from_search:
|
|
logger.info(f"Running fallback file search. Candidates from global search: {potential_repos_from_search}")
|
|
file_search_results = _search_huggingface_repo_files(query, limit, extra_repos=potential_repos_from_search)
|
|
if file_search_results:
|
|
results.extend(file_search_results)
|
|
|
|
# De-duplicate results by URL, keeping the one with the highest score
|
|
seen_urls = {}
|
|
deduped_results = []
|
|
for r in sorted(results, key=lambda x: x.get("match_score", 0), reverse=True):
|
|
url = r.get("url")
|
|
if url not in seen_urls:
|
|
seen_urls[url] = r
|
|
deduped_results.append(r)
|
|
|
|
# Sort by match score (highest first) and filter out poor matches
|
|
deduped_results.sort(key=lambda x: x.get("match_score", 0), reverse=True)
|
|
final_results = [r for r in deduped_results if r.get("match_score", 0) >= 30]
|
|
|
|
return final_results[:limit]
|
|
|
|
|
|
def search_github(query: str, limit: int = 5) -> List[Dict[str, Any]]:
|
|
"""
|
|
Search for models on GitHub releases.
|
|
|
|
Args:
|
|
query: Model name or repo to search for
|
|
limit: Maximum number of results
|
|
|
|
Returns:
|
|
List of results with same format as search_civitai
|
|
"""
|
|
results = []
|
|
|
|
try:
|
|
# GitHub search API
|
|
params = {
|
|
"q": f"{query} extension:safetensors OR extension:ckpt",
|
|
"per_page": limit
|
|
}
|
|
url = f"https://api.github.com/search/repositories?{urlencode(params)}"
|
|
|
|
# Get GitHub token if available
|
|
token = os.environ.get("GITHUB_TOKEN") or os.environ.get("GH_TOKEN")
|
|
headers = {"Accept": "application/vnd.github.v3+json"}
|
|
if token:
|
|
headers["Authorization"] = f"token {token}"
|
|
|
|
data = _make_request(url, headers)
|
|
items = data.get("items", [])
|
|
|
|
for item in items[:limit]:
|
|
full_name = item.get("full_name", "")
|
|
if not full_name:
|
|
continue
|
|
|
|
# Get releases
|
|
releases_url = f"https://api.github.com/repos/{full_name}/releases/latest"
|
|
release_data = _make_request(releases_url, headers)
|
|
|
|
assets = release_data.get("assets", [])
|
|
if not assets:
|
|
continue
|
|
|
|
# Find model files in assets
|
|
for asset in assets:
|
|
name = asset.get("name", "")
|
|
if not name.endswith((".safetensors", ".ckpt", ".pt", ".bin")):
|
|
continue
|
|
|
|
download_url = asset.get("browser_download_url", "")
|
|
if not download_url:
|
|
continue
|
|
|
|
size = asset.get("size", 0)
|
|
size_mb = int(size / (1024 * 1024)) if size else 0
|
|
|
|
# Guess type from filename
|
|
name_lower = name.lower()
|
|
kind = "checkpoints"
|
|
if "lora" in name_lower:
|
|
kind = "loras"
|
|
elif "vae" in name_lower:
|
|
kind = "vae"
|
|
elif "controlnet" in name_lower:
|
|
kind = "controlnet"
|
|
|
|
# Calculate match score
|
|
full_result_name = f"{full_name} - {name}"
|
|
score, match_level = calculate_match_score(query, full_result_name, name)
|
|
|
|
results.append({
|
|
"platform": "github",
|
|
"name": full_result_name,
|
|
"filename": name,
|
|
"url": download_url,
|
|
"page_url": item.get("html_url", ""),
|
|
"type": kind,
|
|
"version": release_data.get("tag_name", "latest"),
|
|
"size_mb": size_mb,
|
|
"sha256": None,
|
|
"match_score": score,
|
|
"match_level": match_level,
|
|
})
|
|
|
|
if len(results) >= limit:
|
|
break
|
|
|
|
if len(results) >= limit:
|
|
break
|
|
|
|
except Exception as e:
|
|
logger.error("GitHub search failed: %s", str(e))
|
|
|
|
# Sort by match score (highest first) and filter out poor matches
|
|
results.sort(key=lambda x: x.get("match_score", 0), reverse=True)
|
|
results = [r for r in results if r.get("match_score", 0) >= 30]
|
|
|
|
return results[:limit]
|
|
|
|
|
|
def _generate_search_variants(query: str) -> List[str]:
|
|
"""
|
|
Generate search query variants to improve search results.
|
|
Handles long names with special characters and underscores.
|
|
|
|
Returns list of queries to try, ordered by specificity.
|
|
Strategy: Try multiple normalized forms for better matching.
|
|
"""
|
|
variants = []
|
|
|
|
# Remove file extension but keep the rest
|
|
query_clean = re.sub(r'\.(safetensors|ckpt|pt|pth|bin)$', '', query, flags=re.IGNORECASE)
|
|
|
|
# 1. Original query (cleaned)
|
|
variants.append(query_clean)
|
|
|
|
# 2. Replace underscores and hyphens with spaces for better matching
|
|
# Example: "Qwen_Rapid_AIO-NSFW-v5.3" → "Qwen Rapid AIO NSFW v5.3"
|
|
query_with_spaces = re.sub(r'[-_]+', ' ', query_clean)
|
|
if query_with_spaces != query_clean:
|
|
variants.append(query_with_spaces)
|
|
|
|
# 3. Replace underscores with hyphens (some APIs prefer hyphens)
|
|
# Example: "stable_diffusion_v1_5" → "stable-diffusion-v1-5"
|
|
query_with_hyphens = re.sub(r'_+', '-', query_clean)
|
|
if query_with_hyphens != query_clean and query_with_hyphens not in variants:
|
|
variants.append(query_with_hyphens)
|
|
|
|
# Normalize for pattern matching
|
|
normalized = query_clean.lower().strip()
|
|
|
|
# 4. For very long names (>40 chars), try extracting key parts
|
|
if len(query_clean) > 40:
|
|
# Remove technical suffixes common in model names
|
|
# Example: "model-name-v1.5-pruned-emaonly-fp16" → "model-name-v1.5"
|
|
compact = re.sub(r'[-_](pruned|ema|emaonly|fp16|fp32|inpainting|training|diffusers)[-_]?', '-',
|
|
query_clean, flags=re.IGNORECASE)
|
|
compact = re.sub(r'[-_]+', '-', compact).strip('-')
|
|
if compact != query_clean and len(compact) >= 10:
|
|
variants.append(compact)
|
|
|
|
# 5. Extract main model name (first significant part before version/variant)
|
|
# Example: "ModelName-v5.3-NSFW-AIO" → "ModelName"
|
|
# But keep at least 2-3 meaningful parts for specificity
|
|
parts = re.split(r'[-_\s]+', query_clean)
|
|
if len(parts) >= 3:
|
|
# Keep first 2-3 meaningful parts (skip very short parts like v1, v2, etc.)
|
|
meaningful_parts = [p for p in parts if len(p) >= 3 or re.match(r'v\d+', p.lower())]
|
|
if len(meaningful_parts) >= 2:
|
|
base_name = ' '.join(meaningful_parts[:3]) # Take first 3 meaningful parts
|
|
if base_name not in variants and len(base_name) >= 10:
|
|
variants.append(base_name)
|
|
|
|
# 6. Check if this is a known Stable Diffusion model
|
|
if re.match(r'^v\d+-\d+', normalized):
|
|
# Official SD model like "v1-5-pruned-emaonly"
|
|
match = re.match(r'^v(\d+)-(\d+)', normalized)
|
|
if match:
|
|
variants.append(f"stable diffusion {match.group(1)}.{match.group(2)}")
|
|
|
|
# Remove duplicates while preserving order
|
|
seen = set()
|
|
unique_variants = []
|
|
for v in variants:
|
|
v_clean = v.strip()
|
|
# Normalize for comparison (case-insensitive, space-normalized)
|
|
v_normalized = re.sub(r'\s+', ' ', v_clean.lower())
|
|
if v_clean and v_normalized not in seen and len(v_clean) >= 3:
|
|
seen.add(v_normalized)
|
|
unique_variants.append(v_clean)
|
|
|
|
# Return up to 4 variants for thorough search
|
|
return unique_variants[:4]
|
|
|
|
|
|
def _extract_model_info_from_url(url: str, query: str) -> Optional[Dict[str, Any]]:
|
|
"""
|
|
Extract model information from Hugging Face or GitHub URL.
|
|
|
|
Args:
|
|
url: The URL to extract info from
|
|
query: Original search query for match scoring
|
|
|
|
Returns:
|
|
Dict with model info or None if extraction fails
|
|
"""
|
|
try:
|
|
# Hugging Face URL pattern: https://huggingface.co/{owner}/{model_name}
|
|
hf_match = re.match(r'https?://(?:www\.)?huggingface\.co/([^/]+)/([^/?#]+)', url)
|
|
if hf_match:
|
|
owner, model_name = hf_match.groups()
|
|
full_name = f"{owner}/{model_name}"
|
|
|
|
# Try to get model files
|
|
api_url = f"https://huggingface.co/api/models/{owner}/{model_name}"
|
|
headers = {"User-Agent": "ComfyUI-Majoor-Downloader"}
|
|
|
|
try:
|
|
model_data = _make_request(api_url, headers)
|
|
tags = model_data.get("tags", [])
|
|
|
|
# Get files
|
|
files_url = f"https://huggingface.co/api/models/{owner}/{model_name}/tree/main"
|
|
files_data = _make_request(files_url, headers)
|
|
|
|
# Find model files
|
|
if isinstance(files_data, list):
|
|
for file_info in files_data:
|
|
path = file_info.get("path", "")
|
|
if path.endswith((".safetensors", ".ckpt", ".pt", ".bin")):
|
|
filename = path.split("/")[-1]
|
|
score, match_level = calculate_match_score(query, full_name, filename)
|
|
|
|
# Only return if score >= 80%
|
|
if score >= 80:
|
|
return {
|
|
"platform": "huggingface",
|
|
"name": full_name,
|
|
"filename": filename,
|
|
"url": f"https://huggingface.co/{owner}/{model_name}/resolve/main/{path}",
|
|
"page_url": url,
|
|
"type": "checkpoints",
|
|
"version": "main",
|
|
"size_mb": 0,
|
|
"sha256": None,
|
|
"match_score": score,
|
|
"match_level": f"{match_level} (Web discovered)",
|
|
}
|
|
except Exception as e:
|
|
logger.debug(f"Failed to fetch HF model details for {full_name}: {e}")
|
|
|
|
# GitHub URL pattern: https://github.com/{owner}/{repo}
|
|
gh_match = re.match(r'https?://(?:www\.)?github\.com/([^/]+)/([^/?#]+)', url)
|
|
if gh_match:
|
|
owner, repo = gh_match.groups()
|
|
full_name = f"{owner}/{repo}"
|
|
|
|
# Try to get latest release
|
|
api_url = f"https://api.github.com/repos/{owner}/{repo}/releases/latest"
|
|
headers = {"Accept": "application/vnd.github.v3+json"}
|
|
|
|
try:
|
|
release_data = _make_request(api_url, headers)
|
|
assets = release_data.get("assets", [])
|
|
|
|
for asset in assets:
|
|
name = asset.get("name", "")
|
|
if name.endswith((".safetensors", ".ckpt", ".pt", ".bin")):
|
|
score, match_level = calculate_match_score(query, full_name, name)
|
|
|
|
# Only return if score >= 80%
|
|
if score >= 80:
|
|
return {
|
|
"platform": "github",
|
|
"name": f"{full_name} - {name}",
|
|
"filename": name,
|
|
"url": asset.get("browser_download_url", ""),
|
|
"page_url": url,
|
|
"type": "checkpoints",
|
|
"version": release_data.get("tag_name", "latest"),
|
|
"size_mb": int(asset.get("size", 0) / (1024 * 1024)),
|
|
"sha256": None,
|
|
"match_score": score,
|
|
"match_level": f"{match_level} (Web discovered)",
|
|
}
|
|
except Exception as e:
|
|
logger.debug(f"Failed to fetch GitHub release for {full_name}: {e}")
|
|
|
|
except Exception as e:
|
|
logger.debug(f"Failed to extract model info from {url}: {e}")
|
|
|
|
return None
|
|
|
|
|
|
def search_web_first(query: str, limit: int = 10) -> List[Dict[str, Any]]:
|
|
"""
|
|
Web-first search (DISABLED).
|
|
|
|
The previous implementation scraped DuckDuckGo HTML results. This is fragile and can violate
|
|
search engine Terms of Service. It has been disabled intentionally.
|
|
|
|
Args:
|
|
query: Model name to search for
|
|
limit: Maximum number of results to process
|
|
|
|
Returns:
|
|
Always returns an empty list.
|
|
"""
|
|
_ = (query, limit)
|
|
logger.info("Web-first search is disabled; using API search only.")
|
|
return []
|
|
|
|
|
|
def search_all_platforms(query: str, limit_per_platform: int = 3) -> Dict[str, Any]:
|
|
"""
|
|
Search all platforms using official APIs with community registry and aliases.
|
|
|
|
Improved search strategy:
|
|
1. Search in community registry first (instant, verified sources)
|
|
2. Use alias system to expand query variants
|
|
3. Search official APIs in parallel (CivitAI, HuggingFace, GitHub, ModelScope)
|
|
4. Deduplicate and rank by score
|
|
|
|
Args:
|
|
query: Model name to search for
|
|
limit_per_platform: Maximum results per platform
|
|
|
|
Returns:
|
|
{
|
|
"query": str,
|
|
"total_results": int,
|
|
"platforms": {
|
|
"community_registry": [...],
|
|
"civitai": [...],
|
|
"huggingface": [...],
|
|
"github": [...],
|
|
"modelscope": [...]
|
|
}
|
|
}
|
|
"""
|
|
import concurrent.futures
|
|
|
|
results = {
|
|
"query": query,
|
|
"total_results": 0,
|
|
"platforms": {
|
|
"community_registry": [],
|
|
"civitai": [],
|
|
"huggingface": [],
|
|
"github": [],
|
|
"modelscope": []
|
|
}
|
|
}
|
|
|
|
# STEP 1: Search in community registry first (instant!)
|
|
try:
|
|
from .model_registry import get_model_registry
|
|
registry = get_model_registry()
|
|
registry_results = registry.search(query, limit=limit_per_platform)
|
|
results["platforms"]["community_registry"] = registry_results
|
|
|
|
# If we have excellent results (95+) from registry, we might skip APIs
|
|
excellent_registry = [r for r in registry_results if r.get("match_score", 0) >= 95]
|
|
if len(excellent_registry) >= 2:
|
|
logger.info(f"Found {len(excellent_registry)} excellent matches in community registry")
|
|
# Still search APIs but with lower priority
|
|
except Exception as e:
|
|
logger.warning(f"Community registry search failed: {e}")
|
|
|
|
# STEP 2: Generate search variants using alias system
|
|
try:
|
|
from .model_aliases import resolve_aliases
|
|
alias_variants = resolve_aliases(query)
|
|
# Combine with original variants
|
|
original_variants = _generate_search_variants(query)
|
|
# Merge and deduplicate
|
|
all_variants = list(set(list(alias_variants) + original_variants))
|
|
# Limit to top 4 most promising variants
|
|
search_queries = all_variants[:4]
|
|
logger.info(f"Searching with {len(search_queries)} variants (including aliases): {search_queries[:3]}...")
|
|
except Exception as e:
|
|
logger.warning(f"Alias resolution failed, using original variants: {e}")
|
|
search_queries = _generate_search_variants(query)
|
|
|
|
all_results_by_platform = {
|
|
"civitai": [],
|
|
"huggingface": [],
|
|
"github": [],
|
|
"modelscope": []
|
|
}
|
|
|
|
# STEP 3: Official API search in parallel
|
|
hf_limit = limit_per_platform * 2
|
|
gh_limit = limit_per_platform * 2
|
|
civitai_limit = limit_per_platform
|
|
modelscope_limit = limit_per_platform
|
|
|
|
for search_query in search_queries:
|
|
# Search 4 platforms in parallel
|
|
with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor:
|
|
future_hf = executor.submit(search_huggingface, search_query, hf_limit)
|
|
future_gh = executor.submit(search_github, search_query, gh_limit)
|
|
future_civitai = executor.submit(search_civitai, search_query, civitai_limit)
|
|
|
|
# Add ModelScope search
|
|
try:
|
|
from .model_search_modelscope import search_modelscope
|
|
future_modelscope = executor.submit(search_modelscope, search_query, modelscope_limit)
|
|
except ImportError:
|
|
logger.debug("ModelScope search not available")
|
|
future_modelscope = None
|
|
|
|
# Collect results
|
|
def _resolve_future(future, platform: str) -> List[Dict[str, Any]]:
|
|
try:
|
|
return future.result()
|
|
except Exception as exc:
|
|
logger.warning("%s search failed for '%s': %s", platform, search_query, exc)
|
|
return []
|
|
|
|
hf_results = _resolve_future(future_hf, "huggingface")
|
|
gh_results = _resolve_future(future_gh, "github")
|
|
civitai_results = _resolve_future(future_civitai, "civitai")
|
|
modelscope_results = _resolve_future(future_modelscope, "modelscope") if future_modelscope else []
|
|
|
|
# Filter API results to only keep score >= 80%
|
|
hf_results = [r for r in hf_results if r.get("match_score", 0) >= 80]
|
|
gh_results = [r for r in gh_results if r.get("match_score", 0) >= 80]
|
|
civitai_results = [r for r in civitai_results if r.get("match_score", 0) >= 80]
|
|
modelscope_results = [r for r in modelscope_results if r.get("match_score", 0) >= 80]
|
|
|
|
# Apply platform priority bonus to scores
|
|
for r in hf_results:
|
|
r["match_score"] = min(100.0, r.get("match_score", 0) + 5)
|
|
r["match_level"] = f"{r.get('match_level', '')} [API+HF Priority]"
|
|
|
|
for r in gh_results:
|
|
r["match_score"] = min(100.0, r.get("match_score", 0) + 3)
|
|
r["match_level"] = f"{r.get('match_level', '')} [API+GH Priority]"
|
|
|
|
for r in modelscope_results:
|
|
r["match_score"] = min(100.0, r.get("match_score", 0) + 2)
|
|
r["match_level"] = f"{r.get('match_level', '')} [ModelScope]"
|
|
|
|
all_results_by_platform["huggingface"].extend(hf_results)
|
|
all_results_by_platform["github"].extend(gh_results)
|
|
all_results_by_platform["civitai"].extend(civitai_results)
|
|
all_results_by_platform["modelscope"].extend(modelscope_results)
|
|
|
|
# Check if we have enough good results to stop searching
|
|
all_good_results = [
|
|
r
|
|
for r in (
|
|
all_results_by_platform["huggingface"] +
|
|
all_results_by_platform["github"] +
|
|
all_results_by_platform["civitai"] +
|
|
all_results_by_platform["modelscope"]
|
|
)
|
|
if r.get("match_score", 0) >= 80
|
|
]
|
|
if len(all_good_results) >= 5:
|
|
logger.info(f"Found enough good results (>= 80%) with query: {search_query}")
|
|
break
|
|
|
|
# Deduplicate results by URL and sort by score
|
|
# IMPORTANT: Only keep results with score >= 80%
|
|
def deduplicate_and_filter(results_list, max_results, min_score=80):
|
|
seen_urls = set()
|
|
unique = []
|
|
for r in sorted(results_list, key=lambda x: x.get("match_score", 0), reverse=True):
|
|
score = r.get("match_score", 0)
|
|
url = r.get("url", "")
|
|
|
|
# Filter: only keep results with score >= min_score
|
|
if score < min_score:
|
|
continue
|
|
|
|
if url and url not in seen_urls:
|
|
seen_urls.add(url)
|
|
unique.append(r)
|
|
return unique[:max_results]
|
|
|
|
# STEP 4: Deduplicate and filter all results (score >= 80%)
|
|
results["platforms"]["huggingface"] = deduplicate_and_filter(all_results_by_platform["huggingface"], hf_limit, min_score=80)
|
|
results["platforms"]["github"] = deduplicate_and_filter(all_results_by_platform["github"], gh_limit, min_score=80)
|
|
results["platforms"]["civitai"] = deduplicate_and_filter(all_results_by_platform["civitai"], civitai_limit, min_score=80)
|
|
results["platforms"]["modelscope"] = deduplicate_and_filter(all_results_by_platform["modelscope"], modelscope_limit, min_score=80)
|
|
|
|
# Community registry already filtered in step 1
|
|
# results["platforms"]["community_registry"] already set
|
|
|
|
# Combine all results with priority order:
|
|
# 1. Community registry (highest priority - verified sources)
|
|
# 2. HuggingFace
|
|
# 3. CivitAI
|
|
# 4. GitHub
|
|
# 5. ModelScope
|
|
all_results = []
|
|
all_results.extend(results["platforms"]["community_registry"])
|
|
all_results.extend(results["platforms"]["huggingface"])
|
|
all_results.extend(results["platforms"]["civitai"])
|
|
all_results.extend(results["platforms"]["github"])
|
|
all_results.extend(results["platforms"]["modelscope"])
|
|
|
|
# Log filtering results
|
|
logger.info(f"After filtering (score >= 80%): {len(all_results)} results total")
|
|
logger.info(f" Community Registry: {len(results['platforms']['community_registry'])}")
|
|
logger.info(f" HuggingFace: {len(results['platforms']['huggingface'])}")
|
|
logger.info(f" CivitAI: {len(results['platforms']['civitai'])}")
|
|
logger.info(f" GitHub: {len(results['platforms']['github'])}")
|
|
logger.info(f" ModelScope: {len(results['platforms']['modelscope'])}")
|
|
|
|
# Sort globally by match score (already filtered to >= 80%)
|
|
all_results.sort(key=lambda x: x.get("match_score", 0), reverse=True)
|
|
|
|
results["total_results"] = len(all_results)
|
|
results["sorted_results"] = all_results # Add globally sorted results
|
|
|
|
# Add helpful links
|
|
results["google_search_url"] = f"https://www.google.com/search?q={quote(query)}+download+safetensors+OR+ckpt"
|
|
results["civitai_search_url"] = f"https://civitai.com/search/models?query={quote(query)}"
|
|
results["huggingface_search_url"] = f"https://huggingface.co/models?search={quote(query)}"
|
|
|
|
# Add stats
|
|
results["search_stats"] = {
|
|
"query_variants_used": len(search_queries),
|
|
"platforms_searched": 5, # registry, hf, civitai, github, modelscope
|
|
"total_matches_found": len(all_results),
|
|
"excellent_matches": len([r for r in all_results if r.get("match_score", 0) >= 95]),
|
|
"good_matches": len([r for r in all_results if 85 <= r.get("match_score", 0) < 95]),
|
|
}
|
|
|
|
return results
|