Add URN (ComfyUI-AIR) resolution via CivitAI

Detect ComfyUI-AIR URNs in workflows, resolve them to expected filenames via the CivitAI API, and prefer those filenames for matching.

Key changes:
- core/workflow_analyzer.py: added URN_REGEX, URN_TYPE_MAP, is_model_filename() now recognizes URNs, and get_node_model_info() emits structured URN info (is_urn, urn fields, category, exists=False).
- core/sources/civitai.py: added resolve_urn(model_id, version_id) with caching, API fetch of model/version and primary file, and result shaping (model_name, version_name, expected_filename, files).
- core/linker.py: import resolve_urn and call it for missing URNs, attach civitai_info/expected_filename, and prefer expected_filename as matching target; cleaned up deduplication logic.
- web/linker.js: UI shows URN original plus resolved CivitAI model/version labels where available.
- TODO.md: added plan/notes for URN detection and next steps.

This improves detection and fuzzy matching for models referenced by URNs so missing model resolution is more accurate. Further matcher and updater enhancements remain TODOs.

TODO.md

@@ -0,0 +1,50 @@
+# ComfyUI Model Linker - Fix URN Detection Plan
+
+## Overview
+Fix plugin to detect ComfyUI-AIR URNs as missing models, resolve to expected filenames via CivitAI API, fuzzy match locals.
+
+**Status: In Progress**
+
+## Steps
+
+### 1. [x] URN Parsing in workflow_analyzer.py
+- Added URN_REGEX, URN_TYPE_MAP
+- Updated is_model_filename() to detect URNs
+- Modified get_node_model_info(): parse URN → 'urn' dict, category mapping, is_urn=True, exists=False
+- Add URN regex to `is_model_filename()`
+- Parse components: base, type→category, provider, model_id, version_id
+- Treat URNs as always "missing" (no folder_paths check)
+
+### 2. [x] CivitAI URN Resolver in sources/civitai.py
+- Added _urn_cache, resolve_urn(model_id, version_id) → model_name, expected_filename (primary file)
+- API: /models/{model_id}?modelVersionId={version_id}
+- Full files list, logging, error handling
+- `resolve_urn(model_id, version_id)` → API https://civitai.com/api/v1/models/{model_id}?modelVersionId={version_id}
+- Extract modelVersions[0].files[primary].name as expected_filename
+- Cache results
+
+### 3. [x] Integrate URN Handling in linker.py
+- Added import .sources.civitai.resolve_urn
+- In analyze_and_find_matches(): resolve_urn() for is_urn=True → civitai_info, expected_filename
+- Matching: use expected_filename for URNs (better fuzzy), else original_path
+- In analyze_and_find_matches(): For URN refs, fetch expected_filename
+- Add to missing_models: 'urn_original', 'expected_filename'
+
+### 4. [ ] Enhance matcher.py
+- Prioritize matching on expected_filename if present
+
+### 5. [ ] Workflow Updater Support (Optional)
+- Replace URN → relative_path in update_model_path()
+
+### 6. [ ] Testing
+- Create sample workflow.json with URNs
+- Run plugin, verify detection/matches
+- Restart ComfyUI
+
+### 7. [ ] Polish
+- UI updates (web/linker.js)
+- Error handling (API fails)
+- Rate limit/cache
+
+**Next: Step 2 - Add resolver to core/sources/civitai.py**
+

core/linker.py

@@ -15,6 +15,7 @@ from .scanner import get_model_files
 from .workflow_analyzer import analyze_workflow_models, identify_missing_models
 from .matcher import find_matches
 from .workflow_updater import update_workflow_nodes
+from .sources.civitai import resolve_urn
 
 logger = logging.getLogger(__name__)
 
@@ -224,17 +225,32 @@ def analyze_and_find_matches(
             missing['workflow_directory'] = url_info.get('directory', '')
             missing['url_source'] = url_info.get('source', '')
     
+    # Handle URNs: fetch expected filename from CivitAI
+    for missing in missing_models:
+        if missing.get('is_urn'):
+            urn = missing['urn']
+            model_info = resolve_urn(urn['model_id'], urn['version_id'])
+            if model_info:
+                missing['civitai_info'] = {
+                    'model_name': model_info.get('model_name'),
+                    'version_name': model_info.get('version_name')
+                }
+            if model_info:
+                missing['civitai_info'] = model_info
+                missing['expected_filename'] = model_info['expected_filename']
+                logger.debug(f"URN {missing['original_path']} → {missing['expected_filename']}")
+    
     # Find matches for each missing model
     missing_with_matches = []
     for missing in missing_models:
-        original_path = missing.get('original_path', '')
+        target_for_matching = missing.get('original_path', '')
+        
+        # For URNs, prefer expected_filename for matching
+        if missing.get('is_urn') and missing.get('expected_filename'):
+            target_for_matching = missing['expected_filename']
         
         # Filter available models by category if known
-        # IMPORTANT: If category is 'unknown', we still try to find the right category
-        # by using node type hints
         category = missing.get('category')
-        
-        # If category is unknown, try to use node type to infer category
         if not category or category == 'unknown':
             from .workflow_analyzer import NODE_TYPE_TO_CATEGORY_HINTS
             node_type = missing.get('node_type', '')
@@ -242,41 +258,32 @@ def analyze_and_find_matches(
         
         candidates = available_models
         if category and category != 'unknown':
-            # Prioritize models from the same category
             candidates = [m for m in available_models if m.get('category') == category]
-            # Also include other categories as fallback
             candidates.extend([m for m in available_models if m.get('category') != category])
         
         # Find matches
         matches = find_matches(
-            original_path,
+            target_for_matching,
             candidates,
             threshold=similarity_threshold,
             max_results=max_matches_per_model
         )
         
-        # Deduplicate matches by absolute path - same physical file should only appear once
-        # This handles cases where the same file exists in multiple base directories
-        # or has different relative_paths but is the same file
+        # Deduplicate matches by absolute path
         seen_absolute_paths = {}
         deduplicated_matches = []
         for match in matches:
             model_dict = match['model']
             absolute_path = model_dict.get('path', '')
-            
-            # Normalize absolute path for comparison
             if absolute_path:
                 absolute_path = os.path.normpath(absolute_path)
             
-            # If we haven't seen this absolute path, add it
             if absolute_path not in seen_absolute_paths:
                 seen_absolute_paths[absolute_path] = match
                 deduplicated_matches.append(match)
             else:
-                # If we've seen this absolute path before, replace with better match if confidence is higher
                 existing_match = seen_absolute_paths[absolute_path]
                 if match['confidence'] > existing_match['confidence']:
-                    # Replace with better match
                     idx = deduplicated_matches.index(existing_match)
                     deduplicated_matches[idx] = match
                     seen_absolute_paths[absolute_path] = match

core/sources/civitai.py

@@ -15,8 +15,9 @@ logger = logging.getLogger(__name__)
 
 CIVITAI_API_URL = "https://civitai.com/api/v1"
 
-# Cache for search results
+# Cache for search results and URN resolutions
 _search_cache: Dict[str, Any] = {}
+_urn_cache: Dict[tuple[int, int], Dict[str, Any]] = {}
 
 
 def parse_civitai_url(url: str) -> Optional[Dict[str, Any]]:
@@ -307,3 +308,92 @@ def search_civitai_by_hash(
         logger.error(f"CivitAI hash lookup error: {e}")
     
     return None
+
+
+def resolve_urn(
+    model_id: int, 
+    version_id: int, 
+    api_key: Optional[str] = None
+) -> Optional[Dict[str, Any]]:
+    """
+    Resolve URN model_id/version_id to model info and expected filename.
+    
+    Args:
+        model_id: CivitAI model ID
+        version_id: CivitAI version ID  
+        api_key: Optional API key
+        
+    Returns:
+        Dict with model name and primary filename, or None
+    """
+    cache_key = (model_id, version_id)
+    if cache_key in _urn_cache:
+        return _urn_cache[cache_key]
+    
+    try:
+        headers = {}
+        if api_key:
+            headers['Authorization'] = f'Bearer {api_key}'
+        
+        url = f"{CIVITAI_API_URL}/models/{model_id}"
+        params = {'modelVersionId': version_id}
+        
+        response = requests.get(url, headers=headers, params=params, timeout=15)
+        
+        if response.status_code != 200:
+            logger.warning(f"CivitAI URN resolve failed: {response.status_code}")
+            _urn_cache[cache_key] = None
+            return None
+        
+        data = response.json()
+        versions = data.get('modelVersions', [])
+        
+        if not versions:
+            logger.warning(f"No versions found for model {model_id}/version {version_id}")
+            _urn_cache[cache_key] = None
+            return None
+        
+        # Get specific version
+        target_version = None
+        for v in versions:
+            if v.get('id') == version_id:
+                target_version = v
+                break
+        
+        if not target_version:
+            logger.warning(f"Version {version_id} not found in model {model_id}")
+            _urn_cache[cache_key] = None
+            return None
+        
+        files = target_version.get('files', [])
+        primary_file = None
+        
+        # Prefer primary file or type=='Model'
+        for f in files:
+            if f.get('primary') or f.get('type') == 'Model':
+                primary_file = f
+                break
+        
+        if not primary_file and files:
+            primary_file = files[0]  # Fallback to first
+        
+        if not primary_file:
+            logger.warning(f"No files found for version {version_id}")
+            _urn_cache[cache_key] = None
+            return None
+        
+        result = {
+            'model_name': data.get('name', 'Unknown'),
+            'version_name': target_version.get('name', 'Unknown'),
+            'expected_filename': primary_file.get('name', 'Unknown'),
+            'files': [{'name': f.get('name'), 'size': f.get('sizeKB', 0)*1024} for f in files]
+        }
+        
+        _urn_cache[cache_key] = result
+        logger.info(f"Resolved URN model {model_id}@{version_id} → {result['expected_filename']}")
+        return result
+        
+    except Exception as e:
+        logger.error(f"CivitAI URN resolve error for {model_id}@{version_id}: {e}")
+        _urn_cache[cache_key] = None
+        return None

core/workflow_analyzer.py

@@ -19,6 +19,24 @@ except ImportError:
 # Common model file extensions
 MODEL_EXTENSIONS = {'.ckpt', '.pt', '.pt2', '.bin', '.pth', '.safetensors', '.pkl', '.sft', '.onnx'}
 
+import re
+
+# URN Type to ComfyUI category mapping
+URN_TYPE_MAP = {
+    'checkpoint': 'checkpoints',
+    'lora': 'loras', 
+    'vae': 'vae',
+    'upscaler': 'upscale_models',
+    'upscale_model': 'upscale_models',
+    'embedding': 'embeddings',
+    'hypernetwork': 'hypernetworks',
+    'controlnet': 'controlnet',
+    'clip': 'clip',
+    'clip_vision': 'clip_vision'
+}
+
+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')
@@ -41,20 +59,24 @@ NODE_TYPE_TO_CATEGORY_HINTS = {
 
 def is_model_filename(value: Any) -> bool:
     """
-    Check if a value looks like a model filename.
+    Check if a value looks like a model filename or URN.
     
     Args:
         value: The value to check
         
     Returns:
-        True if it looks like a model filename
+        True if it looks like a model filename or URN
     """
     if not isinstance(value, str):
         return False
     
-    # Check if it ends with a model extension
+    # Check model extension
     _, ext = os.path.splitext(value.lower())
-    return ext in MODEL_EXTENSIONS
+    if ext in MODEL_EXTENSIONS:
+        return True
+    
+    # Check URN format
+    return bool(URN_REGEX.match(value.strip()))
 
 
 def try_resolve_model_path(value: str, categories: List[str] = None) -> Optional[tuple[str, str]]:
@@ -137,20 +159,46 @@ def get_node_model_info(node: Dict[str, Any]) -> List[Dict[str, Any]]:
     category_hint = NODE_TYPE_TO_CATEGORY_HINTS.get(node_type)
     categories_to_try = [category_hint] if category_hint else None
     
-    # For each widget value, check if it looks like a model file
+    # For each widget value, check if it looks like a model file or URN
     for idx, value in enumerate(widgets_values):
         if not is_model_filename(value):
             continue
         
-        # Try to resolve the model path
-        resolved = try_resolve_model_path(value, categories_to_try)
+        value_str = str(value).strip()
+        
+        # Check if URN
+        urn_match = URN_REGEX.match(value_str)
+        if urn_match:
+            base, typ, provider, model_id, version_id = urn_match.groups()
+            category = URN_TYPE_MAP.get(typ.lower(), 'unknown')
+            
+            model_refs.append({
+                'node_id': node_id,
+                'node_type': node_type,
+                'widget_index': idx,
+                'original_path': value_str,
+                'urn': {
+                    'full': value_str,
+                    'base': base,
+                    'type': typ,
+                    'provider': provider,
+                    'model_id': int(model_id),
+                    'version_id': int(version_id)
+                },
+                'category': category,
+                'full_path': None,
+                'exists': False,  # URNs always need resolution
+                'is_urn': True
+            })
+            continue
+        
+        # Existing logic for local filenames
+        resolved = try_resolve_model_path(value_str, categories_to_try)
         
         if resolved:
             category, full_path = resolved
             exists = os.path.exists(full_path)
         else:
-            # If we can't resolve it, check if it at least looks like a model filename
-            # This might be a missing model or a custom node's model
             category = category_hint or 'unknown'
             full_path = None
             exists = False
@@ -159,10 +207,11 @@ def get_node_model_info(node: Dict[str, Any]) -> List[Dict[str, Any]]:
             'node_id': node_id,
             'node_type': node_type,
             'widget_index': idx,
-            'original_path': value,
+            'original_path': value_str,
             'category': category,
             'full_path': full_path,
-            'exists': exists
+            'exists': exists,
+            'is_urn': False
         })
     
     return model_refs

web/linker.js

@@ -1029,7 +1029,23 @@ class LinkerManagerDialog extends ComfyDialog {
         
         // Card Header: Filename as headline + node chip
         html += `<div class="ml-card-header">`;
-        html += `<h3 class="ml-card-title" title="${missingFilename.full}">${missingFilename.display}</h3>`;
+        // Show URN + CivitAI model name + version name if available
+        let titleHtml = `<span title="${missingFilename.full}">${missingFilename.display}</span>`;
+        if (missing.is_urn && missing.civitai_info) {
+            let civitaiLabel = '';
+            if (missing.civitai_info.model_name) {
+                civitaiLabel += `"${missing.civitai_info.model_name}"`;
+            }
+            if (missing.civitai_info.version_name && missing.civitai_info.version_name !== missing.civitai_info.model_name) {
+                civitaiLabel += ` ${missing.civitai_info.version_name}`;
+            }
+            if (civitaiLabel) {
+                titleHtml += ` <span style="color: var(--ml-text-muted); font-size: 14px;">${civitaiLabel}</span>`;
+            }
+        } else if (missing.civitai_info && missing.civitai_info.model_name) {
+            titleHtml += ` <span style="color: var(--ml-text-muted); font-size: 14px;">(${missing.civitai_info.model_name})</span>`;
+        }
+        html += `<h3 class="ml-card-title">${titleHtml}</h3>`;
         html += `<div style="display: flex; align-items: center; gap: 6px;">`;
         if (missing.category) {
             html += `<span class="ml-category-chip">${missing.category}</span>`;