nodes/huchukato-ComfyUI-HuggingFace
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
Files
ec9ed4e6d64bb0aeaa422dd86fc65098e0a98fee
huchukato-ComfyUI-HuggingFace/utils/helpers.py
T
huchukato e90f2c104c Fix diffusion_models path resolving to unet alias in folder_paths. 
ComfyUI may register diffusion_models as an alias of unet in
folder_paths, so get_folder_paths("diffusion_models") returns the
unet directory. This caused downloads to be saved in models/unet
even when the user selected diffusion_models.

Now get_model_dir prefers the physical folder
models/<normalized_type> when it exists, and verifies that the path
returned by folder_paths actually matches the requested type name
before using it. If folder_paths returns an alias, a dedicated
folder is created under models/ instead.
2026-08-27 15:00:04 +02:00

233 lines
9.3 KiB
Python
Raw Blame History

# ================================================
# File: utils/helpers.py
# ================================================
import os
import urllib.parse
import re
from pathlib import Path
from typing import Optional, List, Dict, Any
import folder_paths
# Import config values needed here
from ..config import PLUGIN_ROOT, MODEL_TYPE_DIRS
# Canonical aliases for model type/folder names. Values are preferred folder names.
MODEL_TYPE_ALIASES = {
"checkpoint": "checkpoints",
"checkpoints": "checkpoints",
"diffusionmodel": "diffusion_models",
"diffusionmodels": "diffusion_models",
"diffusion_model": "diffusion_models",
"diffusion_models": "diffusion_models", # Wan 2.2 and similar go here
"diffusers": "diffusers",
"unet": "unet", # GGUF models go here
"lora": "loras",
"loras": "loras",
"locon": "loras",
"lycoris": "loras",
"vae": "vae",
"embedding": "embeddings",
"embeddings": "embeddings",
"textualinversion": "embeddings",
"hypernetwork": "hypernetworks",
"hypernetworks": "hypernetworks",
"controlnet": "controlnet",
"upscaler": "upscale_models",
"upscalers": "upscale_models",
"upscale_model": "upscale_models",
"upscale_models": "upscale_models",
"motionmodule": "motion_models",
"motionmodules": "motion_models",
"motion_model": "motion_models",
"motion_models": "motion_models",
}
MODEL_TYPE_ALIASES_COMPACT = {
re.sub(r'[^a-z0-9]', '', k): v for k, v in MODEL_TYPE_ALIASES.items()
}
def _normalize_model_type(model_type: str) -> str:
"""Normalize model type string to canonical form."""
if not model_type:
return "other"
normalized = model_type.lower().strip()
# Try exact match first
if normalized in MODEL_TYPE_ALIASES:
return MODEL_TYPE_ALIASES[normalized]
# Try compact match (remove non-alphanumeric chars)
compact = re.sub(r'[^a-z0-9]', '', normalized)
if compact in MODEL_TYPE_ALIASES_COMPACT:
return MODEL_TYPE_ALIASES_COMPACT[compact]
return "other"
def get_model_dir(model_type: str, explicit_save_root: str = "", selected_subdir: str = "") -> Optional[str]:
"""Get the directory path for a given model type."""
try:
# Normalize the model type
normalized_type = _normalize_model_type(model_type)
# Use explicit root if provided
if explicit_save_root:
base_path = explicit_save_root
# If selected_subdir is provided, append it
if selected_subdir:
base_path = os.path.join(base_path, selected_subdir)
return base_path
# Try ComfyUI's folder_paths first
try:
if normalized_type in ["checkpoints", "loras", "vae", "embeddings", "hypernetworks", "controlnet", "upscale_models"]:
return folder_paths.get_folder_paths(normalized_type)[0]
elif normalized_type in ["diffusion_models", "motion_models", "unet", "diffusers"]:
# ComfyUI may alias diffusion_models -> unet (or vice versa) in
# folder_paths, returning the wrong physical folder. Prefer the
# folder that actually matches the requested type name.
expected_folder_name = normalized_type
models_dir = os.path.join(folder_paths.base_path, "models")
physical_dir = os.path.join(models_dir, expected_folder_name)
# If the physical folder exists, use it directly
if os.path.isdir(physical_dir):
return physical_dir
# Otherwise try folder_paths, but verify the returned path
# actually contains the expected folder name (not an alias)
try:
paths = folder_paths.get_folder_paths(normalized_type)
if paths:
candidate = os.path.abspath(paths[0])
candidate_name = os.path.basename(candidate)
if candidate_name == expected_folder_name:
return candidate
# folder_paths returned an alias (e.g. unet for
# diffusion_models). Create the correct folder instead.
print(f"[HuggingFace] folder_paths returned '{candidate}' for "
f"'{normalized_type}' (expected '{expected_folder_name}'). "
f"Creating dedicated folder.")
except Exception:
pass
# Fallback: create a folder named after the normalized type
# under the ComfyUI models directory.
fallback_dir = physical_dir
try:
os.makedirs(fallback_dir, exist_ok=True)
except Exception as mke:
print(f"[HuggingFace] Warning: could not create '{fallback_dir}': {mke}")
return fallback_dir
else:
# For other types, use our extension directory instead of custom_nodes
from ..config import PLUGIN_ROOT
return os.path.join(PLUGIN_ROOT, "other_models")
except:
# Fallback to base_path + type
return os.path.join(folder_paths.base_path, normalized_type)
except Exception as e:
print(f"Error getting model directory for {model_type}: {e}")
return None
def sanitize_filename(filename: str) -> str:
"""Sanitize filename for safe file system usage."""
if not filename:
return "unnamed_file"
# Remove invalid characters
sanitized = re.sub(r'[<>:"/\\|?*]', '_', filename)
# Remove control characters
sanitized = re.sub(r'[\x00-\x1f\x7f]', '', sanitized)
# Limit length
if len(sanitized) > 255:
name, ext = os.path.splitext(sanitized)
sanitized = name[:255-len(ext)] + ext
return sanitized.strip()
def parse_huggingface_input(url_or_id: str) -> tuple[str | None, str | None]:
"""
Parses HuggingFace URL or ID string.
Returns: (model_id, filename) tuple. Both can be None.
Handles URLs like:
- https://huggingface.co/FX-FeiHou/wan2.2-Remix/resolve/main/NSFW/Wan2.2_Remix_NSFW_i2v_14b_high_lighting_fp8_e4m3fn_v2.1.safetensors
- FX-FeiHou/wan2.2-Remix
"""
if not url_or_id:
return None, None
url_or_id = str(url_or_id).strip()
model_id: str | None = None
filename: str | None = None
# Check if it's a direct download URL
if "/resolve/main/" in url_or_id:
try:
parsed_url = urllib.parse.urlparse(url_or_id)
print(f"[DEBUG] Parsed URL: {parsed_url}")
print(f"[DEBUG] Path parts: {parsed_url.path.split('/')}")
# Extract model ID from path
path_parts = parsed_url.path.split('/')
print(f"[DEBUG] Path parts length: {len(path_parts)}")
print(f"[DEBUG] Path parts[2]: {path_parts[2] if len(path_parts) > 2 else 'None'}")
if len(path_parts) >= 4 and path_parts[3] == "resolve":
# URL format: /FX-FeiHou/wan2.2-Remix/resolve/main/NSFW/file.safetensors
model_id = "/".join(path_parts[1:3]) # FX-FeiHou/wan2.2-Remix
filename = "/".join(path_parts[4:]) # main/NSFW/file.safetensors
# Remove "main/" from filename if it's duplicated
if filename.startswith("main/"):
filename = filename[5:] # Remove "main/" prefix
print(f"[DEBUG] Extracted model_id: {model_id}")
print(f"[DEBUG] Extracted filename: {filename}")
print(f"Parsed HF download URL - Model: {model_id}, File: {filename}")
return model_id, filename
else:
print(f"[DEBUG] Path does not match expected format")
except Exception as e:
print(f"Warning: Could not parse HF download URL '{url_or_id}': {e}")
return None, None
# Check if it's a simple model ID (username/repo format)
if "/" in url_or_id and not url_or_id.startswith("http"):
parts = url_or_id.split("/")
if len(parts) >= 2:
model_id = "/".join(parts[0:2])
print(f"Parsed HF model ID: {model_id}")
return model_id, None
# Try parsing as full URL
try:
parsed_url = urllib.parse.urlparse(url_or_id)
if "huggingface.co" in parsed_url.netloc.lower():
path_parts = parsed_url.path.split('/')
if len(path_parts) >= 3:
model_id = "/".join(path_parts[1:3]) # Extract username/repo
print(f"Parsed HF URL - Model: {model_id}")
return model_id, None
except Exception as e:
print(f"Warning: Could not parse HF URL '{url_or_id}': {e}")
print(f"Input '{url_or_id}' is not a recognizable HF model ID or URL.")
return None, None
def get_model_folder_paths(model_type: str) -> List[str]:
"""Get all folder paths for a given model type."""
try:
normalized_type = _normalize_model_type(model_type)
return folder_paths.get_folder_paths(normalized_type)
except:
return []
def get_model_type_folder_name(model_type: str) -> str:
"""Get the standard folder name for a model type."""
return _normalize_model_type(model_type)
Reference in New Issue View Git Blame Copy Permalink