Update main.py

main.py

@@ -130,11 +130,13 @@ CHAT_ROLES = {
         "Habla de forma fluida y creativa. Si el usuario ya compartió algo antes, úsalo como base."
     ),
     "image_agent": (
-        "Eres ImageAgent, un especialista en imágenes y arte visual. "
-        "En este chat puedes describir imágenes, sugerir estilos visuales, hablar de fotografía, diseño y arte. "
-        "Si el usuario pide que GENERES o BUSQUES una imagen, responde con este JSON exacto y nada más: "
-        '{"action":"generate_image","queries":["english term 1","english term 2","english term 3"]} '
-        "Para todo lo demás, conversa normalmente como un experto en imágenes y artes visuales."
+        "Eres ImageAgent, especialista en imágenes y arte visual con IA. "
+        "REGLA CRÍTICA: Si el usuario pide generar, crear, buscar, dibujar o mostrar UNA imagen o foto de CUALQUIER cosa, "
+        "responde ÚNICAMENTE con este JSON sin ningún texto adicional: "
+        '{"action":"generate_image","queries":["detailed english prompt 1","detailed english prompt 2","english prompt 3"]} '
+        "Los queries deben ser DESCRIPTIVOS en inglés (colores, estilo, composición). Ej: "
+        '{"action":"generate_image","queries":["flying cat with angel wings pastel colors","cute cat flying through clouds digital art","cat with wings fantasy illustration"]} '
+        "Para preguntas sobre fotografía, diseño o arte SIN pedir imágenes, conversa normalmente."
     ),
 }
 
@@ -314,15 +316,37 @@ async def fetch_pexels(q):
     except: pass
     return None
 
-async def gen_hf_image(p):
-    if not HF_API_KEY: return None
-    try:
-        async with httpx.AsyncClient(timeout=60) as c:
-            r = await c.post("https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-xl-base-1.0",
-                headers={"Authorization":f"Bearer {HF_API_KEY}"},json={"inputs":p,"parameters":{"width":512,"height":384}})
-            if r.status_code==200 and "image" in r.headers.get("content-type",""):
-                return r.content
-    except: pass
+async def gen_hf_image(prompt: str) -> bytes | None:
+    """Try HuggingFace free image generation models. FLUX first (best quality)."""
+    if not HF_API_KEY:
+        return None
+    # Ordered by quality/availability — FLUX.1 schnell is fastest free model
+    models = [
+        ("black-forest-labs/FLUX.1-schnell",          {"inputs": prompt}),
+        ("black-forest-labs/FLUX.1-dev",               {"inputs": prompt}),
+        ("stabilityai/stable-diffusion-xl-base-1.0",  {"inputs": prompt, "parameters": {"width": 512, "height": 512}}),
+        ("stabilityai/stable-diffusion-2-1",           {"inputs": prompt, "parameters": {"width": 512, "height": 512}}),
+    ]
+    headers = {"Authorization": f"Bearer {HF_API_KEY}"}
+    for model_id, payload in models:
+        try:
+            async with httpx.AsyncClient(timeout=120) as c:
+                r = await c.post(
+                    f"https://api-inference.huggingface.co/models/{model_id}",
+                    headers=headers,
+                    json=payload,
+                )
+                ct = r.headers.get("content-type", "")
+                if r.status_code == 200 and ("image" in ct or r.content[:4] in (b"\xff\xd8\xff\xe0", b"\x89PNG")):
+                    return r.content
+                # 503/loading → try next model immediately
+                if r.status_code in (503, 500):
+                    continue
+                # 429 rate limit → stop trying HF
+                if r.status_code == 429:
+                    break
+        except Exception:
+            continue
     return None
 
 CREATIVE_KEYWORDS = [
@@ -338,9 +362,11 @@ CREATIVE_KEYWORDS = [
 def is_creative(q: str) -> bool:
     return any(k in q.lower() for k in CREATIVE_KEYWORDS)
 
-async def get_image(q: str):
-    # Creative/fantastical → AI first (Pexels won't have flying cats)
-    if is_creative(q):
+async def get_image(q: str, force_generate: bool = False):
+    """Fetch or generate an image for query q.
+    force_generate=True: skip Pexels, go straight to HF image generation.
+    """
+    if force_generate or is_creative(q):
         result = await gen_hf_image(q)
         if result: return result
     # Real-world → Pexels first, AI fallback
@@ -693,19 +719,20 @@ async def chat_with_agent(request: Request):
             parsed = _json2.loads(response.strip())
             if isinstance(parsed, dict) and parsed.get("action") == "generate_image":
                 queries = parsed.get("queries", [message])[:3]
-                _safe = re.sub(r'[^\w]', '_', message[:30])
-                _ts   = datetime.now().strftime('%Y%m%d_%H%M%S')
-                imgs  = await asyncio.gather(*[get_image(q) for q in queries])
+                _safe = re.sub(r'[^\w]', '_', message[:28])
+                _ts   = datetime.now().strftime('%H%M%S')
+                _base = f"chat_{_safe}_{_ts}"
+                imgs  = await asyncio.gather(*[get_image(q, force_generate=True) for q in queries])
                 for i, img in enumerate(imgs):
                     if img:
-                        fname = f"{_safe}_chat_img{i+1}.jpg"
+                        fname = f"{_base}_img{i+1}.jpg"
                         (DOCS_DIR / fname).write_bytes(img)
                         img_files.append(fname)
                 if img_files:
-                    response = f"Encontré {len(img_files)} imagen(es) para: {', '.join(queries[:3])}"
-                    img_result = {"img_base": f"{_safe}_chat", "img_count": len(img_files)}
+                    response = f"✓ {len(img_files)} imagen(es) generada(s) — haz clic en las miniaturas para verlas"
+                    img_result = {"img_base": _base, "img_count": len(img_files)}
                 else:
-                    response = "Lo siento, no pude encontrar/generar imágenes en este momento. Prueba con una descripción diferente."
+                    response = "No pude generar la imagen. Verifica HF_TOKEN y PEXELS_API_KEY en los Secrets del Space."
         except (ValueError, TypeError, KeyError):
             pass  # not JSON — normal text response