import gradio as gr
from PIL import Image
import torch
from diffusers import StableDiffusionPipeline

# Cargar el modelo Stable Diffusion
device = "cuda" if torch.cuda.is_available() else "cpu"
model_id = "CompVis/stable-diffusion-v1-4"  # Usa el modelo adecuado según la versión disponible
pipe = StableDiffusionPipeline.from_pretrained(model_id).to(device)

def cartoonize_image(image):
    prompt = "A Pixar-style cartoon character"
    
    # Preprocesar imagen
    image = image.convert("RGB")
    image = image.resize((512, 512))  # Ajusta el tamaño según sea necesario

    # Generar imagen con Stable Diffusion
    with torch.no_grad():
        result = pipe(prompt=prompt, init_image=image, strength=0.75).images[0]
    
    return result

# Configuración de la interfaz de Gradio
def main():
    with gr.Blocks() as demo:
        gr.Markdown("## Conversión de Imagen a Caricatura Estilo Pixar")
        with gr.Row():
            with gr.Column():
                input_image = gr.Image(type="pil", label="Sube tu imagen")
                output_image = gr.Image(type="pil", label="Imagen Caricaturizada")
        
        input_image.change(
            fn=cartoonize_image,
            inputs=input_image,
            outputs=output_image
        )

    demo.launch()

if __name__ == "__main__":
    main()