mirror of
https://git.datalinker.icu/comfyanonymous/ComfyUI
synced 2026-11-12 11:20:22 +08:00
Create Fix missing except and finally clauses in try blocks
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.
This commit is contained in:
Rahul Vadisetty
GitHub
parent
6730f3e1a3
commit
594e8cb368
257
Fix missing except and finally clauses in try blocks
Normal file
257
Fix missing except and finally clauses in try blocks
Normal file
@ -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
|
||||
Loading…
x
Reference in New Issue
Block a user