Update app.py

app.py

@@ -1165,9 +1165,7 @@ def generate_image(prompt: str):
         if len(clean_prompt) < 3:
             return None, None
 
-        # Add timeout and better logging
-        st.write(f"Connecting to image API at {IMAGE_API_URL}...")
-        
+        # 여기서 API 호출 시도
         try:
             res = Client(IMAGE_API_URL).predict(
                 prompt      = clean_prompt,
@@ -1178,25 +1176,41 @@ def generate_image(prompt: str):
                 seed        = 3,
                 do_img2img  = False,
                 api_name    = "/generate_image",
-                timeout     = 30  # Add timeout
+                timeout     = 30
             )
         except Exception as api_err:
-            st.error(f"Image API connection failed: {api_err}")
-            logging.error(f"Image API connection error: {api_err}")
+            st.error(f"이미지 API 연결 실패: {api_err}")
+            logging.error(f"이미지 API 연결 오류: {api_err}")
             return None, None
 
         if not res:
-            st.warning("Image API returned empty response")
-            logging.warning("Empty response from image API")
+            logging.warning("이미지 API가 빈 응답을 반환함")
             return None, None
-
-        # Debug the response
-        st.write(f"API Response type: {type(res)}")
-        logging.info(f"Image API response: {res}")
-        
+            
         raw = res[0] if isinstance(res, list) else res
-        # Rest of your function...
-
+        
+        # 이미지 데이터 처리 (base64 인 경우)
+        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
+        
+        # URL인 경우
+        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