Update app.py

app.py

@@ -13,28 +13,61 @@ from transformers import pipeline as hf_pipeline
 ##############################################################################
 # 1) ZeroGPU 환경 처리 + device, dtype 설정
 ##############################################################################
-# ZeroGPU 환경이면 무조건 CPU 사용, 아니면 cuda 가용 여부에 따라 결정
-if os.getenv("ZERO_GPU"):
-    device = "cpu"
-else:
+# ZeroGPU 초기화 시도
+try:
+    import zerogpu
+    zerogpu.init()
     device = "cuda" if torch.cuda.is_available() else "cpu"
+except ImportError:
+    # ZeroGPU가 설치되지 않은 경우
+    if os.getenv("ZERO_GPU"):
+        print("ZeroGPU environment variable is set but zerogpu package is not installed.")
+        device = "cpu"
+    else:
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+except Exception as e:
+    print(f"Error initializing ZeroGPU: {e}")
+    device = "cpu"
 
 # GPU일 때만 bfloat16, 그 외에는 float32
 dtype = torch.bfloat16 if device == "cuda" else torch.float32
 
+print(f"Using device: {device}, dtype: {dtype}")
+
 ##############################################################################
 # 2) 모델 로드: 번역 모델, DiffusionPipeline
 ##############################################################################
-translator = hf_pipeline(
-    "translation",
-    model="Helsinki-NLP/opus-mt-ko-en",
-    device=0 if device == "cuda" else -1
-)
+try:
+    translator = hf_pipeline(
+        "translation",
+        model="Helsinki-NLP/opus-mt-ko-en",
+        device=0 if device == "cuda" else -1
+    )
 
-pipe = DiffusionPipeline.from_pretrained(
-    "black-forest-labs/FLUX.1-schnell",
-    torch_dtype=dtype
-).to(device)
+    pipe = DiffusionPipeline.from_pretrained(
+        "black-forest-labs/FLUX.1-schnell",
+        torch_dtype=dtype
+    ).to(device)
+    
+    print("Models loaded successfully")
+except Exception as e:
+    print(f"Error loading models: {e}")
+    # 모델 로드 에러 처리를 위한 더미 함수들
+    def dummy_translator(text):
+        return [{'translation_text': text}]
+    
+    class DummyPipe:
+        def __call__(self, **kwargs):
+            from PIL import Image
+            import numpy as np
+            dummy_img = Image.fromarray(np.zeros((512, 512, 3), dtype=np.uint8))
+            class DummyResult:
+                def __init__(self, img):
+                    self.images = [img]
+            return DummyResult(dummy_img)
+    
+    translator = dummy_translator
+    pipe = DummyPipe()
 
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 2048