glonlas-ComfyUI-image-profile/nodes.py

"""Node definitions for ComfyUI-Image-profile."""

from __future__ import annotations

import json
import math
from typing import Any

import torch
import comfy.model_management

DIM_MULTIPLE = 8
DIM_MIN = 8
DIM_MAX = 16384
STEPS_MIN = 1
STEPS_MAX = 150

DEFAULT_PROFILES: list[dict[str, Any]] = [
    {
        "id": "default-landscape-low",
        "name": "Landscape Low resolution",
        "width": 404,
        "height": 204,
        "steps": 5,
    },
    {
        "id": "default-portrait-low",
        "name": "Portrait Low resolution",
        "width": 204,
        "height": 404,
        "steps": 5,
    },
    {
        "id": "default-landscape-high",
        "name": "Landscape High resolution",
        "width": 1152,
        "height": 864,
        "steps": 8,
    },
    {
        "id": "default-portrait-high",
        "name": "Portrait High resolution",
        "width": 864,
        "height": 1152,
        "steps": 8,
    },
]


class ComfyUIImageProfile:
    """Emit a latent and step count from a selected image profile."""

    def __init__(self) -> None:
        self.device = comfy.model_management.intermediate_device()

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "profiles_json": ("STRING", {"default": "[]", "multiline": False}),
                "selected_profile_id": ("STRING", {"default": ""}),
                "selected_width": ("INT", {"default": 1152, "min": DIM_MIN, "max": DIM_MAX, "step": DIM_MULTIPLE}),
                "selected_height": ("INT", {"default": 864, "min": DIM_MIN, "max": DIM_MAX, "step": DIM_MULTIPLE}),
                "selected_steps": ("INT", {"default": 8, "min": STEPS_MIN, "max": STEPS_MAX, "step": 1}),
            },
            "optional": {
                # When connected, the image's aspect ratio overrides the
                # selected profile's ratio; the profile still sets the
                # resolution (pixel budget) and step count.
                "image": ("IMAGE",),
            },
        }

    RETURN_TYPES = ("LATENT", "INT")
    RETURN_NAMES = ("Latent", "Steps")
    FUNCTION = "generate"
    CATEGORY = "Latent/Profile"

    @staticmethod
    def _clamp(value: int, minimum: int, maximum: int) -> int:
        return max(minimum, min(maximum, value))

    @classmethod
    def _sanitize_dimension(cls, value: Any) -> int:
        try:
            integer = int(value)
        except (TypeError, ValueError):
            integer = DIM_MIN

        integer = cls._clamp(integer, DIM_MIN, DIM_MAX)
        lower = (integer // DIM_MULTIPLE) * DIM_MULTIPLE
        upper = lower + DIM_MULTIPLE

        if lower < DIM_MIN:
            lower = DIM_MIN
        if upper > DIM_MAX:
            upper = DIM_MAX

        if integer - lower >= upper - integer:
            return upper
        return lower

    @classmethod
    def _sanitize_steps(cls, value: Any) -> int:
        try:
            integer = int(value)
        except (TypeError, ValueError):
            integer = STEPS_MIN
        return cls._clamp(integer, STEPS_MIN, STEPS_MAX)

    @classmethod
    def _sanitize_profile(cls, profile: dict[str, Any]) -> dict[str, Any] | None:
        profile_id = str(profile.get("id", "")).strip()
        name = str(profile.get("name", "")).strip()
        if not profile_id or not name:
            return None

        return {
            "id": profile_id,
            "name": name,
            "width": cls._sanitize_dimension(profile.get("width")),
            "height": cls._sanitize_dimension(profile.get("height")),
            "steps": cls._sanitize_steps(profile.get("steps")),
        }

    @classmethod
    def _load_profiles(cls, profiles_json: str) -> list[dict[str, Any]]:
        try:
            parsed = json.loads(profiles_json)
        except (TypeError, ValueError, json.JSONDecodeError):
            parsed = []

        if not isinstance(parsed, list):
            parsed = []

        valid_profiles: list[dict[str, Any]] = []
        for candidate in parsed:
            if not isinstance(candidate, dict):
                continue
            profile = cls._sanitize_profile(candidate)
            if profile is not None:
                valid_profiles.append(profile)

        if valid_profiles:
            return valid_profiles

        return [p for profile in DEFAULT_PROFILES if (p := cls._sanitize_profile(profile)) is not None]

    @staticmethod
    def _image_dimensions(image: Any) -> tuple[int, int] | None:
        """Return (width, height) from a ComfyUI IMAGE tensor, or None.

        ComfyUI images are tensors shaped [batch, height, width, channels]
        (a bare [height, width, channels] is also tolerated).
        """
        shape = getattr(image, "shape", None)
        if shape is None:
            return None
        try:
            dims = tuple(int(value) for value in shape)
        except (TypeError, ValueError):
            return None

        if len(dims) >= 4:
            height, width = dims[1], dims[2]
        elif len(dims) == 3:
            height, width = dims[0], dims[1]
        else:
            return None

        if width <= 0 or height <= 0:
            return None
        return width, height

    @classmethod
    def _fit_ratio_to_area(
        cls,
        image: Any,
        target_width: int,
        target_height: int,
    ) -> tuple[int, int]:
        """Scale the image's aspect ratio to the target pixel budget.

        Keeps roughly the same number of pixels as ``target_width x
        target_height`` (the selected profile's resolution) while adopting the
        image's ratio. Falls back to the target dimensions when the image has
        no usable shape.
        """
        dimensions = cls._image_dimensions(image)
        if dimensions is None:
            return target_width, target_height

        image_width, image_height = dimensions
        area = max(1, target_width * target_height)
        ratio = image_width / image_height
        raw_width = math.sqrt(area * ratio)
        raw_height = math.sqrt(area / ratio)
        return cls._sanitize_dimension(raw_width), cls._sanitize_dimension(raw_height)

    def _resolve_selected_profile(
        self,
        profiles: list[dict[str, Any]],
        selected_profile_id: str,
        selected_width: int,
        selected_height: int,
        selected_steps: int,
    ) -> dict[str, Any]:
        selected = next((p for p in profiles if p["id"] == selected_profile_id), None)
        if selected is not None:
            return selected

        return {
            "id": profiles[0]["id"],
            "name": "Selected",
            "width": self._sanitize_dimension(selected_width),
            "height": self._sanitize_dimension(selected_height),
            "steps": self._sanitize_steps(selected_steps),
        }

    def generate(
        self,
        profiles_json: str,
        selected_profile_id: str,
        selected_width: int,
        selected_height: int,
        selected_steps: int,
        image: Any = None,
    ):
        profiles = self._load_profiles(profiles_json)
        selected = self._resolve_selected_profile(
            profiles,
            selected_profile_id=selected_profile_id,
            selected_width=selected_width,
            selected_height=selected_height,
            selected_steps=selected_steps,
        )

        width = selected["width"]
        height = selected["height"]
        steps = selected["steps"]

        if image is not None:
            # Adopt the connected image's aspect ratio while keeping the
            # profile's resolution (pixel budget) and step count.
            width, height = self._fit_ratio_to_area(image, width, height)

        latent = torch.zeros([1, 4, height // 8, width // 8], device=self.device)
        return ({"samples": latent}, steps)


NODE_CLASS_MAPPINGS = {
    "ComfyUIImageProfile": ComfyUIImageProfile,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "ComfyUIImageProfile": "Image Profile",
}