Update app.py

app.py

@@ -34,7 +34,10 @@ class GlobalVars:
 
 g = GlobalVars()
 
-
+# 파일 상단에 추가
+torch.backends.cudnn.benchmark = False  # 메모리 사용량 감소
+torch.backends.cudnn.deterministic = True
+torch.cuda.set_per_process_memory_fraction(0.7)  # GPU 메모리 사용량 제한
         
 def initialize_models(device):
     try:
@@ -85,10 +88,6 @@ def initialize_models(device):
         print(f"Error during model initialization: {str(e)}")
         raise
     
-# CUDA 메모리 관리 설정
-torch.cuda.empty_cache()
-torch.backends.cuda.matmul.allow_tf32 = True
-torch.backends.cudnn.benchmark = True
 
 # 환경 변수 설정
 # 환경 변수 설정
@@ -104,6 +103,13 @@ os.environ['PYTORCH_NO_CUDA_MEMORY_CACHING'] = '1'
 # CUDA 초기화 방지
 torch.set_grad_enabled(False)
 
+def periodic_cleanup():
+    """주기적으로 실행될 메모리 정리 함수"""
+    clear_gpu_memory()
+    return None
+
+# Gradio 인터페이스에 주기적 정리 추가
+demo.load(periodic_cleanup, every=5)  # 5초마다 정리
 
 # Hugging Face 토큰 설정
 HF_TOKEN = os.getenv("HF_TOKEN")
@@ -209,123 +215,110 @@ def unpack_state(state: dict) -> Tuple[Gaussian, edict, str]:
     return gs, mesh, state['trial_id']
 
 @spaces.GPU
-def image_to_3d(trial_id: str, seed: int, randomize_seed: bool, ss_guidance_strength: float, ss_sampling_steps: int, slat_guidance_strength: float, slat_sampling_steps: int) -> Tuple[dict, str]:
-    print(f"Starting image_to_3d with trial_id: {trial_id}")
-    
-    if not trial_id or trial_id.strip() == "":
-        print("Error: No trial_id provided")
-        return None, None
-        
+def image_to_3d(trial_id: str, seed: int, randomize_seed: bool, ss_guidance_strength: float, 
+                ss_sampling_steps: int, slat_guidance_strength: float, slat_sampling_steps: int) -> Tuple[dict, str]:
     try:
-        # CUDA 메모리 초기화
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
-            torch.cuda.synchronize()
-        gc.collect()
+        # 초기 메모리 정리
+        clear_gpu_memory()
         
-        if randomize_seed:
-            seed = np.random.randint(0, MAX_SEED)
+        if not trial_id or trial_id.strip() == "":
+            return None, None
             
         image_path = f"{TMP_DIR}/{trial_id}.png"
-        print(f"Looking for image at: {image_path}")
-        
         if not os.path.exists(image_path):
-            print(f"Error: Image file not found at {image_path}")
             return None, None
             
         image = Image.open(image_path)
-        print(f"Successfully loaded image with size: {image.size}")
         
-        # 이미지 크기 제한
-        max_size = 512
+        # 이미지 크기 제한 강화
+        max_size = 384  # 더 작은 크기로 제한
         if max(image.size) > max_size:
             ratio = max_size / max(image.size)
             new_size = tuple(int(dim * ratio) for dim in image.size)
             image = image.resize(new_size, Image.LANCZOS)
-            print(f"Resized image to: {image.size}")
         
-        # GPU 작업 시작
         with torch.inference_mode():
             try:
-                # 모델을 GPU로 이동
+                # 파이프라인을 GPU로 이동
                 g.trellis_pipeline.to('cuda')
-                torch.cuda.synchronize()
                 
-                # 3D 생성
+                # 배치 크기 제한
                 outputs = g.trellis_pipeline.run(
                     image,
                     seed=seed,
                     formats=["gaussian", "mesh"],
                     preprocess_image=False,
                     sparse_structure_sampler_params={
-                        "steps": min(ss_sampling_steps, 12),
+                        "steps": min(ss_sampling_steps, 8),  # 스텝 수 제한
                         "cfg_strength": ss_guidance_strength,
+                        "batch_size": 1  # 배치 크기 명시적 제한
                     },
                     slat_sampler_params={
-                        "steps": min(slat_sampling_steps, 12),
+                        "steps": min(slat_sampling_steps, 8),  # 스텝 수 제한
                         "cfg_strength": slat_guidance_strength,
+                        "batch_size": 1
                     },
                 )
-                torch.cuda.synchronize()
                 
-                # 비디오 렌더링
+                # 중간 메모리 정리
+                clear_gpu_memory()
+                
+                # 비디오 렌더링 최적화
                 video = render_utils.render_video(
                     outputs['gaussian'][0],
-                    num_frames=60,
-                    resolution=512
+                    num_frames=30,  # 프레임 수 감소
+                    resolution=384  # 해상도 제한
                 )['color']
-                torch.cuda.synchronize()
                 
                 video_geo = render_utils.render_video(
                     outputs['mesh'][0],
-                    num_frames=60,
-                    resolution=512
+                    num_frames=30,
+                    resolution=384
                 )['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]
+                # CPU로 데이터 이동 및 메모리 정리
+                video = [v.cpu().numpy() for v in video]
+                video_geo = [v.cpu().numpy() for v in video_geo]
+                clear_gpu_memory()
                 
+                # 나머지 처리
                 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')
-                if torch.cuda.is_available():
-                    torch.cuda.empty_cache()
-                    torch.cuda.synchronize()
-                gc.collect()
-        
+                clear_gpu_memory()
+                
     except Exception as e:
         print(f"Error in image_to_3d: {str(e)}")
-        if hasattr(g.trellis_pipeline, 'to'):
-            g.trellis_pipeline.to('cpu')
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
-            torch.cuda.synchronize()
-        gc.collect()
+        g.trellis_pipeline.to('cpu')
+        clear_gpu_memory()
         return None, None
 
 def clear_gpu_memory():
-    """GPU 메모리를 정리하는 유틸리티 함수"""
+    """GPU 메모리를 더 철저하게 정리하는 함수"""
     try:
         if torch.cuda.is_available():
-            with torch.cuda.device('cuda'):
-                torch.cuda.empty_cache()
-                torch.cuda.synchronize()
+            # 모든 GPU 캐시 정리
+            torch.cuda.empty_cache()
+            torch.cuda.synchronize()
+            
+            # 사용하지 않는 캐시된 메모리 해제
+            for i in range(torch.cuda.device_count()):
+                with torch.cuda.device(i):
+                    torch.cuda.empty_cache()
+                    torch.cuda.ipc_collect()
+                    
+        # Python 가비지 컬렉터 실행
         gc.collect()
     except Exception as e:
-        print(f"Error clearing GPU memory: {e}")
+        print(f"Error in clear_gpu_memory: {e}")
 
 def move_to_device(model, device):
     """모델을 안전하게 디바이스로 이동하는 함수"""