Update app.py

app.py

@@ -82,12 +82,13 @@ torch.backends.cudnn.benchmark = True
 
 # 환경 변수 설정
 # 환경 변수 설정
-os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:512"
+os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:512,garbage_collection_threshold:0.6"
 os.environ['SPCONV_ALGO'] = 'native'
 os.environ['SPARSE_BACKEND'] = 'native'
 os.environ['CUDA_LAUNCH_BLOCKING'] = '1'
 os.environ['XFORMERS_FORCE_DISABLE_TRITON'] = '1'
 os.environ['XFORMERS_ENABLE_FLASH_ATTENTION'] = '1'
+os.environ['TORCH_CUDA_MEMORY_ALLOCATOR'] = 'native'
 
 # CUDA 초기화 방지
 torch.set_grad_enabled(False)
@@ -205,6 +206,9 @@ def image_to_3d(trial_id: str, seed: int, randomize_seed: bool, ss_guidance_stre
         return None, None
         
     try:
+        # CUDA 메모리 초기화
+        torch.cuda.empty_cache()
+        
         if randomize_seed:
             seed = np.random.randint(0, MAX_SEED)
             
@@ -220,57 +224,82 @@ def image_to_3d(trial_id: str, seed: int, randomize_seed: bool, ss_guidance_stre
         
         # GPU 작업 시작
         with torch.cuda.device(0):
-            # 모델을 GPU로 이동
-            g.trellis_pipeline.to('cuda')
-            
-            with torch.inference_mode(), torch.cuda.amp.autocast():
-                outputs = g.trellis_pipeline.run(
-                    image,
-                    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,
-                    },
-                )
-
-                # 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']
-                
-                # CPU로 데이터 이동
-                video = [v.cpu().numpy() if torch.is_tensor(v) else v for v in video]
-                video_geo = [v.cpu().numpy() if torch.is_tensor(v) else v for v in video_geo]
-                
-                video = [np.concatenate([video[i], video_geo[i]], axis=1) for i in range(len(video))]
-                new_trial_id = str(uuid.uuid4())
-                video_path = f"{TMP_DIR}/{new_trial_id}.mp4"
-                os.makedirs(os.path.dirname(video_path), exist_ok=True)
-                imageio.mimsave(video_path, video, fps=15)
-                
-                # 상태 저장 전에 CPU로 이동
-                state = pack_state(outputs['gaussian'][0], outputs['mesh'][0], new_trial_id)
+            try:
+                # 모델을 GPU로 이동
+                g.trellis_pipeline.to('cuda')
+                torch.cuda.synchronize()
                 
-                # 모델을 다시 CPU로
+                with torch.inference_mode():
+                    # 첫 번째 단계: 3D 생성
+                    outputs = g.trellis_pipeline.run(
+                        image,
+                        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,
+                        },
+                    )
+                    torch.cuda.synchronize()
+                    
+                    # 두 번째 단계: 비디오 렌더링
+                    video = render_utils.render_video(outputs['gaussian'][0], num_frames=120)['color']
+                    torch.cuda.synchronize()
+                    
+                    video_geo = render_utils.render_video(outputs['mesh'][0], num_frames=120)['normal']
+                    torch.cuda.synchronize()
+                    
+                    # CPU로 데이터 이동 및 후처리
+                    video = [v.cpu().numpy() if torch.is_tensor(v) else v for v in video]
+                    video_geo = [v.cpu().numpy() if torch.is_tensor(v) else v for v in video_geo]
+                    
+                    video = [np.concatenate([video[i], video_geo[i]], axis=1) for i in range(len(video))]
+                    new_trial_id = str(uuid.uuid4())
+                    video_path = f"{TMP_DIR}/{new_trial_id}.mp4"
+                    os.makedirs(os.path.dirname(video_path), exist_ok=True)
+                    imageio.mimsave(video_path, video, fps=15)
+                    
+                    # 상태 저장
+                    state = pack_state(outputs['gaussian'][0], outputs['mesh'][0], new_trial_id)
+                    
+                    return state, video_path
+                    
+            finally:
+                # 정리 작업
                 g.trellis_pipeline.to('cpu')
                 torch.cuda.empty_cache()
-                
-                return state, video_path
+                torch.cuda.synchronize()
         
     except Exception as e:
         print(f"Error in image_to_3d: {str(e)}")
-        # 에러 발생 시에도 모델을 CPU로 되돌림
+        # 에러 발생 시 정리
         if hasattr(g.trellis_pipeline, 'to'):
             g.trellis_pipeline.to('cpu')
         torch.cuda.empty_cache()
+        torch.cuda.synchronize()
         return None, None
 
-
+def clear_gpu_memory():
+    """GPU 메모리를 정리하는 유틸리티 함수"""
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+        torch.cuda.synchronize()
+        
+def move_to_device(model, device):
+    """모델을 안전하게 디바이스로 이동하는 함수"""
+    try:
+        if hasattr(model, 'to'):
+            model.to(device)
+            if device == 'cuda':
+                torch.cuda.synchronize()
+    except Exception as e:
+        print(f"Error moving model to {device}: {str(e)}")
+        
 @spaces.GPU
 def extract_glb(state: dict, mesh_simplify: float, texture_size: int) -> Tuple[str, str]:
     gs, mesh, trial_id = unpack_state(state)