add more output

comfyanonymous · Nov 26, 2024 · e783772 · e783772
1 parent 2dacd69
commit e783772
Showing 1 changed file with 14 additions and 6 deletions.
diff --git a/comfy_extras/nodes_load_3d.py b/comfy_extras/nodes_load_3d.py
@@ -28,17 +28,21 @@ def INPUT_TYPES(s):
             "up_direction": (["original", "-x", "+x", "-y", "+y", "-z", "+z"],),
         }}
 
-    RETURN_TYPES = ("IMAGE",)
+    RETURN_TYPES = ("IMAGE", "MASK", "STRING")
+    RETURN_NAMES = ("image", "mask", "mesh_path")
+
     FUNCTION = "process"
 
     CATEGORY = "3d"
 
-    def process(self, image, **kwargs):
+    def process(self, model_file, image, **kwargs):
         imagepath = folder_paths.get_annotated_filepath(image)
 
         load_image_node = nodes.LoadImage()
 
-        return load_image_node.load_image(image=imagepath)
+        output_image, output_mask = load_image_node.load_image(image=imagepath)
+
+        return output_image, output_mask, model_file,
 
 class Load3DAnimation():
     @classmethod
@@ -64,17 +68,21 @@ def INPUT_TYPES(s):
             "animation_speed": (["0.1", "0.5", "1", "1.5", "2"], {"default": "1"}),
         }}
 
-    RETURN_TYPES = ("IMAGE",)
+    RETURN_TYPES = ("IMAGE", "MASK", "STRING")
+    RETURN_NAMES = ("image", "mask", "mesh_path")
+
     FUNCTION = "process"
 
     CATEGORY = "3d"
 
-    def process(self, image, **kwargs):
+    def process(self, model_file, image, **kwargs):
         imagepath = folder_paths.get_annotated_filepath(image)
 
         load_image_node = nodes.LoadImage()
 
-        return load_image_node.load_image(image=imagepath)
+        output_image, output_mask = load_image_node.load_image(image=imagepath)
+
+        return output_image, output_mask, model_file,
 
 class Preview3D():
     @classmethod