diff --git a/comfy/ldm/modules/diffusionmodules/mmdit.py b/comfy/ldm/modules/diffusionmodules/mmdit.py index b052c48e0..e70f4431f 100644 --- a/comfy/ldm/modules/diffusionmodules/mmdit.py +++ b/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 - - self.proj = operations.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size, bias=bias, dtype=dtype, device=device) + if conv3d: + 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) diff --git a/comfy_extras/nodes_load_3d.py b/comfy_extras/nodes_load_3d.py new file mode 100644 index 000000000..90da9fd67 --- /dev/null +++ b/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" +} \ No newline at end of file diff --git a/folder_paths.py b/folder_paths.py index 577a7bc64..61de51202 100644 --- a/folder_paths.py +++ b/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: - folder_names_and_paths[folder_name][0].insert(0, full_folder_path) + paths, _exts = folder_names_and_paths[folder_name] + 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()) diff --git a/main.py b/main.py index dd6526e13..9d1632633 100644 --- a/main.py +++ b/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) diff --git a/nodes.py b/nodes.py index df4124544..aab31bfcb 100644 --- a/nodes.py +++ b/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 = [] diff --git a/tests-unit/comfy_test/folder_path_test.py b/tests-unit/comfy_test/folder_path_test.py index 0bbec593b..55613505e 100644 --- a/tests-unit/comfy_test/folder_path_test.py +++ b/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(): - folder_paths.add_model_folder_path("test_folder", "/test/path") - assert "/test/path" in folder_paths.get_folder_paths("test_folder") +def test_add_model_folder_path_append(clear_folder_paths): + folder_paths.add_model_folder_path("test_folder", "/default/path", is_default=True) + 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"))