Spaces:
Sleeping
Sleeping
File size: 4,378 Bytes
d677097 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 | # agents/specialized/image_agent.py
import os
import re
import json
import base64
import requests
from pathlib import Path
from datetime import datetime
from agents.specialized.base import build_result, call_llm
from agents.registry import AGENT_REGISTRY
ROLE = """
Eres el agente de imágenes.
Cuando se te pida generar, crear, mostrar o visualizar una imagen, responde SOLO con un JSON:
{"image_queries": ["término en inglés 1", "término en inglés 2", "término en inglés 3"]}
Los términos deben ser específicos, detallados y en inglés para obtener mejores resultados.
Si la tarea no es sobre imágenes, responde: {"skip":"no image task"}
"""
DOCS_DIR = Path("data/docs")
async def generate_with_gemini(queries: list, model_name: str) -> list:
# ⚠️ pendiente de implementar con Gemini imagen real
return []
async def generate_with_hf(queries: list) -> list:
urls = []
for prompt in queries[:2]:
response = requests.post(
"https://router.huggingface.co/v1/images/generations",
headers={
"Authorization": f"Bearer {os.getenv('HF_API_TOKEN')}",
"Content-Type": "application/json",
},
json={
"model": "black-forest-labs/FLUX.1-schnell",
"prompt": prompt,
"size": "1024x1024",
},
)
if response.status_code == 200:
data = response.json()
# Algunos modelos devuelven b64_json, otros url directa
item = data["data"][0]
if "url" in item:
urls.append(item["url"])
elif "b64_json" in item:
image_bytes = base64.b64decode(item["b64_json"])
DOCS_DIR.mkdir(parents=True, exist_ok=True)
file_name = f"hf_image_{datetime.now().timestamp()}.png"
file_path = DOCS_DIR / file_name
with open(file_path, "wb") as f:
f.write(image_bytes)
urls.append(f"/docs/{file_name}")
else:
print("HF image error:", response.text)
return urls
async def run(task: str, context: dict = None) -> dict:
result = build_result("image_agent")
try:
response = await call_llm("image_agent", ROLE, task, context)
result["response"] = response
agent_config = AGENT_REGISTRY.get("image_agent", {})
provider = agent_config.get("provider")
model_name = agent_config.get("models", [None])[0]
print("USANDO PROVIDER:", provider)
print("USANDO MODELO:", model_name)
match = re.search(r'\{.*"image_queries".*\}', response, re.DOTALL)
if match:
try:
data = json.loads(match.group(0))
queries = data.get("image_queries", [])
result["queries"] = queries
result["response"] = "Ideas de imágenes generadas:\n" + "\n".join(queries)
image_urls = []
# Intento 1: provider principal
try:
if provider == "gemini":
image_urls = await generate_with_gemini(queries, model_name)
elif provider == "huggingface":
image_urls = await generate_with_hf(queries)
except Exception as e:
print("Provider principal falló:", str(e))
# Fallback a HF
if not image_urls:
try:
print("Intentando fallback HF...")
image_urls = await generate_with_hf(queries)
except Exception as e:
print("Fallback HF falló:", str(e))
# Fallback visual final
if not image_urls:
image_urls = ["https://picsum.photos/400/300"]
result["image_urls"] = image_urls
result["response"] += "\n\nImagen generada ✔" if image_urls else "\n\n⚠️ No se pudo generar imagen"
except Exception as e:
result["response"] = f"Error procesando JSON: {str(e)}"
else:
result["response"] = "No se detectaron instrucciones claras para imágenes."
except Exception as e:
result["success"] = False
result["error"] = str(e)
return result
|