Update app.py

app.py

@@ -32,39 +32,26 @@ 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 파이프라인 초기화 (CPU 모드로)
+        # Trellis 파이프라인 초기화 (기본 CPU 모드)
         pipeline = TrellisImageTo3DPipeline.from_pretrained(
-            "JeffreyXiang/TRELLIS-image-large",
-            low_cpu_mem_usage=True
+            "JeffreyXiang/TRELLIS-image-large"
         )
         
         # 번역기 초기화
         translator = translation_pipeline(
             "translation",
             model="Helsinki-NLP/opus-mt-ko-en",
-            device="cpu",
-            model_kwargs={"low_cpu_mem_usage": True}
+            device="cpu"
         )
         
         # Flux 파이프라인 초기화
         flux_pipe = FluxPipeline.from_pretrained(
             "black-forest-labs/FLUX.1-dev",
-            low_cpu_mem_usage=True
+            torch_dtype=torch.float32  # CPU 모드
         )
         
         print("Models initialized successfully")
@@ -74,6 +61,42 @@ def initialize_models():
         print(f"Model initialization error: {str(e)}")
         return False
 
+def free_memory():
+    """메모리를 정리하는 유틸리티 함수"""
+    import gc
+    gc.collect()
+
+@spaces.GPU
+def setup_gpu_model(model):
+    """GPU 설정이 필요한 모델을 처리하는 함수"""
+    if torch.cuda.is_available():
+        model = model.to("cuda")
+    return model
+
+if __name__ == "__main__":
+    # CPU 메모리만 정리
+    free_memory()
+    
+    # 모델 초기화
+    if not initialize_models():
+        print("Failed to initialize models")
+        exit(1)
+    
+    try:
+        # rembg 사전 로드 시도 (작은 이미지로)
+        test_image = Image.fromarray(np.ones((64, 64, 3), dtype=np.uint8) * 255)
+        pipeline.preprocess_image(test_image)
+    except Exception as e:
+        print(f"Warning: Failed to preload rembg: {str(e)}")
+    
+    # Gradio 앱 실행
+    demo.queue(max_size=5).launch(
+        share=True,
+        max_threads=2,
+        show_error=True,
+        cache_examples=False
+    )
+
 def translate_if_korean(text):
     if any(ord('가') <= ord(char) <= ord('힣') for char in text):
         translated = translator(text)[0]['translation_text']
@@ -143,10 +166,6 @@ 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:
-        if torch.cuda.is_available():
-            pipeline.to("cuda")
-            pipeline.to(torch.float16)
-        
         if randomize_seed:
             seed = np.random.randint(0, MAX_SEED)
             
@@ -162,22 +181,24 @@ def image_to_3d(trial_id: str, seed: int, randomize_seed: bool, ss_guidance_stre
                 Image.LANCZOS
             )
         
-        with torch.cuda.amp.autocast():
-            with torch.no_grad():
-                outputs = pipeline.run(
-                    input_image,
-                    seed=seed,
-                    formats=["gaussian", "mesh"],
-                    preprocess_image=False,
-                    sparse_structure_sampler_params={
-                        "steps": min(ss_sampling_steps, 15),
-                        "cfg_strength": ss_guidance_strength,
-                    },
-                    slat_sampler_params={
-                        "steps": min(slat_sampling_steps, 15),
-                        "cfg_strength": slat_guidance_strength,
-                    }
-                )
+        if torch.cuda.is_available():
+            pipeline.to("cuda")
+        
+        with torch.no_grad():
+            outputs = pipeline.run(
+                input_image,
+                seed=seed,
+                formats=["gaussian", "mesh"],
+                preprocess_image=False,
+                sparse_structure_sampler_params={
+                    "steps": min(ss_sampling_steps, 15),
+                    "cfg_strength": ss_guidance_strength,
+                },
+                slat_sampler_params={
+                    "steps": min(slat_sampling_steps, 15),
+                    "cfg_strength": slat_guidance_strength,
+                }
+            )
         
         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']
@@ -190,14 +211,15 @@ 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)
         
-        # CPU 모드로 돌아가기
-        pipeline.to("cpu")
+        if torch.cuda.is_available():
+            pipeline.to("cpu")
         
         return state, video_path
         
     except Exception as e:
         print(f"Error in image_to_3d: {str(e)}")
-        pipeline.to("cpu")
+        if torch.cuda.is_available():
+            pipeline.to("cpu")
         raise e
 
 @spaces.GPU