Update app.py

app.py

@@ -1,35 +1,32 @@
 import os
 import sys
 import subprocess
+import gc
+import shutil
+from typing import *
 
-# [AUTO-INSTALL] accelerate 라이브러리 자동 설치
+# [AUTO-INSTALL] accelerate
 try:
     import accelerate
 except ImportError:
-    print("⚠️ accelerate not found. Installing now...")
     subprocess.check_call([sys.executable, "-m", "pip", "install", "accelerate"])
-    print("✅ accelerate installed.")
 
-# [중요] OOM 방지 설정
+# [중요] OOM 방지 및 알고리즘 설정
 os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
 os.environ['SPCONV_ALGO'] = 'native'
 
-import gradio as gr
-from gradio_litmodel3d import LitModel3D
-import shutil
-from typing import *
 import torch
 import torch.nn as nn
 import numpy as np
 import imageio
-import gc
 from easydict import EasyDict as edict
 from PIL import Image
+import gradio as gr
+from gradio_litmodel3d import LitModel3D
+
 from trellis.pipelines import TrellisVGGTTo3DPipeline
 from trellis.representations import Gaussian, MeshExtractResult
 from trellis.utils import render_utils, postprocessing_utils
-
-# [수정] infer_auto_device_map 추가 임포트
 from accelerate import dispatch_model, infer_auto_device_map
 
 MAX_SEED = np.iinfo(np.int32).max
@@ -88,7 +85,6 @@ def pack_state(gs: Gaussian, mesh: MeshExtractResult) -> dict:
 
 def unpack_state(state: dict) -> Tuple[Gaussian, edict, str]:
     device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
-    
     gs = Gaussian(
         aabb=state['gaussian']['aabb'],
         sh_degree=state['gaussian']['sh_degree'],
@@ -150,7 +146,7 @@ def generate_and_extract_glb(
             )
     except Exception as e:
         torch.cuda.empty_cache()
-        raise RuntimeError(f"Generation Failed: {str(e)}\n(Try reducing image size or restart space)")
+        raise RuntimeError(f"Generation Failed: {str(e)}")
 
     video = render_utils.render_video(outputs['gaussian'][0], num_frames=120)['color']
     video_geo = render_utils.render_video(outputs['mesh'][0], num_frames=120)['normal']
@@ -177,7 +173,6 @@ def extract_gaussian(state: dict, req: gr.Request) -> Tuple[str, str]:
     gs, _ = unpack_state(state)
     gaussian_path = os.path.join(user_dir, 'sample.ply')
     gs.save_ply(gaussian_path)
-    
     del gs
     torch.cuda.empty_cache()
     return gaussian_path, gaussian_path
@@ -222,9 +217,7 @@ demo = gr.Blocks(
     """
 )
 with demo:
-    gr.Markdown("""
-    # 💻 ReconViaGen (Fixed Multi-GPU)
-    """)
+    gr.Markdown("# 💻 ReconViaGen (Corrected Multi-GPU)")
     
     with gr.Row():
         with gr.Column():
@@ -311,34 +304,41 @@ with demo:
 if __name__ == "__main__":
     print("🚀 Initializing Pipeline...")
     pipeline = TrellisVGGTTo3DPipeline.from_pretrained("esther11/trellis-vggt-v0-2")
+    
+    # [핵심 수정] 파이프라인 전체를 먼저 CUDA:0으로 이동
+    # 이렇게 해야 pipeline.device가 "cuda:0"으로 설정되고, 
+    # preprocess_image 호출 시 입력 텐서가 자동으로 GPU로 이동하여 birefnet(GPU0)과 만납니다.
+    pipeline.cuda() 
 
     gpu_count = torch.cuda.device_count()
     print(f"⚡ Detected {gpu_count} GPUs.")
 
     if gpu_count > 1:
-        print("⚡ Multi-GPU Mode Activated.")
+        print("⚡ Multi-GPU Mode: Distributing VGGT model.")
         
-        # [수정] slat_model 관련 코드 삭제 (존재하지 않음)
-        # VGGT_model만 분산 처리 (5GB 대용량 모델)
-        print("   - Calculating Device Map for VGGT Model...")
-        try:
-            # infer_auto_device_map을 통해 VGGT 모델의 레이어를 4개 GPU에 나눌 계획표(Map)를 짭니다.
-            vggt_map = infer_auto_device_map(pipeline.VGGT_model, max_memory={i: "22GiB" for i in range(gpu_count)})
-            
-            # 계획표대로 모델을 찢어서 각 GPU에 올립니다.
-            pipeline.VGGT_model = dispatch_model(pipeline.VGGT_model, device_map=vggt_map)
-            print("✅ VGGT Model dispatched successfully.")
-            
-        except Exception as e:
-            print(f"⚠️ Failed to dispatch VGGT model: {e}")
-            print("Falling back to single GPU for VGGT (Risky).")
-            pipeline.VGGT_model.cuda()
-
-        # 가벼운 모델은 0번 고정
-        pipeline.birefnet_model.to("cuda:0")
+        # [핵심 수정] VGGT 모델만 다시 CPU로 내려서 맵을 계산하거나, 
+        # 그냥 infer_auto_device_map에 맡깁니다.
+        # 이미 GPU0에 다 올라가 있으므로, 'max_memory' 제한을 주지 않으면 "GPU0에 다 들어가네?" 하고 분산을 안 할 수 있습니다.
+        # 따라서 VGGT를 강제로 찢기 위해 max_memory를 설정합니다.
+        
+        # 예: 각 GPU당 12GB 정도만 쓰도록 유도하여 분산을 강제함 (VGGT 5GB + 활성 메모리 고려)
+        # 혹은 infer_auto_device_map이 현재 할당된 메모리를 고려하도록 합니다.
+        
+        pipeline.VGGT_model.cpu() # 맵 계산을 위해 잠시 CPU로 이동 (확실한 분산을 위함)
+        
+        print("   - Calculating Device Map for VGGT...")
+        # 모든 GPU를 골고루 쓰도록 유도
+        device_map = infer_auto_device_map(
+            pipeline.VGGT_model, 
+            max_memory={i: "10GiB" for i in range(gpu_count)}, # 강제 분할을 위한 메모리 제한
+            no_split_module_classes=["Block", "ResnetBlock"] # 레이어 중간이 짤리지 않도록 보호
+        )
+        
+        pipeline.VGGT_model = dispatch_model(pipeline.VGGT_model, device_map=device_map)
+        print("✅ VGGT Model dispatched.")
         
     else:
         print("⚠️ Warning: Only 1 GPU detected.")
-        pipeline.cuda()
+        # 이미 위에서 pipeline.cuda() 했으므로 추가 작업 불필요
 
     demo.launch()