Update app.py

app.py

@@ -20,11 +20,12 @@ MAX_SEED = np.iinfo(np.int32).max
 TMP_DIR = "/tmp/Trellis-demo"
 os.makedirs(TMP_DIR, exist_ok=True)
 
-# GPU 메모리 관련 환경 변수 추가
-os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:512'
+# GPU 메모리 관련 환경 변수
+os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:128'  # 더 작은 값으로 설정
 os.environ['CUDA_VISIBLE_DEVICES'] = '0'
-os.environ['CUDA_LAUNCH_BLOCKING'] = '1'  # 디버깅을 위해 1로 설정
+os.environ['CUDA_LAUNCH_BLOCKING'] = '1'
 os.environ['PYTORCH_NO_CUDA_MEMORY_CACHING'] = '1'
+os.environ['CUDA_CACHE_DISABLE'] = '1'
 
 def initialize_models():
     global pipeline, translator, flux_pipe
@@ -38,13 +39,15 @@ def initialize_models():
         
         print("Initializing Trellis pipeline...")
         pipeline = TrellisImageTo3DPipeline.from_pretrained(
-            "JeffreyXiang/TRELLIS-image-large",
-            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
+            "JeffreyXiang/TRELLIS-image-large"
         )
         
         if torch.cuda.is_available():
             pipeline = pipeline.to("cuda")
-            pipeline.enable_model_cpu_offload()
+            # 메모리 최적화를 위한 설정
+            for param in pipeline.parameters():
+                if param.dtype == torch.float32:
+                    param.data = param.data.to(torch.float16)
         
         print("Initializing translator...")
         translator = translation_pipeline(
@@ -448,21 +451,15 @@ if __name__ == "__main__":
     # 메모리 정리
     free_memory()
     
-    # 모델 초기화 - 여러 번 시도
-    max_attempts = 3
-    for attempt in range(max_attempts):
-        if initialize_models():
-            break
-        print(f"Attempt {attempt + 1} failed, retrying...")
-        free_memory()
-    else:
-        print("Failed to initialize models after multiple attempts")
+    # 모델 초기화
+    if not initialize_models():
+        print("Failed to initialize models")
         exit(1)
     
     # Gradio 앱 실행
-    demo.queue(max_size=2).launch(
+    demo.queue(max_size=1).launch(
         share=True,
-        max_threads=4,
+        max_threads=2,  # 스레드 수 감소
         show_error=True,
         server_port=7860,
         server_name="0.0.0.0"