prediction / app.py
redondoandre's picture
Update app.py
2ffd03c verified
import gradio as gr
from transformers import pipeline
from diffusers import StableDiffusionPipeline
import torch
# Cargar modelo para clasificar im谩genes de perros y gatos
classifier = pipeline("image-classification", model="microsoft/resnet-50")
# Cargar modelo para generar im谩genes desde texto
pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5")
pipe.to("cpu")
# Funci贸n para procesar imagen o texto
def process_input(image, description):
if image is not None: # Si el usuario sube una imagen
results = classifier(image)
return results[0]["label"], None # Devuelve la raza detectada
elif description: # Si el usuario escribe un texto
image = pipe(description).images[0]
return None, image # Devuelve la imagen generada
return "Por favor, sube una imagen o escribe una descripci贸n", None
# Interfaz de usuario con Gradio
iface = gr.Interface(
fn=process_input,
inputs=[gr.Image(type="pil"), gr.Textbox(label="Describe un gato o perro")],
outputs=[gr.Textbox(label="Predicci贸n"), gr.Image(label="Imagen generada")],
title="Detector y Generador de Razas de Gatos y Perros",
description="Sube una imagen para conocer la raza o escribe una descripci贸n para generar una imagen."
)
iface.launch()