Update app.py

app.py

@@ -23,45 +23,39 @@ os.makedirs(TMP_DIR, exist_ok=True)
 def initialize_models():
     global pipeline, translator, flux_pipe
     
-    # CUDA 메모리 초기화
-    if torch.cuda.is_available():
-        torch.cuda.empty_cache()
-    
     try:
+        # GPU 메모리 초기화
+        torch.cuda.empty_cache()
+        
+        # GPU 사용 가능 여부 확인
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        
         # Trellis 파이프라인 초기화
-        pipeline = TrellisImageTo3DPipeline.from_pretrained(
-            "JeffreyXiang/TRELLIS-image-large",
-            device_map="auto"  # Zero GPU 환경에 맞게 자동 device 매핑
-        )
+        pipeline = TrellisImageTo3DPipeline.from_pretrained("JeffreyXiang/TRELLIS-image-large")
+        pipeline.to(device)
         
-        # 번역기 초기화 
+        # 번역기 초기화
         translator = translation_pipeline(
             "translation",
             model="Helsinki-NLP/opus-mt-ko-en",
-            device_map="auto"
+            device=0 if device=="cuda" else -1
         )
         
-        # Flux 파이프라인 초기화
+        # Flux 파이프라인 초기화 
         flux_pipe = FluxPipeline.from_pretrained(
             "black-forest-labs/FLUX.1-dev",
-            torch_dtype=torch.float16,  # bfloat16 대신 float16 사용
-            device_map="auto"
+            torch_dtype=torch.float16 if device=="cuda" else torch.float32
         )
         
-        # LoRA 가중치 로드
-        flux_pipe.load_lora_weights(
-            "gokaygokay/Flux-Game-Assets-LoRA-v2",
-            device_map="auto"
-        )
-        flux_pipe.fuse_lora(lora_scale=1.0)
+        if device == "cuda":
+            flux_pipe.enable_model_cpu_offload()
+        
+        return True
         
     except Exception as e:
-        print(f"Error initializing models: {str(e)}")
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
-        raise e
-
-
+        print(f"Model initialization error: {str(e)}")
+        torch.cuda.empty_cache()
+        return False
 
 def translate_if_korean(text):
     if any(ord('가') <= ord(char) <= ord('힣') for char in text):
@@ -116,33 +110,32 @@ 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, 
+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():
-            torch.cuda.empty_cache()
-            
+        torch.cuda.empty_cache()
+        
         if randomize_seed:
             seed = np.random.randint(0, MAX_SEED)
             
         input_image = Image.open(f"{TMP_DIR}/{trial_id}.png")
         
-        # 메모리 최적화를 위한 컨텍스트 매니저 사용
-        with torch.cuda.amp.autocast(enabled=True):
-            outputs = pipeline.run(
-                input_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,
-                }
-            )
+        with torch.cuda.amp.autocast(enabled=torch.cuda.is_available()):
+            with torch.no_grad():
+                outputs = pipeline.run(
+                    input_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,
+                    }
+                )
             
         # 비디오 렌더링
         video = render_utils.render_video(outputs['gaussian'][0], num_frames=120)['color']
@@ -163,8 +156,7 @@ def image_to_3d(trial_id: str, seed: int, randomize_seed: bool, ss_guidance_stre
         
     except Exception as e:
         print(f"Error in image_to_3d: {str(e)}")
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
+        torch.cuda.empty_cache()
         raise e
 
 @spaces.GPU
@@ -334,17 +326,25 @@ with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css) as demo:
     )
 
 if __name__ == "__main__":
-    # CUDA 메모리 캐시 초기화
-    torch.cuda.empty_cache()
+    # 초기 GPU 메모리 정리
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
     
-    # 모델 초기화
-    initialize_models()
+    # 모델 초기화 확인
+    if not initialize_models():
+        print("Failed to initialize models")
+        exit(1)
     
     try:
-        # rembg 사전 로드
-        pipeline.preprocess_image(Image.fromarray(np.zeros((512, 512, 3), dtype=np.uint8)))
-    except:
-        pass
+        # rembg 사전 로드 시도
+        test_image = Image.fromarray(np.zeros((256, 256, 3), dtype=np.uint8))
+        pipeline.preprocess_image(test_image)
+    except Exception as e:
+        print(f"Warning: Failed to preload rembg: {str(e)}")
     
     # Gradio 앱 실행
-    demo.launch(share=True)  # share=True 추가
+    demo.queue(concurrency_count=1).launch(
+        share=True,
+        enable_queue=True,
+        max_threads=1
+    )