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

# Load the model
model_id = "lavaman131/cartoonify"
pipe = StableDiffusionImg2ImgPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
pipe.to("cpu")

# Define the function that processes the image
def cartoonify(image):
    image = image.convert("RGB").resize((512, 512))
    prompt = "A cartoon version of this image, pixar disney style"
    output = pipe(prompt=prompt, image=image, strength=0.5, guidance_scale=6).images[0]
    return output

# Set up Gradio interface
interface = gr.Interface(
    fn=cartoonify,
    inputs=gr.Image(type="pil", label="Upload your photo"),
    outputs=gr.Image(label="Cartoonified image"),
    title="Cartoonify Your Image!",
    description="Upload a photo and get a cartoon version in a Pixar/Disney style."
)

# Launch the app
interface.launch()