From 25a3220e6f01acb60a60c8afa494c53050a908f0 Mon Sep 17 00:00:00 2001
From: Pheat-AI <info@pheat.ai>
Date: Thu, 5 Sep 2024 22:49:35 -0700
Subject: [PATCH] Create batch_enlarged_overlay.py

---
 batch_enlarged_overlay.py | 84 +++++++++++++++++++++++++++++++++++++++
 1 file changed, 84 insertions(+)
 create mode 100644 batch_enlarged_overlay.py

diff --git a/batch_enlarged_overlay.py b/batch_enlarged_overlay.py
new file mode 100644
index 0000000..b70375e
--- /dev/null
+++ b/batch_enlarged_overlay.py
@@ -0,0 +1,84 @@
+import torch
+import numpy as np
+from PIL import Image, ImageOps
+
+# BATCH ENLARGED OVERLAY NODE
+
+class REMADE_Batch_Enlarged_Overlay:
+    def __init__(self):
+        pass 
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "gen_images": ("IMAGE",),  # A batch of generated images
+                "original_image": ("IMAGE",),  # A single original image
+                "enlarge_factor": ("FLOAT", {
+                    "default": 1.05,  # Default enlarge factor
+                    "min": 1.0, 
+                    "max": 2.0, 
+                    "step": 0.01,
+                }),
+            },
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "batch_enlarged_overlay"
+
+    CATEGORY = "REMADE/Image"
+
+    def batch_enlarged_overlay(self, gen_images, original_image, enlarge_factor):
+        # List to hold the result tensors for each image in the batch
+        overlaid_images = []
+
+        # Convert the original image to PIL format and apply the enlargement
+        original_image = tensor2pil(original_image).convert("RGBA")
+        width, height = original_image.size
+        new_size = (int(width * enlarge_factor), int(height * enlarge_factor))
+
+        # Resize the original image (slightly enlarge)
+        enlarged_original = original_image.resize(new_size, Image.LANCZOS)
+
+        # Center the enlarged original image on a 1024x1024 canvas
+        enlarged_image_on_canvas = Image.new("RGBA", (width, height), (0, 0, 0, 0))  # Transparent background
+        offset_x = (width - new_size[0]) // 2
+        offset_y = (height - new_size[1]) // 2
+        enlarged_image_on_canvas.paste(enlarged_original, (offset_x, offset_y))
+
+        # Iterate through each generated image in the batch
+        for gen_image in gen_images:
+            gen_image = tensor2pil(gen_image).convert("RGBA")
+
+            # Overlay the enlarged original image onto the generated image
+            result_image = Image.alpha_composite(gen_image, enlarged_image_on_canvas)
+
+            # Convert the result back to tensor
+            result_tensor = pil2tensor(result_image)
+
+            # Append the resulting tensor to the batch result list
+            overlaid_images.append(result_tensor)
+        
+        # Convert the list of tensors into a single batch tensor
+        result = torch.cat(overlaid_images, dim=0)
+
+        return (result, )
+
+# PIL to Tensor
+def pil2tensor(image):
+    return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0)
+
+# Tensor to PIL
+def tensor2pil(image):
+    return Image.fromarray(np.clip(255. * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8))
+
+# A dictionary that contains all nodes you want to export with their names
+# NOTE: names should be globally unique
+NODE_CLASS_MAPPINGS = {
+    "Batch Enlarged Overlay": REMADE_Batch_Enlarged_Overlay
+}
+
+# A dictionary that contains the friendly/humanly readable titles for the nodes
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "Batch Enlarged Overlay": "Batch Enlarged Overlay"
+}