Update app.py

app.py

@@ -28,13 +28,15 @@ device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")  # 명
 
 ###--------------ZERO GPU 필수/ 메모리 관리 공통 --------------------###
 def clear_memory():
+    """메모리 정리 함수"""
     gc.collect()
     if torch.cuda.is_available():
         try:
-            torch.cuda.empty_cache()
-            torch.cuda.synchronize()
-        except:
-            pass
+            with torch.cuda.device('cuda:0'):
+                torch.cuda.empty_cache()
+                torch.cuda.synchronize()
+        except Exception as e:
+            print(f"Warning: Could not clear CUDA memory: {e}")
 
 ###---------------------------------------------------------------
 
@@ -92,21 +94,33 @@ gd_model = GroundingDinoForObjectDetection.from_pretrained(gd_model_path, torch_
 gd_model = gd_model.to(device=device)
 assert isinstance(gd_model, GroundingDinoForObjectDetection)
 
-# 파이프라인 초기화 부분도 수정
+# 파이프라인 초기화 및 최적화 설정
 pipe = FluxPipeline.from_pretrained(
     "black-forest-labs/FLUX.1-dev",
     torch_dtype=torch.float16,
     use_auth_token=HF_TOKEN
 )
 
-# 메모리 및 성능 최적화 설정
-pipe.enable_attention_slicing()
-pipe.enable_vae_slicing()
-pipe.enable_xformers_memory_efficient_attention()  # xformers가 설치되어 있다면
+# 메모리 최적화 설정 - FluxPipeline에서 지원하는 메서드만 사용
+pipe.enable_attention_slicing(slice_size="auto")
+
+# xformers 최적화 (설치되어 있는 경우에만)
+try:
+    import xformers
+    pipe.enable_xformers_memory_efficient_attention()
+except ImportError:
+    print("xformers is not installed. Skipping memory efficient attention.")
 
+# GPU 설정
 if torch.cuda.is_available():
-    pipe = pipe.to("cuda")
-    pipe.enable_model_cpu_offload()  # 메모리 효율을 위한 CPU 오프로드
+    try:
+        pipe = pipe.to("cuda:0")
+        # CPU 오프로딩이 지원되는 경우에만 활성화
+        if hasattr(pipe, 'enable_model_cpu_offload'):
+            pipe.enable_model_cpu_offload()
+    except Exception as e:
+        print(f"Warning: Could not move pipeline to CUDA: {str(e)}")
+
     
 
 # LoRA 가중치 로드