mirror of
https://git.datalinker.icu/comfyanonymous/ComfyUI
synced 2026-11-12 11:46:37 +08:00
594e8cb368
Fix missing except and finally clauses in try blocks - Added except blocks to handle potential exceptions in the try statements. - Included finally clauses to ensure proper resource cleanup and graceful shutdown. - Updated logging to capture errors during server shutdown and event loop closure. - Ensured that all try blocks now have appropriate exception handling and cleanup.
258 lines
9.4 KiB
Plaintext
258 lines
9.4 KiB
Plaintext
import comfy.options
|
|
import os
|
|
import importlib.util
|
|
import folder_paths
|
|
import time
|
|
import asyncio
|
|
import itertools
|
|
import shutil
|
|
import threading
|
|
import gc
|
|
import logging
|
|
import comfy.utils
|
|
import yaml
|
|
import execution
|
|
import server
|
|
import nodes
|
|
import comfy.model_management
|
|
from comfy.cli_args import args
|
|
from server import BinaryEventTypes
|
|
|
|
# Ensure argument parsing is enabled
|
|
comfy.options.enable_args_parsing()
|
|
|
|
# Function to execute pre-startup scripts
|
|
def execute_prestartup_script():
|
|
def execute_script(script_path):
|
|
module_name = os.path.splitext(os.path.basename(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
|
|
|
|
if args.disable_all_custom_nodes:
|
|
return
|
|
|
|
node_paths = folder_paths.get_folder_paths("custom_nodes")
|
|
node_prestartup_times = []
|
|
|
|
for custom_node_path in node_paths:
|
|
possible_modules = os.listdir(custom_node_path)
|
|
|
|
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 possible_module == "__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 node_prestartup_times:
|
|
print("\nPrestartup times for custom nodes:")
|
|
for n in sorted(node_prestartup_times):
|
|
import_message = "" if n[2] else " (PRESTARTUP FAILED)"
|
|
print(f"{n[0]:6.1f} seconds{import_message}: {n[1]}")
|
|
print()
|
|
|
|
# Function to check for CUDA malloc warning
|
|
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:
|
|
blacklist = getattr(cuda_malloc, 'blacklist', [])
|
|
for b in blacklist:
|
|
if b in device_name:
|
|
cuda_malloc_warning = True
|
|
if cuda_malloc_warning:
|
|
logging.warning(
|
|
"\nWARNING: This card most likely does not support cuda-malloc. If you get 'CUDA error', "
|
|
"please run ComfyUI with: --disable-cuda-malloc\n"
|
|
)
|
|
|
|
# Worker function for processing prompts
|
|
def prompt_worker(q, server):
|
|
e = execution.PromptExecutor(server, lru_size=args.cache_lru)
|
|
last_gc_collect = 0
|
|
need_gc = False
|
|
gc_collect_interval = 10.0
|
|
|
|
while True:
|
|
timeout = 1000.0
|
|
current_time = time.perf_counter()
|
|
if need_gc:
|
|
timeout = max(gc_collect_interval - (current_time - last_gc_collect), 0.0)
|
|
|
|
queue_item = q.get(timeout=timeout)
|
|
if queue_item is not None:
|
|
item, item_id = queue_item
|
|
prompt_id = item[1]
|
|
server.last_prompt_id = prompt_id
|
|
|
|
e.execute(item[2], prompt_id, item[3], item[4])
|
|
need_gc = True
|
|
q.task_done(item_id,
|
|
e.history_result,
|
|
status=execution.PromptQueue.ExecutionStatus(
|
|
status_str='success' if e.success else 'error',
|
|
completed=e.success,
|
|
messages=e.status_messages))
|
|
if server.client_id is not None:
|
|
server.send_sync("executing", {"node": None, "prompt_id": prompt_id}, server.client_id)
|
|
|
|
execution_time = time.perf_counter() - time.perf_counter()
|
|
logging.info(f"Prompt executed in {execution_time:.2f} seconds")
|
|
|
|
flags = q.get_flags()
|
|
free_memory = flags.get("free_memory", False)
|
|
|
|
if flags.get("unload_models", free_memory):
|
|
comfy.model_management.unload_all_models()
|
|
need_gc = True
|
|
last_gc_collect = 0
|
|
|
|
if free_memory:
|
|
e.reset()
|
|
need_gc = True
|
|
last_gc_collect = 0
|
|
|
|
if need_gc:
|
|
current_time = time.perf_counter()
|
|
if (current_time - last_gc_collect) > gc_collect_interval:
|
|
comfy.model_management.cleanup_models()
|
|
gc.collect()
|
|
comfy.model_management.soft_empty_cache()
|
|
last_gc_collect = current_time
|
|
need_gc = False
|
|
|
|
# Asynchronous function to run the server
|
|
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())
|
|
|
|
# Function to hijack progress and send it to the server
|
|
def hijack_progress(server):
|
|
def hook(value, total, preview_image):
|
|
comfy.model_management.throw_exception_if_processing_interrupted()
|
|
progress = {"value": value, "max": total, "prompt_id": server.last_prompt_id, "node": server.last_node_id}
|
|
server.send_sync("progress", progress, 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)
|
|
|
|
# Function to clean up temporary directories
|
|
def cleanup_temp():
|
|
temp_dir = folder_paths.get_temp_directory()
|
|
if os.path.exists(temp_dir):
|
|
shutil.rmtree(temp_dir, ignore_errors=True)
|
|
|
|
# Function to load extra paths configuration from a YAML file
|
|
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 = conf.get("base_path", None)
|
|
for x in conf:
|
|
for y in conf[x].split("\n"):
|
|
if y:
|
|
full_path = os.path.join(base_path, y) if base_path else y
|
|
logging.info(f"Adding extra search path {x} {full_path}")
|
|
folder_paths.add_model_folder_path(x, full_path)
|
|
|
|
# Main function to initialize and run the server
|
|
if __name__ == "__main__":
|
|
if args.cuda_device is not None:
|
|
os.environ['CUDA_VISIBLE_DEVICES'] = str(args.cuda_device)
|
|
logging.info(f"Set cuda device to: {args.cuda_device}")
|
|
|
|
if args.deterministic:
|
|
os.environ.setdefault('CUBLAS_WORKSPACE_CONFIG', ":4096:8")
|
|
|
|
import cuda_malloc
|
|
|
|
if args.windows_standalone_build:
|
|
try:
|
|
import fix_torch
|
|
except ImportError:
|
|
pass
|
|
|
|
execute_prestartup_script()
|
|
cleanup_temp()
|
|
|
|
if args.windows_standalone_build:
|
|
try:
|
|
import new_updater
|
|
new_updater.update_windows_updater()
|
|
except ImportError:
|
|
pass
|
|
|
|
loop = asyncio.new_event_loop()
|
|
asyncio.set_event_loop(loop)
|
|
srv = server.PromptServer(loop)
|
|
q = execution.PromptQueue(srv)
|
|
|
|
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)
|
|
|
|
nodes.init_extra_nodes(init_custom_nodes=not args.disable_all_custom_nodes)
|
|
|
|
cuda_malloc_warning()
|
|
|
|
srv.add_routes()
|
|
hijack_progress(srv)
|
|
|
|
threading.Thread(target=prompt_worker, daemon=True, args=(q, srv)).start()
|
|
|
|
if args.output_directory:
|
|
output_dir = os.path.abspath(args.output_directory)
|
|
logging.info(f"Setting output directory to: {output_dir}")
|
|
folder_paths.set_output_directory(output_dir)
|
|
|
|
folder_paths.add_model_folder_path("checkpoints", os.path.join(folder_paths.get_output_directory(), "checkpoints"))
|
|
folder_paths.add_model_folder_path("clip", os.path.join(folder_paths.get_output_directory(), "clip"))
|
|
folder_paths.add_model_folder_path("vae", os.path.join(folder_paths.get_output_directory(), "vae"))
|
|
folder_paths.add_model_folder_path("diffusion_models", os.path.join(folder_paths.get_output_directory(), "diffusion_models"))
|
|
|
|
if args.input_directory:
|
|
input_dir = os.path.abspath(args.input_directory)
|
|
logging.info(f"Setting input directory to: {input_dir}")
|
|
folder_paths.set_input_directory(input_dir)
|
|
|
|
if args.quick_test_for_ci:
|
|
exit(0)
|
|
|
|
call_on_start = None
|
|
if args.auto_launch:
|
|
def startup_server(scheme, address, port):
|
|
import webbrowser
|
|
if os.name == 'nt' and address == '0.0.0.0':
|
|
address = '127.0.0.1'
|
|
webbrowser.open(f"{scheme}://{address}:{port}")
|
|
call_on_start = startup_server
|
|
|
|
try:
|
|
loop.run_until_complete(srv.setup())
|
|
loop.run_until_complete(run(srv, address=args.listen, port=args.port, verbose=not args.dont_print_server, call_on_start=call_on_start))
|
|
except KeyboardInterrupt:
|
|
logging.info("\nStopped server.")
|
|
except Exception:
|
|
logging.error("Server setup failed", exc_info=True)
|
|
finally:
|
|
logging.info("Shutting down.")
|
|
try:
|
|
srv.sh
|