Update app.py

app.py

@@ -25,34 +25,32 @@ from PIL import Image, ImageDraw, ImageFont
 def clear_memory():
     """메모리 정리 함수"""
     gc.collect()
-    torch.cuda.empty_cache()
-    
-    # 사용하지 않는 캐시 정리
-    if hasattr(torch.cuda, 'empty_cache'):
-        torch.cuda.empty_cache()
-    if hasattr(torch.cuda, 'ipc_collect'):
-        torch.cuda.ipc_collect()
+    if torch.cuda.is_available():
+        try:
+            with torch.cuda.device('cuda'):
+                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, translator
     
     try:
-        clear_memory()
-        
-        # CPU에서만 실행되는 번역 모델 
-        model = AutoModelForSeq2SeqLM.from_pretrained(model_name, 
-            low_cpu_mem_usage=True).to('cpu')
+        # CPU에서 실행되는 번역 모델
+        model = AutoModelForSeq2SeqLM.from_pretrained(
+            model_name, 
+            low_cpu_mem_usage=True
+        ).to('cpu')
         tokenizer = AutoTokenizer.from_pretrained(model_name)
-        translator = pipeline("translation", model=model, tokenizer=tokenizer, 
-            device=-1)
+        translator = pipeline("translation", model=model, tokenizer=tokenizer, device=-1)
         
         # GroundingDINO 모델
         gd_processor = GroundingDinoProcessor.from_pretrained(gd_model_path)
         gd_model = GroundingDinoForObjectDetection.from_pretrained(
             gd_model_path,
             torch_dtype=torch.float16,
-            device_map='auto',
-            low_cpu_mem_usage=True
+            device_map=None  # device_map을 None으로 설정
         )
         
         # Segmenter
@@ -62,11 +60,10 @@ def initialize_models():
         pipe = FluxPipeline.from_pretrained(
             "black-forest-labs/FLUX.1-dev",
             torch_dtype=torch.float16,
-            device_map='auto',
+            device_map=None,  # device_map을 None으로 설정
             low_cpu_mem_usage=True
         )
         pipe.enable_attention_slicing()
-        pipe.enable_model_cpu_offload()
         
     except Exception as e:
         print(f"Model initialization error: {str(e)}")
@@ -601,8 +598,7 @@ if __name__ == "__main__":
             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)
     
@@ -615,7 +611,8 @@ if __name__ == "__main__":
         css=css,
         analytics_enabled=False,
         cache_examples=False
-    ) as demo:
+    ) as demo:    
+
         # HTML 헤더
         gr.HTML("""
             <div class="main-title">