app.py

import gradio as gr
import torch
from diffusers import StableDiffusionImg2ImgPipeline
from torchvision import transforms
from PIL import Image
import io, base64, spaces

# Ghibli model
pipe_ghibli = StableDiffusionImg2ImgPipeline.from_pretrained(
    "nitrosocke/Ghibli-Diffusion", torch_dtype=torch.float16
).to("cuda")

# CartoonGAN model via torch.hub
cartoon_model = torch.hub.load(
    'znxlwm/pytorch-CartoonGAN', 'CartoonGAN',
    pretrained=True, trust_repo=True
).to("cuda").eval()

# Helpers: PIL ↔ Base64
def pil_to_b64(img): buf=io.BytesIO(); img.save(buf,"PNG"); return base64.b64encode(buf.getvalue()).decode()
def b64_to_pil(b): return Image.open(io.BytesIO(base64.b64decode(b))).convert("RGB")

# CartoonGAN processor
def apply_cartoon(img):
    t=transforms.Compose([transforms.Resize((256,256)), transforms.ToTensor()])
    x=t(img).unsqueeze(0).to("cuda")
    with torch.no_grad(): y=cartoon_model(x)[0].clamp(0,1).cpu()
    return transforms.ToPILImage()(y)

# Unified image processor
def process_image(img, effect):
    if effect=="ghibli":
        return pipe_ghibli(prompt="ghibli style", image=img, strength=0.5, guidance_scale=7.5).images[0]
    return apply_cartoon(img)

@spaces.GPU
def process_base64(b64, effect):
    img = b64_to_pil(b64)
    out = process_image(img, effect)
    return pil_to_b64(out)

# Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("# 🎨 Ghibli & CartoonGAN Effects (ZeroGPU)")
    with gr.Tab("Web UI"):
        inp = gr.Image(type="pil")
        eff = gr.Radio(["ghibli","cartoon"], label="Effect")
        btn = gr.Button("Apply")
        out_img = gr.Image()
        btn.click(process_image, [inp, eff], out_img)
    with gr.Tab("Base64 API"):
        in_b64 = gr.Textbox(lines=5)
        eff2 = gr.Radio(["ghibli","cartoon"], label="Effect")
        btn2 = gr.Button("Run API")
        out_b64 = gr.Textbox(lines=5)
        btn2.click(process_base64, [in_b64, eff2], out_b64)

demo.launch()