import gradio as gr
import torch
from PIL import Image
from transformers import pipeline
from diffusers import DiffusionPipeline

# Carga de modelos
modelo_caption = pipeline("image-to-text", model="Salesforce/blip-image-captioning-base")
modelo_clasificador = pipeline("text-classification", model="TheBritishLibrary/bl-books-genre")
modelo_difusion = DiffusionPipeline.from_pretrained(
    "stabilityai/stable-diffusion-2-1", torch_dtype=torch.float32
)
modelo_difusion.to("cuda" if torch.cuda.is_available() else "cpu")

# Función principal
def procesar_imagen(imagen):
    descripcion = modelo_caption(Image.fromarray(imagen))[0]['generated_text']
    genero = modelo_clasificador(descripcion)[0]
    imagen_generada = modelo_difusion(descripcion).images[0]
    
    resultado_genero = f"Género literario: {genero['label']} (confianza: {genero['score']:.2f})"
    
    return descripcion, resultado_genero, imagen_generada

# Interfaz mejorada con Gradio Blocks
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🖼️ IA Generadora de Textos e Imágenes")
    gr.Markdown("Sube una imagen, genera una descripción automática, clasifica su género literario y crea una nueva imagen con IA.")

    with gr.Row():
        with gr.Column():
            imagen_input = gr.Image(type="numpy", label="Imagen de entrada")
            boton = gr.Button("Procesar con IA")

        with gr.Column():
            descripcion_output = gr.Textbox(label="Descripción generada", interactive=False)
            genero_output = gr.Textbox(label="Género literario", interactive=False)
            imagen_output = gr.Image(label="Imagen generada")

    boton.click(fn=procesar_imagen, inputs=imagen_input, outputs=[descripcion_output, genero_output, imagen_output])

demo.launch()