Apache0ne-ComfyUI-DiffiT/utils.py

"""Utility helpers for the ComfyUI DiffiT node pack."""

from __future__ import annotations

import json
import logging
import math
import re
from difflib import SequenceMatcher
from functools import lru_cache
from pathlib import Path
from typing import Any

import torch

LOGGER = logging.getLogger("ComfyUI-DiffiT")

OFFICIAL_DIFFIT_CHECKPOINTS = (
    "diffit_256.safetensors",
    "diffit_512.safetensors",
)

OFFICIAL_CFG_BY_SIZE = {
    256: 4.4,
    512: 1.49,
}


def get_repo_root() -> Path:
    return Path(__file__).resolve().parent


def get_comfy_root() -> Path:
    return get_repo_root().parent.parent


def get_folder_paths_module():
    try:
        import folder_paths  # type: ignore

        return folder_paths
    except Exception:
        return None


def dedupe_paths(paths: list[Path]) -> list[Path]:
    seen: set[str] = set()
    unique: list[Path] = []
    for path in paths:
        normalized = str(path.resolve()) if path.exists() else str(path)
        if normalized not in seen:
            seen.add(normalized)
            unique.append(path)
    return unique


def find_stability_matrix_root() -> Path | None:
    comfy_root = get_comfy_root()
    for parent in comfy_root.parents:
        if parent.name == "Packages" and (parent.parent / "Models").exists():
            return parent.parent
    return None


def get_stability_matrix_diffusion_dir() -> Path | None:
    root = find_stability_matrix_root()
    if root is None:
        return None
    return root / "Models" / "DiffusionModels"


def get_comfy_diffusion_dirs() -> list[Path]:
    comfy_root = get_comfy_root()
    dirs = [
        comfy_root / "models" / "diffusion_models",
        comfy_root / "models" / "unet",
    ]
    folder_paths = get_folder_paths_module()
    if folder_paths is not None:
        try:
            dirs.extend(Path(path) for path in folder_paths.get_folder_paths("diffusion_models"))
        except Exception:
            pass
    return dedupe_paths(dirs)


def get_diffit_search_dirs() -> list[Path]:
    dirs: list[Path] = []
    stability_matrix_dir = get_stability_matrix_diffusion_dir()
    if stability_matrix_dir is not None:
        dirs.append(stability_matrix_dir)
    dirs.extend(get_comfy_diffusion_dirs())

    expanded: list[Path] = []
    for path in dedupe_paths(dirs):
        expanded.append(path)
        expanded.append(path / "DiffiT")
    return dedupe_paths(expanded)


def get_preferred_diffit_download_dir() -> Path:
    stability_matrix_dir = get_stability_matrix_diffusion_dir()
    if stability_matrix_dir is not None:
        stability_matrix_dir.mkdir(parents=True, exist_ok=True)
        return stability_matrix_dir

    comfy_dir = get_comfy_root() / "models" / "diffusion_models"
    comfy_dir.mkdir(parents=True, exist_ok=True)
    return comfy_dir


def get_hf_cache_dir(cache_name: str) -> Path:
    stability_matrix_root = find_stability_matrix_root()
    if stability_matrix_root is not None:
        cache_dir = stability_matrix_root / "Models" / ".cache" / "huggingface" / cache_name
    else:
        cache_dir = get_comfy_root() / "models" / ".cache" / "huggingface" / cache_name
    cache_dir.mkdir(parents=True, exist_ok=True)
    return cache_dir


def validate_image_size(image_size: int) -> int:
    if image_size not in OFFICIAL_CFG_BY_SIZE:
        raise ValueError(
            f"Unsupported image_size={image_size}. DiffiT is only validated for 256 and 512."
        )
    return image_size


def latent_size_for_image(image_size: int) -> int:
    return validate_image_size(image_size) // 8


def default_cfg_scale_for_image_size(image_size: int) -> float:
    return OFFICIAL_CFG_BY_SIZE[validate_image_size(image_size)]


def list_available_diffit_checkpoints() -> list[str]:
    discovered = set(OFFICIAL_DIFFIT_CHECKPOINTS)
    for directory in get_diffit_search_dirs():
        if not directory.exists():
            continue
        for pattern in ("diffit_*.safetensors", "DiffiT/diffit_*.safetensors"):
            for path in directory.glob(pattern):
                if path.is_file():
                    discovered.add(path.name)
    return sorted(discovered)


def resolve_diffit_checkpoint_path(checkpoint_name: str, manual_path: str = "") -> Path | None:
    if manual_path.strip():
        candidate = Path(manual_path.strip())
        if candidate.exists():
            return candidate

    candidate = Path(checkpoint_name)
    if candidate.exists():
        return candidate

    folder_paths = get_folder_paths_module()
    if folder_paths is not None:
        try:
            resolved = folder_paths.get_full_path("diffusion_models", checkpoint_name)
            if resolved:
                return Path(resolved)
        except Exception:
            pass

    for directory in get_diffit_search_dirs():
        direct = directory / checkpoint_name
        if direct.is_file():
            return direct

    return None


def infer_image_size_from_state_dict(state_dict: dict[str, Any]) -> int:
    pos_embed = state_dict.get("pos_embed")
    if isinstance(pos_embed, torch.Tensor) and pos_embed.ndim == 3:
        num_patches = int(pos_embed.shape[1])
        grid = int(round(math.sqrt(num_patches)))
        if grid * grid == num_patches:
            patch_size = 2
            latent_size = grid * patch_size
            image_size = latent_size * 8
            if image_size in OFFICIAL_CFG_BY_SIZE:
                return image_size

    bias_table = state_dict.get("blocks.0.attn.relative_position_bias_table")
    if isinstance(bias_table, torch.Tensor) and bias_table.ndim == 2:
        entries = int(bias_table.shape[0])
        side = int(round(math.sqrt(entries)))
        if side * side == entries:
            window_size = (side + 1) // 2
            latent_size = window_size * 2
            image_size = latent_size * 8
            if image_size in OFFICIAL_CFG_BY_SIZE:
                return image_size

    raise ValueError(
        "Unable to infer DiffiT image size from checkpoint. Use image_size_override=256 or 512."
    )


def normalize_state_dict(state_dict: dict[str, Any]) -> dict[str, Any]:
    if "state_dict" in state_dict and isinstance(state_dict["state_dict"], dict):
        state_dict = state_dict["state_dict"]
    if "model" in state_dict and isinstance(state_dict["model"], dict):
        state_dict = state_dict["model"]

    prefixes = ("module.", "_orig_mod.")
    normalized: dict[str, Any] = {}
    for key, value in state_dict.items():
        for prefix in prefixes:
            if key.startswith(prefix):
                key = key[len(prefix) :]
                break
        normalized[key] = value
    return normalized


def load_state_dict_file(path: Path) -> dict[str, Any]:
    folder_paths = get_folder_paths_module()
    if folder_paths is not None:
        try:
            import comfy.utils  # type: ignore

            loaded = comfy.utils.load_torch_file(str(path), safe_load=True)
            if isinstance(loaded, dict):
                return normalize_state_dict(loaded)
        except Exception:
            pass

    if path.suffix.lower() == ".safetensors":
        from safetensors.torch import load_file

        return normalize_state_dict(load_file(str(path), device="cpu"))

    loaded = torch.load(str(path), map_location="cpu")
    if not isinstance(loaded, dict):
        raise TypeError(f"Unsupported checkpoint format at {path}.")
    return normalize_state_dict(loaded)


def pick_device(preference: str = "auto") -> torch.device:
    pref = preference.lower()
    if pref == "cuda":
        if not torch.cuda.is_available():
            raise RuntimeError("CUDA was requested for DiffiT but no CUDA device is available.")
        return torch.device("cuda")
    if pref == "cpu":
        return torch.device("cpu")
    if torch.cuda.is_available():
        return torch.device("cuda")
    LOGGER.warning("CUDA is unavailable. Falling back to CPU for DiffiT inference.")
    return torch.device("cpu")


def make_generator(device: torch.device, seed: int) -> torch.Generator | None:
    normalized_seed = int(seed) & 0xFFFFFFFFFFFFFFFF
    try:
        generator = torch.Generator(device=device)
        generator.manual_seed(normalized_seed)
        return generator
    except Exception:
        try:
            generator = torch.Generator(device=device.type)
            generator.manual_seed(normalized_seed)
            return generator
        except Exception:
            LOGGER.warning(
                "Falling back to global RNG seeding for DiffiT sampler. Local generator creation failed on %s.",
                device,
            )
            torch.manual_seed(normalized_seed)
            if torch.cuda.is_available():
                torch.cuda.manual_seed_all(normalized_seed)
            return None


def to_comfy_image(decoded: torch.Tensor) -> torch.Tensor:
    image = ((decoded + 1.0) / 2.0).clamp(0.0, 1.0)
    image = image.permute(0, 2, 3, 1).contiguous()
    return image.detach().cpu().float()


@lru_cache(maxsize=1)
def load_imagenet_labels() -> list[str]:
    asset_path = get_repo_root() / "assets" / "optional_imagenet_labels.json"
    if asset_path.is_file():
        labels = json.loads(asset_path.read_text(encoding="utf-8"))
        if isinstance(labels, list) and len(labels) == 1000:
            return [str(item) for item in labels]

    try:
        from torchvision.models._meta import _IMAGENET_CATEGORIES

        if len(_IMAGENET_CATEGORIES) == 1000:
            return list(_IMAGENET_CATEGORIES)
    except Exception:
        pass

    return [f"class_{index}" for index in range(1000)]


def get_imagenet_label_options() -> list[str]:
    labels = load_imagenet_labels()
    return ["<use class_id>"] + [f"{index:03d}: {label}" for index, label in enumerate(labels)]


def normalize_search_text(text: str) -> str:
    normalized = re.sub(r"[^a-z0-9]+", " ", text.lower()).strip()
    return " ".join(normalized.split())


@lru_cache(maxsize=1)
def get_normalized_imagenet_labels() -> list[tuple[int, str, str]]:
    return [
        (index, label, normalize_search_text(label))
        for index, label in enumerate(load_imagenet_labels())
    ]


def build_prompt_phrases(prompt: str) -> list[str]:
    normalized = normalize_search_text(prompt)
    if not normalized:
        return []

    phrases = {normalized}
    tokens = normalized.split()
    max_window = min(4, len(tokens))
    for window in range(1, max_window + 1):
        for start in range(0, len(tokens) - window + 1):
            phrases.add(" ".join(tokens[start : start + window]))
    return sorted(phrases, key=len, reverse=True)


def score_prompt_phrase(prompt_phrase: str, label_text: str) -> float:
    if not prompt_phrase or not label_text:
        return 0.0
    if prompt_phrase == label_text:
        return 1.0

    prompt_tokens = set(prompt_phrase.split())
    label_tokens = set(label_text.split())
    overlap = len(prompt_tokens & label_tokens)
    union = len(prompt_tokens | label_tokens)
    jaccard = overlap / union if union else 0.0
    subset = 1.0 if prompt_tokens and prompt_tokens.issubset(label_tokens) else 0.0
    contains = 1.0 if prompt_phrase in label_text or label_text in prompt_phrase else 0.0
    ratio = SequenceMatcher(None, prompt_phrase, label_text).ratio()

    return max(
        ratio,
        0.75 * jaccard + 0.25 * subset,
        0.65 * contains + 0.35 * jaccard,
    )


def match_prompt_to_imagenet_label(prompt: str, top_k: int = 5) -> list[tuple[float, int, str]]:
    phrases = build_prompt_phrases(prompt)
    if not phrases:
        return []

    scored: list[tuple[float, int, str]] = []
    for index, label, normalized_label in get_normalized_imagenet_labels():
        best_score = max(score_prompt_phrase(phrase, normalized_label) for phrase in phrases)
        scored.append((best_score, index, label))

    scored.sort(key=lambda item: (item[0], -item[1]), reverse=True)
    return scored[: max(1, int(top_k))]