From e8939d4bf1c9befbeba5fd4ae79e9451b212580a Mon Sep 17 00:00:00 2001 From: Dariusz L Date: Fri, 17 Apr 2026 18:55:55 +0200 Subject: [PATCH] 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 | 50 +++++++++++++++++++++ core/linker.py | 39 ++++++++++------- core/sources/civitai.py | 92 ++++++++++++++++++++++++++++++++++++++- core/workflow_analyzer.py | 71 +++++++++++++++++++++++++----- web/linker.js | 18 +++++++- 5 files changed, 241 insertions(+), 29 deletions(-) create mode 100644 TODO.md diff --git a/TODO.md b/TODO.md new file mode 100644 index 0000000..e5b220f --- /dev/null +++ b/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** + diff --git a/core/linker.py b/core/linker.py index 2176f9d..423e2cb 100644 --- a/core/linker.py +++ b/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 diff --git a/core/sources/civitai.py b/core/sources/civitai.py index 5eaf3a4..c34aaf7 100644 --- a/core/sources/civitai.py +++ b/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 diff --git a/core/workflow_analyzer.py b/core/workflow_analyzer.py index ef15ae3..f4fb5c4 100644 --- a/core/workflow_analyzer.py +++ b/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 diff --git a/web/linker.js b/web/linker.js index 94bc263..2da04df 100644 --- a/web/linker.js +++ b/web/linker.js @@ -1029,7 +1029,23 @@ class LinkerManagerDialog extends ComfyDialog { // Card Header: Filename as headline + node chip html += `
`; - html += `

${missingFilename.display}

`; + // Show URN + CivitAI model name + version name if available + let titleHtml = `${missingFilename.display}`; + 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 += ` ${civitaiLabel}`; + } + } else if (missing.civitai_info && missing.civitai_info.model_name) { + titleHtml += ` (${missing.civitai_info.model_name})`; + } + html += `

${titleHtml}

`; html += `
`; if (missing.category) { html += `${missing.category}`;