diff --git a/__init__.py b/__init__.py index b344491..e4f62cc 100644 --- a/__init__.py +++ b/__init__.py @@ -189,6 +189,7 @@ class ModelResolverExtension: format_size_bytes, normalize_sha256, extract_sha256_from_metadata, + fetch_remote_file_size_cached, ) from .core.settings import ( TEMPLATE_KEY_ALIASES, @@ -426,24 +427,8 @@ class ModelResolverExtension: else: source = "workflow" - # Try to get file size with HEAD request (non-blocking, timeout quickly) - file_size = None - try: - import requests - - head_response = requests.head( - workflow_url, - allow_redirects=True, - timeout=5, - ) - if head_response.status_code == 200: - file_size = int( - head_response.headers.get( - "content-length", 0 - ) - ) - except Exception: - pass # Size unknown is fine + # Try to get file size using cached remote helper + file_size = fetch_remote_file_size_cached(workflow_url, timeout=5) missing["download_source"] = { "source": source, diff --git a/core/downloader.py b/core/downloader.py index 9059037..37a6229 100644 --- a/core/downloader.py +++ b/core/downloader.py @@ -24,7 +24,7 @@ log = create_module_logger(__name__) from .resolver import normalize_sha256 from .path_utils import is_path_within, get_path_identity, write_json_atomic, read_json_safe, get_comfy_root_path, calculate_file_sha256 as _calculate_file_sha256, get_filename_from_path -from .type_utils import as_dict, as_list, first_non_empty, format_size_bytes as format_bytes +from .type_utils import as_dict, as_list, first_non_empty, format_size_bytes as format_bytes, DEFAULT_BROWSER_USER_AGENT, extract_response_file_size try: import folder_paths @@ -51,11 +51,7 @@ CHUNK_SIZE = 1024 * 1024 # 1MB chunks for faster downloads CLI_LOG_INTERVAL = 5 # Log progress to CLI every N seconds ARIA2_RPC_TIMEOUT = (2, 5) # local JSON-RPC should respond quickly ARIA2_IDLE_STOP_SECONDS = 5 * 60 -DOWNLOAD_USER_AGENT = ( - "Mozilla/5.0 (Windows NT 10.0; Win64; x64) " - "AppleWebKit/537.36 (KHTML, like Gecko) " - "Chrome/126.0.0.0 Safari/537.36" -) +DOWNLOAD_USER_AGENT = DEFAULT_BROWSER_USER_AGENT from .settings import CATEGORY_MAP, load_settings, normalize_download_backend, normalize_download_category @@ -1464,7 +1460,7 @@ def download_file( response.raise_for_status() # Get total size - total_size = int(response.headers.get("content-length", 0)) + total_size = extract_response_file_size(response) or 0 total_size_str = format_bytes(total_size) if total_size > 0 else "unknown" log.info(f"Size: {total_size_str}") diff --git a/core/matcher.py b/core/matcher.py index 43c9327..260559d 100644 --- a/core/matcher.py +++ b/core/matcher.py @@ -342,13 +342,9 @@ def calculate_filename_confidence(target_filename: str, candidate_filename: str) if target_norm == candidate_norm: return 100.0 - similarity = calculate_similarity_with_normalization( + best_similarity = calculate_similarity_with_normalization( target_filename, candidate_filename ) - similarity_no_ext = calculate_similarity_with_normalization( - os.path.splitext(target_filename)[0], os.path.splitext(candidate_filename)[0] - ) - best_similarity = max(similarity, similarity_no_ext) # Cap similarity at 0.999 for non-exact matches to prevent false 100% scores if best_similarity >= 0.999: @@ -375,12 +371,7 @@ def calculate_archived_model_confidence( if query_norm == candidate_norm: return 100.0 - similarity = calculate_similarity_with_normalization(query, candidate) - similarity_no_ext = calculate_similarity_with_normalization( - os.path.splitext(query)[0], - os.path.splitext(candidate)[0], - ) - score = max(similarity, similarity_no_ext) + score = calculate_similarity_with_normalization(query, candidate) if query_norm and candidate_norm and ( query_norm in candidate_norm or candidate_norm in query_norm ): @@ -401,3 +392,32 @@ def calculate_candidate_rank( return matches, confidence + score +def build_filename_search_queries(filename: str) -> List[str]: + """ + Build a list of candidate search queries from a filename by stripping + extensions and common precision/format suffixes (e.g. fp16, bf16, pruned). + """ + basename = os.path.basename(filename or "").strip() + if not basename: + return [] + + stem = os.path.splitext(basename)[0].strip() + stem_lower = stem.lower() + + # Pattern to match common precision/format suffixes + pattern = r"[-_]?(fp16|fp32|fp8|fp4|bf16|e4m3fn|mixed|scaled|pruned|emaonly|q4|q8)" + + # 1. Strip suffix followed by any other characters (like clean_filename_for_search) + cleaned = re.sub(pattern + r".*$", "", stem, flags=re.IGNORECASE).strip(" -_") + cleaned_lower = cleaned.lower() + + # 2. Strip suffix at the end of the stem (like simplified in search queries) + simplified = re.sub(pattern + r"$", "", stem_lower, flags=re.IGNORECASE).strip(" -_") + + queries = [] + for query in [basename, stem, stem_lower, cleaned, cleaned_lower, simplified]: + query_val = (query or "").strip(" -_") + if query_val and query_val not in queries: + queries.append(query_val) + + return queries diff --git a/core/sources/civarchive.py b/core/sources/civarchive.py index 7f9ab56..488b713 100644 --- a/core/sources/civarchive.py +++ b/core/sources/civarchive.py @@ -22,6 +22,7 @@ from ..matcher import ( base_model_matches as _base_model_matches, base_model_score as _base_model_score, calculate_candidate_rank, + build_filename_search_queries, ) from ..path_utils import get_filename_from_path from ..type_utils import ( @@ -2076,20 +2077,8 @@ def _is_hash_verified_exact_match(result: Optional[Dict[str, Any]], confidence: def _build_search_queries(filename: str) -> List[str]: - basename = get_filename_from_path(filename or "").strip() - stem = os.path.splitext(basename)[0].strip() - simplified = re.sub( - r"[-_]?(fp16|fp32|fp8|fp4|bf16|e4m3fn|scaled|pruned|emaonly|mixed|q4|q8)$", - "", - stem, - flags=re.IGNORECASE, - ).strip(" -_") - - queries = [] - for query in [basename, stem, simplified]: - if query and query not in queries: - queries.append(query) - return queries + # Wrap the unified query builder helper + return build_filename_search_queries(filename) def search_civarchive( diff --git a/core/sources/huggingface.py b/core/sources/huggingface.py index b3315a5..8abe5e0 100644 --- a/core/sources/huggingface.py +++ b/core/sources/huggingface.py @@ -19,6 +19,7 @@ log = create_module_logger(__name__) from ..path_utils import write_json_atomic, METADATA_DIR, read_json_safe from ..type_utils import check_credential_http, parse_size_header as _parse_size_header, fetch_remote_file_size, fetch_remote_file_size_cached, clear_remote_size_cache +from ..matcher import build_filename_search_queries HF_API_URL = "https://huggingface.co/api" HF_AUTHOR_FALLBACKS = ["Comfy-Org"] @@ -558,22 +559,8 @@ def _find_matching_file_in_author_index( def _build_huggingface_search_queries(filename: str) -> List[str]: - filename_base = os.path.splitext(os.path.basename(filename))[0].lower() - cleaned_base = clean_filename_for_search(filename).lower() - simplified_base = re.sub( - r"[-_]?(fp16|fp32|fp8|fp4|bf16|e4m3fn|mixed|scaled|pruned|emaonly|q4|q8)$", - "", - filename_base, - flags=re.IGNORECASE, - ) - - queries = [] - for query in [filename_base, cleaned_base, simplified_base]: - query = (query or "").strip(" -_") - if query and query not in queries: - queries.append(query) - - return queries + # Wrap the unified query builder helper + return build_filename_search_queries(filename) def _get_repos_by_author( diff --git a/core/workflow_analyzer.py b/core/workflow_analyzer.py index 46ffbf9..8b14b60 100644 --- a/core/workflow_analyzer.py +++ b/core/workflow_analyzer.py @@ -28,74 +28,7 @@ URN_REGEX = re.compile(r"^urn:air:([^:]+):([^:]+):([^:]+):(\d+)@(\d+)$") # Mapping of common node types to their expected model category # This is used as hints but we don't rely solely on this # UNETLoader uses 'diffusion_models' category (folder_paths maps 'unet' to 'diffusion_models') -NODE_TYPE_TO_CATEGORY_HINTS = { - "CheckpointLoaderSimple": "checkpoints", - "CheckpointLoader": "checkpoints", - "DiffusersLoader": "diffusers", - "unCLIPCheckpointLoader": "checkpoints", - "ImageOnlyCheckpointLoader": "checkpoints", - "VAELoader": "vae", - "VAELoaderKJ": "vae", - "LoraLoader": "loras", - "LoraLoaderModelOnly": "loras", - "LoraLoaderBypass": "loras", - "LoraLoaderBypassModelOnly": "loras", - "LoraLoaderV2": "loras", - "Lora Loader (LoraManager)": "loras", # LoraManager custom node - "Lora Stacker (LoraManager)": "loras", # LoraManager Stacker node - "Power Lora Loader (rgthree)": "loras", # rgthree's Power Lora Loader - "CreateHookLora": "loras", - "CreateHookLoraModelOnly": "loras", - "CreateHookModelAsLora": "checkpoints", - "CreateHookModelAsLoraModelOnly": "checkpoints", - "UNETLoader": "diffusion_models", - "LoaderGGUF": "diffusion_models", - "LoaderGGUFAdvanced": "diffusion_models", - "UnetLoaderGGUF": "diffusion_models", - "UnetLoaderGGUFAdvanced": "diffusion_models", - "LatentUpscaleModelLoader": "latent_upscale_models", - "CLIPLoader": "text_encoders", - "DualCLIPLoader": "text_encoders", - "CLIPLoaderGGUF": "text_encoders", - "ClipLoaderGGUF": "text_encoders", - "DualCLIPLoaderGGUF": "text_encoders", - "DualClipLoaderGGUF": "text_encoders", - "TripleCLIPLoader": "text_encoders", - "TripleClipLoader": "text_encoders", - "TripleCLIPLoaderGGUF": "text_encoders", - "TripleClipLoaderGGUF": "text_encoders", - "QuadrupleCLIPLoader": "text_encoders", - "QuadrupleClipLoader": "text_encoders", - "QuadrupleCLIPLoaderGGUF": "text_encoders", - "QuadrupleClipLoaderGGUF": "text_encoders", - "ControlNetLoader": "controlnet", - "DiffControlNetLoader": "controlnet", - "ControlNetLoaderAdvanced": "controlnet", - "ACN_ControlNetLoaderAdvanced": "controlnet", - "ACN_DiffControlNetLoaderAdvanced": "controlnet", - "CLIPVisionLoader": "clip_vision", - "StyleModelLoader": "style_models", - "GLIGENLoader": "gligen", - "UpscaleModelLoader": "upscale_models", - "SAMLoader": "sams", - "UltralyticsDetectorProvider": "ultralytics", - "AudioEncoderLoader": "audio_encoders", - "LoadBackgroundRemovalModel": "background_removal", - "LoadDA3Model": "geometry_estimation", - "FrameInterpolationModelLoader": "frame_interpolation", - "LoadMediaPipeFaceLandmarker": "detection", - "ModelPatchLoader": "model_patches", - "LoadMoGeModel": "geometry_estimation", - "PhotoMakerLoader": "photomaker", - "OpticalFlowLoader": "optical_flow", - "HypernetworkLoader": "hypernetworks", - "EmbeddingLoader": "embeddings", - # LTX-Video nodes - "LTXVAudioVAELoader": "checkpoints", - "LowVRAMAudioVAELoader": "checkpoints", - "LTXVGemmaCLIPModelLoader": "text_encoders", - "LTXAVTextEncoderLoader": "text_encoders", -} +NODE_TYPE_TO_CATEGORY_HINTS = {} # Workflow widget_values do not always include file extensions. ComfyUI still # validates these combo widgets against folder_paths by exact value at queue time. @@ -199,6 +132,22 @@ NODE_TYPE_MODEL_WIDGET_CATEGORIES = { "LTXAVTextEncoderLoader": {0: "text_encoders", 1: "checkpoints"}, } +# Mapping of common node types to their expected model category. +# Derived dynamically from NODE_TYPE_MODEL_WIDGET_CATEGORIES. +NODE_TYPE_TO_CATEGORY_HINTS = { + node_type: widget_map[min(widget_map.keys())] + for node_type, widget_map in NODE_TYPE_MODEL_WIDGET_CATEGORIES.items() + if widget_map +} + +# Fallbacks and overrides for custom nodes that do not use standard model widget indices +NODE_TYPE_TO_CATEGORY_HINTS.update({ + "LoraLoaderV2": "loras", + "Lora Loader (LoraManager)": "loras", + "Lora Stacker (LoraManager)": "loras", + "Power Lora Loader (rgthree)": "loras", +}) + # Model category hints by widget/input name. Workflow JSON does not always preserve # widget names, but when it does this catches custom loaders without a node-type entry. MODEL_WIDGET_NAME_TO_CATEGORY = { diff --git a/web/resolver/actions/resolve_download_methods.js b/web/resolver/actions/resolve_download_methods.js index c4a84ef..74fc041 100644 --- a/web/resolver/actions/resolve_download_methods.js +++ b/web/resolver/actions/resolve_download_methods.js @@ -4,6 +4,7 @@ import { $el } from "../../../../../scripts/ui.js"; import { createModuleLogger } from "../../log_system/log_funcs.js"; import { getSvgIcon } from "../../utils/icon_utils.js"; import { getModelCardUrl } from "../utils/url_utils.js"; +import { getCivitaiModelUrl } from "../globals.js"; const log = createModuleLogger('resolve_download_methods'); @@ -2245,7 +2246,7 @@ export const resolveDownloadMethods = { const civitai = data.civitai; const labelHtml = this.renderVersionedModelNameHtml(civitai.name, civitai.version_name) || this.escapeHtml(civitai.filename || 'Model'); - const url = modelUrl || `https://civitai.com/models/${modelId}?modelVersionId=${versionId}`; + const url = modelUrl || getCivitaiModelUrl(modelId, versionId); loadingEl.innerHTML = `${labelHtml}`; } else if (loadingEl) { loadingEl.textContent = 'Not found'; @@ -2684,7 +2685,7 @@ export const resolveDownloadMethods = { } if (civitaiResult && civitaiResult.download_url) { - const modelUrl = civitaiResult.url || (civitaiResult.model_id ? `https://civitai.com/models/${civitaiResult.model_id}${civitaiResult.version_id ? `?modelVersionId=${civitaiResult.version_id}` : ''}` : ''); + const modelUrl = civitaiResult.url || getCivitaiModelUrl(civitaiResult.model_id, civitaiResult.version_id); const downloadFilename = civitaiResult.filename || missing.civitai_info?.expected_filename || civitaiResult.name; const modelName = civitaiResult.name || missing.civitai_info?.model_name || downloadFilename || 'Model'; const civitaiCategory = this.getSourceResultDownloadCategory?.( diff --git a/web/resolver/globals.js b/web/resolver/globals.js index 7115474..cde166b 100644 --- a/web/resolver/globals.js +++ b/web/resolver/globals.js @@ -1,4 +1,4 @@ -export function registerGlobalHelpers() { +export function registerGlobalHelpers() { // Global helper functions for inline onclick handlers window.MLToggleHidden = function(id, trigger, collapsedText, expandedText) { const element = document.getElementById(id); @@ -130,3 +130,8 @@ } }; } + +export function getCivitaiModelUrl(modelId, versionId) { + if (!modelId) return ''; + return `https://civitai.com/models/${encodeURIComponent(modelId)}${versionId ? `?modelVersionId=${encodeURIComponent(versionId)}` : ''}`; +} diff --git a/web/resolver/search/search_panel.js b/web/resolver/search/search_panel.js index ca727bc..8905723 100644 --- a/web/resolver/search/search_panel.js +++ b/web/resolver/search/search_panel.js @@ -3,6 +3,7 @@ import { api } from "../../../../../scripts/api.js"; import { $el } from "../../../../../scripts/ui.js"; import { getSvgIcon } from "../../utils/icon_utils.js"; import { getModelCardUrl } from "../utils/url_utils.js"; +import { getCivitaiModelUrl } from "../globals.js"; export const searchPanelMethods = { /** * Build stable cache key for a missing model entry @@ -2250,7 +2251,7 @@ export const searchPanelMethods = { size: civitai.size, model_id: civitai.model_id || ids.modelId, version_id: civitai.version_id || ids.versionId, - model_url: civitai.url || `https://civitai.com/models/${ids.modelId}?modelVersionId=${ids.versionId}`, + model_url: civitai.url || getCivitaiModelUrl(ids.modelId, ids.versionId), base_model: civitai.base_model, tags: civitai.tags || [], searchedAt diff --git a/web/resolver/views/missing_browser_methods.js b/web/resolver/views/missing_browser_methods.js index d69d849..7eae6a3 100644 --- a/web/resolver/views/missing_browser_methods.js +++ b/web/resolver/views/missing_browser_methods.js @@ -5,6 +5,7 @@ import { getSvgIcon } from "../../utils/icon_utils.js"; import { createFloatingTreePicker } from "../utils/tree_picker.js"; import { showNotification as showNotificationUtils } from "../../utils/notification_utils.js"; import { startSplitterDrag } from "../utils/splitter_drag.js"; +import { getCivitaiModelUrl } from "../globals.js"; export const missingBrowserMethods = { getMissingFilename(missing = {}) { @@ -2021,7 +2022,7 @@ export const missingBrowserMethods = { const modelId = missing.urn_model_id || missing.urn?.model_id; const versionId = missing.urn_version_id || missing.urn?.version_id; - const modelUrl = missing.is_urn && modelId ? `https://civitai.com/models/${modelId}${versionId ? '?modelVersionId=' + versionId : ''}` : ''; + const modelUrl = missing.is_urn && modelId ? getCivitaiModelUrl(modelId, versionId) : ''; const urnLoadingId = `urn-loading-${missing.node_id}-${missing.widget_index}`; if (missing.is_urn) { diff --git a/web/resolver/views/model_info_methods.js b/web/resolver/views/model_info_methods.js index ca07346..04e0ae1 100644 --- a/web/resolver/views/model_info_methods.js +++ b/web/resolver/views/model_info_methods.js @@ -4,6 +4,7 @@ import { $el } from "../../../../../scripts/ui.js"; import { getSvgIcon } from "../../utils/icon_utils.js"; import { escapeHtml, escapeJsString, sanitizeDescriptionHtml } from "../utils/html_utils.js"; import { getModelCardUrl } from "../utils/url_utils.js"; +import { getCivitaiModelUrl } from "../globals.js"; export const modelInfoMethods = { /** * Handle click outside context menu to hide it @@ -270,7 +271,7 @@ export const modelInfoMethods = { const modelId = model?.model_id || model?.modelId || model?.civitai_model_id || model?.civitaiModelId || model?.id || ''; const versionId = model?.version_id || model?.versionId || model?.modelVersionId || model?.civitai_model_version_id || model?.civitaiModelVersionId || ''; if (sourceKey === 'civitai' && modelId) { - return `https://civitai.com/models/${modelId}${versionId ? `?modelVersionId=${versionId}` : ''}`; + return getCivitaiModelUrl(modelId, versionId); } if (sourceKey === 'civarchive' && modelId) { return `https://civarchive.com/models/${modelId}${versionId ? `?modelVersionId=${versionId}` : ''}`; @@ -1080,7 +1081,7 @@ export const modelInfoMethods = { const modelId = model.model_id || model.modelId; const versionId = model.version_id || model.versionId; if (modelId) { - return `https://civitai.com/models/${modelId}${versionId ? `?modelVersionId=${versionId}` : ''}`; + return getCivitaiModelUrl(modelId, versionId); } return ''; }, @@ -2366,7 +2367,7 @@ export const modelInfoMethods = { const versionId = resource.modelVersionId || resource.versionId || resource.model_version_id; const url = resource.url || resource.modelUrl - || (modelId ? `https://civitai.com/models/${modelId}${versionId ? `?modelVersionId=${versionId}` : ''}` : ''); + || getCivitaiModelUrl(modelId, versionId); return { name: resource.name || resource.modelName || resource.model || resource.hash || 'Resource', version: resource.versionName || resource.modelVersionName || resource.version || '', @@ -2811,7 +2812,7 @@ export const modelInfoMethods = { || contextModel.civitaiModelVersionId; if (modelId) { - return `https://civitai.com/models/${encodeURIComponent(modelId)}${versionId ? `?modelVersionId=${encodeURIComponent(versionId)}` : ''}`; + return getCivitaiModelUrl(modelId, versionId); } const explicitUrl = selectedVersion.platform_url