app.py

import gradio as gr
import torch
from PIL import Image
import numpy as np
from clip_interrogator import Config, Interrogator
import logging
import os
import warnings
from datetime import datetime
import gc
import spaces

# Suprimir warnings
warnings.filterwarnings("ignore", category=FutureWarning)
warnings.filterwarnings("ignore", category=UserWarning)
os.environ["TOKENIZERS_PARALLELISM"] = "false"

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Detectar dispositivo
def get_device():
    if torch.cuda.is_available():
        return "cuda"
    elif torch.backends.mps.is_available():
        return "mps" 
    else:
        return "cpu"

DEVICE = get_device()

CLIP_MODELS = {
    "general": "ViT-L-14/openai",
    "stable_diffusion": "ViT-L-14/openai", 
    "midjourney": "ViT-H-14/laion2b_s32b_b79k",
    "flux": "ViT-L-14/openai"
}

MODES = {
    "fast": "⚡ Rápido",
    "classic": "⚖️ Clásico", 
    "best": "⭐ Mejor"
}

class ImagePromptGenerator:
    def __init__(self):
        self.interrogator = None
        self.usage_count = 0
        self.device = DEVICE
        self.is_initialized = False
    
    def initialize_model(self, progress_callback=None):
        if self.is_initialized:
            return True
            
        try:
            if progress_callback:
                progress_callback("🔄 Cargando CLIP Interrogator...")
            
            config = Config(
                clip_model_name="ViT-L-14/openai",
                download_cache=True,
                chunk_size=2048,
                quiet=True,
                device=self.device
            )
            
            self.interrogator = Interrogator(config)
            self.is_initialized = True
            
            if self.device == "cpu":
                gc.collect()
            else:
                torch.cuda.empty_cache()
                
            return True
            
        except Exception as e:
            logger.error(f"Error: {e}")
            return False
    
    def optimize_image(self, image):
        if image is None:
            return None
            
        if isinstance(image, np.ndarray):
            image = Image.fromarray(image)
        elif not isinstance(image, Image.Image):
            image = Image.open(image)
        
        if image.mode != 'RGB':
            image = image.convert('RGB')
        
        max_size = 768 if self.device != "cpu" else 512
        if image.size[0] > max_size or image.size[1] > max_size:
            image.thumbnail((max_size, max_size), Image.Resampling.LANCZOS)
        
        return image
    
    @spaces.GPU
    def generate_prompt(self, image, model_type="general", mode="best", progress_callback=None):
        try:
            if not self.is_initialized:
                if not self.initialize_model(progress_callback):
                    return "❌ Error inicializando modelo.", ""
            
            if image is None:
                return "❌ Sube una imagen.", ""
            
            self.usage_count += 1
            
            if progress_callback:
                progress_callback("🖼️ Procesando imagen...")
            
            image = self.optimize_image(image)
            if image is None:
                return "❌ Error procesando imagen.", ""
            
            if progress_callback:
                progress_callback("🧠 Generando prompt...")
            
            start_time = datetime.now()
            
            try:
                if mode == "fast":
                    prompt = self.interrogator.interrogate_fast(image)
                elif mode == "classic":
                    prompt = self.interrogator.interrogate_classic(image)
                else:
                    prompt = self.interrogator.interrogate(image)
                    
            except Exception as e:
                prompt = self.interrogator.interrogate_fast(image)
            
            end_time = datetime.now()
            duration = (end_time - start_time).total_seconds()
            
            if self.device == "cpu":
                gc.collect()
            else:
                torch.cuda.empty_cache()
            
            gpu_status = "🚀 ZeroGPU" if torch.cuda.is_available() else "🖥️ CPU"
            
            info = f"""
**✅ Prompt generado - Pariente AI**

{gpu_status} **|** {model_type.replace('_', ' ').title()} **|** {MODES.get(mode, mode)} **|** {duration:.1f}s
**Uso #{self.usage_count}** - {datetime.now().strftime('%H:%M')}

*"Porque cuando no tienes nada en la cabeza, te preocupas de la tipografía"*
"""
            
            return prompt, info
            
        except Exception as e:
            return f"❌ Error: {str(e)}", "💡 Intenta con modo rápido o imagen más pequeña"

generator = ImagePromptGenerator()

@spaces.GPU  
def process_image_with_progress(image, model_type, mode):
    def progress_callback(message):
        return message
    
    yield "🚀 Activando ZeroGPU...", """
**🚀 Pariente AI - Procesando**

⚡ ZeroGPU gratuito activado
🎯 Código abierto honesto
💡 Research real, no marketing
"""
    
    prompt, info = generator.generate_prompt(image, model_type, mode, progress_callback)
    yield prompt, info

def clear_outputs():
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
    return "", ""

def create_interface():
    css = """
    .gradio-container {
        max-width: 1400px !important;
        font-family: 'Inter', system-ui, sans-serif;
    }
    .prompt-output {
        font-family: 'SF Mono', 'Monaco', monospace !important;
        font-size: 14px !important;
        line-height: 1.6 !important;
        background: linear-gradient(135deg, #f8f9fa 0%, #e9ecef 100%) !important;
        border-radius: 12px !important;
        padding: 20px !important;
        border: 1px solid #dee2e6 !important;
    }
    .main-title {
        text-align: center;
        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
        -webkit-background-clip: text;
        -webkit-text-fill-color: transparent;
        font-size: 3em !important;
        font-weight: 800 !important;
        margin-bottom: 0.3em !important;
    }
    .subtitle {
        text-align: center;
        font-style: italic;
        color: #6c757d;
        font-size: 1.2em;
        margin-bottom: 2em;
    }
    """
    
    with gr.Blocks(theme=gr.themes.Soft(), title="Pariente AI - Image to Prompt", css=css) as interface:
        
        gr.HTML("""
        <div class="main-title">🤖 Pariente AI</div>
        <div class="subtitle">"Porque cuando no tienes nada en la cabeza, te preocupas de la tipografía"</div>
        """)
        
        gr.Markdown("### 🎨 Image to Prompt - Research real, no marketing")
        
        with gr.Row():
            with gr.Column(scale=1):
                gr.Markdown("## 📤 Imagen")
                image_input = gr.Image(
                    label="Sube tu imagen",
                    type="pil",
                    height=300
                )
                
                gr.Markdown("## ⚙️ Config")
                model_selector = gr.Dropdown(
                    choices=["general", "stable_diffusion", "midjourney", "flux"],
                    value="general",
                    label="Modelo objetivo"
                )
                
                mode_selector = gr.Dropdown(
                    choices=list(MODES.keys()),
                    value="best", 
                    label="Modo"
                )
                
                generate_btn = gr.Button("🚀 Generar Prompt", variant="primary", size="lg")
            
            with gr.Column(scale=1):
                gr.Markdown("## 📝 Resultado")
                prompt_output = gr.Textbox(
                    label="Prompt generado",
                    placeholder="Tu prompt aparecerá aquí...",
                    lines=10,
                    elem_classes=["prompt-output"],
                    show_copy_button=True
                )
                
                info_output = gr.Markdown(value="")
                
                with gr.Row():
                    clear_btn = gr.Button("🗑️ Limpiar", size="sm")
        
        gr.Markdown("""
---
### 💡 Realidad vs Marketing

**¿Ves esto? ZeroGPU gratuito. ¿Tu startup cobra $50/mes por lo mismo? Si lo pagas es señal de que eres gilipollas**

**🔬 Pariente AI hace research real:**
- Creamos modelos desde cero
- Publicamos papers de verdad  
- Mostramos el código siempre
- Innovamos, no copiamos

**🤡 Startup típica hace marketing:**
- Copia código de GitHub
- Lo envuelve en CSS bonito
- Cobra como "innovación"
- Busca inversores con PowerPoints

---

**⚡ Powered by Pariente AI** - *Research real, no bullshit*
        """)
        
        generate_btn.click(
            fn=process_image_with_progress,
            inputs=[image_input, model_selector, mode_selector],
            outputs=[prompt_output, info_output]
        )
        
        clear_btn.click(
            fn=clear_outputs,
            outputs=[prompt_output, info_output]
        )
    
    return interface

if __name__ == "__main__":
    logger.info("🚀 Iniciando Pariente AI")
    interface = create_interface()
    interface.launch(
        server_name="0.0.0.0",
        server_port=7860,
        show_error=True
    )