Update app.py

app.py

@@ -22,9 +22,19 @@ from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
 import gc
 from PIL import Image, ImageDraw, ImageFont
 
-# 모델 초기화 부분 수정
+def clear_memory():
+    """메모리 정리 함수 - Spaces GPU 환경에 맞게 수정"""
+    gc.collect()
+    if torch.cuda.is_available():
+        try:
+            with torch.cuda.device('cuda:0'):  # 명시적으로 cuda:0 지정
+                torch.cuda.empty_cache()
+        except Exception as e:
+            print(f"GPU memory management warning: {e}")
+
 def initialize_models():
-    global segmenter, gd_model, gd_processor, pipe
+    """모델 초기화 함수 - Spaces GPU 환경에 맞게 수정"""
+    global segmenter, gd_model, gd_processor, pipe, translator
     
     try:
         # GPU 메모리 정리
@@ -39,14 +49,12 @@ def initialize_models():
         gd_processor = GroundingDinoProcessor.from_pretrained(gd_model_path)
         gd_model = GroundingDinoForObjectDetection.from_pretrained(
             gd_model_path,
-            torch_dtype=torch.float16,  # float32 대신 float16 사용
-            device_map="auto"  # 자동 디바이스 매핑
+            torch_dtype=torch.float16,
+            device_map='cuda:0'  # 명시적으로 cuda:0 지정
         )
         
         # Segmenter 초기화
-        segmenter = BoxSegmenter(device="cpu")
-        if torch.cuda.is_available():
-            segmenter.to(device)
+        segmenter = BoxSegmenter(device='cuda:0')  # 명시적으로 cuda:0 지정
         
         # FLUX 파이프라인 초기화
         pipe = FluxPipeline.from_pretrained(
@@ -61,27 +69,17 @@ def initialize_models():
             hf_hub_download(
                 "ByteDance/Hyper-SD",
                 "Hyper-FLUX.1-dev-8steps-lora.safetensors",
-                use_auth_token=HF_TOKEN
+                token=HF_TOKEN
             )
         )
         pipe.fuse_lora(lora_scale=0.125)
         
         if torch.cuda.is_available():
-            pipe = pipe.to("cuda:0")
+            pipe = pipe.to('cuda:0')  # 명시적으로 cuda:0 지정
             
     except Exception as e:
         print(f"Model initialization error: {str(e)}")
         raise
-        
-def clear_memory():
-    """메모리 정리 강화 함수"""
-    gc.collect()
-    torch.cuda.empty_cache()
-    
-    if torch.cuda.is_available():
-        with torch.cuda.device(0):
-            torch.cuda.reset_peak_memory_stats()
-            torch.cuda.empty_cache()
 # GPU 설정
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")  # 명시적으로 cuda:0 지정
 
@@ -605,8 +603,20 @@ def update_process_button(img, prompt):
         interactive=bool(img and prompt),
         variant="primary" if bool(img and prompt) else "secondary"
     )
-
 if __name__ == "__main__":
+    # CUDA 설정
+    if torch.cuda.is_available():
+        try:
+            torch.cuda.set_device('cuda:0')  # 명시적으로 cuda:0 설정
+            torch.backends.cudnn.benchmark = True
+            torch.backends.cuda.matmul.allow_tf32 = True
+        except Exception as e:
+            print(f"CUDA setup warning: {e}")
+    
+    # HF 토큰 설정
+    if HF_TOKEN:
+        login(token=HF_TOKEN, add_to_git_credential=False)
+    
     # 모델 초기화
     initialize_models()
     
@@ -814,7 +824,6 @@ if __name__ == "__main__":
             queue=True
         )
 
-    # Gradio 앱 실행
     demo.queue(max_size=3)
     demo.launch(
         server_name="0.0.0.0",