main.py 5.76 KB
Newer Older
1
2
3
import os
import importlib.util
import folder_paths
4
import time
5
6
7

def execute_prestartup_script():
    def execute_script(script_path):
8
9
10
11
12
13
14
15
16
        module_name = os.path.splitext(script_path)[0]
        try:
            spec = importlib.util.spec_from_file_location(module_name, script_path)
            module = importlib.util.module_from_spec(spec)
            spec.loader.exec_module(module)
            return True
        except Exception as e:
            print(f"Failed to execute startup-script: {script_path} / {e}")
        return False
17
18
19
20

    node_paths = folder_paths.get_folder_paths("custom_nodes")
    for custom_node_path in node_paths:
        possible_modules = os.listdir(custom_node_path)
21
        node_prestartup_times = []
22
23
24
25
26
27
28

        for possible_module in possible_modules:
            module_path = os.path.join(custom_node_path, possible_module)
            if os.path.isfile(module_path) or module_path.endswith(".disabled") or module_path == "__pycache__":
                continue

            script_path = os.path.join(module_path, "prestartup_script.py")
29
30
31
32
33
34
35
36
37
38
39
40
41
            if os.path.exists(script_path):
                time_before = time.perf_counter()
                success = execute_script(script_path)
                node_prestartup_times.append((time.perf_counter() - time_before, module_path, success))
    if len(node_prestartup_times) > 0:
        print("\nPrestartup times for custom nodes:")
        for n in sorted(node_prestartup_times):
            if n[2]:
                import_message = ""
            else:
                import_message = " (PRESTARTUP FAILED)"
            print("{:6.1f} seconds{}:".format(n[0], import_message), n[1])
        print()
42
43
44
45
46

execute_prestartup_script()


# Main code
EllangoK's avatar
EllangoK committed
47
import asyncio
48
import itertools
49
import shutil
comfyanonymous's avatar
comfyanonymous committed
50
import threading
51
import gc
52

53
from comfy.cli_args import args
comfyanonymous's avatar
comfyanonymous committed
54

pythongosssss's avatar
pythongosssss committed
55
56
57
58
if os.name == "nt":
    import logging
    logging.getLogger("xformers").addFilter(lambda record: 'A matching Triton is not available' not in record.getMessage())

comfyanonymous's avatar
comfyanonymous committed
59
if __name__ == "__main__":
EllangoK's avatar
EllangoK committed
60
61
62
63
    if args.cuda_device is not None:
        os.environ['CUDA_VISIBLE_DEVICES'] = str(args.cuda_device)
        print("Set cuda device to:", args.cuda_device)

64
    import cuda_malloc
EllangoK's avatar
EllangoK committed
65

66
import comfy.utils
EllangoK's avatar
EllangoK committed
67
import yaml
68

69
import execution
EllangoK's avatar
EllangoK committed
70
import server
space-nuko's avatar
space-nuko committed
71
from server import BinaryEventTypes
EllangoK's avatar
EllangoK committed
72
from nodes import init_custom_nodes
73
import comfy.model_management
74

pythongosssss's avatar
pythongosssss committed
75
def prompt_worker(q, server):
76
    e = execution.PromptExecutor(server)
comfyanonymous's avatar
comfyanonymous committed
77
    while True:
78
        item, item_id = q.get()
79
80
81
        execution_start_time = time.perf_counter()
        prompt_id = item[1]
        e.execute(item[2], prompt_id, item[3], item[4])
82
        q.task_done(item_id, e.outputs_ui)
83
84
        if server.client_id is not None:
            server.send_sync("executing", { "node": None, "prompt_id": prompt_id }, server.client_id)
comfyanonymous's avatar
comfyanonymous committed
85

86
87
88
        print("Prompt executed in {:.2f} seconds".format(time.perf_counter() - execution_start_time))
        gc.collect()
        comfy.model_management.soft_empty_cache()
reaper47's avatar
reaper47 committed
89

90
91
async def run(server, address='', port=8188, verbose=True, call_on_start=None):
    await asyncio.gather(server.start(address, port, verbose, call_on_start), server.publish_loop())
comfyanonymous's avatar
comfyanonymous committed
92

reaper47's avatar
reaper47 committed
93

pythongosssss's avatar
pythongosssss committed
94
def hijack_progress(server):
95
    def hook(value, total, preview_image):
reaper47's avatar
reaper47 committed
96
        server.send_sync("progress", {"value": value, "max": total}, server.client_id)
97
98
        if preview_image is not None:
            server.send_sync(BinaryEventTypes.UNENCODED_PREVIEW_IMAGE, preview_image, server.client_id)
99
    comfy.utils.set_progress_bar_global_hook(hook)
comfyanonymous's avatar
comfyanonymous committed
100

reaper47's avatar
reaper47 committed
101

102
103
104
def cleanup_temp():
    temp_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), "temp")
    if os.path.exists(temp_dir):
105
        shutil.rmtree(temp_dir, ignore_errors=True)
106

reaper47's avatar
reaper47 committed
107

108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
def load_extra_path_config(yaml_path):
    with open(yaml_path, 'r') as stream:
        config = yaml.safe_load(stream)
    for c in config:
        conf = config[c]
        if conf is None:
            continue
        base_path = None
        if "base_path" in conf:
            base_path = conf.pop("base_path")
        for x in conf:
            for y in conf[x].split("\n"):
                if len(y) == 0:
                    continue
                full_path = y
                if base_path is not None:
                    full_path = os.path.join(base_path, full_path)
                print("Adding extra search path", x, full_path)
                folder_paths.add_model_folder_path(x, full_path)

reaper47's avatar
reaper47 committed
128

comfyanonymous's avatar
comfyanonymous committed
129
if __name__ == "__main__":
130
131
    cleanup_temp()

132
133
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
pythongosssss's avatar
pythongosssss committed
134
    server = server.PromptServer(loop)
135
    q = execution.PromptQueue(server)
136

137
138
139
140
141
142
143
144
    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):
        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):
            load_extra_path_config(config_path)

145
146
    init_custom_nodes()
    server.add_routes()
pythongosssss's avatar
pythongosssss committed
147
148
    hijack_progress(server)

reaper47's avatar
reaper47 committed
149
    threading.Thread(target=prompt_worker, daemon=True, args=(q, server,)).start()
150

EllangoK's avatar
EllangoK committed
151
152
    if args.output_directory:
        output_dir = os.path.abspath(args.output_directory)
153
        print(f"Setting output directory to: {output_dir}")
154
155
        folder_paths.set_output_directory(output_dir)

EllangoK's avatar
EllangoK committed
156
    if args.quick_test_for_ci:
157
158
        exit(0)

159
    call_on_start = None
EllangoK's avatar
EllangoK committed
160
    if args.auto_launch:
161
162
        def startup_server(address, port):
            import webbrowser
reaper47's avatar
reaper47 committed
163
            webbrowser.open(f"http://{address}:{port}")
164
165
        call_on_start = startup_server

reaper47's avatar
reaper47 committed
166
    try:
EllangoK's avatar
EllangoK committed
167
        loop.run_until_complete(run(server, address=args.listen, port=args.port, verbose=not args.dont_print_server, call_on_start=call_on_start))
reaper47's avatar
reaper47 committed
168
169
    except KeyboardInterrupt:
        print("\nStopped server")
comfyanonymous's avatar
comfyanonymous committed
170

171
    cleanup_temp()