speedyrulz-comfyui-resolution-selector/resolution_selector.py

"""Resolution Selector node for ComfyUI.

Pick a megapixel budget, an aspect ratio, and an edge multiple; the node
computes the closest usable resolution and outputs an empty latent, the
width/height as INTs, and (optionally) the input image resized with
lanczos then center-cropped to that resolution.
"""

import math

import torch

import comfy.utils

try:
    import comfy.model_management

    def _latent_device():
        return comfy.model_management.intermediate_device()
except Exception:  # pragma: no cover - very old ComfyUI
    def _latent_device():
        return torch.device("cpu")


LATENT_TYPES = {
    "SD / SDXL": 4,
    "Flux / SD3": 16,
}


def compute_dimensions(megapixels, ratio, multiple):
    """Find width/height on the `multiple` grid closest to the requested
    megapixel total while staying as close as possible to `ratio` (w/h).

    Mirrored in web/js/resolutionSelector.js - keep both in sync.
    """
    multiple = max(1, int(multiple))
    target = max(megapixels, 0.01) * 1_000_000.0
    ratio = max(ratio, 1e-6)

    ideal_h = math.sqrt(target / ratio)
    ideal_w = ideal_h * ratio

    best = None
    for sw in (math.floor(ideal_w / multiple), math.ceil(ideal_w / multiple)):
        for sh in (math.floor(ideal_h / multiple), math.ceil(ideal_h / multiple)):
            w = max(multiple, sw * multiple)
            h = max(multiple, sh * multiple)
            px_err = abs(w * h - target) / target
            ratio_err = abs(math.log((w / h) / ratio))
            score = px_err + ratio_err
            if best is None or score < best[0]:
                best = (score, w, h)
    return best[1], best[2]


class ResolutionSelector:
    """Megapixels + aspect ratio -> empty latent / ints / resized image."""

    CATEGORY = "utils/resolution"
    FUNCTION = "run"
    RETURN_TYPES = ("LATENT", "INT", "INT", "IMAGE")
    RETURN_NAMES = ("latent", "width", "height", "image")
    DESCRIPTION = (
        "Selects the resolution closest to the requested megapixels and "
        "aspect ratio, snapped to a multiple. Outputs an empty latent, "
        "width/height, and the input image resized (lanczos) and "
        "center-cropped to fit."
    )

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "megapixels": ("FLOAT", {
                    "default": 1.05, "min": 0.01, "max": 64.0,
                    "step": 0.01, "round": 0.01,
                    "tooltip": "Target size in megapixels, e.g. 4.15",
                }),
                "aspect_width": ("INT", {
                    "default": 3, "min": 1, "max": 8192,
                    "tooltip": "Aspect ratio width term (e.g. 16 in 16:9)",
                }),
                "aspect_height": ("INT", {
                    "default": 4, "min": 1, "max": 8192,
                    "tooltip": "Aspect ratio height term (e.g. 9 in 16:9)",
                }),
                "multiple": ("INT", {
                    "default": 32, "min": 1, "max": 1024,
                    "tooltip": "Each edge is snapped to a multiple of this",
                }),
                "use_image": ("BOOLEAN", {
                    "default": False,
                    "tooltip": "Off = ignore the input image entirely",
                }),
                "use_image_ratio": ("BOOLEAN", {
                    "default": False,
                    "tooltip": "Use the input image's aspect ratio instead "
                               "of aspect_width / aspect_height",
                }),
                "latent_type": (list(LATENT_TYPES.keys()), {
                    "default": "Flux / SD3",
                    "tooltip": "Latent channel format for the empty latent",
                }),
            },
            "optional": {
                "image": ("IMAGE",),
            },
        }

    def run(self, megapixels, aspect_width, aspect_height, multiple,
            use_image, use_image_ratio, latent_type, image=None):
        ratio = aspect_width / max(1, aspect_height)
        if use_image and use_image_ratio and image is not None:
            ratio = image.shape[2] / image.shape[1]  # [B, H, W, C]

        width, height = compute_dimensions(megapixels, ratio, multiple)

        channels = LATENT_TYPES.get(latent_type, 4)
        latent = torch.zeros(
            [1, channels, height // 8, width // 8], device=_latent_device())

        out_image = None
        if use_image and image is not None:
            out_image = self._resize_and_crop(image, width, height)

        return ({"samples": latent}, width, height, out_image)

    @staticmethod
    def _resize_and_crop(image, target_w, target_h):
        """Lanczos-resize to cover the target, then center-crop the rest."""
        s = image.movedim(-1, 1)  # [B, H, W, C] -> [B, C, H, W]
        in_h, in_w = s.shape[2], s.shape[3]

        scale = max(target_w / in_w, target_h / in_h)
        new_w = max(target_w, round(in_w * scale))
        new_h = max(target_h, round(in_h * scale))

        s = comfy.utils.common_upscale(s, new_w, new_h, "lanczos", "disabled")

        x = (new_w - target_w) // 2
        y = (new_h - target_h) // 2
        s = s[:, :, y:y + target_h, x:x + target_w]

        return s.movedim(1, -1)


NODE_CLASS_MAPPINGS = {
    "ResolutionSelectorMP": ResolutionSelector,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "ResolutionSelectorMP": "Resolution Selector (MP + Ratio)",
}