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()