import gradio as gr
import torch
from diffusers import StableDiffusionPipeline

device = "cuda" if torch.cuda.is_available() else "cpu"
model_id = "nitrosocke/Ghibli-Diffusion"

# Load the model once and keep it in memory
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16 if device == "cuda" else torch.float32)
pipe.to(device)
pipe.enable_attention_slicing()  # Optimize memory usage

def generate_ghibli_style(image):
    prompt = "ghibli style portrait"
    with torch.inference_mode():  # Disables gradient calculations for faster inference
        result = pipe(prompt, image=image, strength=0.6, guidance_scale=6.5, num_inference_steps=25).images[0]  # Reduced steps & optimized scale
    return result

iface = gr.Interface(
    fn=generate_ghibli_style,
    inputs=gr.Image(type="pil"),
    outputs=gr.Image(),
    title="Studio Ghibli Portrait Generator",
    description="Upload a photo to generate a Ghibli-style portrait!"
)

iface.launch()