From 50bf66e5c44fe3637f29999034c10a0c083c7600 Mon Sep 17 00:00:00 2001
From: comfyanonymous <comfyanonymous@protonmail.com>
Date: Wed, 14 Aug 2024 02:49:25 -0400
Subject: [PATCH] Disable cuda malloc by default.

---
 comfy/cli_args.py | 4 ++--
 cuda_malloc.py    | 8 ++++++--
 2 files changed, 8 insertions(+), 4 deletions(-)

diff --git a/comfy/cli_args.py b/comfy/cli_args.py
index 2397de3d..a19d98e2 100644
--- a/comfy/cli_args.py
+++ b/comfy/cli_args.py
@@ -51,8 +51,8 @@ parser.add_argument("--auto-launch", action="store_true", help="Automatically la
 parser.add_argument("--disable-auto-launch", action="store_true", help="Disable auto launching the browser.")
 parser.add_argument("--cuda-device", type=int, default=None, metavar="DEVICE_ID", help="Set the id of the cuda device this instance will use.")
 cm_group = parser.add_mutually_exclusive_group()
-cm_group.add_argument("--cuda-malloc", action="store_true", help="Enable cudaMallocAsync (enabled by default for torch 2.0 and up).")
-cm_group.add_argument("--disable-cuda-malloc", action="store_true", help="Disable cudaMallocAsync.")
+cm_group.add_argument("--cuda-malloc", action="store_true", help="Enable cudaMallocAsync.")
+cm_group.add_argument("--disable-cuda-malloc", action="store_true", help="Disable cudaMallocAsync (The current default).")
 
 
 fp_group = parser.add_mutually_exclusive_group()
diff --git a/cuda_malloc.py b/cuda_malloc.py
index eb2857c5..04d42f91 100644
--- a/cuda_malloc.py
+++ b/cuda_malloc.py
@@ -2,6 +2,7 @@ import os
 import importlib.util
 from comfy.cli_args import args
 import subprocess
+import logging
 
 #Can't use pytorch to get the GPU names because the cuda malloc has to be set before the first import.
 def get_gpu_names():
@@ -63,7 +64,7 @@ def cuda_malloc_supported():
     return True
 
 
-if not args.cuda_malloc:
+if args.cuda_malloc:
     try:
         version = ""
         torch_spec = importlib.util.find_spec("torch")
@@ -74,8 +75,11 @@ if not args.cuda_malloc:
                 module = importlib.util.module_from_spec(spec)
                 spec.loader.exec_module(module)
                 version = module.__version__
+        supported = False
         if int(version[0]) >= 2: #enable by default for torch version 2.0 and up
-            args.cuda_malloc = cuda_malloc_supported()
+            supported = cuda_malloc_supported()
+        if not supported:
+            logging.warning("WARNING: cuda malloc enabled but not supported.")
     except:
         pass