diff --git a/Fix missing except and finally clauses in try blocks b/Fix missing except and finally clauses in try blocks new file mode 100644 index 000000000..0c8ae1c0c --- /dev/null +++ b/Fix missing except and finally clauses in try blocks @@ -0,0 +1,257 @@ +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