Update app.py

app.py

@@ -32,47 +32,46 @@ MAX_SEED = np.iinfo(np.int32).max
 TMP_DIR = "/tmp/Trellis-demo"
 os.makedirs(TMP_DIR, exist_ok=True)
 
+def free_memory():
+    """메모리를 정리하는 유틸리티 함수"""
+    import gc
+    gc.collect()
+
+@spaces.GPU
+def free_gpu_memory():
+    """GPU 메모리를 정리하는 유틸리티 함수"""
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+
 def initialize_models():
     global pipeline, translator, flux_pipe
     
     try:
-        # Trellis 파이프라인 초기화 (더 강화된 메모리 최적화)
+        # Trellis 파이프라인 초기화 (CPU 모드로)
         pipeline = TrellisImageTo3DPipeline.from_pretrained(
             "JeffreyXiang/TRELLIS-image-large",
-            device_map="auto",
-            low_cpu_mem_usage=True,
-            torch_dtype=torch.float16  # 반정밀도 사용
+            low_cpu_mem_usage=True
         )
         
-        # 번역기 초기화 (더 작은 모델 사용)
+        # 번역기 초기화
         translator = translation_pipeline(
             "translation",
             model="Helsinki-NLP/opus-mt-ko-en",
             device="cpu",
-            model_kwargs={
-                "low_cpu_mem_usage": True,
-                "torch_dtype": torch.float16
-            }
+            model_kwargs={"low_cpu_mem_usage": True}
         )
         
-        # Flux 파이프라인 초기화 (메모리 최적화)
+        # Flux 파이프라인 초기화
         flux_pipe = FluxPipeline.from_pretrained(
             "black-forest-labs/FLUX.1-dev",
-            device_map="auto",
-            low_cpu_mem_usage=True,
-            torch_dtype=torch.float16,
-            variant="fp16"
+            low_cpu_mem_usage=True
         )
         
-        # 불필요한 캐시 정리
-        free_memory()
-        
         print("Models initialized successfully")
         return True
         
     except Exception as e:
         print(f"Model initialization error: {str(e)}")
-        free_memory()
         return False
 
 def translate_if_korean(text):
@@ -144,16 +143,15 @@ def unpack_state(state: dict) -> Tuple[Gaussian, edict, str]:
 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):
     try:
-        free_memory()
+        if torch.cuda.is_available():
+            pipeline.to("cuda")
+            pipeline.to(torch.float16)
         
         if randomize_seed:
             seed = np.random.randint(0, MAX_SEED)
             
         input_image = Image.open(f"{TMP_DIR}/{trial_id}.png")
         
-        # GPU 메모리 사용량 제한
-        torch.cuda.set_per_process_memory_fraction(0.6)
-        
         # 더 작은 이미지 크기 사용
         max_size = 512
         if max(input_image.size) > max_size:
@@ -181,7 +179,6 @@ def image_to_3d(trial_id: str, seed: int, randomize_seed: bool, ss_guidance_stre
                     }
                 )
         
-        # 더 적은 프레임으로 비디오 생성
         video = render_utils.render_video(outputs['gaussian'][0], num_frames=30)['color']
         video_geo = render_utils.render_video(outputs['mesh'][0], num_frames=30)['normal']
         video = [np.concatenate([video[i], video_geo[i]], axis=1) for i in range(len(video))]
@@ -193,12 +190,14 @@ def image_to_3d(trial_id: str, seed: int, randomize_seed: bool, ss_guidance_stre
         
         state = pack_state(outputs['gaussian'][0], outputs['mesh'][0], trial_id)
         
-        free_memory()
+        # CPU 모드로 돌아가기
+        pipeline.to("cpu")
+        
         return state, video_path
         
     except Exception as e:
         print(f"Error in image_to_3d: {str(e)}")
-        free_memory()
+        pipeline.to("cpu")
         raise e
 
 @spaces.GPU
@@ -256,23 +255,7 @@ footer {
     visibility: hidden;
 }
 """
-def free_memory():
-    """메모리를 정리하는 강화된 유틸리티 함수"""
-    import gc
-    import psutil
-    
-    # Python 가비지 컬렉션 강제 실행
-    gc.collect()
-    
-    # CUDA 메모리 정리
-    if torch.cuda.is_available():
-        torch.cuda.empty_cache()
-        torch.cuda.synchronize()
-    
-    # RAM 캐시 정리 시도
-    if psutil.POSIX:
-        import os
-        os.system('sync')
+
 # Gradio 인터페이스 정의
 with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css) as demo:
     gr.Markdown("""
@@ -390,6 +373,7 @@ with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css) as demo:
     )
 
 if __name__ == "__main__":
+    # CPU 메모리만 정리
     free_memory()
     
     # 모델 초기화
@@ -398,7 +382,7 @@ if __name__ == "__main__":
         exit(1)
     
     try:
-        # 최소 크기 이미지로 rembg 테스트
+        # rembg 사전 로드 시도 (작은 이미지로)
         test_image = Image.fromarray(np.ones((64, 64, 3), dtype=np.uint8) * 255)
         pipeline.preprocess_image(test_image)
     except Exception as e:
@@ -410,7 +394,5 @@ if __name__ == "__main__":
         max_threads=2,
         show_error=True,
         cache_examples=False,
-        enable_queue=True,
-        server_port=7860,
-        server_name="0.0.0.0"
+        enable_queue=True
     )