Initial commit

.gitignore

@@ -0,0 +1 @@
+__pycache__

__init__.py

@@ -0,0 +1,37 @@
+import os
+import subprocess
+import importlib.util
+import sys
+
+import __main__
+
+python = sys.executable
+
+def is_installed(package, package_overwrite=None):
+    try:
+        spec = importlib.util.find_spec(package)
+    except ModuleNotFoundError:
+        pass
+
+    package = package_overwrite or package
+
+    if spec is None:
+        print(f"Installing {package}...")
+        command = f'"{python}" -m pip install {package}'
+  
+        result = subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True, env=os.environ)
+
+        if result.returncode != 0:
+            print(f"Couldn't install\nCommand: {command}\nError code: {result.returncode}")
+
+from .image_fitting_node import ImageSquareFittingNode
+
+NODE_CLASS_MAPPINGS = {
+    "ImageSquareAdapterNode": ImageSquareAdapterNode
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "ImageSquareAdapterNode": "Image Square Adapter Node"
+}
+
+print('\033[34mNimbus Nodes: \033[92mLoaded\033[0m')

image_fitting_node.py

@@ -0,0 +1,99 @@
+from PIL import Image
+import torch
+import numpy as np
+
+from .utils import pil2tensor, tensor2pil
+
+
+class ImageSquareFittingNode:
+    """
+    A custom node for ComfyUI to fit an image into a square frame,
+    resizing and padding it as necessary, with options for resampling, supersampling,
+    and various fitting modes.
+    """
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "image": ("IMAGE",),
+                "target_size": ("INT", {"default": 224, "min": 1, "max": 10000, "step": 1}),
+                "fill_color": ("STRING", {"default": "255,255,255"}),
+                "resampling": (["lanczos", "nearest", "bilinear", "bicubic"], {"default": "lanczos"}),
+                "supersample": (["true", "false"], {"default": "false"}),
+                "fitting_mode": (["none", "top", "bottom", "center"], {"default": "none"}),
+            }
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "image_fit_in_square"
+
+    CATEGORY = "Nimbus-Pack/Image"
+
+    def image_fit_in_square(self, image, target_size=224, fill_color='255,255,255', resampling='lanczos',
+                            supersample='false', fitting_mode='none'):
+        scaled_images = []
+        for img in image:
+            scaled_images.append(
+                self.apply_fit_image(tensor2pil(img), target_size, fill_color, resampling, supersample, fitting_mode))
+
+        scaled_images = torch.cat(scaled_images, dim=0)
+
+        return (scaled_images,)
+
+    def apply_fit_image(self, image: Image.Image, target_size: int, fill_color: str, resample: str, supersample: str,
+                        fitting_mode: str):
+        # Convert fill_color string to tuple
+        fill_color = tuple(map(int, fill_color.split(',')))
+
+        # Define a dictionary of resampling filters
+        resample_filters = {
+            'nearest': Image.NEAREST,
+            'bilinear': Image.BILINEAR,
+            'bicubic': Image.BICUBIC,
+            'lanczos': Image.LANCZOS
+        }
+
+        # Calculate scaling factor and new size
+        scaling_factor = target_size / float(max(image.size))
+        new_size = tuple([int(x * scaling_factor) for x in image.size])
+
+        # Apply supersample if needed
+        if supersample == 'true':
+            image = image.resize((new_size[0] * 8, new_size[1] * 8), resample=resample_filters[resample])
+
+        # Resize the image
+        image = image.resize(new_size, resample=resample_filters[resample])
+
+        # Adjust image fitting based on the mode
+        if fitting_mode == 'none':
+            # Current behavior - centering the image
+            new_img = Image.new("RGB", (target_size, target_size), fill_color)
+            position = ((target_size - new_size[0]) // 2, (target_size - new_size[1]) // 2)
+            new_img.paste(image, position)
+        else:
+            # Resize width to target size, adjust height placement based on the fitting_mode
+            width, height = image.size
+            new_height = int(height * (target_size / float(width)))
+            image = image.resize((target_size, new_height), resample=resample_filters[resample])
+            new_img = Image.new("RGB", (target_size, target_size), fill_color)
+
+            if fitting_mode == 'top':
+                position = (0, 0)
+            elif fitting_mode == 'bottom':
+                position = (0, target_size - new_height)
+            elif fitting_mode == 'center':
+                position = (0, (target_size - new_height) // 2)
+
+            new_img.paste(image, position)
+
+        return pil2tensor(new_img)
+
+
+NODE_CLASS_MAPPINGS = {
+    "ImageSquareAdapterNode": ImageSquareAdapterNode
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "ImageSquareAdapterNode": "Image Square Adapter Node"
+}

utils.py

@@ -0,0 +1,11 @@
+from PIL import Image
+import torch
+import numpy as np
+
+# Tensor to PIL
+def tensor2pil(img):
+    return Image.fromarray(np.clip(255. * img.cpu().numpy().squeeze(), 0, 255).astype(np.uint8))
+
+# PIL to Tensor
+def pil2tensor(img):
+    return torch.from_numpy(np.array(img).astype(np.float32) / 255.0).unsqueeze(0)