From 0cc81464bb9482195941baa42b3c88963345bfe3 Mon Sep 17 00:00:00 2001
From: Won-Kyu Park <wkpark@gmail.com>
Date: Sat, 12 Oct 2024 22:52:40 +0900
Subject: [PATCH] lint, add init_cuda_malloc()

---
 cuda_malloc.py | 28 ++++++++++++++--------------
 webui.py       |  3 ++-
 2 files changed, 16 insertions(+), 15 deletions(-)

diff --git a/cuda_malloc.py b/cuda_malloc.py
index 41bd1368ee8..fae5d73cfdc 100644
--- a/cuda_malloc.py
+++ b/cuda_malloc.py
@@ -39,9 +39,9 @@ def enum_display_devices():
     else:
         gpu_names = set()
         out = subprocess.check_output(['nvidia-smi', '-L'])
-        for l in out.split(b'\n'):
-            if len(l) > 0:
-                gpu_names.add(l.decode('utf-8').split(' (UUID')[0])
+        for line in out.split(b'\n'):
+            if len(line) > 0:
+                gpu_names.add(line.decode('utf-8').split(' (UUID')[0])
         return gpu_names
 
 blacklist = {"GeForce GTX TITAN X", "GeForce GTX 980", "GeForce GTX 970", "GeForce GTX 960", "GeForce GTX 950", "GeForce 945M",
@@ -55,7 +55,7 @@ def enum_display_devices():
 def cuda_malloc_supported():
     try:
         names = get_gpu_names()
-    except:
+    except Exception:
         names = set()
     for x in names:
         if "NVIDIA" in x:
@@ -82,16 +82,16 @@ def cuda_malloc_supported():
                 version = module.__version__
         if int(version[0]) >= 2: #enable by default for torch version 2.0 and up
             args.cuda_malloc = cuda_malloc_supported()
-    except:
+    except Exception:
         pass
 
+def init_cuda_malloc():
+    if args.cuda_malloc and not args.disable_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"
 
-if args.cuda_malloc and not args.disable_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
-    print(f"Setup environment PYTORCH_CUDA_ALLOC_CONF={env_var}")
+        os.environ['PYTORCH_CUDA_ALLOC_CONF'] = env_var
+        print(f"Setup environment PYTORCH_CUDA_ALLOC_CONF={env_var}")
diff --git a/webui.py b/webui.py
index 0567668d8b4..aef977f7c20 100644
--- a/webui.py
+++ b/webui.py
@@ -10,7 +10,8 @@
 startup_timer = timer.startup_timer
 startup_timer.record("launcher")
 
-import cuda_malloc
+from cuda_malloc import init_cuda_malloc
+init_cuda_malloc()
 startup_timer.record("cuda_malloc")
 
 initialize.imports()