Update app.py

app.py

@@ -24,18 +24,23 @@ import gc
 def clear_memory():
     """메모리 정리 함수"""
     gc.collect()
-    if torch.cuda.is_available():
-        try:
-            torch.cuda.empty_cache()
-        except:
-            pass
+    try:
+        if torch.cuda.is_available():
+            with torch.cuda.device(0):  # 명시적으로 device 0 사용
+                torch.cuda.empty_cache()
+    except:
+        pass
+
+# GPU 설정
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")  # 명시적으로 cuda:0 지정
 
 # GPU 설정을 try-except로 감싸기
 if torch.cuda.is_available():
     try:
-        torch.cuda.empty_cache()
-        torch.backends.cudnn.benchmark = True
-        torch.backends.cuda.matmul.allow_tf32 = True
+        with torch.cuda.device(0):
+            torch.cuda.empty_cache()
+            torch.backends.cudnn.benchmark = True
+            torch.backends.cuda.matmul.allow_tf32 = True
     except:
         print("Warning: Could not configure CUDA settings")
 
@@ -62,8 +67,6 @@ BoundingBox = tuple[int, int, int, int]
 pillow_heif.register_heif_opener()
 pillow_heif.register_avif_opener()
 
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
 # HF 토큰 설정
 HF_TOKEN = os.getenv("HF_TOKEN")
 if HF_TOKEN is None:
@@ -106,9 +109,9 @@ pipe.fuse_lora(lora_scale=0.125)
 # GPU 설정을 try-except로 감싸기
 try:
     if torch.cuda.is_available():
-        pipe.to("cuda")
-except:
-    print("Warning: Could not move pipeline to CUDA")
+        pipe = pipe.to("cuda:0")  # 명시적으로 cuda:0 지정
+except Exception as e:
+    print(f"Warning: Could not move pipeline to CUDA: {str(e)}")
 
 class timer:
     def __init__(self, method_name="timed process"):