ComfyUI-Lora-Manager/py/utils/utils.py

from difflib import SequenceMatcher
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


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()

        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 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