Working in integration with other ComfyUI workflows

__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):

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