diff --git a/core/linker.py b/core/linker.py index fbe1f45..0ebc9a3 100644 --- a/core/linker.py +++ b/core/linker.py @@ -23,7 +23,7 @@ logger = logging.getLogger(__name__) URL_PATTERN = re.compile(r'(https?://(?:huggingface\.co|civitai\.com)[^\s"\'<>\)\\]+)') # Model file extensions to look for -MODEL_EXTENSIONS = (".safetensors", ".ckpt", ".pt", ".pth", ".bin", ".onnx") +MODEL_EXTENSIONS = (".safetensors", ".ckpt", ".pt", ".pth", ".bin", ".onnx", ".gguf") def extract_workflow_urls(workflow_json: Dict[str, Any]) -> Dict[str, Dict[str, Any]]: @@ -85,7 +85,7 @@ def extract_workflow_urls(workflow_json: Dict[str, Any]) -> Dict[str, Dict[str, # 3. Extract model filenames via regex model_pattern = re.compile( - r"([\w\-\.%]+\.(?:safetensors|ckpt|pt|pth|bin|onnx))", re.IGNORECASE + r"([\w\-\.%]+\.(?:safetensors|ckpt|pt|pth|bin|onnx|gguf))", re.IGNORECASE ) model_files_raw = model_pattern.findall(workflow_str) @@ -334,7 +334,8 @@ def apply_resolution( 'widget_index': widget index, 'resolved_path': absolute path to resolved model, 'category': model category (optional), - 'resolved_model': model dict from scanner (optional) + 'resolved_model': model dict from scanner (optional), + 'nested_key': nested key for dict-type widgets (optional) } Returns: @@ -359,6 +360,9 @@ def apply_resolution( "original_lora_name": resolution.get( "original_lora_name" ), # Original lora name for LoraManager replacement + "nested_key": resolution.get( + "nested_key" + ), # For dict-type widgets (e.g. Power Lora Loader) } # If resolved_model provided, extract path if needed diff --git a/core/scanner.py b/core/scanner.py index 4dd3ef3..1347e1c 100644 --- a/core/scanner.py +++ b/core/scanner.py @@ -17,41 +17,55 @@ except ImportError: # Model file extensions to look for # This matches folder_paths.supported_pt_extensions -MODEL_EXTENSIONS = {'.ckpt', '.pt', '.pt2', '.bin', '.pth', '.safetensors', '.pkl', '.sft', '.onnx'} +MODEL_EXTENSIONS = { + ".ckpt", + ".pt", + ".pt2", + ".bin", + ".pth", + ".safetensors", + ".pkl", + ".sft", + ".onnx", + ".gguf", +} def get_model_directories() -> Dict[str, Tuple[List[str], set]]: """ Get all configured model directories from folder_paths. - + Returns: Dictionary mapping category name to a tuple. ComfyUI may provide either: - (paths, extensions), or - (paths, extensions, recursive_flag) """ global folder_paths - + if folder_paths is None: # Try to import again try: import folder_paths as fp + folder_paths = fp except ImportError: logging.error("Model Linker: folder_paths still not available") return {} - + return folder_paths.folder_names_and_paths.copy() -def scan_directory(directory: str, extensions: set, category: str) -> List[Dict[str, str]]: +def scan_directory( + directory: str, extensions: set, category: str +) -> List[Dict[str, str]]: """ Recursively scan a single directory for model files. - + Args: directory: Absolute path to directory to scan extensions: Set of file extensions to look for category: Model category name (e.g., 'checkpoints', 'loras') - + Returns: List of dictionaries with model information: { @@ -63,33 +77,37 @@ def scan_directory(directory: str, extensions: set, category: str) -> List[Dict[ } """ models = [] - + if not os.path.exists(directory) or not os.path.isdir(directory): logging.debug(f"Directory does not exist or is not accessible: {directory}") return models - + try: # Get absolute path and normalize base_directory = os.path.abspath(directory) - + # Walk through directory recursively for root, dirs, files in os.walk(base_directory, followlinks=True): # Skip hidden directories - dirs[:] = [d for d in dirs if not d.startswith('.')] - + dirs[:] = [d for d in dirs if not d.startswith(".")] + for filename in files: # Check if file has a model extension file_ext = os.path.splitext(filename)[1].lower() - + # For categories with empty extension set, accept all files # Otherwise, check if extension matches # Accept if: # - no explicit extensions configured, or # - matches configured extensions, or # - matches our known model extensions - if len(extensions or set()) == 0 or file_ext in extensions or file_ext in MODEL_EXTENSIONS: + if ( + len(extensions or set()) == 0 + or file_ext in extensions + or file_ext in MODEL_EXTENSIONS + ): full_path = os.path.join(root, filename) - + # Calculate relative path from base directory # IMPORTANT: Use OS-native path separators (backslashes on Windows) # This matches ComfyUI's recursive_search format for get_filename_list @@ -101,33 +119,35 @@ def scan_directory(directory: str, extensions: set, category: str) -> List[Dict[ except ValueError: # If paths are on different drives (Windows), use filename only relative_path = filename - - models.append({ - 'filename': filename, - 'path': full_path, - 'relative_path': relative_path, - 'category': category, - 'base_directory': base_directory - }) + + models.append( + { + "filename": filename, + "path": full_path, + "relative_path": relative_path, + "category": category, + "base_directory": base_directory, + } + ) except (OSError, PermissionError) as e: logging.warning(f"Error scanning directory {directory}: {e}") - + return models def scan_all_directories() -> List[Dict[str, str]]: """ Scan all configured model directories and return list of available models. - + Returns: List of dictionaries with model information (same format as scan_directory) """ all_models = [] directories = get_model_directories() - + for category, value in directories.items(): # Skip categories that aren't typically model directories - if category in ['custom_nodes', 'configs']: + if category in ["custom_nodes", "configs"]: continue # Unpack folder_paths value flexibly: (paths, extensions) or (paths, extensions, recursive) @@ -143,11 +163,13 @@ def scan_all_directories() -> List[Dict[str, str]]: paths = list(value) raw_exts = [] elif isinstance(value, dict): - paths = value.get('paths') or value.get('path') or [] - raw_exts = value.get('extensions') or [] + paths = value.get("paths") or value.get("path") or [] + raw_exts = value.get("extensions") or [] else: # Unknown format; skip category - logging.debug(f"Unexpected folder_paths format for category {category}: {type(value)}") + logging.debug( + f"Unexpected folder_paths format for category {category}: {type(value)}" + ) continue # Normalize extensions to a set[str] @@ -156,26 +178,32 @@ def scan_all_directories() -> List[Dict[str, str]]: elif raw_exts: extensions = {str(raw_exts).lower()} except Exception as e: - logging.warning(f"Error interpreting folder_paths entry for {category}: {e}") + logging.warning( + f"Error interpreting folder_paths entry for {category}: {e}" + ) continue for directory_path in paths: try: models = scan_directory(directory_path, extensions, category) all_models.extend(models) - logging.debug(f"Found {len(models)} models in {category}/{directory_path}") + logging.debug( + f"Found {len(models)} models in {category}/{directory_path}" + ) except Exception as e: - logging.warning(f"Error scanning {category} directory {directory_path}: {e}") - + logging.warning( + f"Error scanning {category} directory {directory_path}: {e}" + ) + return all_models def get_model_files() -> List[Dict[str, str]]: """ Get list of all available model files with metadata. - + This is the main entry point for getting model files. - + Returns: List of model dictionaries (same format as scan_directory) """ diff --git a/core/sources/huggingface.py b/core/sources/huggingface.py index fdcacf6..dbee763 100644 --- a/core/sources/huggingface.py +++ b/core/sources/huggingface.py @@ -22,36 +22,33 @@ _search_cache: Dict[str, Any] = {} def parse_huggingface_url(url: str) -> Optional[Dict[str, str]]: """ Parse a HuggingFace URL to extract repo and filename. - + Handles formats: - https://huggingface.co/user/repo/resolve/main/file.safetensors - https://huggingface.co/user/repo/blob/main/file.safetensors - hf://user/repo/file.safetensors - + Returns: Dictionary with 'repo' and 'filename' keys, or None if not HF URL """ - if url.startswith('hf://'): - parts = url[5:].split('/', 2) + if url.startswith("hf://"): + parts = url[5:].split("/", 2) if len(parts) >= 3: - return { - 'repo': f"{parts[0]}/{parts[1]}", - 'filename': parts[2] - } + return {"repo": f"{parts[0]}/{parts[1]}", "filename": parts[2]} return None - + parsed = urlparse(url) - if 'huggingface.co' not in parsed.netloc: + if "huggingface.co" not in parsed.netloc: return None - - match = re.match(r'^/([^/]+/[^/]+)/(resolve|blob)/([^/]+)/(.+)$', parsed.path) + + match = re.match(r"^/([^/]+/[^/]+)/(resolve|blob)/([^/]+)/(.+)$", parsed.path) if match: return { - 'repo': match.group(1), - 'branch': match.group(3), - 'filename': match.group(4) + "repo": match.group(1), + "branch": match.group(3), + "filename": match.group(4), } - + return None @@ -67,111 +64,122 @@ def clean_filename_for_search(filename: str) -> str: """ base = os.path.splitext(filename)[0] # Remove common precision/format suffixes - base = re.sub(r'[-_]?(fp16|fp32|fp8|bf16|e4m3fn|scaled|pruned|emaonly|q4|q8).*$', '', base, flags=re.IGNORECASE) + base = re.sub( + r"[-_]?(fp16|fp32|fp8|bf16|e4m3fn|scaled|pruned|emaonly|q4|q8).*$", + "", + base, + flags=re.IGNORECASE, + ) return base def search_huggingface_for_file( - filename: str, - token: Optional[str] = None, - exact_only: bool = False + filename: str, token: Optional[str] = None, exact_only: bool = False ) -> Optional[Dict[str, Any]]: """ Search HuggingFace for a specific model file. Returns the first repo that contains a matching file. - + Args: filename: Filename to search for token: Optional HF token exact_only: If True, only return exact filename matches (for downloads). If False, also try partial matching (for local file resolution). - + Returns: Dict with url, repo, filename if found, None otherwise """ global _search_cache - + cache_key = f"hf_{filename}_exact{exact_only}" if cache_key in _search_cache: return _search_cache[cache_key] - + try: # Use filename without extension for search filename_base = os.path.splitext(filename)[0].lower() - + headers = {} if token: - headers['Authorization'] = f'Bearer {token}' - + headers["Authorization"] = f"Bearer {token}" + # Search for repos containing this filename search_url = f"{HF_API_URL}/models?search={quote(filename_base)}&limit=10" - + response = requests.get(search_url, headers=headers, timeout=10) if response.status_code != 200: logger.debug(f"HuggingFace search returned {response.status_code}") return None - + repos = response.json() - + for repo in repos: - repo_id = repo.get('id', '') + repo_id = repo.get("id", "") if not repo_id: continue - + # Check if this repo actually has a matching file files_url = f"{HF_API_URL}/models/{repo_id}/tree/main" - + try: files_response = requests.get(files_url, headers=headers, timeout=10) if files_response.status_code == 200: files = files_response.json() - + for file_info in files: - file_path = file_info.get('path', '') - file_base = os.path.splitext(os.path.basename(file_path))[0].lower() - + file_path = file_info.get("path", "") + file_base = os.path.splitext(os.path.basename(file_path))[ + 0 + ].lower() + # Check for exact match first (always try this) if file_path.endswith(filename): result = { - 'source': 'huggingface', - 'repo_id': repo_id, - 'filename': os.path.basename(file_path), - 'path': file_path, - 'url': get_huggingface_download_url(repo_id, file_path), - 'size': file_info.get('size'), - 'match_type': 'exact' + "source": "huggingface", + "repo_id": repo_id, + "filename": os.path.basename(file_path), + "path": file_path, + "url": get_huggingface_download_url(repo_id, file_path), + "size": file_info.get("size"), + "match_type": "exact", } _search_cache[cache_key] = result logger.info(f"Found {filename} on HuggingFace: {repo_id}") return result - + # Check for partial match (filename_base in file_base or vice versa) # Skip partial matches if exact_only is True - prevents confusing # users with wrong model suggestions for downloads if not exact_only: if filename_base in file_base or file_base in filename_base: - if file_path.endswith('.safetensors') or file_path.endswith('.ckpt'): + if file_path.endswith( + ".safetensors" + ) or file_path.endswith(".ckpt"): result = { - 'source': 'huggingface', - 'repo_id': repo_id, - 'filename': os.path.basename(file_path), - 'path': file_path, - 'url': get_huggingface_download_url(repo_id, file_path), - 'size': file_info.get('size'), - 'match_type': 'partial' + "source": "huggingface", + "repo_id": repo_id, + "filename": os.path.basename(file_path), + "path": file_path, + "url": get_huggingface_download_url( + repo_id, file_path + ), + "size": file_info.get("size"), + "match_type": "partial", } _search_cache[cache_key] = result - logger.info(f"Found similar file for {filename} on HuggingFace: {repo_id}/{file_path}") + logger.info( + f"Found similar file for {filename} on HuggingFace: {repo_id}/{file_path}" + ) return result - + except Exception as e: logger.debug(f"Error checking repo {repo_id}: {e}") continue - + # Not found _search_cache[cache_key] = None return None - + except Exception as e: logger.error(f"HuggingFace search error: {e}") return None @@ -181,85 +189,79 @@ def search_huggingface( query: str, model_type: Optional[str] = None, limit: int = 10, - token: Optional[str] = None + token: Optional[str] = None, ) -> List[Dict[str, Any]]: """ Search HuggingFace Hub for models (general search). Returns repos that might be relevant, not guaranteed to have exact file. """ results = [] - + try: headers = {} if token: - headers['Authorization'] = f'Bearer {token}' - - params = { - 'search': query, - 'limit': limit, - 'full': 'true' - } - + headers["Authorization"] = f"Bearer {token}" + + params = {"search": query, "limit": limit, "full": "true"} + response = requests.get( - f"{HF_API_URL}/models", - params=params, - headers=headers, - timeout=15 + f"{HF_API_URL}/models", params=params, headers=headers, timeout=15 ) - + if response.status_code == 200: models = response.json() - + for model in models: - repo_id = model.get('id', '') - - results.append({ - 'source': 'huggingface', - 'repo': repo_id, - 'name': model.get('modelId', repo_id), - 'downloads': model.get('downloads', 0), - 'likes': model.get('likes', 0), - 'url': f"https://huggingface.co/{repo_id}" - }) - + repo_id = model.get("id", "") + + results.append( + { + "source": "huggingface", + "repo": repo_id, + "name": model.get("modelId", repo_id), + "downloads": model.get("downloads", 0), + "likes": model.get("likes", 0), + "url": f"https://huggingface.co/{repo_id}", + } + ) + except Exception as e: logger.error(f"HuggingFace search error: {e}") - + return results -def get_repo_files( - repo: str, - token: Optional[str] = None -) -> List[Dict[str, Any]]: +def get_repo_files(repo: str, token: Optional[str] = None) -> List[Dict[str, Any]]: """Get list of model files in a HuggingFace repo.""" files = [] - + try: headers = {} if token: - headers['Authorization'] = f'Bearer {token}' - + headers["Authorization"] = f"Bearer {token}" + response = requests.get( - f"{HF_API_URL}/models/{repo}/tree/main", - headers=headers, - timeout=15 + f"{HF_API_URL}/models/{repo}/tree/main", headers=headers, timeout=15 ) - + if response.status_code == 200: items = response.json() - + for item in items: - path = item.get('path', '') - if path.endswith(('.safetensors', '.ckpt', '.pt', '.bin', '.pth', '.onnx')): - files.append({ - 'filename': os.path.basename(path), - 'path': path, - 'url': get_huggingface_download_url(repo, path), - 'size': item.get('size') - }) - + path = item.get("path", "") + if path.endswith( + (".safetensors", ".ckpt", ".pt", ".bin", ".pth", ".onnx", ".gguf") + ): + files.append( + { + "filename": os.path.basename(path), + "path": path, + "url": get_huggingface_download_url(repo, path), + "size": item.get("size"), + } + ) + except Exception as e: logger.error(f"Error getting repo files: {e}") - + return files diff --git a/core/workflow_analyzer.py b/core/workflow_analyzer.py index 5e17cd4..dbcf380 100644 --- a/core/workflow_analyzer.py +++ b/core/workflow_analyzer.py @@ -27,6 +27,7 @@ MODEL_EXTENSIONS = { ".pkl", ".sft", ".onnx", + ".gguf", } import re @@ -59,6 +60,7 @@ NODE_TYPE_TO_CATEGORY_HINTS = { "LoraLoaderModelOnly": "loras", "LoraLoaderV2": "loras", "Lora Loader (LoraManager)": "loras", # LoraManager custom node + "Power Lora Loader (rgthree)": "loras", # rgthree's Power Lora Loader "UNETLoader": "diffusion_models", "ControlNetLoader": "controlnet", "ControlNetLoaderAdvanced": "controlnet", @@ -66,6 +68,22 @@ NODE_TYPE_TO_CATEGORY_HINTS = { "UpscaleModelLoader": "upscale_models", "HypernetworkLoader": "hypernetworks", "EmbeddingLoader": "embeddings", + # LTX-Video nodes + "LTXVAudioVAELoader": "checkpoints", + "LowVRAMAudioVAELoader": "checkpoints", + "LTXVGemmaCLIPModelLoader": "text_encoders", +} + +# Keys within dict-type widget values that contain model file references. +# Some nodes (e.g. rgthree Power Lora Loader) store model info as objects like +# {"on": true, "lora": "name.safetensors", "strength": 1.0} inside widgets_values. +# Maps nested key name -> category hint. +NESTED_MODEL_KEYS = { + "lora": "loras", + "ckpt_name": "checkpoints", + "checkpoint": "checkpoints", + "vae_name": "vae", + "control_net_name": "controlnet", } @@ -277,6 +295,46 @@ def get_node_model_info(node: Dict[str, Any]) -> List[Dict[str, Any]]: # 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): + # Check for dict-type widget values containing model references (e.g. Power Lora Loader) + # Some nodes store model info as objects like {"on": true, "lora": "name.safetensors", "strength": 1.0} + if isinstance(value, dict): + for nested_key, nested_category_hint in NESTED_MODEL_KEYS.items(): + nested_value = value.get(nested_key) + if ( + not nested_value + or not isinstance(nested_value, str) + or not is_model_filename(nested_value) + ): + continue + + value_str = nested_value.strip() + nested_categories = ( + [nested_category_hint] if nested_category_hint else None + ) + + resolved = try_resolve_model_path(value_str, nested_categories) + if resolved: + category, full_path = resolved + exists = os.path.exists(full_path) + else: + category = nested_category_hint or "unknown" + full_path = None + exists = False + + model_refs.append( + { + "node_id": node_id, + "node_type": node_type, + "widget_index": idx, + "original_path": value_str, + "category": category, + "full_path": full_path, + "exists": exists, + "is_urn": False, + "connected": is_active, + "nested_key": nested_key, # Track nested key for updates + } + ) continue value_str = str(value).strip() diff --git a/core/workflow_updater.py b/core/workflow_updater.py index 4b1e13f..d8c662a 100644 --- a/core/workflow_updater.py +++ b/core/workflow_updater.py @@ -312,9 +312,18 @@ def update_model_path( relative_path = resolved_path # Update the widget value - widgets_values[widget_index] = relative_path - - logging.debug(f"Updated node {node_id}, widget {widget_index} to: {relative_path}") + # Handle nested dict values (e.g. Power Lora Loader with {"on": true, "lora": "name.safetensors", "strength": 1.0}) + nested_key = mapping.get("nested_key") if mapping else None + if nested_key and isinstance(widgets_values[widget_index], dict): + widgets_values[widget_index][nested_key] = relative_path + logging.debug( + f"Updated node {node_id}, widget {widget_index}[{nested_key}] to: {relative_path}" + ) + else: + widgets_values[widget_index] = relative_path + logging.debug( + f"Updated node {node_id}, widget {widget_index} to: {relative_path}" + ) return True diff --git a/web/linker.js b/web/linker.js index 07d15ec..7c056f7 100644 --- a/web/linker.js +++ b/web/linker.js @@ -1069,8 +1069,8 @@ class LinkerManagerDialog extends ComfyDialog { const parts = models.map(model => { const fullName = model.name || model.original_path?.split(/[\/\\]/).pop() || 'Unknown'; let name = fullName; - if (fullName.match(/\.(safetensors|ckpt|pt|pth|bin|pkl|sft|onnx)$/i)) { - name = fullName.replace(/\.(safetensors|ckpt|pt|pth|bin|pkl|sft|onnx)$/i, ''); + if (fullName.match(/\.(safetensors|ckpt|pt|pth|bin|pkl|sft|onnx|gguf)$/i)) { + name = fullName.replace(/\.(safetensors|ckpt|pt|pth|bin|pkl|sft|onnx|gguf)$/i, ''); } const strength = model.strength !== null && model.strength !== undefined ? model.strength.toFixed(2) @@ -1123,7 +1123,7 @@ class LinkerManagerDialog extends ComfyDialog { if (hasActive) { const activeStr = modelsObj.active.map(m => { const fullName = m.name || m.original_path?.split(/[\/\\]/).pop() || 'Unknown'; - let name = fullName.replace(/\.(safetensors|ckpt|pt|pth|bin|pkl|sft|onnx)$/i, ''); + let name = fullName.replace(/\.(safetensors|ckpt|pt|pth|bin|pkl|sft|onnx|gguf)$/i, ''); const strength = m.strength !== null && m.strength !== undefined ? m.strength.toFixed(2) : '1.00'; return `<${displayName}:${name}:${strength}>`; }).join(' '); @@ -1137,7 +1137,7 @@ class LinkerManagerDialog extends ComfyDialog { for (const model of modelsObj.active) { const fullName = model.name || model.original_path?.split(/[\/\\]/).pop() || 'Unknown'; - const name = fullName.replace(/\.(safetensors|ckpt|pt|pth|bin|pkl|sft|onnx)$/i, ''); + const name = fullName.replace(/\.(safetensors|ckpt|pt|pth|bin|pkl|sft|onnx|gguf)$/i, ''); const strength = model.strength !== null && model.strength !== undefined ? model.strength.toFixed(2) : null; html += `${name}${strength !== null ? `${strength}` : ''}`; } @@ -1147,7 +1147,7 @@ class LinkerManagerDialog extends ComfyDialog { if (hasInactive) { const inactiveStr = modelsObj.inactive.map(m => { const fullName = m.name || m.original_path?.split(/[\/\\]/).pop() || 'Unknown'; - let name = fullName.replace(/\.(safetensors|ckpt|pt|pth|bin|pkl|sft|onnx)$/i, ''); + let name = fullName.replace(/\.(safetensors|ckpt|pt|pth|bin|pkl|sft|onnx|gguf)$/i, ''); const strength = m.strength !== null && m.strength !== undefined ? m.strength.toFixed(2) : '1.00'; return `<${displayName}:${name}:${strength}>`; }).join(' '); @@ -1161,7 +1161,7 @@ class LinkerManagerDialog extends ComfyDialog { for (const model of modelsObj.inactive) { const fullName = model.name || model.original_path?.split(/[\/\\]/).pop() || 'Unknown'; - const name = fullName.replace(/\.(safetensors|ckpt|pt|pth|bin|pkl|sft|onnx)$/i, ''); + const name = fullName.replace(/\.(safetensors|ckpt|pt|pth|bin|pkl|sft|onnx|gguf)$/i, ''); const strength = model.strength !== null && model.strength !== undefined ? model.strength.toFixed(2) : null; html += `${name}${strength !== null ? `${strength}` : ''}`; }