diff --git a/__init__.py b/__init__.py index 6aa945f..be6e50c 100644 --- a/__init__.py +++ b/__init__.py @@ -1448,31 +1448,8 @@ class ModelResolverExtension: ) def result_url_looks_like_model_file(url, expected_filename=""): - text = str(url or "").strip() - if not text.startswith(("http://", "https://")): - return False - try: - from urllib.parse import unquote as _unquote, urlparse as _urlparse - parsed = _urlparse(text) - host = parsed.netloc.lower() - path = _unquote(parsed.path or "") - except Exception: - host = "" - path = text - if host.endswith("huggingface.co") and path.startswith("/spaces/"): - return False - if ( - host.endswith("civitai.com") - or host.endswith("civitai.red") - ) and path.startswith("/api/download/"): - return True - - basename = _os.path.basename(path).lower() - expected = _os.path.basename(str(expected_filename or "")).lower() - if expected and basename == expected: - return True - model_extensions = MODEL_EXTENSIONS - return _os.path.splitext(basename)[1].lower() in model_extensions + from core.type_utils import looks_like_model_file + return looks_like_model_file(url, expected_filename) def collect_result_download_urls(result): urls = [] diff --git a/core/resolver.py b/core/resolver.py index 8434184..4b2ad4a 100644 --- a/core/resolver.py +++ b/core/resolver.py @@ -315,17 +315,10 @@ def _metadata_sidecar_candidates(model_path: str) -> List[str]: return unique -def _as_dict(value: Any) -> Dict[str, Any]: - return as_dict(value) - - -def _as_list(value: Any) -> List[Any]: - return as_list(value) - def _collect_hashes_from_container(value: Any) -> List[str]: - container = _as_dict(value) - hashes = _as_dict(container.get("hashes")) + container = as_dict(value) + hashes = as_dict(container.get("hashes")) return [ normalize_sha256(candidate) for candidate in ( @@ -382,11 +375,11 @@ def _extract_model_sha256_from_metadata( nested_file_lists = [ metadata.get("files"), - _as_dict(metadata.get("selected_version")).get("files"), - _as_dict(metadata.get("civitai")).get("files"), + as_dict(metadata.get("selected_version")).get("files"), + as_dict(metadata.get("civitai")).get("files"), ] for file_list in nested_file_lists: - for file_info in _as_list(file_list): + for file_info in as_list(file_list): if isinstance(file_info, dict) and _metadata_file_matches_model(file_info, model): values.extend(_collect_hashes_from_container(file_info)) diff --git a/core/sources/civarchive.py b/core/sources/civarchive.py index 7bb7d60..e8cc78b 100644 --- a/core/sources/civarchive.py +++ b/core/sources/civarchive.py @@ -31,10 +31,10 @@ from ..type_utils import ( parse_size_to_bytes, get_version_sort_key, DEFAULT_BROWSER_USER_AGENT, - parse_size_header as _parse_size_header, - parse_content_range_size as _parse_content_range_size, - extract_response_file_size as _response_file_size, parse_civitai_model_path, + fetch_remote_file_size, + looks_like_model_file, + normalize_model_image, ) from ..progress import report_progress, get_progress_reporter from ..log_system.log_funcs import create_module_logger @@ -569,20 +569,8 @@ def _download_url_looks_like_model_file( normalized = _normalize_download_url(url) if not normalized: return False + return looks_like_model_file(normalized, expected_filename) - parsed = urlparse(normalized) - host = parsed.netloc.lower() - path = unquote(parsed.path or "") - if host.endswith("huggingface.co") and path.startswith("/spaces/"): - return False - if normalized.startswith(CIVITAI_DOWNLOAD_URL_PREFIXES): - return True - - basename = get_filename_from_path(path).lower() - expected = get_filename_from_path(str(expected_filename or "")).lower() - if expected and basename == expected: - return True - return _has_known_model_extension(basename) def _archive_link_is_dead(item: Dict[str, Any]) -> bool: @@ -633,43 +621,12 @@ def _fetch_remote_file_size_bytes(url: Any, timeout: int = 15) -> Optional[int]: if probe_url in _download_size_cache: return _download_size_cache[probe_url] - size = None - try: - response = requests.head( - probe_url, - headers=REQUEST_HEADERS, - allow_redirects=True, - timeout=timeout, - ) - try: - if response.status_code < 400: - size = _response_file_size(response) - finally: - response.close() - except Exception as e: - log.debug(f"CivArchive size HEAD failed: url={probe_url}, error={e}") - - if not size: - try: - response = requests.get( - probe_url, - headers={**REQUEST_HEADERS, "Range": "bytes=0-0"}, - allow_redirects=True, - stream=True, - timeout=timeout, - ) - try: - if response.status_code < 400: - size = _response_file_size(response) - finally: - response.close() - except Exception as e: - log.debug(f"CivArchive size range request failed: url={probe_url}, error={e}") - + size = fetch_remote_file_size(probe_url, headers=REQUEST_HEADERS, timeout=timeout) _download_size_cache[probe_url] = size return size + def _resolve_file_size_bytes( file_info: Dict[str, Any], download_urls: Optional[List[str]] = None, @@ -972,21 +929,8 @@ def _normalize_archive_mirrors(file_info: Dict[str, Any]) -> List[Dict[str, Any] def _normalize_archive_image(image: Dict[str, Any]) -> Dict[str, Any]: - meta = image.get("meta") if isinstance(image.get("meta"), dict) else {} - return { - "url": image.get("url") or image.get("imageUrl") or image.get("src"), - "civitaiUrl": image.get("civitaiUrl") or image.get("postUrl"), - "seed": image.get("seed") or meta.get("seed"), - "steps": image.get("steps") or meta.get("steps"), - "cfg": image.get("cfg") or meta.get("cfg"), - "sampler": image.get("sampler") or meta.get("sampler"), - "model": image.get("model") or meta.get("model") or meta.get("Model"), - "positive": image.get("positive") or image.get("prompt") or meta.get("prompt"), - "negative": image.get("negative") or meta.get("negative_prompt") or meta.get("negativePrompt"), - "width": image.get("width") or meta.get("width"), - "height": image.get("height") or meta.get("height"), - "metadata": meta, - } + return normalize_model_image(image) + def _normalize_archive_file(file_info: Dict[str, Any], model_id: Optional[int], version_id: Optional[int]) -> Dict[str, Any]: @@ -1197,8 +1141,6 @@ def get_civarchive_model_details( -def _version_sort_key(version: Dict[str, Any]) -> tuple: - return get_version_sort_key(version) def _collect_normalized_download_urls(file_info: Dict[str, Any]) -> List[str]: @@ -1374,7 +1316,7 @@ def _find_model_title_match_in_model_details( if not versions: return None - versions = sorted(versions, key=_version_sort_key, reverse=True) + versions = sorted(versions, key=get_version_sort_key, reverse=True) rejected_by_base_model = False for version in versions: diff --git a/core/sources/civitai.py b/core/sources/civitai.py index 46e92d3..84ee028 100644 --- a/core/sources/civitai.py +++ b/core/sources/civitai.py @@ -29,6 +29,7 @@ from ..type_utils import ( check_credential_http, DEFAULT_BROWSER_USER_AGENT, parse_civitai_model_path, + normalize_model_image, ) from ..progress import report_progress, get_progress_reporter from ..path_utils import calculate_file_sha256, get_filename_from_path @@ -155,12 +156,6 @@ def _build_civitai_result_from_version( -def _version_sort_key(version: Dict[str, Any]) -> tuple: - return get_version_sort_key(version) - - -def _select_primary_model_file(files: List[Dict[str, Any]]) -> Optional[Dict[str, Any]]: - return select_primary_model_file(files) def _find_model_title_match_in_model( @@ -199,10 +194,10 @@ def _find_model_title_match_in_model( ) return None - versions = sorted(versions, key=_version_sort_key, reverse=True) + versions = sorted(versions, key=get_version_sort_key, reverse=True) for version in versions: - file_info = _select_primary_model_file(version.get("files") or []) + file_info = select_primary_model_file(version.get("files") or []) if not file_info: continue @@ -637,40 +632,6 @@ def _find_civitai_file_in_model( return None -def _extract_civitai_image_id(image_url: str) -> Optional[str]: - """ - Extract a CivitAI image ID from an image CDN URL. - Example: - https://image.civitai.com/.../width=1800/1917130.jpeg -> 1917130 - """ - if not image_url: - return None - - match = re.search(r"/(\d+)(?:\.[A-Za-z0-9]+)?(?:[?#].*)?$", image_url) - if match: - return match.group(1) - - return None - - -def _build_civitai_image_url(img: Dict[str, Any]) -> str: - """ - Build a stable CivitAI image page URL from available image metadata. - """ - civitai_url = img.get("civitaiUrl") - if civitai_url: - return civitai_url - - image_id = img.get("id") - if image_id is not None: - return f"https://civitai.com/images/{image_id}" - - extracted_id = _extract_civitai_image_id(img.get("url", "")) - if extracted_id: - return f"https://civitai.com/images/{extracted_id}" - - return "" - def parse_civitai_url(url: str) -> Optional[Dict[str, Any]]: """ @@ -1275,17 +1236,6 @@ def resolve_urn( return None -def _get_sha256_hash(file_path: str) -> Optional[str]: - """ - Compute sha256 hash of a file by reading it in chunks. - - Args: - file_path: Full path to the file - - Returns: - SHA256 hash as hex string, or None if file doesn't exist - """ - return calculate_file_sha256(file_path) def get_model_info_by_hash( @@ -1499,48 +1449,8 @@ def _extract_model_images(version_info: Dict[str, Any]) -> List[Dict[str, Any]]: if not isinstance(img, dict): continue - # Get the image URL - img_url = img.get("url", "") - if not img_url: - continue - - # Get metadata from image info (may be nested in 'meta' object) - meta = img.get("meta", {}) - if not isinstance(meta, dict): - meta = {} - - img_info = { - "url": img_url, - "civitaiUrl": _build_civitai_image_url(img), - "seed": img.get("seed") or meta.get("seed"), - "steps": img.get("steps") or meta.get("steps"), - "cfg": img.get("cfg") or meta.get("cfg") or meta.get("cfgScale"), - "denoise": img.get("denoise") or meta.get("denoise"), - "scheduler": img.get("scheduler") or meta.get("scheduler"), - "sampler": img.get("sampler") or meta.get("sampler"), - "model": img.get("model") or meta.get("model") or meta.get("Model"), - "positive": img.get("positive") or meta.get("prompt"), - "negative": ( - img.get("negative") - or meta.get("negative_prompt") - or meta.get("negativePrompt") - or meta.get("Negative prompt") - ), - "clip_skip": img.get("clipSkip") or meta.get("Clip skip") or meta.get("clipSkip"), - "width": img.get("width") or meta.get("width"), - "height": img.get("height") or meta.get("height"), - "resources": ( - img.get("resources") - or img.get("additionalResources") - or meta.get("resources") - or meta.get("additionalResources") - or [] - ), - "metadata": meta, - } - - # Only add if we have at least a URL - if img_info["url"]: + img_info = normalize_model_image(img) + if img_info.get("url"): images.append(img_info) return images @@ -1721,37 +1631,9 @@ def _metadata_to_model_info(metadata: Dict[str, Any]) -> Dict[str, Any]: or [] ): if isinstance(img, dict) and img.get("url"): - img_meta = img.get("meta") or {} - images.append( - { - "url": img.get("url", ""), - "civitaiUrl": _build_civitai_image_url(img), - "seed": img_meta.get("seed"), - "steps": img_meta.get("steps"), - "cfg": img_meta.get("cfg") or img_meta.get("cfgScale"), - "denoise": img_meta.get("denoise"), - "scheduler": img_meta.get("scheduler"), - "sampler": img_meta.get("sampler"), - "model": img_meta.get("model") or img_meta.get("Model"), - "positive": img_meta.get("prompt"), - "negative": ( - img_meta.get("negative_prompt") - or img_meta.get("negativePrompt") - or img_meta.get("Negative prompt") - ), - "clip_skip": img_meta.get("Clip skip") or img_meta.get("clipSkip"), - "width": img.get("width") or img_meta.get("width"), - "height": img.get("height") or img_meta.get("height"), - "resources": ( - img.get("resources") - or img.get("additionalResources") - or img_meta.get("resources") - or img_meta.get("additionalResources") - or [] - ), - "metadata": img_meta, - } - ) + img_info = normalize_model_image(img) + if img_info.get("url"): + images.append(img_info) # Get trained words from common sidecar shapes, including LoRA Manager metadata. trained_words = _normalize_metadata_trained_words( @@ -1961,7 +1843,7 @@ def get_model_info_for_file( return result # If no metadata file, compute hash and look up on CivitAI - file_hash = _get_sha256_hash(file_path) + file_hash = calculate_file_sha256(file_path) if not file_hash: return None diff --git a/core/sources/huggingface.py b/core/sources/huggingface.py index 3e790b9..5806dad 100644 --- a/core/sources/huggingface.py +++ b/core/sources/huggingface.py @@ -18,7 +18,7 @@ from ..log_system.log_funcs import create_module_logger log = create_module_logger(__name__) from ..path_utils import write_json_atomic, METADATA_DIR -from ..type_utils import check_credential_http, parse_size_header as _parse_size_header, parse_content_range_size as _parse_content_range_size, extract_response_file_size as _response_file_size +from ..type_utils import check_credential_http, parse_size_header as _parse_size_header, fetch_remote_file_size HF_API_URL = "https://huggingface.co/api" HF_AUTHOR_FALLBACKS = ["Comfy-Org"] @@ -383,43 +383,12 @@ def _fetch_remote_file_size_bytes( if cache_key in _download_size_cache: return _download_size_cache[cache_key] - size = None - try: - response = requests.head( - probe_url, - headers=request_headers, - allow_redirects=True, - timeout=timeout, - ) - try: - if response.status_code < 400: - size = _response_file_size(response) - finally: - response.close() - except Exception as e: - log.debug(f"HuggingFace size HEAD failed: url={probe_url}, error={e}") - - if not size: - try: - response = requests.get( - probe_url, - headers={**request_headers, "Range": "bytes=0-0"}, - allow_redirects=True, - stream=True, - timeout=timeout, - ) - try: - if response.status_code < 400: - size = _response_file_size(response) - finally: - response.close() - except Exception as e: - log.debug(f"HuggingFace size range request failed: url={probe_url}, error={e}") - + size = fetch_remote_file_size(probe_url, headers=request_headers, timeout=timeout) _download_size_cache[cache_key] = size return size + def clean_filename_for_search(filename: str) -> str: """ Clean up filename for better search results. diff --git a/core/sources/lora_manager_archive.py b/core/sources/lora_manager_archive.py index d3c2fcd..16e3b59 100644 --- a/core/sources/lora_manager_archive.py +++ b/core/sources/lora_manager_archive.py @@ -171,9 +171,6 @@ def _extract_search_tokens(query: str) -> Dict[str, List[str]]: } -def _extract_primary_file(files: List[Dict[str, Any]]) -> Optional[Dict[str, Any]]: - return select_primary_model_file(files) - def _load_version_files( conn: sqlite3.Connection, version_id: int @@ -428,7 +425,7 @@ def _build_result_from_row( model_id = row["model_id"] version_id = row["version_id"] files = _load_version_files(conn, version_id) - primary_file = _extract_primary_file(files) + primary_file = select_primary_model_file(files) filename = ( primary_file.get("name") if primary_file diff --git a/core/sources/model_list.py b/core/sources/model_list.py index edbb103..16dbbd4 100644 --- a/core/sources/model_list.py +++ b/core/sources/model_list.py @@ -19,7 +19,7 @@ from ..log_system.log_funcs import create_module_logger log = create_module_logger(__name__) -from ..path_utils import METADATA_DIR +from ..path_utils import METADATA_DIR, read_json_safe, write_json_atomic MODEL_LIST_FILE = os.path.join(METADATA_DIR, "model-list.json") MODEL_LIST_META_FILE = os.path.join(METADATA_DIR, "model-list.meta.json") MODEL_LIST_SOURCE_URL = ( @@ -50,23 +50,14 @@ def _fetch_json_url(url: str) -> Dict[str, Any]: return json.loads(response.read().decode("utf-8")) -def _read_json_file(path: str, default: Any) -> Any: - return read_json_safe(path, default) - - -from ..path_utils import write_json_atomic, read_json_safe - - -def _write_json_file_atomic(path: str, data: Any): - write_json_atomic(path, data, indent=2) def _read_model_list_file() -> Dict[str, Any]: - return _read_json_file(MODEL_LIST_FILE, {"models": []}) + return read_json_safe(MODEL_LIST_FILE, {"models": []}) def _read_model_list_meta() -> Dict[str, Any]: - return _read_json_file(MODEL_LIST_META_FILE, {}) + return read_json_safe(MODEL_LIST_META_FILE, {}) def _get_local_model_list_sha() -> str: @@ -164,7 +155,7 @@ def update_model_list_from_remote() -> Dict[str, Any]: if os.path.exists(MODEL_LIST_FILE): shutil.copy2(MODEL_LIST_FILE, f"{MODEL_LIST_FILE}.bak") - _write_json_file_atomic(MODEL_LIST_FILE, data) + write_json_atomic(MODEL_LIST_FILE, data, indent=2) meta = { "source": "Comfy-Org/ComfyUI-Manager", "source_url": MODEL_LIST_SOURCE_URL, @@ -175,7 +166,7 @@ def update_model_list_from_remote() -> Dict[str, Any]: "model_count": len(models), "updated_at": _utc_now_iso(), } - _write_json_file_atomic(MODEL_LIST_META_FILE, meta) + write_json_atomic(MODEL_LIST_META_FILE, meta, indent=2) reload_model_list() return { diff --git a/core/type_utils.py b/core/type_utils.py index 7cc4d1a..5f98ae9 100644 --- a/core/type_utils.py +++ b/core/type_utils.py @@ -526,3 +526,144 @@ def parse_civitai_model_path(path: str, query_string: str) -> Optional[Dict[str, return result return None + +def fetch_remote_file_size( + url: str, + headers: Optional[Dict[str, str]] = None, + timeout: int = 15, +) -> Optional[int]: + """ + Sprawdza rozmiar zdalnego pliku wykonując zapytanie HEAD i opcjonalnie GET Range=bytes=0-0. + """ + import requests + request_headers = {**(headers or {}), "Accept-Encoding": "identity"} + + size = None + try: + response = requests.head( + url, + headers=request_headers, + allow_redirects=True, + timeout=timeout, + ) + try: + if response.status_code < 400: + size = extract_response_file_size(response) + finally: + response.close() + except Exception: + pass + + if not size: + try: + response = requests.get( + url, + headers={**request_headers, "Range": "bytes=0-0"}, + allow_redirects=True, + stream=True, + timeout=timeout, + ) + try: + if response.status_code < 400: + size = extract_response_file_size(response) + finally: + response.close() + except Exception: + pass + + return size + + +def looks_like_model_file(url: str, expected_filename: str = "") -> bool: + """ + Walidacja, czy dany link URL wygląda na poprawny plik modelu. + """ + import os + from urllib.parse import unquote, urlparse + + text = str(url or "").strip() + if not text.startswith(("http://", "https://")): + return False + try: + parsed = urlparse(text) + host = parsed.netloc.lower() + path = unquote(parsed.path or "") + except Exception: + host = "" + path = text + + if host.endswith("huggingface.co") and path.startswith("/spaces/"): + return False + + # Civitai/CivArchive api/download url prefix check + if ( + host.endswith("civitai.com") + or host.endswith("civitai.red") + or host.endswith("civarchive.com") + ) and "/api/download/" in path: + return True + + basename = os.path.basename(path).lower() + expected = os.path.basename(str(expected_filename or "")).lower() + if expected and basename == expected: + return True + + ext = os.path.splitext(basename)[1].lower() + return ext in MODEL_EXTENSIONS + + +def normalize_model_image(image_data: Dict[str, Any], default_civitai_url: str = "") -> Dict[str, Any]: + """ + Normalizuje metadane obrazu pochodzące z różnych API do jednolitego słownika. + """ + if not isinstance(image_data, dict): + return {} + + meta = image_data.get("meta") + if not isinstance(meta, dict): + meta = {} + + url = image_data.get("url") or image_data.get("imageUrl") or image_data.get("src") or "" + + civitai_url = image_data.get("civitaiUrl") or image_data.get("postUrl") or default_civitai_url + if not civitai_url and isinstance(image_data.get("id"), (int, str)): + civitai_url = f"https://civitai.com/images/{image_data.get('id')}" + + if not civitai_url and url: + import re + match = re.search(r"/(\d+)(?:\.[A-Za-z0-9]+)?(?:[?#].*)?$", str(url)) + if match: + civitai_url = f"https://civitai.com/images/{match.group(1)}" + + + return { + "url": url, + "civitaiUrl": civitai_url, + "seed": image_data.get("seed") or meta.get("seed"), + "steps": image_data.get("steps") or meta.get("steps"), + "cfg": image_data.get("cfg") or meta.get("cfg") or meta.get("cfgScale"), + "denoise": image_data.get("denoise") or meta.get("denoise"), + "scheduler": image_data.get("scheduler") or meta.get("scheduler"), + "sampler": image_data.get("sampler") or meta.get("sampler"), + "model": image_data.get("model") or meta.get("model") or meta.get("Model"), + "positive": image_data.get("positive") or image_data.get("prompt") or meta.get("prompt"), + "negative": ( + image_data.get("negative") + or meta.get("negative_prompt") + or meta.get("negativePrompt") + or meta.get("Negative prompt") + ), + "clip_skip": image_data.get("clipSkip") or meta.get("Clip skip") or meta.get("clipSkip"), + "width": image_data.get("width") or meta.get("width"), + "height": image_data.get("height") or meta.get("height"), + "resources": ( + image_data.get("resources") + or image_data.get("additionalResources") + or meta.get("resources") + or meta.get("additionalResources") + or [] + ), + "metadata": meta, + } + +