Update app.py

app.py

@@ -208,7 +208,11 @@ def image_to_3d(trial_id: str, seed: int, randomize_seed: bool, ss_guidance_stre
         return None, None
         
     try:
-        clear_gpu_memory()
+        # CUDA 메모리 초기화
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+            torch.cuda.synchronize()
+        gc.collect()
         
         if randomize_seed:
             seed = np.random.randint(0, MAX_SEED)
@@ -231,58 +235,77 @@ def image_to_3d(trial_id: str, seed: int, randomize_seed: bool, ss_guidance_stre
             image = image.resize(new_size, Image.LANCZOS)
             print(f"Resized image to: {image.size}")
         
-        with spaces.GPU(), torch.inference_mode():
-            # 3D 생성
-            g.trellis_pipeline.to('cuda')
-            outputs = g.trellis_pipeline.run(
-                image,
-                seed=seed,
-                formats=["gaussian", "mesh"],
-                preprocess_image=False,
-                sparse_structure_sampler_params={
-                    "steps": min(ss_sampling_steps, 12),
-                    "cfg_strength": ss_guidance_strength,
-                },
-                slat_sampler_params={
-                    "steps": min(slat_sampling_steps, 12),
-                    "cfg_strength": slat_guidance_strength,
-                },
-            )
-            
-            # 비디오 렌더링
-            video = render_utils.render_video(
-                outputs['gaussian'][0],
-                num_frames=60,
-                resolution=512
-            )['color']
-            
-            video_geo = render_utils.render_video(
-                outputs['mesh'][0],
-                num_frames=60,
-                resolution=512
-            )['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)
-            
-            state = pack_state(outputs['gaussian'][0], outputs['mesh'][0], new_trial_id)
-            
-            return state, video_path
-            
+        # GPU 작업 시작
+        with torch.inference_mode():
+            try:
+                # 모델을 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),
+                        "cfg_strength": ss_guidance_strength,
+                    },
+                    slat_sampler_params={
+                        "steps": min(slat_sampling_steps, 12),
+                        "cfg_strength": slat_guidance_strength,
+                    },
+                )
+                torch.cuda.synchronize()
+                
+                # 비디오 렌더링
+                video = render_utils.render_video(
+                    outputs['gaussian'][0],
+                    num_frames=60,
+                    resolution=512
+                )['color']
+                torch.cuda.synchronize()
+                
+                video_geo = render_utils.render_video(
+                    outputs['mesh'][0],
+                    num_frames=60,
+                    resolution=512
+                )['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')
+                if torch.cuda.is_available():
+                    torch.cuda.empty_cache()
+                    torch.cuda.synchronize()
+                gc.collect()
+        
     except Exception as e:
         print(f"Error in image_to_3d: {str(e)}")
-        return None, None
-    finally:
         if hasattr(g.trellis_pipeline, 'to'):
             g.trellis_pipeline.to('cpu')
-        clear_gpu_memory()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+            torch.cuda.synchronize()
+        gc.collect()
+        return None, None
 
 def clear_gpu_memory():
     """GPU 메모리를 정리하는 유틸리티 함수"""
@@ -326,7 +349,11 @@ def deactivate_button() -> gr.Button:
 @spaces.GPU
 def text_to_image(prompt: str, height: int, width: int, steps: int, scales: float, seed: int) -> Image.Image:
     try:
-        clear_gpu_memory()
+        # CUDA 메모리 초기화
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+            torch.cuda.synchronize()
+        gc.collect()
         
         # 한글 감지 및 번역
         def contains_korean(text):
@@ -343,7 +370,7 @@ def text_to_image(prompt: str, height: int, width: int, steps: int, scales: floa
         width = min(width, 512)
         steps = min(steps, 12)
         
-        with spaces.GPU(), torch.inference_mode():
+        with torch.inference_mode():
             generated_image = g.flux_pipe(
                 prompt=[formatted_prompt],
                 generator=torch.Generator('cuda').manual_seed(int(seed)),
@@ -368,7 +395,10 @@ def text_to_image(prompt: str, height: int, width: int, steps: int, scales: floa
         print(f"Error in image generation: {str(e)}")
         return None
     finally:
-        clear_gpu_memory()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+            torch.cuda.synchronize()
+        gc.collect()
 
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
     gr.Markdown("""## Craft3D""")