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

# Initialize pipeline
dtype = torch.float16 if torch.cuda.is_available() else torch.float32
pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
    "nitrosocke/Ghibli-Diffusion",
    torch_dtype=dtype
).to("cuda" if torch.cuda.is_available() else "cpu")

def process_image(input_img):
    if input_img is None:
        return None
    
    # Convert and resize input
    input_img = input_img.convert("RGB").resize((512, 512))
    
    # Generate output
    result = pipe(
        prompt="ghibli style, studio ghibli, anime art",
        image=input_img,
        strength=0.7,
        guidance_scale=10
    ).images[0]
    
    return result

# Create Gradio interface
demo = gr.Interface(
    fn=process_image,
    inputs=gr.Image(type="pil"),
    outputs=gr.Image(type="pil"),
    title="🎨 Ghibli Style Transfer",
    description="Upload an image to transform it into Studio Ghibli style artwork"
)

demo.launch()