wangbo/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2026-01-10 06:10:50 +08:00
Code Issues Actions 10 Packages Projects Releases Wiki Activity

wip

Browse Source
This commit is contained in:
Benjamin Berman 2023-08-02 17:59:10 -07:00
parent 82d0edf121
commit 3d3c5ae344
13 changed files with 60 additions and 65 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

0
cuda_malloc.py → comfy/cmd/cuda_malloc.py
View File

0
execution.py → comfy/cmd/execution.py
View File

8
folder_paths.py → comfy/cmd/folder_paths.py
View File

@ -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 = {}

3
latent_preview.py → comfy/cmd/latent_preview.py
View File

@ -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

26
main.py → comfy/cmd/main.py
View File

@ -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")

10
server.py → comfy/cmd/server.py
View File

@ -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)

6
comfy/diffusers_load.py
View File

@ -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

39
nodes.py → comfy/nodes/base_nodes.py
View File

@ -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"

17
comfy/nodes/package.py
View File

@ -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

4
comfy_extras/nodes/nodes_hypernetwork.py
View File

@ -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",)

2
comfy_extras/nodes/nodes_model_merging.py
View File

@ -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

6
comfy_extras/nodes/nodes_upscale_model.py
View File

@ -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"

4
setup.py
View File

@ -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'
],
},
)