FilesExpand file tree

master

/

main.py

Copy path

Blame

More file actions

Blame

More file actions

Latest commit

 

History

History

History

192 lines (155 loc) · 6.65 KB

master

/

main.py

Copy path

Top

File metadata and controls

• Code
• Blame

192 lines (155 loc) · 6.65 KB

Raw

Copy raw file

Download raw file

Open symbols panel

Edit and raw actions

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

16

17

18

19

20

21

22

23

24

25

26

27

28

29

30

31

32

33

34

35

36

37

38

39

40

41

42

43

44

45

46

47

48

49

50

51

52

53

54

55

56

57

58

59

60

61

62

63

64

65

66

67

68

69

70

71

72

73

74

75

76

77

78

79

80

81

82

83

84

85

86

87

88

89

90

91

92

93

94

95

96

97

98

99

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

131

132

133

134

135

136

137

138

139

140

141

142

143

144

145

146

147

148

149

150

151

152

153

154

155

156

157

158

159

160

161

162

163

164

165

166

167

168

169

170

171

172

173

174

175

176

177

178

179

180

181

182

183

184

185

186

187

188

189

190

191

192

import os

import importlib.util

import folder_paths

import time

def execute_prestartup_script():

def execute_script(script_path):

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

node_paths = folder_paths.get_folder_paths("custom_nodes")

for custom_node_path in node_paths:

possible_modules = os.listdir(custom_node_path)

node_prestartup_times = []

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

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

execute_prestartup_script()

# Main code

import asyncio

import itertools

import shutil

import threading

import gc

from comfy.cli_args import args

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

import comfy.utils

import yaml

import execution

import server

from server import BinaryEventTypes

from nodes import init_custom_nodes

import comfy.model_management

def cuda_malloc_warning():

device = comfy.model_management.get_torch_device()

device_name = comfy.model_management.get_torch_device_name(device)

cuda_malloc_warning = False

if "cudaMallocAsync" in device_name:

for b in cuda_malloc.blacklist:

if b in device_name:

cuda_malloc_warning = True

if cuda_malloc_warning:

print("\nWARNING: this card most likely does not support cuda-malloc, if you get \"CUDA error\" please run ComfyUI with: --disable-cuda-malloc\n")

def prompt_worker(q, server):

e = execution.PromptExecutor(server)

while True:

item, item_id = q.get()

execution_start_time = time.perf_counter()

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)

print("Prompt executed in {:.2f} seconds".format(time.perf_counter() - execution_start_time))

gc.collect()

comfy.model_management.soft_empty_cache()

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

def hijack_progress(server):

def hook(value, total, preview_image):

comfy.model_management.throw_exception_if_processing_interrupted()

server.send_sync("progress", {"value": value, "max": total}, server.client_id)

if preview_image is not None:

server.send_sync(BinaryEventTypes.UNENCODED_PREVIEW_IMAGE, preview_image, server.client_id)

comfy.utils.set_progress_bar_global_hook(hook)

def cleanup_temp():

temp_dir = folder_paths.get_temp_directory()

if os.path.exists(temp_dir):

shutil.rmtree(temp_dir, ignore_errors=True)

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)

if __name__ == "__main__":

if args.temp_directory:

temp_dir = os.path.join(os.path.abspath(args.temp_directory), "temp")

print(f"Setting temp directory to: {temp_dir}")

folder_paths.set_temp_directory(temp_dir)

cleanup_temp()

loop = asyncio.new_event_loop()

asyncio.set_event_loop(loop)

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

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)

init_custom_nodes()

cuda_malloc_warning()

server.add_routes()

hijack_progress(server)

threading.Thread(target=prompt_worker, daemon=True, args=(q, server,)).start()

if args.output_directory:

output_dir = os.path.abspath(args.output_directory)

print(f"Setting output directory to: {output_dir}")

folder_paths.set_output_directory(output_dir)

if args.quick_test_for_ci:

exit(0)

call_on_start = None

if args.auto_launch:

def startup_server(address, port):

import webbrowser

if os.name == 'nt' and address == '0.0.0.0':

address = '127.0.0.1'

webbrowser.open(f"http://{address}:{port}")

call_on_start = startup_server

try:

loop.run_until_complete(run(server, address=args.listen, port=args.port, verbose=not args.dont_print_server, call_on_start=call_on_start))

except KeyboardInterrupt:

print("\nStopped server")

cleanup_temp()