Update app.py

app.py

@@ -29,8 +29,8 @@ from trellis.pipelines import TrellisVGGTTo3DPipeline
 from trellis.representations import Gaussian, MeshExtractResult
 from trellis.utils import render_utils, postprocessing_utils
 
-# accelerate 로드 (위에서 설치했으므로 통과됨)
-from accelerate import dispatch_model
+# [수정] infer_auto_device_map 추가 임포트
+from accelerate import dispatch_model, infer_auto_device_map
 
 MAX_SEED = np.iinfo(np.int32).max
 TMP_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'tmp')
@@ -87,7 +87,9 @@ 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'],
@@ -149,8 +151,10 @@ def generate_and_extract_glb(
             )
     except Exception as e:
         torch.cuda.empty_cache()
-        raise e
+        # 에러 메시지에 메모리 팁 추가
+        raise RuntimeError(f"Generation Failed: {str(e)}\n(Try reducing image size or restart space)")
 
+    # 렌더링은 CPU 혹은 0번 GPU에서 수행 (메모리 절약)
     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']
     video = [np.concatenate([video[i], video_geo[i]], axis=1) for i in range(len(video))]
@@ -222,7 +226,7 @@ demo = gr.Blocks(
 )
 with demo:
     gr.Markdown("""
-    # 💻 ReconViaGen (Multi-GPU Auto-Install)
+    # 💻 ReconViaGen (Fixed Multi-GPU)
     """)
     
     with gr.Row():
@@ -315,21 +319,24 @@ if __name__ == "__main__":
     print(f"⚡ Detected {gpu_count} GPUs.")
 
     if gpu_count > 1:
-        print("⚡ Multi-GPU Mode Activated: Distributing model across all available GPUs.")
+        print("⚡ Multi-GPU Mode Activated.")
         
-        pipeline.VGGT_model = dispatch_model(
-            pipeline.VGGT_model, 
-            device_map="balanced" 
-        )
+        # [수정] infer_auto_device_map으로 맵(Dictionary)을 먼저 생성해야 함
+        # "balanced"는 문자열이 아니라 내부 동작 방식이므로, infer_auto_device_map을 통해 
+        # 실제 레이어별 GPU 할당표(dict)를 받아와야 dispatch_model이 알아먹습니다.
         
-        pipeline.slat_model = dispatch_model(
-            pipeline.slat_model,
-            device_map="balanced"
-        )
+        print("   - Calculating Device Map for VGGT Model...")
+        vggt_map = infer_auto_device_map(pipeline.VGGT_model)
+        pipeline.VGGT_model = dispatch_model(pipeline.VGGT_model, device_map=vggt_map)
+        
+        print("   - Calculating Device Map for SLAT Model...")
+        slat_map = infer_auto_device_map(pipeline.slat_model)
+        pipeline.slat_model = dispatch_model(pipeline.slat_model, device_map=slat_map)
         
+        # 가벼운 모델은 0번 고정
         pipeline.birefnet_model.to("cuda:0")
         
-        print("✅ Model dispatched successfully via Accelerate.")
+        print("✅ Models dispatched successfully.")
         
     else:
         print("⚠️ Warning: Only 1 GPU detected.")