image_tool_reve.py

"""
SISTEMA DE IMÁGENES REVE MEJORADO
- Generación por lotes
- Sello automático
- Optimización de parámetros
"""
import os
import time
import base64
import json
import asyncio
from typing import List, Tuple, Optional, Dict, Any
from io import BytesIO
from pathlib import Path
from datetime import datetime

import aiohttp
import requests
from PIL import Image, ImageDraw, ImageFont
import concurrent.futures

# ===================== CONFIGURACIÓN =====================
REVE_API_KEY = os.getenv("REVE_API_KEY")
REVE_URL = "https://api.reve.com/v1/image/create"

# SOLUCIÓN: Verificación no bloqueante (Línea 29)
# En lugar de lanzar un error que detiene la app, definimos un estado.
REVE_AVAILABLE = REVE_API_KEY is not None and REVE_API_KEY != ""

# ===================== SELLO BATUTO AUTOMATIZADO =====================
class BatutoWatermark:
    def __init__(self):
        self.watermark_text = "BATUTO-ART"
        self.default_font = None
        
    def _get_font(self, image_width: int):
        """Obtiene fuente apropiada para el tamaño de imagen"""
        font_size = max(24, int(image_width * 0.035))
        
        try:
            # Intentar cargar fuentes comunes
            font_paths = [
                "arial.ttf",
                "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf",
                "/System/Library/Fonts/Helvetica.ttc"
            ]
            
            for path in font_paths:
                if os.path.exists(path):
                    return ImageFont.truetype(path, font_size)
        except:
            pass
        
        # Fallback a fuente por defecto
        return ImageFont.load_default()
    
    def apply_watermark(self, image: Image.Image) -> Image.Image:
        """Aplica sello BATUTO a la imagen"""
        try:
            img = image.convert("RGBA")
            draw = ImageDraw.Draw(img)
            
            font = self._get_font(img.width)
            
            # Colores del sello (dorado con sombra)
            color_primary = (212, 175, 55, 255)
            color_shadow = (0, 0, 0, 180)
            
            # Posición (esquina inferior derecha)
            margin_x = int(img.width * 0.02)
            margin_y = int(img.height * 0.02)
            shadow_offset = 2
            
            # Calcular tamaño del texto
            text_bbox = draw.textbbox((0, 0), self.watermark_text, font=font)
            text_width = text_bbox[2] - text_bbox[0]
            text_height = text_bbox[3] - text_bbox[1]
            
            # Posicionar en esquina inferior derecha
            x = img.width - text_width - margin_x
            y = img.height - text_height - margin_y
            
            # Dibujar sombra
            draw.text(
                (x + shadow_offset, y + shadow_offset),
                self.watermark_text,
                font=font,
                fill=color_shadow
            )
            
            # Dibujar texto principal
            draw.text(
                (x, y),
                self.watermark_text,
                font=font,
                fill=color_primary
            )
            
            return img.convert("RGB")
            
        except Exception as e:
            print(f"⚠️ Error aplicando sello: {e}")
            return image

# ===================== GENERADOR DE IMÁGENES ASÍNCRONO =====================
class AsyncReveGenerator:
    def __init__(self, api_key: str, max_workers: int = 4):
        if not api_key:
            raise ValueError("API Key no proporcionada para AsyncReveGenerator")
        self.api_key = api_key
        self.max_workers = max_workers
        self.watermarker = BatutoWatermark()
        self.session = None
        
    async def __aenter__(self):
        self.session = aiohttp.ClientSession(
            headers={
                "Authorization": f"Bearer {self.api_key}",
                "Content-Type": "application/json"
            },
            timeout=aiohttp.ClientTimeout(total=120)
        )
        return self
    
    async def __aexit__(self, exc_type, exc_val, exc_tb):
        if self.session:
            await self.session.close()
    
    async def generate_single(
        self,
        prompt: str,
        ratio: str = "1:1",
        version: str = "latest",
        index: int = 0
    ) -> Optional[Tuple[Image.Image, Dict]]:
        """Genera una sola imagen asíncronamente"""
        try:
            payload = {
                "prompt": prompt,
                "aspect_ratio": ratio,
                "version": version
            }
            
            async with self.session.post(REVE_URL, json=payload) as response:
                if response.status == 200:
                    data = await response.json()
                    
                    if "image" in data:
                        # Decodificar imagen base64
                        img_data = base64.b64decode(data["image"])
                        img = Image.open(BytesIO(img_data)).convert("RGB")
                        
                        # Aplicar sello
                        img = self.watermarker.apply_watermark(img)
                        
                        metadata = {
                            "index": index,
                            "prompt": prompt,
                            "ratio": ratio,
                            "version": version,
                            "timestamp": datetime.now().isoformat(),
                            "api_response": data.get("id", "unknown")
                        }
                        
                        return img, metadata
                    else:
                        print(f"❌ No image in API response for prompt {index}")
                        return None
                else:
                    error_text = await response.text()
                    print(f"❌ API Error {response.status}: {error_text}")
                    return None
                    
        except Exception as e:
            print(f"❌ Error generando imagen {index}: {e}")
            return None
    
    async def generate_batch(
        self,
        prompts: List[str],
        ratio: str = "1:1",
        version: str = "latest",
        max_concurrent: int = None
    ) -> List[Tuple[Image.Image, Dict]]:
        """Genera múltiples imágenes concurrentemente"""
        if not self.session:
            raise RuntimeError("Session not initialized. Use async context manager.")
        
        concurrent_limit = max_concurrent or self.max_workers
        semaphore = asyncio.Semaphore(concurrent_limit)
        
        async def limited_generate(prompt, idx):
            async with semaphore:
                return await self.generate_single(prompt, ratio, version, idx)
        
        tasks = [limited_generate(prompt, i) for i, prompt in enumerate(prompts)]
        results = await asyncio.gather(*tasks, return_exceptions=True)
        
        # Filtrar resultados exitosos
        successful = []
        for result in results:
            if isinstance(result, tuple) and result[0] is not None:
                successful.append(result)
            elif isinstance(result, Exception):
                print(f"⚠️ Exception in generation: {result}")
        
        return successful

# ===================== FUNCIONES PÚBLICAS =====================
def generar_imagen_reve(
    prompt: str,
    ratio: str = "1:1",
    version: str = "latest",
    apply_watermark: bool = True,
    save_path: Optional[str] = None
) -> Image.Image:
    """
    Genera una imagen usando REVE API (versión síncrona)
    """
    # VERIFICACIÓN SEGURA: La app no se detiene aquí
    if not REVE_AVAILABLE:
        raise RuntimeError("""
        🔧 Configuración de REVE API Requerida
        
        Para usar la generación de imágenes, configura tu clave:
        1. Ve a 'Settings' en tu Space de Hugging Face.
        2. Haz clic en 'Repository secrets'.
        3. Agrega: REVE_API_KEY = tu_clave_secreta
        4. Guarda y reinicia el Space.
        
        La aplicación seguirá funcionando para otras funcionalidades.
        """)
    
    payload = {
        "prompt": prompt,
        "aspect_ratio": ratio,
        "version": version
    }
    
    headers = {
        "Authorization": f"Bearer {REVE_API_KEY}",
        "Content-Type": "application/json"
    }
    
    try:
        response = requests.post(
            REVE_URL,
            headers=headers,
            json=payload,
            timeout=120
        )
        
        if response.status_code != 200:
            raise RuntimeError(f"❌ Error REVE API: {response.status_code} - {response.text}")
        
        data = response.json()
        
        if "image" not in data:
            raise RuntimeError("❌ No se encontró imagen en la respuesta de REVE")
        
        # Decodificar y cargar imagen
        img_bytes = base64.b64decode(data["image"])
        img = Image.open(BytesIO(img_bytes)).convert("RGB")
        
        # Aplicar sello si está habilitado
        if apply_watermark:
            watermarker = BatutoWatermark()
            img = watermarker.apply_watermark(img)
        
        # Guardar si se especifica ruta
        if save_path:
            os.makedirs(os.path.dirname(save_path), exist_ok=True)
            img.save(save_path)
            print(f"✅ Imagen guardada en: {save_path}")
        
        return img
        
    except requests.exceptions.Timeout:
        raise RuntimeError("⏰ Timeout de REVE API (120s)")
    except Exception as e:
        raise RuntimeError(f"❌ Error generando imagen: {str(e)}")

async def generar_imagenes_concurrentes(
    prompts: List[str],
    ratio: str = "1:1",
    version: str = "latest",
    max_concurrent: int = 4
) -> List[Image.Image]:
    """
    Genera múltiples imágenes concurrentemente
    """
    if not REVE_AVAILABLE:
        raise RuntimeError("REVE_API_KEY no configurada. Verifica los Secretos del Space.")
    
    async with AsyncReveGenerator(REVE_API_KEY, max_concurrent) as generator:
        results = await generator.generate_batch(prompts, ratio, version, max_concurrent)
        return [img for img, _ in results]

def generar_imagen_con_metadatos(
    prompt: str,
    ratio: str = "1:1",
    version: str = "latest"
) -> Dict[str, Any]:
    """
    Genera imagen con metadatos completos
    """
    start_time = time.time()
    
    # Verificación consistente
    if not REVE_AVAILABLE:
        return {
            "image": None,
            "metadata": {
                "success": False,
                "error": "REVE_API_KEY no configurada en los Secretos del Space.",
                "generation_time": 0.0,
                "timestamp": datetime.now().isoformat()
            }
        }
    
    try:
        img = generar_imagen_reve(prompt, ratio, version, apply_watermark=True)
        
        # Crear metadatos
        metadata = {
            "success": True,
            "prompt": prompt,
            "aspect_ratio": ratio,
            "version": version,
            "generation_time": time.time() - start_time,
            "image_size": img.size,
            "image_mode": img.mode,
            "timestamp": datetime.now().isoformat(),
            "watermark_applied": True
        }
        
        return {
            "image": img,
            "metadata": metadata
        }
        
    except Exception as e:
        return {
            "image": None,
            "metadata": {
                "success": False,
                "error": str(e),
                "generation_time": time.time() - start_time,
                "timestamp": datetime.now().isoformat()
            }
        }