# agents/specialized/image_agent.py # # Cascade de generación de imágenes: # 1. HuggingFace FLUX.1-schnell (genera imagen) # → sube a HF Dataset vfven/mission-control-images (persistente) # 2. Gemini Imagen API (pendiente, requiere billing) # 3. Picsum Photos (fallback visual garantizado) import os import re import json import asyncio from pathlib import Path from datetime import datetime from io import BytesIO 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") HF_MODEL = "black-forest-labs/FLUX.1-schnell" HF_DATASET = "vfven/mission-control-images" IMAGE_PROVIDERS = ["fal-ai", "hf-inference", "replicate", "together"] # ───────────────────────────────────────────────────────────────────────────── # Subir imagen a HF Dataset (persistente) # ───────────────────────────────────────────────────────────────────────────── async def _upload_to_dataset(image, file_name: str) -> str | None: """ Sube una imagen PIL al dataset vfven/mission-control-images. Devuelve la URL pública permanente o None si falla. El dataset se crea automáticamente si no existe. """ try: from huggingface_hub import HfApi import inspect token = os.getenv("HF_API_TOKEN") api = HfApi(token=token) # Crear dataset si no existe try: api.create_repo( repo_id=HF_DATASET, repo_type="dataset", private=False, exist_ok=True, ) print(f"[Dataset] repo {HF_DATASET} listo") except Exception as e: print(f"[Dataset] advertencia al crear repo: {e}") # Serializar imagen PIL → bytes en memoria buf = BytesIO() await asyncio.to_thread(image.save, buf, format="PNG") buf.seek(0) # Subir al dataset path_in_repo = f"images/{file_name}" # upload_file acepta file-like object kwargs = dict( path_or_fileobj=buf, path_in_repo=path_in_repo, repo_id=HF_DATASET, repo_type="dataset", token=token, ) # commit_message solo en versiones que lo soporten if "commit_message" in inspect.signature(api.upload_file).parameters: kwargs["commit_message"] = f"Add {file_name}" await asyncio.to_thread(api.upload_file, **kwargs) # URL pública del dataset url = f"https://huggingface.co/datasets/{HF_DATASET}/resolve/main/{path_in_repo}" print(f"[Dataset] subida OK → {url}") return url except Exception as e: print(f"[Dataset] error subiendo imagen: {e}") return None # ───────────────────────────────────────────────────────────────────────────── # 1. HuggingFace FLUX # ───────────────────────────────────────────────────────────────────────────── async def _generate_with_hf(queries: list) -> list: import inspect from huggingface_hub import InferenceClient token = os.getenv("HF_API_TOKEN") urls = [] supports_provider = "provider" in inspect.signature(InferenceClient.__init__).parameters for prompt in queries[:2]: generated = False if supports_provider: for prov in IMAGE_PROVIDERS: try: print(f"[HF] provider={prov} → '{prompt[:45]}...'") client = InferenceClient(provider=prov, api_key=token) image = await asyncio.to_thread( client.text_to_image, prompt, model=HF_MODEL ) file_name = f"hf_{datetime.now().strftime('%Y%m%d_%H%M%S_%f')}.png" # Intentar subir al dataset (persistente) dataset_url = await _upload_to_dataset(image, file_name) if dataset_url: urls.append(dataset_url) print(f"[HF] imagen en dataset → {dataset_url}") else: # Fallback: guardar local (efímero) DOCS_DIR.mkdir(parents=True, exist_ok=True) file_path = DOCS_DIR / file_name await asyncio.to_thread(image.save, str(file_path)) urls.append(f"/docs/{file_name}") print(f"[HF] imagen local (efímera) → {file_name}") generated = True break except Exception as e: print(f"[HF] {prov} falló: {e}") else: try: print(f"[HF] versión antigua sin provider → '{prompt[:45]}...'") client = InferenceClient(token=token) image = await asyncio.to_thread(client.text_to_image, prompt, model=HF_MODEL) file_name = f"hf_{datetime.now().strftime('%Y%m%d_%H%M%S_%f')}.png" dataset_url = await _upload_to_dataset(image, file_name) if dataset_url: urls.append(dataset_url) else: DOCS_DIR.mkdir(parents=True, exist_ok=True) await asyncio.to_thread(image.save, str(DOCS_DIR / file_name)) urls.append(f"/docs/{file_name}") generated = True except Exception as e: print(f"[HF] excepción: {e}") if not generated: print(f"[HF] todos los providers fallaron para este prompt") return urls # ───────────────────────────────────────────────────────────────────────────── # 2. Gemini Imagen API (pendiente) # ───────────────────────────────────────────────────────────────────────────── async def _generate_with_gemini(queries: list) -> list: print("[Gemini] imagen API no habilitada aún, saltando...") return [] # ───────────────────────────────────────────────────────────────────────────── # Cascade: Gemini → HF → picsum # ───────────────────────────────────────────────────────────────────────────── async def _generate_images(queries: list, provider: str) -> list: image_urls = [] if provider == "gemini": try: image_urls = await _generate_with_gemini(queries) except Exception as e: print(f"[cascade] Gemini falló: {e}") if not image_urls: try: image_urls = await _generate_with_hf(queries) except Exception as e: print(f"[cascade] HF falló: {e}") if not image_urls: print("[cascade] usando picsum como fallback final") image_urls = [ f"https://picsum.photos/seed/{abs(hash(q)) % 9999}/400/300" for q in queries[:2] ] return image_urls # ───────────────────────────────────────────────────────────────────────────── # Entry point # ───────────────────────────────────────────────────────────────────────────── 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", "huggingface") print(f"[image_agent] provider={provider}") match = re.search(r'\{.*"image_queries".*\}', response, re.DOTALL) if match: try: data = json.loads(match.group(0)) queries = data.get("image_queries", []) if not queries: result["response"] = "El agente no generó queries de imagen." return result result["queries"] = queries result["response"] = "Generando imágenes para:\n" + "\n".join(f"• {q}" for q in queries) image_urls = await _generate_images(queries, provider) result["image_urls"] = image_urls result["response"] += "\n\nImagen generada ✔" if image_urls else "\n\n⚠️ No se pudo generar imagen" except json.JSONDecodeError as e: result["response"] = f"Error parseando JSON del agente: {e}" else: result["response"] = response except Exception as e: result["success"] = False result["error"] = str(e) print(f"[image_agent] excepción general: {e}") return result