From eb702ddcff28dae669c1d662700d85e165d776ec Mon Sep 17 00:00:00 2001
From: EnragedAntelope <176548309+EnragedAntelope@users.noreply.github.com>
Date: Fri, 24 Jan 2025 13:42:40 -0500
Subject: [PATCH] Initial code.

---
 __init__.py |   3 ++
 nodes.py    | 102 ++++++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 105 insertions(+)
 create mode 100644 __init__.py
 create mode 100644 nodes.py

diff --git a/__init__.py b/__init__.py
new file mode 100644
index 0000000..50fb1bb
--- /dev/null
+++ b/__init__.py
@@ -0,0 +1,3 @@
+from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
+
+__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
\ No newline at end of file
diff --git a/nodes.py b/nodes.py
new file mode 100644
index 0000000..1f010b5
--- /dev/null
+++ b/nodes.py
@@ -0,0 +1,102 @@
+import numpy as np
+from typing import Tuple
+
+class ConstrainResolution:
+    """
+    A ComfyUI node that analyzes images and suggests optimal dimensions 
+    while preserving aspect ratio.
+    """
+    
+    @classmethod
+    def INPUT_TYPES(cls):
+        """Define input types with tooltips"""
+        return {
+            "required": {
+                "image": ("IMAGE", {
+                    "tooltip": "Input image to analyze for resolution constraints"
+                }),
+                "min_res": ("INT", {
+                    "default": 704,
+                    "min": 1,
+                    "tooltip": "Minimum resolution in pixels for both width and height"
+                }),
+                "max_res": ("INT", {
+                    "default": 1280,
+                    "min": 1,
+                    "tooltip": "Maximum resolution in pixels for both width and height"
+                }),
+            }
+        }
+
+    RETURN_TYPES = ("IMAGE", "INT", "INT", "FLOAT", "FLOAT")
+    RETURN_NAMES = ("Image", "Constrained Width", "Constrained Height", "Constrained Aspect Ratio", "Original Aspect Ratio")
+    
+    # Add tooltips for outputs
+    RETURN_TYPES_TOOLTIPS = {
+        "Image": "The unchanged input image (passed through)",
+        "Constrained Width": "Suggested width that fits within min/max constraints while preserving aspect ratio",
+        "Constrained Height": "Suggested height that fits within min/max constraints while preserving aspect ratio",
+        "Constrained Aspect Ratio": "The aspect ratio of the constrained dimensions",
+        "Original Aspect Ratio": "The original aspect ratio of the input image"
+    }
+    
+    FUNCTION = "process"
+    CATEGORY = "image/resolution"
+
+    @staticmethod
+    def calculate_aspect_ratio(width: int, height: int) -> float:
+        return round(width / height, 4)
+
+    @staticmethod
+    def calculate_optimal_dimensions(width: int, height: int, min_res: int, max_res: int) -> Tuple[int, int]:
+        aspect_ratio = width / height
+        
+        if aspect_ratio > 1:  # Landscape
+            if width > max_res:
+                width = max_res
+                height = int(width / aspect_ratio)
+            if height < min_res:
+                height = min_res
+                width = int(height * aspect_ratio)
+        else:  # Portrait
+            if height > max_res:
+                height = max_res
+                width = int(height * aspect_ratio)
+            if width < min_res:
+                width = min_res
+                height = int(width / aspect_ratio)
+
+        # Final bounds check
+        width = min(max(width, min_res), max_res)
+        height = min(max(height, min_res), max_res)
+        
+        return width, height
+
+    def process(self, image, min_res, max_res):
+        # Get image dimensions
+        if len(image.shape) == 4:
+            height = image.shape[1]
+            width = image.shape[2]
+        else:
+            height = image.shape[0]
+            width = image.shape[1]
+
+        # Calculate original aspect ratio
+        original_aspect_ratio = self.calculate_aspect_ratio(width, height)
+
+        # Calculate optimal dimensions
+        suggested_width, suggested_height = self.calculate_optimal_dimensions(width, height, min_res, max_res)
+        
+        # Calculate final aspect ratio
+        aspect_ratio = self.calculate_aspect_ratio(suggested_width, suggested_height)
+
+        # Return image first, followed by other outputs
+        return (image, suggested_width, suggested_height, aspect_ratio, original_aspect_ratio)
+
+NODE_CLASS_MAPPINGS = {
+    "ConstrainResolution": ConstrainResolution
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "ConstrainResolution": "Constrain Resolution"
+} 
\ No newline at end of file