Update app.py

app.py

@@ -1,14 +1,17 @@
 import gradio as gr
 import torch
-import io, base64
-from PIL import Image
 from diffusers import StableDiffusionImg2ImgPipeline
+from torchvision import transforms
+from PIL import Image
+import io
+import base64
 import spaces
 
 # Load Ghibli model
 pipe_ghibli = StableDiffusionImg2ImgPipeline.from_pretrained(
     "nitrosocke/Ghibli-Diffusion",
-    torch_dtype=torch.float16
+    torch_dtype=torch.float16,
+    use_safetensors=True
 ).to("cuda")
 
 # Load AnimeGANv2 model
@@ -18,6 +21,7 @@ animegan = torch.hub.load(
     pretrained="face_paint_512_v2"
 ).to("cuda").eval()
 
+# Base64 conversion helpers
 def pil_to_b64(img: Image.Image) -> str:
     buf = io.BytesIO()
     img.save(buf, format="PNG")
@@ -26,20 +30,21 @@ def pil_to_b64(img: Image.Image) -> str:
 def b64_to_pil(b64: str) -> Image.Image:
     return Image.open(io.BytesIO(base64.b64decode(b64))).convert("RGB")
 
+# AnimeGAN processor
 def apply_animegan(img: Image.Image) -> Image.Image:
-    img = img.resize((512,512))
-    tensor = torch.from_numpy(
-        (torch.FloatTensor(torch.ByteTensor(torch.ByteStorage.from_buffer(img.tobytes())))
-         .float()/255.)
-    ).permute(2,0,1).unsqueeze(0).to("cuda")
+    transform = transforms.Compose([
+        transforms.Resize((512, 512)),
+        transforms.ToTensor()
+    ])
+    img_tensor = transform(img).unsqueeze(0).to("cuda")
     with torch.no_grad():
-        out = animegan(tensor)[0].clamp(0,1).cpu()
-    return Image.fromarray((out.permute(1,2,0).numpy()*255).astype("uint8"))
+        out = animegan(img_tensor)[0].clamp(0, 1).cpu()
+    return transforms.ToPILImage()(out)
 
+# Unified image processor
 def process_image(img: Image.Image, effect: str) -> Image.Image:
     if effect == "ghibli":
-        return pipe_ghibli(prompt="ghibli style", image=img,
-                            strength=0.5, guidance_scale=7.5).images[0]
+        return pipe_ghibli(prompt="ghibli style", image=img, strength=0.5, guidance_scale=7.5).images[0]
     else:
         return apply_animegan(img)
 
@@ -49,21 +54,22 @@ def process_base64(b64: str, effect: str) -> str:
     out = process_image(img, effect)
     return pil_to_b64(out)
 
+# Gradio UI
 with gr.Blocks() as demo:
-    gr.Markdown("# 🎨 Ghibli & AnimeGANv2 (Cartoon) Effects")
+    gr.Markdown("# 🎨 Ghibli & AnimeGAN Effects (ZeroGPU Compatible)")
 
     with gr.Tab("Web UI"):
-        inp = gr.Image(type="pil")
-        eff = gr.Radio(["ghibli", "anime"], label="Effect")
+        inp = gr.Image(type="pil", label="Upload Image")
+        eff = gr.Radio(["ghibli", "anime"], label="Select Effect")
         btn = gr.Button("Apply Effect")
-        out_img = gr.Image()
+        out_img = gr.Image(label="Output Image")
         btn.click(process_image, [inp, eff], out_img)
 
     with gr.Tab("Base64 API"):
-        b64_in = gr.Textbox(lines=5)
-        eff2 = gr.Radio(["ghibli", "anime"], label="Effect")
+        b64_in = gr.Textbox(label="Input Image (Base64)", lines=5)
+        eff2 = gr.Radio(["ghibli", "anime"], label="Select Effect")
         btn2 = gr.Button("Run API")
-        b64_out = gr.Textbox(lines=5)
+        b64_out = gr.Textbox(label="Output Image (Base64)", lines=5)
         btn2.click(process_base64, [b64_in, eff2], b64_out)
 
 demo.launch()