310 lines
12 KiB
Python
310 lines
12 KiB
Python
# ================================================
|
|
# File: utils/helpers.py
|
|
# ================================================
|
|
import json
|
|
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",
|
|
"clip": "clip",
|
|
"clip_vision": "clip_vision",
|
|
"clipvision": "clip_vision",
|
|
"text_encoder": "text_encoders",
|
|
"text_encoders": "text_encoders",
|
|
"textencoder": "text_encoders",
|
|
"textencoders": "text_encoders",
|
|
"ipadapter": "ipadapter",
|
|
"ip_adapter": "ipadapter",
|
|
"style_models": "style_models",
|
|
"stylemodel": "style_models",
|
|
"gligen": "gligen",
|
|
"photomaker": "photomaker",
|
|
"inpaint": "inpaint",
|
|
}
|
|
|
|
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 _is_within(path: str, base: str) -> bool:
|
|
"""True if path resolves to base or a location inside it."""
|
|
try:
|
|
return os.path.commonpath((os.path.realpath(base), os.path.realpath(path))) == os.path.realpath(base)
|
|
except (ValueError, OSError):
|
|
return False
|
|
|
|
def _join_within(base: str, subdir: str) -> Optional[str]:
|
|
"""Join a client-supplied subdir onto base, rejecting escapes."""
|
|
if not subdir:
|
|
return base
|
|
joined = os.path.realpath(os.path.join(base, subdir))
|
|
if not _is_within(joined, base):
|
|
print(f"[HuggingFace] Rejected subdir escaping base directory: {subdir}")
|
|
return None
|
|
return joined
|
|
|
|
def _allowed_save_roots() -> List[str]:
|
|
"""Roots a client-provided save_root is allowed to point into."""
|
|
roots = set()
|
|
roots.add(os.path.realpath(os.path.join(folder_paths.base_path, "models")))
|
|
# All ComfyUI-registered model paths (covers extra_model_paths.yaml)
|
|
try:
|
|
for entry in folder_paths.folder_names_and_paths.values():
|
|
for p in entry[0]:
|
|
roots.add(os.path.realpath(p))
|
|
except Exception:
|
|
pass
|
|
# Plugin-registered custom roots
|
|
try:
|
|
roots_file = os.path.join(PLUGIN_ROOT, "custom_roots.json")
|
|
if os.path.exists(roots_file):
|
|
with open(roots_file, "r", encoding="utf-8") as f:
|
|
data = json.load(f)
|
|
if isinstance(data, dict):
|
|
for value in data.values():
|
|
for p in (value or []):
|
|
if isinstance(p, str) and p.strip():
|
|
roots.add(os.path.realpath(p))
|
|
except Exception:
|
|
pass
|
|
return sorted(roots)
|
|
|
|
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, but only if it stays inside a known
|
|
# model root (ComfyUI models dir, registered folder_paths, custom roots)
|
|
if explicit_save_root:
|
|
resolved_root = os.path.realpath(explicit_save_root)
|
|
if not any(_is_within(resolved_root, root) for root in _allowed_save_roots()):
|
|
print(f"[HuggingFace] Rejected save_root outside allowed model roots: {explicit_save_root}")
|
|
return None
|
|
return _join_within(resolved_root, selected_subdir)
|
|
|
|
# First: check if the raw model_type (as passed by the UI dropdown) exists
|
|
# as a physical folder under ComfyUI/models/. The dropdown is populated
|
|
# dynamically from these folders, so they should be used as-is.
|
|
raw_type = model_type.lower().strip() if model_type else ""
|
|
models_dir = os.path.join(folder_paths.base_path, "models")
|
|
if raw_type:
|
|
physical_dir = os.path.join(models_dir, raw_type)
|
|
if _is_within(physical_dir, models_dir) and os.path.isdir(physical_dir):
|
|
return _join_within(physical_dir, selected_subdir)
|
|
|
|
# Try ComfyUI's folder_paths
|
|
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"]:
|
|
expected_folder_name = normalized_type
|
|
physical_dir = os.path.join(models_dir, expected_folder_name)
|
|
|
|
if os.path.isdir(physical_dir):
|
|
return physical_dir
|
|
|
|
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
|
|
print(f"[HuggingFace] folder_paths returned '{candidate}' for "
|
|
f"'{normalized_type}' (expected '{expected_folder_name}'). "
|
|
f"Creating dedicated folder.")
|
|
except Exception:
|
|
pass
|
|
|
|
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
|
|
elif normalized_type != "other":
|
|
# Normalized type is known but not in the special cases above;
|
|
# try folder_paths for it
|
|
try:
|
|
paths = folder_paths.get_folder_paths(normalized_type)
|
|
if paths:
|
|
return paths[0]
|
|
except Exception:
|
|
pass
|
|
# Fall through to create under models/
|
|
fallback_dir = os.path.join(models_dir, normalized_type)
|
|
os.makedirs(fallback_dir, exist_ok=True)
|
|
return fallback_dir
|
|
else:
|
|
# Truly unknown type -- create a folder under models/ using the
|
|
# raw name so the user gets the folder they selected in the UI.
|
|
if raw_type:
|
|
fallback_dir = os.path.join(models_dir, raw_type)
|
|
if _is_within(fallback_dir, models_dir):
|
|
os.makedirs(fallback_dir, exist_ok=True)
|
|
return fallback_dir
|
|
print(f"[HuggingFace] Rejected model_type escaping models dir: {model_type}")
|
|
return None
|
|
return os.path.join(PLUGIN_ROOT, "other_models")
|
|
|
|
except:
|
|
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)
|