main.py 5 KB
Newer Older
comfyanonymous's avatar
comfyanonymous committed
1
2
import os
import sys
3
import shutil
4

comfyanonymous's avatar
comfyanonymous committed
5
import threading
6
import asyncio
comfyanonymous's avatar
comfyanonymous committed
7

pythongosssss's avatar
pythongosssss committed
8
9
10
11
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
12
13
14
15
16
17
18
if __name__ == "__main__":
    if '--help' in sys.argv:
        print("Valid Command line Arguments:")
        print("\t--listen\t\t\tListen on 0.0.0.0 so the UI can be accessed from other computers.")
        print("\t--port 8188\t\t\tSet the listen port.")
        print("\t--dont-upcast-attention\t\tDisable upcasting of attention \n\t\t\t\t\tcan boost speed but increase the chances of black images.\n")
        print("\t--use-split-cross-attention\tUse the split cross attention optimization instead of the sub-quadratic one.\n\t\t\t\t\tIgnored when xformers is used.")
19
        print("\t--use-pytorch-cross-attention\tUse the new pytorch 2.0 cross attention function.")
20
        print("\t--disable-xformers\t\tdisables xformers")
comfyanonymous's avatar
comfyanonymous committed
21
        print()
22
        print("\t--highvram\t\t\tBy default models will be unloaded to CPU memory after being used.\n\t\t\t\t\tThis option keeps them in GPU memory.\n")
23
        print("\t--normalvram\t\t\tUsed to force normal vram use if lowvram gets automatically enabled.")
24
25
26
        print("\t--lowvram\t\t\tSplit the unet in parts to use less vram.")
        print("\t--novram\t\t\tWhen lowvram isn't enough.")
        print()
27
        print("\t--cpu\t\t\tTo use the CPU for everything (slow).")
comfyanonymous's avatar
comfyanonymous committed
28
29
        exit()

pythongosssss's avatar
pythongosssss committed
30
31
32
    if '--dont-upcast-attention' in sys.argv:
        print("disabling upcasting of attention")
        os.environ['ATTN_PRECISION'] = "fp16"
33

34
35
import execution
import server
36
37
import folder_paths
import yaml
38

pythongosssss's avatar
pythongosssss committed
39
def prompt_worker(q, server):
40
    e = execution.PromptExecutor(server)
comfyanonymous's avatar
comfyanonymous committed
41
    while True:
42
        item, item_id = q.get()
comfyanonymous's avatar
comfyanonymous committed
43
        e.execute(item[-2], item[-1])
pythongosssss's avatar
pythongosssss committed
44
        q.task_done(item_id, e.outputs)
comfyanonymous's avatar
comfyanonymous committed
45

46
47
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
48

pythongosssss's avatar
pythongosssss committed
49
50
51
52
53
54
55
56
57
def hijack_progress(server):
    from tqdm.auto import tqdm
    orig_func = getattr(tqdm, "update")
    def wrapped_func(*args, **kwargs):
        pbar = args[0]
        v = orig_func(*args, **kwargs)
        server.send_sync("progress", { "value": pbar.n, "max": pbar.total}, server.client_id)            
        return v
    setattr(tqdm, "update", wrapped_func)
comfyanonymous's avatar
comfyanonymous committed
58

59
60
61
def cleanup_temp():
    temp_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), "temp")
    if os.path.exists(temp_dir):
62
        shutil.rmtree(temp_dir, ignore_errors=True)
63

64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
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)

comfyanonymous's avatar
comfyanonymous committed
84
if __name__ == "__main__":
85
86
    cleanup_temp()

87
88
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
pythongosssss's avatar
pythongosssss committed
89
    server = server.PromptServer(loop)
90
    q = execution.PromptQueue(server)
91

pythongosssss's avatar
pythongosssss committed
92
93
94
    hijack_progress(server)

    threading.Thread(target=prompt_worker, daemon=True, args=(q,server,)).start()
95
96
97
98
    if '--listen' in sys.argv:
        address = '0.0.0.0'
    else:
        address = '127.0.0.1'
99

comfyanonymous's avatar
comfyanonymous committed
100
101
102
103
    dont_print = False
    if '--dont-print-server' in sys.argv:
        dont_print = True

104
105
106
107
108
109
110
111
112
    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 '--extra-model-paths-config' in sys.argv:
        indices = [(i + 1) for i in range(len(sys.argv) - 1) if sys.argv[i] == '--extra-model-paths-config']
        for i in indices:
            load_extra_path_config(sys.argv[i])

113
114
115
116
117
118
119
    port = 8188
    try:
        p_index = sys.argv.index('--port')
        port = int(sys.argv[p_index + 1])
    except:
        pass

120
121
122
    if '--quick-test-for-ci' in sys.argv:
        exit(0)

123
124
125
126
127
128
129
    call_on_start = None
    if "--windows-standalone-build" in sys.argv:
        def startup_server(address, port):
            import webbrowser
            webbrowser.open("http://{}:{}".format(address, port))
        call_on_start = startup_server

pythongosssss's avatar
pythongosssss committed
130
131
    if os.name == "nt":
        try:
132
            loop.run_until_complete(run(server, address=address, port=port, verbose=not dont_print, call_on_start=call_on_start))
pythongosssss's avatar
pythongosssss committed
133
134
135
        except KeyboardInterrupt:
            pass
    else:
136
        loop.run_until_complete(run(server, address=address, port=port, verbose=not dont_print, call_on_start=call_on_start))
comfyanonymous's avatar
comfyanonymous committed
137

138
    cleanup_temp()