Update app.py

app.py

@@ -38,24 +38,34 @@ def initialize_models():
             torch.backends.cudnn.allow_tf32 = True
         
         print("Initializing Trellis pipeline...")
-        pipeline = TrellisImageTo3DPipeline.from_pretrained(
-            "JeffreyXiang/TRELLIS-image-large"
-        )
-        
-        if torch.cuda.is_available():
-            pipeline = pipeline.to("cuda")
-            # 메모리 최적화를 위한 설정
-            for param in pipeline.parameters():
-                if param.dtype == torch.float32:
-                    param.data = param.data.to(torch.float16)
+        try:
+            pipeline = TrellisImageTo3DPipeline.from_pretrained(
+                "JeffreyXiang/TRELLIS-image-large"
+            )
+            
+            if pipeline is None:
+                raise ValueError("Pipeline initialization returned None")
+                
+            if torch.cuda.is_available():
+                pipeline = pipeline.to("cuda")
+                # Half precision으로 변환
+                pipeline = pipeline.half()
+                
+        except Exception as e:
+            print(f"Error initializing Trellis pipeline: {str(e)}")
+            raise
         
         print("Initializing translator...")
-        translator = translation_pipeline(
-            "translation",
-            model="Helsinki-NLP/opus-mt-ko-en",
-            device=0 if torch.cuda.is_available() else -1
-        )
-        
+        try:
+            translator = translation_pipeline(
+                "translation",
+                model="Helsinki-NLP/opus-mt-ko-en",
+                device=0 if torch.cuda.is_available() else -1
+            )
+        except Exception as e:
+            print(f"Error initializing translator: {str(e)}")
+            raise
+            
         flux_pipe = None
         
         print("Models initialized successfully")
@@ -63,6 +73,7 @@ def initialize_models():
         
     except Exception as e:
         print(f"Model initialization error: {str(e)}")
+        free_memory()
         return False
 
 def get_flux_pipe():
@@ -94,10 +105,11 @@ def free_memory():
     # CUDA 메모리 정리
     if torch.cuda.is_available():
         torch.cuda.empty_cache()
+        torch.cuda.synchronize()
     
     # 임시 파일 정리
     tmp_dirs = ['/tmp/transformers_cache', '/tmp/torch_home', 
-                '/tmp/huggingface', '/tmp/cache']
+                '/tmp/huggingface', '/tmp/cache', TMP_DIR]
     
     for dir_path in tmp_dirs:
         if os.path.exists(dir_path):
@@ -444,6 +456,9 @@ if __name__ == "__main__":
     if torch.cuda.is_available():
         print(f"Using GPU: {torch.cuda.get_device_name()}")
         print(f"Available GPU memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.2f} GB")
+        
+        # CUDA 메모리 설정
+        torch.cuda.set_per_process_memory_fraction(0.8)  # GPU 메모리 사용량 제한
     
     # 디렉토리 생성
     os.makedirs(TMP_DIR, exist_ok=True)
@@ -459,8 +474,9 @@ if __name__ == "__main__":
     # Gradio 앱 실행
     demo.queue(max_size=1).launch(
         share=True,
-        max_threads=2,  # 스레드 수 감소
+        max_threads=2,
         show_error=True,
         server_port=7860,
-        server_name="0.0.0.0"
+        server_name="0.0.0.0",
+        enable_queue=True
     )