Update app.py

app.py

@@ -1165,40 +1165,37 @@ def generate_image(prompt: str):
         if len(clean_prompt) < 3:
             return None, None
 
-        res = Client(IMAGE_API_URL).predict(
-            prompt      = clean_prompt,
-            width       = 768,
-            height      = 768,
-            guidance    = 3.5,
-            inference_steps = 30,
-            seed        = 3,
-            do_img2img  = False,
-            api_name    = "/generate_image",
-        )
+        # Add timeout and better logging
+        st.write(f"Connecting to image API at {IMAGE_API_URL}...")
+        
+        try:
+            res = Client(IMAGE_API_URL).predict(
+                prompt      = clean_prompt,
+                width       = 768,
+                height      = 768,
+                guidance    = 3.5,
+                inference_steps = 30,
+                seed        = 3,
+                do_img2img  = False,
+                api_name    = "/generate_image",
+                timeout     = 30  # Add timeout
+            )
+        except Exception as api_err:
+            st.error(f"Image API connection failed: {api_err}")
+            logging.error(f"Image API connection error: {api_err}")
+            return None, None
 
-        # ② 그레이디오/파스트API 구현에 따라 res 가 1-element일 수 있으므로 len 체크 제거
         if not res:
+            st.warning("Image API returned empty response")
             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: {e}", exc_info=True)
-        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...
 
 
 # ───────────────────────────── Kaggle API ─────────────────────────────
@@ -2343,22 +2340,42 @@ def process_input(prompt: str, uploaded_files):
 
             status.update(label="Invention ideas created!", state="complete")
 
-     
-# ── Auto image generation ─────────────────────────────
+# ── Auto image generation (Improved) ─────────────────────────────
             img_data = img_caption = None
             if st.session_state.generate_image and full_response:
-    # ① header(###)가 있든 없든, 한/영 표기를 모두 잡아낸다
-                match = re.search(
+                # Use multiple regex patterns to catch different image prompt formats
+                prompt_patterns = [
+                    # Original pattern
                     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…"):
-                        img_data, img_caption = generate_image(raw_prompt)
-                    if img_data:
-                        st.image(img_data, caption=f"Visualized Concept – {img_caption}")
+                    # Alternative formats
+                    r"(?:#{0,3}\s*)?(?:Image|이미지)(?:\s*prompt)?(?:[:：\-]\s*)?\n*([^\n]+)",
+                    r"\*\*(?:Image|이미지)(?:\s*prompt)?\*\*(?:[:：\-]\s*)?\n*([^\n]+)",
+                    r"[Ii]mage\s+[Pp]rompt:\s*([^\n]+)"
+                ]
+                
+                raw_prompt = None
+                for pattern in prompt_patterns:
+                    match = re.search(pattern, full_response, re.I)
+                    if match:
+                        raw_prompt = re.sub(r'[\r\n"\'\\]', " ", match.group(1)).strip()
+                        break
+                
+                if raw_prompt:
+                    with st.spinner(f"Generating image for concept: {raw_prompt[:50]}..."):
+                        try:
+                            img_data, img_caption = generate_image(raw_prompt)
+                            if img_data:
+                                st.image(img_data, caption=f"Visualized Concept – {img_caption}")
+                            else:
+                                st.warning("⚠️ Image generation produced no results. Using text-only output.")
+                                logging.warning(f"Image generation failed for prompt: {raw_prompt}")
+                        except Exception as img_err:
+                            st.error(f"⚠️ Image generation error: {str(img_err)[:100]}...")
+                            logging.error(f"Image generation error: {img_err}", exc_info=True)
+                else:
+                    logging.warning("No image prompt found in the response")
+                    if "debug" in st.session_state and st.session_state.debug:
+                        st.info("DEBUG: No image prompt pattern matched in the response.")
 
 
             answer_msg = {"role": "assistant", "content": full_response}