Merge branch 'comfyanonymous:master' into master

comfy/ldm/modules/diffusionmodules/mmdit.py

@@ -71,45 +71,33 @@ class PatchEmbed(nn.Module):
             strict_img_size: bool = True,
             dynamic_img_pad: bool = True,
             padding_mode='circular',
+            conv3d=False,
             dtype=None,
             device=None,
             operations=None,
     ):
         super().__init__()
-        self.patch_size = (patch_size, patch_size)
+        try:
+            len(patch_size)
+            self.patch_size = patch_size
+        except:
+            if conv3d:
+                self.patch_size = (patch_size, patch_size, patch_size)
+            else:
+                self.patch_size = (patch_size, patch_size)
         self.padding_mode = padding_mode
-        if img_size is not None:
-            self.img_size = (img_size, img_size)
-            self.grid_size = tuple([s // p for s, p in zip(self.img_size, self.patch_size)])
-            self.num_patches = self.grid_size[0] * self.grid_size[1]
-        else:
-            self.img_size = None
-            self.grid_size = None
-            self.num_patches = None
 
         # flatten spatial dim and transpose to channels last, kept for bwd compat
         self.flatten = flatten
         self.strict_img_size = strict_img_size
         self.dynamic_img_pad = dynamic_img_pad
-
+        if conv3d:
-        self.proj = operations.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size, bias=bias, dtype=dtype, device=device)
+            self.proj = operations.Conv3d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size, bias=bias, dtype=dtype, device=device)
+        else:
+            self.proj = operations.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size, bias=bias, dtype=dtype, device=device)
         self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()
 
     def forward(self, x):
-        # B, C, H, W = x.shape
-        # if self.img_size is not None:
-        #     if self.strict_img_size:
-        #         _assert(H == self.img_size[0], f"Input height ({H}) doesn't match model ({self.img_size[0]}).")
-        #         _assert(W == self.img_size[1], f"Input width ({W}) doesn't match model ({self.img_size[1]}).")
-        #     elif not self.dynamic_img_pad:
-        #         _assert(
-        #             H % self.patch_size[0] == 0,
-        #             f"Input height ({H}) should be divisible by patch size ({self.patch_size[0]})."
-        #         )
-        #         _assert(
-        #             W % self.patch_size[1] == 0,
-        #             f"Input width ({W}) should be divisible by patch size ({self.patch_size[1]})."
-        #         )
         if self.dynamic_img_pad:
             x = comfy.ldm.common_dit.pad_to_patch_size(x, self.patch_size, padding_mode=self.padding_mode)
         x = self.proj(x)

comfy_extras/nodes_load_3d.py

@@ -0,0 +1,95 @@
+import nodes
+import folder_paths
+import os
+
+def normalize_path(path):
+    return path.replace('\\', '/')
+
+class Load3D():
+    @classmethod
+    def INPUT_TYPES(s):
+        input_dir = os.path.join(folder_paths.get_input_directory(), "3d")
+
+        os.makedirs(input_dir, exist_ok=True)
+
+        files = [normalize_path(os.path.join("3d", f)) for f in os.listdir(input_dir) if f.endswith(('.gltf', '.glb', '.obj', '.mtl', '.fbx', '.stl'))]
+
+        return {"required": {
+            "model_file": (sorted(files), {"file_upload": True}),
+            "image": ("LOAD_3D", {}),
+            "width": ("INT", {"default": 1024, "min": 1, "max": 4096, "step": 1}),
+            "height": ("INT", {"default": 1024, "min": 1, "max": 4096, "step": 1}),
+            "show_grid": ([True, False],),
+            "camera_type": (["perspective", "orthographic"],),
+            "view": (["front", "right", "top", "isometric"],),
+            "material": (["original", "normal", "wireframe", "depth"],),
+            "bg_color": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFF, "step": 1, "display": "color"}),
+            "light_intensity": ("INT", {"default": 10, "min": 1, "max": 20, "step": 1}),
+            "up_direction": (["original", "-x", "+x", "-y", "+y", "-z", "+z"],),
+        }}
+
+    RETURN_TYPES = ("IMAGE", "MASK", "STRING")
+    RETURN_NAMES = ("image", "mask", "mesh_path")
+
+    FUNCTION = "process"
+
+    CATEGORY = "3d"
+
+    def process(self, model_file, image, **kwargs):
+        imagepath = folder_paths.get_annotated_filepath(image)
+
+        load_image_node = nodes.LoadImage()
+
+        output_image, output_mask = load_image_node.load_image(image=imagepath)
+
+        return output_image, output_mask, model_file,
+
+class Load3DAnimation():
+    @classmethod
+    def INPUT_TYPES(s):
+        input_dir = os.path.join(folder_paths.get_input_directory(), "3d")
+
+        os.makedirs(input_dir, exist_ok=True)
+
+        files = [normalize_path(os.path.join("3d", f)) for f in os.listdir(input_dir) if f.endswith(('.gltf', '.glb', '.fbx'))]
+
+        return {"required": {
+            "model_file": (sorted(files), {"file_upload": True}),
+            "image": ("LOAD_3D_ANIMATION", {}),
+            "width": ("INT", {"default": 1024, "min": 1, "max": 4096, "step": 1}),
+            "height": ("INT", {"default": 1024, "min": 1, "max": 4096, "step": 1}),
+            "show_grid": ([True, False],),
+            "camera_type": (["perspective", "orthographic"],),
+            "view": (["front", "right", "top", "isometric"],),
+            "material": (["original", "normal", "wireframe", "depth"],),
+            "bg_color": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFF, "step": 1, "display": "color"}),
+            "light_intensity": ("INT", {"default": 10, "min": 1, "max": 20, "step": 1}),
+            "up_direction": (["original", "-x", "+x", "-y", "+y", "-z", "+z"],),
+            "animation_speed": (["0.1", "0.5", "1", "1.5", "2"], {"default": "1"}),
+        }}
+
+    RETURN_TYPES = ("IMAGE", "MASK", "STRING")
+    RETURN_NAMES = ("image", "mask", "mesh_path")
+
+    FUNCTION = "process"
+
+    CATEGORY = "3d"
+
+    def process(self, model_file, image, **kwargs):
+        imagepath = folder_paths.get_annotated_filepath(image)
+
+        load_image_node = nodes.LoadImage()
+
+        output_image, output_mask = load_image_node.load_image(image=imagepath)
+
+        return output_image, output_mask, model_file,
+
+NODE_CLASS_MAPPINGS = {
+    "Load3D": Load3D,
+    "Load3DAnimation": Load3DAnimation
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "Load3D": "Load 3D",
+    "Load3DAnimation": "Load 3D - Animation"
+}

folder_paths.py

@@ -200,10 +200,17 @@ def add_model_folder_path(folder_name: str, full_folder_path: str, is_default: b
     global folder_names_and_paths
     folder_name = map_legacy(folder_name)
     if folder_name in folder_names_and_paths:
-        if is_default:
+        paths, _exts = folder_names_and_paths[folder_name]
-            folder_names_and_paths[folder_name][0].insert(0, full_folder_path)
+        if full_folder_path in paths:
+            if is_default and paths[0] != full_folder_path:
+                # If the path to the folder is not the first in the list, move it to the beginning.
+                paths.remove(full_folder_path)
+                paths.insert(0, full_folder_path)
         else:
-            folder_names_and_paths[folder_name][0].append(full_folder_path)
+            if is_default:
+                paths.insert(0, full_folder_path)
+            else:
+                paths.append(full_folder_path)
     else:
         folder_names_and_paths[folder_name] = ([full_folder_path], set())
 

main.py

@@ -7,6 +7,9 @@ import folder_paths
 import time
 from comfy.cli_args import args
 from app.logger import setup_logger
+import itertools
+import utils.extra_config
+import logging
 
 if __name__ == "__main__":
     #NOTE: These do not do anything on core ComfyUI which should already have no communication with the internet, they are for custom nodes.
@@ -16,6 +19,40 @@ if __name__ == "__main__":
 
 setup_logger(log_level=args.verbose)
 
+def apply_custom_paths():
+    # extra model paths
+    extra_model_paths_config_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), "extra_model_paths.yaml")
+    if os.path.isfile(extra_model_paths_config_path):
+        utils.extra_config.load_extra_path_config(extra_model_paths_config_path)
+
+    if args.extra_model_paths_config:
+        for config_path in itertools.chain(*args.extra_model_paths_config):
+            utils.extra_config.load_extra_path_config(config_path)
+
+    # --output-directory, --input-directory, --user-directory
+    if args.output_directory:
+        output_dir = os.path.abspath(args.output_directory)
+        logging.info(f"Setting output directory to: {output_dir}")
+        folder_paths.set_output_directory(output_dir)
+
+    # These are the default folders that checkpoints, clip and vae models will be saved to when using CheckpointSave, etc.. nodes
+    folder_paths.add_model_folder_path("checkpoints", os.path.join(folder_paths.get_output_directory(), "checkpoints"))
+    folder_paths.add_model_folder_path("clip", os.path.join(folder_paths.get_output_directory(), "clip"))
+    folder_paths.add_model_folder_path("vae", os.path.join(folder_paths.get_output_directory(), "vae"))
+    folder_paths.add_model_folder_path("diffusion_models",
+                                       os.path.join(folder_paths.get_output_directory(), "diffusion_models"))
+    folder_paths.add_model_folder_path("loras", os.path.join(folder_paths.get_output_directory(), "loras"))
+
+    if args.input_directory:
+        input_dir = os.path.abspath(args.input_directory)
+        logging.info(f"Setting input directory to: {input_dir}")
+        folder_paths.set_input_directory(input_dir)
+
+    if args.user_directory:
+        user_dir = os.path.abspath(args.user_directory)
+        logging.info(f"Setting user directory to: {user_dir}")
+        folder_paths.set_user_directory(user_dir)
+
 
 def execute_prestartup_script():
     def execute_script(script_path):
@@ -57,18 +94,16 @@ def execute_prestartup_script():
             print("{:6.1f} seconds{}:".format(n[0], import_message), n[1])
         print()
 
+apply_custom_paths()
 execute_prestartup_script()
 
 
 # Main code
 import asyncio
-import itertools
 import shutil
 import threading
 import gc
 
-import logging
-import utils.extra_config
 
 if os.name == "nt":
     logging.getLogger("xformers").addFilter(lambda record: 'A matching Triton is not available' not in record.getMessage())
@@ -209,14 +244,6 @@ if __name__ == "__main__":
     server = server.PromptServer(loop)
     q = execution.PromptQueue(server)
 
-    extra_model_paths_config_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), "extra_model_paths.yaml")
-    if os.path.isfile(extra_model_paths_config_path):
-        utils.extra_config.load_extra_path_config(extra_model_paths_config_path)
-
-    if args.extra_model_paths_config:
-        for config_path in itertools.chain(*args.extra_model_paths_config):
-            utils.extra_config.load_extra_path_config(config_path)
-
     nodes.init_extra_nodes(init_custom_nodes=not args.disable_all_custom_nodes)
 
     cuda_malloc_warning()
@@ -226,28 +253,6 @@ if __name__ == "__main__":
 
     threading.Thread(target=prompt_worker, daemon=True, args=(q, server,)).start()
 
-    if args.output_directory:
-        output_dir = os.path.abspath(args.output_directory)
-        logging.info(f"Setting output directory to: {output_dir}")
-        folder_paths.set_output_directory(output_dir)
-
-    #These are the default folders that checkpoints, clip and vae models will be saved to when using CheckpointSave, etc.. nodes
-    folder_paths.add_model_folder_path("checkpoints", os.path.join(folder_paths.get_output_directory(), "checkpoints"))
-    folder_paths.add_model_folder_path("clip", os.path.join(folder_paths.get_output_directory(), "clip"))
-    folder_paths.add_model_folder_path("vae", os.path.join(folder_paths.get_output_directory(), "vae"))
-    folder_paths.add_model_folder_path("diffusion_models", os.path.join(folder_paths.get_output_directory(), "diffusion_models"))
-    folder_paths.add_model_folder_path("loras", os.path.join(folder_paths.get_output_directory(), "loras"))
-
-    if args.input_directory:
-        input_dir = os.path.abspath(args.input_directory)
-        logging.info(f"Setting input directory to: {input_dir}")
-        folder_paths.set_input_directory(input_dir)
-    
-    if args.user_directory:
-        user_dir = os.path.abspath(args.user_directory)
-        logging.info(f"Setting user directory to: {user_dir}")
-        folder_paths.set_user_directory(user_dir)
-
     if args.quick_test_for_ci:
         exit(0)
 

nodes.py

@@ -2150,6 +2150,7 @@ def init_builtin_extra_nodes():
         "nodes_mahiro.py",
         "nodes_lt.py",
         "nodes_hooks.py",
+        "nodes_load_3d.py",
     ]
 
     import_failed = []

tests-unit/comfy_test/folder_path_test.py

@@ -7,6 +7,14 @@ from unittest.mock import patch
 
 import folder_paths
 
+@pytest.fixture()
+def clear_folder_paths():
+    # Clear the global dictionary before each test to ensure isolation
+    original = folder_paths.folder_names_and_paths.copy()
+    folder_paths.folder_names_and_paths.clear()
+    yield
+    folder_paths.folder_names_and_paths = original
+
 @pytest.fixture
 def temp_dir():
     with tempfile.TemporaryDirectory() as tmpdirname:
@@ -30,9 +38,33 @@ def test_get_annotated_filepath():
     assert folder_paths.get_annotated_filepath("test.txt", default_dir) == os.path.join(default_dir, "test.txt")
     assert folder_paths.get_annotated_filepath("test.txt [output]") == os.path.join(folder_paths.get_output_directory(), "test.txt")
 
-def test_add_model_folder_path():
+def test_add_model_folder_path_append(clear_folder_paths):
-    folder_paths.add_model_folder_path("test_folder", "/test/path")
+    folder_paths.add_model_folder_path("test_folder", "/default/path", is_default=True)
-    assert "/test/path" in folder_paths.get_folder_paths("test_folder")
+    folder_paths.add_model_folder_path("test_folder", "/test/path", is_default=False)
+    assert folder_paths.get_folder_paths("test_folder") == ["/default/path", "/test/path"]
+
+
+def test_add_model_folder_path_insert(clear_folder_paths):
+    folder_paths.add_model_folder_path("test_folder", "/test/path", is_default=False)
+    folder_paths.add_model_folder_path("test_folder", "/default/path", is_default=True)
+    assert folder_paths.get_folder_paths("test_folder") == ["/default/path", "/test/path"]
+
+
+def test_add_model_folder_path_re_add_existing_default(clear_folder_paths):
+    folder_paths.add_model_folder_path("test_folder", "/test/path", is_default=False)
+    folder_paths.add_model_folder_path("test_folder", "/old_default/path", is_default=True)
+    assert folder_paths.get_folder_paths("test_folder") == ["/old_default/path", "/test/path"]
+    folder_paths.add_model_folder_path("test_folder", "/test/path", is_default=True)
+    assert folder_paths.get_folder_paths("test_folder") == ["/test/path", "/old_default/path"]
+
+
+def test_add_model_folder_path_re_add_existing_non_default(clear_folder_paths):
+    folder_paths.add_model_folder_path("test_folder", "/test/path", is_default=False)
+    folder_paths.add_model_folder_path("test_folder", "/default/path", is_default=True)
+    assert folder_paths.get_folder_paths("test_folder") == ["/default/path", "/test/path"]
+    folder_paths.add_model_folder_path("test_folder", "/test/path", is_default=False)
+    assert folder_paths.get_folder_paths("test_folder") == ["/default/path", "/test/path"]
+
 
 def test_recursive_search(temp_dir):
     os.makedirs(os.path.join(temp_dir, "subdir"))