mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-09-23 12:04:09 -03:00
Martial Michel
e9aff35957
Add Load Image Metadata (LoraManager) to extract reusable prompts, model references, LoRA stacks, and sampling settings from images. Prefer saved A1111-style parameters by default, with optional workflow and subgraph sampler selection. Resolve local model and LoRA names, report missing resources, and recover extraction failures with explicit defaults and readable diagnostics. Include parser, resource-resolution, and node regression tests, plus usage documentation.
747 lines
27 KiB
Python
747 lines
27 KiB
Python
from difflib import SequenceMatcher
|
|
import logging
|
|
import os
|
|
import re
|
|
from typing import Any, Dict, List, Optional
|
|
from ..services.service_registry import ServiceRegistry
|
|
from ..config import config
|
|
from ..services.settings_manager import get_settings_manager
|
|
import asyncio
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def get_lora_info(lora_name):
|
|
"""Get the lora path and trigger words from cache"""
|
|
|
|
async def _get_lora_info_async():
|
|
scanner = await ServiceRegistry.get_lora_scanner()
|
|
cache = await scanner.get_cached_data()
|
|
|
|
lora_name_normalized = lora_name.replace("\\", "/")
|
|
lora_name_no_ext = lora_name_normalized
|
|
for ext in (".safetensors", ".ckpt", ".pt", ".bin"):
|
|
if lora_name_no_ext.lower().endswith(ext):
|
|
lora_name_no_ext = lora_name_no_ext[: -len(ext)]
|
|
break
|
|
|
|
has_path = "/" in lora_name_no_ext
|
|
basename = os.path.basename(lora_name_no_ext) if has_path else lora_name_no_ext
|
|
best_fallback = None
|
|
|
|
for item in cache.raw_data:
|
|
file_name = item.get("file_name", "")
|
|
folder = item.get("folder", "")
|
|
file_name_no_ext = file_name
|
|
for ext in (".safetensors", ".ckpt", ".pt", ".bin"):
|
|
if file_name_no_ext.lower().endswith(ext):
|
|
file_name_no_ext = file_name_no_ext[: -len(ext)]
|
|
break
|
|
path_name = f"{folder}/{file_name_no_ext}".replace("\\", "/") if folder else file_name_no_ext
|
|
|
|
if lora_name_no_ext not in (file_name_no_ext, path_name):
|
|
if has_path and file_name_no_ext == basename:
|
|
if folder and lora_name_no_ext.startswith(folder.replace("\\", "/") + "/"):
|
|
best_fallback = item
|
|
elif best_fallback is None:
|
|
best_fallback = item
|
|
continue
|
|
|
|
file_path = item.get("file_path")
|
|
if not file_path:
|
|
continue
|
|
|
|
all_roots = list(config.loras_roots or []) + list(
|
|
config.extra_loras_roots or []
|
|
)
|
|
for root in all_roots:
|
|
root = root.replace(os.sep, "/")
|
|
if file_path.startswith(root):
|
|
relative_path = os.path.relpath(file_path, root).replace(
|
|
os.sep, "/"
|
|
)
|
|
civitai = item.get("civitai", {})
|
|
trigger_words = (
|
|
civitai.get("trainedWords", []) if civitai else []
|
|
)
|
|
return relative_path, trigger_words
|
|
civitai = item.get("civitai", {})
|
|
trigger_words = civitai.get("trainedWords", []) if civitai else []
|
|
return file_path, trigger_words
|
|
|
|
if best_fallback:
|
|
file_path = best_fallback.get("file_path")
|
|
if file_path:
|
|
civitai = best_fallback.get("civitai", {})
|
|
trigger_words = civitai.get("trainedWords", []) if civitai else []
|
|
return file_path, trigger_words
|
|
|
|
return lora_name, []
|
|
|
|
try:
|
|
# Check if we're already in an event loop
|
|
loop = asyncio.get_running_loop()
|
|
# If we're in a running loop, we need to use a different approach
|
|
# Create a new thread to run the async code
|
|
import concurrent.futures
|
|
|
|
def run_in_thread():
|
|
new_loop = asyncio.new_event_loop()
|
|
asyncio.set_event_loop(new_loop)
|
|
try:
|
|
return new_loop.run_until_complete(_get_lora_info_async())
|
|
finally:
|
|
new_loop.close()
|
|
|
|
with concurrent.futures.ThreadPoolExecutor() as executor:
|
|
future = executor.submit(run_in_thread)
|
|
return future.result()
|
|
|
|
except RuntimeError:
|
|
# No event loop is running, we can use asyncio.run()
|
|
return asyncio.run(_get_lora_info_async())
|
|
|
|
|
|
def get_lora_info_absolute(lora_name):
|
|
"""Get the absolute lora path and trigger words from cache
|
|
|
|
Returns:
|
|
tuple: (absolute_path, trigger_words) where absolute_path is the full
|
|
file system path to the LoRA file, or original lora_name if not found
|
|
"""
|
|
|
|
async def _get_lora_info_absolute_async():
|
|
scanner = await ServiceRegistry.get_lora_scanner()
|
|
cache = await scanner.get_cached_data()
|
|
|
|
# Stack producers can resolve an exact business path. Preserve it even
|
|
# when several indexed LoRAs share the same basename.
|
|
if os.path.isabs(lora_name):
|
|
for item in cache.raw_data:
|
|
file_path = item.get("file_path")
|
|
if file_path and os.path.abspath(file_path) == os.path.abspath(lora_name):
|
|
civitai = item.get("civitai") or {}
|
|
return file_path, civitai.get("trainedWords", [])
|
|
return lora_name, []
|
|
|
|
lora_name_normalized = lora_name.replace("\\", "/")
|
|
lora_name_no_ext = lora_name_normalized
|
|
for ext in (".safetensors", ".ckpt", ".pt", ".bin"):
|
|
if lora_name_no_ext.lower().endswith(ext):
|
|
lora_name_no_ext = lora_name_no_ext[: -len(ext)]
|
|
break
|
|
|
|
has_path = "/" in lora_name_no_ext
|
|
basename = os.path.basename(lora_name_no_ext) if has_path else lora_name_no_ext
|
|
best_fallback = None
|
|
|
|
for item in cache.raw_data:
|
|
file_name = item.get("file_name", "")
|
|
folder = item.get("folder", "")
|
|
file_name_no_ext = file_name
|
|
for ext in (".safetensors", ".ckpt", ".pt", ".bin"):
|
|
if file_name_no_ext.lower().endswith(ext):
|
|
file_name_no_ext = file_name_no_ext[: -len(ext)]
|
|
break
|
|
path_name = f"{folder}/{file_name_no_ext}".replace("\\", "/") if folder else file_name_no_ext
|
|
|
|
if lora_name_no_ext == file_name_no_ext:
|
|
file_path = item.get("file_path")
|
|
if file_path:
|
|
civitai = item.get("civitai", {})
|
|
trigger_words = civitai.get("trainedWords", []) if civitai else []
|
|
return file_path, trigger_words
|
|
|
|
if lora_name_no_ext == path_name:
|
|
file_path = item.get("file_path")
|
|
if file_path:
|
|
civitai = item.get("civitai", {})
|
|
trigger_words = civitai.get("trainedWords", []) if civitai else []
|
|
return file_path, trigger_words
|
|
|
|
if has_path and file_name_no_ext == basename:
|
|
if folder and lora_name_no_ext.startswith(folder.replace("\\", "/") + "/"):
|
|
best_fallback = item
|
|
elif best_fallback is None:
|
|
best_fallback = item
|
|
|
|
if best_fallback:
|
|
file_path = best_fallback.get("file_path")
|
|
if file_path:
|
|
civitai = best_fallback.get("civitai", {})
|
|
trigger_words = civitai.get("trainedWords", []) if civitai else []
|
|
return file_path, trigger_words
|
|
|
|
return lora_name, []
|
|
|
|
try:
|
|
# Check if we're already in an event loop
|
|
loop = asyncio.get_running_loop()
|
|
# If we're in a running loop, we need to use a different approach
|
|
# Create a new thread to run the async code
|
|
import concurrent.futures
|
|
|
|
def run_in_thread():
|
|
new_loop = asyncio.new_event_loop()
|
|
asyncio.set_event_loop(new_loop)
|
|
try:
|
|
return new_loop.run_until_complete(_get_lora_info_absolute_async())
|
|
finally:
|
|
new_loop.close()
|
|
|
|
with concurrent.futures.ThreadPoolExecutor() as executor:
|
|
future = executor.submit(run_in_thread)
|
|
return future.result()
|
|
|
|
except RuntimeError:
|
|
# No event loop is running, we can use asyncio.run()
|
|
return asyncio.run(_get_lora_info_absolute_async())
|
|
|
|
|
|
def get_checkpoint_info_absolute(checkpoint_name):
|
|
"""Get the absolute checkpoint path and metadata from cache
|
|
|
|
Supports ComfyUI-style model names (e.g., "folder/model_name.ext")
|
|
|
|
Args:
|
|
checkpoint_name: The model name, can be:
|
|
- ComfyUI format: "folder/model_name.safetensors"
|
|
- Simple name: "model_name"
|
|
|
|
Returns:
|
|
tuple: (absolute_path, metadata) where absolute_path is the full
|
|
file system path to the checkpoint file, or original checkpoint_name if not found,
|
|
metadata is the full model metadata dict or None
|
|
"""
|
|
|
|
async def _get_checkpoint_info_absolute_async():
|
|
from ..services.service_registry import ServiceRegistry
|
|
|
|
scanner = await ServiceRegistry.get_checkpoint_scanner()
|
|
cache = await scanner.get_cached_data()
|
|
|
|
# Get model roots for matching
|
|
model_roots = scanner.get_model_roots()
|
|
|
|
# Normalize the checkpoint name
|
|
normalized_name = checkpoint_name.replace(os.sep, "/")
|
|
|
|
for item in cache.raw_data:
|
|
file_path = item.get("file_path", "")
|
|
if not file_path:
|
|
continue
|
|
|
|
# Format the stored path as ComfyUI-style name
|
|
formatted_name = _format_model_name_for_comfyui(file_path, model_roots)
|
|
|
|
# Match by formatted name (normalize separators for robust comparison)
|
|
if formatted_name.replace(os.sep, "/") == normalized_name or formatted_name == checkpoint_name:
|
|
return file_path, item
|
|
|
|
# Also try matching by basename only (for backward compatibility)
|
|
file_name = item.get("file_name", "")
|
|
if (
|
|
file_name == checkpoint_name
|
|
or file_name == os.path.splitext(normalized_name)[0]
|
|
):
|
|
return file_path, item
|
|
|
|
return checkpoint_name, None
|
|
|
|
try:
|
|
# Check if we're already in an event loop
|
|
loop = asyncio.get_running_loop()
|
|
# If we're in a running loop, we need to use a different approach
|
|
# Create a new thread to run the async code
|
|
import concurrent.futures
|
|
|
|
def run_in_thread():
|
|
new_loop = asyncio.new_event_loop()
|
|
asyncio.set_event_loop(new_loop)
|
|
try:
|
|
return new_loop.run_until_complete(
|
|
_get_checkpoint_info_absolute_async()
|
|
)
|
|
finally:
|
|
new_loop.close()
|
|
|
|
with concurrent.futures.ThreadPoolExecutor() as executor:
|
|
future = executor.submit(run_in_thread)
|
|
return future.result()
|
|
|
|
except RuntimeError:
|
|
# No event loop is running, we can use asyncio.run()
|
|
return asyncio.run(_get_checkpoint_info_absolute_async())
|
|
|
|
|
|
def _format_model_name_for_comfyui(file_path: str, model_roots: list[str]) -> str:
|
|
"""Format file path to ComfyUI-style model name (relative path with extension)
|
|
|
|
Example: /path/to/checkpoints/Illustrious/model.safetensors -> Illustrious/model.safetensors
|
|
|
|
Args:
|
|
file_path: Absolute path to the model file
|
|
model_roots: List of model root directories
|
|
|
|
Returns:
|
|
ComfyUI-style model name with relative path and extension
|
|
"""
|
|
# Find the matching root and get relative path
|
|
for root in model_roots:
|
|
try:
|
|
# Normalize paths for comparison
|
|
norm_file = os.path.normcase(os.path.abspath(file_path))
|
|
norm_root = os.path.normcase(os.path.abspath(root))
|
|
|
|
# Add trailing separator for prefix check
|
|
if not norm_root.endswith(os.sep):
|
|
norm_root += os.sep
|
|
|
|
if norm_file.startswith(norm_root):
|
|
# Use os.path.relpath to get relative path with OS-native separator
|
|
return os.path.relpath(file_path, root)
|
|
except (ValueError, TypeError):
|
|
continue
|
|
|
|
# If no root matches, just return the basename with extension
|
|
return os.path.basename(file_path)
|
|
|
|
|
|
def model_patcher_to_name(model_patcher: Any) -> Optional[str]:
|
|
"""Extract a ComfyUI-style model name from a MODEL (ModelPatcher) object.
|
|
|
|
Core ComfyUI loaders record the absolute weight file path on the patcher's
|
|
``cached_patcher_init`` attribute:
|
|
- load_checkpoint_guess_config -> (fn, (ckpt_path, ...), index)
|
|
- load_diffusion_model -> (fn, (unet_path, model_options))
|
|
Patcher clones (LoRA loaders, model merges, ...) preserve the attribute,
|
|
so the name is recoverable anywhere downstream of a core loader — including
|
|
from LoRA Manager's own loaders (CheckpointLoaderLM / UNETLoaderLM), which
|
|
call the same core load functions.
|
|
|
|
The absolute path is converted to the ComfyUI-style relative name used by
|
|
the metadata pipeline (covering standard ComfyUI roots and LoRA Manager
|
|
extra folder paths).
|
|
|
|
Returns None when the path cannot be recovered (e.g. third-party loaders
|
|
that never set ``cached_patcher_init``).
|
|
"""
|
|
init = getattr(model_patcher, "cached_patcher_init", None)
|
|
if not isinstance(init, (tuple, list)) or len(init) < 2:
|
|
return None
|
|
args = init[1]
|
|
abs_path = args[0] if args else None
|
|
if not isinstance(abs_path, str) or not abs_path:
|
|
return None
|
|
return _abs_model_path_to_name(abs_path)
|
|
|
|
|
|
def sampler_object_to_name(sampler: Any) -> Optional[str]:
|
|
"""Extract a ComfyUI-style sampler name from a SAMPLER (KSAMPLER) object.
|
|
|
|
Standard outputs (KSamplerSelect, most built-in sampler nodes) round-trip
|
|
losslessly via the underlying sampler function's ``__name__``
|
|
(``sample_euler`` -> ``euler``). A few edge cases need special-casing
|
|
because the function name diverges from the ``SAMPLER_NAMES`` entry:
|
|
|
|
- ``dpm_fast`` / ``dpm_adaptive`` are local closures inside
|
|
``comfy.samplers.ksampler`` (``dpm_fast_function`` / ``dpm_adaptive_function``)
|
|
- ``uni_pc`` / ``uni_pc_bh2`` use ``sample_unipc`` / ``sample_unipc_bh2``
|
|
|
|
``ddim`` is constructed by ComfyUI as ``euler`` with random inpaint, so
|
|
the original ``ddim`` name is unrecoverable (extracts as ``euler``).
|
|
Custom sampler nodes that pass non-``sample_*`` functions return None.
|
|
|
|
Returns None when the name cannot be recovered.
|
|
"""
|
|
sampler_function = getattr(sampler, "sampler_function", None)
|
|
func_name = getattr(sampler_function, "__name__", None)
|
|
if not isinstance(func_name, str) or not func_name:
|
|
return None
|
|
if func_name == "dpm_fast_function":
|
|
return "dpm_fast"
|
|
if func_name == "dpm_adaptive_function":
|
|
return "dpm_adaptive"
|
|
if func_name.startswith("sample_"):
|
|
name = func_name[len("sample_"):]
|
|
if name == "unipc":
|
|
return "uni_pc"
|
|
if name == "unipc_bh2":
|
|
return "uni_pc_bh2"
|
|
return name or None
|
|
return None
|
|
|
|
|
|
def _abs_model_path_to_name(abs_path: str) -> str:
|
|
"""Convert an absolute model path to a ComfyUI-style relative name.
|
|
|
|
Tries standard ComfyUI model roots plus LoRA Manager extra folder paths;
|
|
falls back to the bare filename.
|
|
"""
|
|
try:
|
|
roots: List[str] = list(config.base_models_roots or [])
|
|
roots.extend(config.extra_checkpoints_roots or [])
|
|
roots.extend(config.extra_unet_roots or [])
|
|
formatted = _format_model_name_for_comfyui(abs_path, roots)
|
|
if formatted:
|
|
return formatted
|
|
except Exception:
|
|
pass
|
|
return os.path.basename(abs_path)
|
|
|
|
|
|
def fuzzy_match(text: str, pattern: str, threshold: float = 0.85) -> bool:
|
|
"""
|
|
Check if text matches pattern using fuzzy matching.
|
|
Returns True if similarity ratio is above threshold.
|
|
"""
|
|
if not pattern or not text:
|
|
return False
|
|
|
|
# Convert both to lowercase for case-insensitive matching
|
|
text = text.lower()
|
|
pattern = pattern.lower()
|
|
|
|
# Split pattern into words
|
|
search_words = pattern.split()
|
|
|
|
# Check each word
|
|
for word in search_words:
|
|
# First check if word is a substring (faster)
|
|
if word in text:
|
|
continue
|
|
|
|
# If not found as substring, try fuzzy matching
|
|
# Check if any part of the text matches this word
|
|
found_match = False
|
|
for text_part in text.split():
|
|
ratio = SequenceMatcher(None, text_part, word).ratio()
|
|
if ratio >= threshold:
|
|
found_match = True
|
|
break
|
|
|
|
if not found_match:
|
|
return False
|
|
|
|
# All words found either as substrings or fuzzy matches
|
|
return True
|
|
|
|
|
|
def sanitize_folder_name(name: str, replacement: str = "_") -> str:
|
|
"""Sanitize a folder name by removing or replacing invalid characters.
|
|
|
|
Args:
|
|
name: The original folder name.
|
|
replacement: The character to use when replacing invalid characters.
|
|
|
|
Returns:
|
|
A sanitized folder name safe to use across common filesystems.
|
|
"""
|
|
|
|
if not name:
|
|
return ""
|
|
|
|
# Replace invalid characters commonly restricted on Windows and POSIX
|
|
invalid_chars_pattern = r'[<>:"/\\|?*\x00-\x1f]'
|
|
sanitized = re.sub(invalid_chars_pattern, replacement, name)
|
|
|
|
# Trim whitespace introduced during sanitization
|
|
sanitized = sanitized.strip()
|
|
|
|
# Collapse repeated replacement characters to a single instance
|
|
if replacement:
|
|
sanitized = re.sub(f"{re.escape(replacement)}+", replacement, sanitized)
|
|
# Combine stripping to be idempotent:
|
|
# Right side: strip replacement, space, and dot (Windows restriction)
|
|
# Left side: strip replacement and space (leading dots are allowed)
|
|
sanitized = sanitized.rstrip(" ." + replacement).lstrip(" " + replacement)
|
|
else:
|
|
# If no replacement, just strip spaces and dots from right, spaces from left
|
|
sanitized = sanitized.rstrip(" .").lstrip(" ")
|
|
|
|
if not sanitized:
|
|
return "unnamed"
|
|
|
|
return sanitized
|
|
|
|
|
|
def calculate_recipe_fingerprint(loras):
|
|
"""
|
|
Calculate a unique fingerprint for a recipe based on its LoRAs.
|
|
|
|
The fingerprint is created by sorting LoRA hashes, filtering invalid entries,
|
|
normalizing strength values to 2 decimal places, and joining in format:
|
|
hash1:strength1|hash2:strength2|...
|
|
|
|
Args:
|
|
loras (list): List of LoRA dictionaries with hash and strength values
|
|
|
|
Returns:
|
|
str: The calculated fingerprint
|
|
"""
|
|
if not loras:
|
|
return ""
|
|
|
|
valid_loras = []
|
|
for lora in loras:
|
|
if lora.get("exclude", False):
|
|
continue
|
|
|
|
hash_value = lora.get("hash", "")
|
|
if isinstance(hash_value, str):
|
|
hash_value = hash_value.lower()
|
|
else:
|
|
hash_value = str(hash_value).lower() if hash_value else ""
|
|
if not hash_value and lora.get("modelVersionId"):
|
|
hash_value = str(lora.get("modelVersionId"))
|
|
|
|
if not hash_value:
|
|
continue
|
|
|
|
# Normalize strength to 2 decimal places (check both strength and weight fields)
|
|
strength_val = lora.get("strength", lora.get("weight", 1.0))
|
|
try:
|
|
strength = round(float(strength_val), 2)
|
|
except (ValueError, TypeError):
|
|
strength = 1.0
|
|
|
|
valid_loras.append((hash_value, strength))
|
|
|
|
# Sort by hash
|
|
valid_loras.sort()
|
|
|
|
# Join in format hash1:strength1|hash2:strength2|...
|
|
fingerprint = "|".join(
|
|
[f"{hash_value}:{strength}" for hash_value, strength in valid_loras]
|
|
)
|
|
|
|
return fingerprint
|
|
|
|
|
|
def normalize_prompt_for_dedup(prompt) -> str:
|
|
"""Normalize a positive prompt for duplicate recipe matching.
|
|
|
|
Applies casefolding, collapses whitespace runs into single spaces, and
|
|
trims leading/trailing whitespace. Missing or non-string prompts
|
|
normalize to an empty string.
|
|
|
|
Args:
|
|
prompt: The positive prompt text (or None)
|
|
|
|
Returns:
|
|
str: The normalized prompt
|
|
"""
|
|
if not prompt or not isinstance(prompt, str):
|
|
return ""
|
|
return re.sub(r"\s+", " ", prompt).strip().casefold()
|
|
|
|
|
|
def calculate_relative_path_for_model(
|
|
model_data: Dict[str, Any], model_type: str = "lora"
|
|
) -> str:
|
|
"""Calculate relative path for existing model using template from settings
|
|
|
|
Args:
|
|
model_data: Model data from scanner cache
|
|
model_type: Type of model ('lora', 'checkpoint', 'embedding')
|
|
|
|
Returns:
|
|
Relative path string (empty string for flat structure)
|
|
"""
|
|
# Get path template from settings for specific model type
|
|
settings_manager = get_settings_manager()
|
|
path_template = settings_manager.get_download_path_template(model_type)
|
|
|
|
# If template is empty, return empty path (flat structure)
|
|
if not path_template:
|
|
return ""
|
|
|
|
# Get base model name from model metadata
|
|
civitai_data = model_data.get("civitai", {})
|
|
|
|
# For CivitAI models, prefer civitai data only if 'id' exists; for non-CivitAI models, use model_data directly
|
|
if civitai_data and civitai_data.get("id") is not None:
|
|
base_model = model_data.get("base_model", "")
|
|
# Get author from civitai creator data
|
|
creator_info = civitai_data.get("creator") or {}
|
|
author = creator_info.get("username") or "Anonymous"
|
|
else:
|
|
# Fallback to model_data fields for non-CivitAI models
|
|
base_model = model_data.get("base_model", "")
|
|
author = "Anonymous" # Default for non-CivitAI models
|
|
|
|
model_tags = model_data.get("tags", [])
|
|
|
|
# Apply mapping if available
|
|
base_model_mappings = settings_manager.get("base_model_path_mappings", {})
|
|
mapped_base_model = base_model_mappings.get(base_model, base_model)
|
|
|
|
# Convert all tags to lowercase to avoid case sensitivity issues on Windows
|
|
lowercase_tags = [tag.lower() for tag in model_tags if isinstance(tag, str)]
|
|
first_tag = settings_manager.resolve_priority_tag_for_model(
|
|
lowercase_tags, model_type
|
|
)
|
|
|
|
if not first_tag:
|
|
first_tag = "no tags" # Default if no tags available
|
|
|
|
# Format the template with available data
|
|
model_name = sanitize_folder_name(model_data.get("model_name", ""))
|
|
version_name = ""
|
|
|
|
if isinstance(civitai_data, dict):
|
|
version_name = sanitize_folder_name(civitai_data.get("name") or "")
|
|
|
|
formatted_path = path_template
|
|
formatted_path = formatted_path.replace("{base_model}", mapped_base_model)
|
|
formatted_path = formatted_path.replace("{first_tag}", first_tag)
|
|
formatted_path = formatted_path.replace("{author}", author)
|
|
formatted_path = formatted_path.replace("{model_name}", model_name)
|
|
formatted_path = formatted_path.replace("{version_name}", version_name)
|
|
|
|
if model_type == "embedding":
|
|
formatted_path = formatted_path.replace(" ", "_")
|
|
|
|
# Sanitize the resolved path to prevent path traversal
|
|
formatted_path = formatted_path.lstrip("/")
|
|
while "//" in formatted_path:
|
|
formatted_path = formatted_path.replace("//", "/")
|
|
formatted_path = formatted_path.rstrip("/")
|
|
|
|
return formatted_path
|
|
|
|
|
|
def calculate_filename_for_model(
|
|
model_data: Dict[str, Any], model_type: str = "lora"
|
|
) -> str:
|
|
"""Calculate the filename stem for a model using the filename template.
|
|
|
|
Mirrors the data extraction of :func:`calculate_relative_path_for_model`
|
|
but renders a single filename (no path segments). Missing values resolve
|
|
to empty segments instead of the path-oriented defaults ("Anonymous" /
|
|
"no tags") so templates degrade gracefully.
|
|
|
|
Args:
|
|
model_data: Model data from scanner cache
|
|
model_type: Type of model ('lora', 'checkpoint', 'embedding')
|
|
|
|
Returns:
|
|
Sanitized filename stem without extension, or an empty string when no
|
|
template is configured, the template is invalid, or the rendered name
|
|
is empty.
|
|
"""
|
|
settings_manager = get_settings_manager()
|
|
template = settings_manager.get_download_filename_template(model_type)
|
|
|
|
if not template:
|
|
return ""
|
|
|
|
# A filename template must render a single name, never folder segments.
|
|
if "/" in template or "\\" in template:
|
|
logger.warning(
|
|
"Filename template for %s contains a path separator and is ignored: %r",
|
|
model_type,
|
|
template,
|
|
)
|
|
return ""
|
|
|
|
civitai_data = model_data.get("civitai", {})
|
|
|
|
author = ""
|
|
if isinstance(civitai_data, dict) and civitai_data.get("id") is not None:
|
|
creator_info = civitai_data.get("creator") or {}
|
|
author = creator_info.get("username") or ""
|
|
|
|
base_model = model_data.get("base_model", "")
|
|
base_model_mappings = settings_manager.get("base_model_path_mappings", {})
|
|
mapped_base_model = base_model_mappings.get(base_model, base_model)
|
|
|
|
lowercase_tags = [
|
|
tag.lower() for tag in model_data.get("tags", []) if isinstance(tag, str)
|
|
]
|
|
first_tag = settings_manager.resolve_priority_tag_for_model(
|
|
lowercase_tags, model_type
|
|
)
|
|
|
|
model_name = model_data.get("model_name", "")
|
|
version_name = ""
|
|
if isinstance(civitai_data, dict):
|
|
version_name = civitai_data.get("name") or ""
|
|
|
|
sha256 = model_data.get("sha256") or ""
|
|
hash_short = sha256[:10].lower() if isinstance(sha256, str) else ""
|
|
|
|
file_path = model_data.get("file_path") or ""
|
|
if isinstance(file_path, str) and file_path:
|
|
original_name = os.path.splitext(os.path.basename(file_path))[0]
|
|
else:
|
|
original_name = os.path.splitext(str(model_data.get("file_name", "")))[0]
|
|
|
|
def _sanitize_value(value: Any) -> str:
|
|
# sanitize_folder_name falls back to "unnamed" for empty input; for
|
|
# templates an empty value must stay empty so segments collapse.
|
|
text = str(value) if value else ""
|
|
return sanitize_folder_name(text) if text else ""
|
|
|
|
replacements = {
|
|
"{model_name}": _sanitize_value(model_name),
|
|
"{version_name}": _sanitize_value(version_name),
|
|
"{base_model}": _sanitize_value(mapped_base_model),
|
|
"{author}": _sanitize_value(author),
|
|
"{first_tag}": _sanitize_value(first_tag),
|
|
"{hash_short}": hash_short,
|
|
"{original_name}": _sanitize_value(original_name),
|
|
}
|
|
|
|
result = template
|
|
for placeholder, value in replacements.items():
|
|
result = result.replace(placeholder, value)
|
|
|
|
if model_type == "embedding":
|
|
result = result.replace(" ", "_")
|
|
|
|
# Strip characters that are illegal in filenames on common filesystems.
|
|
result = re.sub(r'[:*?"<>|]', "", result)
|
|
# Collapse runs of identical separators introduced by empty substitutions.
|
|
result = re.sub(r"([-_. ])\1+", r"\1", result)
|
|
# Drop separators left dangling next to each other ("- -" -> "-").
|
|
result = re.sub(r" ?([-_.]) (?=[-_.])", r"\1", result)
|
|
# A stem must not start or end with separators, spaces or dots.
|
|
result = result.strip("-_. ")
|
|
|
|
return result
|
|
|
|
|
|
def remove_empty_dirs(path):
|
|
"""Recursively remove empty directories starting from the given path.
|
|
|
|
Args:
|
|
path (str): Root directory to start cleaning from
|
|
|
|
Returns:
|
|
int: Number of empty directories removed
|
|
"""
|
|
removed_count = 0
|
|
|
|
if not os.path.isdir(path):
|
|
return removed_count
|
|
|
|
# List all files in directory
|
|
files = os.listdir(path)
|
|
|
|
# Process all subdirectories first
|
|
for file in files:
|
|
full_path = os.path.join(path, file)
|
|
if os.path.isdir(full_path):
|
|
removed_count += remove_empty_dirs(full_path)
|
|
|
|
# Check if directory is now empty (after processing subdirectories)
|
|
if not os.listdir(path):
|
|
try:
|
|
os.rmdir(path)
|
|
removed_count += 1
|
|
except OSError:
|
|
pass
|
|
|
|
return removed_count
|