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

# Modelos
modeloObtenerTextoImagen = pipeline(
    "image-to-text", 
    model="Salesforce/blip-image-captioning-base"
)

modeloGenerarImagen = DiffusionPipeline.from_pretrained(
    "sd-legacy/stable-diffusion-v1-5", 
    torch_dtype=torch.float32
)

def obtenerDescripcion(imagen):
    """
    Función que recibe una imagen, genera una descripción de la misma
    y luego crea una nueva imagen basada en esa descripción.
    """
    # Obtener descripción de la imagen
    descripcion = modeloObtenerTextoImagen(Image.fromarray(imagen))
    texto_generado = descripcion[0]['generated_text']
    
    # Generar imagen a partir de la descripción
    imagen_generada = modeloGenerarImagen(texto_generado).images[0]
    
    return imagen_generada, texto_generado

# Layout con Gradio Blocks
with gr.Blocks(title="Asistente Visual para Personas con Discapacidad") as demo:
    
    gr.Markdown(
        """
        # Asistente Visual
        Esta herramienta ayuda a personas con discapacidad visual a comprender imágenes.
        1. Sube una imagen.
        2. La IA generará una descripción de la imagen y una nueva imagen interpretativa.
        """
    )
    
    with gr.Row():
        with gr.Column():
            imagen_input = gr.Image(label="Sube tu imagen aquí", type="numpy")
            boton_generar = gr.Button("Generar descripción e imagen")
        
        with gr.Column():
            imagen_output = gr.Image(label="Imagen generada")
            descripcion_output = gr.Textbox(label="Descripción generada por IA", lines=4)

    # Conectar botón a función
    boton_generar.click(
        fn=obtenerDescripcion, 
        inputs=imagen_input, 
        outputs=[imagen_output, descripcion_output]
    )

# Lanzar app
demo.launch(share=True)