Spaces:
Running
Running
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() | |