Update main.py

main.py

@@ -1,3 +1,16 @@
+# GPU 패치: NVIDIA 드라이버가 없으면 torch.cuda 관련 호출이 CPU 디바이스로 전환되도록 처리
+import torch
+if not torch.cuda.is_available():
+    # torch.cuda.current_device()가 호출되면 0을 반환하도록 패치하고,
+    # torch.device(정수)를 요청받으면 "cpu" 장치로 반환하도록 변경합니다.
+    torch.cuda.current_device = lambda: 0
+    original_torch_device = torch.device
+    def patched_torch_device(arg):
+        if isinstance(arg, int):
+            return original_torch_device("cpu")
+        return original_torch_device(arg)
+    torch.device = patched_torch_device
+
 import comfy.options
 comfy.options.enable_args_parsing()
 
@@ -8,10 +21,8 @@ import time
 from comfy.cli_args import args
 from app.logger import setup_logger
 
-
 setup_logger(log_level=args.verbose)
 
-
 def execute_prestartup_script():
     def execute_script(script_path):
         module_name = os.path.splitext(script_path)[0]
@@ -28,28 +39,41 @@ def execute_prestartup_script():
         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)
-        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__":
+            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 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])
+                elapsed = time.perf_counter() - time_before
+                node_prestartup_times.append((elapsed, module_path, success))
+    
+    if node_prestartup_times:
+        # Rich 라이브러리가 있다면 테이블 형식으로 출력합니다.
+        try:
+            from rich.console import Console
+            from rich.table import Table
+            console = Console()
+            table = Table(title="Prestartup Times for Custom Nodes")
+            table.add_column("Time (s)", justify="right")
+            table.add_column("Status")
+            table.add_column("Custom Node Path")
+            for elapsed, module_path, success in sorted(node_prestartup_times, key=lambda x: x[0]):
+                status = "[green]Success[/green]" if success else "[red]Failed[/red]"
+                table.add_row(f"{elapsed:.1f}", status, module_path)
+            console.print(table)
+        except ImportError:
+            # Rich 미설치 시 기존 print 방식 사용
+            print("\nPrestartup times for custom nodes:")
+            for elapsed, module_path, success in sorted(node_prestartup_times, key=lambda x: x[0]):
+                import_message = "" if success else " (PRESTARTUP FAILED)"
+                print("{:6.1f} seconds{}:".format(elapsed, import_message), module_path)
         print()
 
 execute_prestartup_script()
@@ -83,7 +107,7 @@ if __name__ == "__main__":
 if args.windows_standalone_build:
     try:
         import fix_torch
-    except:
+    except Exception:
         pass
 
 import comfy.utils
@@ -94,16 +118,20 @@ from server import BinaryEventTypes
 import nodes
 import comfy.model_management
 
+# NVIDIA GPU 드라이버가 없는 경우 경고 메시지를 남깁니다.
+if not torch.cuda.is_available():
+    logging.warning("No NVIDIA GPU driver found. Running in CPU mode. Performance may be degraded.")
+
 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
+    warning_needed = 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:
-            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")
+                warning_needed = True
+        if warning_needed:
+            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")
 
 def prompt_worker(q, server):
     e = execution.PromptExecutor(server, lru_size=args.cache_lru)
@@ -114,7 +142,7 @@ def prompt_worker(q, server):
     while True:
         timeout = 1000.0
         if need_gc:
-            timeout = max(gc_collect_interval - (current_time - last_gc_collect), 0.0)
+            timeout = max(gc_collect_interval - (time.perf_counter() - last_gc_collect), 0.0)
 
         queue_item = q.get(timeout=timeout)
         if queue_item is not None:
@@ -166,7 +194,6 @@ async def run(server, address='', port=8188, verbose=True, call_on_start=None):
         addresses.append((addr, port))
     await asyncio.gather(server.start_multi_address(addresses, call_on_start), server.publish_loop())
 
-
 def hijack_progress(server):
     def hook(value, total, preview_image):
         comfy.model_management.throw_exception_if_processing_interrupted()
@@ -177,13 +204,11 @@ def hijack_progress(server):
             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)
 
-
 if __name__ == "__main__":
     if args.temp_directory:
         temp_dir = os.path.join(os.path.abspath(args.temp_directory), "temp")
@@ -195,7 +220,7 @@ if __name__ == "__main__":
         try:
             import new_updater
             new_updater.update_windows_updater()
-        except:
+        except Exception:
             pass
 
     loop = asyncio.new_event_loop()
@@ -225,7 +250,7 @@ if __name__ == "__main__":
         logging.info(f"Setting output directory to: {output_dir}")
         folder_paths.set_output_directory(output_dir)
 
-    #These are the default folders that checkpoints, clip and vae models will be saved to when using CheckpointSave, etc.. nodes
+    # 기본 모델 저장 폴더 경로 설정
     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"))