Update app.py

app.py

@@ -1,9 +1,9 @@
 import gradio as gr
 import spaces
 from gradio_litmodel3d import LitModel3D
-
 import os
 import shutil
+
 os.environ['SPCONV_ALGO'] = 'native'
 from typing import *
 import torch
@@ -15,22 +15,19 @@ from trellis.pipelines import TrellisVGGTTo3DPipeline
 from trellis.representations import Gaussian, MeshExtractResult
 from trellis.utils import render_utils, postprocessing_utils
 
-
-
 MAX_SEED = np.iinfo(np.int32).max
 TMP_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'tmp')
-# TMP_DIR = "tmp/Trellis-demo"
-# os.environ['GRADIO_TEMP_DIR'] = 'tmp'
 os.makedirs(TMP_DIR, exist_ok=True)
 
 def start_session(req: gr.Request):
     user_dir = os.path.join(TMP_DIR, str(req.session_hash))
     os.makedirs(user_dir, exist_ok=True)
-    
-    
+
 def end_session(req: gr.Request):
     user_dir = os.path.join(TMP_DIR, str(req.session_hash))
-    shutil.rmtree(user_dir)
+    # 폴더가 존재할 때만 삭제하도록 수정 (FileNotFoundError 방지)
+    if os.path.exists(user_dir):
+        shutil.rmtree(user_dir)
 
 @spaces.GPU
 def preprocess_image(image: Image.Image) -> Image.Image:
@@ -436,29 +433,32 @@ with demo:
     )
     
 
-# 파일 최하단부 수정
 if __name__ == "__main__":
-    # 1. 파이프라인 로드
+    # 1. 모델 로드
     pipeline = TrellisVGGTTo3DPipeline.from_pretrained("Stable-X/trellis-vggt-v0-2")
     
-    # 2. 사용 가능한 GPU 개수 확인
-    num_gpus = torch.cuda.device_count()
-    print(f"시스템에서 감지된 GPU 개수: {num_gpus}")
-
-    if num_gpus > 1:
-        print("멀티 GPU 모드로 실행합니다.")
-        pipeline.cuda()
-        # 주요 모델들을 DataParallel로 감싸서 모든 GPU에 분산
-        pipeline.VGGT_model = torch.nn.DataParallel(pipeline.VGGT_model)
-        pipeline.birefnet_model = torch.nn.DataParallel(pipeline.birefnet_model)
-        # pipeline.dreamsim_model이 있는 경우 아래 주석 해제
-        # if hasattr(pipeline, 'dreamsim_model'):
-        #     pipeline.dreamsim_model = torch.nn.DataParallel(pipeline.dreamsim_model)
-    else:
-        print("단일 GPU 모드로 실행합니다.")
-        pipeline.cuda()
-        pipeline.VGGT_model.cuda()
-        pipeline.birefnet_model.cuda()
-
+    # 2. 멀티 GPU 강제 설정
+    if torch.cuda.is_available():
+        device = torch.device("cuda:0")
+        # 모든 하위 모델들을 먼저 0번 GPU로 확실히 보냅니다
+        pipeline.to(device)
+        
+        num_gpus = torch.cuda.device_count()
+        if num_gpus > 1:
+            print(f"--- 멀티 GPU 활성화: {num_gpus}개의 GPU를 사용합니다 ---")
+            # 에러가 났던 birefnet과 주요 모델들을 DataParallel로 래핑
+            try:
+                if hasattr(pipeline, 'VGGT_model'):
+                    pipeline.VGGT_model = torch.nn.DataParallel(pipeline.VGGT_model).cuda()
+                if hasattr(pipeline, 'birefnet_model'):
+                    pipeline.birefnet_model = torch.nn.DataParallel(pipeline.birefnet_model).cuda()
+                if hasattr(pipeline, 'sparse_structure_decoder'):
+                    pipeline.sparse_structure_decoder = torch.nn.DataParallel(pipeline.sparse_structure_decoder).cuda()
+                if hasattr(pipeline, 'slat_decoder'):
+                    pipeline.slat_decoder = torch.nn.DataParallel(pipeline.slat_decoder).cuda()
+            except Exception as e:
+                print(f"멀티 GPU 설정 중 경고 발생(단일 GPU로 전환): {e}")
+                pipeline.to(device)
+    
     # 3. 앱 실행
     demo.launch()