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}`;