Update app.py

app.py

@@ -313,39 +313,52 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
     )
 
 if __name__ == "__main__":
-    # 3D 생성 파이프라인
-    pipeline = TrellisImageTo3DPipeline.from_pretrained(
-        "JeffreyXiang/TRELLIS-image-large",
-        use_auth_token=HF_TOKEN
-    )
-    pipeline.to("cuda")  # .cuda() 대신 .to("cuda") 사용
-    
-    # 이미지 생성 파이프라인
-    pipe = FluxPipeline.from_pretrained(
-        "black-forest-labs/FLUX.1-dev",
-        torch_dtype=torch.bfloat16,
-        use_auth_token=HF_TOKEN,
-        device_map="auto"  # device_map 추가
-    )
-    
-    # Hyper-SD LoRA 로드
-    pipe.load_lora_weights(
-        hf_hub_download(
+    # CUDA 사용 가능 여부 확인
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"Using device: {device}")
+
+    try:
+        # 3D 생성 파이프라인
+        pipeline = TrellisImageTo3DPipeline.from_pretrained(
+            "JeffreyXiang/TRELLIS-image-large",
+            token=HF_TOKEN  # use_auth_token 대신 token 사용
+        )
+        pipeline.to(device)
+        
+        # 이미지 생성 파이프라인
+        pipe = FluxPipeline.from_pretrained(
+            "black-forest-labs/FLUX.1-dev",
+            torch_dtype=torch.bfloat16,
+            token=HF_TOKEN,  # use_auth_token 대신 token 사용
+            device_map="auto"
+        )
+        
+        # Hyper-SD LoRA 로드
+        lora_path = hf_hub_download(
             "ByteDance/Hyper-SD",
             "Hyper-FLUX.1-dev-8steps-lora.safetensors",
-            use_auth_token=HF_TOKEN
+            token=HF_TOKEN  # use_auth_token 대신 token 사용
         )
-    )
-    pipe.fuse_lora(lora_scale=0.125)
-    
-    # 번역기 초기화 (GPU 설정 추가)
-    translator = pipeline("translation", 
-                        model="Helsinki-NLP/opus-mt-ko-en",
-                        device=0 if torch.cuda.is_available() else -1)
-    
-    try:
-        pipeline.preprocess_image(Image.fromarray(np.zeros((512, 512, 3), dtype=np.uint8)))
-    except:
-        pass
+        pipe.load_lora_weights(lora_path)
+        pipe.fuse_lora(lora_scale=0.125)
+        
+        # 번역기 초기화
+        translator = pipeline(
+            "translation", 
+            model="Helsinki-NLP/opus-mt-ko-en",
+            device=0 if torch.cuda.is_available() else -1
+        )
+        
+        # 초기 이미지 전처리 테스트
+        try:
+            test_image = Image.fromarray(np.zeros((512, 512, 3), dtype=np.uint8))
+            pipeline.preprocess_image(test_image)
+        except Exception as e:
+            print(f"Warning: Initial preprocessing test failed: {e}")
+            
+        # Gradio 인터페이스 실행
+        demo.launch(allowed_paths=[PERSISTENT_DIR])
         
-    demo.launch(allowed_paths=[PERSISTENT_DIR])
+    except Exception as e:
+        print(f"Error during initialization: {e}")
+        raise