diff --git a/cuda_malloc.py b/comfy/cmd/cuda_malloc.py similarity index 100% rename from cuda_malloc.py rename to comfy/cmd/cuda_malloc.py diff --git a/execution.py b/comfy/cmd/execution.py similarity index 100% rename from execution.py rename to comfy/cmd/execution.py diff --git a/folder_paths.py b/comfy/cmd/folder_paths.py similarity index 96% rename from folder_paths.py rename to comfy/cmd/folder_paths.py index eb7d39b88..8e3e50abd 100644 --- a/folder_paths.py +++ b/comfy/cmd/folder_paths.py @@ -7,7 +7,7 @@ supported_pt_extensions = set(['.ckpt', '.pt', '.bin', '.pth', '.safetensors']) folder_names_and_paths = {} base_path = os.path.dirname(os.path.realpath(__file__)) -models_dir = os.path.join(base_path, "models") +models_dir = os.path.join(base_path, "../../models") folder_names_and_paths["checkpoints"] = ([os.path.join(models_dir, "checkpoints")], supported_ckpt_extensions) folder_names_and_paths["configs"] = ([os.path.join(models_dir, "configs")], [".yaml"]) @@ -26,13 +26,13 @@ folder_names_and_paths["gligen"] = ([os.path.join(models_dir, "gligen")], suppor folder_names_and_paths["upscale_models"] = ([os.path.join(models_dir, "upscale_models")], supported_pt_extensions) -folder_names_and_paths["custom_nodes"] = ([os.path.join(base_path, "custom_nodes")], []) +folder_names_and_paths["custom_nodes"] = ([os.path.join(base_path, "../../custom_nodes")], []) folder_names_and_paths["hypernetworks"] = ([os.path.join(models_dir, "hypernetworks")], supported_pt_extensions) -output_directory = os.path.join(os.path.dirname(os.path.realpath(__file__)), "output") +output_directory = os.path.realpath(os.path.join(os.path.dirname(os.path.realpath(__file__)), "../../output")) temp_directory = os.path.join(os.path.dirname(os.path.realpath(__file__)), "temp") -input_directory = os.path.join(os.path.dirname(os.path.realpath(__file__)), "input") +input_directory = os.path.realpath(os.path.join(os.path.dirname(os.path.realpath(__file__)), "../../input")) filename_list_cache = {} diff --git a/latent_preview.py b/comfy/cmd/latent_preview.py similarity index 98% rename from latent_preview.py rename to comfy/cmd/latent_preview.py index 30c1d1317..eff000c1f 100644 --- a/latent_preview.py +++ b/comfy/cmd/latent_preview.py @@ -1,10 +1,9 @@ import torch from PIL import Image -import struct import numpy as np from comfy.cli_args import args, LatentPreviewMethod from comfy.taesd.taesd import TAESD -import folder_paths +from ..cmd import folder_paths MAX_PREVIEW_RESOLUTION = 512 diff --git a/main.py b/comfy/cmd/main.py similarity index 90% rename from main.py rename to comfy/cmd/main.py index b7700b34a..c4288ace1 100644 --- a/main.py +++ b/comfy/cmd/main.py @@ -1,6 +1,6 @@ import os import importlib.util -import folder_paths +from ..cmd import folder_paths import time def execute_prestartup_script(): @@ -56,23 +56,21 @@ if os.name == "nt": import logging logging.getLogger("xformers").addFilter(lambda record: 'A matching Triton is not available' not in record.getMessage()) -if __name__ == "__main__": - if args.cuda_device is not None: - os.environ['CUDA_VISIBLE_DEVICES'] = str(args.cuda_device) - print("Set cuda device to:", args.cuda_device) - import cuda_malloc +if args.cuda_device is not None: + os.environ['CUDA_VISIBLE_DEVICES'] = str(args.cuda_device) + print("Set cuda device to:", args.cuda_device) import comfy.utils import yaml -import execution -import server -from server import BinaryEventTypes +from ..cmd import execution +from ..cmd import server as server_module +from .server import BinaryEventTypes import comfy.model_management -def prompt_worker(q: execution.PromptQueue, _server: server.PromptServer): +def prompt_worker(q: execution.PromptQueue, _server: server_module.PromptServer): e = execution.PromptExecutor(_server) while True: item, item_id = q.get() @@ -80,8 +78,8 @@ def prompt_worker(q: execution.PromptQueue, _server: server.PromptServer): prompt_id = item[1] e.execute(item[2], prompt_id, item[3], item[4]) q.task_done(item_id, e.outputs_ui) - if server.client_id is not None: - server.send_sync("executing", { "node": None, "prompt_id": prompt_id }, server.client_id) + if _server.client_id is not None: + _server.send_sync("executing", { "node": None, "prompt_id": prompt_id }, _server.client_id) print("Prompt executed in {:.2f} seconds".format(time.perf_counter() - execution_start_time)) gc.collect() @@ -126,12 +124,12 @@ def load_extra_path_config(yaml_path): folder_paths.add_model_folder_path(x, full_path) -if __name__ == "__main__": +def main(): cleanup_temp() loop = asyncio.new_event_loop() asyncio.set_event_loop(loop) - server = server.PromptServer(loop) + server = server_module.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") diff --git a/server.py b/comfy/cmd/server.py similarity index 99% rename from server.py rename to comfy/cmd/server.py index 631abbcc5..31fa2420c 100644 --- a/server.py +++ b/comfy/cmd/server.py @@ -15,8 +15,8 @@ import aiofiles import aiohttp from aiohttp import web -import execution -import folder_paths +from ..cmd import execution +from ..cmd import folder_paths import mimetypes from comfy.digest import digest @@ -89,7 +89,7 @@ class PromptServer(): middlewares=middlewares) self.sockets = dict() self.web_root = os.path.join(os.path.dirname( - os.path.realpath(__file__)), "web") + os.path.realpath(__file__)), "../../web") routes = web.RouteTableDef() self.routes = routes self.last_node_id = None @@ -140,7 +140,7 @@ class PromptServer(): def get_dir_by_type(dir_type=None): type_dir = "" if dir_type is None: - dir_type = "input" + dir_type = "../../input" if dir_type == "input": type_dir = folder_paths.get_input_directory() @@ -732,7 +732,7 @@ class PromptServer(): @classmethod def get_upload_dir(cls) -> str: - upload_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), "input") + upload_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), "../../input") if not os.path.exists(upload_dir): os.makedirs(upload_dir) diff --git a/comfy/diffusers_load.py b/comfy/diffusers_load.py index 11d94c340..26c2ceee6 100644 --- a/comfy/diffusers_load.py +++ b/comfy/diffusers_load.py @@ -1,13 +1,11 @@ import json -import os import yaml -import folder_paths +from .cmd import folder_paths from comfy.sd import load_checkpoint import os.path as osp -import re import torch -from safetensors.torch import load_file, save_file +from safetensors.torch import load_file from . import diffusers_convert diff --git a/nodes.py b/comfy/nodes/base_nodes.py similarity index 99% rename from nodes.py rename to comfy/nodes/base_nodes.py index 965589994..fcdc7b2fd 100644 --- a/nodes.py +++ b/comfy/nodes/base_nodes.py @@ -1,12 +1,9 @@ import torch import os -import sys import json import hashlib -import traceback import math -import time import random from PIL import Image, ImageOps @@ -25,10 +22,7 @@ import comfy.clip_vision import comfy.model_management from comfy.cli_args import args -import importlib - -import folder_paths -import latent_preview +from comfy.cmd import folder_paths, latent_preview from comfy.nodes.common import MAX_RESOLUTION @@ -410,8 +404,8 @@ class LoadLatent: class CheckpointLoader: @classmethod def INPUT_TYPES(s): - return {"required": { "config_name": (folder_paths.get_filename_list("configs"), ), - "ckpt_name": (folder_paths.get_filename_list("checkpoints"), )}} + return {"required": { "config_name": (folder_paths.get_filename_list("configs"),), + "ckpt_name": (folder_paths.get_filename_list("checkpoints"),)}} RETURN_TYPES = ("MODEL", "CLIP", "VAE") FUNCTION = "load_checkpoint" @@ -425,7 +419,7 @@ class CheckpointLoader: class CheckpointLoaderSimple: @classmethod def INPUT_TYPES(s): - return {"required": { "ckpt_name": (folder_paths.get_filename_list("checkpoints"), ), + return {"required": { "ckpt_name": (folder_paths.get_filename_list("checkpoints"),), }} RETURN_TYPES = ("MODEL", "CLIP", "VAE") FUNCTION = "load_checkpoint" @@ -467,7 +461,7 @@ class DiffusersLoader: class unCLIPCheckpointLoader: @classmethod def INPUT_TYPES(s): - return {"required": { "ckpt_name": (folder_paths.get_filename_list("checkpoints"), ), + return {"required": { "ckpt_name": (folder_paths.get_filename_list("checkpoints"),), }} RETURN_TYPES = ("MODEL", "CLIP", "VAE", "CLIP_VISION") FUNCTION = "load_checkpoint" @@ -503,7 +497,7 @@ class LoraLoader: def INPUT_TYPES(s): return {"required": { "model": ("MODEL",), "clip": ("CLIP", ), - "lora_name": (folder_paths.get_filename_list("loras"), ), + "lora_name": (folder_paths.get_filename_list("loras"),), "strength_model": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01}), "strength_clip": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01}), }} @@ -536,7 +530,7 @@ class LoraLoader: class VAELoader: @classmethod def INPUT_TYPES(s): - return {"required": { "vae_name": (folder_paths.get_filename_list("vae"), )}} + return {"required": { "vae_name": (folder_paths.get_filename_list("vae"),)}} RETURN_TYPES = ("VAE",) FUNCTION = "load_vae" @@ -551,7 +545,7 @@ class VAELoader: class ControlNetLoader: @classmethod def INPUT_TYPES(s): - return {"required": { "control_net_name": (folder_paths.get_filename_list("controlnet"), )}} + return {"required": { "control_net_name": (folder_paths.get_filename_list("controlnet"),)}} RETURN_TYPES = ("CONTROL_NET",) FUNCTION = "load_controlnet" @@ -567,7 +561,7 @@ class DiffControlNetLoader: @classmethod def INPUT_TYPES(s): return {"required": { "model": ("MODEL",), - "control_net_name": (folder_paths.get_filename_list("controlnet"), )}} + "control_net_name": (folder_paths.get_filename_list("controlnet"),)}} RETURN_TYPES = ("CONTROL_NET",) FUNCTION = "load_controlnet" @@ -660,7 +654,7 @@ class ControlNetApplyAdvanced: class UNETLoader: @classmethod def INPUT_TYPES(s): - return {"required": { "unet_name": (folder_paths.get_filename_list("unet"), ), + return {"required": { "unet_name": (folder_paths.get_filename_list("unet"),), }} RETURN_TYPES = ("MODEL",) FUNCTION = "load_unet" @@ -675,7 +669,7 @@ class UNETLoader: class CLIPLoader: @classmethod def INPUT_TYPES(s): - return {"required": { "clip_name": (folder_paths.get_filename_list("clip"), ), + return {"required": { "clip_name": (folder_paths.get_filename_list("clip"),), }} RETURN_TYPES = ("CLIP",) FUNCTION = "load_clip" @@ -690,8 +684,9 @@ class CLIPLoader: class DualCLIPLoader: @classmethod def INPUT_TYPES(s): - return {"required": { "clip_name1": (folder_paths.get_filename_list("clip"), ), "clip_name2": (folder_paths.get_filename_list("clip"), ), - }} + return {"required": { "clip_name1": (folder_paths.get_filename_list("clip"),), "clip_name2": ( + folder_paths.get_filename_list("clip"),), + }} RETURN_TYPES = ("CLIP",) FUNCTION = "load_clip" @@ -706,7 +701,7 @@ class DualCLIPLoader: class CLIPVisionLoader: @classmethod def INPUT_TYPES(s): - return {"required": { "clip_name": (folder_paths.get_filename_list("clip_vision"), ), + return {"required": { "clip_name": (folder_paths.get_filename_list("clip_vision"),), }} RETURN_TYPES = ("CLIP_VISION",) FUNCTION = "load_clip" @@ -736,7 +731,7 @@ class CLIPVisionEncode: class StyleModelLoader: @classmethod def INPUT_TYPES(s): - return {"required": { "style_model_name": (folder_paths.get_filename_list("style_models"), )}} + return {"required": { "style_model_name": (folder_paths.get_filename_list("style_models"),)}} RETURN_TYPES = ("STYLE_MODEL",) FUNCTION = "load_style_model" @@ -801,7 +796,7 @@ class unCLIPConditioning: class GLIGENLoader: @classmethod def INPUT_TYPES(s): - return {"required": { "gligen_name": (folder_paths.get_filename_list("gligen"), )}} + return {"required": { "gligen_name": (folder_paths.get_filename_list("gligen"),)}} RETURN_TYPES = ("GLIGEN",) FUNCTION = "load_gligen" diff --git a/comfy/nodes/package.py b/comfy/nodes/package.py index f66a8e074..e90f6d163 100644 --- a/comfy/nodes/package.py +++ b/comfy/nodes/package.py @@ -5,9 +5,13 @@ import pkgutil import time import types -import nodes as base_nodes +from comfy.nodes import base_nodes as base_nodes from comfy_extras import nodes as comfy_extras_nodes -import custom_nodes + +try: + import custom_nodes +except: + custom_nodes = None from comfy.nodes.package_typing import ExportedNodes from functools import reduce @@ -68,11 +72,12 @@ def import_all_nodes_in_workspace() -> ExportedNodes: comfy_extras_nodes ]), ExportedNodes()) + custom_nodes_mappings = ExportedNodes() + if custom_nodes is not None: + custom_nodes_mappings = _import_and_enumerate_nodes_in_module(custom_nodes, print_import_times=True) - custom_nodes_mappings = _import_and_enumerate_nodes_in_module(custom_nodes, print_import_times=True) - - # don't allow custom nodes to overwrite base nodes - custom_nodes_mappings -= base_and_extra + # don't allow custom nodes to overwrite base nodes + custom_nodes_mappings -= base_and_extra _comfy_nodes.update(base_and_extra + custom_nodes_mappings) return _comfy_nodes diff --git a/comfy_extras/nodes/nodes_hypernetwork.py b/comfy_extras/nodes/nodes_hypernetwork.py index d16c49aeb..15a311b64 100644 --- a/comfy_extras/nodes/nodes_hypernetwork.py +++ b/comfy_extras/nodes/nodes_hypernetwork.py @@ -1,5 +1,5 @@ import comfy.utils -import folder_paths +from comfy.cmd import folder_paths import torch def load_hypernetwork_patch(path, strength): @@ -88,7 +88,7 @@ class HypernetworkLoader: @classmethod def INPUT_TYPES(s): return {"required": { "model": ("MODEL",), - "hypernetwork_name": (folder_paths.get_filename_list("hypernetworks"), ), + "hypernetwork_name": (folder_paths.get_filename_list("hypernetworks"),), "strength": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01}), }} RETURN_TYPES = ("MODEL",) diff --git a/comfy_extras/nodes/nodes_model_merging.py b/comfy_extras/nodes/nodes_model_merging.py index bce4b3dd0..3e161555a 100644 --- a/comfy_extras/nodes/nodes_model_merging.py +++ b/comfy_extras/nodes/nodes_model_merging.py @@ -2,7 +2,7 @@ import comfy.sd import comfy.utils import comfy.model_base -import folder_paths +from comfy.cmd import folder_paths import json import os diff --git a/comfy_extras/nodes/nodes_upscale_model.py b/comfy_extras/nodes/nodes_upscale_model.py index abd182e6e..1c1730e4b 100644 --- a/comfy_extras/nodes/nodes_upscale_model.py +++ b/comfy_extras/nodes/nodes_upscale_model.py @@ -1,14 +1,14 @@ -import os from comfy_extras.chainner_models import model_loading from comfy import model_management import torch import comfy.utils -import folder_paths +from comfy.cmd import folder_paths + class UpscaleModelLoader: @classmethod def INPUT_TYPES(s): - return {"required": { "model_name": (folder_paths.get_filename_list("upscale_models"), ), + return {"required": { "model_name": (folder_paths.get_filename_list("upscale_models"),), }} RETURN_TYPES = ("UPSCALE_MODEL",) FUNCTION = "load_model" diff --git a/setup.py b/setup.py index 9c832a36c..5ae69e46e 100644 --- a/setup.py +++ b/setup.py @@ -151,8 +151,8 @@ setup( setup_requires=["pip", "wheel"], entry_points={ 'console_scripts': [ - # todo: eventually migrate main here - 'comfyui-openapi-gen = comfy.cmd.openapi_gen:main' + 'comfyui-openapi-gen = comfy.cmd.openapi_gen:main', + 'comfyui = comfy.cmd.main:main' ], }, )