Update app.py

app.py

@@ -87,7 +87,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(
@@ -151,10 +150,8 @@ 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)")
 
-    # 렌더링은 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))]
@@ -321,23 +318,25 @@ if __name__ == "__main__":
     if gpu_count > 1:
         print("⚡ Multi-GPU Mode Activated.")
         
-        # [수정] infer_auto_device_map으로 맵(Dictionary)을 먼저 생성해야 함
-        # "balanced"는 문자열이 아니라 내부 동작 방식이므로, infer_auto_device_map을 통해 
-        # 실제 레이어별 GPU 할당표(dict)를 받아와야 dispatch_model이 알아먹습니다.
-        
+        # [수정] slat_model 관련 코드 삭제 (존재하지 않음)
+        # VGGT_model만 분산 처리 (5GB 대용량 모델)
         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)
-        
+        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")
         
-        print("✅ Models dispatched successfully.")
-        
     else:
         print("⚠️ Warning: Only 1 GPU detected.")
         pipeline.cuda()