Update app.py

app.py

@@ -1164,31 +1164,43 @@ def generate_image(prompt: str):
         clean_prompt = prompt.strip("\"'").strip()
         if len(clean_prompt) < 3:
             return None, None
-        logging.info(f"Sending image generation request with prompt: {clean_prompt}")
 
         res = Client(IMAGE_API_URL).predict(
-            prompt=clean_prompt,
-            width=768,
-            height=768,
-            guidance=3.5,
-            inference_steps=30,
-            seed=3,
-            do_img2img=False,
-            init_image=None,
-            image2image_strength=0.8,
-            resize_img=True,
-            api_name="/generate_image"
+            prompt      = clean_prompt,
+            width       = 768,
+            height      = 768,
+            guidance    = 3.5,
+            inference_steps = 30,
+            seed        = 3,
+            do_img2img  = False,
+            api_name    = "/generate_image",
         )
-        if res and len(res) >= 2 and res[0]:
-            logging.info("Successfully received image data")
-            return res[0], clean_prompt
-        else:
-            logging.warning(f"Invalid response format from image API: {res}")
+
+        # ② 그레이디오/파스트API 구현에 따라 res 가 1-element일 수 있으므로 len 체크 제거
+        if not res:
+            logging.warning("Empty response from image API")
             return None, None
+
+        raw = res[0]  # URL, 파일경로, 또는 base64
+        if raw.startswith("http"):
+            return raw, clean_prompt
+
+        # base64일 경우 디코딩
+        if 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 decode failed: {e}")
+            return None, None
+
     except Exception as e:
-        logging.error(f"Image generation error: {str(e)}", exc_info=True)
+        logging.error(f"Image generation error: {e}", exc_info=True)
         return None, None
 
+
 # ───────────────────────────── Kaggle API ─────────────────────────────
 @st.cache_resource
 def check_kaggle_availability():
@@ -2331,12 +2343,16 @@ def process_input(prompt: str, uploaded_files):
 
             status.update(label="Invention ideas created!", state="complete")
 
-            # Auto image generation
+     
+# ── Auto image generation ─────────────────────────────
             img_data = img_caption = None
             if st.session_state.generate_image and full_response:
-                match = re.search(r"###\s*이미지\s*프롬프트\s*\n+([^\n]+)", full_response, re.I)
-                if not match:
-                    match = re.search(r"Image\s+Prompt\s*[:\-]\s*([^\n]+)", full_response, re.I)
+    # ① header(###)가 있든 없든, 한/영 표기를 모두 잡아낸다
+                match = re.search(
+                    r"(?:#{0,3}\s*)?(?:Image|이미지)\s*프롬프트\s*(?:[:：\-]\s*)?\n+([^\n]+)",
+                    full_response,
+                    re.I,
+                )
                 if match:
                     raw_prompt = re.sub(r'[\r\n"\'\\]', " ", match.group(1)).strip()
                     with st.spinner("Generating illustrative image…"):
@@ -2344,6 +2360,7 @@ def process_input(prompt: str, uploaded_files):
                     if img_data:
                         st.image(img_data, caption=f"Visualized Concept – {img_caption}")
 
+
             answer_msg = {"role": "assistant", "content": full_response}
             if img_data:
                 answer_msg["image"] = img_data