From 88e05cad696b05366ec09be702e24d16eaa4f4f6 Mon Sep 17 00:00:00 2001
From: Anson Kao <ansonkao@ansons-mbp.home>
Date: Wed, 29 Nov 2023 21:45:27 -0500
Subject: [PATCH] Working in integration with other ComfyUI workflows

---
 __init__.py                 | 10 +++++-----
 comfyui_discopixel_nodes.py |  9 +++++++++
 2 files changed, 14 insertions(+), 5 deletions(-)

diff --git a/__init__.py b/__init__.py
index 3433576..b53e22e 100644
--- a/__init__.py
+++ b/__init__.py
@@ -1,19 +1,19 @@
 """
 @author: Anson Kao
-@title: ComfyUI Geometry
-@nickname: ComfyUI Geometry
-@description: A small collection of custom nodes for use with ComfyUI, for geometry calculations
+@title: ComfyUI Discopixel
+@nickname: ComfyUI Discopixel
+@description: A small collection of custom nodes for use with ComfyUI, by Discopixel
 """
 
 import importlib
 
-print(f"Loading ComfyUI Geometry!")
+print(f"Loading ComfyUI Discopixel nodes!")
 
 NODE_CLASS_MAPPINGS = {}
 NODE_DISPLAY_NAME_MAPPINGS = {}
 
 NODE_MODULES = [
-    ".comfyui_geometry_nodes",
+    ".comfyui_discopixel_nodes",
 ]
 
 def load_nodes(module_name: str):
diff --git a/comfyui_discopixel_nodes.py b/comfyui_discopixel_nodes.py
index 8838700..039d844 100644
--- a/comfyui_discopixel_nodes.py
+++ b/comfyui_discopixel_nodes.py
@@ -195,6 +195,15 @@ class TransformTemplateOntoFaceMask:
         )
 
     def transform_template(self, face_mask, template_image, template_mask):
+        # Ensure batch dimension for consistent handling
+        if face_mask.dim() == 2:
+            face_mask = face_mask.unsqueeze(0)
+        if template_image.dim() == 3:
+            template_image = template_image.unsqueeze(0)
+        if template_mask.dim() == 2:
+            template_mask = template_mask.unsqueeze(0)
+
+        # Calculate the transformation
         centroid_x, centroid_y, ellipse_length, ellipse_width, rotation, matrix = self.calculate_transformation(face_mask, template_mask)
 
         # Create a blank canvas the same size as the mask