From 1679abd86d944521cad8a94a09d30fd5e238ae22 Mon Sep 17 00:00:00 2001
From: comfyanonymous <comfyanonymous@protonmail.com>
Date: Mon, 17 Jul 2023 11:00:14 -0400
Subject: [PATCH] Add a command line argument to enable backend:cudaMallocAsync

---
 comfy/cli_args.py         |  1 +
 comfy/model_management.py |  2 +-
 main.py                   | 10 +++++++++-
 3 files changed, 11 insertions(+), 2 deletions(-)

diff --git a/comfy/cli_args.py b/comfy/cli_args.py
index bef1868b9e6..9a388b03b58 100644
--- a/comfy/cli_args.py
+++ b/comfy/cli_args.py
@@ -40,6 +40,7 @@ def __call__(self, parser, namespace, values, option_string=None):
 parser.add_argument("--output-directory", type=str, default=None, help="Set the ComfyUI output directory.")
 parser.add_argument("--auto-launch", action="store_true", help="Automatically launch ComfyUI in the default 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.")
+parser.add_argument("--cuda-malloc", action="store_true", help="Enable cudaMallocAsync.")
 parser.add_argument("--dont-upcast-attention", action="store_true", help="Disable upcasting of attention. Can boost speed but increase the chances of black images.")
 
 fp_group = parser.add_mutually_exclusive_group()
diff --git a/comfy/model_management.py b/comfy/model_management.py
index 92c8ac842fb..69542cc371c 100644
--- a/comfy/model_management.py
+++ b/comfy/model_management.py
@@ -204,7 +204,7 @@ def is_nvidia():
 def get_torch_device_name(device):
     if hasattr(device, 'type'):
         if device.type == "cuda":
-            return "{} {}".format(device, torch.cuda.get_device_name(device))
+            return "{} {} : {}".format(device, torch.cuda.get_device_name(device), torch.cuda.get_allocator_backend())
         else:
             return "{}".format(device.type)
     else:
diff --git a/main.py b/main.py
index 802e4bfe4c4..a22545573ca 100644
--- a/main.py
+++ b/main.py
@@ -51,7 +51,6 @@ def execute_script(script_path):
 import gc
 
 from comfy.cli_args import args
-import comfy.utils
 
 if os.name == "nt":
     import logging
@@ -62,7 +61,16 @@ def execute_script(script_path):
         os.environ['CUDA_VISIBLE_DEVICES'] = str(args.cuda_device)
         print("Set cuda device to:", args.cuda_device)
 
+    if args.cuda_malloc:
+        env_var = os.environ.get('PYTORCH_CUDA_ALLOC_CONF', None)
+        if env_var is None:
+            env_var = "backend:cudaMallocAsync"
+        else:
+            env_var += ",backend:cudaMallocAsync"
+
+        os.environ['PYTORCH_CUDA_ALLOC_CONF'] = env_var
 
+import comfy.utils
 import yaml
 
 import execution