Update app_fine.py

app_fine.py

@@ -840,14 +840,31 @@ with demo:
     )
     
 
-# Launch the Gradio app
+# 파일 최하단 if __name__ == "__main__": 내부를 다음과 같이 수정하세요.
+
 if __name__ == "__main__":
     pipeline = TrellisVGGTTo3DPipeline.from_pretrained("Stable-X/trellis-vggt-v0-2")
-    # pipeline = TrellisVGGTTo3DPipeline.from_pretrained("weights/trellis-vggt-v0-1")
-    pipeline.cuda()
-    pipeline.VGGT_model.cuda()
-    pipeline.birefnet_model.cuda()
-    pipeline.dreamsim_model.cuda()
-    mast3r_model = AsymmetricMASt3R.from_pretrained("naver/MASt3R_ViTLarge_BaseDecoder_512_catmlpdpt_metric").cuda().eval() 
-    # mast3r_model = AsymmetricMASt3R.from_pretrained("weights/MAST3R/MASt3R_ViTLarge_BaseDecoder_512_catmlpdpt_metric.pth").cuda().eval()   
-    demo.launch()
+    
+    # --- 멀티 GPU 설정 시작 ---
+    num_gpus = torch.cuda.device_count()
+    print(f"사용 가능한 GPU 개수: {num_gpus}")
+    
+    # 메인 파이프라인 및 모델들을 멀티 GPU에 분산
+    if num_gpus > 1:
+        pipeline.cuda() # 우선 메인 로드
+        # 내부의 큰 모델들을 DataParallel로 감싸기
+        pipeline.VGGT_model = torch.nn.DataParallel(pipeline.VGGT_model)
+        pipeline.birefnet_model = torch.nn.DataParallel(pipeline.birefnet_model)
+        # pipeline.dreamsim_model = torch.nn.DataParallel(pipeline.dreamsim_model) # 필요시 해제
+    else:
+        pipeline.cuda()
+    
+    # MAST3R 모델 멀티 GPU 설정
+    mast3r_raw = AsymmetricMASt3R.from_pretrained("naver/MASt3R_ViTLarge_BaseDecoder_512_catmlpdpt_metric")
+    if num_gpus > 1:
+        mast3r_model = torch.nn.DataParallel(mast3r_raw).cuda().eval()
+    else:
+        mast3r_model = mast3r_raw.cuda().eval()
+    # --- 멀티 GPU 설정 끝 ---
+
+    demo.launch()