From e5787df940a0a42a032740bc97ec0fdf6706574f Mon Sep 17 00:00:00 2001
From: "Bug, Ltd." <194549623+bugltd@users.noreply.github.com>
Date: Mon, 12 May 2025 16:19:19 +0200
Subject: [PATCH] add node ImageDownscaleToTotalPixels (#20)

---
 node_list.json            |  3 +-
 src/python/nodes/utils.py | 58 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 60 insertions(+), 1 deletion(-)

diff --git a/node_list.json b/node_list.json
index 48916f5..6b29e22 100644
--- a/node_list.json
+++ b/node_list.json
@@ -32,5 +32,6 @@
   "Save Text File (lab)": "Save text content to a file",
   "Sleep (lab)": "Delay execution for the given amount of seconds",
   "Convert to Any (lab)": "Convert any value to type Any ('*)",
-  "Resolution to Dimensions (lab)": "Split a resolution (as string, ex. '1024x1024') to dimensions, optionally multiplied by scale factor"
+  "Resolution to Dimensions (lab)": "Split a resolution (as string, ex. '1024x1024') to dimensions, optionally multiplied by scale factor",
+  "Image: Downscale to Total Pixels (lab)": "Downscale an image to match target size in megapuxels"
 }
\ No newline at end of file
diff --git a/src/python/nodes/utils.py b/src/python/nodes/utils.py
index 191de3a..88d7c40 100644
--- a/src/python/nodes/utils.py
+++ b/src/python/nodes/utils.py
@@ -5,8 +5,10 @@ from PIL import Image
 import numpy as np
 import torch  # type: ignore
 import re
+import math
 
 import folder_paths  # type: ignore
+import comfy.utils
 
 from ..collection.register_nodes import register_node
 from ..shared.utils import ANY_TYPE, pillow_to_tensor
@@ -339,3 +341,59 @@ class ResolutionToDims:
             int(int(m.group(1)) * scale_factor),
             int(int(m.group(2)) * scale_factor),
         )
+
+
+@register_node('Image: Downscale to Total Pixels', 'utils')
+class ImageDownscaleToTotalPixels:
+    upscale_methods = ['nearest-exact', 'bilinear', 'area', 'bicubic', 'lanczos']
+
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            'required': {
+                'image': ('IMAGE',),
+                'upscale_method': (
+                    s.upscale_methods,
+                    {'default': 'lanczos', 'tooltip': 'upscale method'},
+                ),
+                'megapixels': (
+                    'FLOAT',
+                    {
+                        'default': '1',
+                        'min': 0,
+                        'max': 16,
+                        'tooltip': 'target size in megapixels',
+                    },
+                ),
+            }
+        }
+
+    FUNCTION = 'run'
+    RETURN_TYPES = ('IMAGE', 'INT', 'INT')
+    RETURN_NAMES = ('image', 'width', 'height')
+    OUTPUT_TOOLTIPS = ('downscaled image', 'width', 'height')
+    DESCRIPTION = 'Downscale an image to match target size in megapuxels.\nIf the image is smaller than the target size, it is kept as is.'
+
+    def run(self, image, upscale_method, megapixels):
+        orig_samples = image.movedim(-1, 1)
+        orig_width = orig_samples.shape[3]
+        orig_height = orig_samples.shape[2]
+        orig_total = orig_width * orig_height
+        target_total = int(megapixels * 1024 * 1024)
+
+        if orig_total <= target_total:
+            return (image, orig_width, orig_height)
+
+        scale_by = math.sqrt(target_total / orig_total)
+        target_width = round(orig_width * scale_by)
+        # we calculate height differently, to respect aspect ratio as much as possible
+        target_height = round(orig_height * target_width / orig_width)
+
+        s = comfy.utils.common_upscale(
+            orig_samples, target_width, target_height, upscale_method, 'disabled'
+        )
+        s = s.movedim(1, -1)
+        return (s, target_width, target_height)