Update app.py

app.py

@@ -1158,48 +1158,63 @@ def process_uploaded_files(uploaded_files):
 
 # ───────────────────────────── Image Generation ──────────────────────
 def generate_image(prompt: str):
+    """이미지 생성 함수"""
     if not prompt:
         return None, None
+    
     try:
         clean_prompt = prompt.strip("\"'").strip()
         if len(clean_prompt) < 3:
             return None, None
 
-        # Client 객체 생성 시 timeout 설정 (선택사항)
-        client = Client(IMAGE_API_URL, timeout=30)  # 여기서 timeout 설정
+        # Gradio 클라이언트 생성 (timeout은 클라이언트 생성 시 설정)
+        client = Client(IMAGE_API_URL)
         
         try:
-            # predict 함수에서는 timeout 인자 제거
             res = client.predict(
-                prompt      = clean_prompt,
-                width       = 768,
-                height      = 768,
-                guidance    = 3.5,
+                prompt = clean_prompt,
+                width = 768,
+                height = 768,
+                guidance = 3.5,
                 inference_steps = 30,
-                seed        = 3,
-                do_img2img  = False,
-                api_name    = "/generate_image"
-                # timeout 매개변수 제거함
+                seed = 3,
+                do_img2img = False,
+                api_name = "/generate_image"
             )
         except Exception as api_err:
             st.error(f"이미지 API 연결 실패: {api_err}")
             logging.error(f"이미지 API 연결 오류: {api_err}")
             return None, None
 
-
+        if not res:
+            logging.warning("이미지 API가 빈 응답을 반환함")
+            return None, None
+        
+        # 응답 처리
+        raw = res[0] if isinstance(res, list) else res
         
-        # URL인 경우
+        # 응답 타입에 따른 처리
+        if isinstance(raw, str) and raw.startswith("data:image"):
+            raw = raw.split(",", 1)[1]
+            try:
+                import base64
+                img_bytes = base64.b64decode(raw)
+                return img_bytes, clean_prompt
+            except Exception as e:
+                logging.error(f"base64 디코딩 실패: {e}")
+                return None, None
         elif isinstance(raw, str) and raw.startswith("http"):
             return raw, clean_prompt
-        
-        # 바이너리 데이터인 경우
         else:
+            # 바이너리 데이터이거나 다른 형식일 경우
             return raw, clean_prompt
             
     except Exception as e:
         logging.error(f"이미지 생성 오류: {e}", exc_info=True)
         return None, None
 
+
+
 # ───────────────────────────── Kaggle API ─────────────────────────────
 @st.cache_resource
 def check_kaggle_availability():