refactor: centralize repeated filename/search/URL/HTTP helpers

Replace duplicated filename-to-query logic with a single \uild_filename_search_queries\ helper in matcher/civarchive/huggingface, standardize Civitai URL generation in the frontend via \getCivitaiModelUrl\, unify browser user-agent and response file-size extraction in the downloader path, and derive workflow node-type category hints from widget categories rather than maintaining a large static mapping.
2026-07-02 22:56:21 +02:00
parent 8ce70c30e6
commit ce29f88c6d
11 changed files with 78 additions and 143 deletions
@@ -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,
@@ -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}")

@@ -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
@@ -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(
@@ -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(
@@ -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 = {
@@ -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 = `<a href="${url}" target="_blank" class="mr-inline-civitai-link">${labelHtml}</a>`;
                } 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?.(
@@ -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)}` : ''}`;
+}
@@ -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
@@ -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) {
@@ -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