Update app.py

app.py

@@ -1,28 +1,29 @@
+import os
+# [중요] 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 os
 import shutil
-os.environ['SPCONV_ALGO'] = 'native'
 from typing import *
 import torch
-import torch.nn as nn  # nn 모듈 추가
 import numpy as np
 import imageio
+import gc  # 가비지 컬렉션 추가
 from easydict import EasyDict as edict
 from PIL import Image
 from trellis.pipelines import TrellisVGGTTo3DPipeline
 from trellis.representations import Gaussian, MeshExtractResult
 from trellis.utils import render_utils, postprocessing_utils
 
-# --- [Fix] DataParallel 속성 접근 오류 해결을 위한 래퍼 클래스 ---
-class CustomDataParallel(nn.DataParallel):
-    def __getattr__(self, name):
-        try:
-            return super().__getattr__(name)
-        except AttributeError:
-            return getattr(self.module, name)
-# -----------------------------------------------------------
+# [중요] 모델 분산을 위한 accelerate 라이브러리 체크
+try:
+    from accelerate import dispatch_model
+    ACCELERATE_AVAILABLE = True
+except ImportError:
+    ACCELERATE_AVAILABLE = False
 
 MAX_SEED = np.iinfo(np.int32).max
 TMP_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'tmp')
@@ -37,6 +38,9 @@ def end_session(req: gr.Request):
     user_dir = os.path.join(TMP_DIR, str(req.session_hash))
     if os.path.exists(user_dir):
         shutil.rmtree(user_dir)
+    # 세션 종료 시 메모리 정리
+    gc.collect()
+    torch.cuda.empty_cache()
 
 def preprocess_image(image: Image.Image) -> Image.Image:
     """
@@ -118,6 +122,9 @@ def pack_state(gs: Gaussian, mesh: MeshExtractResult) -> dict:
     
     
 def unpack_state(state: dict) -> Tuple[Gaussian, edict, str]:
+    # 언팩 시 바로 CUDA로 올리면 메모리 튈 수 있으므로 상황에 맞게 device 설정
+    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    
     gs = Gaussian(
         aabb=state['gaussian']['aabb'],
         sh_degree=state['gaussian']['sh_degree'],
@@ -126,15 +133,15 @@ def unpack_state(state: dict) -> Tuple[Gaussian, edict, str]:
         opacity_bias=state['gaussian']['opacity_bias'],
         scaling_activation=state['gaussian']['scaling_activation'],
     )
-    gs._xyz = torch.tensor(state['gaussian']['_xyz'], device='cuda')
-    gs._features_dc = torch.tensor(state['gaussian']['_features_dc'], device='cuda')
-    gs._scaling = torch.tensor(state['gaussian']['_scaling'], device='cuda')
-    gs._rotation = torch.tensor(state['gaussian']['_rotation'], device='cuda')
-    gs._opacity = torch.tensor(state['gaussian']['_opacity'], device='cuda')
+    gs._xyz = torch.tensor(state['gaussian']['_xyz'], device=device)
+    gs._features_dc = torch.tensor(state['gaussian']['_features_dc'], device=device)
+    gs._scaling = torch.tensor(state['gaussian']['_scaling'], device=device)
+    gs._rotation = torch.tensor(state['gaussian']['_rotation'], device=device)
+    gs._opacity = torch.tensor(state['gaussian']['_opacity'], device=device)
     
     mesh = edict(
-        vertices=torch.tensor(state['mesh']['vertices'], device='cuda'),
-        faces=torch.tensor(state['mesh']['faces'], device='cuda'),
+        vertices=torch.tensor(state['mesh']['vertices'], device=device),
+        faces=torch.tensor(state['mesh']['faces'], device=device),
     )
     
     return gs, mesh
@@ -191,25 +198,31 @@ def generate_and_extract_glb(
         str: The path to the extracted GLB file.
         str: The path to the extracted GLB file (for download).
     """
+    # [수정] 추론 시작 전 가비지 컬렉션 수행
+    gc.collect()
+    torch.cuda.empty_cache()
+    
     user_dir = os.path.join(TMP_DIR, str(req.session_hash))
     image_files = [image[0] for image in multiimages]
 
     # Generate 3D model
-    outputs, _, _ = pipeline.run(
-        image=image_files,
-        seed=seed,
-        formats=["gaussian", "mesh"],
-        preprocess_image=False,
-        sparse_structure_sampler_params={
-            "steps": ss_sampling_steps,
-            "cfg_strength": ss_guidance_strength,
-        },
-        slat_sampler_params={
-            "steps": slat_sampling_steps,
-            "cfg_strength": slat_guidance_strength,
-        },
-        mode=multiimage_algo,
-    )
+    # [수정] torch.no_grad()를 사용하여 불필요한 그라디언트 메모리 사용 방지
+    with torch.no_grad():
+        outputs, _, _ = pipeline.run(
+            image=image_files,
+            seed=seed,
+            formats=["gaussian", "mesh"],
+            preprocess_image=False,
+            sparse_structure_sampler_params={
+                "steps": ss_sampling_steps,
+                "cfg_strength": ss_guidance_strength,
+            },
+            slat_sampler_params={
+                "steps": slat_sampling_steps,
+                "cfg_strength": slat_guidance_strength,
+            },
+            mode=multiimage_algo,
+        )
 
     # Render video
     # import uuid
@@ -234,7 +247,11 @@ def generate_and_extract_glb(
     # Pack state for optional Gaussian extraction
     state = pack_state(gs, mesh)
     
+    # [수정] 사용 끝난 변수 명시적 삭제 및 메모리 정리
+    del outputs, gs, mesh, glb
+    gc.collect()
     torch.cuda.empty_cache()
+    
     return state, video_path, glb_path, glb_path
 
 
@@ -257,12 +274,16 @@ 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
 
 
 def prepare_multi_example() -> List[Image.Image]:
-    # assets 경로 체크 추가 (에러 방지용)
+    # 에러 방지용 경로 체크
     if not os.path.exists("assets/example_multi_image"):
         return []
     multi_case = list(set([i.split('_')[0] for i in os.listdir("assets/example_multi_image")]))
@@ -445,19 +466,26 @@ if __name__ == "__main__":
     print("Initializing Pipeline...")
     pipeline = TrellisVGGTTo3DPipeline.from_pretrained("esther11/trellis-vggt-v0-2")
     pipeline.cuda()
-    pipeline.VGGT_model.cuda()
-    pipeline.birefnet_model.cuda()
 
-    # --- [Fix] Multi-GPU Logic ---
+    # [수정] 멀티 GPU 처리 로직 (Model Parallelism)
+    # 기존 DataParallel은 모델을 복제하여 VRAM을 2배로 쓰므로 OOM 발생함.
+    # 대신 accelerate의 dispatch_model을 사용하여 모델을 여러 GPU에 분할 적재해야 함.
     if torch.cuda.device_count() > 1:
-        print(f"⚡ Multi-GPU Detected: {torch.cuda.device_count()} GPUs found.")
-        print("Wrapping VGGT_model with CustomDataParallel to handle attributes correctly.")
-        
-        # VGGT_model을 CustomDataParallel로 감싸서 분산 처리
-        # CustomDataParallel은 'aggregator' 같은 속성을 module 내부에서 찾아줌
-        pipeline.VGGT_model = CustomDataParallel(pipeline.VGGT_model)
+        if ACCELERATE_AVAILABLE:
+            print(f"⚡ Accelerate detected: {torch.cuda.device_count()} GPUs found.")
+            print("Applying 'device_map=auto' to VGGT_model to split layers across GPUs (Memory Efficient).")
+            # VGGT_model이 가장 무거우므로 이를 여러 GPU에 쪼개서 올림
+            pipeline.VGGT_model = dispatch_model(pipeline.VGGT_model, device_map="auto")
+        else:
+            print("⚠️ 'accelerate' library not found. Cannot split model across GPUs.")
+            print("Installing 'accelerate' (`pip install accelerate`) is highly recommended for multi-GPU inference.")
+            # 라이브러리가 없으면 기본 동작(단일 GPU or 기존 상태) 유지
     else:
         print(f"Running on Single GPU: {torch.cuda.get_device_name(0)}")
-    # -----------------------------
 
+    # 나머지 모델들도 CUDA로 이동 (accelerate 적용 안된 경우)
+    if not ACCELERATE_AVAILABLE or torch.cuda.device_count() <= 1:
+        pipeline.VGGT_model.cuda()
+    pipeline.birefnet_model.cuda()
+    
     demo.launch()